From 18e641c41c08fceb091180343f578e1e118f8901 Mon Sep 17 00:00:00 2001 From: Gal Hubara Agam <96368689+galagam@users.noreply.github.com> Date: Wed, 24 Jun 2026 12:21:12 -0700 Subject: [PATCH 1/3] [https://nvbugs/6316981][fix] AutoDeploy: zero-init MTP replay work-items buffer The SuperV3 MTP replay path allocated the shared `replay_work_items` metadata buffer with `torch.empty` in `CachedSequenceInterface._create_mamba_hybrid_cache_manager`. During CUDA graph capture/warmup the replay SSM kernel runs before `prepare_replay_metadata()` has populated this buffer, so in `persistent_main` mode the kernel reads uninitialized cache-slot indices and performs an out-of-bounds SSM state-cache access -- an illegal memory access that crashed executor init (e.g. test_mtp[nvfp4_ws4_80gb-trtllm]). Allocate the buffer with `torch.zeros` so every work item maps to slot 0 / pnat 0 (in-bounds) until real metadata is written per step. This mirrors the PyTorch backend, which already zero-inits `replay_work_items` (mamba2_metadata.py). Add a regression test that drives the real replay op (flashinfer_cached_ssm -> replay_selective_state_update) at the production SuperV3 shape (persistent_main mode) with the buffer allocated through the production path and left unprepared (capture scenario), asserting no IMA. The test deterministically faults without the fix and passes with it. Unwaive the SuperV3 test_mtp ws4 cases (bf16 / fp8 / nvfp4) under nvbugs 6316981 and 6336682, now that the capture-time IMA is fixed. Register the previously-uncovered nvfp4_ws8 trtllm MTP case on DGX-B200 (existing 8-GPU post-merge section). Raise the SuperV3 test_mtp acceptance-rate gate to a single 0.50 for all dtypes (observed >=53% on bf16/fp8/nvfp4 ws4). Signed-off-by: Gal Hubara Agam <96368689+galagam@users.noreply.github.com> --- .../_torch/auto_deploy/shim/interface.py | 8 +- .../defs/accuracy/test_llm_api_autodeploy.py | 5 +- .../test_lists/test-db/l0_dgx_b200.yml | 1 + tests/integration/test_lists/waives.txt | 6 - .../mamba/test_flashinfer_mamba_cached_op.py | 243 +++++++++++++++++- 5 files changed, 251 insertions(+), 12 deletions(-) diff --git a/tensorrt_llm/_torch/auto_deploy/shim/interface.py b/tensorrt_llm/_torch/auto_deploy/shim/interface.py index 30ec0b8ae616..78b90c2458ae 100644 --- a/tensorrt_llm/_torch/auto_deploy/shim/interface.py +++ b/tensorrt_llm/_torch/auto_deploy/shim/interface.py @@ -904,7 +904,13 @@ def _create_and_assign_state_views( self._caches[buf_name] = global_tensor if replay_work_items: - self._replay_work_items = torch.empty( + # Zero-init (not torch.empty): CUDA graph capture/warmup runs the + # replay SSM kernel before prepare_replay_metadata() populates this + # buffer, so uninitialized values would feed garbage cache slots to + # the kernel and trigger an out-of-bounds state-cache access (IMA). + # Zeros map every work item to slot 0 / pnat 0, which is in-bounds. + # Mirrors the PyTorch backend (mamba2_metadata.py). + self._replay_work_items = torch.zeros( self.info.max_num_state_slots, REPLAY_WORK_ITEM_WIDTH, device=self.info.device, diff --git a/tests/integration/defs/accuracy/test_llm_api_autodeploy.py b/tests/integration/defs/accuracy/test_llm_api_autodeploy.py index 902f7addf53a..420316ed193c 100644 --- a/tests/integration/defs/accuracy/test_llm_api_autodeploy.py +++ b/tests/integration/defs/accuracy/test_llm_api_autodeploy.py @@ -825,10 +825,7 @@ def test_mtp(self, world_size, attn_backend, model_id): task = GSM8K(self.MODEL_NAME) task.evaluate(llm) - # bf16 acceptance is stable; fp8/nvfp4 have higher variance due to - # arithmetic rounding, so use a lower threshold for quantized models. - min_rate = 0.50 if model_id == "bf16" else 0.40 - self.check_acceptance_rate(llm, min_acceptance_rate=min_rate) + self.check_acceptance_rate(llm, min_acceptance_rate=0.50) print_memory_usage("after evaluation") diff --git a/tests/integration/test_lists/test-db/l0_dgx_b200.yml b/tests/integration/test_lists/test-db/l0_dgx_b200.yml index 879311d807df..e060f8a8f8bc 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_b200.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_b200.yml @@ -408,5 +408,6 @@ l0_dgx_b200: - accuracy/test_llm_api_autodeploy.py::TestModelRegistryAccuracy::test_autodeploy_from_registry[deepseek-ai_DeepSeek-R1-0528-True] - accuracy/test_llm_api_autodeploy.py::TestQwen3_5_397B_MoE::test_nvfp4[8] - accuracy/test_llm_api_autodeploy.py::TestNemotronUltraV3::test_accuracy[nvfp4-8] + - accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_mtp[nvfp4_ws8_80gb-trtllm] # ------------- AutoDeploy Perf Sanity --------------- - perf/test_perf_sanity.py::test_e2e[aggr_upload-deepseek_r1_fp8_ad_blackwell-r1_fp8_ad_ws8_1k1k] TIMEOUT (120) diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index afb767029b1d..ab16d4c38e78 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -6,7 +6,6 @@ accuracy/test_llm_api.py::TestLlama3_1_8BInstruct::test_guided_decoding_4gpus[xg accuracy/test_llm_api_autodeploy.py::TestNemotronNanoV3::test_accuracy[bf16-4-attn_dp_off-trtllm] SKIP (https://nvbugs/6367792) accuracy/test_llm_api_autodeploy.py::TestNemotronNanoV3::test_accuracy[fp8-4-attn_dp_off-trtllm] SKIP (https://nvbugs/6367792) accuracy/test_llm_api_autodeploy.py::TestNemotronNanoV3::test_accuracy[nvfp4-4-attn_dp_off-trtllm] SKIP (https://nvbugs/6367792) -accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_mtp[fp8_ws4_80gb-trtllm] SKIP (https://nvbugs/6336682) accuracy/test_llm_api_pytorch.py::TestDeepSeekR1::test_fp8_blockscale[throughput_mtp] SKIP (https://nvbugs/6281818) accuracy/test_llm_api_pytorch.py::TestDeepSeekR1::test_fp8_blockscale[throughput_mtp_trtllm] SKIP (https://nvbugs/6281818) accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_nvfp4_multi_gpus_chunked_prefill[latency] SKIP (https://nvbugs/6276981) @@ -157,20 +156,15 @@ full:B300/accuracy/test_llm_api_pytorch.py::TestKimiK25::test_nvfp4[ep8] SKIP (h full:B300/disaggregated/test_disaggregated.py::test_disaggregated_ctxpp2_genpp2[TinyLlama-1.1B-Chat-v1.0] SKIP (https://nvbugs/6322073) full:B300/unittest/_torch/modules/moe/test_moe_backend.py::test_moe_backend -k "TRTLLM" SKIP (https://nvbugs/6165866) full:DGX_B200/unittest/_torch/modules/moe/test_moe_backend.py::test_moe_backend -k "TRTLLM" SKIP (https://nvbugs/6165866) -full:DGX_H100/accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_mtp[fp8_ws4_80gb-trtllm] SKIP (https://nvbugs/6336682) full:GB200/accuracy/test_dwdp_disaggregated_serving.py::TestDwdpDeepSeekV3Lite::test_dwdp_accuracy SKIP (https://nvbugs/6276923) full:GB200/accuracy/test_dwdp_disaggregated_serving.py::TestDwdpDeepSeekV3Lite::test_dwdp_accuracy_contention_opt SKIP (https://nvbugs/6276923) full:GB200/accuracy/test_dwdp_disaggregated_serving.py::TestDwdpDeepSeekV3Lite::test_dwdp_accuracy_mode_b_overlap SKIP (https://nvbugs/6276923) -full:GB200/accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_mtp[fp8_ws4_80gb-trtllm] SKIP (https://nvbugs/6316981) full:GB200/accuracy/test_llm_api_pytorch.py::TestQwen3_5_35B_A3B::test_fp8_moe_dflash SKIP (https://nvbugs/6316985) full:GB200/accuracy/test_llm_api_pytorch.py::TestQwen3_5_4B::test_dflash SKIP (https://nvbugs/6344883) full:GB200/accuracy/test_llm_api_pytorch_multimodal.py::TestQwen2_5_VL_7B::test_auto_dtype SKIP (https://nvbugs/6316983) full:GB200/disaggregated/test_ad_disagg.py::test_async_eagle3_full_model_handoff SKIP (https://nvbugs/6369254) full:GB300/accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_auto_dtype[True-False-False-False] SKIP (https://nvbugs/6316984) full:GB300/accuracy/test_llm_api_autodeploy.py::TestNemotronNanoV3::test_accuracy[nvfp4-1-attn_dp_off-trtllm] SKIP (https://nvbugs/6329165) -full:GB300/accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_mtp[bf16_ws4_180gb-trtllm] SKIP (https://nvbugs/6316981) -full:GB300/accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_mtp[fp8_ws4_80gb-trtllm] SKIP (https://nvbugs/6316981) -full:GB300/accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_mtp[nvfp4_ws4_80gb-trtllm] SKIP (https://nvbugs/6316981) full:GB300/accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_4gpus[v2_kv_cache-ep4-trtllm-fp8] SKIP (https://nvbugs/6316980) full:GB300/accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_4gpus[v2_kv_cache_no_reuse-tp4-trtllm-fp8] SKIP (https://nvbugs/6316980) full:GB300/accuracy/test_llm_api_pytorch.py::TestQwen3_5_35B_A3B::test_fp8_moe_dflash SKIP (https://nvbugs/6316985) diff --git a/tests/unittest/auto_deploy/singlegpu/custom_ops/mamba/test_flashinfer_mamba_cached_op.py b/tests/unittest/auto_deploy/singlegpu/custom_ops/mamba/test_flashinfer_mamba_cached_op.py index 89beba153c45..026cf49fc678 100644 --- a/tests/unittest/auto_deploy/singlegpu/custom_ops/mamba/test_flashinfer_mamba_cached_op.py +++ b/tests/unittest/auto_deploy/singlegpu/custom_ops/mamba/test_flashinfer_mamba_cached_op.py @@ -19,7 +19,22 @@ from test_triton_mamba_cached_op import _random_params import tensorrt_llm._torch.auto_deploy # noqa: F401 -from tensorrt_llm._torch.auto_deploy.custom_ops.attention_interface import BatchInfo +from tensorrt_llm._torch.auto_deploy._compat import KvCacheConfig +from tensorrt_llm._torch.auto_deploy.custom_ops.attention_interface import ( + BatchInfo, + CausalConvResourceHandler, + IntermediateConvStateHandler, + ReplayCacheBufIdxHandler, + ReplayNWritesHandler, + ReplayOldBHandler, + ReplayOldDAcumsumHandler, + ReplayOldDtHandler, + ReplayOldXHandler, + ReplayPrevNumAcceptedHandler, + ReplayWorkItemsHandler, + SSMResourceHandler, +) +from tensorrt_llm._torch.auto_deploy.shim.interface import CachedSequenceInterface from tensorrt_llm._torch.modules.mamba.mamba2_metadata import ( REPLAY_WORK_CACHE_BUF_IDX, REPLAY_WORK_CACHE_SLOT, @@ -27,6 +42,7 @@ REPLAY_WORK_PNAT, REPLAY_WORK_POSITION_IN_DECODE_BATCH, ) +from tensorrt_llm._torch.modules.mamba.replay_selective_state_update import _resolve_tuning from tensorrt_llm._torch.modules.mamba.replay_selective_state_update import ( replay_selective_state_update as _real_replay_selective_state_update, ) @@ -265,3 +281,228 @@ def test_flashinfer_extend_replay_calls_replay_kernel(mamba_env, head_dim): ) assert out.shape == hidden_states.shape assert torch.isfinite(out).all() + + +class _SpecDecModeForReplayTest: + def use_one_engine(self): + return True + + +class _SpecConfigForReplayTest: + def __init__(self, max_draft_len: int): + self.max_draft_len = max_draft_len + self.tokens_per_gen_step = max_draft_len + 1 + self.spec_dec_mode = _SpecDecModeForReplayTest() + + +def _build_interface_with_replay_buffers(num_heads, head_dim, d_state, n_groups, max_batch_size): + """Allocate replay buffers through the real production path (CachedSequenceInterface). + + Registers the Mamba + replay-buffer resource bundle for one layer and runs + initialize_resources(), which is where the cache-manager-bound replay + work-items buffer (interface._replay_work_items -- the tensor the replay SSM + kernel actually reads) is allocated. + """ + conv_dim = head_dim * num_heads + 2 * n_groups * d_state + interface = CachedSequenceInterface( + max_seq_len=128, + max_batch_size=max_batch_size, + max_num_tokens=(128 + 1) * max_batch_size, + device="cuda", + kv_cache_config=KvCacheConfig( + tokens_per_block=32, max_tokens=1024, free_gpu_memory_fraction=0.0 + ), + spec_config=_SpecConfigForReplayTest(max_draft_len=2), + ) + interface.add_resource( + "ssm_state_0", + SSMResourceHandler( + num_heads=num_heads, head_dim=head_dim, d_state=d_state, dtype=torch.bfloat16 + ), + ) + interface.add_resource( + "conv_state_0", CausalConvResourceHandler(conv_dim=conv_dim, d_conv=4, dtype=torch.float32) + ) + interface.add_resource( + "intermediate_conv_state_0", + IntermediateConvStateHandler(conv_dim=conv_dim, d_conv=4, dtype=torch.float32), + ) + interface.add_resource( + "replay_old_x_0", + ReplayOldXHandler(num_heads=num_heads, head_dim=head_dim, dtype=torch.bfloat16), + ) + interface.add_resource( + "replay_old_B_0", + ReplayOldBHandler(n_groups=n_groups, d_state=d_state, dtype=torch.bfloat16), + ) + interface.add_resource("replay_old_dt_0", ReplayOldDtHandler(num_heads=num_heads)) + interface.add_resource("replay_old_dA_cumsum_0", ReplayOldDAcumsumHandler(num_heads=num_heads)) + interface.add_resource("replay_cache_buf_idx_0", ReplayCacheBufIdxHandler()) + interface.add_resource("replay_prev_num_accepted_0", ReplayPrevNumAcceptedHandler()) + interface.add_resource("replay_work_items_0", ReplayWorkItemsHandler()) + interface.add_resource("replay_n_writes_0", ReplayNWritesHandler()) + return interface + + +def test_extend_replay_no_ima_with_capture_time_work_items(mamba_env, monkeypatch): + """The replay SSM op must not fault when its work-items buffer is unprepared. + + Reproduces the SuperV3 MTP illegal-memory-access (PR #14203 regression): during + CUDA graph capture/warmup the replay SSM kernel runs *before* + prepare_replay_metadata() has populated the work-items buffer, so the buffer + holds only whatever its allocation produced. The kernel reads the per-row + cache_slot field and indexes the SSM state cache with it -- unchecked -- so a + garbage slot index is an out-of-bounds access (IMA). + + This drives the real op (flashinfer_cached_ssm -> replay_selective_state_update) + with the work-items buffer allocated through the production path + (CachedSequenceInterface), exactly as at capture time (no prepare_replay_metadata + call), and asserts the kernel completes with no CUDA fault. + + Determinism: torch.empty is monkeypatched to poison the work-items buffer with + an out-of-bounds slot index, mirroring the garbage that triggered the original + IMA. With the fix (the buffer is allocated via torch.zeros) the poison does not + apply and every slot is in-bounds; a regression to torch.empty yields the + poisoned out-of-bounds slot and this test faults. + """ + device = mamba_env["device"] + dtype = mamba_env["dtype"] + + # Production SuperV3 Mamba2 shape (AutoDeploy replicates mamba -> full heads/groups). + # The replay kernel selects its mode from effective batch = raw_batch * nheads; this + # config (8 * 128 = 1024) lands in persistent_main, which reads the cache slot from + # the work-items buffer (persistent_dynamic, used at small batch, reads it from + # state_batch_indices instead and would not exercise the bug). Guard the mode below. + num_extend = 8 + tokens_per_extend = 7 # num_nextn_predict_layers (6) + 1 + num_heads = 128 + head_dim = 64 + n_groups, ssm_state_size = 8, 128 + max_batch_size = 8 + + mode, _ = _resolve_tuning(num_extend, num_heads, "bf16", "RN") + assert mode == "persistent_main", ( + f"test must exercise persistent_main (reads work_items); resolved {mode!r}. " + "Re-size num_extend/num_heads if the tuning table changed." + ) + + interface = _build_interface_with_replay_buffers( + num_heads, head_dim, ssm_state_size, n_groups, max_batch_size + ) + n_slots = interface.info.max_num_state_slots + + # Simulate capture-time allocator garbage deterministically: poison only the + # work-items buffer shape so a torch.empty regression produces an out-of-bounds + # cache slot. The torch.zeros fix is not intercepted and stays in-bounds. + real_empty = torch.empty + + def poisoned_empty(*args, **kwargs): + t = real_empty(*args, **kwargs) + if ( + kwargs.get("dtype") == torch.int32 + and len(args) == 2 + and args[0] == n_slots + and args[1] == REPLAY_WORK_ITEM_WIDTH + ): + t.fill_(0x7BADBEEF) # out-of-bounds cache slot (cache has max_batch_size slots) + return t + + monkeypatch.setattr(torch, "empty", poisoned_empty) + + interface.initialize_resources() + + # The buffers the kernel reads, exactly as bound at capture time. Crucially we do + # NOT call prepare_replay_metadata(), matching the warmup/capture path. + replay_work_items = interface._replay_work_items + replay_n_writes = interface._replay_n_writes + assert replay_work_items is not None and replay_n_writes is not None + + # Per-token inputs and the remaining replay caches (independent of the buffer + # under test) are built directly, as in test_flashinfer_extend_replay_calls_replay_kernel. + (hidden_states, A, B, C, D, dt, dt_bias, time_step_limit, chunk_size) = _random_params( + device, dtype, num_extend, tokens_per_extend, num_heads, head_dim, n_groups, ssm_state_size + ) + ssm_state_cache = torch.zeros( + max_batch_size, num_heads, head_dim, ssm_state_size, device=device, dtype=dtype + ) + # One distinct, in-bounds cache slot per extend sequence. + slot_idx = torch.arange(num_extend, device=device, dtype=torch.int32) + + replay_history_size = 16 + replay_old_x = torch.zeros( + max_batch_size, + 2, + replay_history_size, + num_heads, + head_dim, + device=device, + dtype=torch.bfloat16, + ) + replay_old_b = torch.zeros( + max_batch_size, + 2, + replay_history_size, + n_groups, + ssm_state_size, + device=device, + dtype=torch.bfloat16, + ) + replay_old_dt = torch.zeros( + max_batch_size, 2, num_heads, replay_history_size, device=device, dtype=torch.float32 + ) + replay_old_da_cumsum = torch.zeros( + max_batch_size, 2, num_heads, replay_history_size, device=device, dtype=torch.float32 + ) + replay_cache_buf_idx = torch.zeros(max_batch_size, device=device, dtype=torch.int32) + replay_prev_num_accepted = torch.zeros(max_batch_size, device=device, dtype=torch.int32) + + _bi = BatchInfo() + _bi.update([0, 0, num_extend, num_extend * tokens_per_extend, 0, 0]) + _bi.update_use_replay(True) + batch_info_host = _bi.serialize() + cu_seqlen = torch.arange( + 0, (num_extend + 1) * tokens_per_extend, tokens_per_extend, device=device, dtype=torch.int32 + ) + use_initial_states = torch.zeros(num_extend, device=device, dtype=torch.bool) + any_prefill_use_initial_states_host = torch.tensor([False], device=device, dtype=torch.bool) + + out = torch.ops.auto_deploy.flashinfer_cached_ssm( + hidden_states, + A, + B, + C, + D, + dt, + dt_bias, + # STANDARD METADATA + batch_info_host, + cu_seqlen, + slot_idx, + use_initial_states, + any_prefill_use_initial_states_host, + # EXTRA METADATA + None, + None, + None, # chunk_indices, chunk_offsets, seq_idx_prefill + # CACHES + ssm_state_cache, + None, # intermediate_ssm_state_cache (None in replay mode) + replay_old_x, + replay_old_b, + replay_old_dt, + replay_old_da_cumsum, + replay_cache_buf_idx, + replay_prev_num_accepted, + replay_work_items, + replay_n_writes, + # CONSTANTS + time_step_limit, + chunk_size, + ) + + # Force kernel completion: a torch.empty regression feeds an out-of-bounds cache + # slot to the kernel here and raises a CUDA illegal memory access. The torch.zeros + # fix keeps every slot in-bounds, so this completes cleanly. + torch.cuda.synchronize() + assert out.shape == hidden_states.shape + assert torch.isfinite(out).all() From 77d1245a0ab99666103ecd801b564711c807d2dc Mon Sep 17 00:00:00 2001 From: Gal Hubara Agam <96368689+galagam@users.noreply.github.com> Date: Wed, 24 Jun 2026 22:35:24 -0700 Subject: [PATCH 2/3] [https://nvbugs/6316981][fix] AutoDeploy: prepare MTP replay metadata during cudagraph capture Make replay-metadata preparation symmetric across runtime and cudagraph capture by registering prepare_replay_metadata as a host-prepare hook, so it runs inside nest_sequences() on every path -- exactly like all other AD attention metadata. Previously prepare_replay_metadata was a side call only in ADEngine._prepare_inputs (runtime). The cudagraph capture path (SequenceInfo.set_capture_batch) synthesizes every other metadata buffer (slot_idx, KV page metadata, in-graph prepare ops) but never invoked it, so replay_work_items was the lone metadata buffer left unprepared at capture -- the root cause of the SuperV3 MTP capture-time IMA. The hook is registered only when the replay buffers were actually allocated (ssm_replay enabled), so non-replay models pay no per-forward cost and never touch the unallocated buffers. The explicit runtime call is removed (the hook now covers it). The zero-init of replay_work_items is kept as defense-in-depth (reframed in comment), mirroring the PyTorch backend which also keeps the zeros guard alongside its per-forward prepare. Verified on 4xGB200: with the buffer deliberately poisoned at allocation, the hook overwrites it at capture and every runtime step -- test_mtp[nvfp4_ws4_80gb-trtllm] passes (acceptance 54.5%, no IMA), confirming zero-init is now optional. Signed-off-by: Gal Hubara Agam <96368689+galagam@users.noreply.github.com> --- .../_torch/auto_deploy/shim/ad_executor.py | 1 - .../_torch/auto_deploy/shim/interface.py | 20 +++- .../mamba/test_flashinfer_mamba_cached_op.py | 108 ++++++------------ 3 files changed, 48 insertions(+), 81 deletions(-) diff --git a/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py b/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py index 9e3387d26cb3..b272740deab4 100644 --- a/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py +++ b/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py @@ -993,7 +993,6 @@ def _prepare_inputs( _ungathered_new_lens=new_tokens_lens, **extra_args, ) - self.cache_seq_interface.prepare_replay_metadata() self.iter_states["num_ctx_requests"] = num_prefill self.iter_states["num_ctx_tokens"] = num_prefill_tokens diff --git a/tensorrt_llm/_torch/auto_deploy/shim/interface.py b/tensorrt_llm/_torch/auto_deploy/shim/interface.py index 78b90c2458ae..48ca5a3bb78f 100644 --- a/tensorrt_llm/_torch/auto_deploy/shim/interface.py +++ b/tensorrt_llm/_torch/auto_deploy/shim/interface.py @@ -904,12 +904,11 @@ def _create_and_assign_state_views( self._caches[buf_name] = global_tensor if replay_work_items: - # Zero-init (not torch.empty): CUDA graph capture/warmup runs the - # replay SSM kernel before prepare_replay_metadata() populates this - # buffer, so uninitialized values would feed garbage cache slots to - # the kernel and trigger an out-of-bounds state-cache access (IMA). - # Zeros map every work item to slot 0 / pnat 0, which is in-bounds. - # Mirrors the PyTorch backend (mamba2_metadata.py). + # Zero-init as an extra precaution (not torch.empty). The + # prepare_replay_metadata host-prepare hook (registered in + # initialize_resources) populates this buffer on every + # nest_sequences -- runtime and cudagraph capture alike -- so the + # replay SSM kernel never reads it unprepared. self._replay_work_items = torch.zeros( self.info.max_num_state_slots, REPLAY_WORK_ITEM_WIDTH, @@ -1423,6 +1422,15 @@ def initialize_resources(self) -> int: f"max_tokens={s['max_tokens']}" ) + if self._replay_work_items is not None: + # Wrapper takes **kwargs to satisfy the host-prepare callable protocol; + # prepare_replay_metadata reads everything it needs from self, so no + # graph-input args are requested (empty arg list). + def _replay_metadata_hook(**_sequence_info_args) -> None: + self.prepare_replay_metadata() + + self.info.register_host_prepare_for_attention_forward(_replay_metadata_hook, []) + return len(self._caches) def _requires_token_estimate(self) -> bool: diff --git a/tests/unittest/auto_deploy/singlegpu/custom_ops/mamba/test_flashinfer_mamba_cached_op.py b/tests/unittest/auto_deploy/singlegpu/custom_ops/mamba/test_flashinfer_mamba_cached_op.py index 026cf49fc678..de8f614ca773 100644 --- a/tests/unittest/auto_deploy/singlegpu/custom_ops/mamba/test_flashinfer_mamba_cached_op.py +++ b/tests/unittest/auto_deploy/singlegpu/custom_ops/mamba/test_flashinfer_mamba_cached_op.py @@ -42,7 +42,6 @@ REPLAY_WORK_PNAT, REPLAY_WORK_POSITION_IN_DECODE_BATCH, ) -from tensorrt_llm._torch.modules.mamba.replay_selective_state_update import _resolve_tuning from tensorrt_llm._torch.modules.mamba.replay_selective_state_update import ( replay_selective_state_update as _real_replay_selective_state_update, ) @@ -344,102 +343,64 @@ def _build_interface_with_replay_buffers(num_heads, head_dim, d_state, n_groups, return interface -def test_extend_replay_no_ima_with_capture_time_work_items(mamba_env, monkeypatch): - """The replay SSM op must not fault when its work-items buffer is unprepared. +def test_extend_replay_no_ima(mamba_env): + """The replay path must not cause an out-of-bounds access on the replay buffers. - Reproduces the SuperV3 MTP illegal-memory-access (PR #14203 regression): during - CUDA graph capture/warmup the replay SSM kernel runs *before* - prepare_replay_metadata() has populated the work-items buffer, so the buffer - holds only whatever its allocation produced. The kernel reads the per-row - cache_slot field and indexes the SSM state cache with it -- unchecked -- so a - garbage slot index is an out-of-bounds access (IMA). + Behavioral guard for the replay path: the replay work-items buffer is filled with + garbage (an out-of-bounds cache slot, simulating uninitialized memory), then the + production metadata-prep path runs and the real replay op executes; the test asserts + no CUDA fault. With the fix, prep populates the buffer before the kernel reads it, so + the garbage never reaches the kernel; without it the out-of-bounds slot survives and + faults. - This drives the real op (flashinfer_cached_ssm -> replay_selective_state_update) - with the work-items buffer allocated through the production path - (CachedSequenceInterface), exactly as at capture time (no prepare_replay_metadata - call), and asserts the kernel completes with no CUDA fault. - - Determinism: torch.empty is monkeypatched to poison the work-items buffer with - an out-of-bounds slot index, mirroring the garbage that triggered the original - IMA. With the fix (the buffer is allocated via torch.zeros) the poison does not - apply and every slot is in-bounds; a regression to torch.empty yields the - poisoned out-of-bounds slot and this test faults. + Filling the buffer directly keeps the poison confined to this one buffer and makes + the failure deterministic: fresh CUDA memory is often benign, so a poison-free run + cannot reliably reproduce the bug. """ device = mamba_env["device"] dtype = mamba_env["dtype"] - # Production SuperV3 Mamba2 shape (AutoDeploy replicates mamba -> full heads/groups). - # The replay kernel selects its mode from effective batch = raw_batch * nheads; this - # config (8 * 128 = 1024) lands in persistent_main, which reads the cache slot from - # the work-items buffer (persistent_dynamic, used at small batch, reads it from - # state_batch_indices instead and would not exercise the bug). Guard the mode below. + # Production SuperV3 Mamba2 shape (AutoDeploy replicates mamba -> full heads/groups), + # large enough that the replay kernel runs its persistent_main path, which reads the + # cache slot from the replay work-items buffer. num_extend = 8 tokens_per_extend = 7 # num_nextn_predict_layers (6) + 1 num_heads = 128 head_dim = 64 n_groups, ssm_state_size = 8, 128 - max_batch_size = 8 - - mode, _ = _resolve_tuning(num_extend, num_heads, "bf16", "RN") - assert mode == "persistent_main", ( - f"test must exercise persistent_main (reads work_items); resolved {mode!r}. " - "Re-size num_extend/num_heads if the tuning table changed." - ) interface = _build_interface_with_replay_buffers( - num_heads, head_dim, ssm_state_size, n_groups, max_batch_size + num_heads, head_dim, ssm_state_size, n_groups, max_batch_size=num_extend ) - n_slots = interface.info.max_num_state_slots - - # Simulate capture-time allocator garbage deterministically: poison only the - # work-items buffer shape so a torch.empty regression produces an out-of-bounds - # cache slot. The torch.zeros fix is not intercepted and stays in-bounds. - real_empty = torch.empty - - def poisoned_empty(*args, **kwargs): - t = real_empty(*args, **kwargs) - if ( - kwargs.get("dtype") == torch.int32 - and len(args) == 2 - and args[0] == n_slots - and args[1] == REPLAY_WORK_ITEM_WIDTH - ): - t.fill_(0x7BADBEEF) # out-of-bounds cache slot (cache has max_batch_size slots) - return t - - monkeypatch.setattr(torch, "empty", poisoned_empty) - interface.initialize_resources() - # The buffers the kernel reads, exactly as bound at capture time. Crucially we do - # NOT call prepare_replay_metadata(), matching the warmup/capture path. + # Fill the work-items buffer with an out-of-bounds cache slot, simulating garbage / + # uninitialized memory. The production metadata-prep below must overwrite it before + # the kernel reads it; if prep is missing (the bug) this poison survives and faults. + interface._replay_work_items.fill_(0x7FFFFFFF) # int32-max: out-of-bounds slot + + # Drive the production metadata-prep path -- the same one cudagraph capture uses -- + # so the replay work-items / n-writes buffers are populated exactly as in real runs + # (set_capture_batch -> nest_sequences -> prepare_replay_metadata host-prepare hook). + interface.info.set_capture_batch(max_draft_len=tokens_per_extend - 1, batch_size=num_extend) replay_work_items = interface._replay_work_items replay_n_writes = interface._replay_n_writes - assert replay_work_items is not None and replay_n_writes is not None - # Per-token inputs and the remaining replay caches (independent of the buffer - # under test) are built directly, as in test_flashinfer_extend_replay_calls_replay_kernel. + # Per-token inputs and the remaining replay caches for the same extend batch. (hidden_states, A, B, C, D, dt, dt_bias, time_step_limit, chunk_size) = _random_params( device, dtype, num_extend, tokens_per_extend, num_heads, head_dim, n_groups, ssm_state_size ) ssm_state_cache = torch.zeros( - max_batch_size, num_heads, head_dim, ssm_state_size, device=device, dtype=dtype + num_extend, num_heads, head_dim, ssm_state_size, device=device, dtype=dtype ) - # One distinct, in-bounds cache slot per extend sequence. slot_idx = torch.arange(num_extend, device=device, dtype=torch.int32) replay_history_size = 16 replay_old_x = torch.zeros( - max_batch_size, - 2, - replay_history_size, - num_heads, - head_dim, - device=device, - dtype=torch.bfloat16, + num_extend, 2, replay_history_size, num_heads, head_dim, device=device, dtype=torch.bfloat16 ) replay_old_b = torch.zeros( - max_batch_size, + num_extend, 2, replay_history_size, n_groups, @@ -448,13 +409,13 @@ def poisoned_empty(*args, **kwargs): dtype=torch.bfloat16, ) replay_old_dt = torch.zeros( - max_batch_size, 2, num_heads, replay_history_size, device=device, dtype=torch.float32 + num_extend, 2, num_heads, replay_history_size, device=device, dtype=torch.float32 ) replay_old_da_cumsum = torch.zeros( - max_batch_size, 2, num_heads, replay_history_size, device=device, dtype=torch.float32 + num_extend, 2, num_heads, replay_history_size, device=device, dtype=torch.float32 ) - replay_cache_buf_idx = torch.zeros(max_batch_size, device=device, dtype=torch.int32) - replay_prev_num_accepted = torch.zeros(max_batch_size, device=device, dtype=torch.int32) + replay_cache_buf_idx = torch.zeros(num_extend, device=device, dtype=torch.int32) + replay_prev_num_accepted = torch.zeros(num_extend, device=device, dtype=torch.int32) _bi = BatchInfo() _bi.update([0, 0, num_extend, num_extend * tokens_per_extend, 0, 0]) @@ -500,9 +461,8 @@ def poisoned_empty(*args, **kwargs): chunk_size, ) - # Force kernel completion: a torch.empty regression feeds an out-of-bounds cache - # slot to the kernel here and raises a CUDA illegal memory access. The torch.zeros - # fix keeps every slot in-bounds, so this completes cleanly. + # Synchronize so any out-of-bounds access on the replay buffers surfaces here as a + # CUDA error rather than asynchronously later. torch.cuda.synchronize() assert out.shape == hidden_states.shape assert torch.isfinite(out).all() From 27f1dd78e7ac06964170f7be32698584024a9f0c Mon Sep 17 00:00:00 2001 From: Gal Hubara Agam <96368689+galagam@users.noreply.github.com> Date: Thu, 25 Jun 2026 02:06:10 -0700 Subject: [PATCH 3/3] refactor Signed-off-by: Gal Hubara Agam <96368689+galagam@users.noreply.github.com> --- .../_torch/auto_deploy/shim/interface.py | 2 +- .../mamba/test_flashinfer_mamba_cached_op.py | 30 ++++++++++--------- 2 files changed, 17 insertions(+), 15 deletions(-) diff --git a/tensorrt_llm/_torch/auto_deploy/shim/interface.py b/tensorrt_llm/_torch/auto_deploy/shim/interface.py index 48ca5a3bb78f..d3250137e532 100644 --- a/tensorrt_llm/_torch/auto_deploy/shim/interface.py +++ b/tensorrt_llm/_torch/auto_deploy/shim/interface.py @@ -1422,7 +1422,7 @@ def initialize_resources(self) -> int: f"max_tokens={s['max_tokens']}" ) - if self._replay_work_items is not None: + if self.info.batch_info.is_use_replay(): # Wrapper takes **kwargs to satisfy the host-prepare callable protocol; # prepare_replay_metadata reads everything it needs from self, so no # graph-input args are requested (empty arg list). diff --git a/tests/unittest/auto_deploy/singlegpu/custom_ops/mamba/test_flashinfer_mamba_cached_op.py b/tests/unittest/auto_deploy/singlegpu/custom_ops/mamba/test_flashinfer_mamba_cached_op.py index de8f614ca773..f88d89728734 100644 --- a/tests/unittest/auto_deploy/singlegpu/custom_ops/mamba/test_flashinfer_mamba_cached_op.py +++ b/tests/unittest/auto_deploy/singlegpu/custom_ops/mamba/test_flashinfer_mamba_cached_op.py @@ -343,19 +343,19 @@ def _build_interface_with_replay_buffers(num_heads, head_dim, d_state, n_groups, return interface -def test_extend_replay_no_ima(mamba_env): +def test_extend_replay_init_buffers(mamba_env): """The replay path must not cause an out-of-bounds access on the replay buffers. - Behavioral guard for the replay path: the replay work-items buffer is filled with - garbage (an out-of-bounds cache slot, simulating uninitialized memory), then the - production metadata-prep path runs and the real replay op executes; the test asserts - no CUDA fault. With the fix, prep populates the buffer before the kernel reads it, so - the garbage never reaches the kernel; without it the out-of-bounds slot survives and - faults. + Behavioral guard for the replay path: every buffer the prepare hook populates (the + work-items buffer and the n-writes count) is filled with garbage (out-of-bounds + values, simulating uninitialized memory), then the production metadata-prep path runs + and the real replay op executes; the test asserts no CUDA fault. With the fix, prep + populates the buffers before the kernel reads them, so the garbage never reaches the + kernel; without it the out-of-bounds values survive and fault. - Filling the buffer directly keeps the poison confined to this one buffer and makes - the failure deterministic: fresh CUDA memory is often benign, so a poison-free run - cannot reliably reproduce the bug. + Filling the buffers directly keeps the poison confined to them and makes the failure + deterministic: fresh CUDA memory is often benign, so a poison-free run cannot + reliably reproduce the bug. """ device = mamba_env["device"] dtype = mamba_env["dtype"] @@ -374,10 +374,12 @@ def test_extend_replay_no_ima(mamba_env): ) interface.initialize_resources() - # Fill the work-items buffer with an out-of-bounds cache slot, simulating garbage / - # uninitialized memory. The production metadata-prep below must overwrite it before - # the kernel reads it; if prep is missing (the bug) this poison survives and faults. - interface._replay_work_items.fill_(0x7FFFFFFF) # int32-max: out-of-bounds slot + # Poison every buffer the prepare hook populates -- the work-items buffer and the + # n-writes count -- with out-of-bounds values, simulating garbage / uninitialized + # memory. The production metadata-prep below must overwrite them before the kernel + # reads them; if prep is missing (the bug) the poison survives and faults. + interface._replay_work_items.fill_(0x7FFFFFFF) # int32-max: out-of-bounds cache slot + interface._replay_n_writes.fill_(0x7FFFFFFF) # int32-max: out-of-bounds write count # Drive the production metadata-prep path -- the same one cudagraph capture uses -- # so the replay work-items / n-writes buffers are populated exactly as in real runs