[None][feat] DSA: adaptive indexer prefill chunk size for long sequences#15459
Conversation
The indexer's fp8_fp4_mqa_logits / fp8_mqa_logits activation memory scales with indexer_max_chunk_size * K_compressed (the compressed KV length of the current request). For very long sequences this can OOM: e.g. a ~500K-token request with the default 32K chunk size needs ~16GB of activation memory. The previous workaround was to uniformly lower indexer_max_chunk_size (e.g. to 8K), but that costs prefill throughput for the common case where the vast majority of requests are far below these lengths. Instead, select the indexer prefill chunk size per-batch based on the largest compressed KV length among the context requests: - max K_compressed > 512K -> 8K chunk - 256K <= max K_compressed <= 512K -> 16K chunk - max K_compressed < 256K -> configured chunk size (unchanged) The heuristic only ever reduces the configured chunk size (never increases it), so the common case keeps its larger, higher-throughput chunk. It is applied only on the indexer's own chunking path (not the MLA chunked-prefill path, which already bounds chunk size) and is safe to vary per-batch because the prefill path does not use CUDA graphs. Signed-off-by: Fanrong Li <lfr-0531@users.noreply.github.com>
|
/bot run |
|
PR_Github #54853 [ run ] triggered by Bot. Commit: |
|
PR_Github #54853 [ run ] completed with state
|
|
/bot run |
|
PR_Github #55123 [ run ] triggered by Bot. Commit: |
|
PR_Github #55123 [ run ] completed with state
|
|
/bot run |
|
PR_Github #55185 [ run ] triggered by Bot. Commit: |
|
/bot kill |
|
PR_Github #55350 [ kill ] triggered by Bot. Commit: |
|
PR_Github #55185 [ run ] completed with state |
|
PR_Github #55350 [ kill ] completed with state |
The Pre-commit Check CI runs `bandit -r tensorrt_llm`, which flags 8 B105 (hardcoded_password_string) false positives: string literals compared against `token_back_*` variables (the name contains "token", which bandit treats as a secret). These are plain mode/schedule comparisons, not secrets. Annotate the 8 lines with `# nosec B105`. scripts/bandit.yaml is intentionally left untouched (its header asks not to modify it). Signed-off-by: Fanrong Li <lfr-0531@users.noreply.github.com>
|
/bot run |
|
PR_Github #55371 [ run ] triggered by Bot. Commit: |
|
PR_Github #55371 [ run ] completed with state
|
|
/bot run |
|
PR_Github #55380 [ run ] triggered by Bot. Commit: |
|
/bot run --disable-fail-fast |
|
PR_Github #55393 [ run ] triggered by Bot. Commit: |
|
PR_Github #55393 [ run ] completed with state
|
…ock pre-commit The repo-wide `bandit -r tensorrt_llm` gate in scripts/release_check.py flags B105 (hardcoded_password_string) on token_back_mode string comparisons (e.g. token_back_mode != "epi_warps") in the megamoe files because the variable name contains "token". These are false positives. Add targeted `# nosec B105` suppressions (the test-scoped form, which the gate permits since it does not increment "Total lines skipped (#nosec)"). Mirrors the same suppressions in PR NVIDIA#15459 so this PR's Pre-commit Check passes without waiting for that PR to land. Signed-off-by: Yukun He <23156053+hyukn@users.noreply.github.com>
|
/bot run --disable-fail-fast |
|
PR_Github #55505 [ run ] triggered by Bot. Commit: |
|
PR_Github #55505 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #55633 [ run ] triggered by Bot. Commit: |
|
PR_Github #55633 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #55717 [ run ] triggered by Bot. Commit: |
|
PR_Github #55717 [ run ] completed with state
|
|
/bot run |
|
PR_Github #55934 [ run ] triggered by Bot. Commit: |
|
PR_Github #55934 [ run ] completed with state
|
Work around the gb200-x4 CI infra issue by requesting the -split node pool for the GB200-4_GPUs-PyTorch-DS-1 test, mirroring PR NVIDIA#15548 (which did the same for the GB200 perf-sanity stages). Signed-off-by: Fanrong Li <lfr-0531@users.noreply.github.com>
|
/bot run |
|
PR_Github #56016 [ run ] triggered by Bot. Commit: |
|
PR_Github #56016 [ run ] completed with state |
…ces (NVIDIA#15459) Signed-off-by: Fanrong Li <lfr-0531@users.noreply.github.com> Co-authored-by: Fanrong Li <lfr-0531@users.noreply.github.com> Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
…ces (NVIDIA#15459) Signed-off-by: Fanrong Li <lfr-0531@users.noreply.github.com> Co-authored-by: Fanrong Li <lfr-0531@users.noreply.github.com> Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
2 participants
Background / Motivation
The DSA indexer's
fp8_fp4_mqa_logits/fp8_mqa_logitsactivation memoryscales with
indexer_max_chunk_size * K_compressed(the compressed KV lengthof the current request). For very long sequences this can OOM — e.g. a
~500K-token request with the default 32K chunk size needs ~16GB of activation
memory.
The current workaround is to uniformly lower
indexer_max_chunk_size(e.g. to8K). That avoids the OOM but costs prefill throughput for the common case,
where the vast majority of requests are far below these lengths — a blunt
one-size-fits-all setting for what is really a long-tail problem.
Summary
Select the indexer prefill chunk size per-batch based on the largest
compressed KV length (
K_compressed) among the context requests, instead ofalways using the statically-configured value:
> 512K[256K, 512K]< 256KImplementation:
select_indexer_chunk_size(configured_chunk_size, max_k_compressed)in
dsa.py. It only ever reduces the configured chunk size (neverincreases it), so any explicitly-configured value still acts as an upper
bound.
Indexer.prepare_for_chunked_prefill, on the indexer's ownchunking path only. The MLA chunked-prefill path is untouched (it already
bounds the chunk to the MLA chunk).
max K_compressedis read fromindexer_params.kv_lens(already host-sideduring prefill prepare).
Impact
avoids the OOM without manual tuning.
higher-throughput chunk size.
indexer_max_chunk_sizecontinues to act asthe upper bound.
Notes
sizes (8K / 16K / 32K) are centralized in
_INDEXER_CHUNK_SIZE_HEURISTICforeasy tuning.
select_indexer_chunk_sizeboundaries and validate end-to-end memory/perf on a long-sequence workload.
🤖 Generated with Claude Code