Skip to content

vllm.models.beam.nvidia.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/nvidia/dflash2_conv_ops.py 
56
57
58
59
60
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
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
    if num_warps is None:
        num_warps = 16 if rows <= 256 else (8 if rows < 4096 else 4)
    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,
            triton.next_power_of_2(width),
            num_warps=num_warps,
            enable_fp_fusion=False,
        )
    return output, residual