From 60ad65f659ebc2575a71812a8a325489267a4cf3 Mon Sep 17 00:00:00 2001 From: Gal Hubara Agam <96368689+galagam@users.noreply.github.com> Date: Mon, 4 May 2026 08:47:18 -0700 Subject: [PATCH 01/17] [None][perf] AutoDeploy: replay SSM kernel for MTP extend path MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Integrate replay_selective_state_update (tl.dot fast-forward) into the AD flashinfer_ssm backend as an optional path controlled by: insert_cached_ssm_attention: backend: flashinfer_ssm ssm_replay: true The replay kernel replaces the FlashInfer horizontal SSM scan in the MTP extend path with the same two-kernel algorithm PT uses: - _replay_precompute_kernel: computes CB_scaled/decay_vec from cached B,dt,A - _replay_state_update_kernel: fast-forwards SSM state via tl.dot on cached values, then computes output using new x/C and precomputed tensors Performance (B200, ws4, NVFP4, conc=32, ISL/OSL=1k): SSM Other: 3.51ms -> 2.05ms replay_state_update_kernel: 1.72ms (vs PT target 0.414ms) Total iter: 20.58ms -> 18.19ms (vs PT target 13.47ms) While this is an improvement over existing kernel, it still needs tuning to reach the full potential. ── Infrastructure change: get_kwarg_cache_keys() ── The AD transform wires custom-op arguments positionally via call_function(args=(*qkv, *meta, *cache_in_nodes, *constants)) Relying solely on _InsertCachedOperator's existing positional routing — always including all resource nodes in cache_in_nodes — would add 7 dead placeholder nodes to every flashinfer_ssm graph, even for non-MTP engines where replay is disabled. Instead, replay args (and intermediate_ssm_state_cache) are routed via a new descriptor hook: get_kwarg_cache_keys(). Keys in that set are emitted as call_function(kwargs={key: node}) rather than appended to cache_in_nodes, keeping non-replay graphs free of unused graph inputs. In replay mode, intermediate_ssm_state_cache is simply absent from get_cache_initializers — it uses the function's default (None) and no graph node is created. In non-replay mode it appears as a kwarg with the real spec tensor, same as the six replay cache args. ── Testing ── Unit smoke (tests/unittest/auto_deploy/singlegpu/smoke/test_ad_speculative_decoding.py): - test_super_mtp_smoke: existing triton_ssm path (unchanged, all pass) - test_super_mtp_ssm_replay_smoke: flashinfer_ssm + ssm_replay=True with mamba_head_dim=64 / ssm_state_size=64 - (~57s) Accuracy (B200, ws4_180gb, trtllm backend): nvfp4: PASSED (GSM8K ~92%, acceptance_rate ~45%) fp8: PASSED (GSM8K ~92%, acceptance_rate ~45%) Signed-off-by: Gal Hubara Agam <96368689+galagam@users.noreply.github.com> [None][perf] AutoDeploy: skip .contiguous() in flashinfer_ssm replay path Signed-off-by: Gal Hubara Agam <96368689+galagam@users.noreply.github.com> [None][fix] AutoDeploy: fix flashinfer_cached_ssm crash in non-MTP non-replay path Running AutoDeploy with super_v3.yaml (non-MTP, flashinfer_ssm backend) crashed during setup_engine() with IndexError in PyTorch's adinplaceorview_impl: File "torch/_library/custom_ops.py", in adinplaceorview_impl increment_version(args[idx]) IndexError: tuple index out of range Root cause: the op was registered with mutates_args=("ssm_state_cache", "intermediate_ssm_state_cache", "replay_old_x", ...) PyTorch's adinplaceorview_impl computes mutated_idxs from the schema and calls args[idx] for each. intermediate_ssm_state_cache sits at schema index 18, but in the non-MTP FX graph it is routed as a kwarg (via get_kwarg_cache_keys()), so args only has 18 positional elements (indices 0-17) and args[18] raises IndexError. Fix: remove all optional kwarg-routed tensors from mutates_args, keeping only ssm_state_cache (which is always positional, at schema index 15). The removed tensors — intermediate_ssm_state_cache and the replay buffers — are inference-only caches that don't participate in autograd; their mutations are correctly handled by the CUDA kernels without needing PyTorch version tracking. Unit smoke (all 5 pass): test_super_mtp_smoke: PASS test_super_mtp_ssm_replay_smoke: PASS test_kv_cache_extra_seq_len_*: PASS test_mtp_autodeploy_uses_eagle_*: PASS test_detect_hidden_states_capture_*: PASS Signed-off-by: Gal Hubara Agam <96368689+galagam@users.noreply.github.com> [None][perf] AutoDeploy: wire PDL flags for conv1d→precompute→state_update chain Wire the kernel-level PDL flags for the full conv1d → precompute → state_update PDL chain, matching the PyTorch backend (#13453): triton_backend_causal_conv.py (extend path): causal_conv1d_update(..., launch_dependent_kernels=True) Kernel emits gdc_launch_dependents() at block start. flashinfer_backend_mamba.py (replay path): replay_selective_state_update(..., launch_with_pdl=True) Precompute kernel: gdc_wait() for conv1d signal, gdc_launch_dependents() for state_update. State_update kernel: gdc_wait() for precompute. Both functions silently disable PDL on sm < 90 (Ampere/B200 only). Signed-off-by: Gal Hubara Agam <96368689+galagam@users.noreply.github.com> [None][perf] AutoDeploy: reorder causal conv ops to minimize gap before replay precompute Profile analysis revealed that between the extend causal_conv1d_update kernel and the PDL-dependent replay precompute kernel, there is a ~15.8 µs gap containing four kernels from unrelated model components. Two those kernels (torch.arange + slot_idx.to(int32)) were from the same op body and ran after the extend conv kernel unnecessarily. Changes: 1. Pre-hoist slot_idx.to(torch.int32) to top of function body so the int conversion kernel runs before any path-specific processing. 2. Process DECODE before EXTEND so decode-related kernels (slot conversion, scatter) land before the extend conv kernel, not after. 3. Process zero_fill before EXTEND for the same reason. 4. Use copy_() for extend scatter-back: single efficient kernel instead of transpose+view+assign (which could expand to multiple kernels with some layouts). The two moved kernels (arange + slot_idx.to) now correctly appear before extend conv in the trace. The remaining two kernels (BF16 copy, SiLU) are from other model that also run on the same stream. For PDL to fire, we might require either a dedicated stream for conv+replay or a fused op. Signed-off-by: Gal Hubara Agam <96368689+galagam@users.noreply.github.com> [None][perf] AutoDeploy: fuse silu activation into Triton causal conv1d fuse_causal_conv_activation was hardcoded to match only the CUDA-backed causal conv wrapper (cuda_cached_causal_conv1d_wrapper), leaving the Triton- backed wrapper (triton_cached_causal_conv1d_wrapper) unfused. In NemotronH/ SuperV3, the SSM mixer applies a SiLU activation to the full conv output before splitting it into x/B/C: hidden_states_B_C = self.act( # ← self.act = ACT2FN["silu"] torch.ops.auto_deploy.torch_causal_conv1d(...) ) With the Triton backend this act(conv_out) remained as a separate FX node, producing a 7.0 µs silu_kernel in stream 3512 between _causal_conv1d_update_kernel and _replay_precompute_kernel — blocking the PDL chain. Fix: - Extend _match_causal_conv_activation_pattern to accept a list of target ops. - In FuseCausalConvActivation._apply, pass both the CUDA and Triton wrappers as candidates so the silu is absorbed into the conv op via the existing activation= constant argument. - Replacement now uses conv_node.target (the matched op) instead of the hardcoded CUDA wrapper, so both backends produce the correct fused call. match_swiglu_pattern and fuse_silu_mul were not applicable (0 matches): this activation is a standalone silu, not a SwiGLU silu(x)*y pair. Signed-off-by: Gal Hubara Agam <96368689+galagam@users.noreply.github.com> config adjutments: max_seq_len: 16k, max_num_tokens: 8k Signed-off-by: Gal Hubara Agam <96368689+galagam@users.noreply.github.com> --- .../model_registry/configs/super_v3_mtp.yaml | 6 +- .../custom_ops/attention_interface.py | 153 +++++++++++++ .../mamba/flashinfer_backend_mamba.py | 216 ++++++++++++++---- .../mamba/triton_backend_causal_conv.py | 76 +++--- .../_torch/auto_deploy/shim/interface.py | 131 ++++++++++- .../transform/library/fuse_causal_conv.py | 21 +- .../auto_deploy/transform/library/kvcache.py | 20 +- .../transform/library/ssm_cache.py | 39 +++- .../_torch/pyexecutor/mamba_cache_manager.py | 60 +++++ .../mamba/test_flashinfer_mamba_cached_op.py | 1 - .../smoke/test_ad_speculative_decoding.py | 64 ++++++ 11 files changed, 690 insertions(+), 97 deletions(-) diff --git a/examples/auto_deploy/model_registry/configs/super_v3_mtp.yaml b/examples/auto_deploy/model_registry/configs/super_v3_mtp.yaml index e4ea9ab6ee0a..cf1528a882d6 100644 --- a/examples/auto_deploy/model_registry/configs/super_v3_mtp.yaml +++ b/examples/auto_deploy/model_registry/configs/super_v3_mtp.yaml @@ -3,13 +3,14 @@ runtime: trtllm compile_backend: torch-cudagraph max_batch_size: 128 -max_seq_len: 65536 # tunable +max_seq_len: 16384 +max_num_tokens: 8192 enable_chunked_prefill: false # TODO: Support chunked prefill w/ spec dec attn_backend: trtllm model_factory: AutoModelForCausalLM skip_loading_weights: false cuda_graph_config: - batch_sizes: [1, 2, 4, 8, 16, 24, 32, 64, 128] + max_batch_size: 64 kv_cache_config: # tunable mamba cache dtype # --> use float32 for accuracy and default (auto) for speed @@ -54,6 +55,7 @@ transforms: # for speculative Mamba state caching in this configuration. insert_cached_ssm_attention: backend: flashinfer_ssm + ssm_replay: true insert_cached_causal_conv: backend: triton_causal_conv fuse_nvfp4_moe: diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py index 50d163209f80..2e9b26cea69f 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py @@ -1833,6 +1833,149 @@ def from_base(cls, base: Optional["SSMResourceHandler"]) -> Optional["SpecSSMRes ) +class ReplayOldXHandler(StateResourceHandler): + """Per-layer old_x cache for the replay SSM kernel (single-buffered, bf16). + + Shape: (max_batch, T, num_heads, head_dim) — T is determined by the manager's + spec_config (max_draft_len + 1), not by this handler. Acts as a type marker. + Routes to MambaHybridCacheManager via get_replay_old_x(layer_idx). + """ + + def __init__(self, num_heads: int, head_dim: int, dtype: torch.dtype) -> None: + self.num_heads = num_heads + self.head_dim = head_dim + super().__init__(dtype=dtype) + + @property + def state_shape(self) -> Tuple[int, int]: + return (self.num_heads, self.head_dim) + + def allocate(self, sequence_info) -> torch.Tensor: + raise RuntimeError("ReplayOldXHandler must be backed by MambaHybridCacheManager") + + def __eq__(self, other) -> bool: + return ( + isinstance(other, ReplayOldXHandler) + and self.num_heads == other.num_heads + and self.head_dim == other.head_dim + and self.dtype == other.dtype + ) + + +class ReplayOldBHandler(StateResourceHandler): + """Per-layer old_B cache for the replay SSM kernel (double-buffered, bf16). + + Shape: (max_batch, 2, T, n_groups, d_state) — T from manager. + Routes to MambaHybridCacheManager via get_replay_old_B(layer_idx). + """ + + def __init__(self, n_groups: int, d_state: int, dtype: torch.dtype) -> None: + self.n_groups = n_groups + self.d_state = d_state + super().__init__(dtype=dtype) + + @property + def state_shape(self) -> Tuple[int, int]: + return (self.n_groups, self.d_state) + + def allocate(self, sequence_info) -> torch.Tensor: + raise RuntimeError("ReplayOldBHandler must be backed by MambaHybridCacheManager") + + def __eq__(self, other) -> bool: + return ( + isinstance(other, ReplayOldBHandler) + and self.n_groups == other.n_groups + and self.d_state == other.d_state + and self.dtype == other.dtype + ) + + +class ReplayOldDtHandler(StateResourceHandler): + """Per-layer old_dt cache for the replay SSM kernel (double-buffered, fp32). + + Shape: (max_batch, 2, num_heads, T) — T from manager. + Routes to MambaHybridCacheManager via get_replay_old_dt(layer_idx). + """ + + def __init__(self, num_heads: int) -> None: + self.num_heads = num_heads + super().__init__(dtype=torch.float32) + + @property + def state_shape(self) -> Tuple[int]: + return (self.num_heads,) + + def allocate(self, sequence_info) -> torch.Tensor: + raise RuntimeError("ReplayOldDtHandler must be backed by MambaHybridCacheManager") + + def __eq__(self, other) -> bool: + return isinstance(other, ReplayOldDtHandler) and self.num_heads == other.num_heads + + +class ReplayOldDAcumsumHandler(StateResourceHandler): + """Per-layer old_dA_cumsum cache for the replay SSM kernel (double-buffered, fp32). + + Shape: (max_batch, 2, num_heads, T) — T from manager. + Routes to MambaHybridCacheManager via get_replay_old_dA_cumsum(layer_idx). + """ + + def __init__(self, num_heads: int) -> None: + self.num_heads = num_heads + super().__init__(dtype=torch.float32) + + @property + def state_shape(self) -> Tuple[int]: + return (self.num_heads,) + + def allocate(self, sequence_info) -> torch.Tensor: + raise RuntimeError("ReplayOldDAcumsumHandler must be backed by MambaHybridCacheManager") + + def __eq__(self, other) -> bool: + return isinstance(other, ReplayOldDAcumsumHandler) and self.num_heads == other.num_heads + + +class ReplayCacheBufIdxHandler(StateResourceHandler): + """Global cache_buf_idx tensor for the replay SSM kernel (shared across all layers, int32). + + Shape: (max_batch,). Routes to MambaHybridCacheManager.get_replay_cache_buf_idx(). + The same tensor is returned for every SSM layer — the FX graph has one constant-address + input per layer, all pointing to the same buffer. + """ + + def __init__(self) -> None: + super().__init__(dtype=torch.int32) + + @property + def state_shape(self) -> Tuple[()]: + return () + + def allocate(self, sequence_info) -> torch.Tensor: + raise RuntimeError("ReplayCacheBufIdxHandler must be backed by MambaHybridCacheManager") + + def __eq__(self, other) -> bool: + return isinstance(other, ReplayCacheBufIdxHandler) + + +class ReplayPrevNumAcceptedHandler(StateResourceHandler): + """Global prev_num_accepted_tokens tensor for the replay SSM kernel (int32, shared). + + Shape: (max_batch,). Routes to MambaHybridCacheManager.get_replay_prev_num_accepted_tokens(). + """ + + def __init__(self) -> None: + super().__init__(dtype=torch.int32) + + @property + def state_shape(self) -> Tuple[()]: + return () + + def allocate(self, sequence_info) -> torch.Tensor: + raise RuntimeError("ReplayPrevNumAcceptedHandler must be backed by MambaHybridCacheManager") + + def __eq__(self, other) -> bool: + return isinstance(other, ReplayPrevNumAcceptedHandler) + + class SpecCausalConvResourceHandler(StateResourceHandler): """Intermediate conv state cache descriptor for speculative decoding. @@ -2035,6 +2178,16 @@ def get_constants(cls, source_attn_node: Node) -> List[Constant]: """ return [] + @classmethod + def get_kwarg_cache_keys(cls) -> Set[str]: + """Return cache resource keys that should be passed as kwargs (not positional args). + + By default all cache resources are positional. Override to route specific cache keys + (e.g. replay tensors that must appear after constants in the function signature) to + the call_function kwargs dict instead. + """ + return set() + @classmethod def get_layer_idx(cls, source_attn_node: Node) -> Optional[int]: """Return the logical layer index associated with a source attention node, if any.""" diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/flashinfer_backend_mamba.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/flashinfer_backend_mamba.py index 563cfe34c9b9..25defab12f5d 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/flashinfer_backend_mamba.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/flashinfer_backend_mamba.py @@ -19,11 +19,22 @@ from flashinfer.mamba import selective_state_update as _flashinfer_ssm_update from torch.fx import Node +from tensorrt_llm._torch.modules.mamba.replay_selective_state_update import ( + replay_selective_state_update, +) +from tensorrt_llm._utils import get_sm_version + from ..._compat import KvCacheConfig from ..attention_interface import ( AttentionRegistry, BatchInfo, MHACallable, + ReplayCacheBufIdxHandler, + ReplayOldBHandler, + ReplayOldDAcumsumHandler, + ReplayOldDtHandler, + ReplayOldXHandler, + ReplayPrevNumAcceptedHandler, ResourceHandlerDict, SpecSSMResourceHandler, ) @@ -49,7 +60,7 @@ def _fi_align(t: torch.Tensor) -> torch.Tensor: @torch.library.custom_op( "auto_deploy::flashinfer_cached_ssm", - mutates_args=("ssm_state_cache", "intermediate_ssm_state_cache"), + mutates_args=("ssm_state_cache",), ) def _flashinfer_cached_ssm( # INPUTS (dense but may be flattened across sequences) @@ -72,12 +83,21 @@ def _flashinfer_cached_ssm( seq_idx_prefill: torch.Tensor, # [1, num_prefill_tokens] # CACHES ssm_state_cache: torch.Tensor, # [max_batch_size, num_heads, head_dim, ssm_state_size] + # CONSTANTS — Optional defaults so infer_schema accepts them; always supplied by the transform. + # Positions 16-17: constants come right after the single positional cache node (ssm_state_cache). + time_step_limit: Optional[List[float]] = None, + chunk_size: int = 256, + # SPEC / REPLAY CACHES — all kwargs, absent in replay mode or non-replay mode respectively. + # Lowercase names to match PyTorch FX node name normalization. intermediate_ssm_state_cache: Optional[ torch.Tensor - ], # [spec_state_size, max_draft_len+1, num_heads, head_dim, d_state] - # CONSTANTS - time_step_limit: List[float], - chunk_size: int, + ] = None, # kwarg; present in non-replay, absent in replay (uses None default) + replay_old_x: Optional[torch.Tensor] = None, # [max_batch, T, nheads, head_dim] + replay_old_b: Optional[torch.Tensor] = None, # [max_batch, 2, T, ngroups, dstate] + replay_old_dt: Optional[torch.Tensor] = None, # [max_batch, 2, nheads, T] fp32 + replay_old_da_cumsum: Optional[torch.Tensor] = None, # [max_batch, 2, nheads, T] fp32 + replay_cache_buf_idx: Optional[torch.Tensor] = None, # [max_batch] int32 + replay_prev_num_accepted: Optional[torch.Tensor] = None, # [max_batch] int32 out: Optional[torch.Tensor] = None, ) -> torch.Tensor: b, s, num_heads, head_dim, bs, hs_flat, B_flat, C_flat, dt_flat = _flatten_ssm_inputs( @@ -89,6 +109,15 @@ def _flashinfer_cached_ssm( num_prefill_tokens, num_extend_tokens, num_decode_tokens = batch_info.get_num_tokens() num_total_tokens = num_prefill_tokens + num_extend_tokens + num_decode_tokens + # Pre-hoist index type conversions so their kernels run before the extend + # causal_conv1d_update kernel rather than landing in the gap between it + # and the PDL-dependent replay precompute kernel. + slot_idx_i32 = slot_idx.to(torch.int32) + # NOTE: intermediate_state_indices_extend (torch.arange) is NOT hoisted here. + # It is only needed for the non-replay FlashInfer extend path. Hoisting it + # unconditionally would create an unnecessary arange kernel in the gap for + # the replay path (where replay_old_x is not None and the arange is unused). + if out is not None: preallocated_ssm_out = out.view(bs, num_heads, head_dim) else: @@ -116,7 +145,7 @@ def _flashinfer_cached_ssm( chunk_offsets, seq_idx_prefill, ssm_state_cache, - time_step_limit, + time_step_limit or [], chunk_size, preallocated_ssm_out[:num_prefill_tokens].unsqueeze(0), ) @@ -130,7 +159,7 @@ def _flashinfer_cached_ssm( A, D, dt_bias, - slot_idx, + slot_idx_i32, # already int32 — slice is a no-op, no extra kernel seq_start=num_prefill, token_start=num_prefill_tokens, num_seq=num_extend, @@ -142,14 +171,8 @@ def _flashinfer_cached_ssm( if extend_inputs is not None: tokens_per_extend = num_extend_tokens // num_extend - if intermediate_ssm_state_cache.size(1) < tokens_per_extend: - raise RuntimeError( - "flashinfer_cached_ssm: intermediate_ssm_state_cache is too small " - f"for extend branch (size1={intermediate_ssm_state_cache.size(1)}, " - f"tokens_per_extend={tokens_per_extend})" - ) ( - slot_idx_extend, + slot_idx_extend, # already int32 (sliced from slot_idx_i32) x_extend, B_extend, C_extend, @@ -163,27 +186,66 @@ def _flashinfer_cached_ssm( num_prefill_tokens : num_prefill_tokens + num_extend_tokens ].view(num_extend, tokens_per_extend, num_heads, head_dim) - intermediate_state_indices = torch.arange( - num_extend, dtype=torch.int32, device=slot_idx_extend.device - ) - _flashinfer_ssm_update( - ssm_state_cache, - x_extend, - dt_extend, - A_full, - B_extend, - C_extend, - D_full, - z=None, - dt_bias=dt_bias_hp, - dt_softplus=True, - state_batch_indices=slot_idx_extend.to(torch.int32), - out=preallocated_ssm_out_e, - disable_state_update=True, - intermediate_states_buffer=intermediate_ssm_state_cache, - cache_steps=tokens_per_extend, - intermediate_state_indices=intermediate_state_indices, - ) + slot_idx_extend_i32 = slot_idx_extend # already int32; no .to() kernel + + use_replay = replay_old_x is not None + if use_replay: + # Replay path: fast-forward SSM state via tl.dot on cached values. + # State is updated in-place; no disable_state_update needed. + # x_extend/B_extend/C_extend are non-contiguous views from the CUDA graph's + # preallocated buffer (stride_T > nheads × head_dim). The Triton kernel + # handles arbitrary strides via pointer arithmetic; skipping .contiguous() + # avoids a captured BF16 copy kernel and reduces GPU time by ~5%. + replay_selective_state_update( + ssm_state_cache, + replay_old_x, + replay_old_b, + replay_old_dt, + replay_old_da_cumsum, + replay_cache_buf_idx, + replay_prev_num_accepted, + x_extend, + dt_extend, + A_full, + B_extend, + C_extend, + out=preallocated_ssm_out_e, + D=D_full, + dt_bias=dt_bias_hp, + dt_softplus=True, + state_batch_indices=slot_idx_extend_i32, + launch_with_pdl=True, # PDL chain: triton_causal_conv extend → precompute → main + ) + else: + if intermediate_ssm_state_cache.size(1) < tokens_per_extend: + raise RuntimeError( + "flashinfer_cached_ssm: intermediate_ssm_state_cache is too small " + f"for extend branch (size1={intermediate_ssm_state_cache.size(1)}, " + f"tokens_per_extend={tokens_per_extend})" + ) + # Only needed for the non-replay FlashInfer path; not hoisted to avoid + # an unnecessary arange kernel in the stream gap for the replay path. + intermediate_state_indices = torch.arange( + num_extend, dtype=torch.int32, device=slot_idx_extend_i32.device + ) + _flashinfer_ssm_update( + ssm_state_cache, + x_extend.contiguous(), + dt_extend, + A_full, + B_extend.contiguous(), + C_extend.contiguous(), + D_full, + z=None, + dt_bias=dt_bias_hp, + dt_softplus=True, + state_batch_indices=slot_idx_extend_i32, + out=preallocated_ssm_out_e, + disable_state_update=True, + intermediate_states_buffer=intermediate_ssm_state_cache, + cache_steps=tokens_per_extend, + intermediate_state_indices=intermediate_state_indices, + ) # DECODE: single-token autoregressive path decode_inputs = _prepare_ssm_decode_inputs( @@ -194,7 +256,7 @@ def _flashinfer_cached_ssm( A, D, dt_bias, - slot_idx, + slot_idx_i32, # already int32 — no .to() kernel at decode time num_prefill + num_extend, num_prefill_tokens + num_extend_tokens, num_decode, @@ -206,7 +268,7 @@ def _flashinfer_cached_ssm( if decode_inputs is not None: ( - slot_idx_decode, + slot_idx_decode, # already int32 (sliced from slot_idx_i32) x_decode, B_decode, C_decode, @@ -220,7 +282,7 @@ def _flashinfer_cached_ssm( B_decode = _fi_align(B_decode) C_decode = _fi_align(C_decode) - slot_idx_decode_i32 = slot_idx_decode.to(torch.int32) + slot_idx_decode_i32 = slot_idx_decode # already int32; no .to() kernel y_decode = _flashinfer_ssm_update( ssm_state_cache, x_decode, @@ -270,12 +332,20 @@ def _flashinfer_cached_ssm_fake( seq_idx_prefill: torch.Tensor, # [1, num_prefill_tokens] # CACHES ssm_state_cache: torch.Tensor, # [max_batch_size, num_heads, head_dim, ssm_state_size] + # CONSTANTS + time_step_limit: Optional[List[float]] = None, + chunk_size: int = 256, + # SPEC / REPLAY CACHES — all kwargs, absent in replay mode or non-replay mode respectively. + # Lowercase names to match PyTorch FX node name normalization. intermediate_ssm_state_cache: Optional[ torch.Tensor - ], # [spec_state_size, max_draft_len+1, num_heads, head_dim, d_state] - # CONSTANTS - time_step_limit: List[float], - chunk_size: int, + ] = None, # [spec_state_size, max_draft_len+1, num_heads, head_dim, d_state] + replay_old_x: Optional[torch.Tensor] = None, # [max_batch, T, nheads, head_dim] + replay_old_b: Optional[torch.Tensor] = None, # [max_batch, 2, T, ngroups, dstate] + replay_old_dt: Optional[torch.Tensor] = None, # [max_batch, 2, nheads, T] fp32 + replay_old_da_cumsum: Optional[torch.Tensor] = None, # [max_batch, 2, nheads, T] fp32 + replay_cache_buf_idx: Optional[torch.Tensor] = None, # [max_batch] int32 + replay_prev_num_accepted: Optional[torch.Tensor] = None, # [max_batch] int32 out: Optional[torch.Tensor] = None, ): if out is not None: @@ -294,6 +364,30 @@ def _flashinfer_cached_ssm_fake( @AttentionRegistry.register("flashinfer_ssm") class FlashinferBackendSSM(BaseBackendSSM): + # When ssm_replay=True, use the replay kernel (tl.dot fast-forward) instead of FlashInfer. + # Disabled automatically when: block reuse enabled, tree attention, or SM < 80. + ssm_replay: bool = False + + # Cache keys always passed as kwargs (follow constants at positions 16-17). + # intermediate_ssm_state_cache is always a kwarg: present in non-replay, absent in replay + # (uses function default None). Replay keys are present only when ssm_replay=True. + # All keys lowercase to match PyTorch FX node name normalization. + _KWARG_CACHE_KEYS = frozenset( + { + "intermediate_ssm_state_cache", + "replay_old_x", + "replay_old_b", + "replay_old_dt", + "replay_old_da_cumsum", + "replay_cache_buf_idx", + "replay_prev_num_accepted", + } + ) + + @classmethod + def get_kwarg_cache_keys(cls): + return cls._KWARG_CACHE_KEYS + @classmethod def get_cached_attention_op(cls) -> MHACallable: return torch.ops.auto_deploy.flashinfer_cached_ssm.default @@ -304,15 +398,43 @@ def get_cache_initializers( ) -> ResourceHandlerDict: ret = super().get_cache_initializers(source_attn_node, cache_config) - # check head_dim is supported by flashinfer - if ret["ssm_state_cache"].head_dim not in FLASHINFER_SUPPORTED_HEAD_DIMS: + # In replay mode we never call FlashInfer, so head_dim can be anything. + if ( + not cls.ssm_replay + and ret["ssm_state_cache"].head_dim not in FLASHINFER_SUPPORTED_HEAD_DIMS + ): raise ValueError( f"flashinfer_ssm only supports head_dim in {FLASHINFER_SUPPORTED_HEAD_DIMS}. " f"Got head_dim={ret['ssm_state_cache'].head_dim}. " "Consider using 'triton_ssm' backend instead." ) - ret["intermediate_ssm_state_cache"] = SpecSSMResourceHandler.from_base( - ret["ssm_state_cache"] - ) + if cls.ssm_replay and get_sm_version() >= 80: + # Replay mode: intermediate_ssm_state_cache is absent from the graph — it uses the + # function's default (None) via kwarg routing. Only replay buffers are registered. + # T is determined by MambaHybridCacheManager from spec_config at construction time. + ssm_h = ret["ssm_state_cache"] + x_dtype = torch.bfloat16 + + ret["replay_old_x"] = ReplayOldXHandler( + num_heads=ssm_h.num_heads, head_dim=ssm_h.head_dim, dtype=x_dtype + ) + # Derive n_groups from B tensor shape at the source node. + # Use lowercase keys: FX lowercases node names, so "replay_old_B" → + # node.name "replay_old_b". Use lowercase to keep _caches / node names consistent. + B_fake = source_attn_node.args[2].meta["val"] + n_groups = B_fake.shape[-2] if B_fake.ndim >= 4 else 1 + ret["replay_old_b"] = ReplayOldBHandler( + n_groups=n_groups, d_state=ssm_h.d_state, dtype=x_dtype + ) + ret["replay_old_dt"] = ReplayOldDtHandler(num_heads=ssm_h.num_heads) + ret["replay_old_da_cumsum"] = ReplayOldDAcumsumHandler(num_heads=ssm_h.num_heads) + ret["replay_cache_buf_idx"] = ReplayCacheBufIdxHandler() + ret["replay_prev_num_accepted"] = ReplayPrevNumAcceptedHandler() + else: + # Non-replay mode: include intermediate_ssm_state_cache as a kwarg resource. + # It is absent in replay mode, where the function default (None) takes over. + ret["intermediate_ssm_state_cache"] = SpecSSMResourceHandler.from_base( + ret["ssm_state_cache"] + ) return ret diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/triton_backend_causal_conv.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/triton_backend_causal_conv.py index e90245e566ab..a3f4581d7e89 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/triton_backend_causal_conv.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/triton_backend_causal_conv.py @@ -100,6 +100,10 @@ def _triton_cached_causal_conv1d( else: w2d = weight + # Pre-hoist index type conversions so they run as early kernels, not between + # the extend conv1d update and the downstream replay precompute kernel. + slot_idx_i32 = slot_idx.to(torch.int32) + # PREFILL: concatenate all prefill tokens and run one varlen forward if num_prefill > 0: # x_varlen: (dim, cu_seq_len) @@ -117,7 +121,7 @@ def _triton_cached_causal_conv1d( conv_state_cache, prefill_cu_seqlen, seq_lens_cpu, - cache_indices=slot_idx[:num_prefill].to(torch.int32), + cache_indices=slot_idx_i32[:num_prefill], has_initial_state=use_initial_states[:num_prefill], activation=activation, pad_slot_id=PAD_SLOT_ID, @@ -125,7 +129,40 @@ def _triton_cached_causal_conv1d( # Scatter outputs back to input buffer inp_flat[:num_prefill_tokens] = y_varlen.transpose(0, 1) + # DECODE: process decode before extend so decode-related kernels (slot index + # conversions, scatter-back) don't land in the gap between the extend + # causal_conv1d_update kernel and the downstream PDL-dependent replay precompute + # kernel. Decode and extend access disjoint slot indices so reordering is safe. + if num_decode > 0: + x_decode = inp_flat[ + num_prefill_tokens + num_extend_tokens : num_total_tokens + ] # [num_decode_tokens, C_in] + + # Note: Triton causal_conv1d_update returns a new tensor (not in-place like CUDA version) + # so we need to capture the output and write it back + y_decode = causal_conv1d_update( + x_decode, # [batch, dim] + conv_state_cache, + w2d, + bias, + activation=activation, + cache_seqlens=None, + conv_state_indices=slot_idx_i32[num_prefill + num_extend : num_seq], + pad_slot_id=PAD_SLOT_ID, + ) + inp_flat[num_prefill_tokens + num_extend_tokens : num_total_tokens] = y_decode + + # Zero padding positions beyond valid tokens (for piecewise CUDA graph). + # Run before extend so the zero_() kernel doesn't land after extend conv. + if num_total_tokens < bs: + inp_flat[num_total_tokens:].zero_() + # EXTEND: use the update kernel so extend tokens write the intermediate state cache. + # This block runs last so that: + # 1. All non-extend kernels (decode, zero_()) have already executed. + # 2. After causal_conv1d_update emits gdc_launch_dependents(), the very next + # kernel in this stream is the extend scatter-back — just one copy_ — before + # the downstream flashinfer_ssm precompute kernel starts via PDL. if num_extend > 0: # num_extend_tokens == num_extend * (max_draft_len + 1) for static draft lengths # (dynamic lengths not supported) @@ -136,7 +173,7 @@ def _triton_cached_causal_conv1d( "that is too small for the extend branch" ) - slot_idx_extend = slot_idx[num_prefill : num_prefill + num_extend].to(torch.int32) + slot_idx_extend = slot_idx_i32[num_prefill : num_prefill + num_extend] # The intermediate state cache will be stored in these indices and read by the mamba_cache_manager, # which expects them in the indices arange(num_extend). They are not used across requests, so we @@ -161,34 +198,15 @@ def _triton_cached_causal_conv1d( intermediate_conv_window=intermediate_conv_state_cache, intermediate_state_indices=intermediate_state_indices, pad_slot_id=PAD_SLOT_ID, + launch_dependent_kernels=True, # PDL signal to downstream replay precompute ) - inp_flat[num_prefill_tokens : num_prefill_tokens + num_extend_tokens] = y_extend.transpose( - 1, 2 - ).view(-1, inp_flat.shape[1]) - - # DECODE: batch update for single-token sequences - if num_decode > 0: - x_decode = inp_flat[ - num_prefill_tokens + num_extend_tokens : num_total_tokens - ] # [num_decode_tokens, C_in] - - # Note: Triton causal_conv1d_update returns a new tensor (not in-place like CUDA version) - # so we need to capture the output and write it back - y_decode = causal_conv1d_update( - x_decode, # [batch, dim] - conv_state_cache, - w2d, - bias, - activation=activation, - cache_seqlens=None, - conv_state_indices=slot_idx[num_prefill + num_extend : num_seq].to(torch.int32), - pad_slot_id=PAD_SLOT_ID, - ) - inp_flat[num_prefill_tokens + num_extend_tokens : num_total_tokens] = y_decode - - # Zero padding positions beyond valid tokens (for piecewise CUDA graph) - if num_total_tokens < bs: - inp_flat[num_total_tokens:].zero_() + # Single copy_ scatter-back: y_extend [N,C,T] → inp_flat slice [N*T,C]. + # Uses one efficient non-contiguous-to-contiguous copy instead of + # transpose+view+assign (previously 2+ kernels). This is now the only kernel + # between causal_conv1d_update and the PDL-dependent precompute. + inp_flat[num_prefill_tokens : num_prefill_tokens + num_extend_tokens].view( + num_extend, tokens_per_extend, -1 + ).copy_(y_extend.permute(0, 2, 1)) @_triton_cached_causal_conv1d.register_fake diff --git a/tensorrt_llm/_torch/auto_deploy/shim/interface.py b/tensorrt_llm/_torch/auto_deploy/shim/interface.py index 4bfc0cf00f68..749bd0dd2a54 100644 --- a/tensorrt_llm/_torch/auto_deploy/shim/interface.py +++ b/tensorrt_llm/_torch/auto_deploy/shim/interface.py @@ -50,6 +50,12 @@ from ..custom_ops.attention_interface import ( CausalConvResourceHandler, KVPagedResourceHandler, + ReplayCacheBufIdxHandler, + ReplayOldBHandler, + ReplayOldDAcumsumHandler, + ReplayOldDtHandler, + ReplayOldXHandler, + ReplayPrevNumAcceptedHandler, ResourceHandler, ResourceHandlerDict, SequenceInfo, @@ -400,19 +406,69 @@ def _identify_managed_state_resources( and handler == SpecCausalConvResourceHandler.from_base(conv_ref) ] + # Replay SSM buffers — per-layer (old_x, old_B, old_dt, old_dA_cumsum) + # and global (cache_buf_idx, prev_num_accepted_tokens). + replay_old_x = [ + (name, handler) + for name, handler in self._resource_lookup.items() + if isinstance(handler, ReplayOldXHandler) + ] + replay_old_B = [ + (name, handler) + for name, handler in self._resource_lookup.items() + if isinstance(handler, ReplayOldBHandler) + ] + replay_old_dt = [ + (name, handler) + for name, handler in self._resource_lookup.items() + if isinstance(handler, ReplayOldDtHandler) + ] + replay_old_dA_cumsum = [ + (name, handler) + for name, handler in self._resource_lookup.items() + if isinstance(handler, ReplayOldDAcumsumHandler) + ] + replay_cache_buf_idx = [ + (name, handler) + for name, handler in self._resource_lookup.items() + if isinstance(handler, ReplayCacheBufIdxHandler) + ] + replay_prev_num_accepted = [ + (name, handler) + for name, handler in self._resource_lookup.items() + if isinstance(handler, ReplayPrevNumAcceptedHandler) + ] + # When speculative decoding is enabled, the backend must supply matching spec buffers. # When it is not enabled, spec buffers may still be registered by the backend (e.g. # triton_ssm always registers intermediate_ssm_state_cache) but will not be bound. + # Exception: in replay mode intermediate_ssm_state_cache is omitted entirely from + # get_cache_initializers and uses the function's default (None) via kwarg routing. + use_replay = len(replay_old_x) > 0 if self._spec_config is not None: - assert len(ssm_spec) == len(ssm_managed), ( - f"Mismatched SSM spec layer count: expected {len(ssm_managed)}, got {len(ssm_spec)}" - ) + if not use_replay: + assert len(ssm_spec) == len(ssm_managed), ( + f"Mismatched SSM spec layer count: expected {len(ssm_managed)}, got {len(ssm_spec)}" + ) assert len(conv_spec) == len(conv_managed), ( f"Mismatched Conv spec layer count: expected {len(conv_managed)}, " f"got {len(conv_spec)}" ) - return ssm_ref, ssm_managed, ssm_spec, conv_ref, conv_managed, conv_spec + return ( + ssm_ref, + ssm_managed, + ssm_spec, + conv_ref, + conv_managed, + conv_spec, + replay_old_x, + replay_old_B, + replay_old_dt, + replay_old_dA_cumsum, + replay_cache_buf_idx, + replay_prev_num_accepted, + ) def _prepare_kv_cache_config( self, @@ -541,6 +597,12 @@ def _create_and_assign_state_views( conv_ref: Optional[CausalConvResourceHandler], conv_managed: list, conv_spec: list, + replay_old_x: list = (), + replay_old_B: list = (), + replay_old_dt: list = (), + replay_old_dA_cumsum: list = (), + replay_cache_buf_idx: list = (), + replay_prev_num_accepted: list = (), ) -> Tuple[MambaHybridCacheManager, int]: """Create MambaHybridCacheManager and assign views for state resources. @@ -556,20 +618,26 @@ def _create_and_assign_state_views( conv_ref: Reference Conv handler or None. conv_managed: List of base Conv resources. conv_spec: List of speculative Conv resources. + replay_old_x/B/dt/dA_cumsum: Per-layer replay cache resource lists. + replay_cache_buf_idx/prev_num_accepted: Global replay resource lists. Returns: Tuple of (manager, num_managed_mamba_layers). """ + # Detect replay mode from presence of ReplayOldXHandler resources. + use_replay = len(replay_old_x) > 0 + # Mamba state params can be derived from reference handlers and number of managed (non-speculative) resources. mamba_params = self._get_mamba_state_params( ssm_ref, len(ssm_managed), conv_ref, len(conv_managed) ) num_managed_mamba_layers = mamba_params["mamba_num_layers"] - # Create the hybrid cache manager + # Create the hybrid cache manager, enabling replay path if detected. manager = MixedMambaHybridCacheManager( **mamba_params, **kv_cache_kwargs, + use_replay_state_update=use_replay, ) # Retrieve and assign views for Mamba-managed resources (up to num_managed_mamba_layers). @@ -605,6 +673,36 @@ def _create_and_assign_state_views( assert spec_view.is_contiguous(), f"Non-contiguous state {spec_conv_name}" self._caches[spec_conv_name] = spec_view + # Replay per-layer buffers (only when replay mode is active and spec dec enabled). + if use_replay and self._spec_config is not None: + for buf_list, getter, label in [ + (replay_old_x, manager.get_replay_old_x, "old_x"), + (replay_old_B, manager.get_replay_old_B, "old_B"), + (replay_old_dt, manager.get_replay_old_dt, "old_dt"), + (replay_old_dA_cumsum, manager.get_replay_old_dA_cumsum, "old_dA_cumsum"), + ]: + if buf_list: + buf_name = buf_list[layer_idx][0] + buf_view = getter(layer_idx) + if buf_view is None: + raise RuntimeError( + f"Replay buffer {label} not allocated for layer {layer_idx}. " + "Is use_replay_state_update=True in the cache manager?" + ) + self._caches[buf_name] = buf_view + + # Replay global buffers — same tensor for every layer; bind once per resource entry. + if use_replay and self._spec_config is not None: + global_buf_map = [ + (replay_cache_buf_idx, manager.get_replay_cache_buf_idx()), + (replay_prev_num_accepted, manager.get_replay_prev_num_accepted_tokens()), + ] + for buf_list, global_tensor in global_buf_map: + if global_tensor is None: + continue + for buf_name, _ in buf_list: + self._caches[buf_name] = global_tensor + return manager, num_managed_mamba_layers def _assign_kv_cache_views(self, kv_managed: Dict[str, KVPagedResourceHandler]) -> int: @@ -701,9 +799,20 @@ def _create_kv_cache_manager(self, max_tokens: Optional[int] = None) -> Dict: """ # 1. Identify managed resources kv_ref, kv_managed = self._identify_managed_kv_resources() - ssm_ref, ssm_managed, ssm_spec, conv_ref, conv_managed, conv_spec = ( - self._identify_managed_state_resources() - ) + ( + ssm_ref, + ssm_managed, + ssm_spec, + conv_ref, + conv_managed, + conv_spec, + replay_old_x, + replay_old_B, + replay_old_dt, + replay_old_dA_cumsum, + replay_cache_buf_idx, + replay_prev_num_accepted, + ) = self._identify_managed_state_resources() # 2. Prepare configuration kv_cache_config = self._prepare_kv_cache_config(max_tokens, kv_managed) @@ -722,6 +831,12 @@ def _create_kv_cache_manager(self, max_tokens: Optional[int] = None) -> Dict: conv_ref, conv_managed, conv_spec, + replay_old_x=replay_old_x, + replay_old_B=replay_old_B, + replay_old_dt=replay_old_dt, + replay_old_dA_cumsum=replay_old_dA_cumsum, + replay_cache_buf_idx=replay_cache_buf_idx, + replay_prev_num_accepted=replay_prev_num_accepted, ) else: # No typed state resources - use pure KVCacheManager diff --git a/tensorrt_llm/_torch/auto_deploy/transform/library/fuse_causal_conv.py b/tensorrt_llm/_torch/auto_deploy/transform/library/fuse_causal_conv.py index bce8476e6ec4..e1998e901dd1 100644 --- a/tensorrt_llm/_torch/auto_deploy/transform/library/fuse_causal_conv.py +++ b/tensorrt_llm/_torch/auto_deploy/transform/library/fuse_causal_conv.py @@ -22,6 +22,7 @@ from torch.fx import GraphModule, Node from ...custom_ops.mamba.cuda_backend_causal_conv import cuda_cached_causal_conv1d_wrapper +from ...custom_ops.mamba.triton_backend_causal_conv import triton_cached_causal_conv1d_wrapper from ...models.factory import ModelFactory from ...shim.interface import CachedSequenceInterface from ...utils.node_utils import is_op @@ -36,20 +37,21 @@ def _match_causal_conv_activation_pattern( Match the causal_conv + activation pattern in the graph. The pattern corresponds to: - conv_out = cuda_cached_causal_conv1d(...) + conv_out = cuda_cached_causal_conv1d(...) OR triton_cached_causal_conv1d(...) out = activation(conv_out) Args: graph: The graph module to search - target_op: The target causal conv op to match + target_op: The target causal conv op to match (or a tuple/list of ops) Returns: A list of tuples (conv_node, activation_node, activation_name) for each match """ + target_ops = target_op if isinstance(target_op, (list, tuple)) else [target_op] matches = [] for node in graph.nodes: - if not is_op(node, target_op): + if not any(is_op(node, op) for op in target_ops): continue # Check if this node has exactly one user and it's an activation @@ -101,22 +103,29 @@ def _apply( ) -> Tuple[GraphModule, TransformInfo]: graph = gm.graph - target_op = cuda_cached_causal_conv1d_wrapper + # Match both the CUDA-backed and Triton-backed causal conv wrappers. + # fuse_causal_conv_activation previously only matched the CUDA wrapper, + # leaving the Triton wrapper's activation (e.g. silu(conv_out)) as a + # separate FX node that landed between the conv and SSM kernels in the + # CUDA graph stream, blocking the PDL chain. + target_ops = [cuda_cached_causal_conv1d_wrapper, triton_cached_causal_conv1d_wrapper] # Step 1: Identify causal_conv + activation pattern matches = _match_causal_conv_activation_pattern( graph, - target_op=target_op, + target_op=target_ops, ) # Step 2: Replace matched patterns with fused version for conv_node, activation_node, activation_name in matches: + # Determine which wrapper was matched so the replacement uses the same op + conv_target_op = conv_node.target with graph.inserting_after(conv_node): # Create new call with fused activation # Replace the last arg (activation=None) with activation_name new_args = list(conv_node.args[:-1]) + [activation_name] fused_node = graph.call_function( - target_op, + conv_target_op, args=tuple(new_args), ) diff --git a/tensorrt_llm/_torch/auto_deploy/transform/library/kvcache.py b/tensorrt_llm/_torch/auto_deploy/transform/library/kvcache.py index 2b32b1987d0e..467a484cad97 100644 --- a/tensorrt_llm/_torch/auto_deploy/transform/library/kvcache.py +++ b/tensorrt_llm/_torch/auto_deploy/transform/library/kvcache.py @@ -211,6 +211,7 @@ def _insert_cached_attn_node( cache_nodes: List[Node], constants: List[Constant], prepared_attn_mask: Optional[Node] = None, + cache_kwarg_nodes: Optional[Dict[str, Node]] = None, ): """Insert a cached attention node into the graph.""" with gm.graph.inserting_before(attn_node): @@ -225,7 +226,14 @@ def _insert_cached_attn_node( all_args = (*all_args, prepared_attn_mask) cached_attn_node = gm.graph.call_function( cached_attn_op, - args=all_args, + args=( + *qkv_nodes, + *meta_nodes_std, + *meta_nodes_extra, + *cache_nodes, + *constants, + ), + kwargs=cache_kwarg_nodes or {}, ) attn_node.replace_all_uses_with(cached_attn_node) gm.graph.erase_node(attn_node) @@ -295,11 +303,18 @@ def _apply( "Each non-shared attention layer must own exactly one cache." ) cache_in_nodes = [] + cache_kwarg_nodes = {} + kwarg_keys = attn_descriptor.get_kwarg_cache_keys() for k, resource_handler in attn_descriptor.get_cache_initializers( attn_node, cm.kv_cache_config ).items(): resource_name = cm.add_resource(k, resource_handler) - cache_in_nodes.append(self._process_cache_node(gm, resource_name)) + node = self._process_cache_node(gm, resource_name) + if k in kwarg_keys: + # Pass as named kwarg to call_function (e.g. replay args after constants) + cache_kwarg_nodes[k] = node + else: + cache_in_nodes.append(node) if layer_idx is not None: cache_nodes_by_layer_idx[layer_idx] = cache_in_nodes @@ -327,6 +342,7 @@ def _apply( cache_in_nodes, constants, prepared_mask, + cache_kwarg_nodes, ) num_cached_attn_replacements += 1 diff --git a/tensorrt_llm/_torch/auto_deploy/transform/library/ssm_cache.py b/tensorrt_llm/_torch/auto_deploy/transform/library/ssm_cache.py index 36f2b8cb7e76..7a6b2be901b6 100644 --- a/tensorrt_llm/_torch/auto_deploy/transform/library/ssm_cache.py +++ b/tensorrt_llm/_torch/auto_deploy/transform/library/ssm_cache.py @@ -14,8 +14,28 @@ # limitations under the License. """A set of transforms to handle SSM cache transforms.""" -from ..interface import TransformRegistry -from .kvcache import _InsertCachedOperator +from typing import Type + +from pydantic import Field + +from ..interface import TransformConfig, TransformRegistry +from .kvcache import InsertCachedAttentionConfig, _InsertCachedOperator + + +class SSMCacheTransformConfig(InsertCachedAttentionConfig): + """Configuration for insert_cached_ssm_attention. + + Extends the base attention config with SSM-specific options. + """ + + ssm_replay: bool = Field( + default=False, + description=( + "Enable the replay SSM kernel (tl.dot fast-forward) for the MTP extend path. " + "Requires SM >= 80 (Ampere+). Falls back to FlashInfer when disabled or " + "when incompatible features are active (block reuse, tree attention)." + ), + ) # TODO: think about separating valid attention backends per transform better in the future @@ -23,6 +43,21 @@ class SSMCacheTransform(_InsertCachedOperator): """A transform to handle SSM cache operations.""" + config: SSMCacheTransformConfig + + @classmethod + def get_config_class(cls) -> Type[TransformConfig]: + return SSMCacheTransformConfig + + @property + def attn_descriptor(self): + descriptor = super().attn_descriptor + if self.config.ssm_replay and hasattr(descriptor, "ssm_replay"): + # Return a thin subclass with ssm_replay enabled so get_cache_initializers + # allocates replay buffers instead of intermediate_ssm_state_cache. + return type(descriptor.__name__ + "_Replay", (descriptor,), {"ssm_replay": True}) + return descriptor + @TransformRegistry.register("insert_cached_causal_conv") class InitializeCausalConvCache(_InsertCachedOperator): diff --git a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py index 5aeeeadfe861..74c130862902 100644 --- a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py @@ -609,6 +609,41 @@ def get_intermediate_conv_states(self, return self.mamba_cache.at_layer_idx( layer_offset).intermediate_conv_window + def get_replay_old_x(self, layer_idx: int) -> Optional[torch.Tensor]: + if not self._use_replay_state_update: + return None + layer_offset = self.mamba_layer_offsets[layer_idx] + return self.mamba_cache.old_x[layer_offset] + + def get_replay_old_B(self, layer_idx: int) -> Optional[torch.Tensor]: + if not self._use_replay_state_update: + return None + layer_offset = self.mamba_layer_offsets[layer_idx] + return self.mamba_cache.old_B[layer_offset] + + def get_replay_old_dt(self, layer_idx: int) -> Optional[torch.Tensor]: + if not self._use_replay_state_update: + return None + layer_offset = self.mamba_layer_offsets[layer_idx] + return self.mamba_cache.old_dt[layer_offset] + + def get_replay_old_dA_cumsum(self, + layer_idx: int) -> Optional[torch.Tensor]: + if not self._use_replay_state_update: + return None + layer_offset = self.mamba_layer_offsets[layer_idx] + return self.mamba_cache.old_dA_cumsum[layer_offset] + + def get_replay_cache_buf_idx(self) -> Optional[torch.Tensor]: + if not self._use_replay_state_update: + return None + return self.mamba_cache.cache_buf_idx + + def get_replay_prev_num_accepted_tokens(self) -> Optional[torch.Tensor]: + if not self._use_replay_state_update: + return None + return self.mamba_cache.prev_num_accepted_tokens + def is_speculative(self) -> bool: return isinstance(self.mamba_cache, self.SpeculativeState) @@ -813,6 +848,31 @@ def get_intermediate_conv_states(self, assert not self._use_cpp, "get_intermediate_conv_states is not supported in CppMambaCacheManager" return self._impl.get_intermediate_conv_states(layer_idx) + def get_replay_old_x(self, layer_idx: int) -> Optional[torch.Tensor]: + assert not self._use_cpp, "get_replay_old_x is not supported in CppMambaCacheManager" + return self._impl.get_replay_old_x(layer_idx) + + def get_replay_old_B(self, layer_idx: int) -> Optional[torch.Tensor]: + assert not self._use_cpp, "get_replay_old_B is not supported in CppMambaCacheManager" + return self._impl.get_replay_old_B(layer_idx) + + def get_replay_old_dt(self, layer_idx: int) -> Optional[torch.Tensor]: + assert not self._use_cpp, "get_replay_old_dt is not supported in CppMambaCacheManager" + return self._impl.get_replay_old_dt(layer_idx) + + def get_replay_old_dA_cumsum(self, + layer_idx: int) -> Optional[torch.Tensor]: + assert not self._use_cpp, "get_replay_old_dA_cumsum is not supported in CppMambaCacheManager" + return self._impl.get_replay_old_dA_cumsum(layer_idx) + + def get_replay_cache_buf_idx(self) -> Optional[torch.Tensor]: + assert not self._use_cpp, "get_replay_cache_buf_idx is not supported in CppMambaCacheManager" + return self._impl.get_replay_cache_buf_idx() + + def get_replay_prev_num_accepted_tokens(self) -> Optional[torch.Tensor]: + assert not self._use_cpp, "get_replay_prev_num_accepted_tokens is not supported in CppMambaCacheManager" + return self._impl.get_replay_prev_num_accepted_tokens() + def is_speculative(self) -> bool: return self._impl.is_speculative() 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 82614a5e1395..8a29a0167217 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 @@ -104,7 +104,6 @@ def test_flashinfer_decode_matches_triton(mamba_env): None, # seq_idx_prefill # CACHES ssm_state_cache_flashinfer, - None, # intermediate_ssm_state_cache (not used in decode-only path) # CONSTANTS time_step_limit, chunk_size, diff --git a/tests/unittest/auto_deploy/singlegpu/smoke/test_ad_speculative_decoding.py b/tests/unittest/auto_deploy/singlegpu/smoke/test_ad_speculative_decoding.py index 37f838a50294..5ceb84ed62fe 100644 --- a/tests/unittest/auto_deploy/singlegpu/smoke/test_ad_speculative_decoding.py +++ b/tests/unittest/auto_deploy/singlegpu/smoke/test_ad_speculative_decoding.py @@ -92,6 +92,70 @@ def test_super_mtp_smoke(): assert len(prompts_and_outputs) == 1 +def test_super_mtp_ssm_replay_smoke(): + """Test one-model MTP with flashinfer_ssm + ssm_replay=True. + + Uses mamba_head_dim=64 (required by FlashInfer) and ssm_state_size=64 + (>= 16 for the replay kernel's tl.dot constraint), actually executing + _replay_precompute_kernel and _replay_state_update_kernel. + """ + test_prompt = "What is the capital of France?" + model_hub_id = "nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-BF16" + model_path = hf_id_to_local_model_dir(model_hub_id) + + # Get base small-model config and update SSM dims for FlashInfer + replay. + # hidden_size must equal mamba_num_heads × mamba_head_dim (4 × 64 = 256). + experiment_config = get_small_model_config( + model_hub_id, + transforms={ + "insert_cached_causal_conv": {"backend": "triton_causal_conv"}, + "insert_cached_ssm_attention": {"backend": "flashinfer_ssm", "ssm_replay": True}, + }, + ) + experiment_config["args"]["model_kwargs"].update( + { + "hidden_size": 256, + "intermediate_size": 256, + "mamba_num_heads": 4, + "mamba_head_dim": 64, + "ssm_state_size": 64, + "moe_intermediate_size": 128, + "moe_shared_expert_intermediate_size": 128, + "moe_latent_size": 64, + } + ) + experiment_config["args"]["model"] = model_path + experiment_config["args"]["runtime"] = "trtllm" + experiment_config["args"]["world_size"] = 1 + experiment_config["args"]["speculative_config"] = MTPDecodingConfig( + num_nextn_predict_layers=3, + mtp_eagle_one_model=True, + speculative_model=model_path, + ) + experiment_config["args"]["speculative_model_kwargs"] = experiment_config["args"][ + "model_kwargs" + ] + experiment_config["args"]["attn_backend"] = "flashinfer" + experiment_config["args"]["disable_overlap_scheduler"] = True + experiment_config["args"]["compile_backend"] = "torch-simple" + experiment_config["args"]["max_num_tokens"] = 256 + experiment_config["prompt"]["batch_size"] = 1 + experiment_config["prompt"]["queries"] = test_prompt + + cfg = ExperimentConfig(**experiment_config) + cfg.prompt.sp_kwargs = { + "max_tokens": 20, + "top_k": None, + "temperature": 0.0, + "seed": 42, + } + + results = main(cfg) + + prompts_and_outputs = results["prompts_and_outputs"] + assert len(prompts_and_outputs) == 1 + + def test_kv_cache_extra_seq_len_for_spec_dec(): """Test that get_extra_seq_len_for_kv_cache computes correct extra capacity.""" from tensorrt_llm._torch.auto_deploy.llm_args import LlmArgs From 8be13b2fc402afbcfa006c3ee2e2aa91031a29fe Mon Sep 17 00:00:00 2001 From: Gal Hubara Agam <96368689+galagam@users.noreply.github.com> Date: Sun, 17 May 2026 05:18:20 -0700 Subject: [PATCH 02/17] fix low_memory_overrides in integ test Signed-off-by: Gal Hubara Agam <96368689+galagam@users.noreply.github.com> --- .../defs/accuracy/test_llm_api_autodeploy.py | 10 ++-------- 1 file changed, 2 insertions(+), 8 deletions(-) diff --git a/tests/integration/defs/accuracy/test_llm_api_autodeploy.py b/tests/integration/defs/accuracy/test_llm_api_autodeploy.py index 632ee879409b..e279416f6768 100644 --- a/tests/integration/defs/accuracy/test_llm_api_autodeploy.py +++ b/tests/integration/defs/accuracy/test_llm_api_autodeploy.py @@ -148,19 +148,14 @@ def low_memory_overrides(config, max_batch_size=32, free_gpu_memory_fraction=0.4, max_seq_len=8192, - max_num_tokens=8192, - cuda_graph_batch_sizes=None): + max_num_tokens=8192): """Update and return config that reduce memory footprint for unquantized (bf16) runs.""" - if cuda_graph_batch_sizes is None: - cuda_graph_batch_sizes = [ - s for s in [1, 2, 4, 8, 16, 32, 64, 128] if s <= max_batch_size - ] config.update({ "max_batch_size": max_batch_size, "max_seq_len": max_seq_len, "max_num_tokens": max_num_tokens, "cuda_graph_config": { - "batch_sizes": cuda_graph_batch_sizes + "max_batch_size": max_batch_size, }, }) kv_cache_config = config.setdefault("kv_cache_config", {}) @@ -769,7 +764,6 @@ def test_mtp(self, world_size, attn_backend, model_id): low_memory_overrides( kwargs, max_batch_size=8, - cuda_graph_batch_sizes=[1, 2, 4, 8], ) kwargs["attn_backend"] = attn_backend From ddc7596fee007c4d1ef0178ecd8d3dad3b338db0 Mon Sep 17 00:00:00 2001 From: Gal Hubara Agam <96368689+galagam@users.noreply.github.com> Date: Sun, 17 May 2026 06:50:42 -0700 Subject: [PATCH 03/17] [None][fix] adp_router: restore get_relax_value fix lost during rebase bb1affe76e fixed get_relax_value to guard against None scheduling_params and None attention_dp_rank, using bool() to safely handle Optional[bool]. Commit 23cd0c3eae (rebase from main) re-introduced the local get_relax_value definitions without the fix, causing TypeError when sorted() compared two None values from attention_dp_relax at higher concurrencies. Signed-off-by: Gal Hubara Agam <96368689+galagam@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py b/tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py index d031109ef6e4..b717fc38ca1e 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py @@ -303,7 +303,7 @@ def route_requests( def get_relax_value(req_item): scheduling_params = getattr(req_item.request, "py_scheduling_params", None) - if scheduling_params is None: + if scheduling_params is None or scheduling_params.attention_dp_rank is None: return True val = scheduling_params.attention_dp_relax return True if val is None else val @@ -587,7 +587,7 @@ def route_requests( def get_relax_value(req_item): scheduling_params = getattr(req_item.request, "py_scheduling_params", None) - if scheduling_params is None: + if scheduling_params is None or scheduling_params.attention_dp_rank is None: return True val = scheduling_params.attention_dp_relax return True if val is None else val From be37bec7ca57c36b206f57d455f2785cc8ad4818 Mon Sep 17 00:00:00 2001 From: Gal Hubara Agam <96368689+galagam@users.noreply.github.com> Date: Sun, 17 May 2026 06:08:31 -0700 Subject: [PATCH 04/17] [None][perf] AutoDeploy: add WS=4 MTP config and unify mamba_ssm_cache_dtype Add super_v3_mtp_ws4.yaml (AD, EP+BMM sharding, flashinfer_ssm + replay) and super_v3_mtp_ws{1,4}.yaml (PT curated configs). Align all 4 configs to: max_batch_size: 64 max_seq_len: 8192 max_num_tokens: 2048 cuda_graph_config: enable_padding: true max_batch_size: 64 and no attention DP Signed-off-by: Gal Hubara Agam <96368689+galagam@users.noreply.github.com> --- .../model_registry/configs/super_v3_mtp.yaml | 9 ++- .../configs/super_v3_mtp_ws4.yaml | 67 +++++++++++++++++++ .../configs/curated/super_v3_mtp_ws1.yaml | 18 +++++ .../configs/curated/super_v3_mtp_ws4.yaml | 19 ++++++ 4 files changed, 108 insertions(+), 5 deletions(-) create mode 100644 examples/auto_deploy/model_registry/configs/super_v3_mtp_ws4.yaml create mode 100644 examples/configs/curated/super_v3_mtp_ws1.yaml create mode 100644 examples/configs/curated/super_v3_mtp_ws4.yaml diff --git a/examples/auto_deploy/model_registry/configs/super_v3_mtp.yaml b/examples/auto_deploy/model_registry/configs/super_v3_mtp.yaml index cf1528a882d6..376b9e65493a 100644 --- a/examples/auto_deploy/model_registry/configs/super_v3_mtp.yaml +++ b/examples/auto_deploy/model_registry/configs/super_v3_mtp.yaml @@ -2,18 +2,17 @@ # runtime: trtllm compile_backend: torch-cudagraph -max_batch_size: 128 -max_seq_len: 16384 -max_num_tokens: 8192 +max_batch_size: 64 +max_seq_len: 8192 +max_num_tokens: 2048 enable_chunked_prefill: false # TODO: Support chunked prefill w/ spec dec attn_backend: trtllm model_factory: AutoModelForCausalLM skip_loading_weights: false cuda_graph_config: max_batch_size: 64 + enable_padding: true kv_cache_config: - # tunable mamba cache dtype - # --> use float32 for accuracy and default (auto) for speed mamba_ssm_cache_dtype: auto speculative_config: decoding_type: MTP diff --git a/examples/auto_deploy/model_registry/configs/super_v3_mtp_ws4.yaml b/examples/auto_deploy/model_registry/configs/super_v3_mtp_ws4.yaml new file mode 100644 index 000000000000..a3294c8a516d --- /dev/null +++ b/examples/auto_deploy/model_registry/configs/super_v3_mtp_ws4.yaml @@ -0,0 +1,67 @@ +# Config for SuperV3 with MTP speculative decoding, WS=4 (4-GPU EP+BMM sharding). +# Derived from super_v3_mtp.yaml with world_size: 4 to enable sharding. +# +runtime: trtllm +compile_backend: torch-cudagraph +world_size: 4 +max_batch_size: 64 +max_seq_len: 8192 +max_num_tokens: 2048 +enable_chunked_prefill: false # TODO: Support chunked prefill w/ spec dec +attn_backend: trtllm +model_factory: AutoModelForCausalLM +skip_loading_weights: false +cuda_graph_config: + max_batch_size: 64 + enable_padding: true +kv_cache_config: + # tunable mamba cache dtype + # --> use float32 for accuracy and default (auto) for speed + mamba_ssm_cache_dtype: auto +speculative_config: + decoding_type: MTP + num_nextn_predict_layers: 6 + mtp_eagle_one_model: true +transforms: + detect_sharding: + allreduce_strategy: NCCL + # NOTE: add 'tp' to sharding dims only for high-throughput runs + # For low-latency, keep mamba and attention replicated + sharding_dims: ['ep', 'bmm'] + # NOTE: sharding_source applies only to TP sharding + sharding_source: ['manual'] + manual_config: + head_dim: 128 + tp_plan: + # mamba SSM layer + "in_proj": "mamba" + "out_proj": "rowwise" + # attention layer + "q_proj": "colwise" + "k_proj": "colwise" + "v_proj": "colwise" + "o_proj": "rowwise" + # moe layer: SHARED experts + "up_proj": "colwise" + "down_proj": "rowwise" + # MoLE: latent projections: simple shard + "fc1_latent_proj": "gather" + "fc2_latent_proj": "gather" + # multi_stream_moe disabled - creates NaNs in target's sampled logits. TODO: investigate + fix + multi_stream_moe: + stage: compile + enabled: false + fuse_mamba_a_log: + stage: post_load_fusion + enabled: true + # FlashInfer SSM is used for the MTP extend path; Triton causal conv is still required + # for speculative Mamba state caching in this configuration. + insert_cached_ssm_attention: + backend: flashinfer_ssm + ssm_replay: true + insert_cached_causal_conv: + backend: triton_causal_conv + fuse_nvfp4_moe: + backend: trtllm_gen + fuse_fp8_moe: + allow_different_input_scales: true diff --git a/examples/configs/curated/super_v3_mtp_ws1.yaml b/examples/configs/curated/super_v3_mtp_ws1.yaml new file mode 100644 index 000000000000..61f47b3b92b7 --- /dev/null +++ b/examples/configs/curated/super_v3_mtp_ws1.yaml @@ -0,0 +1,18 @@ +# Nemotron-3-Super with MTP speculative decoding, WS=1 (single GPU). +max_batch_size: 64 +max_seq_len: 8192 +max_num_tokens: 2048 +tensor_parallel_size: 1 +moe_expert_parallel_size: 1 +trust_remote_code: true +enable_chunked_prefill: false +cuda_graph_config: + enable_padding: true + max_batch_size: 64 +kv_cache_config: + enable_block_reuse: false + mamba_ssm_cache_dtype: auto +num_postprocess_workers: 4 +speculative_config: + decoding_type: MTP + num_nextn_predict_layers: 6 diff --git a/examples/configs/curated/super_v3_mtp_ws4.yaml b/examples/configs/curated/super_v3_mtp_ws4.yaml new file mode 100644 index 000000000000..2d2a9b4390f5 --- /dev/null +++ b/examples/configs/curated/super_v3_mtp_ws4.yaml @@ -0,0 +1,19 @@ +# Nemotron-3-Super with MTP speculative decoding, WS=4 (4-GPU TP+EP). +max_batch_size: 64 +max_seq_len: 8192 +max_num_tokens: 2048 +tensor_parallel_size: 4 +moe_expert_parallel_size: 4 +trust_remote_code: true +enable_attention_dp: false +enable_chunked_prefill: false +cuda_graph_config: + enable_padding: true + max_batch_size: 64 +kv_cache_config: + enable_block_reuse: false + mamba_ssm_cache_dtype: auto +num_postprocess_workers: 4 +speculative_config: + decoding_type: MTP + num_nextn_predict_layers: 6 From 5ee80f218f6f401214675aacac515d8f43e6198c Mon Sep 17 00:00:00 2001 From: Gal Hubara Agam <96368689+galagam@users.noreply.github.com> Date: Fri, 22 May 2026 03:44:20 -0700 Subject: [PATCH 05/17] [None][fix] narrow FlashInfer SSD eligibility gate The FlashInfer SSD CuTe kernel (ssd_kernel.py) was written specifically for Nemotron-H (chunk_size=128, dstate=128) and contains a bug in tma_partition_for_mma_b_operand: it always tiles the B SSM operand with tile_shape_mnk_intra2[1:] = (headdim, chunk_size) instead of the correct tile_shape_mnk_intra1[1:] = (chunk_size, dstate). This makes gmem K-mode = chunk_size/16 mismatch smem K-mode = dstate/16, causing CuTe JIT compilation failure when dstate != chunk_size. The kernel only compiled for the reference Nemotron-H config because dstate == chunk_size == 128 masked the bug. Add `dstate == chunk_size` guard to _flashinfer_ssd_supported so the kernel is only attempted when the bug is masked (dstate == chunk_size, both in the valid set). This gates the only affected path: Blackwell (is_sm_100f) + flashinfer_ssm backend + non-NemotronH dstate (e.g. the smoke test with ssm_state_size=64, chunk_size=128). Production Nemotron-H (dstate=128, chunk_size=128) is unaffected. Other tests use triton_ssm backend or dstate=8 (never reaches this gate). See https://github.com/flashinfer-ai/flashinfer/issues/3397 for upstream fix. Also add test_ad_speculative_decoding.py to l0_b200.yml so this Blackwell- specific path is covered by CI going forward. Signed-off-by: Gal Hubara Agam <96368689+galagam@users.noreply.github.com> --- tensorrt_llm/_torch/modules/mamba/ssd_combined.py | 10 +++++++++- tests/integration/test_lists/test-db/l0_b200.yml | 1 + 2 files changed, 10 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/modules/mamba/ssd_combined.py b/tensorrt_llm/_torch/modules/mamba/ssd_combined.py index efb02364d291..0a2e7d0f1057 100644 --- a/tensorrt_llm/_torch/modules/mamba/ssd_combined.py +++ b/tensorrt_llm/_torch/modules/mamba/ssd_combined.py @@ -47,9 +47,17 @@ def _flashinfer_ssd_supported(chunk_size, dstate, headdim): + # The kernel was written for Nemotron-H (chunk_size=128, dstate=128) and has a + # bug in tma_partition_for_mma_b_operand: it tiles the B SSM operand with + # tile_shape_mnk_intra2[1:] = (headdim, chunk_size) instead of the correct + # tile_shape_mnk_intra1[1:] = (chunk_size, dstate), making gmem K-mode = + # chunk_size/16 mismatch smem K-mode = dstate/16. Only compiles when + # dstate == chunk_size (masked the bug for NemotronH). See + # https://github.com/flashinfer-ai/flashinfer/issues/3397 for upstream fix. return (chunk_size in _FLASHINFER_SSD_VALID_M_MODES and dstate in _FLASHINFER_SSD_VALID_M_MODES - and headdim in _FLASHINFER_SSD_VALID_HEAD_DIMS) + and headdim in _FLASHINFER_SSD_VALID_HEAD_DIMS + and dstate == chunk_size) @functools.cache diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 3b04ba55c708..459fa3073bda 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -334,6 +334,7 @@ l0_b200: - unittest/auto_deploy/singlegpu/custom_ops/moe/test_trtllm_moe.py - unittest/auto_deploy/singlegpu/custom_ops/quantization/test_quant.py - unittest/auto_deploy/singlegpu/smoke/test_ad_build_small_single.py -k "Nemotron-3-Nano-30B-A3B-FP8 or Nemotron-Nano-3-30B-A3.5B-dev or Llama-4-Scout" + - unittest/auto_deploy/singlegpu/smoke/test_ad_speculative_decoding.py - unittest/auto_deploy/singlegpu/transformations/library/test_fuse_relu2_quant_nvfp4.py - unittest/auto_deploy/singlegpu/transformations/library/test_moe_fusion.py - unittest/auto_deploy/singlegpu/transformations/library/test_nvfp4_swiglu.py From ff4235611307f55243c39abeff7162f2c5ddb577 Mon Sep 17 00:00:00 2001 From: Gal Hubara Agam <96368689+galagam@users.noreply.github.com> Date: Sun, 24 May 2026 01:26:02 -0700 Subject: [PATCH 06/17] [None][refactor] remove kwargs routing from flashinfer_ssm; make all caches positional The get_kwarg_cache_keys() hook in FlashinferBackendSSM routed 7 cache args (intermediate_ssm_state_cache + 6 replay caches) as FX call_function kwargs to work around their position after constants. Replace with the simpler positional pattern used by triton_cached_ssm: - All caches now come before constants in the function signature, matching the triton_cached_ssm pattern (ssm_state_cache, optional caches..., time_step_limit, chunk_size). - In non-replay mode: intermediate_ssm_state_cache is a real graph node, replay caches are literal None in the args tuple. - In replay mode: intermediate_ssm_state_cache is literal None, replay caches are real graph nodes. - kvcache.py: None resource handlers append literal None to cache_in_nodes instead of registering a graph placeholder, removing the need for any kwarg routing infrastructure. - Remove get_kwarg_cache_keys() from AttentionDescriptor and its override in FlashinferBackendSSM. Signed-off-by: Gal Hubara Agam <96368689+galagam@users.noreply.github.com> --- .../custom_ops/attention_interface.py | 10 -- .../mamba/flashinfer_backend_mamba.py | 103 +++++++----------- .../auto_deploy/transform/library/kvcache.py | 30 ++--- .../mamba/test_flashinfer_mamba_cached_op.py | 7 ++ 4 files changed, 56 insertions(+), 94 deletions(-) diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py index 2e9b26cea69f..fd7cd15c0102 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py @@ -2178,16 +2178,6 @@ def get_constants(cls, source_attn_node: Node) -> List[Constant]: """ return [] - @classmethod - def get_kwarg_cache_keys(cls) -> Set[str]: - """Return cache resource keys that should be passed as kwargs (not positional args). - - By default all cache resources are positional. Override to route specific cache keys - (e.g. replay tensors that must appear after constants in the function signature) to - the call_function kwargs dict instead. - """ - return set() - @classmethod def get_layer_idx(cls, source_attn_node: Node) -> Optional[int]: """Return the logical layer index associated with a source attention node, if any.""" diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/flashinfer_backend_mamba.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/flashinfer_backend_mamba.py index 25defab12f5d..97819c56ef71 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/flashinfer_backend_mamba.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/flashinfer_backend_mamba.py @@ -60,7 +60,7 @@ def _fi_align(t: torch.Tensor) -> torch.Tensor: @torch.library.custom_op( "auto_deploy::flashinfer_cached_ssm", - mutates_args=("ssm_state_cache",), + mutates_args=("ssm_state_cache", "intermediate_ssm_state_cache"), ) def _flashinfer_cached_ssm( # INPUTS (dense but may be flattened across sequences) @@ -83,21 +83,20 @@ def _flashinfer_cached_ssm( seq_idx_prefill: torch.Tensor, # [1, num_prefill_tokens] # CACHES ssm_state_cache: torch.Tensor, # [max_batch_size, num_heads, head_dim, ssm_state_size] - # CONSTANTS — Optional defaults so infer_schema accepts them; always supplied by the transform. - # Positions 16-17: constants come right after the single positional cache node (ssm_state_cache). - time_step_limit: Optional[List[float]] = None, - chunk_size: int = 256, - # SPEC / REPLAY CACHES — all kwargs, absent in replay mode or non-replay mode respectively. - # Lowercase names to match PyTorch FX node name normalization. intermediate_ssm_state_cache: Optional[ torch.Tensor - ] = None, # kwarg; present in non-replay, absent in replay (uses None default) - replay_old_x: Optional[torch.Tensor] = None, # [max_batch, T, nheads, head_dim] - replay_old_b: Optional[torch.Tensor] = None, # [max_batch, 2, T, ngroups, dstate] - replay_old_dt: Optional[torch.Tensor] = None, # [max_batch, 2, nheads, T] fp32 - replay_old_da_cumsum: Optional[torch.Tensor] = None, # [max_batch, 2, nheads, T] fp32 - replay_cache_buf_idx: Optional[torch.Tensor] = None, # [max_batch] int32 - replay_prev_num_accepted: Optional[torch.Tensor] = None, # [max_batch] int32 + ], # [spec_state_size, max_draft_len+1, num_heads, head_dim, d_state]; None in replay mode + replay_old_x: Optional[torch.Tensor], # [max_batch, T, nheads, head_dim]; None in non-replay + replay_old_b: Optional[torch.Tensor], # [max_batch, 2, T, ngroups, dstate]; None in non-replay + replay_old_dt: Optional[torch.Tensor], # [max_batch, 2, nheads, T] fp32; None in non-replay + replay_old_da_cumsum: Optional[ + torch.Tensor + ], # [max_batch, 2, nheads, T] fp32; None in non-replay + replay_cache_buf_idx: Optional[torch.Tensor], # [max_batch] int32; None in non-replay + replay_prev_num_accepted: Optional[torch.Tensor], # [max_batch] int32; None in non-replay + # CONSTANTS + time_step_limit: List[float], + chunk_size: int, out: Optional[torch.Tensor] = None, ) -> torch.Tensor: b, s, num_heads, head_dim, bs, hs_flat, B_flat, C_flat, dt_flat = _flatten_ssm_inputs( @@ -145,7 +144,7 @@ def _flashinfer_cached_ssm( chunk_offsets, seq_idx_prefill, ssm_state_cache, - time_step_limit or [], + time_step_limit, chunk_size, preallocated_ssm_out[:num_prefill_tokens].unsqueeze(0), ) @@ -332,20 +331,20 @@ def _flashinfer_cached_ssm_fake( seq_idx_prefill: torch.Tensor, # [1, num_prefill_tokens] # CACHES ssm_state_cache: torch.Tensor, # [max_batch_size, num_heads, head_dim, ssm_state_size] - # CONSTANTS - time_step_limit: Optional[List[float]] = None, - chunk_size: int = 256, - # SPEC / REPLAY CACHES — all kwargs, absent in replay mode or non-replay mode respectively. - # Lowercase names to match PyTorch FX node name normalization. intermediate_ssm_state_cache: Optional[ torch.Tensor - ] = None, # [spec_state_size, max_draft_len+1, num_heads, head_dim, d_state] - replay_old_x: Optional[torch.Tensor] = None, # [max_batch, T, nheads, head_dim] - replay_old_b: Optional[torch.Tensor] = None, # [max_batch, 2, T, ngroups, dstate] - replay_old_dt: Optional[torch.Tensor] = None, # [max_batch, 2, nheads, T] fp32 - replay_old_da_cumsum: Optional[torch.Tensor] = None, # [max_batch, 2, nheads, T] fp32 - replay_cache_buf_idx: Optional[torch.Tensor] = None, # [max_batch] int32 - replay_prev_num_accepted: Optional[torch.Tensor] = None, # [max_batch] int32 + ], # [spec_state_size, max_draft_len+1, num_heads, head_dim, d_state]; None in replay mode + replay_old_x: Optional[torch.Tensor], # [max_batch, T, nheads, head_dim]; None in non-replay + replay_old_b: Optional[torch.Tensor], # [max_batch, 2, T, ngroups, dstate]; None in non-replay + replay_old_dt: Optional[torch.Tensor], # [max_batch, 2, nheads, T] fp32; None in non-replay + replay_old_da_cumsum: Optional[ + torch.Tensor + ], # [max_batch, 2, nheads, T] fp32; None in non-replay + replay_cache_buf_idx: Optional[torch.Tensor], # [max_batch] int32; None in non-replay + replay_prev_num_accepted: Optional[torch.Tensor], # [max_batch] int32; None in non-replay + # CONSTANTS + time_step_limit: List[float], + chunk_size: int, out: Optional[torch.Tensor] = None, ): if out is not None: @@ -368,26 +367,6 @@ class FlashinferBackendSSM(BaseBackendSSM): # Disabled automatically when: block reuse enabled, tree attention, or SM < 80. ssm_replay: bool = False - # Cache keys always passed as kwargs (follow constants at positions 16-17). - # intermediate_ssm_state_cache is always a kwarg: present in non-replay, absent in replay - # (uses function default None). Replay keys are present only when ssm_replay=True. - # All keys lowercase to match PyTorch FX node name normalization. - _KWARG_CACHE_KEYS = frozenset( - { - "intermediate_ssm_state_cache", - "replay_old_x", - "replay_old_b", - "replay_old_dt", - "replay_old_da_cumsum", - "replay_cache_buf_idx", - "replay_prev_num_accepted", - } - ) - - @classmethod - def get_kwarg_cache_keys(cls): - return cls._KWARG_CACHE_KEYS - @classmethod def get_cached_attention_op(cls) -> MHACallable: return torch.ops.auto_deploy.flashinfer_cached_ssm.default @@ -409,21 +388,19 @@ def get_cache_initializers( "Consider using 'triton_ssm' backend instead." ) + ssm_h = ret["ssm_state_cache"] + + # All 7 optional caches are always registered positionally (None = unused in this mode). + # intermediate_ssm_state_cache: real in non-replay, None in replay. + # replay_old_*: real in replay mode (SM80+), None otherwise. if cls.ssm_replay and get_sm_version() >= 80: - # Replay mode: intermediate_ssm_state_cache is absent from the graph — it uses the - # function's default (None) via kwarg routing. Only replay buffers are registered. - # T is determined by MambaHybridCacheManager from spec_config at construction time. - ssm_h = ret["ssm_state_cache"] + ret["intermediate_ssm_state_cache"] = None x_dtype = torch.bfloat16 - + B_fake = source_attn_node.args[2].meta["val"] + n_groups = B_fake.shape[-2] if B_fake.ndim >= 4 else 1 ret["replay_old_x"] = ReplayOldXHandler( num_heads=ssm_h.num_heads, head_dim=ssm_h.head_dim, dtype=x_dtype ) - # Derive n_groups from B tensor shape at the source node. - # Use lowercase keys: FX lowercases node names, so "replay_old_B" → - # node.name "replay_old_b". Use lowercase to keep _caches / node names consistent. - B_fake = source_attn_node.args[2].meta["val"] - n_groups = B_fake.shape[-2] if B_fake.ndim >= 4 else 1 ret["replay_old_b"] = ReplayOldBHandler( n_groups=n_groups, d_state=ssm_h.d_state, dtype=x_dtype ) @@ -432,9 +409,11 @@ def get_cache_initializers( ret["replay_cache_buf_idx"] = ReplayCacheBufIdxHandler() ret["replay_prev_num_accepted"] = ReplayPrevNumAcceptedHandler() else: - # Non-replay mode: include intermediate_ssm_state_cache as a kwarg resource. - # It is absent in replay mode, where the function default (None) takes over. - ret["intermediate_ssm_state_cache"] = SpecSSMResourceHandler.from_base( - ret["ssm_state_cache"] - ) + ret["intermediate_ssm_state_cache"] = SpecSSMResourceHandler.from_base(ssm_h) + ret["replay_old_x"] = None + ret["replay_old_b"] = None + ret["replay_old_dt"] = None + ret["replay_old_da_cumsum"] = None + ret["replay_cache_buf_idx"] = None + ret["replay_prev_num_accepted"] = None return ret diff --git a/tensorrt_llm/_torch/auto_deploy/transform/library/kvcache.py b/tensorrt_llm/_torch/auto_deploy/transform/library/kvcache.py index 467a484cad97..601c02fc1277 100644 --- a/tensorrt_llm/_torch/auto_deploy/transform/library/kvcache.py +++ b/tensorrt_llm/_torch/auto_deploy/transform/library/kvcache.py @@ -211,11 +211,10 @@ def _insert_cached_attn_node( cache_nodes: List[Node], constants: List[Constant], prepared_attn_mask: Optional[Node] = None, - cache_kwarg_nodes: Optional[Dict[str, Node]] = None, ): """Insert a cached attention node into the graph.""" with gm.graph.inserting_before(attn_node): - all_args = ( + args = ( *qkv_nodes, *meta_nodes_std, *meta_nodes_extra, @@ -223,18 +222,8 @@ def _insert_cached_attn_node( *constants, ) if prepared_attn_mask is not None: - all_args = (*all_args, prepared_attn_mask) - cached_attn_node = gm.graph.call_function( - cached_attn_op, - args=( - *qkv_nodes, - *meta_nodes_std, - *meta_nodes_extra, - *cache_nodes, - *constants, - ), - kwargs=cache_kwarg_nodes or {}, - ) + args = (*args, prepared_attn_mask) + cached_attn_node = gm.graph.call_function(cached_attn_op, args=args) attn_node.replace_all_uses_with(cached_attn_node) gm.graph.erase_node(attn_node) @@ -303,17 +292,15 @@ def _apply( "Each non-shared attention layer must own exactly one cache." ) cache_in_nodes = [] - cache_kwarg_nodes = {} - kwarg_keys = attn_descriptor.get_kwarg_cache_keys() for k, resource_handler in attn_descriptor.get_cache_initializers( attn_node, cm.kv_cache_config ).items(): - resource_name = cm.add_resource(k, resource_handler) - node = self._process_cache_node(gm, resource_name) - if k in kwarg_keys: - # Pass as named kwarg to call_function (e.g. replay args after constants) - cache_kwarg_nodes[k] = node + if resource_handler is None: + # None sentinel: pass literal None positionally, no resource allocated. + cache_in_nodes.append(None) else: + resource_name = cm.add_resource(k, resource_handler) + node = self._process_cache_node(gm, resource_name) cache_in_nodes.append(node) if layer_idx is not None: cache_nodes_by_layer_idx[layer_idx] = cache_in_nodes @@ -342,7 +329,6 @@ def _apply( cache_in_nodes, constants, prepared_mask, - cache_kwarg_nodes, ) num_cached_attn_replacements += 1 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 8a29a0167217..285f2ccdcc6f 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 @@ -104,6 +104,13 @@ def test_flashinfer_decode_matches_triton(mamba_env): None, # seq_idx_prefill # CACHES ssm_state_cache_flashinfer, + None, # intermediate_ssm_state_cache + None, # replay_old_x + None, # replay_old_b + None, # replay_old_dt + None, # replay_old_da_cumsum + None, # replay_cache_buf_idx + None, # replay_prev_num_accepted # CONSTANTS time_step_limit, chunk_size, From 142d75016bcb26fb68add93f0774e6fb9da98534 Mon Sep 17 00:00:00 2001 From: Gal Hubara Agam <96368689+galagam@users.noreply.github.com> Date: Sun, 24 May 2026 03:13:08 -0700 Subject: [PATCH 07/17] [None][refactor] move SSM replay mode flag into BatchInfo Replace `use_replay = replay_old_x is not None` with a dedicated BatchInfo slot (slot 13) so runtime branching in _flashinfer_cached_ssm is driven by an explicit flag rather than by the presence/absence of an Optional cache tensor. Consistent with how num_extend, gather_required, and max_draft_len already live in BatchInfo. - BatchInfo: add slot 13 (use_replay), update_use_replay(), is_use_replay() - CachedSequenceInterface: initialize slot to False in __init__; set it to len(replay_old_x) > 0 in initialize_cache after resources are identified - flashinfer_backend_mamba: read batch_info.is_use_replay() instead of inspecting replay_old_x; update ssm_replay comment to clarify it is a compile-time allocation flag only Signed-off-by: Gal Hubara Agam <96368689+galagam@users.noreply.github.com> --- .../custom_ops/attention_interface.py | 17 +++++++++++++++-- .../mamba/flashinfer_backend_mamba.py | 9 +++++---- .../_torch/auto_deploy/shim/interface.py | 6 ++++++ 3 files changed, 26 insertions(+), 6 deletions(-) diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py index fd7cd15c0102..80a9cfd69559 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py @@ -379,7 +379,7 @@ class BatchInfo: Args: batch_info_host: The batch info tensor on the host. - The information is stored in a 12-element batch_info_host tensor as follows: + The information is stored in a 14-element batch_info_host tensor as follows: Slots 0-5 (batch composition): - [0] num_prefill: number of prefill requests @@ -402,10 +402,13 @@ class BatchInfo: Slot 12 (spec-decoding info, set once at runtime init): - [12] max_draft_len: maximum draft length per request (0 when not spec dec) + Slot 13 (replay mode flag, set once at runtime init): + - [13] use_replay: 1 if SSM replay state-update path is active, 0 otherwise + All fields can be accessed and updated with the convenience functions below. """ - _NUM_ELEMENTS = 13 + _NUM_ELEMENTS = 14 def __init__(self, batch_info_host: Optional[torch.Tensor] = None): if batch_info_host is None: @@ -528,6 +531,16 @@ def update_max_draft_len(self, max_draft_len: int) -> None: def get_max_draft_len(self) -> int: return int(self._batch_info[12]) + # --- replay mode flag (slot 13) writer --- + + def update_use_replay(self, use_replay: bool) -> None: + self._batch_info[13] = int(use_replay) + + # --- replay mode flag (slot 13) reader --- + + def is_use_replay(self) -> bool: + return bool(self._batch_info[13]) + class SequenceInfo: """An interface to hold information about how the sequence is laid out and stored in cache. diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/flashinfer_backend_mamba.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/flashinfer_backend_mamba.py index 97819c56ef71..6e7a35eb4b64 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/flashinfer_backend_mamba.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/flashinfer_backend_mamba.py @@ -115,7 +115,7 @@ def _flashinfer_cached_ssm( # NOTE: intermediate_state_indices_extend (torch.arange) is NOT hoisted here. # It is only needed for the non-replay FlashInfer extend path. Hoisting it # unconditionally would create an unnecessary arange kernel in the gap for - # the replay path (where replay_old_x is not None and the arange is unused). + # the replay path (where use_replay=True and the arange is unused). if out is not None: preallocated_ssm_out = out.view(bs, num_heads, head_dim) @@ -187,7 +187,7 @@ def _flashinfer_cached_ssm( slot_idx_extend_i32 = slot_idx_extend # already int32; no .to() kernel - use_replay = replay_old_x is not None + use_replay = batch_info.is_use_replay() if use_replay: # Replay path: fast-forward SSM state via tl.dot on cached values. # State is updated in-place; no disable_state_update needed. @@ -363,8 +363,9 @@ def _flashinfer_cached_ssm_fake( @AttentionRegistry.register("flashinfer_ssm") class FlashinferBackendSSM(BaseBackendSSM): - # When ssm_replay=True, use the replay kernel (tl.dot fast-forward) instead of FlashInfer. - # Disabled automatically when: block reuse enabled, tree attention, or SM < 80. + # ssm_replay=True allocates replay buffers at compile time (SM80+), which causes + # BatchInfo.is_use_replay() to return True at runtime. Never check ssm_replay at + # runtime — use batch_info.is_use_replay() instead. ssm_replay: bool = False @classmethod diff --git a/tensorrt_llm/_torch/auto_deploy/shim/interface.py b/tensorrt_llm/_torch/auto_deploy/shim/interface.py index 749bd0dd2a54..b9aece63792d 100644 --- a/tensorrt_llm/_torch/auto_deploy/shim/interface.py +++ b/tensorrt_llm/_torch/auto_deploy/shim/interface.py @@ -151,6 +151,8 @@ def __init__( self.info.batch_info.update_max_draft_len( spec_config.max_draft_len if spec_config is not None else 0 ) + # Default to non-replay; updated in initialize_cache once resources are identified. + self.info.batch_info.update_use_replay(False) @property def args(self) -> Tuple[torch.Tensor, ...]: @@ -814,6 +816,10 @@ def _create_kv_cache_manager(self, max_tokens: Optional[int] = None) -> Dict: replay_prev_num_accepted, ) = self._identify_managed_state_resources() + # Propagate replay mode into BatchInfo so SSM backends can branch on it + # without inspecting the presence/absence of Optional cache tensors. + self.info.batch_info.update_use_replay(len(replay_old_x) > 0) + # 2. Prepare configuration kv_cache_config = self._prepare_kv_cache_config(max_tokens, kv_managed) kv_cache_kwargs = self._build_kv_cache_kwargs(kv_ref, kv_managed, kv_cache_config) From 019205983c27b3918b059ff69afa83d5c7bd3e0a Mon Sep 17 00:00:00 2001 From: Gal Hubara Agam <96368689+galagam@users.noreply.github.com> Date: Sun, 24 May 2026 03:53:09 -0700 Subject: [PATCH 08/17] [None][fix] flashinfer_ssm: fix head-dim validation and replay buffer registration Two related issues in FlashinferBackendSSM.get_cache_initializers: 1. Head-dim validation was gated on `not cls.ssm_replay`, so it was skipped whenever ssm_replay=True even if SM < 80 caused the replay path to be inactive and FlashInfer would still be used at runtime. 2. The cache registration branch checked `cls.ssm_replay and SM >= 80` while the validation used a different condition, creating an inconsistency. Fix: compute `use_replay = cls.ssm_replay and get_sm_version() >= 80` once and drive both the head-dim check (`not use_replay`) and the replay-vs-FlashInfer buffer registration (`if use_replay`) from the same boolean, so unsupported head dims can never reach _flashinfer_ssm_update. Also add replay buffers (replay_old_x/b/dt/da_cumsum) to mutates_args since the precompute and main kernels write to them in-place. Signed-off-by: Gal Hubara Agam <96368689+galagam@users.noreply.github.com> --- .../mamba/flashinfer_backend_mamba.py | 24 +++++++++++++------ 1 file changed, 17 insertions(+), 7 deletions(-) diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/flashinfer_backend_mamba.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/flashinfer_backend_mamba.py index 6e7a35eb4b64..a78062876731 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/flashinfer_backend_mamba.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/flashinfer_backend_mamba.py @@ -60,7 +60,16 @@ def _fi_align(t: torch.Tensor) -> torch.Tensor: @torch.library.custom_op( "auto_deploy::flashinfer_cached_ssm", - mutates_args=("ssm_state_cache", "intermediate_ssm_state_cache"), + mutates_args=( + "ssm_state_cache", + "intermediate_ssm_state_cache", + # replay buffers: written in-place by the precompute and main kernels + # (double-buffered B/dt/dA_cumsum, single-buffered x); None in non-replay mode + "replay_old_x", + "replay_old_b", + "replay_old_dt", + "replay_old_da_cumsum", + ), ) def _flashinfer_cached_ssm( # INPUTS (dense but may be flattened across sequences) @@ -378,11 +387,12 @@ def get_cache_initializers( ) -> ResourceHandlerDict: ret = super().get_cache_initializers(source_attn_node, cache_config) - # In replay mode we never call FlashInfer, so head_dim can be anything. - if ( - not cls.ssm_replay - and ret["ssm_state_cache"].head_dim not in FLASHINFER_SUPPORTED_HEAD_DIMS - ): + # Replay requires both ssm_replay=True and SM >= 80 (Ampere+). + # Only when use_replay=True do we skip FlashInfer at runtime; otherwise + # _flashinfer_ssm_update is called and head_dim must be supported. + use_replay = cls.ssm_replay and get_sm_version() >= 80 + + if not use_replay and ret["ssm_state_cache"].head_dim not in FLASHINFER_SUPPORTED_HEAD_DIMS: raise ValueError( f"flashinfer_ssm only supports head_dim in {FLASHINFER_SUPPORTED_HEAD_DIMS}. " f"Got head_dim={ret['ssm_state_cache'].head_dim}. " @@ -394,7 +404,7 @@ def get_cache_initializers( # All 7 optional caches are always registered positionally (None = unused in this mode). # intermediate_ssm_state_cache: real in non-replay, None in replay. # replay_old_*: real in replay mode (SM80+), None otherwise. - if cls.ssm_replay and get_sm_version() >= 80: + if use_replay: ret["intermediate_ssm_state_cache"] = None x_dtype = torch.bfloat16 B_fake = source_attn_node.args[2].meta["val"] From c397065292e79964a1aa1785e39624a4c6d05d99 Mon Sep 17 00:00:00 2001 From: Gal Hubara Agam <96368689+galagam@users.noreply.github.com> Date: Sun, 24 May 2026 04:00:04 -0700 Subject: [PATCH 09/17] [None][fix] assert replay bundle completeness in _identify_managed_state_resources Replay handlers are registered as an all-or-none bundle per SSM layer. Previously only replay_old_x was checked (use_replay = len(replay_old_x) > 0) with no verification that the remaining buffers matched. If any per-layer list (old_B, old_dt, old_dA_cumsum) or global list (cache_buf_idx, prev_num_accepted) had fewer entries than ssm_managed, _create_and_assign_state_views would index past the shorter list causing an IndexError. Add assertions after use_replay is set: each of the six replay lists must have length == len(ssm_managed), catching partial or mismatched registrations early with a clear error message. Signed-off-by: Gal Hubara Agam <96368689+galagam@users.noreply.github.com> --- .../_torch/auto_deploy/shim/interface.py | 20 +++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/tensorrt_llm/_torch/auto_deploy/shim/interface.py b/tensorrt_llm/_torch/auto_deploy/shim/interface.py index b9aece63792d..859140b4dd5e 100644 --- a/tensorrt_llm/_torch/auto_deploy/shim/interface.py +++ b/tensorrt_llm/_torch/auto_deploy/shim/interface.py @@ -447,6 +447,26 @@ def _identify_managed_state_resources( # Exception: in replay mode intermediate_ssm_state_cache is omitted entirely from # get_cache_initializers and uses the function's default (None) via kwarg routing. use_replay = len(replay_old_x) > 0 + if use_replay: + n = len(ssm_managed) + for lst, name in [ + (replay_old_x, "replay_old_x"), + (replay_old_B, "replay_old_B"), + (replay_old_dt, "replay_old_dt"), + (replay_old_dA_cumsum, "replay_old_dA_cumsum"), + ]: + assert len(lst) == n, ( + f"Replay bundle mismatch: {name} has {len(lst)} entries, " + f"expected {n} (== len(ssm_managed))" + ) + assert len(replay_cache_buf_idx) == n, ( + f"Replay bundle mismatch: replay_cache_buf_idx has {len(replay_cache_buf_idx)} " + f"entries, expected {n} (== len(ssm_managed))" + ) + assert len(replay_prev_num_accepted) == n, ( + f"Replay bundle mismatch: replay_prev_num_accepted has " + f"{len(replay_prev_num_accepted)} entries, expected {n} (== len(ssm_managed))" + ) if self._spec_config is not None: if not use_replay: assert len(ssm_spec) == len(ssm_managed), ( From 8911587fa4d5583d61b8d78d6b48394c28d73419 Mon Sep 17 00:00:00 2001 From: Gal Hubara Agam <96368689+galagam@users.noreply.github.com> Date: Sun, 24 May 2026 06:19:04 -0700 Subject: [PATCH 10/17] [None][test] add extend+replay unit test for flashinfer_cached_ssm Add test_flashinfer_extend_replay_calls_replay_kernel, parametrized over head_dim={32, 64}: constructs an extend-only batch with BatchInfo.is_use_replay()=True and all six replay buffers allocated, then uses patch(wraps=) to spy on replay_selective_state_update and asserts it is called exactly once. head_dim=32 (not in FLASHINFER_SUPPORTED_HEAD_DIMS) proves the replay kernel bypasses the FlashInfer head-dim constraint; head_dim=64 covers production-like configs. Update test_super_mtp_ssm_replay_smoke with an accurate docstring: the test covers compilation and full-stack init (transforms, cache manager with replay buffers, MTP Eagle inference); the runtime replay kernel path is now covered by the new unit test. Signed-off-by: Gal Hubara Agam <96368689+galagam@users.noreply.github.com> --- .../mamba/test_flashinfer_mamba_cached_op.py | 113 ++++++++++++++++++ .../smoke/test_ad_speculative_decoding.py | 13 +- 2 files changed, 121 insertions(+), 5 deletions(-) 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 285f2ccdcc6f..b8db6132dc0b 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 @@ -12,12 +12,17 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +from unittest.mock import patch + import pytest import torch 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.modules.mamba.replay_selective_state_update import ( + replay_selective_state_update as _real_replay_selective_state_update, +) @pytest.fixture @@ -126,3 +131,111 @@ def test_flashinfer_decode_matches_triton(mamba_env): assert torch.allclose( after_flashinfer.to(after_triton.dtype), after_triton, atol=atol, rtol=rtol ) + + +@pytest.mark.parametrize( + "head_dim", + [ + pytest.param(32, id="unsupported_flashinfer_headdim"), # not in {64, 128} + pytest.param(64, id="supported_flashinfer_headdim"), # in {64, 128} + ], +) +def test_flashinfer_extend_replay_calls_replay_kernel(mamba_env, head_dim): + """Verify replay_selective_state_update is invoked on the extend+replay path. + + Parametrized over head_dim to cover both a dim that FlashInfer does not support + (proving replay is independent of the FlashInfer head-dim constraint) and one + that it does support (ensuring the path works for production-like configs too). + """ + device = mamba_env["device"] + dtype = mamba_env["dtype"] + + num_extend = 1 + tokens_per_extend = 2 # draft tokens per MTP verification step + num_heads = 4 + n_groups, ssm_state_size = 2, 16 # ssm_state_size >= 16 for replay tl.dot + max_batch_size = 2 # cache slots + + (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 + ) + slot_idx = torch.tensor([0], device=device, dtype=torch.int32) + + # Replay buffers: all zeros (first step, nothing cached yet; kernel still runs). + replay_old_x = torch.zeros( + max_batch_size, tokens_per_extend, num_heads, head_dim, device=device, dtype=torch.bfloat16 + ) + replay_old_b = torch.zeros( + max_batch_size, + 2, + tokens_per_extend, + n_groups, + ssm_state_size, + device=device, + dtype=torch.bfloat16, + ) + replay_old_dt = torch.zeros( + max_batch_size, 2, num_heads, tokens_per_extend, device=device, dtype=torch.float32 + ) + replay_old_da_cumsum = torch.zeros( + max_batch_size, 2, num_heads, tokens_per_extend, 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) + + # Extend-only batch with replay mode enabled. + _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.tensor([0, 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) + + _TARGET = ( + "tensorrt_llm._torch.auto_deploy.custom_ops.mamba" + ".flashinfer_backend_mamba.replay_selective_state_update" + ) + with patch(_TARGET, wraps=_real_replay_selective_state_update) as mock_replay: + 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, + # CONSTANTS + time_step_limit, + chunk_size, + ) + + assert mock_replay.call_count == num_extend, ( + f"Expected {num_extend} replay call(s), got {mock_replay.call_count}" + ) + assert out.shape == hidden_states.shape + assert torch.isfinite(out).all() diff --git a/tests/unittest/auto_deploy/singlegpu/smoke/test_ad_speculative_decoding.py b/tests/unittest/auto_deploy/singlegpu/smoke/test_ad_speculative_decoding.py index 5ceb84ed62fe..a589f8261499 100644 --- a/tests/unittest/auto_deploy/singlegpu/smoke/test_ad_speculative_decoding.py +++ b/tests/unittest/auto_deploy/singlegpu/smoke/test_ad_speculative_decoding.py @@ -93,11 +93,14 @@ def test_super_mtp_smoke(): def test_super_mtp_ssm_replay_smoke(): - """Test one-model MTP with flashinfer_ssm + ssm_replay=True. - - Uses mamba_head_dim=64 (required by FlashInfer) and ssm_state_size=64 - (>= 16 for the replay kernel's tl.dot constraint), actually executing - _replay_precompute_kernel and _replay_state_update_kernel. + """Smoke test: MTP Eagle one-model with flashinfer_ssm + ssm_replay=True compiles and runs. + + Verifies that the full pipeline — transforms, cache manager init with replay buffers, + and MTP inference — completes without error. The AD SSM custom ops are not directly + invoked at runtime in this configuration (Eagle3OneModelSampler drives its own forward + loop); the replay kernel path is covered by test_flashinfer_extend_replay_calls_replay_kernel. + Uses mamba_head_dim=64 and ssm_state_size=64 to satisfy FlashInfer constraints on the + decode path (which IS called in this config). """ test_prompt = "What is the capital of France?" model_hub_id = "nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-BF16" From 58c3bec828e42cfa303ace2d71587520fcb643a3 Mon Sep 17 00:00:00 2001 From: Gal Hubara Agam <96368689+galagam@users.noreply.github.com> Date: Sun, 24 May 2026 07:09:17 -0700 Subject: [PATCH 11/17] [None][cleanup] drop redundant slot_idx alias comments in flashinfer_backend_mamba Signed-off-by: Gal Hubara Agam <96368689+galagam@users.noreply.github.com> --- .../mamba/flashinfer_backend_mamba.py | 19 ++++++++----------- 1 file changed, 8 insertions(+), 11 deletions(-) diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/flashinfer_backend_mamba.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/flashinfer_backend_mamba.py index a78062876731..b1b79acef781 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/flashinfer_backend_mamba.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/flashinfer_backend_mamba.py @@ -167,7 +167,7 @@ def _flashinfer_cached_ssm( A, D, dt_bias, - slot_idx_i32, # already int32 — slice is a no-op, no extra kernel + slot_idx_i32, seq_start=num_prefill, token_start=num_prefill_tokens, num_seq=num_extend, @@ -180,7 +180,7 @@ def _flashinfer_cached_ssm( if extend_inputs is not None: tokens_per_extend = num_extend_tokens // num_extend ( - slot_idx_extend, # already int32 (sliced from slot_idx_i32) + slot_idx_extend, x_extend, B_extend, C_extend, @@ -194,8 +194,6 @@ def _flashinfer_cached_ssm( num_prefill_tokens : num_prefill_tokens + num_extend_tokens ].view(num_extend, tokens_per_extend, num_heads, head_dim) - slot_idx_extend_i32 = slot_idx_extend # already int32; no .to() kernel - use_replay = batch_info.is_use_replay() if use_replay: # Replay path: fast-forward SSM state via tl.dot on cached values. @@ -221,7 +219,7 @@ def _flashinfer_cached_ssm( D=D_full, dt_bias=dt_bias_hp, dt_softplus=True, - state_batch_indices=slot_idx_extend_i32, + state_batch_indices=slot_idx_extend, launch_with_pdl=True, # PDL chain: triton_causal_conv extend → precompute → main ) else: @@ -234,7 +232,7 @@ def _flashinfer_cached_ssm( # Only needed for the non-replay FlashInfer path; not hoisted to avoid # an unnecessary arange kernel in the stream gap for the replay path. intermediate_state_indices = torch.arange( - num_extend, dtype=torch.int32, device=slot_idx_extend_i32.device + num_extend, dtype=torch.int32, device=slot_idx_extend.device ) _flashinfer_ssm_update( ssm_state_cache, @@ -247,7 +245,7 @@ def _flashinfer_cached_ssm( z=None, dt_bias=dt_bias_hp, dt_softplus=True, - state_batch_indices=slot_idx_extend_i32, + state_batch_indices=slot_idx_extend, out=preallocated_ssm_out_e, disable_state_update=True, intermediate_states_buffer=intermediate_ssm_state_cache, @@ -264,7 +262,7 @@ def _flashinfer_cached_ssm( A, D, dt_bias, - slot_idx_i32, # already int32 — no .to() kernel at decode time + slot_idx_i32, num_prefill + num_extend, num_prefill_tokens + num_extend_tokens, num_decode, @@ -276,7 +274,7 @@ def _flashinfer_cached_ssm( if decode_inputs is not None: ( - slot_idx_decode, # already int32 (sliced from slot_idx_i32) + slot_idx_decode, x_decode, B_decode, C_decode, @@ -290,7 +288,6 @@ def _flashinfer_cached_ssm( B_decode = _fi_align(B_decode) C_decode = _fi_align(C_decode) - slot_idx_decode_i32 = slot_idx_decode # already int32; no .to() kernel y_decode = _flashinfer_ssm_update( ssm_state_cache, x_decode, @@ -302,7 +299,7 @@ def _flashinfer_cached_ssm( z=None, dt_bias=dt_bias_hp, dt_softplus=True, - state_batch_indices=slot_idx_decode_i32, + state_batch_indices=slot_idx_decode, ) preallocated_ssm_out[num_prefill_tokens + num_extend_tokens : num_total_tokens].copy_( y_decode From 534079d3e4875b2a7c43282c58f619607190d072 Mon Sep 17 00:00:00 2001 From: Gal Hubara Agam <96368689+galagam@users.noreply.github.com> Date: Sun, 24 May 2026 09:54:37 -0700 Subject: [PATCH 12/17] [None][test] add flashinfer_mamba_cached_op to B200 CI The replay kernel (replay_selective_state_update) requires SM>=80; H100 CI already covers this file via the singlegpu/custom_ops directory entry, but B200 only lists individual files so the test was missing. Add it explicitly so the Blackwell-specific replay path is covered on B200. Signed-off-by: Gal Hubara Agam <96368689+galagam@users.noreply.github.com> --- tests/integration/test_lists/test-db/l0_b200.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 459fa3073bda..a56f4c0d7ec6 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -330,6 +330,7 @@ l0_b200: # custom_ops, smoke, transformations files) and pure-FP8 tests are covered # on Hopper (l0_h100.yml) and not duplicated here. - unittest/auto_deploy/singlegpu/custom_ops/attention/test_triton_paged_attention.py::TestSDPADispatch + - unittest/auto_deploy/singlegpu/custom_ops/mamba/test_flashinfer_mamba_cached_op.py - unittest/auto_deploy/singlegpu/custom_ops/moe/test_ad_moe_op.py - unittest/auto_deploy/singlegpu/custom_ops/moe/test_trtllm_moe.py - unittest/auto_deploy/singlegpu/custom_ops/quantization/test_quant.py From ad62b701bdf8337065d37480caaea609174e9084 Mon Sep 17 00:00:00 2001 From: Gal Hubara Agam <96368689+galagam@users.noreply.github.com> Date: Sun, 24 May 2026 10:54:34 -0700 Subject: [PATCH 13/17] [None][fix] low_memory_overrides: set batch_sizes and max_batch_size in cuda_graph_config Setting only max_batch_size conflicted with YAMLs that already specify batch_sizes (e.g. super_v3.yaml has [1..384]): Pydantic requires max_batch_size == max(batch_sizes) when both are present. Override both fields so the filtered batch_sizes list replaces the YAML's list and max_batch_size stays consistent with it. Signed-off-by: Gal Hubara Agam <96368689+galagam@users.noreply.github.com> --- tests/integration/defs/accuracy/test_llm_api_autodeploy.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/integration/defs/accuracy/test_llm_api_autodeploy.py b/tests/integration/defs/accuracy/test_llm_api_autodeploy.py index e279416f6768..35697cd05fb5 100644 --- a/tests/integration/defs/accuracy/test_llm_api_autodeploy.py +++ b/tests/integration/defs/accuracy/test_llm_api_autodeploy.py @@ -150,11 +150,15 @@ def low_memory_overrides(config, max_seq_len=8192, max_num_tokens=8192): """Update and return config that reduce memory footprint for unquantized (bf16) runs.""" + cuda_graph_batch_sizes = [ + s for s in [1, 2, 4, 8, 16, 32, 64, 128] if s <= max_batch_size + ] config.update({ "max_batch_size": max_batch_size, "max_seq_len": max_seq_len, "max_num_tokens": max_num_tokens, "cuda_graph_config": { + "batch_sizes": cuda_graph_batch_sizes, "max_batch_size": max_batch_size, }, }) From f9c7c87f854d3ad3c40ebb0a0b6212186b7ebedc Mon Sep 17 00:00:00 2001 From: Gal Hubara Agam <96368689+galagam@users.noreply.github.com> Date: Thu, 28 May 2026 05:11:02 -0700 Subject: [PATCH 14/17] drop changes committed by mistake Signed-off-by: Gal Hubara Agam <96368689+galagam@users.noreply.github.com> --- .../configs/curated/super_v3_mtp_ws1.yaml | 18 ------------------ .../configs/curated/super_v3_mtp_ws4.yaml | 19 ------------------- 2 files changed, 37 deletions(-) delete mode 100644 examples/configs/curated/super_v3_mtp_ws1.yaml delete mode 100644 examples/configs/curated/super_v3_mtp_ws4.yaml diff --git a/examples/configs/curated/super_v3_mtp_ws1.yaml b/examples/configs/curated/super_v3_mtp_ws1.yaml deleted file mode 100644 index 61f47b3b92b7..000000000000 --- a/examples/configs/curated/super_v3_mtp_ws1.yaml +++ /dev/null @@ -1,18 +0,0 @@ -# Nemotron-3-Super with MTP speculative decoding, WS=1 (single GPU). -max_batch_size: 64 -max_seq_len: 8192 -max_num_tokens: 2048 -tensor_parallel_size: 1 -moe_expert_parallel_size: 1 -trust_remote_code: true -enable_chunked_prefill: false -cuda_graph_config: - enable_padding: true - max_batch_size: 64 -kv_cache_config: - enable_block_reuse: false - mamba_ssm_cache_dtype: auto -num_postprocess_workers: 4 -speculative_config: - decoding_type: MTP - num_nextn_predict_layers: 6 diff --git a/examples/configs/curated/super_v3_mtp_ws4.yaml b/examples/configs/curated/super_v3_mtp_ws4.yaml deleted file mode 100644 index 2d2a9b4390f5..000000000000 --- a/examples/configs/curated/super_v3_mtp_ws4.yaml +++ /dev/null @@ -1,19 +0,0 @@ -# Nemotron-3-Super with MTP speculative decoding, WS=4 (4-GPU TP+EP). -max_batch_size: 64 -max_seq_len: 8192 -max_num_tokens: 2048 -tensor_parallel_size: 4 -moe_expert_parallel_size: 4 -trust_remote_code: true -enable_attention_dp: false -enable_chunked_prefill: false -cuda_graph_config: - enable_padding: true - max_batch_size: 64 -kv_cache_config: - enable_block_reuse: false - mamba_ssm_cache_dtype: auto -num_postprocess_workers: 4 -speculative_config: - decoding_type: MTP - num_nextn_predict_layers: 6 From 684f72f8d13babd797f728bcbd64bd04b0e8f683 Mon Sep 17 00:00:00 2001 From: Gal Hubara Agam <96368689+galagam@users.noreply.github.com> Date: Thu, 28 May 2026 09:37:30 -0700 Subject: [PATCH 15/17] fix merge Signed-off-by: Gal Hubara Agam <96368689+galagam@users.noreply.github.com> --- tensorrt_llm/_torch/auto_deploy/shim/interface.py | 4 ++-- tensorrt_llm/_torch/auto_deploy/transform/library/kvcache.py | 3 --- 2 files changed, 2 insertions(+), 5 deletions(-) diff --git a/tensorrt_llm/_torch/auto_deploy/shim/interface.py b/tensorrt_llm/_torch/auto_deploy/shim/interface.py index e446eb185e15..53662854af51 100644 --- a/tensorrt_llm/_torch/auto_deploy/shim/interface.py +++ b/tensorrt_llm/_torch/auto_deploy/shim/interface.py @@ -997,8 +997,8 @@ def _create_kv_cache_manager(self, max_tokens: Optional[int] = None) -> Dict: 1. the final number of tokens is synced (min) across ranks 2. rounding for getting a multiple of tokens_per_block """ - # 1. Identify managed resources - kv_ref, kv_managed = self._identify_managed_kv_resources() + # 1. Identify managed resources and per-pool configurations + kv_managed, pool_configurations = self._identify_managed_kv_resources() ( ssm_ref, ssm_managed, diff --git a/tensorrt_llm/_torch/auto_deploy/transform/library/kvcache.py b/tensorrt_llm/_torch/auto_deploy/transform/library/kvcache.py index 70be521330f0..760fd46378c0 100644 --- a/tensorrt_llm/_torch/auto_deploy/transform/library/kvcache.py +++ b/tensorrt_llm/_torch/auto_deploy/transform/library/kvcache.py @@ -425,9 +425,6 @@ def _apply( ) cache_in_nodes = [] group_idx = 0 - for k, resource_handler in attn_descriptor.get_cache_initializers( - attn_node, cm.kv_cache_config - ).items(): for k, resource_handler in attn_descriptor.get_cache_initializers( attn_node, cm.kv_cache_config ).items(): From fa5c865b98add89b64c5277a362235e90bb1612f Mon Sep 17 00:00:00 2001 From: Gal Hubara Agam <96368689+galagam@users.noreply.github.com> Date: Mon, 1 Jun 2026 02:14:24 -0700 Subject: [PATCH 16/17] improved comments Signed-off-by: Gal Hubara Agam <96368689+galagam@users.noreply.github.com> --- examples/auto_deploy/model_registry/configs/super_v3_mtp.yaml | 2 +- .../auto_deploy/custom_ops/mamba/flashinfer_backend_mamba.py | 2 ++ 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/examples/auto_deploy/model_registry/configs/super_v3_mtp.yaml b/examples/auto_deploy/model_registry/configs/super_v3_mtp.yaml index 376b9e65493a..c3b5911ff956 100644 --- a/examples/auto_deploy/model_registry/configs/super_v3_mtp.yaml +++ b/examples/auto_deploy/model_registry/configs/super_v3_mtp.yaml @@ -3,7 +3,7 @@ runtime: trtllm compile_backend: torch-cudagraph max_batch_size: 64 -max_seq_len: 8192 +max_seq_len: 8192 # tunable max_num_tokens: 2048 enable_chunked_prefill: false # TODO: Support chunked prefill w/ spec dec attn_backend: trtllm diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/flashinfer_backend_mamba.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/flashinfer_backend_mamba.py index b1b79acef781..076d1fb4ba4a 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/flashinfer_backend_mamba.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/flashinfer_backend_mamba.py @@ -120,6 +120,8 @@ def _flashinfer_cached_ssm( # Pre-hoist index type conversions so their kernels run before the extend # causal_conv1d_update kernel rather than landing in the gap between it # and the PDL-dependent replay precompute kernel. + # slot_idx is a per-request cache slot index in [0, max_batch_size], so the + # int32 cast can never overflow. slot_idx_i32 = slot_idx.to(torch.int32) # NOTE: intermediate_state_indices_extend (torch.arange) is NOT hoisted here. # It is only needed for the non-replay FlashInfer extend path. Hoisting it From ae7c388a7f4130589035721ddc0ba8406382ab1f Mon Sep 17 00:00:00 2001 From: Gal Hubara Agam <96368689+galagam@users.noreply.github.com> Date: Tue, 2 Jun 2026 01:00:16 -0700 Subject: [PATCH 17/17] remove accidentally committed file Signed-off-by: Gal Hubara Agam <96368689+galagam@users.noreply.github.com> --- .../configs/super_v3_mtp_ws4.yaml | 67 ------------------- 1 file changed, 67 deletions(-) delete mode 100644 examples/auto_deploy/model_registry/configs/super_v3_mtp_ws4.yaml diff --git a/examples/auto_deploy/model_registry/configs/super_v3_mtp_ws4.yaml b/examples/auto_deploy/model_registry/configs/super_v3_mtp_ws4.yaml deleted file mode 100644 index a3294c8a516d..000000000000 --- a/examples/auto_deploy/model_registry/configs/super_v3_mtp_ws4.yaml +++ /dev/null @@ -1,67 +0,0 @@ -# Config for SuperV3 with MTP speculative decoding, WS=4 (4-GPU EP+BMM sharding). -# Derived from super_v3_mtp.yaml with world_size: 4 to enable sharding. -# -runtime: trtllm -compile_backend: torch-cudagraph -world_size: 4 -max_batch_size: 64 -max_seq_len: 8192 -max_num_tokens: 2048 -enable_chunked_prefill: false # TODO: Support chunked prefill w/ spec dec -attn_backend: trtllm -model_factory: AutoModelForCausalLM -skip_loading_weights: false -cuda_graph_config: - max_batch_size: 64 - enable_padding: true -kv_cache_config: - # tunable mamba cache dtype - # --> use float32 for accuracy and default (auto) for speed - mamba_ssm_cache_dtype: auto -speculative_config: - decoding_type: MTP - num_nextn_predict_layers: 6 - mtp_eagle_one_model: true -transforms: - detect_sharding: - allreduce_strategy: NCCL - # NOTE: add 'tp' to sharding dims only for high-throughput runs - # For low-latency, keep mamba and attention replicated - sharding_dims: ['ep', 'bmm'] - # NOTE: sharding_source applies only to TP sharding - sharding_source: ['manual'] - manual_config: - head_dim: 128 - tp_plan: - # mamba SSM layer - "in_proj": "mamba" - "out_proj": "rowwise" - # attention layer - "q_proj": "colwise" - "k_proj": "colwise" - "v_proj": "colwise" - "o_proj": "rowwise" - # moe layer: SHARED experts - "up_proj": "colwise" - "down_proj": "rowwise" - # MoLE: latent projections: simple shard - "fc1_latent_proj": "gather" - "fc2_latent_proj": "gather" - # multi_stream_moe disabled - creates NaNs in target's sampled logits. TODO: investigate + fix - multi_stream_moe: - stage: compile - enabled: false - fuse_mamba_a_log: - stage: post_load_fusion - enabled: true - # FlashInfer SSM is used for the MTP extend path; Triton causal conv is still required - # for speculative Mamba state caching in this configuration. - insert_cached_ssm_attention: - backend: flashinfer_ssm - ssm_replay: true - insert_cached_causal_conv: - backend: triton_causal_conv - fuse_nvfp4_moe: - backend: trtllm_gen - fuse_fp8_moe: - allow_different_input_scales: true