Skip to content

vllm.models.minimax_m3.amd.ops.nvfp4_index_prefill.host

Launch wrapper of the FlyDSL NVFP4 prefill index scorers.

Functions:

minimax_m3_nvfp4_index_score_prefill(q_idx, index_cache, score, block_table, cu_seqlens_q, seq_lens, init_blocks, local_blocks, max_query_len, max_seq_len)

Block scores of every prefill row for 1, 2 or 4 index heads.

Source code in vllm/models/minimax_m3/amd/ops/nvfp4_index_prefill/host.py 
 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
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
def minimax_m3_nvfp4_index_score_prefill(
    q_idx: torch.Tensor,  # [total_q, H, 128] e4m3
    index_cache: torch.Tensor,  # [num_blocks, 128, 72] uint8
    score: torch.Tensor,  # [H, total_q, >= cdiv(max_seq_len, 128)] fp32
    block_table: torch.Tensor,  # [batch, >= max_blocks] int32
    cu_seqlens_q: torch.Tensor,  # [batch + 1] int32
    seq_lens: torch.Tensor,  # [batch] int32
    init_blocks: int,
    local_blocks: int,
    max_query_len: int,
    max_seq_len: int,
) -> None:
    """Block scores of every prefill row for 1, 2 or 4 index heads."""
    batch = cu_seqlens_q.shape[0] - 1
    total_q, heads, _ = q_idx.shape
    if total_q == 0 or block_table.shape[0] == 0:
        return
    assert _supported(q_idx, index_cache, score, block_table, cu_seqlens_q, seq_lens)
    max_blocks = cdiv(max_seq_len, _kernel.BLK)
    assert max_blocks <= score.shape[2]
    work = max_blocks * cdiv(max_query_len, 256) * batch * heads
    small = max_query_len * heads <= 128 or work < SMALL_WORK
    shared = small if PIN_SHARED is None else PIN_SHARED
    mq = PIN_MQ or (1 if small else 2)
    target = SMALL_TARGET if shared and work < MID_WORK else TARGET_WORKGROUPS
    tile_q = _kernel_shared().tile_q(mq) if shared else _kernel.TILE_Q
    qt = cdiv(max_query_len * heads, tile_q)
    min_seg = 1 if shared else _kernel.GROUP
    max_nseg = max(1, min(MAX_SEGMENTS, max_blocks // min_seg))
    nseg = max(1, min(max_nseg, target // max(1, qt * batch)))
    slots = WORKGROUPS_PER_CU * cu_count(q_idx.device)
    while nseg < max_nseg:
        rounds = qt * batch * nseg / slots
        if rounds / math.ceil(rounds) >= MIN_ROUND_FILL:
            break
        nseg += 1
    lut = () if shared else (nvfp4_e4m3_lut(q_idx.device).data_ptr(),)
    run_compiled(
        _get_kernel_shared(heads, mq) if shared else _get_kernel(heads),
        q_idx.data_ptr(),
        index_cache.data_ptr(),
        *lut,
        score.data_ptr(),
        block_table.data_ptr(),
        cu_seqlens_q.data_ptr(),
        seq_lens.data_ptr(),
        int(index_cache.stride(0)),
        int(block_table.stride(0)),
        int(score.stride(1)),
        int(nseg),
        int(qt),
        int(total_q),
        int(batch),
        int(init_blocks),
        int(local_blocks),
        int(score.stride(0)),
        int(cdiv(qt * nseg * batch, 8) * 8),
        torch.cuda.current_stream(),
    )

nvfp4_e4m3_lut(device)

[256] int64: byte j of entry s is e4m3_rn_satfinite(m_j * e4m3(s)) for the E2M1 magnitudes m = 0, 0.5, 1, 1.5, 2, 3, 4, 6.

Source code in vllm/models/minimax_m3/amd/ops/nvfp4_index_prefill/host.py 
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
def nvfp4_e4m3_lut(device: torch.device) -> torch.Tensor:
    """[256] int64: byte j of entry s is e4m3_rn_satfinite(m_j * e4m3(s)) for the
    E2M1 magnitudes m = 0, 0.5, 1, 1.5, 2, 3, 4, 6."""
    lut = _NVFP4_E4M3_LUT.get(device)
    if lut is None:
        # Device ops only, so a first call under graph capture is fine.
        scales = torch.arange(256, dtype=torch.uint8, device=device)
        scales = scales.view(torch.float8_e4m3fn).float()
        code = torch.arange(8, device=device)
        exp, man = code // 2, (code % 2).float()
        mags = torch.where(
            exp == 0, man * 0.5, torch.exp2((exp - 1).float()) * (1.0 + man * 0.5)
        )
        # e2m1 * e4m3 is exact in fp32; saturate, then round once to E4M3.
        values = (scales[:, None] * mags[None, :]).clamp(-448.0, 448.0)
        lut = values.to(torch.float8_e4m3fn).view(torch.int64).reshape(256)
        _NVFP4_E4M3_LUT[device] = lut
    return lut