vllm.models.deepseek_v4.common.ops.fused_indexer_q

Functions:

• fused_indexer_q_rope_quant –

  Fused RoPE + quantize Q for the sparse indexer.

_indexer_weights_out_dtypes(vllm_config)

Weights dtypes the model's indexer layers ask for: fp32 for the dense scoring kernels, plus bf16 when the DeepSeek V4.1 sparse-logits indexer (SparseMQAIndexer) is enabled.

Source code in vllm/models/deepseek_v4/common/ops/fused_indexer_q.py

def _indexer_weights_out_dtypes(vllm_config: Any) -> tuple[torch.dtype, ...]:
    """Weights dtypes the model's indexer layers ask for: fp32 for the dense
    scoring kernels, plus bf16 when the DeepSeek V4.1 sparse-logits indexer
    (`SparseMQAIndexer`) is enabled."""
    if vllm_config.attention_config.indexer_sparse_logits:
        return (torch.float32, torch.bfloat16)
    return (torch.float32,)

_quantize_mxfp4_pair(x_lo, x_hi)

Quantize a block of MXFP4_BLOCK_SIZE fp32 values given as two interleaved halves (x_lo = values at even positions in the block, x_hi = values at odd positions). Returns: - packed : uint8[BLOCK/2] (low nibble = quant(x_lo), high = quant(x_hi)) - ue8m0 : scalar uint8 (block scale = 2^(ue8m0 - 127))

Source code in vllm/models/deepseek_v4/common/ops/fused_indexer_q.py

@triton.jit
def _quantize_mxfp4_pair(x_lo, x_hi):
    """Quantize a block of MXFP4_BLOCK_SIZE fp32 values given as two
    interleaved halves (x_lo = values at even positions in the block,
    x_hi = values at odd positions). Returns:
        - packed : uint8[BLOCK/2]  (low nibble = quant(x_lo), high = quant(x_hi))
        - ue8m0  : scalar uint8    (block scale = 2^(ue8m0 - 127))
    """
    amax = tl.maximum(tl.max(tl.abs(x_lo)), tl.max(tl.abs(x_hi)))
    # 6 * 2^-126 is from https://huggingface.co/deepseek-ai/DeepSeek-V4-Pro/blob/main/inference/kernel.py#L163
    amax = tl.maximum(amax, 6.0 * (2**-126))
    # ue8m0 block scale: 2^ceil(log2(amax/6.0)).
    log2_ratio = tl.math.ceil(tl.math.log2(amax * (1.0 / 6.0)))
    log2_ratio = tl.minimum(tl.maximum(log2_ratio, -127.0), 127.0)
    scale = tl.math.exp2(log2_ratio)
    ue8m0 = (log2_ratio + 127.0).to(tl.uint8)

    inv_scale = 1.0 / scale
    packed = _fp32x2_to_fp4x2(x_lo * inv_scale, x_hi * inv_scale)
    return packed, ue8m0

fused_indexer_q_rope_quant(positions, index_q, index_q_cos_sin_cache, index_weights, index_weights_softmax_scale, index_weights_head_scale, use_fp4=False, weights_out_dtype=torch.float32)

Fused RoPE + quantize Q for the sparse indexer.

weights_out_dtype is the dtype the downstream scoring kernel takes: fp32 for the dense MQA-logits kernels, bf16 for DeepGEMM's sparse MQA-logits kernels (CUDA MXFP4 path only).

Weight-fold semantics (important — the two paths differ):

FP8 path (use_fp4=False, default): q_fp8 : (T, H, HEAD_DIM) platform fp8 (e4m3fnuz on gfx942, e4m3fn elsewhere); per-token-per-head scalar scale (NOT stored — folded into weights below) weights_out = weights * q_scale * softmax_scale * head_scale Rationale: a single per-token q_scale is a scalar the downstream FP8 logits kernel would otherwise multiply in. Folding it into weights avoids emitting a separate tensor and is free for the logits kernel.

MXFP4 path (use_fp4=True): q_packed : (T, H, HEAD_DIM // 2) uint8 (2 E2M1 nibbles per byte) q_scale : (T, H, HEAD_DIM // MXFP4_BLOCK_SIZE) uint8 ue8m0 bytes weights_out = weights * softmax_scale * head_scale Rationale: MXFP4 has PER-BLOCK (32-element) scales that live with the Q values — they cannot be folded into a per-token weight scalar, so weights carries only the softmax and head scales.

Returns (q_quant, weights_out) where q_quant is either a Tensor (FP8) or a (values, scales) tuple (MXFP4). This matches the union type accepted by SparseAttnIndexer.forward_*.

Source code in vllm/models/deepseek_v4/common/ops/fused_indexer_q.py

def fused_indexer_q_rope_quant(
    positions: torch.Tensor,
    index_q: torch.Tensor,
    index_q_cos_sin_cache: torch.Tensor,
    # Index weights
    index_weights: torch.Tensor,
    index_weights_softmax_scale: float,
    index_weights_head_scale: float,
    use_fp4: bool = False,
    weights_out_dtype: torch.dtype = torch.float32,
) -> tuple[
    torch.Tensor | tuple[torch.Tensor, torch.Tensor],
    torch.Tensor,
]:
    """Fused RoPE + quantize Q for the sparse indexer.

    ``weights_out_dtype`` is the dtype the downstream scoring kernel takes:
    fp32 for the dense MQA-logits kernels, bf16 for DeepGEMM's sparse
    MQA-logits kernels (CUDA MXFP4 path only).

    Weight-fold semantics (important — the two paths differ):

    FP8 path (use_fp4=False, default):
        q_fp8      : (T, H, HEAD_DIM) platform fp8 (e4m3fnuz on gfx942,
                     e4m3fn elsewhere); per-token-per-head scalar scale
                     (NOT stored — folded into weights below)
        weights_out = weights * q_scale * softmax_scale * head_scale
        Rationale: a single per-token q_scale is a scalar the downstream FP8
        logits kernel would otherwise multiply in. Folding it into `weights`
        avoids emitting a separate tensor and is free for the logits kernel.

    MXFP4 path (use_fp4=True):
        q_packed   : (T, H, HEAD_DIM // 2) uint8 (2 E2M1 nibbles per byte)
        q_scale    : (T, H, HEAD_DIM // MXFP4_BLOCK_SIZE) uint8 ue8m0 bytes
        weights_out = weights * softmax_scale * head_scale
        Rationale: MXFP4 has PER-BLOCK (32-element) scales that live with
        the Q values — they cannot be folded into a per-token weight
        scalar, so `weights` carries only the softmax and head scales.

    Returns (q_quant, weights_out) where q_quant is either a Tensor (FP8) or
    a (values, scales) tuple (MXFP4). This matches the union type accepted
    by `SparseAttnIndexer.forward_*`.
    """
    assert positions.ndim == 1
    assert index_q.ndim == 3
    assert index_q_cos_sin_cache.ndim == 2

    num_tokens = positions.shape[0]
    num_index_q_heads = index_q.shape[1]
    index_q_head_dim = index_q.shape[2]

    index_weights_out = torch.empty_like(index_weights, dtype=weights_out_dtype)

    if use_fp4:
        assert index_q_head_dim % MXFP4_BLOCK_SIZE == 0, (
            f"head_dim={index_q_head_dim} must be a multiple of MXFP4 block "
            f"size {MXFP4_BLOCK_SIZE}"
        )
        num_scale_blocks = index_q_head_dim // MXFP4_BLOCK_SIZE
        index_q_packed = torch.empty(
            (num_tokens, num_index_q_heads, index_q_head_dim // 2),
            dtype=torch.uint8,
            device=index_q.device,
        )
        index_q_scale = torch.empty(
            (num_tokens, num_index_q_heads, num_scale_blocks),
            dtype=torch.uint8,
            device=index_q.device,
        )
        if has_cutedsl():
            # lazily import, otherwise some tests fail due to CUDA driver init failure.
            from vllm.models.deepseek_v4.nvidia.ops.fused_indexer_q_cutedsl import (
                _INDEXER_Q_MXFP4_KERNEL,
            )

            _INDEXER_Q_MXFP4_KERNEL(
                positions=positions,
                q=index_q,
                cos_sin_cache=index_q_cos_sin_cache,
                weights=index_weights,
                weights_softmax_scale=index_weights_softmax_scale,
                weights_head_scale=index_weights_head_scale,
                q_packed=index_q_packed,
                q_scale=index_q_scale,
                weights_out=index_weights_out,
            )
        elif current_platform.is_xpu():
            assert weights_out_dtype == torch.float32, weights_out_dtype
            torch.ops.vllm.xpu_deepseek_fused_indexer_q_rope_mxfp4(
                index_q,
                positions,
                index_q_cos_sin_cache,
                index_weights,
                index_weights_softmax_scale,
                index_weights_head_scale,
                index_q_packed,
                index_q_scale,
                index_weights_out,
            )
        else:
            _FUSED_INDEXER_Q_ROPE_MXFP4_TRITON_KERNEL(
                positions,
                index_q,
                index_q_cos_sin_cache,
                index_weights,
                index_weights_softmax_scale,
                index_weights_head_scale,
                index_q_packed,
                index_q_scale,
                index_weights_out,
            )

        # Values stay uint8 (2 E2M1 nibbles per byte). Scales are 4 ue8m0
        # bytes per (token, head) reinterpreted as one int32, then squeezed
        # from (T, H, 1) to (T, H) to match DeepGEMM's expected q_sf rank
        # (prefill wants 2-D (seq_len, num_heads); decode reshapes this to
        # 3-D (batch, next_n, num_heads)).
        return (
            index_q_packed,
            index_q_scale.view(torch.int32).squeeze(-1),
        ), index_weights_out

    assert weights_out_dtype == torch.float32, weights_out_dtype
    fp8_dtype = current_platform.fp8_dtype()
    use_fnuz = fp8_dtype == torch.float8_e4m3fnuz
    fp8_max = 224.0 if use_fnuz else 448.0
    index_q_fp8 = torch.empty_like(index_q, dtype=fp8_dtype)
    if has_cutedsl():
        # lazily import, otherwise some tests fail due to CUDA driver init failure.
        from vllm.models.deepseek_v4.nvidia.ops.fused_indexer_q_cutedsl import (
            _INDEXER_Q_FP8_KERNEL,
        )

        _INDEXER_Q_FP8_KERNEL(
            positions=positions,
            q=index_q,
            cos_sin_cache=index_q_cos_sin_cache,
            weights=index_weights,
            weights_softmax_scale=index_weights_softmax_scale,
            weights_head_scale=index_weights_head_scale,
            q_fp8=index_q_fp8.view(torch.uint8),
            weights_out=index_weights_out,
        )
    elif current_platform.is_xpu():
        torch.ops.vllm.xpu_deepseek_fused_indexer_q_rope_fp8(
            index_q,
            positions,
            index_q_cos_sin_cache,
            index_weights,
            index_weights_softmax_scale,
            index_weights_head_scale,
            index_q_fp8,
            index_weights_out,
        )
    elif current_platform.is_cpu():
        from vllm._custom_ops import fused_indexer_q_rope_quant_cpu

        fused_indexer_q_rope_quant_cpu(
            positions.contiguous(),
            index_q.to(torch.float32),
            index_q_cos_sin_cache.to(torch.float32).contiguous(),
            index_q_fp8,
            index_weights.to(torch.float32),
            index_weights_softmax_scale,
            index_weights_head_scale,
            index_weights_out,
        )
    else:
        _FUSED_INDEXER_Q_ROPE_QUANT_TRITON_KERNEL(
            positions,
            index_q,
            index_q_cos_sin_cache,
            index_weights,
            index_weights_softmax_scale,
            index_weights_head_scale,
            index_q_fp8,
            index_weights_out,
            fp8_max=fp8_max,
            use_fnuz=use_fnuz,
        )
    return index_q_fp8, index_weights_out