[None][fix] DSv4: bound indexer MQA-logits transient via query tiling#15569
Conversation
|
/bot run --disable-fail-fast |
|
PR_Github #55374 [ run ] triggered by Bot. Commit: |
|
PR_Github #55374 [ run ] completed with state
|
On the disagg CTX (prefill) worker, the DSA indexer's fp8_mqa_logits allocates a [q_chunk x kv] logits output whose KV dimension is the full (compressed) context and is unbounded. With max_num_tokens=8192 on a long context this single torch.empty reaches ~18-22 GB; under PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True it can stall indefinitely in cuMemCreate on the longest-context (attention_dp laggard) rank (GPU idle, host sched_yield-spin) while peer ranks spin in the MoE EP kernel -> 300s watchdog -> whole-engine hang. With the native allocator the same allocation fails as a (recoverable) OOM. Tile the query dimension of the indexer MQA-logits + top-k so each fp8_mqa_logits call allocates at most q_tile x kv, where q_tile = min(local_q, BUDGET // num_k_tokens). Results are identical (each query row's logits/top-k are independent; the KV is gathered once per chunk and unchanged across tiles), so the per-call allocation is the same size and the caching allocator reuses one block (peak ~= one tile, no extra sync). The budget is tunable via TLLM_INDEXER_MQA_LOGITS_ELEM_BUDGET (default 1<<31 elems). Verified in isolation on GB200: logits bit-identical and top-k set identical (full vs tiled), peak ~40GB -> ~4GB, no perf regression (short contexts use a single tile). Signed-off-by: Yukun He <23156053+hyukn@users.noreply.github.com>
hyukn
force-pushed
the
fix/dsv4-indexer-mqa-logits-qtiling
branch
from
615c0d6 to
2d6f662
Compare
…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 |
|
/bot run --disable-fail-fast |
|
PR_Github #55524 [ run ] triggered by Bot. Commit: |
|
PR_Github #55524 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #55639 [ run ] triggered by Bot. Commit: |
|
PR_Github #55639 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #55673 [ run ] triggered by Bot. Commit: |
|
PR_Github #55673 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #55774 [ run ] triggered by Bot. Commit: |
|
PR_Github #55774 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #55819 [ run ] triggered by Bot. Commit: |
|
PR_Github #55819 [ run ] completed with state
|
|
/bot run |
|
PR_Github #55937 [ run ] triggered by Bot. Commit: |
|
PR_Github #55937 [ run ] completed with state
|
|
/bot run |
|
PR_Github #56068 [ run ] triggered by Bot. Commit: |
|
PR_Github #56068 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #56081 [ run ] triggered by Bot. Commit: |
|
PR_Github #56081 [ run ] completed with state
|
|
/bot run |
|
PR_Github #56098 [ run ] triggered by Bot. Commit: |
|
PR_Github #56098 [ run ] completed with state |
…NVIDIA#15569) Signed-off-by: Yukun He <23156053+hyukn@users.noreply.github.com> Co-authored-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com> Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
…NVIDIA#15569) Signed-off-by: Yukun He <23156053+hyukn@users.noreply.github.com> Co-authored-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com> Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
3 participants
@coderabbitai summary
Description
On the disaggregated context (prefill) worker, the DeepSeek-V4 DSA sparse-MLA indexer can deadlock the whole engine on long-context requests.
Root cause. In
Indexer.sparse_attn_indexer(prefill path),fp8_mqa_logitsallocates its[q_chunk, kv]logits output viatorch.empty, where the KV dimension is the full (compressed) context and is not bounded per call. Withmax_num_tokens=8192on a long context this single transient reaches ~18–22 GB and keeps growing with context length. UnderPYTORCH_CUDA_ALLOC_CONF=expandable_segments:True, an allocation that large can stall indefinitely insidecuMemCreate(the CUDA VMM driver call) on the longest-context (attention_dplaggard) rank: that rank's GPU goes idle while its host thread spins in the driver, the peer ranks block at the next MoE all-to-all waiting for it, and the 300s hang watchdog eventually trips and tears the engine down. With the native (non-expandable) caching allocator, the same oversized request instead surfaces as a (recoverable) CUDA OOM.A symbolized backtrace of a live hung rank pinpoints the stall to the indexer's logits allocation:
Fix. Tile the query dimension of the indexer MQA-logits + top-k so that each
fp8_mqa_logitscall allocates at mostq_tile * kv, whereq_tile = min(local_q, BUDGET // num_k_tokens)andBUDGETis a tunable element budget (TLLM_INDEXER_MQA_LOGITS_ELEM_BUDGET, default1 << 31). This is mathematically identical to the original computation: each query row's logits and top-k are independent of other rows, and the KV side (chunk_k) is gathered once per chunk and is unchanged across tiles. Because the per-call allocations are therefore the same size, the stream-ordered caching allocator reuses a single block, so the peak transient is ~one tile with no added synchronization and no kernel/API change. Short contexts are unaffected (they run as a single tile). Capping the per-call transient fixes both theexpandable_segmentshang and the native-allocator OOM at anymax_num_tokens, while keepingmax_num_tokens=8192— unlike themax_num_tokens <= 4096workaround, this introduces no batch-size or throughput regression.The change is confined to the chunked-prefill path that hangs. The single-pass fallback branch (
else: _call_mqa_logits over [:num_ctx_tokens]) is intentionally left as-is: it is only reached for small, non-chunked prefills where the allocation is already bounded.Test Coverage
Isolation tests on GB200 (SM100), with
expandable_segments:True:max|diff| = 0.0) and identical top-k index sets. The only per-element index differences are pre-existing top-k tie-ordering non-determinism that is also present full-vs-full on identical inputs.End-to-end on the failing configuration (disaggregated DEP4: context tp4 / ep4 with
attention_dp,max_num_tokens=8192,free_gpu_memory_fraction=0.6,expandable_segments:True, ISL up to ~801K, ~1.15M compressed-KV tokens observed):Hang detected after 300 seconds), with the benchmark frozen at ~700 completed requests.PR Checklist
Please review the following before submitting your PR:
q_tilesizing helper if reviewers prefer.GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.