[None][fix] Stabilize Gemma4 FA2 CUDA Graph decode on Hopper

tensorrt_llm/_torch/attention_backend/flashinfer.py

@@ -37,6 +37,7 @@
 from tensorrt_llm.functional import AttentionMaskType
 from tensorrt_llm.logger import logger
 from tensorrt_llm.models.modeling_utils import QuantConfig
+from tensorrt_llm.runtime.kv_cache_manager_v2._common import BAD_PAGE_INDEX
 
 from ..metadata import KVCacheParams
 from ..utils import get_global_attrs, get_model_extra_attrs, torch_multi_arange
@@ -232,6 +233,8 @@ class MLAPlanParams:
 @dataclass(kw_only=True)
 class FlashInferWrappers:
     is_planned: bool
+    fa2_plan_num_blocks: Optional[tuple[int, ...]] = field(default=None,
+                                                           repr=False)
     decode_wrapper: Optional[
         flashinfer.BatchDecodeWithPagedKVCacheWrapper] = None
     prefill_wrapper: Optional[
@@ -579,6 +582,20 @@ def get_paged_kv_indices_for_layer(self, layer_idx: int) -> torch.Tensor:
         total_blocks = self.num_generation_blocks + self.num_context_blocks
         return self._vswa_pool_indices_cache[pool_id][:total_blocks]
 
+    def _sanitize_swa_page_indices(self, page_indices: torch.Tensor,
+                                   layer_idx: int) -> None:
+        """Replace evicted SWA pages with a safe in-range page index."""
+        window_vec = getattr(self.kv_cache_manager, 'max_attention_window_vec',
+                             None)
+        if not window_vec or window_vec[layer_idx % len(window_vec)] is None:
+            return
+
+        # KVCacheManagerV2 marks evicted out-of-window pages with -1.
+        # FlashInfer may dereference page IDs before applying window_left, so
+        # keep masked positions in range. The SWA mask excludes their values
+        # from the attention result.
+        page_indices.masked_fill_(page_indices == BAD_PAGE_INDEX, 0)
+
     def swap_paged_kv_indices_for_layer(self, layer_idx: int) -> None:
         """Copy pool-specific page indices into the shared buffer.
 
@@ -1425,12 +1442,32 @@ def _clean_cached_plans(self, *, defer_plan: bool):
             # corresponding forward pass. So, flush them out here as they won't be relevant for
             # subsequent forward calls.
             if plan_params.attention_mask_data is None and plan_params.multi_item_params is None:
-                self._plan_params_to_wrappers[plan_params].is_planned = False
+                wrappers = self._plan_params_to_wrappers[plan_params]
+                if wrappers.fa2_plan_num_blocks is not None:
+                    continue
+                wrappers.is_planned = False
                 if not defer_plan:
                     self._plan_with_params(plan_params)
             else:
                 del self._plan_params_to_wrappers[plan_params]
 
+    def _refresh_fa2_cuda_graph_plans(self) -> None:
+        """Refresh captured FA2 schedules after page metadata is finalized."""
+        num_blocks = tuple(self.num_blocks[self.num_contexts:])
+        for plan_params, wrappers in self._plan_params_to_wrappers.items():
+            if (plan_params.attention_mask_data is not None
+                    or plan_params.multi_item_params is not None
+                    or wrappers.fa2_plan_num_blocks is None):
+                continue
+            if not num_blocks:
+                wrappers.fa2_plan_num_blocks = None
+            elif wrappers.fa2_plan_num_blocks != num_blocks:
+                # Graph replay does not re-enter forward_impl. Each wrapper
+                # owns its persistent integer plan workspace, while the shared
+                # float workspace is run scratch.
+                wrappers.is_planned = False
+                self._plan_with_params(plan_params)
+
     def prepare(self) -> None:
 
         def _to_int32_tensor(arr: np.ndarray) -> torch.Tensor:
@@ -1542,8 +1579,18 @@ def _to_int32_tensor(arr: np.ndarray) -> torch.Tensor:
         self.num_generation_blocks = sum(self.num_blocks[self.num_contexts:])
 
         # indices of used cache blocks for each sequence
+        primary_layer_idx = None
+        if self._vswa_layer_to_pool is not None:
+            primary_pool_id = self._vswa_layer_to_pool.get(0, 0)
+            primary_layer_idx = self._vswa_pool_to_rep_layer[primary_pool_id]
+        else:
+            layer_offsets = getattr(self.kv_cache_manager, 'layer_offsets', {})
+            primary_layer_idx = next(iter(layer_offsets), None)
+
         paged_kv_indices = self.kv_cache_manager.get_batch_cache_indices_flat(
-            self.request_ids, self.num_blocks)
+            self.request_ids, self.num_blocks, layer_idx=primary_layer_idx)
+        if primary_layer_idx is not None:
+            self._sanitize_swa_page_indices(paged_kv_indices, primary_layer_idx)
 
         self._paged_kv_indices[:paged_kv_indices.size(0)].copy_(
             paged_kv_indices, non_blocking=True)
@@ -1579,6 +1626,7 @@ def _to_int32_tensor(arr: np.ndarray) -> torch.Tensor:
                 pool_indices = \
                     self.kv_cache_manager.get_batch_cache_indices_flat(
                         self.request_ids, self.num_blocks, layer_idx=rep_layer)
+                self._sanitize_swa_page_indices(pool_indices, rep_layer)
                 buf = getattr(self, f'_vswa_pool_buf_{pool_id}')
                 buf[:pool_indices.size(0)].copy_(pool_indices,
                                                  non_blocking=True)
@@ -1656,19 +1704,6 @@ def _to_int32_tensor(arr: np.ndarray) -> torch.Tensor:
             self._positions[:positions.size(0)].copy_(positions,
                                                       non_blocking=True)
 
-        # Multi-wrapper case (Gemma4 hybrid: different head_dim per layer)
-        # shares one workspace_buffer; eager plan() would overwrite earlier
-        # wrappers' workspace, so defer plan() to forward_impl. Single-wrapper
-        # case (e.g., Llama, Gemma3 uniform head_dim) needs eager plan() here
-        # because forward_impl cannot plan() during cuda-graph stream capture.
-        active_wrappers = [
-            pp for pp in self._plan_params_to_wrappers
-            if pp.attention_mask_data is None
-        ]
-        defer_plan = len(active_wrappers) > 1
-        if not (self._is_separate_kv_draft_view and self.is_cuda_graph):
-            self._clean_cached_plans(defer_plan=defer_plan)
-
         # Re-plan MLA wrappers outside of forward/capture using the params
         # cached by prior warmup forwards. Forward still handles first-use or
         # dtype/shape changes by syncing only on a plan cache miss.
@@ -1798,6 +1833,20 @@ def _to_int32_tensor(arr: np.ndarray) -> torch.Tensor:
                                                              non_blocking=True)
                     if self.num_generations < batch_size:
                         kv_lens_buf[self.num_generations:batch_size].zero_()
+
+        # Refresh captured FA2 schedules only after all page metadata updates.
+        # Defer ordinary multi-wrapper plans to forward_impl; single-wrapper
+        # models still plan eagerly because forward_impl cannot plan during
+        # graph capture.
+        active_wrappers = [
+            pp for pp in self._plan_params_to_wrappers
+            if pp.attention_mask_data is None
+        ]
+        defer_plan = len(active_wrappers) > 1
+        if not (self._is_separate_kv_draft_view and self.is_cuda_graph):
+            self._refresh_fa2_cuda_graph_plans()
+            self._clean_cached_plans(defer_plan=defer_plan)
+
         if (not self._is_shared_kv_draft_view
                 and not self._is_separate_kv_draft_view
                 and self._draft_metadata is not None):
@@ -1987,8 +2036,12 @@ def prefill_plan():
                 custom_mask=plan_params.attention_mask_data,
             )
 
+        use_graph_tensor_cores = self.is_cuda_graph and plan_params.head_dim > 128
         if wrappers.decode_wrapper is None:
             use_tensor_cores = self._use_tensor_cores(plan_params)
+            # Gemma4's H256/H512 plans need a tensor-core wrapper with a stable
+            # CUDA Graph launch layout. prepare() may refresh its split-K
+            # schedule in the wrapper's fixed workspace as KV pages change.
 
             wrappers.decode_wrapper = \
                 flashinfer.BatchDecodeWithPagedKVCacheWrapper(
@@ -1998,11 +2051,11 @@ def prefill_plan():
                     paged_kv_indptr_buffer=self.paged_kv_indptr_decode,
                     paged_kv_indices_buffer=self._paged_kv_indices,
                     paged_kv_last_page_len_buffer=self._paged_kv_last_page_len,
-                    use_tensor_cores=use_tensor_cores
+                    use_tensor_cores=use_tensor_cores or use_graph_tensor_cores
                     or flashinfer_backend == "trtllm-gen",
                     backend=flashinfer_backend
                     if flashinfer_backend != "fa2" else
-                    ("fa2" if torch.cuda.get_device_capability(0) == (
+                    ("fa2" if torch.cuda.get_device_capability() == (
                         9, 0) else "auto"),
                 )
         decode_wrapper = wrappers.decode_wrapper
@@ -2042,7 +2095,11 @@ def decode_plan():
                 block_tables=block_tables,
                 # Keep FlashInfer's recorded graph shape aligned with the wrapper cache key.
                 q_len_per_req=plan_params.q_len_per_req,
+                disable_split_kv=False,
             )
+            if use_graph_tensor_cores and decode_wrapper._backend == 'fa2':
+                wrappers.fa2_plan_num_blocks = tuple(
+                    self.num_blocks[self.num_contexts:])
             self._publish_decode_wrapper_kv_lens(decode_wrapper)
 
         # Must sync after append_paged_kv_cache and before plan.

tensorrt_llm/_torch/models/modeling_gemma4.py

@@ -1339,11 +1339,10 @@ def get_context_mask(
         """Build context mask with causal + bidirectional for MM tokens.
 
         Returns a [extend_len, prefix_len + extend_len] mask where:
-        - The first `prefix_len` columns (cached/paged history) are True for
-          all rows. SWA window enforcement is delegated to the kernel's
-          window_left clip. Bidirectional MM across the prefix/extend
-          boundary is NOT supported here; callers must ensure chunk
-          boundaries do not split a multimodal block.
+        - The first `prefix_len` columns (cached/paged history) apply the
+          sliding window using absolute token positions. Bidirectional MM
+          across the prefix/extend boundary is NOT supported here; callers
+          must ensure chunk boundaries do not split a multimodal block.
         - The last `extend_len` columns follow the original causal +
           (optional) sliding window + MM-bidirectional logic.
         """
@@ -1360,9 +1359,19 @@ def get_context_mask(
         causal_mask = causal_mask.masked_fill(token_type_mask, True)
 
         if prefix_len > 0:
-            prefix_block = torch.ones(
-                extend_len, prefix_len, dtype=causal_mask.dtype, device=device
-            )
+            if (
+                effective_sliding_window is not None
+                and effective_sliding_window < prefix_len + extend_len
+            ):
+                query_pos = prefix_len + pos
+                prefix_pos = torch.arange(prefix_len, device=device)
+                prefix_block = (
+                    prefix_pos.unsqueeze(0) > query_pos.unsqueeze(1) - effective_sliding_window
+                )
+            else:
+                prefix_block = torch.ones(
+                    extend_len, prefix_len, dtype=causal_mask.dtype, device=device
+                )
             causal_mask = torch.cat([prefix_block, causal_mask], dim=1)
 
         return causal_mask

tests/integration/test_lists/test-db/l0_h100.yml

@@ -68,6 +68,7 @@ l0_h100:
   - unittest/_torch/modeling -k "modeling_mixtral"
   - unittest/_torch/modeling -k "modeling_gemma3"
   - unittest/_torch/modeling -k "modeling_gpt_oss"
+  - unittest/_torch/modeling -k "modeling_gemma4"
   - unittest/_torch/modeling -k "modeling_whisper"  # CPU-only log-mel parity
   - unittest/_torch/modeling/test_modeling_nemotron_h.py::test_nemotron_h_sanity
   # Real-weight Nano CG/overlap and chunked-prefill path smoke (MoE L0 cannot