From 02f0e53d79763abbe1c2f2ad51c375e5af3fcaeb Mon Sep 17 00:00:00 2001 From: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> Date: Fri, 28 Aug 2026 00:47:49 -0700 Subject: [PATCH 1/2] [nvbugs/6669902][fix] Update MiniMax-M3 decode tests to the block-table build_kv_page_indices signature PR #17986 narrowed build_kv_page_indices to take the host block-id table directly instead of (req_to_token, slot_ids, ...), and PR #17842 merged 47 minutes later with tests still calling the four-argument form, so _flat_page_table raised TypeError before either kernel ran. The helper now wants exactly what these tests already hold, so drop the req_to_token reconstruction and pass block_table on the host. The page ids are unchanged: the old form gathered (block_table[b, p] * PAGE_SIZE) // PAGE_SIZE at each page boundary, and slot_ids was an identity map. Both A/B tests now reach their fmha_sm100 comparison and agree with it, so the three waivers for this bug are removed. Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> --- .../minimax_m3_index_decode_score.py | 15 +------------- .../test_minimax_m3_index_decode_score.py | 20 +------------------ .../test_minimax_m3_sparse_attn_decode.py | 20 +------------------ 3 files changed, 3 insertions(+), 52 deletions(-) diff --git a/tests/microbenchmarks/minimax_m3_index_decode_score.py b/tests/microbenchmarks/minimax_m3_index_decode_score.py index 5fbca1f35b2f..520c8f9b8ae4 100644 --- a/tests/microbenchmarks/minimax_m3_index_decode_score.py +++ b/tests/microbenchmarks/minimax_m3_index_decode_score.py @@ -26,19 +26,6 @@ HEAD_DIM = 128 -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, - ) - - def _time_us(fn, warmup: int = 20, iters: int = 100) -> float: for _ in range(warmup): fn() @@ -92,7 +79,7 @@ def run_cutedsl(): qo_lens_cpu = torch.full((batch,), dql, dtype=torch.int32) kv_lens_cpu = torch.full((batch,), seq_len, dtype=torch.int32) - kv_indices = _flat_page_table(block_table, kv_lens_cpu).cuda() + kv_indices = build_kv_page_indices(block_table.cpu(), kv_lens_cpu, PAGE_SIZE).cuda() k_paged = k_cache.unsqueeze(1) def run_msa(): diff --git a/tests/unittest/_torch/attention/sparse/test_minimax_m3_index_decode_score.py b/tests/unittest/_torch/attention/sparse/test_minimax_m3_index_decode_score.py index 2efe611ef2a5..fdc24a2e1e3a 100644 --- a/tests/unittest/_torch/attention/sparse/test_minimax_m3_index_decode_score.py +++ b/tests/unittest/_torch/attention/sparse/test_minimax_m3_index_decode_score.py @@ -32,24 +32,6 @@ pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -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. - - build_kv_page_indices reads those ids out of a token-level slot map, so this - rebuilds the map the block table implies and lets the production helper do - the flattening. - """ - 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, - ) - - def _runner(): """Return the CuTe DSL runner class, skipping the test if unavailable.""" pytest.importorskip("cutlass") @@ -290,7 +272,7 @@ def test_index_decode_score_matches_msa_proxy(dtype): qo_lens_cpu = torch.ones(batch, dtype=torch.int32) kv_lens_cpu = torch.tensor(seq_lens, dtype=torch.int32) - kv_indices = _flat_page_table(block_table, kv_lens_cpu).cuda() + kv_indices = build_kv_page_indices(block_table.cpu(), kv_lens_cpu, PAGE_SIZE).cuda() msa_score = _proxy_max_score( idx_q, k_cache.unsqueeze(1), diff --git a/tests/unittest/_torch/attention/sparse/test_minimax_m3_sparse_attn_decode.py b/tests/unittest/_torch/attention/sparse/test_minimax_m3_sparse_attn_decode.py index 91459251056c..ac974fb00960 100644 --- a/tests/unittest/_torch/attention/sparse/test_minimax_m3_sparse_attn_decode.py +++ b/tests/unittest/_torch/attention/sparse/test_minimax_m3_sparse_attn_decode.py @@ -36,24 +36,6 @@ ) -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. - - build_kv_page_indices reads those ids out of a token-level slot map, so this - rebuilds the map the block table implies and lets the production helper do - the flattening. - """ - 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, - ) - - def _reference_sparse_decode( q: torch.Tensor, k_paged: torch.Tensor, @@ -397,7 +379,7 @@ def test_sparse_decode_matches_msa_kernel(): qo_lens_cpu = torch.ones(batch, dtype=torch.int32) kv_lens_cpu = torch.tensor(seq_lens, dtype=torch.int32) - kv_indices = _flat_page_table(block_table, kv_lens_cpu).cuda() + kv_indices = build_kv_page_indices(block_table.cpu(), kv_lens_cpu, PAGE_SIZE).cuda() msa_out = torch.empty_like(q) run_msa_sparse_gqa( q.to(torch.float8_e4m3fn), From 6e2225b7fcf0393246356cbc9fb43d753e32bae0 Mon Sep 17 00:00:00 2001 From: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> Date: Sat, 29 Aug 2026 00:51:51 -0700 Subject: [PATCH 2/2] [nvbugs/6669902][chore] Remove stale waiver after fix Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> --- tests/integration/test_lists/waives.txt | 3 --- 1 file changed, 3 deletions(-) diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index 72141d14d007..aa2a9fd22bef 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -307,10 +307,7 @@ 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/test_minimax_m3_index_decode_score.py::test_index_decode_score_matches_msa_proxy[dtype0] SKIP (https://nvbugs/6669902) -unittest/_torch/attention/sparse/test_minimax_m3_index_decode_score.py::test_index_decode_score_matches_msa_proxy[dtype1] SKIP (https://nvbugs/6669902) unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py::test_msa_paged_hnd_input_materializes_unaligned_outer_stride SKIP (https://nvbugs/6661846) -unittest/_torch/attention/sparse/test_minimax_m3_sparse_attn_decode.py::test_sparse_decode_matches_msa_kernel SKIP (https://nvbugs/6669902) 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) unittest/_torch/executor/test_overlap_scheduler.py::test_overlap_scheduler_block_reuse_cache_hit SKIP (https://nvbugs/6608387)