From 613b012d5e3fbdf0fcb5bc038f66ca0baa9bd308 Mon Sep 17 00:00:00 2001 From: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> Date: Wed, 9 Sep 2026 14:49:05 -0700 Subject: [PATCH] [https://nvbugs/6670516][fix] Call build_kv_page_indices with the signature it has PR 17986 narrowed build_kv_page_indices from the req_to_token plus slot_ids pair to a single block_ids_cpu table and migrated its production caller. The three-argument callers that landed 47 minutes later in PR 17842 still passed four, so the two unit tests below raised TypeError on every run and were waived. PR 18106 has since repaired both tests, leaving the page-table microbenchmark as the last stale caller and the waivers with nothing to hide. Repoints the microbenchmark at the current helper, which already holds the block table the old code was expanding to token slots and folding back, and drops the three waivers. Unwaived: test_index_decode_score_matches_msa_proxy[dtype0], the same case at [dtype1], and test_sparse_decode_matches_msa_kernel. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> --- tests/integration/test_lists/waives.txt | 3 --- tests/microbenchmarks/minimax_m3_index_decode_score.py | 10 +--------- 2 files changed, 1 insertion(+), 12 deletions(-) diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index 2217a4340135..38216035f855 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -284,9 +284,6 @@ test_e2e.py::test_ptp_quickstart_advanced_deepseek_r1_w4afp8_8gpus[DeepSeek-R1-W test_e2e.py::test_ptp_quickstart_bert[TRTLLM-BertForSequenceClassification-bert/bert-base-uncased-yelp-polarity] SKIP (https://nvbugs/6605819) test_e2e.py::test_ptp_quickstart_bert[VANILLA-BertForSequenceClassification-bert/bert-base-uncased-yelp-polarity] SKIP (bug pending, tracked in PR 17414) unittest/_torch/attention/sparse/dsa/test_req_idx_per_token.py::test_on_update_kv_lens_rebuilds_stale_map SKIP (https://nvbugs/6574939) -unittest/_torch/attention/sparse/msa/test_minimax_m3_index_decode_score.py::test_index_decode_score_matches_msa_proxy[dtype0] SKIP (https://nvbugs/6669902) -unittest/_torch/attention/sparse/msa/test_minimax_m3_index_decode_score.py::test_index_decode_score_matches_msa_proxy[dtype1] SKIP (https://nvbugs/6669902) -unittest/_torch/attention/sparse/msa/test_minimax_m3_sparse_attn_decode.py::test_sparse_decode_matches_msa_kernel SKIP (https://nvbugs/6669902) unittest/_torch/attention/sparse/msa/test_msa_backend.py::test_msa_paged_hnd_input_materializes_unaligned_outer_stride SKIP (https://nvbugs/6661846) unittest/_torch/attention/test_attention_backends.py::test_attention_backend[exaone_moe_gqa_swa128-ctx-bf16-HND-p32-v1] SKIP (https://nvbugs/6668773) unittest/_torch/attention/test_attention_backends.py::test_attention_backend[qwen2_0_5b_gqa_hd64-ctx-bf16-HND-p32-v1] SKIP (https://nvbugs/6641268) diff --git a/tests/microbenchmarks/minimax_m3_index_decode_score.py b/tests/microbenchmarks/minimax_m3_index_decode_score.py index eb1d344a552a..f58d288cedbf 100644 --- a/tests/microbenchmarks/minimax_m3_index_decode_score.py +++ b/tests/microbenchmarks/minimax_m3_index_decode_score.py @@ -28,15 +28,7 @@ def _flat_page_table(block_table: torch.Tensor, kv_lens_cpu: torch.Tensor) -> torch.Tensor: """Flatten a block table into the per-request page ids fmha_sm100 consumes.""" - batch, max_pages = block_table.shape - intra = torch.arange(PAGE_SIZE, dtype=torch.int32) - req_to_token = (block_table.cpu().to(torch.int32) * PAGE_SIZE).unsqueeze(2) + intra - return build_kv_page_indices( - req_to_token.reshape(batch, max_pages * PAGE_SIZE), - torch.arange(batch, dtype=torch.int32), - kv_lens_cpu, - PAGE_SIZE, - ) + return build_kv_page_indices(block_table.cpu().to(torch.int32), kv_lens_cpu, PAGE_SIZE) def _time_us(fn, warmup: int = 20, iters: int = 100) -> float: