[None][feat] DSA: adaptive indexer prefill chunk size for long sequences

[lfr-0531]

@coderabbitai summary

Description

Cherry-pick of the core change from #15459 (already merged to feat/deepseek_v4) onto feat/deepseek_v4_bench. The CI-infra-only changes from that PR (bandit # nosec, the gb200-x4-split node-pool swap) are intentionally not included — this is the dsa.py change only.

The DSA indexer's fp8(_fp4)_mqa_logits activation memory scales with indexer_max_chunk_size * K_compressed (the compressed KV length of the request). For very long sequences this OOMs — e.g. a ~500K-token request at the default 32K chunk size needs ~16GB of activation memory. Uniformly lowering indexer_max_chunk_size (e.g. to 8K) avoids the OOM but sacrifices prefill throughput for the common case, where the vast majority of requests are far below these lengths.

This change makes the indexer prefill chunk size adaptive per-batch, keyed on the largest compressed KV length among the context requests:

max K_compressed in batch	effective chunk size
> 512K	8K
[256K, 512K]	16K
< 256K	configured (default 32K, unchanged)

Implementation notes:

• New helper select_indexer_chunk_size(configured_chunk_size, max_k_compressed) in dsa.py; thresholds centralized in _INDEXER_CHUNK_SIZE_HEURISTIC for easy tuning.
• It only ever reduces the configured chunk size (never increases it), so an explicitly configured indexer_max_chunk_size still acts as an upper bound.
• Wired into Indexer.prepare_for_chunked_prefill, on the indexer's own chunking path only. The MLA chunked-prefill path is untouched (it already bounds the chunk to the MLA chunk size).
• max K_compressed is read from indexer_params.kv_lens (already available host-side during prefill prepare; no extra device sync).
• Safe to vary per-batch because the prefill path does not use CUDA graphs.

Test Coverage

• Covered by the existing DSA indexer prefill / chunked-prefill unit tests under tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py.
• The heuristic only reduces the chunk size for very long sequences (>256K K_compressed); for the common case (<256K) the configured chunk size is used unchanged, so existing behavior and test expectations are preserved.
• Originally validated and CI-merged via [None][feat] DSA: adaptive indexer prefill chunk size for long sequences #15459 on feat/deepseek_v4; this PR is an identical cherry-pick of that dsa.py change.

PR Checklist

Please review the following before submitting your PR:

• PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.

• PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

• Test cases are provided for new code paths (see test instructions)

• Any new dependencies have been scanned for license and vulnerabilities

• CODEOWNERS updated if ownership changes

• Documentation updated as needed

• Update tava architecture diagram if there is a significant design change in PR.

• The reviewers assigned automatically/manually are appropriate for the PR.

• Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.