Skip to content

vllm.models.beam.amd.dflash2_conv_ops

Functions:

  • conv_residual_norm –

    Finish grouped convolution and update the residual before RMSNorm.

conv_residual_norm(hidden, coefficients, base, residual, weight, eps, query_len, group_size, *, num_warps=None)

Finish grouped convolution and update the residual before RMSNorm.

Source code in vllm/models/beam/amd/dflash2_conv_ops.py 
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
def conv_residual_norm(
    hidden: torch.Tensor,
    coefficients: torch.Tensor,
    base: torch.Tensor,
    residual: torch.Tensor,
    weight: torch.Tensor,
    eps: float,
    query_len: int,
    group_size: int,
    *,
    num_warps: int | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Finish grouped convolution and update the residual before RMSNorm."""
    rows, width = hidden.shape
    block = triton.next_power_of_2(width)
    if num_warps is None:
        num_warps = 16 if rows <= 256 else (8 if rows < 4096 else 4)
    # Extra 64-lane wavefronts replicate the block, and a replica can reload
    # RESIDUAL after another wavefront has updated it in place.
    # https://github.com/triton-lang/triton/issues/12616
    num_warps = min(num_warps, max(1, block // 64))
    output = torch.empty_like(hidden)
    if rows:
        _conv_residual_norm_kernel[(rows,)](
            hidden,
            coefficients,
            base,
            residual,
            weight,
            output,
            hidden.stride(0),
            coefficients.stride(0),
            coefficients.stride(1),
            residual.stride(0),
            width,
            group_size,
            base.shape[0],
            query_len,
            eps,
            block,
            num_warps=num_warps,
            enable_fp_fusion=False,
        )
    return output, residual