Skip to content

vllm.models.minimax_m3.nvidia.ops.gather_nvfp4_index_k

Gather NVFP4 index-K pages into compact FP8 pages for the Q8KV8 prefill indexer.

A packed NVFP4 page [128, 72] holds the E2M1 codes of all 128 tokens (64 B per token, low nibble first) followed by their E4M3 16-element block scales (token * 8 + group), with no global scale. Each element dequantizes to e4m3_rn_satfinite(e2m1 * scale), the value the Q8KV4 decode indexer scores, so prefill and decode rank pages from identical keys.

Where Triton can emit it, each four elements are one QMUL4 (mul.rn.satfinite.e4m3x4.e2m1x4.e4m3x4), the instruction the decode kernel dequantizes with; otherwise the product is formed exactly in fp32 and rounded once, which gives the same bytes.

Functions:

_qmul4_supported(device_index) cached

Whether Triton emits PTX that can use QMUL4 on this device.

Triton takes the PTX ISA version from its Blackwell ptxas: the bundled one unless TRITON_PTXAS_BLACKWELL_PATH points at a CUDA 13.4+ ptxas.

Source code in vllm/models/minimax_m3/nvidia/ops/gather_nvfp4_index_k.py 
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
@functools.cache
def _qmul4_supported(device_index: int) -> bool:
    """Whether Triton emits PTX that can use QMUL4 on this device.

    Triton takes the PTX ISA version from its Blackwell ptxas: the bundled one
    unless ``TRITON_PTXAS_BLACKWELL_PATH`` points at a CUDA 13.4+ ptxas.
    """
    capability = current_platform.get_device_capability(device_index)
    if capability is None or capability.to_int() not in (100, 103):
        return False
    try:
        from triton.backends.nvidia.compiler import get_ptxas, ptx_get_version

        ptxas_version = get_ptxas(capability.to_int()).version
        return ptx_get_version(ptxas_version) >= QMUL4_MIN_PTX_VERSION
    except Exception:
        # Triton-internal helpers; without them keep the portable path.
        return False

minimax_m3_gather_nvfp4_index_k_fp8(kv_cache, block_table, num_pages, page_offsets, max_num_pages, out)

Dequantize each request's leading NVFP4 pages into contiguous FP8 pages.

Parameters:

  • kv_cache

    (Tensor) –

    [num_blocks, 128, 72] uint8 packed NVFP4 pages.

  • block_table

    (Tensor) –

    [num_reqs, max_blocks] int32 physical page ids.

  • num_pages

    (Tensor) –

    [num_reqs] int32 pages to gather per request (its first num_pages[r] logical pages).

  • page_offsets

    (Tensor) –

    [num_reqs] int32 first output page of each request.

  • max_num_pages

    (int) –

    Host upper bound of num_pages.

  • out

    (Tensor) –

    [>= sum(num_pages), 128, 128] contiguous float8_e4m3fn.

Returns:

Source code in vllm/models/minimax_m3/nvidia/ops/gather_nvfp4_index_k.py 
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
def minimax_m3_gather_nvfp4_index_k_fp8(
    kv_cache: torch.Tensor,
    block_table: torch.Tensor,
    num_pages: torch.Tensor,
    page_offsets: torch.Tensor,
    max_num_pages: int,
    out: torch.Tensor,
) -> torch.Tensor:
    """Dequantize each request's leading NVFP4 pages into contiguous FP8 pages.

    Args:
        kv_cache: ``[num_blocks, 128, 72]`` uint8 packed NVFP4 pages.
        block_table: ``[num_reqs, max_blocks]`` int32 physical page ids.
        num_pages: ``[num_reqs]`` int32 pages to gather per request (its first
            ``num_pages[r]`` logical pages).
        page_offsets: ``[num_reqs]`` int32 first output page of each request.
        max_num_pages: Host upper bound of ``num_pages``.
        out: ``[>= sum(num_pages), 128, 128]`` contiguous float8_e4m3fn.

    Returns:
        ``out``.

    """
    assert kv_cache.dtype == torch.uint8 and kv_cache.shape[1:] == (
        PAGE_SIZE,
        HEAD_DIM // 2 + HEAD_DIM // 16,
    )
    assert kv_cache.stride(1) == kv_cache.shape[2] and kv_cache.stride(2) == 1
    assert out.dtype == torch.float8_e4m3fn and out.is_contiguous()
    assert out.shape[1:] == (PAGE_SIZE, HEAD_DIM)
    assert block_table.stride(1) == 1
    num_reqs = block_table.shape[0]
    if num_reqs == 0 or max_num_pages == 0:
        return out
    _gather_nvfp4_index_k_fp8_kernel[(num_reqs, max_num_pages)](
        kv_cache,
        block_table,
        num_pages,
        page_offsets,
        out,
        kv_cache.stride(0),
        block_table.stride(0),
        PAGE_SIZE=PAGE_SIZE,
        HEAD_DIM=HEAD_DIM,
        E4M3_MAX=E4M3_MAX,
        USE_QMUL4=_qmul4_supported(kv_cache.device.index or 0),
        num_warps=4,
    )
    return out