diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/msa_utils.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/msa_utils.py index b3aa264d893e..db0312237d8d 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/msa_utils.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/msa_utils.py @@ -23,6 +23,21 @@ MSA_REQUIRED_HEAD_DIM = 128 +def check_decode_span_shape(kernel: str, total_q: int, batch: int, query_len: int) -> None: + """Reject a q that does not cover exactly the batch it was handed. + + The decode kernels derive the request id as token // query_len, so a longer + q reads page table rows and lengths past the batch's last one. The caller + names itself so the error does too, rather than surfacing as an assert + several frames inside a kernel. + """ + if total_q != batch * query_len: + raise ValueError( + f"{kernel}: total_q ({total_q}) must be batch ({batch}) * " + f"decode_query_len ({query_len})." + ) + + def is_msa_layer(attn) -> bool: """Whether this layer's attention is served by the MiniMax-M3 MSA kernels.""" sparse_params = attn.sparse_params @@ -103,6 +118,8 @@ def write_msa_main_kv( out_cache_loc: torch.Tensor, k: torch.Tensor, v: torch.Tensor, + *, + num_live_tokens: int, ) -> None: """Write new-token K and V into the paged main cache at out_cache_loc. @@ -116,10 +133,18 @@ def write_msa_main_kv( head_dim = int(k_view.shape[3]) num_tokens = int(k.shape[0]) write_kv_slots( - k_view, out_cache_loc, k.reshape(num_tokens, num_kv_heads, head_dim), layout="HND" + k_view, + out_cache_loc, + k.reshape(num_tokens, num_kv_heads, head_dim), + layout="HND", + num_live_tokens=num_live_tokens, ) write_kv_slots( - v_view, out_cache_loc, v.reshape(num_tokens, num_kv_heads, head_dim), layout="HND" + v_view, + out_cache_loc, + v.reshape(num_tokens, num_kv_heads, head_dim), + layout="HND", + num_live_tokens=num_live_tokens, ) @@ -163,6 +188,8 @@ def write_msa_phase_kv( metadata.msa_out_cache_loc[token_offset : token_offset + num_tokens], k, v, + # k and v are this phase's own token slice, so every row owns a slot. + num_live_tokens=num_tokens, ) @@ -268,6 +295,7 @@ def select_blocks_from_maxscore( "MSA_REQUIRED_HEAD_DIM", "MSA_REQUIRED_TOPK", "build_kv_page_indices", + "check_decode_span_shape", "msa_package_available", "msa_paged_kv", "per_token_valid_blocks", diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/paged_cache.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/paged_cache.py index 3a20070d56c3..44df6e9ee08c 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/paged_cache.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/paged_cache.py @@ -15,6 +15,7 @@ def write_kv_slots( values: torch.Tensor, *, layout: Literal["NHD", "HND"] = "NHD", + num_live_tokens: int, ) -> None: """Write per-token values into a K, V, or index-K cache at given slots. @@ -28,7 +29,27 @@ def write_kv_slots( mapping satisfies this contract: ``get_block_ids_per_seq`` canonicalizes padded ``BAD_PAGE_INDEX`` entries before ``build_paged_kv_slot_mapping`` selects only the allocated live-token positions. + + `num_live_tokens` is how many leading rows own a real cache slot; the rest + are dropped. It is required rather than defaulted because a caller passing + the padded token extent of a piecewise CUDA graph corrupts the cache + silently: torch wraps a negative index, so the -1 sentinel past the live + count lands in the last page instead of raising. """ + # Trimming by count keeps this sync-free, the sentinel tail being contiguous + # by construction, where masking on the slot values would not. + if num_live_tokens < 0: + raise ValueError(f"num_live_tokens must be non-negative, got {num_live_tokens}") + if out_cache_loc.shape[0] < num_live_tokens or values.shape[0] < num_live_tokens: + raise ValueError( + f"num_live_tokens={num_live_tokens} exceeds the rows supplied " + f"(out_cache_loc={out_cache_loc.shape[0]}, values={values.shape[0]})" + ) + if num_live_tokens == 0: + return + out_cache_loc = out_cache_loc[:num_live_tokens] + values = values[:num_live_tokens] + with torch.no_grad(): if cache.ndim >= 4: token_axis = 2 if layout == "HND" else 1 diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/triton_sparse_decode.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/triton_sparse_decode.py index 4e00b384b857..664c131f7b9c 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/triton_sparse_decode.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/triton_sparse_decode.py @@ -36,6 +36,8 @@ from tensorrt_llm._torch.memory_buffer_utils import get_memory_buffers +from .msa_utils import check_decode_span_shape + # One sparse block is exactly one KV page. SPARSE_BLOCK_SIZE = 128 @@ -315,11 +317,9 @@ def minimax_m3_sparse_attn_decode( """ total_q, num_heads, head_dim = q.shape num_kv_heads = int(k_paged.shape[1]) - if total_q != int(seq_lens.shape[0]) * decode_query_len: - raise ValueError( - f"total_q ({total_q}) must be batch ({int(seq_lens.shape[0])}) * " - f"decode_query_len ({decode_query_len})." - ) + check_decode_span_shape( + "MiniMax-M3 Triton sparse decode", total_q, int(seq_lens.shape[0]), decode_query_len + ) if int(k_paged.shape[2]) != SPARSE_BLOCK_SIZE: raise ValueError( f"MiniMax-M3 sparse decode requires page_size={SPARSE_BLOCK_SIZE}; " diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/trtllm_gen_dense_decode.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/trtllm_gen_dense_decode.py index b477dc06a8d7..c94991b836f6 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/trtllm_gen_dense_decode.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/trtllm_gen_dense_decode.py @@ -27,6 +27,8 @@ from tensorrt_llm._torch.memory_buffer_utils import get_memory_buffers +from .msa_utils import check_decode_span_shape + @functools.lru_cache(maxsize=None) def _counter_size(num_heads: int, max_num_requests: int, device_index: int) -> int: @@ -252,6 +254,15 @@ def minimax_m3_trtllm_gen_dense_decode( """ import flashinfer + # The multi-CTA KV counters are sized against max_num_requests, so a batch + # read out of a longer q would undersize them. + check_decode_span_shape( + "MiniMax-M3 trtllm-gen dense decode", + int(q.shape[0]), + int(seq_lens.shape[0]), + decode_query_len, + ) + kv_pool, subpages_per_slot = kv_cache_manager.get_kv_subpage_pool(layer_idx, "HND") num_heads = int(q.shape[1]) diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/msa_backend.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/msa_backend.py index 5abbedc884aa..87baca8f4f04 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/msa_backend.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/msa_backend.py @@ -182,6 +182,8 @@ class MiniMaxM3MsaSparseAttentionMetadata(TrtllmAttentionMetadata): # _msa_fields_ready marks that the current step's buffers are populated. _msa_buffers_ready: bool = False _msa_fields_ready: bool = False + # Read through msa_live_token_count, which states the contract. + _msa_live_total_q: int = 0 # Sparse geometry the plans need. _msa_params: Optional[MiniMaxM3SparseMetadataParams] = None # This step's fmha_sm100 plans, plain tuples with no graph-stable buffers @@ -974,6 +976,10 @@ def _build_msa_fields(self) -> None: decoding the inputs for on_update_kv_lens are staged as well. """ self._msa_fields_ready = False + # The live count describes the slot mapping staged at the end of this + # method, so clear it with the ready flag. An early return below would + # otherwise leave a count describing a batch that is no longer scheduled. + self._msa_live_total_q = 0 if not self._msa_buffers_ready: return request_ids = self.request_ids @@ -1034,6 +1040,13 @@ def _build_msa_fields(self) -> None: f"smaller than the step's per-request page count ({block_table_cols})." ) + # Piecewise graphs pad token-shaped inputs to the capture bucket, and the + # fused index producer runs over that padded extent: it sits inside the + # captured region, so trimming it to a host-side count would make its + # shape dynamic. A negative slot is what makes those rows cache-write + # no-ops. Fill before the copy, because a prior larger step leaves valid + # slot ids in the tail and those address real pages. + self.msa_out_cache_loc.fill_(-1) self.msa_out_cache_loc[:total_new_tokens].copy_(out_cache_loc, non_blocking=True) # Captured producers also execute padded rows. Invalidate only the # unwritten tail so they cannot reuse the previous step's live slots. @@ -1060,6 +1073,12 @@ def _build_msa_fields(self) -> None: self.msa_subpage_block_table[:batch_size], ) + # Live, unpadded new-token count for this step, staged before either + # exit so the non-dynamic early return below carries it too. Cache + # writers read it through msa_live_token_count to find where + # msa_out_cache_loc stops holding real slots. + self._msa_live_total_q = total_new_tokens + self._msa_kv_lens_dynamic = self._msa_kv_lens_may_change() if not self._msa_kv_lens_dynamic: self._msa_fields_ready = True @@ -1087,6 +1106,21 @@ def _build_msa_fields(self) -> None: self.msa_kv_lens_staged[:batch_size].copy_(kv_lens_cpu, non_blocking=True) self._msa_fields_ready = True + def msa_live_token_count(self) -> int: + """This step's live, unpadded new-token count. + + Cache writers take the padded token extent and need this to find where + msa_out_cache_loc stops holding real slots. Raising when no mapping is + staged keeps an unprepared step from writing against another step's + slots, which a stale count would otherwise allow. + """ + if not self._msa_fields_ready: + raise RuntimeError( + "MiniMax-M3 MSA cache write requires prepared metadata, but " + "prepare() did not stage a slot mapping for this step." + ) + return self._msa_live_total_q + def msa_idx_k_cache(self, layer_idx: int) -> torch.Tensor: """Return the paged index-K cache in the HND layout MSA consumes.""" return self.kv_cache_manager.get_index_k_buffer(layer_idx) @@ -1101,6 +1135,7 @@ def msa_write_idx_k(self, layer_idx: int, idx_k: torch.Tensor) -> None: self.msa_out_cache_loc[:num_tokens], idx_k.reshape(num_tokens, 1, sparse_index_dim), layout="HND", + num_live_tokens=self.msa_live_token_count(), ) def msa_proxy_max_score_view( @@ -1257,17 +1292,22 @@ def write_layer_caches( return num_kv_heads = int(k_view.shape[1]) head_dim = int(k_view.shape[3]) + # The dispatch clips k/v/idx_k to the step's live tokens before this + # runs, so every supplied row owns a real slot: num_tokens is the live + # count write_kv_slots requires. write_kv_slots( k_view, out_cache_loc, k.reshape(num_tokens, num_kv_heads, head_dim), layout="HND", + num_live_tokens=num_tokens, ) write_kv_slots( v_view, out_cache_loc, v.reshape(num_tokens, num_kv_heads, head_dim), layout="HND", + num_live_tokens=num_tokens, ) if idx_k is not None: write_kv_slots( @@ -1275,6 +1315,7 @@ def write_layer_caches( out_cache_loc, idx_k.reshape(num_tokens, 1, int(idx_cache.shape[-1])), layout="HND", + num_live_tokens=num_tokens, ) def run_indexer( diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/triton_backend.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/triton_backend.py index 844fd979469c..a6ce0e0b4608 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/triton_backend.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/triton_backend.py @@ -158,8 +158,11 @@ def _write_main_kv_slots_to_pool( ``out_cache_loc`` is the 1-D ``[num_new_tokens]`` int tensor of flat slot ids to update. ``pool[:, kv_index]`` is a storage-sharing view, so the shared :func:`common.write_kv_slots` propagates the write to the pool. + + Every row owns a slot, so the row count is the live count: the slot mapping + emits one real slot per new token and no sentinel. """ - write_kv_slots(pool[:, kv_index], out_cache_loc, values) + write_kv_slots(pool[:, kv_index], out_cache_loc, values, num_live_tokens=int(values.shape[0])) def _write_main_kv_slots( @@ -173,7 +176,7 @@ def _write_main_kv_slots( flat-slot layout used by focused unit tests and the 4-D paged view of ``kv_pool[:, 0]`` / ``kv_pool[:, 1]``. """ - write_kv_slots(cache, out_cache_loc, values) + write_kv_slots(cache, out_cache_loc, values, num_live_tokens=int(values.shape[0])) def _scatter_topk_to_block_mask( diff --git a/tensorrt_llm/_torch/models/modeling_minimaxm3.py b/tensorrt_llm/_torch/models/modeling_minimaxm3.py index a7722bccfc1e..d8223e46410e 100644 --- a/tensorrt_llm/_torch/models/modeling_minimaxm3.py +++ b/tensorrt_llm/_torch/models/modeling_minimaxm3.py @@ -886,6 +886,41 @@ def _minimax_m3_fused_sparse_qkv_producer_fake( ) +def _dispatch_attention_over_live_tokens( + attn_layer: "MiniMaxM3Attention", + q: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], + idx_q: Optional[torch.Tensor], + idx_k: Optional[torch.Tensor], + attn_metadata: AttentionMetadata, + output: torch.Tensor, +) -> None: + """Run the attention core over the step's live tokens alone. + + A piecewise CUDA graph pads token-shaped inputs up to its capture bucket + without adding requests to go with them (see _get_padding_params in + model_engine), so q can outrun the rows the batch has. The kernels below + read a request out of a token index, so the pad comes off here, once for + both dispatch paths. + + No kernel writes the pad rows of output, so they are zeroed rather than left + holding stale values that nothing can distinguish from a real NaN. + """ + num_tokens = int(attn_metadata.num_tokens) + if num_tokens < int(output.shape[0]): + output[num_tokens:].zero_() + attn_layer._dispatch_attention_backend( + q[:num_tokens], + k[:num_tokens] if k is not None else None, + v[:num_tokens] if v is not None else None, + idx_q[:num_tokens] if idx_q is not None else None, + idx_k[:num_tokens] if idx_k is not None else None, + attn_metadata, + output[:num_tokens], + ) + + @torch.library.custom_op("trtllm::minimax_m3_attn_custom_op_inplace", mutates_args=("output",)) def minimax_m3_attn_custom_op_inplace( q: Optional[torch.Tensor], @@ -923,15 +958,9 @@ def minimax_m3_attn_custom_op_inplace( k = v = idx_k = None if q is None: raise RuntimeError(f"MiniMax-M3 attention layer {layer_idx} received no query tensor.") - attn_layer._dispatch_attention_backend( - q[:num_tokens], - k[:num_tokens] if k is not None else None, - v[:num_tokens] if v is not None else None, - idx_q[:num_tokens] if idx_q is not None else None, - idx_k[:num_tokens] if idx_k is not None else None, - attn_metadata, - output[:num_tokens], - ) + # The live token count is a host value, so the compiled graph above must + # not see it: it would guard on it and recapture per count. + _dispatch_attention_over_live_tokens(attn_layer, q, k, v, idx_q, idx_k, attn_metadata, output) maybe_bcg_minimax_m3_attn_custom_op_inplace = eager_on_graph(minimax_m3_attn_custom_op_inplace) @@ -1843,7 +1872,9 @@ def _forward_attention_core( output, ) else: - self._dispatch_attention_backend(q, k, v, idx_q, idx_k, attn_metadata, output) + # A step that runs here rather than through compile is padded all + # the same, since the bucket is agreed across ranks. + _dispatch_attention_over_live_tokens(self, q, k, v, idx_q, idx_k, attn_metadata, output) return output def _dispatch_attention_backend( diff --git a/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py b/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py index 4fd1c5535b13..eea1190ce693 100644 --- a/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py +++ b/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py @@ -27,6 +27,7 @@ ) from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_utils import ( MSA_REQUIRED_TOPK, + check_decode_span_shape, msa_paged_kv, ) from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.paged_cache import ( @@ -645,6 +646,8 @@ def get_index_k_buffer(self, layer_idx): metadata.kv_cache_manager = manager metadata.msa_out_cache_loc = torch.tensor([2, page_size + 5], dtype=torch.int32) values = torch.arange(2 * head_dim, dtype=torch.float32).reshape(2, 1, head_dim) + metadata._msa_fields_ready = True + metadata._msa_live_total_q = 2 returned = metadata.msa_idx_k_cache(3) metadata.msa_write_idx_k(3, values) @@ -770,6 +773,7 @@ def msa_write_idx_k(self, layer_idx: int, idx_k: torch.Tensor) -> None: self.msa_out_cache_loc, idx_k, layout="HND", + num_live_tokens=int(idx_k.shape[0]), ) def msa_idx_k_cache(self, layer_idx: int) -> torch.Tensor: @@ -1186,6 +1190,178 @@ def test_the_decode_span_of_a_mixed_step_is_its_generation_suffix(): assert metadata.msa_max_kv_len == 40 +def test_decode_span_shape_check_names_the_kernel_that_rejected_the_q(): + """The guard both decode kernels share.""" + # Eleven speculative decode requests of 4 query tokens each. + check_decode_span_shape("kernel", 44, 11, 4) + + # A piecewise CUDA graph's pad folded into the batch: the kernels would read + # 128 page table rows out of a batch that only has 11. + with pytest.raises(ValueError, match=r"kernel: total_q \(512\) must be batch \(11\)"): + check_decode_span_shape("kernel", 512, 11, 4) + + +def test_both_decode_kernels_reject_a_q_that_outruns_the_batch(): + """The guard has to be reached, not merely available. + + Both kernels read their shapes before touching a device, so the refusal + happens at the call and needs no GPU. The dense kernel is handed no cache + manager for the same reason: it must decline before consulting one. + """ + batch, query_len, total_q = 11, 4, 512 + num_heads, head_dim, page_size = 4, 128, 128 + q = torch.empty(total_q, num_heads, head_dim) + output = torch.empty(total_q, num_heads, head_dim) + block_table = torch.zeros(batch, 4, dtype=torch.int32) + seq_lens = torch.zeros(batch, dtype=torch.int32) + + pytest.importorskip("triton") + from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.triton_sparse_decode import ( + minimax_m3_sparse_attn_decode, + ) + + paged = torch.empty(1, 1, page_size, head_dim) + with pytest.raises( + ValueError, match=r"Triton sparse decode: total_q \(512\) must be batch \(11\)" + ): + minimax_m3_sparse_attn_decode( + q, + paged, + paged, + torch.zeros(1, total_q, 64, dtype=torch.int64), + block_table, + seq_lens, + sm_scale=head_dim**-0.5, + output=output, + decode_query_len=query_len, + ) + + pytest.importorskip("flashinfer") + from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.trtllm_gen_dense_decode import ( + minimax_m3_trtllm_gen_dense_decode, + ) + + with pytest.raises( + ValueError, match=r"trtllm-gen dense decode: total_q \(512\) must be batch \(11\)" + ): + minimax_m3_trtllm_gen_dense_decode( + q, + None, + 0, + block_table, + seq_lens, + sm_scale=head_dim**-0.5, + output=output, + decode_query_len=query_len, + max_seq_len=1024, + max_num_requests=batch, + ) + + +def test_a_shrinking_step_leaves_no_live_slot_in_the_padded_tail(): + """The slot-guarded consumers recognize a negative slot and nothing else, so + every row a step does not own has to hold one.""" + block_ids = torch.arange(3 * 4, dtype=torch.int32).reshape(3, 4) + + class FakeCacheManager: + tokens_per_block = 4 + + def get_block_ids_per_seq(self, request_ids): + return block_ids[: len(request_ids)] + + def get_buffers(self, layer_idx, kv_layout="NHD"): + # Only its device is read, which keeps this test off the GPU. + return torch.zeros(1, dtype=torch.bfloat16) + + metadata_cls = MiniMaxM3MsaSparseAttention.Metadata + metadata = metadata_cls.__new__(metadata_cls) + metadata.kv_cache_manager = FakeCacheManager() + metadata.mapping = None + metadata._msa_buffers_ready = True + metadata._msa_params = None + metadata._msa_decode_span = None + metadata.max_num_sequences = 3 + metadata.msa_subpage_block_table = None + metadata._msa_subpages_per_slot = 0 + metadata.msa_out_cache_loc = torch.zeros(16, dtype=torch.int32) + metadata.msa_kv_indices = torch.zeros(3 * 4, dtype=torch.int32) + metadata.msa_block_table = torch.zeros(3, 4, dtype=torch.int32) + metadata.msa_seq_lens_cuda = torch.zeros(3, dtype=torch.int32) + # Non-speculative, so no on_update_kv_lens staging: the slot tail this test + # is about is the same either way, and the staging buffers are not part of + # the fixture. + metadata._msa_kv_lens_may_change = lambda: False + + def build(request_ids, qo_lens, kv_lens): + # The host lengths are read-only properties over these. + metadata.request_ids = request_ids + metadata._msa_qo_lens_cpu = torch.tensor(qo_lens, dtype=torch.int32) + metadata._msa_kv_lens_cpu = torch.tensor(kv_lens, dtype=torch.int32) + metadata._msa_qo_offset_cpu = metadata._msa_kv_lens_cpu - metadata._msa_qo_lens_cpu + metadata._build_msa_fields() + + # A three-request step of six new tokens, then a smaller one of two. + build([0, 1, 2], (4, 1, 1), (9, 3, 5)) + assert metadata.msa_live_token_count() == 6 + wide_tail = metadata.msa_out_cache_loc[2:6].tolist() + assert all(slot >= 0 for slot in wide_tail) + + build([0, 1], (1, 1), (3, 5)) + + assert metadata.msa_live_token_count() == 2 + # The rows the shrunk step does not own, including those the wider one did. + assert metadata.msa_out_cache_loc[2:].tolist() == [-1] * 14 + assert (metadata.msa_out_cache_loc[:2] >= 0).all() + + +def test_the_eager_writer_drops_the_sentinel_tail_it_is_handed(): + """The padded rows of a step must reach no page at all, the wrap target of a + surviving -1 included.""" + num_pages, page_size, head_dim = 4, 8, 16 + cache = torch.zeros(num_pages, 1, page_size, head_dim, dtype=torch.bfloat16) + # Two live tokens, then the -1 tail a capture bucket pads the step out to. + out_cache_loc = torch.tensor([2, page_size + 5, -1, -1], dtype=torch.int32) + values = torch.arange(4 * head_dim, dtype=torch.float32).reshape(4, 1, head_dim) + + write_kv_slots(cache, out_cache_loc, values, layout="HND", num_live_tokens=2) + + torch.testing.assert_close(cache[0, 0, 2], values[0, 0].to(torch.bfloat16)) + torch.testing.assert_close(cache[1, 0, 5], values[1, 0].to(torch.bfloat16)) + # The page a wrapped -1 would have hit. + assert not cache[num_pages - 1].any() + + +def test_the_eager_writer_refuses_a_live_count_it_has_no_rows_for(): + """A count past the rows supplied is a caller bug, not a short write.""" + cache = torch.zeros(4, 1, 8, 16, dtype=torch.bfloat16) + out_cache_loc = torch.tensor([2, 5], dtype=torch.int32) + values = torch.zeros(2, 1, 16) + + with pytest.raises(ValueError, match=r"num_live_tokens=3 exceeds the rows supplied"): + write_kv_slots(cache, out_cache_loc, values, layout="HND", num_live_tokens=3) + + with pytest.raises(ValueError, match="must be non-negative"): + write_kv_slots(cache, out_cache_loc, values, layout="HND", num_live_tokens=-1) + + # A step that scheduled nothing writes nothing rather than erroring. + write_kv_slots(cache, out_cache_loc, values, layout="HND", num_live_tokens=0) + assert not cache.any() + + +def test_an_unprepared_step_has_no_live_count_to_write_against(): + """A stale count would let a write land on another step's slots.""" + metadata_cls = MiniMaxM3MsaSparseAttention.Metadata + metadata = metadata_cls.__new__(metadata_cls) + metadata._msa_fields_ready = False + metadata._msa_live_total_q = 7 + + with pytest.raises(RuntimeError, match="did not stage a slot mapping"): + metadata.msa_live_token_count() + + metadata._msa_fields_ready = True + assert metadata.msa_live_token_count() == 7 + + def test_a_pure_prefill_step_has_no_decode_span(): """A step with no generation row has nothing for the decode kernels, and fmha_sm100 keeps every plan and the page table they read.""" @@ -1808,14 +1984,26 @@ def test_fused_scatter_matches_reference(src_dtype, cache_dtype, with_idx): ref_pool = pool.clone() ref_idx_pool = idx_pool.clone() write_kv_slots( - ref_pool[:, 0], slots, k.reshape(num_tokens, num_kv_heads, head_dim), layout="HND" + ref_pool[:, 0], + slots, + k.reshape(num_tokens, num_kv_heads, head_dim), + layout="HND", + num_live_tokens=num_tokens, ) write_kv_slots( - ref_pool[:, 1], slots, v.reshape(num_tokens, num_kv_heads, head_dim), layout="HND" + ref_pool[:, 1], + slots, + v.reshape(num_tokens, num_kv_heads, head_dim), + layout="HND", + num_live_tokens=num_tokens, ) if with_idx: write_kv_slots( - ref_idx_pool[:, 0], slots, idx_k.reshape(num_tokens, 1, head_dim), layout="HND" + ref_idx_pool[:, 0], + slots, + idx_k.reshape(num_tokens, 1, head_dim), + layout="HND", + num_live_tokens=num_tokens, ) assert fused_write_layer_caches( diff --git a/tests/unittest/_torch/models/test_minimax_m3.py b/tests/unittest/_torch/models/test_minimax_m3.py index 32705c5708e8..59b250952953 100644 --- a/tests/unittest/_torch/models/test_minimax_m3.py +++ b/tests/unittest/_torch/models/test_minimax_m3.py @@ -66,6 +66,7 @@ MiniMaxM3MoE, MiniMaxM3QKVIndexerLinear, _build_swiglu_oai_dense_mlp, + _dispatch_attention_over_live_tokens, _load_qkv_index_proj_weights, _minimax_m3_swiglu_oai, _moe_routed_output_is_global, @@ -457,7 +458,7 @@ def _dispatch_attention_backend(self, q, k, v, idx_q, idx_k, attn_metadata, outp assert layer.producer_shapes == ((2, 5), (1, 2), 2) torch.testing.assert_close(output[:2], packed[:2, :3]) - torch.testing.assert_close(output[2:], torch.full((2, 3), -1.0)) + torch.testing.assert_close(output[2:], torch.zeros((2, 3))) @pytest.mark.cpu_only @@ -996,6 +997,98 @@ def _has_cuda() -> bool: # --------------------------------------------------------------------------- +def test_attention_dispatch_clips_the_piecewise_token_pad(): + """Only the live tokens reach the attention core, and the output's pad rows + come back zeroed rather than as the buffer supplied them.""" + seen = {} + + def capture(q, k, v, idx_q, idx_k, attn_metadata, output): + del attn_metadata + seen["rows"] = [None if t is None else int(t.shape[0]) for t in (q, k, v, idx_q, idx_k)] + seen["out_rows"] = int(output.shape[0]) + output.fill_(7.0) + + attn_layer = SimpleNamespace(_dispatch_attention_backend=capture) + # Eleven speculative decode requests of 4 query tokens, padded to 64. + padded, live, hidden = 64, 44, 8 + q = torch.ones((padded, hidden)) + output = torch.full((padded, hidden), float("nan")) + + _dispatch_attention_over_live_tokens( + attn_layer, + q, + q, + q, + None, + None, + SimpleNamespace(num_tokens=live), + output, + ) + + assert seen["rows"] == [live, live, live, None, None] + assert seen["out_rows"] == live + assert torch.equal(output[:live], torch.full((live, hidden), 7.0)) + assert torch.equal(output[live:], torch.zeros(padded - live, hidden)) + + +def test_attention_dispatch_leaves_an_unpadded_step_alone(): + """No pad, so nothing to clip and nothing to zero.""" + seen = {} + + def capture(q, k, v, idx_q, idx_k, attn_metadata, output): + del q, k, v, idx_q, idx_k, attn_metadata + seen["out"] = output + + output = torch.full((5, 8), float("nan")) + _dispatch_attention_over_live_tokens( + SimpleNamespace(_dispatch_attention_backend=capture), + torch.ones((5, 8)), + None, + None, + None, + None, + SimpleNamespace(num_tokens=5), + output, + ) + + assert seen["out"].shape == (5, 8) + assert output.isnan().all() + + +@pytest.mark.cpu_only +def test_every_attention_dispatch_goes_through_the_live_token_clip(): + """A dispatch path that bypassed the clip would hand the kernels the pad. + + The clipping helper is the only thing standing between a padded step and + kernels that read a request out of a token index, so which callers reach + the backend is itself the invariant. + """ + import ast + import inspect + + from tensorrt_llm._torch.models import modeling_minimaxm3 + + tree = ast.parse(inspect.getsource(modeling_minimaxm3)) + parents = {child: parent for parent in ast.walk(tree) for child in ast.iter_child_nodes(parent)} + + def enclosing_function(node): + while node is not None: + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)): + return node.name + node = parents.get(node) + return None + + callers = { + enclosing_function(node) + for node in ast.walk(tree) + if isinstance(node, ast.Call) + and isinstance(node.func, ast.Attribute) + and node.func.attr == "_dispatch_attention_backend" + } + + assert callers == {"_dispatch_attention_over_live_tokens"} + + def test_is_minimax_m3_vl_config_detects_vl(): assert is_minimax_m3_vl_config(_make_vl_config()) is True diff --git a/tests/unittest/_torch/thop/parallel_hw_agnostic/test_minimax_m3_fp8_indexer.py b/tests/unittest/_torch/thop/parallel_hw_agnostic/test_minimax_m3_fp8_indexer.py index 042ccd858e6a..bbe8cad52171 100644 --- a/tests/unittest/_torch/thop/parallel_hw_agnostic/test_minimax_m3_fp8_indexer.py +++ b/tests/unittest/_torch/thop/parallel_hw_agnostic/test_minimax_m3_fp8_indexer.py @@ -78,6 +78,37 @@ def _strided_cache(num_pages: int, page_size: int = 128, stride_scale: int = 7) return backing[::stride_scale] +def _guarded_cache( + num_pages: int, page_size: int = 128, stride_scale: int = 7, guard_pages: int = 2 +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Allocate a strided cache flanked by zeroed guard regions. + + Returns the cache view plus the backing below and above it. A store from a + negative slot lands below the base and one from an out-of-range slot lands + above the last page, so both stay inspectable rather than depending on a + fault that an out-of-bounds FP8 store may never raise. + """ + total_pages = num_pages + 2 * guard_pages + backing = torch.zeros( + total_pages * stride_scale, + 1, + page_size, + 128, + dtype=torch.float8_e4m3fn, + device="cuda", + ) + first = guard_pages * stride_scale + last = (guard_pages + num_pages) * stride_scale + view = backing[first:last:stride_scale] + assert view.shape[0] == num_pages + return view, backing[:first], backing[last:] + + +def _assert_all_zero(*regions: torch.Tensor) -> None: + for region in regions: + assert torch.count_nonzero(region.reshape(-1).view(torch.uint8)).item() == 0 + + def _run( qk: torch.Tensor, cache: torch.Tensor, @@ -163,6 +194,55 @@ def test_minimax_m3_fp8_indexer_defensively_skips_invalid_direct_op_slots() -> N assert torch.count_nonzero(backing[2].view(torch.uint8)).item() == 0 +@pytest.mark.parametrize("num_live_tokens", [0, 4]) +@pytest.mark.parametrize("tail_slot", [-1, 4 * 128, 4 * 128 + 61, 5 * 128 + 127]) +def test_minimax_m3_fp8_indexer_ignores_a_padded_slot_tail( + tail_slot: int, num_live_tokens: int +) -> None: + """Rows past the live prefix must leave every cache byte untouched. + + A direct caller can hand the kernel a padded token height whose trailing + slots are the -1 sentinel or a stale out-of-range id, which the two guards + have to drop. The failure modes differ: -1 truncates to page 0 at offset + -1, one token below the cache base, while an out-of-range page scatters + above the pool. The tail slots start at the first page past the end and + stay inside the guard region, so a regression trips an assert rather than + faulting. + """ + torch.manual_seed(2468) + num_heads_q = 4 + page_size = 128 + num_pages = 4 + padded_tokens = 17 + + qk = torch.randn(padded_tokens, (num_heads_q + 1) * 128, dtype=torch.bfloat16, device="cuda") + q_weight = torch.randn(128, dtype=torch.bfloat16, device="cuda") + k_weight = torch.randn(128, dtype=torch.bfloat16, device="cuda") + position_ids = torch.arange(padded_tokens, dtype=torch.int32, device="cuda") + 1024 + + cache, below, above = _guarded_cache(num_pages, page_size) + slots = torch.full((padded_tokens,), tail_slot, dtype=torch.int32, device="cuda") + # One page per live row, so a stray tail store cannot be mistaken for one. + pages = torch.arange(num_live_tokens, dtype=torch.int32, device="cuda") + within = (pages * 37) % page_size + slots[:num_live_tokens] = pages * page_size + within + + q_out = _run(qk, cache, slots, q_weight, k_weight, position_ids, num_heads_q) + q_ref, k_ref = _reference(qk, num_heads_q, q_weight, k_weight, position_ids) + + # Only the cache store is slot-gated; index-Q is produced for every row. + _assert_fp8_close(q_out, q_ref) + _assert_all_zero(below, above) + if num_live_tokens: + _assert_fp8_close(cache[pages.long(), 0, within.long()], k_ref[:num_live_tokens]) + + # A tail store landing on a valid page at the wrong offset would clear both + # guard regions, so require every unwritten slot to stay zero as well. + written = torch.zeros(num_pages, page_size, dtype=torch.bool, device="cuda") + written[pages.long(), within.long()] = True + _assert_all_zero(cache[:, 0][~written]) + + def test_minimax_m3_fp8_indexer_cuda_graph_replay_updates_outputs() -> None: torch.manual_seed(5678) num_tokens = 16