[None][fix] Fix Gemma4 illegal memory access when max_seq_len is at most the sliding window size

[Hudayday]

Description

Symptom

Every Gemma4 checkpoint (E2B / E4B / 12B / 26B-A4B / 31B) crashes with
CUDA_ERROR_ILLEGAL_ADDRESS during the warmup prefill on the FlashInfer
attention backend whenever max_seq_len is at most the sliding window
size (512), e.g. max_seq_len=256. Default (large) max_seq_len is
unaffected, which is why regular deployments and accuracy runs never hit
it. This is not a recent regression: it reproduces on months-old builds.

Root cause

The KV cache manager (KVCacheManagerV2) allocates one slot per token
block covering every layer's page, and each layer converts the shared
slot number into a page index of its own buffer with a per-layer
multiplier ("page-index scale"). Layers whose pages have different byte
sizes have different scales: for gemma-4-E2B, the 28 sliding-attention
layers (head_dim 256) have scale 56 while the 7 global-attention layers
(head_dim 512) have scale 14. The core documents that computed page
indices may only be shared between buffers with the same scale
(get_page_index_scale).

When max_seq_len is at most the sliding window size, every attention
window is clamped to max_seq_len, so all layers land in a single KV
pool. The wrapper, however, kept one scale per pool (taken from the
pool's first layer) and shared one converted page-index list across all
layers. Global-attention layers therefore received page ids up to 4x
past the end of their buffer, and append_paged_kv_cache wrote out of
bounds. With larger max_seq_len the sliding and global layers sit in
separate pools, each pool is homogeneous, and the single per-pool scale
happens to be correct — which is why only small max_seq_len crashes.

Fix (Python only; no changes to pool grouping, memory pools, or C++)

• KVCacheManagerV2.get_batch_cache_indices(layer_idx=...) now converts
  with that layer's own page-index scale (new helper
  get_layer_page_index_scale), matching the core's documented contract.
• FlashInferAttentionMetadata now shares one page-index list only among
  layers with the same pool and the same scale, and engages the per-pool
  index swap when windows differ (VSWA) or scales differ (previously
  VSWA only).

Test Coverage

• New regression test
  tests/unittest/_torch/modeling/test_gemma4_e2e_dummy.py::test_e2e_text_e2b_dummy_small_max_seq_len[256|512]
  (real gemma-4-E2B geometry, dummy weights). It fails with the illegal
  memory access before this fix and passes with it. Registered in the
  B200 pre-merge list (l0_b200.yml).
• Manual validation on B200 with real gemma-4-E2B-it weights:
  max_seq_len 64 (batch 8), 256, and 512 crashed before the fix and
  pass with it; 768 and 2048 (above the window size) pass before and
  after (no regression). Generated text is coherent in all passing runs.

Summary by CodeRabbit

• Bug Fixes

  • Improved attention cache handling so layers with different page-index scaling are treated separately, reducing incorrect reuse of cached page-index data.

• Tests

  • Added regression coverage for Gemma-4 dummy end-to-end generation at multiple sequence lengths.
  • Expanded the B200 integration test list to include the new Gemma-4 cases.

[coderabbitai]

📝 Walkthrough

Walkthrough

This PR refines page-index scale handling for KV cache layers with differing head dimensions. A new get_layer_page_index_scale accessor is added to the KV cache manager and used to compute per-layer scales, which then drives a revised VSWA pool grouping key in FlashInfer attention metadata. A new Gemma4 E2B regression test validates the behavior.

Changes

Per-layer page-index scale support

Layer / File(s)	Summary
KV cache manager per-layer scale accessor and propagation tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py	Adds get_layer_page_index_scale(layer_idx), computes per-layer index_scale in get_batch_cache_indices, and extends _get_batch_cache_indices_by_pool_id to accept and use an optional index_scale override, falling back to pool-level index_scales otherwise.
VSWA pool grouping by combined pool/scale key tensorrt_llm/_torch/attention_backend/flashinfer.py	Reworks VSWA setup to group layers by (pool_id, page_index_scale) instead of just pool_id, activating per-layer pool splitting when is_vswa is true or multiple distinct scale groups exist.
Gemma4 E2B dummy regression test and test list entry tests/unittest/_torch/modeling/test_gemma4_e2e_dummy.py, tests/integration/test_lists/test-db/l0_b200.yml	Adds a parametrized E2E test building an E2B dummy config with differing sliding/global head dims and running FLASHINFER-backed dummy inference across max_seq_len values of 256 and 512, plus corresponding entries in the b200 test list.

Estimated code review effort: 3 (Moderate) | ~25 minutes

Sequence Diagram(s)

sequenceDiagram
  participant FlashInferAttentionMetadata
  participant KVCacheManagerV2

  FlashInferAttentionMetadata->>KVCacheManagerV2: get_layer_page_index_scale(layer_idx)
  KVCacheManagerV2-->>FlashInferAttentionMetadata: page_index_scale (per layer)
  FlashInferAttentionMetadata->>FlashInferAttentionMetadata: group layers by (pool_id, page_index_scale)
  alt is_vswa true OR multiple scale groups
    FlashInferAttentionMetadata->>FlashInferAttentionMetadata: build _vswa_layer_to_pool / _vswa_pool_to_rep_layer per group
  else single group and not vswa
    FlashInferAttentionMetadata->>FlashInferAttentionMetadata: use shared page-index list
  end

Loading

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title is clear, concise, and accurately summarizes the main fix using the required ticket/type format.
Description check	✅ Passed	The description clearly explains the symptom, root cause, fix, and test coverage, with only the checklist section omitted.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Actionable comments posted: 1

🧹 Nitpick comments (2)

tests/unittest/_torch/modeling/test_gemma4_e2e_dummy.py (1)

245-254: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick win

Test implicitly depends on the real checkpoint's sliding_window remaining 512.

With shrink_hidden=False, sliding_window is left as-is from the real E2B config.json rather than being explicitly patched. The comment notes "sliding_window must stay 512 so both max_seq_len values are at most the window" — if the upstream checkpoint's config ever changes this value, the test would silently stop covering the intended regression scenario without failing loudly. Consider explicitly asserting/patching sliding_window == 512 in _make_dummy_config_dir (even outside the shrink_hidden branch) for determinism.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/unittest/_torch/modeling/test_gemma4_e2e_dummy.py` around lines 245 -
254, The E2B dummy test setup in _make_dummy_config_dir currently relies on the
upstream config.json to leave sliding_window at 512 when shrink_hidden is false,
which makes the regression coverage non-deterministic. Update the dummy config
generation path used by test_gemma4_e2e_dummy to explicitly assert or patch
sliding_window to 512 regardless of shrink_hidden, so the geometry in the E2B
case stays stable and the test fails loudly if the real checkpoint config
changes.

tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py (1)

2778-2783: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Document layer_idx semantics in the docstring.

get_layer_page_index_scale is a new cross-file public API (consumed from flashinfer.py), but the docstring omits an Args/Returns section clarifying that layer_idx is the global layer index (converted internally via self.layer_offsets), not the local pool offset. Given this entire PR is about layer-index/scale mismatches, documenting this explicitly would reduce future misuse risk.

As per coding guidelines, "Externally called functions should have docstrings, and their arguments should be documented," and "Use Google style docstrings for classes and functions in Python, parseable by Sphinx."

📝 Proposed docstring update

     def get_layer_page_index_scale(self, layer_idx: int) -> int:
-        """Page-index scale of this layer's KV buffer. Layers in one pool can
-        have different scales (e.g. different head_dim), so per-layer callers
-        must not use the pool-level scale."""
+        """Page-index scale of this layer's KV buffer.
+
+        Layers in one pool can have different scales (e.g. different
+        head_dim), so per-layer callers must not use the pool-level scale.
+
+        Args:
+            layer_idx: Global layer index (mapped internally to the local
+                pool offset via ``self.layer_offsets``).
+
+        Returns:
+            The KEY buffer's page-index scale for this layer.
+        """

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py` around lines 2778 -
2783, The public method get_layer_page_index_scale currently lacks Google-style
argument documentation and does not clarify that layer_idx is the global layer
index, not the local pool offset. Update its docstring to include Args and
Returns sections, explicitly stating that layer_idx is translated through
self.layer_offsets before calling self.impl.get_page_index_scale, so callers in
flashinfer.py and other external users do not misuse the API.

Source: Coding guidelines

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@tests/unittest/_torch/modeling/test_gemma4_e2e_dummy.py`:
- Around line 229-266: This test can silently run with TrtllmAttention instead
of FlashInfer, so add an explicit FlashInfer availability guard before
constructing LLM in test_e2e_text_e2b_dummy_small_max_seq_len. Use the existing
availability symbol or decorator pattern already used in tests (for example a
requires_* or skipif check tied to FlashInfer) so the test is skipped when
FlashInfer is unavailable. Keep the current attn_backend="FLASHINFER" setup, but
ensure the test only executes the intended FlashInfer VSWA path.

---

Nitpick comments:
In `@tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py`:
- Around line 2778-2783: The public method get_layer_page_index_scale currently
lacks Google-style argument documentation and does not clarify that layer_idx is
the global layer index, not the local pool offset. Update its docstring to
include Args and Returns sections, explicitly stating that layer_idx is
translated through self.layer_offsets before calling
self.impl.get_page_index_scale, so callers in flashinfer.py and other external
users do not misuse the API.

In `@tests/unittest/_torch/modeling/test_gemma4_e2e_dummy.py`:
- Around line 245-254: The E2B dummy test setup in _make_dummy_config_dir
currently relies on the upstream config.json to leave sliding_window at 512 when
shrink_hidden is false, which makes the regression coverage non-deterministic.
Update the dummy config generation path used by test_gemma4_e2e_dummy to
explicitly assert or patch sliding_window to 512 regardless of shrink_hidden, so
the geometry in the E2B case stays stable and the test fails loudly if the real
checkpoint config changes.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 63582079-4513-4ee8-91a5-898772f82e45

📥 Commits

Reviewing files that changed from the base of the PR and between 73d3abc and 4ac41bb.

📒 Files selected for processing (4)

• tensorrt_llm/_torch/attention_backend/flashinfer.py
• tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py
• tests/integration/test_lists/test-db/l0_b200.yml
• tests/unittest/_torch/modeling/test_gemma4_e2e_dummy.py

[Hudayday]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #58142 [ run ] triggered by Bot. Commit: 4ac41bb Link to invocation

[tensorrt-cicd]

PR_Github #58142 [ run ] completed with state FAILURE. Commit: 4ac41bb
/LLM/main/L0_MergeRequest_PR pipeline #46797 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[Hudayday]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #58223 [ run ] triggered by Bot. Commit: 4ac41bb Link to invocation

[tensorrt-cicd]

PR_Github #58223 [ run ] completed with state FAILURE. Commit: 4ac41bb
/LLM/main/L0_MergeRequest_PR pipeline #46864 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[Hudayday]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #58235 [ run ] triggered by Bot. Commit: 4ac41bb Link to invocation

[tensorrt-cicd]

PR_Github #58235 [ run ] completed with state FAILURE. Commit: 4ac41bb
/LLM/main/L0_MergeRequest_PR pipeline #46876 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[Hudayday]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #58348 [ run ] triggered by Bot. Commit: 4ac41bb Link to invocation

[tensorrt-cicd]

PR_Github #58348 [ run ] completed with state FAILURE. Commit: 4ac41bb

Link to invocation

[Hudayday]

/bot run --disable-fail-fast

[Hudayday]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #58378 [ run ] triggered by Bot. Commit: 4ac41bb Link to invocation

[tensorrt-cicd]

PR_Github #58378 [ run ] completed with state SUCCESS. Commit: 4ac41bb
/LLM/main/L0_MergeRequest_PR pipeline #47002 completed with status: 'SUCCESS'

CI Report

Link to invocation

[nvpohanh]

[by Codex] @lowsfer Could you review this PR? Thanks!

[eopXD]

Fix looks good to me.

A further round-up will be to have the grouping of (slot_group, index_scale) built inside the manager itself, building the partition where index_scales is built. A more schematic approach will be (as Yao mentioned to me in chat), a page index converter from the manager that adapts regarding known models.

For now the fix is enough to unblock the IMA in the FlashInfer backend. A core fix will free us from backend-local mappings (for example TRTLLM will fall for the same IMA, although not reachable right now for Gemma4 since the backend is not supported yet).

[Hudayday]

/bot run --disable-fail-fast

[github-actions]

👎 Promotion blocked, new vulnerability found

Vulnerability report

Component	Vulnerability	Description	Severity
pytorch	CVE-2025-3000	A vulnerability classified as critical has been found in PyTorch 2.6.0. This affects the function torch.jit.script. The manipulation leads to memory corruption. It is possible to launch the attack on the local host. The exploit has been disclosed to the public and may be used.	MEDIUM

[Hudayday]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #59206 [ run ] triggered by Bot. Commit: 991459c Link to invocation

[tensorrt-cicd]

PR_Github #59206 [ run ] completed with state FAILURE. Commit: 991459c
/LLM/main/L0_MergeRequest_PR pipeline #47704 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[yihwang-nv]

Thanks, LGTM!

[Hudayday]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #59347 [ run ] triggered by Bot. Commit: f80bb64 Link to invocation

[tensorrt-cicd]

PR_Github #59347 [ run ] completed with state SUCCESS. Commit: f80bb64
/LLM/main/L0_MergeRequest_PR pipeline #47825 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[Hudayday]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #59385 [ run ] triggered by Bot. Commit: f80bb64 Link to invocation

[tensorrt-cicd]

PR_Github #59385 [ run ] completed with state SUCCESS. Commit: f80bb64
/LLM/main/L0_MergeRequest_PR pipeline #47859 completed with status: 'SUCCESS'

CI Report

Link to invocation