[TRTLLM-16303][refactor] modularize Mamba cache managers

tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/modeling.py

@@ -1870,18 +1870,19 @@ def _bind_step_buffers(self, cm, g: _PoolGeometry, mm, conv0, ssm0) -> dict:
         )
         _require(
             cm,
-            ("get_intermediate_conv_states", "get_intermediate_ssm_states"),
-            "binds its per-step state buffers",
+            ("mamba_layer_cache", "get_replay_state_update_metadata"),
+            "binds its per-step state buffers and replay metadata",
         )
-        # `get_intermediate_*` builds a fresh view per call, so bind them here.
-        int_conv = {i: cm.get_intermediate_conv_states(i) for i in self.gdn_layers}
+        # Bind the per-layer views once for subsequent forwards.
+        layer_caches = {i: cm.mamba_layer_cache(i) for i in self.gdn_layers}
+        int_conv = {i: cache.intermediate_conv_window for i, cache in layer_caches.items()}
         int_idx = getattr(cm, "intermediate_state_indices", None)
         g0 = self.gdn_layers[0]
         check_step_buffers(g, int_conv[g0], conv0, int_idx)
-        int_ssm0 = cm.get_intermediate_ssm_states(g0)
-        if not getattr(cm, "use_replay_state_update", False):
+        int_ssm0 = layer_caches[g0].intermediate_ssm
+        if cm.get_replay_state_update_metadata() is None:
             check_legacy_state(g, int_ssm0, ssm0)
-            int_ssm = {i: cm.get_intermediate_ssm_states(i) for i in self.gdn_layers}
+            int_ssm = {i: cache.intermediate_ssm for i, cache in layer_caches.items()}
             return dict(int_idx=int_idx, int_conv=int_conv, int_ssm=int_ssm)
         assert int_ssm0 is None, (
             "the cache manager runs the cached-replay verify path and also allocated a "
@@ -1893,7 +1894,7 @@ def _bind_replay(self, cm, g: _PoolGeometry, mm, ssm0) -> _Replay:
         """The cached-replay history buffers and bookkeeping of one manager."""
         meta = cm.get_replay_state_update_metadata()
         assert meta is not None, (
-            "use_replay_state_update is True but the manager returns no replay metadata "
+            "the cached-replay binding requires replay metadata "
             "(prev_num_accepted_tokens / cache_buf_idx / widths)"
         )
         history, width = meta.replay_history_size, meta.replay_step_width

tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ssm/fused_recurrent_gated_delta_rule_cached_replay_update.py

@@ -194,7 +194,7 @@ class FusedRecurrentGatedDeltaRuleCachedReplayUpdate(OpWrapper):
       history);
     * after every step, before the next: with work items,
       `commit_gdn_cached_replay_history_layers` over every local GDN layer; then the
-      advance (`_advance_replay_state` in `pyexecutor/kv_cache/mamba_cache_manager.py`)
+      advance (`_advance_replay_state` in `modules/fla/cache_manager.py`)
       with the accepted counts including the bonus token. Advancing as if every
       draft were accepted is silent at the op and wrong from the next step on;
     * the conv state is not this op's: under replay the conv still writes its

tensorrt_llm/_torch/attention/backends/sparse/csa2/metadata.py

@@ -6,6 +6,7 @@
 
 from contextlib import contextmanager, nullcontext
 from dataclasses import dataclass
+from math import prod
 from typing import TYPE_CHECKING
 
 import numpy as np
@@ -22,6 +23,21 @@
 
 _SWA_TILE = 128
 _HEAD_DIM = 512
+_TMA_MIN_BACKING_BYTES = 128 * 1024
+
+
+def _allocate_staging_pool(
+    shape: tuple[int, ...], dtype: torch.dtype, device: torch.device
+) -> torch.Tensor:
+    # Native sparse FMHA describes an unbounded row domain. Blackwell TMA
+    # out-of-bounds fetches require 128 KiB mapped after the pool address,
+    # including when the logical staging pool is smaller.
+    elements = prod(shape)
+    element_bytes = 2 if dtype == torch.bfloat16 else 1
+    storage = torch.empty(
+        max(elements, _TMA_MIN_BACKING_BYTES // element_bytes), dtype=dtype, device=device
+    )
+    return storage[:elements].view(shape)
 
 
 @dataclass
@@ -231,15 +247,15 @@ def _allocate(
             )
         self._csa2_context_geometry = None
         if shared_rows is None:
-            self.swa_pool = torch.empty(
-                capacity, _SWA_TILE, _HEAD_DIM, dtype=staging_dtype, device=device
+            self.swa_pool = _allocate_staging_pool(
+                (capacity, _SWA_TILE, _HEAD_DIM), staging_dtype, device
             )
-            self.extra_pool = torch.empty(
-                capacity, max(extra_capacity, 1), _HEAD_DIM, dtype=staging_dtype, device=device
+            self.extra_pool = _allocate_staging_pool(
+                (capacity, max(extra_capacity, 1), _HEAD_DIM), staging_dtype, device
             )
         else:
-            self.shared_pool = torch.empty(
-                shared_rows, _HEAD_DIM, dtype=staging_dtype, device=device
+            self.shared_pool = _allocate_staging_pool(
+                (shared_rows, _HEAD_DIM), staging_dtype, device
             )
             self._csa2_shared_selected = torch.empty(shared_rows, dtype=torch.int32, device=device)
             self.swa_pool = self.extra_pool = self.shared_pool