From 021cd19ae8317257b3382853c6917e1e62c458c4 Mon Sep 17 00:00:00 2001 From: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> Date: Thu, 8 Oct 2026 23:37:34 -0700 Subject: [PATCH 1/2] perf: rc1_m3_dense_decode_sdpa_per_kv_head_loop (+0.00%) Landed by repair-bot on outer loop 1. Measured gain for this commit: +0.00%. Test case: minimax_m3_mxfp8-bench-pytorch-float8-maxbs:1-maxnt:2048-input_output_len:8000,1000-reqs:10-con:1-gpus:4 NVBug: 6877539 Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> --- .../_torch/models/modeling_minimaxm3.py | 45 ++++++++++--------- 1 file changed, 23 insertions(+), 22 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_minimaxm3.py b/tensorrt_llm/_torch/models/modeling_minimaxm3.py index a7722bccfc1e..bd141860d376 100644 --- a/tensorrt_llm/_torch/models/modeling_minimaxm3.py +++ b/tensorrt_llm/_torch/models/modeling_minimaxm3.py @@ -1788,30 +1788,31 @@ def _sdpa_dense_attention_core( else: # Decode: one query token per request at position seq_lens - 1. valid = kv_positions < seq_lens_dev.unsqueeze(-1) # [batch, max_k] - q_b = q_view.view(batch, 1, self.num_heads, self.head_dim).transpose( - 1, 2 - ) # [batch, H, 1, d] mask_b = valid.unsqueeze(1).unsqueeze(1) # [batch, 1, 1, k] - # Expand K/V one KV head at a time: all heads at once needs an - # O(batch * max_k * num_heads) temporary that overflows the CUDA-graph - # pool under attention DP. - out_b = q.new_empty(batch, self.num_heads, 1, self.head_dim) + # Fold the GQA group into SDPA's query-length axis: the ``group`` Q + # heads sharing KV head h become ``group`` query rows of one SDPA + # head that reads KV head h unexpanded. Exact for decode (every + # query row sits at seq_lens - 1, so one key mask serves them all), + # and it needs neither an O(batch * max_k * num_heads) + # repeat_interleave temporary in the CUDA-graph pool nor one SDPA + # launch per KV head. + num_kv_heads = max(self.num_key_value_heads, 1) + q_b = q_view.view(batch, num_kv_heads, group, self.head_dim) # [batch, KV, group, d] + k_b = k_padded.transpose(1, 2) # [batch, KV, k, d] + v_b = v_padded.transpose(1, 2) # [batch, KV, k, d] with sdpa_kernel(_DENSE_SDPA_BACKENDS): - for h in range(max(self.num_key_value_heads, 1)): - qh = slice(h * group, (h + 1) * group) - k_h = k_padded[:, :, h : h + 1].repeat_interleave(group, dim=2) - v_h = v_padded[:, :, h : h + 1].repeat_interleave(group, dim=2) - out_b[:, qh] = torch.nn.functional.scaled_dot_product_attention( - q_b[:, qh].to(q.dtype), - k_h.transpose(1, 2).to(q.dtype), - v_h.transpose(1, 2).to(q.dtype), - attn_mask=mask_b, - dropout_p=0.0, - is_causal=False, - ) # [batch, group, 1, d] - # Drop the singleton query axis; a transpose(1, 2).reshape would - # scramble (head, head_dim) when H != head_dim. - output.view(batch, self.num_heads, self.head_dim).copy_(out_b.squeeze(2)) + out_b = torch.nn.functional.scaled_dot_product_attention( + q_b.to(q.dtype), + k_b.to(q.dtype), + v_b.to(q.dtype), + attn_mask=mask_b, + dropout_p=0.0, + is_causal=False, + ) # [batch, KV, group, d] + # Q head h * group + g is query row g of SDPA head h, so viewing the + # output as [batch, KV, group, d] keeps (head, head_dim) ordering; a + # transpose(1, 2).reshape would scramble it when H != head_dim. + output.view(batch, num_kv_heads, group, self.head_dim).copy_(out_b) return output From 782c5271a1e2818ff30933e735b723c99f438561 Mon Sep 17 00:00:00 2001 From: chenfeiz0326 Date: Sun, 11 Oct 2026 03:55:06 -0700 Subject: [PATCH 2/2] [https://nvbugs/6877539][test] Add parity test for MiniMax-M3 dense SDPA decode The decode branch of MiniMaxM3Attention._sdpa_dense_attention_core folds each GQA group into SDPA's query-length axis. No unit test exercised it. Add a parity test that runs the real method against a per-request repeat_interleave reference. It covers num_kv_heads in {1, 2, 4} and group in {1, 4}, with a batch of mixed seq_lens. Each request owns scattered pages, and the unused pool slots are filled with large garbage values, so the seq_lens mask must hide them. Validated on B300. The test passes on this PR's code and on the previous per-KV-head loop. It fails on four broken variants: a transpose/reshape output copy, a group-major Q view, a batch-row-0 mask, and no mask. Co-Authored-By: Claude Signed-off-by: chenfeiz0326 --- .../unittest/_torch/models/test_minimax_m3.py | 97 +++++++++++++++++++ 1 file changed, 97 insertions(+) diff --git a/tests/unittest/_torch/models/test_minimax_m3.py b/tests/unittest/_torch/models/test_minimax_m3.py index 32705c5708e8..bb526f2394b4 100644 --- a/tests/unittest/_torch/models/test_minimax_m3.py +++ b/tests/unittest/_torch/models/test_minimax_m3.py @@ -2391,3 +2391,100 @@ def test_minimax_m3_swiglu_oai_fused_matches_reference(dtype): plain_ref = torch.nn.functional.silu(gate_c) * up_c plain = swiglu(gate_up, swiglu_limit=limit) torch.testing.assert_close(plain.float(), plain_ref.float(), atol=atol, rtol=1e-2) + + +def _dense_sdpa_decode_reference(q, k_ctx, v_ctx, seq_lens, num_kv_heads, head_dim): + """Per-request GQA decode via explicit ``repeat_interleave`` (pre-#18872 math). + + ``q`` is ``[batch, num_heads * head_dim]``; ``k_ctx`` / ``v_ctx`` are + ``[batch, max_len, num_kv_heads, head_dim]`` holding each request's full + context, valid for positions ``< seq_lens[b]``. + """ + batch = q.shape[0] + num_heads = q.shape[1] // head_dim + group = num_heads // num_kv_heads + out = torch.empty_like(q, dtype=torch.float32) + for b in range(batch): + n = int(seq_lens[b]) + q_b = q[b].view(num_heads, 1, head_dim).float() # [H, 1, d] + k_b = k_ctx[b, :n].float().repeat_interleave(group, dim=1).transpose(0, 1) # [H, n, d] + v_b = v_ctx[b, :n].float().repeat_interleave(group, dim=1).transpose(0, 1) + o_b = torch.nn.functional.scaled_dot_product_attention(q_b, k_b, v_b) # [H, 1, d] + out[b] = o_b.reshape(-1) + return out + + +@pytest.mark.skipif(not _has_cuda(), reason="dense SDPA decode runs on the CUDA KV pool") +@pytest.mark.parametrize("num_kv_heads", [1, 2, 4]) +@pytest.mark.parametrize("group", [1, 4]) +def test_minimax_m3_sdpa_dense_decode_matches_repeat_interleave_reference(num_kv_heads, group): + """The decode branch of ``_sdpa_dense_attention_core`` matches plain GQA math. + + Covers ``group > 1``, ``num_kv_heads > 1`` and a batch with mixed + ``seq_lens`` (shorter rows padded out to ``max_k``), so a head-order or + mask-broadcast mistake in the GQA-into-query-axis fold shows up here even + though it would not move throughput (NVBug 6877539). + """ + torch.manual_seed(0) + device, dtype = torch.device("cuda"), torch.bfloat16 + head_dim, tokens_per_block = 64, 4 + num_heads = num_kv_heads * group + seq_lens = torch.tensor([5, 1, 11, 8], dtype=torch.int32) # mixed; max 11 -> 3 blocks + batch, max_len = int(seq_lens.numel()), int(seq_lens.max()) + blocks_per_req = -(-max_len // tokens_per_block) + + # Requests own scattered, non-contiguous pages so a gather or slot mix-up + # cannot pass by accident. + num_pages = batch * blocks_per_req + 3 + pages = torch.randperm(num_pages)[: batch * blocks_per_req].view(batch, blocks_per_req) + pos = torch.arange(blocks_per_req * tokens_per_block) + req_to_token = ( + pages[:, pos // tokens_per_block] * tokens_per_block + pos % tokens_per_block + ).to(device=device, dtype=torch.int32) + + k_ctx = torch.randn(batch, max_len, num_kv_heads, head_dim, device=device, dtype=dtype) + v_ctx = torch.randn_like(k_ctx) + # Pool layout [num_pages, kv_factor, tokens_per_block, num_kv_heads, head_dim]. + # Fill every unused slot with garbage: the seq_lens mask must hide it. + kv_pool = ( + torch.randn( + num_pages, 2, tokens_per_block, num_kv_heads, head_dim, device=device, dtype=dtype + ) + * 50 + ) + for b in range(batch): + n = int(seq_lens[b]) + slots = req_to_token[b, : n - 1].long() # context before this decode step + page, within = slots // tokens_per_block, slots % tokens_per_block + kv_pool[page, 0, within] = k_ctx[b, : n - 1] + kv_pool[page, 1, within] = v_ctx[b, : n - 1] + # The new token's K/V arrives through ``k`` / ``v`` and is written by the method. + out_cache_loc = torch.stack( + [req_to_token[b, int(seq_lens[b]) - 1] for b in range(batch)] + ).long() + k_new = k_ctx[torch.arange(batch), seq_lens.long() - 1].reshape(batch, -1) + v_new = v_ctx[torch.arange(batch), seq_lens.long() - 1].reshape(batch, -1) + q = torch.randn(batch, num_heads * head_dim, device=device, dtype=dtype) + + layer = SimpleNamespace( + layer_idx=0, num_heads=num_heads, num_key_value_heads=num_kv_heads, head_dim=head_dim + ) + m3_meta = SimpleNamespace( + slot_ids=torch.arange(batch, device=device, dtype=torch.int32), + req_to_token=req_to_token, + seq_lens=seq_lens.to(device), + max_seqlen_k=max_len, + is_prefill=False, + ) + attn_metadata = SimpleNamespace( + kv_cache_manager=SimpleNamespace(get_buffers=lambda layer_idx: kv_pool), + minimax_m3={"metadata": m3_meta, "out_cache_loc": out_cache_loc}, + is_cuda_graph=False, + max_seq_len=max_len, + ) + output = torch.empty(batch, num_heads * head_dim, device=device, dtype=dtype) + + MiniMaxM3Attention._sdpa_dense_attention_core(layer, q, k_new, v_new, attn_metadata, output) + + ref = _dense_sdpa_decode_reference(q, k_ctx, v_ctx, seq_lens, num_kv_heads, head_dim) + torch.testing.assert_close(output.float(), ref, atol=2e-2, rtol=2e-2)