Skip to content

vllm.models.deepseek_v41.common.ops.query_quant

Functions:

can_fuse_query_quant(linears)

Require both local projections to use the same MXFP8 activation ABI.

Source code in vllm/models/deepseek_v41/common/ops/query_quant.py 
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
def can_fuse_query_quant(linears: list[torch.nn.Module]) -> bool:
    """Require both local projections to use the same MXFP8 activation ABI."""
    if not current_platform.is_cuda():
        return False
    from vllm.model_executor.kernels.linear.mxfp8.flashinfer import (
        FlashInferCutedslMxfp8LinearKernel,
        FlashInferCutlassMxfp8LinearKernel,
    )

    # QuantKey does not encode scale layout; restrict this producer to the
    # consumers that accept F8_128x4 swizzled scales.
    return all(
        getattr(linear, "input_quant_key", None) == kMxfp8Dynamic
        and type(getattr(getattr(linear, "quant_method", None), "kernel", None))
        in (FlashInferCutedslMxfp8LinearKernel, FlashInferCutlassMxfp8LinearKernel)
        for linear in linears
    )

fused_q_kv_rmsnorm_quant(qr, kv, q_weight, kv_weight, eps)

Normalize Q/KV and quantize Q with FlashInfer's swizzled MXFP8 scales.

Source code in vllm/models/deepseek_v41/common/ops/query_quant.py 
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
def fused_q_kv_rmsnorm_quant(
    qr: torch.Tensor,
    kv: torch.Tensor,
    q_weight: torch.Tensor,
    kv_weight: torch.Tensor,
    eps: float,
) -> tuple[QuantizedActivation, torch.Tensor]:
    """Normalize Q/KV and quantize Q with FlashInfer's swizzled MXFP8 scales."""
    assert qr.ndim == kv.ndim == 2 and qr.shape[0] == kv.shape[0]
    assert qr.stride(-1) == kv.stride(-1) == 1
    assert q_weight.is_contiguous() and kv_weight.is_contiguous()
    assert qr.shape[1] % 32 == 0
    tokens, q_size = qr.shape
    kv_size = kv.shape[1]
    qo = torch.empty(qr.shape, dtype=torch.float8_e4m3fn, device=qr.device)
    kvo = torch.empty(kv.shape, dtype=kv.dtype, device=kv.device)
    padded_tokens = triton.cdiv(tokens, 128) * 128
    padded_groups = triton.cdiv(q_size // 32, 4) * 4
    scales = torch.empty(
        padded_tokens * padded_groups, dtype=torch.uint8, device=qr.device
    )
    if tokens:
        block = triton.next_power_of_2(max(q_size, kv_size))
        _q_kv_norm_quant_kernel[(padded_tokens, 2)](
            qr,
            kv,
            q_weight,
            kv_weight,
            qo,
            kvo,
            scales,
            tokens,
            qr.stride(0),
            kv.stride(0),
            eps,
            q_size,
            kv_size,
            block,
            current_platform.is_arch_support_pdl(),
            num_warps=8 if block >= 2048 else 4,
        )
    return QuantizedActivation(qo, scales, qr.dtype, qr.shape, kMxfp8Dynamic), kvo