[None][perf] Add TRTLLM_DSV4_MEM_OPTS master switch for DSv4 memory optimizations

[qiaoxj07]

Summary

Two related categories of changes for DSv4-Pro on GB300 with attention-DP (ADP) + MoE EP:

1. Memory optimizations under a single env switch TRTLLM_DSV4_MEM_OPTS (the original scope of this PR).
2. Autotuner deadlock fixes for DSv4-Pro ADP runs — pre/post tp_barrier around _run_autotuner_warmup and a collective-vote on cache-hit in AutoTuner.choose_one. These are unconditional (not gated on the mem-opts env switch).

A complementary rank-divergent-warmup-batch issue is addressed by upstream PR #14126; the local fix that previously lived in this branch has been reverted to defer to that PR.

With TRTLLM_DSV4_MEM_OPTS off (default), behavior is byte-identical to current source for the mem-opts code paths. The autotuner fixes are unconditional but only affect the autotuner warmup phase, not steady-state inference.

What's in the env switch (TRTLLM_DSV4_MEM_OPTS=1)

1. Rope cos/sin table cap — tensorrt_llm/_torch/attention_backend/interface.py adds RopeParams.from_config auto-cap of max_positions to runtime max_seq_len. Yarn cos/sin at position p depends only on p and yarn constants (which use original_max_position_embeddings, not table size), so truncating to max_seq_len yields identical values for indices < cap. tensorrt_llm/_torch/models/modeling_deepseekv4.py plumbs llm_args.max_seq_len onto the HF PretrainedConfig so from_config can see it.

2. Attention workspace pre-allocation — tensorrt_llm/_torch/attention_backend/trtllm.py sizes the workspace at AttentionMetadata.__init__ time from the live model_config using an empirically validated formula (2 × max_num_tokens × num_heads × head_size × bytes_per_elt × 1.25 for paged-context FMHA, MLA / ADP / FP8-KV aware). Without this, the C++ runtime's lazy workspace_.value().resize_(...) tries a single 0 → ~1.6 GiB jump during warmup at large max_num_tokens, fails under tight memory, and leaves CUDA state degraded enough that the MoE kernel later hits CUDA_ERROR_LAUNCH_FAILED. AttentionMetadata gains an optional model_config field; tensorrt_llm/_torch/pyexecutor/model_engine.py wires it in.

3. MHC pool routing — tensorrt_llm/_torch/modules/mhc/mhc_cuda.py routes _FusedHcWorkspaceCache outputs (residual_cur, post_mix_cur, etc.) through the global Buffers pool so MHC scratch can share backing memory with other layer-transient buffers. For DSv4-Pro ctx-only shapes the resulting alignment is naturally 128-byte friendly so TMA descriptor encoding does not fail.

Autotuner deadlock fixes (always on, gated by tp_size > 1)

4. Pre/post tp_barrier around _run_autotuner_warmup

Symptom: when Cache size after warmup is 0, the autotuner exits near-instantly on ranks that already had cache, while a slower rank is still inside its own autotuner. Faster ranks race ahead into the next collective-issuing _general_warmup step → mismatched collectives → silent hang.

Fix: bracket _run_autotuner_warmup with TP-group barriers so every rank enters and exits the autotuner together.

5. AutoTuner.choose_one collective vote on cache hit

Symptom: with ADP + MoE EP, different ranks see different per-expert routing shapes → is_cache_hit diverges across ranks. The slow path of choose_one ends in _maybe_sync_cache_data, which issues a TP/CP collective (allgather/broadcast). If some ranks early-return at the cache-hit shortcut while others enter the collective, the collective deadlocks.

Fix: insert an unconditional tp_cp_allgather(is_cache_hit) collective vote in choose_one. Every rank participates in the allgather regardless of local cache state. Only short-circuit when all ranks hit.

The collective must be unconditional (not gated on local is_cache_hit) — otherwise hit ranks call the allgather while miss ranks skip it, recreating the deadlock at the vote site.

Related upstream PR

NVIDIA/TensorRT-LLM#14126 ("Fail loudly on asymmetric warmup batch under attention-DP") addresses a separate but related rank-divergent failure mode where _create_warmup_request returns None on some ranks but not others. An earlier version of this branch had a local tp_allgather + min fix on curr_max_num_tokens for the same symptom (silent fix-and-proceed); it has been reverted in 7ff13466b7 to defer to #14126's fail-loudly approach and reduce review surface.

Validation

1-job baseline (single ctx-only run, batch=2, max_seq=8192, TP=4 EP=4, ADP on)

metric	value
bench throughput	50,456 tok/s
gsm8k strict-match (1319 samples)	96.89%
gsm8k flexible-extract	96.89%

5-job parallel (same config, all required env on)

Test	bench tok/s	gsm8k avg	Status
test1	50,681	96.40%	✓
test2	—	—	✗ Issue 5 (see follow-ups)
test3	50,427	96.47%	✓
test4	50,667	96.02%	✓
test5	50,441	96.40%	✓

4/5 = 80% pass rate, up from ~0% without these fixes and ~60% with only expandable_segments.

Required env (for parallel-job runs)

export TRTLLM_DSV4_MEM_OPTS=1                              # mem-opts master switch
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True    # eliminates caching-allocator fragmentation at 16k warmup

(The PR code is correct without expandable_segments; that env reduces an upstream OOM-cascade failure mode in concurrent runs.)

Memory footprint (rank 0, DSv4-Pro GB300 TP=4)

stage	with switch	baseline
memory after model weight load	235.83 GiB	237.32 GiB (-1.49 GiB)
memory at stage_end_model_engine_warmup	250.08 GiB	254.97 GiB (-4.88 GiB)
workspace resize warnings	0 (or 1 small ~20 MiB)	~9

Test plan

• DSv4-Pro ctx-only bench (TP=4 GB300, batch=2, max_seq=8192): 50,456 tok/s (single) / 50.4–50.7k tok/s (parallel) ✓
• trtllm-eval gsm8k (1319 samples, max_batch=64, max_seq=2048): 96.02–96.89% ✓
• 5-job parallel hang-fix validation: 4/5 pass (1 failure is independent MoE PARALLEL strategy issue tracked as follow-up)
• OFF behavior byte-identical for mem-opts paths (verified via PyTorch memory snapshot comparison)
• CI

Files touched

tensorrt_llm/_torch/attention_backend/interface.py | 36 ++
tensorrt_llm/_torch/attention_backend/trtllm.py    | 96 ++
tensorrt_llm/_torch/autotuner.py                   | 17 +
tensorrt_llm/_torch/models/modeling_deepseekv4.py  | 12 +
tensorrt_llm/_torch/modules/mhc/mhc_cuda.py        | 57 +
tensorrt_llm/_torch/pyexecutor/model_engine.py     | 24

Open follow-ups (not in this PR)

• MoE PARALLEL strategy tactic-split desync — in _profile_runners, _maybe_parallelize_tactics distributes tactics across ranks (each rank profiles 1/N). MoE forward has internal NCCL all-to-all that requires all ranks at the same dispatch payload. Different per-rank tactics → mismatched payloads → dispatch timeout. Proposed 1-line fix: change MoERunner.tuning_config.distributed_tuning_strategy from PARALLEL to INDEPENDENT so all ranks profile all tactics in the same order. Validated as the root cause for the 1/5 failure in 5-job parallel; deferred to a follow-up PR to keep this one focused.
• Bind getWorkspaceSizeForContext from C++ so the Python pre-alloc matches the runtime's request exactly (currently the heuristic undersizes by ~20 MiB and the C++ side does one quiet resize).
• Harden Buffers._view_as to pad to a 128-byte aligned offset so the MHC pool is safe for all models, not just DSv4-Pro shapes.

🤖 Generated with Claude Code

[qiaoxj07]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #48162 [ run ] triggered by Bot. Commit: cb66e08 Link to invocation

[tensorrt-cicd]

PR_Github #48162 [ run ] completed with state SUCCESS. Commit: cb66e08
/LLM/main/L0_MergeRequest_PR pipeline #37984 completed with status: 'SUCCESS'

CI Report

Link to invocation

[qiaoxj07]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #48349 [ run ] triggered by Bot. Commit: 7ff1346 Link to invocation

[tensorrt-cicd]

PR_Github #48349 [ run ] completed with state SUCCESS. Commit: 7ff1346
/LLM/main/L0_MergeRequest_PR pipeline #38155 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

[qiaoxj07]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #48407 [ run ] triggered by Bot. Commit: 7ff1346 Link to invocation

[tensorrt-cicd]

PR_Github #48407 [ run ] completed with state SUCCESS. Commit: 7ff1346
/LLM/main/L0_MergeRequest_PR pipeline #38209 completed with status: 'SUCCESS'

CI Report

Link to invocation

[lfr-0531]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #48736 [ run ] triggered by Bot. Commit: 7ff1346 Link to invocation

[tensorrt-cicd]

PR_Github #48736 [ run ] completed with state FAILURE. Commit: 7ff1346
/LLM/main/L0_MergeRequest_PR pipeline #38504 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