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 
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
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(
        get_input_quant_key(linear) == 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 
 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
124
125
126
127
128
129
130
131
132
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
    scales = torch.empty(
        mxfp8_scale_bytes(tokens, q_size), dtype=torch.uint8, device=qr.device
    )
    if tokens:
        block = triton.next_power_of_2(max(q_size, kv_size))
        launch_pdl = current_platform.is_arch_support_pdl()
        _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,
            launch_pdl,
            num_warps=8 if block >= 2048 else 4,
            launch_pdl=launch_pdl,
        )
    return QuantizedActivation(qo, scales, qr.dtype, qr.shape, kMxfp8Dynamic), kvo

mxfp8_scale_bytes(tokens, width)

Bytes of F8_128x4 scales for tokens rows of width MXFP8 values.

Source code in vllm/models/deepseek_v41/common/ops/query_quant.py 
31
32
33
def mxfp8_scale_bytes(tokens: int, width: int) -> int:
    """Bytes of F8_128x4 scales for `tokens` rows of `width` MXFP8 values."""
    return triton.cdiv(tokens, 128) * 128 * triton.cdiv(width // 32, 4) * 4

mxfp8_scale_offset(row, group, GROUPS)

F8_128x4: [row/128, group/4, row%32, row%128/32, group%4].

Source code in vllm/models/deepseek_v41/common/ops/query_quant.py 
24
25
26
27
28
@triton.jit
def mxfp8_scale_offset(row, group, GROUPS: tl.constexpr):
    """F8_128x4: [row/128, group/4, row%32, row%128/32, group%4]."""
    offset = (row // 128 * triton.cdiv(GROUPS, 4) + group // 4) * 512
    return offset + row % 32 * 16 + row % 128 // 32 * 4 + group % 4