From 985d00a751eb8f2682e269aeda153f8d065416e2 Mon Sep 17 00:00:00 2001 From: Yongye Zhu Date: Thu, 18 Jun 2026 03:48:33 +0000 Subject: [PATCH] [Test] M3 indexer: add sm_scale arg to _reference_index_topk Fixes mypy call-arg errors: two fp8 call sites pass sm_scale (by keyword and positionally) that the reference signature lacked. Add sm_scale (default 1.0) and apply it to the score; top-k selection is invariant to a positive scalar, so existing callers are unchanged. Co-Authored-By: Claude Opus 4.8 Signed-off-by: Yongye Zhu (cherry picked from commit 522f262265b6dbcb9ccda12776f0cd9857ceb839) --- tests/kernels/attention/test_minimax_m3.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/kernels/attention/test_minimax_m3.py b/tests/kernels/attention/test_minimax_m3.py index 5e6016ffd53..0340ca9a477 100644 --- a/tests/kernels/attention/test_minimax_m3.py +++ b/tests/kernels/attention/test_minimax_m3.py @@ -134,6 +134,7 @@ def _reference_index_topk( topk: int, init_blocks: int, local_blocks: int, + sm_scale: float = 1.0, ) -> torch.Tensor: total_q, num_idx_heads, _ = idx_q.shape out = torch.full( @@ -149,7 +150,7 @@ def _reference_index_topk( num_blocks = (seq_len + BLOCK_SIZE - 1) // BLOCK_SIZE pages = block_table[req_id, :num_blocks] k = index_kv_cache[pages].reshape(num_blocks * BLOCK_SIZE, -1) - score = torch.einsum("qhd,kd->hqk", q.float(), k.float()) + score = torch.einsum("qhd,kd->hqk", q.float(), k.float()) * sm_scale q_pos = prefix_len + torch.arange(q_len, device=idx_q.device) k_pos = torch.arange(k.shape[0], device=idx_q.device)