From 2e5f1a29757dee18481aa3b9a70634214724939c Mon Sep 17 00:00:00 2001 From: chenfeiz0326 <203214996+chenfeiz0326@users.noreply.github.com> Date: Tue, 13 Oct 2026 21:52:06 -0700 Subject: [PATCH] [None][fix] Guard the CuTe DSL MLA decode batch and add a regression test #20256 limited the page table CuteDslMlaFmha._run_mla_decode passes to the CuTe DSL MLA decode kernel to this batch's generation rows; with max_num_sequences columns the kernel launched one request per column and read q, the output and the sequence lengths past their ends (illegal address in the first Kimi-K3 helix warmup forward on GB300). On top of that: - Slice sequence_lengths to the batch as well, so every operand the kernel receives has the same batch size. - CuteDSLNVMlaDecodeBlackwellRunner.forward now raises if page_table, q or the output disagrees with cache_seqs on the batch dimension, so a mismatched caller fails with a message instead of reading out of bounds. - Add test_cute_dsl_mla_decode_batch_below_max_num_sequences: runs _run_mla_decode with a 32-slot block table and a batch of 1 or 3 at the Kimi-K3 helix verify-group shape (H=96, 7 query tokens, per-token kv_bounds), bf16 and fp8, against a reference. It fails on the pre-#20256 slice and passes now. Signed-off-by: chenfeiz0326 <203214996+chenfeiz0326@users.noreply.github.com> --- .../attention/backends/fmha/cute_dsl_mla.py | 2 +- .../_torch/custom_ops/cute_dsl_custom_ops.py | 11 ++ .../attention/test_cute_dsl_mla_helix.py | 115 ++++++++++++++++++ 3 files changed, 127 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/attention/backends/fmha/cute_dsl_mla.py b/tensorrt_llm/_torch/attention/backends/fmha/cute_dsl_mla.py index 8ab481750855..a58a95ef7792 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/cute_dsl_mla.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/cute_dsl_mla.py @@ -466,7 +466,7 @@ def _run_mla_decode( local_layer_idx = attn.get_local_layer_idx(meta) pool_idx = int(pool_mapping[local_layer_idx, 0]) page_table_layer = block_offsets[pool_idx, :, 0, :] - cache_seqs_base = params.sequence_lengths.to(torch.int32) + cache_seqs_base = params.sequence_lengths[:batch_size].to(torch.int32) # The page-table batch dimension determines the kernel launch batch. page_table = ( page_table_layer[meta.num_contexts : meta.num_contexts + batch_size] diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py index 5d5b34b9da3d..8e90c68d5153 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py @@ -12366,6 +12366,17 @@ def forward( # workspace = lse + split_kv_workspace batch_size = cache_seqs.shape[0] + # The kernel launches one request per page-table column, so a + # wider table would read q, o and cache_seqs past their ends. + if (page_table.shape[1] != batch_size + or q_latent.shape[3] != batch_size + or o.shape[3] != batch_size): + raise RuntimeError( + "CuteDSLNVMlaDecodeBlackwellRunner requires page_table " + "(max_blocks, B), q_latent and o (..., B) to agree with " + f"cache_seqs (B,), got B={batch_size}, page_table=" + f"{tuple(page_table.shape)}, q_latent=" + f"{tuple(q_latent.shape)}, o={tuple(o.shape)}.") d_latent = q_latent.shape[1] softmax_stats_kernel = None if softmax_stats is not None: diff --git a/tests/unittest/_torch/attention/test_cute_dsl_mla_helix.py b/tests/unittest/_torch/attention/test_cute_dsl_mla_helix.py index 4345eb3110de..93a31454f71a 100644 --- a/tests/unittest/_torch/attention/test_cute_dsl_mla_helix.py +++ b/tests/unittest/_torch/attention/test_cute_dsl_mla_helix.py @@ -146,3 +146,118 @@ def test_cute_dsl_mla_helix_stats_and_empty_local_kv( actual_log_partition = softmax_stats[1, :, 0] + softmax_stats[1, :, 1].log() expected_log_partition = torch.logsumexp(scores, dim=-1) torch.testing.assert_close(actual_log_partition, expected_log_partition, rtol=1e-5, atol=1e-5) + + +class _FakeKvCacheManager: + def __init__(self, kv_pool: torch.Tensor): + self._kv_pool = kv_pool + + def get_buffers(self, layer_idx: int) -> torch.Tensor: + return self._kv_pool + + +@pytest.mark.parametrize("input_dtype", [torch.bfloat16, torch.float8_e4m3fn]) +@pytest.mark.parametrize("batch_size", [1, 3]) +def test_cute_dsl_mla_decode_batch_below_max_num_sequences( + input_dtype: torch.dtype, batch_size: int +) -> None: + """A decode batch smaller than the metadata's sequence capacity. + + kv_cache_block_offsets holds one row per sequence slot (max_num_sequences), + while the kernel derives its batch size from the page table it receives. + The FMHA wrapper must hand it only this batch's generation rows; any extra + column launches CTAs for requests past the end of q, the output and the + sequence lengths. Runs the K3 helix verify-group shape (H=96, 1 + 6 draft + tokens, per-token kv_bounds) through CuteDslMlaFmha._run_mla_decode. + """ + from types import SimpleNamespace + + from tensorrt_llm._torch.attention.backends.fmha.cute_dsl_mla import CuteDslMlaFmha + + torch.manual_seed(23) + device = torch.device("cuda") + max_num_sequences, seq_len_q, num_heads = 32, 7, 96 + latent_dim, rope_dim, qk_nope_head_dim, page_size = 512, 64, 128, 64 + max_blocks = 8 + num_pages = max_num_sequences * max_blocks + head_dim = latent_dim + rope_dim + kv_lens = [200, 70, 135][:batch_size] + softmax_scale = 1.0 / math.sqrt(qk_nope_head_dim + rope_dim) + + kv_pool = ( + torch.randn(num_pages, 1, page_size, 1, head_dim, device=device, dtype=torch.bfloat16) * 0.1 + ).to(input_dtype) + # One disjoint page range per sequence slot. Slots past the batch are + # allocated but must never reach the kernel. + block_offsets = ( + torch.arange(num_pages, dtype=torch.int32, device=device) + .view(1, max_num_sequences, 1, max_blocks) + .expand(1, max_num_sequences, 2, max_blocks) + .contiguous() + ) + kv_bounds = torch.cat( + [ + torch.arange(kv_len - seq_len_q + 1, kv_len + 1, dtype=torch.int32, device=device) + for kv_len in kv_lens + ] + ) + meta = SimpleNamespace( + kv_cache_manager=_FakeKvCacheManager(kv_pool), + kv_cache_block_offsets=block_offsets, + host_kv_cache_pool_mapping=torch.zeros(1, 2, dtype=torch.int32), + num_contexts=0, + max_num_requests=max_num_sequences, + helix_position_offsets=torch.zeros( + batch_size * seq_len_q, dtype=torch.int32, device=device + ), + _helix_spec_tokens_valid=True, + helix_kv_bounds=kv_bounds, + ) + attn = SimpleNamespace( + kv_lora_rank=latent_dim, + qk_rope_head_dim=rope_dim, + qk_nope_head_dim=qk_nope_head_dim, + num_heads=num_heads, + q_scaling=1.0, + layer_idx=0, + get_local_layer_idx=lambda _meta: 0, + ) + num_tokens = batch_size * seq_len_q + q = torch.randn(num_tokens, num_heads * head_dim, device=device, dtype=torch.bfloat16) * 0.1 + output = torch.empty(num_tokens, num_heads * latent_dim, device=device, dtype=torch.bfloat16) + softmax_stats = torch.empty(num_tokens, num_heads, 2, device=device, dtype=torch.float32) + # As in production: one sequence length per generation request. + sequence_lengths = torch.tensor(kv_lens, dtype=torch.int32, device=device) + + fmha = object.__new__(CuteDslMlaFmha) + if input_dtype == torch.float8_e4m3fn: + fmha._cute_dsl_fp8_scale = (softmax_scale, 1.0) + workspace = torch.empty( + fmha._required_workspace_size(num_heads, seq_len_q, latent_dim, max_num_sequences), + device=device, + dtype=torch.uint8, + ) + params = SimpleNamespace( + attn=attn, + meta=meta, + num_requests=batch_size, + tokens_per_block=page_size, + sequence_lengths=sequence_lengths, + workspace=workspace, + fwd=SimpleNamespace(quant_q_buffer=None, softmax_stats_tensor=softmax_stats), + ) + + fmha._run_mla_decode(q, output, params, input_dtype) + torch.cuda.synchronize() + + q_rows = q.to(input_dtype).float().view(batch_size, seq_len_q, num_heads, head_dim) + out_rows = output.float().view(batch_size, seq_len_q, num_heads, latent_dim) + tol = 5e-2 if input_dtype == torch.float8_e4m3fn else 3e-2 + for b, kv_len in enumerate(kv_lens): + pages = block_offsets[0, b, 0, : -(-kv_len // page_size)].long() + keys = kv_pool[pages, 0, :, 0].float().reshape(-1, head_dim)[:kv_len] + for t in range(seq_len_q): + bound = int(kv_bounds[b * seq_len_q + t]) + scores = q_rows[b, t] @ keys[:bound].transpose(0, 1) * softmax_scale + expected = torch.softmax(scores, dim=-1) @ keys[:bound, :latent_dim] + torch.testing.assert_close(out_rows[b, t], expected, rtol=tol, atol=tol)