[None][fix] Derive DSA indexer K-cache block stride at use time

[lucifer1004]

DeepseekV4TrtllmAttentionMetadata.__post_init__ calls super().__post_init__() before installing the model's real compress_ratios, so the base class derives _tokens_per_block from the default [1] and freezes it at the raw tokens-per-block — even though the comment right above the derivation warns the indexer slot mapping must use the compressed stride. The indexer K-cache is compressed-native (pages of tokens_per_block // compress_ratio tokens), so the fullkv/chunked-prefill gather computes flat offsets with a stride compress_ratio× too large. Data movement stays self-consistent at low block ids, which hides the bug; once KV block reuse retention (or a sufficiently large pool watermark) pushes block ids past capacity/compress_ratio, indexerKCacheGatherUnifiedKernel reads out of bounds — measured on a DSv4 tp2 setup: block_ids 35427 in-bounds, 36015 → cudaErrorIllegalAddress, against a predicted threshold of numel/stride = 35459. A side effect of the inflated addressing is that the indexer pool's effective capacity is silently 1/compress_ratio of its allocation.

Commit 1 derives IndexerParams.tokens_per_block from the live cache-manager geometry at Indexer.build_indexer_params instead of the stale metadata attribute. The write side (the DSv4 indexer compressor) already uses the manager geometry directly and is untouched; DeepSeek-V3.2-family models (compress divisor 1) are behavior-unchanged. Commit 2 adds a pre-launch bounds assert for the C++ gather matching its Python sibling _gather_k_cache_for_chunk, gated behind TRTLLM_DEBUG_INDEXER_GATHER_BOUNDS=1 (off by default; the check costs a sync on the hot path).

Validated on DSv4 (4× RTX PRO 6000 Blackwell): with block reuse enabled, a 5-round 20-concurrent ~7k-token load that previously faulted within 40–63 requests completes 103/103 with zero crashes; a full closed-loop qualification with reuse enabled passes (gsm8k 0.9447, serving bench green through the concurrency ramp); prefix-resume semantics verified across a mid-block partial-reuse boundary (early-context and near-boundary facts both answered correctly, reuse hits confirmed by 1.89× TTFT); chunked-prefill long-context ladder green at 128k/512k/1M with prefix reuse up to 52.7× TTFT gain at 1M. Reuse-disabled smoke shows no regression.

May be related to the failure family discussed in #15295 (retention-pressure-dependent corruption under block reuse), though that report is on a compress-divisor-1 model where this specific stride mismatch does not apply.

🤖 Generated with Claude Code

Summary by CodeRabbit

• Bug Fixes
  • Improved sparse attention cache indexing to use the current KV-cache configuration, preventing incorrect compressed-token mapping.
  • Added optional debug validation to detect out-of-bounds cache access during index gathering.

[coderabbitai]

📝 Walkthrough

Walkthrough

The indexer now derives compressed-token strides from the live KV-cache configuration and adds an environment-gated bounds assertion before prefill K-cache gathers.

Changes

Indexer cache addressing

Layer / File(s)	Summary
Use live KV-cache stride tensorrt_llm/_torch/attention_backend/sparse/dsa.py	IndexerParams.tokens_per_block is computed from the KV-cache manager’s live tokens_per_block and the effective compression ratio.
Validate prefill gather bounds tensorrt_llm/_torch/attention_backend/sparse/dsa.py	When TRTLLM_DEBUG_INDEXER_GATHER_BOUNDS=1, prefill slot-mapping byte offsets are checked against the indexer K-cache tensor before the gather operation.

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

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title follows the required [None][fix] format and clearly summarizes the main change to derive the DSA indexer stride at use time.
Description check	✅ Passed	It clearly explains the bug, fix, and validation, though it doesn't use the template's section headings.
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]

🧹 Nitpick comments (1)

tensorrt_llm/_torch/attention_backend/sparse/dsa.py (1)

1949-1950: 🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick win

Guard against non-divisible tokens_per_block // compress_ratio.

Floor division silently truncates when tokens_per_block is not an exact multiple of compress_ratio. Since this value becomes the indexer stride feeding slot-mapping addressing, a truncated block size would mis-stride the compressed K-cache pages — the exact class of bug this PR is fixing, but in the opposite direction. An explicit check makes the invariant (raw block size divisible by the compression ratio) fail loudly instead of producing subtly wrong gathers. Also consider reusing the local kv_cache_manager (Line 1923) here for consistency.

♻️ Proposed guard

+        raw_tokens_per_block = kv_cache_manager.tokens_per_block
+        assert raw_tokens_per_block % compress_ratio == 0, (
+            f"tokens_per_block ({raw_tokens_per_block}) must be divisible by "
+            f"compress_ratio ({compress_ratio}) for exact indexer strides")
         indexer_tokens_per_block = (
-            metadata.kv_cache_manager.tokens_per_block // compress_ratio)
+            raw_tokens_per_block // compress_ratio)

🤖 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/attention_backend/sparse/dsa.py` around lines 1949 -
1950, Validate that the local kv_cache_manager.tokens_per_block is exactly
divisible by compress_ratio before computing indexer_tokens_per_block, and fail
explicitly if the invariant is violated instead of relying on floor division.
Then compute the stride from kv_cache_manager rather than
metadata.kv_cache_manager for consistency; update the surrounding logic in the
relevant attention setup method.

🤖 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.

Nitpick comments:
In `@tensorrt_llm/_torch/attention_backend/sparse/dsa.py`:
- Around line 1949-1950: Validate that the local
kv_cache_manager.tokens_per_block is exactly divisible by compress_ratio before
computing indexer_tokens_per_block, and fail explicitly if the invariant is
violated instead of relying on floor division. Then compute the stride from
kv_cache_manager rather than metadata.kv_cache_manager for consistency; update
the surrounding logic in the relevant attention setup method.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 7292f50b-26c8-40ae-87ed-6355f3130de4

📥 Commits

Reviewing files that changed from the base of the PR and between e2ce7a6 and 3bf4490.

📒 Files selected for processing (1)

• tensorrt_llm/_torch/attention_backend/sparse/dsa.py