diff --git a/tests/kernels/attention/test_minimax_m3_fp8_triton_indexer.py b/tests/kernels/attention/test_minimax_m3_fp8_triton_indexer.py index 247faba17f..dcf7a6bc03 100644 --- a/tests/kernels/attention/test_minimax_m3_fp8_triton_indexer.py +++ b/tests/kernels/attention/test_minimax_m3_fp8_triton_indexer.py @@ -107,3 +107,37 @@ def test_fp8_index_query_and_cache_prefill_topk_matches_bf16() -> None: return {int(b) for b in idx[0, 0].tolist() if b >= 0} _check_selection(select(keys), select(keys.to(torch.float8_e4m3fn))) + + +@pytest.mark.skipif(not FP8_CUDA_SUPPORTED, reason="CUDA FP8 support required") +@torch.inference_mode() +def test_bf16_query_against_fp8_cache_decode_topk_matches_bf16() -> None: + """A mixed pair must upcast rather than emit a mixed-dtype fp8 dot. + + Production allocates index_q with the index-K cache dtype, so the scorers + normally see a matched pair and keep the native FP8 MMA. This pins the + other branch: a bf16 query against an e4m3 cache still compiles and still + ranks the planted blocks. + """ + dev = "cuda" + q, keys = _make_inputs(dev) + block_table = torch.arange(NUM_PAGES, dtype=torch.int32, device=dev).unsqueeze(0) + seq_lens = torch.tensor([SEQ_LEN], dtype=torch.int32, device=dev) + + def select(cache: torch.Tensor) -> set[int]: + idx = minimax_m3_index_decode( + q, # bf16 query, deliberately not cast to the cache dtype + cache.view(NUM_PAGES, PAGE, HEAD_DIM).contiguous(), + block_table, + seq_lens, + SEQ_LEN, + TOPK, + 0, + 0, + HEADS, + 1, + 1, + ) + return {int(b) for b in idx[0, 0].tolist() if b >= 0} + + _check_selection(select(keys), select(keys.to(torch.float8_e4m3fn))) diff --git a/vllm/models/minimax_m3/common/ops/index_topk.py b/vllm/models/minimax_m3/common/ops/index_topk.py index 94b74363e7..025c3dd428 100644 --- a/vllm/models/minimax_m3/common/ops/index_topk.py +++ b/vllm/models/minimax_m3/common/ops/index_topk.py @@ -124,9 +124,11 @@ def _index_block_score_kernel( order=(1, 0), ) q = tl.load(q_ptrs, boundary_check=(0,), padding_option="zero") - # FP8 has no mixed-dtype dot with bf16/fp32. Leave bf16 and fp32 loads - # in their stored dtype so fp32 pipeline tests keep full precision. - if q.dtype.is_fp8(): + # A matched fp8 pair (e.g. e4m3 x e4m3) lowers to a native FP8 MMA, so + # upcast only when the operands differ: there is no mixed-dtype fp8 MMA, + # including across fp8 flavours (fp8e4nv vs fp8e4b8). bf16 and fp32 loads + # always stay in their stored dtype. + if q.dtype.is_fp8() and q.dtype != ik_cache_ptr.dtype.element_ty: q = q.to(tl.bfloat16) q_start = prefix_len + pid_q * BLOCK_SIZE_Q @@ -162,7 +164,7 @@ def _index_block_score_kernel( + off_k[None, :] * stride_ik_pos + off_d[:, None] * stride_ik_d, ) - if k.dtype.is_fp8(): + if k.dtype.is_fp8() and k.dtype != q.dtype: k = k.to(tl.bfloat16) qk = tl.dot(q, k, out_dtype=tl.float32) # apply causal mask as needed @@ -378,7 +380,7 @@ def _decode_index_score_kernel( mask=q_mask[None, :], other=0.0, ) # [D,HQ] - if q.dtype.is_fp8(): + if q.dtype.is_fp8() and q.dtype != ik_cache_ptr.dtype.element_ty: q = q.to(tl.bfloat16) for blk in tl.range(chunk_start_block, chunk_end_block): page = tl.load(bt_row + blk).to(tl.int64) @@ -393,10 +395,10 @@ def _decode_index_score_kernel( + off_k[:, None] * stride_ik_pos + off_d * stride_ik_d, ) # [N,D] - # Upcast only FP8 cache loads. BF16 operands keep FP8 dot - # instructions out of the fallback path; FP32 accumulation - # preserves score accuracy. BF16/FP32 loads stay as stored. - if k.dtype.is_fp8(): + # Matched fp8 operands keep the native FP8 MMA; any mismatch upcasts, + # since no mixed-dtype fp8 MMA exists. FP32 accumulation preserves + # score accuracy either way. BF16/FP32 loads stay as stored. + if k.dtype.is_fp8() and k.dtype != q.dtype: k = k.to(tl.bfloat16) kq = tl.dot(k, q, out_dtype=tl.float32) # [N,HQ] kq = tl.where(pos_mask & q_mask[None, :], kq, float("-inf"))