Eg/pwcg tmla

[MrGeva]

Quick primer: PWCG (Piecewise CUDA Graph)

CUDA graphs let you record a sequence of GPU ops once and replay them with almost no launch overhead. The catch: many real ops can't be recorded — paged attention, MoE, anything with data-dependent shapes or CPU control flow.

Piecewise CUDA Graph is the AutoDeploy compile mode that bridges the gap:

1. Take the traced FX graph and split it at "dynamic" op boundaries (attention, MoE, SSM, etc.).
2. Capture the static segments between them as individual CUDA graphs.
3. Run dynamic ops eagerly between replays.
4. Pad token counts up to a fixed bucket (e.g. 256, 512, …, 8192) so static graphs only need capturing at a small set of sizes (piecewise_num_tokens).

Key invariants once that's set up:

• Static replays demand stable input/output addresses → dynamic ops write into pre-allocated out= buffers.
• Bucket padding means the tensor has more rows than real tokens. Most ops happily ignore the extra zero rows. A few don't, and those need to be told the real count.

---

1. 311c02ddf4 — "mmlu 82.7, gsm8k not tested" (the big enablement, ~1000 lines)

This is where DeepSeek-R1 actually gets PWCG turned on, and it's the bulk of the engineering work.

a. Config flip in examples/auto_deploy/model_registry/configs/deepseek-r1.yaml

Turns piecewise on, sets buckets, raises max_num_tokens to 15360:

compile_model:
  piecewise_enabled: true
  piecewise_num_tokens: [256, 512, 1024, 2048, 3072, 4096, 5120, 6144, 8192]

b. Token-count-aware MoE — new dynamic-op category

Touches piecewise_utils.py, piecewise_runner.py, torch_cudagraph.py, trtllm_moe.py.

The problem: trtllm_quant_finegrained_fp8_moe_fused routes per token. If the input has 4096 rows but only 3000 are real, the 1096 padded rows still get routed to experts and skew load/output.

The fix introduces a third category of dynamic op (_TOKEN_COUNT_DYNAMIC_OPS) on top of the existing inplace / metadata-prep / persistent-buffer categories:

• A new classifier needs_token_count_arg(submod) walks the FX submodule and returns True if it contains one of these ops.
• _inject_out_param is extended (inject_num_tokens=True) to add a num_tokens placeholder to the FX submodule and wire it as a kwarg to the op call.
• DynamicOpWrapper gains pass_real_num_tokens, so at runtime it sets kwargs["num_tokens"] = real_nt (the unpadded count tracked via a new class-level _current_real_num_tokens).
• The MoE op itself slices x, selected_experts, routing_weights to [:effective_num_tokens] before the kernel, then a new _finalize_token_sliced_moe_output zero-pads the result back to bucket size so the next static segment sees the captured shape.

c. Trailing static partition runs eagerly (TrailingEagerStaticWrapper)

The last static piece in the split (typically lm_head + logits) is intentionally not captured — it's small and the cost of capturing isn't worth it.

But this creates a shape mismatch: during replay, that trailing piece may receive a real-token-sized tensor (because dynamic ops produce real-token output) alongside bucket-sized integer scalars that FX baked in as args to view/reshape.

The new wrapper, at prepare time, statically inspects the trailing FX graph to find:

• Placeholder indices used as shape arguments to view/reshape ops → integers to rewrite.
• Placeholder indices that hold token-shaped activations (e.g. residual branches) → tensors to narrow.

At runtime, when real_nt < bucket_nt, it swaps any int placeholder equal to bucket_nt with real_nt, and narrows token-tensor placeholders to real_nt. The trailing eager code then runs on the correct real-token shapes.

d. MLA gets an out= buffer (trtllm_mla.py)

trtllm_mla_with_cache (and its impl/fake) now accept out: Optional[torch.Tensor]. When provided:

• The output tensor is sized to the bucket capacity (capacity-checked at runtime).
• Real positions are filled, the padded tail is zeroed.
• Returns a 0-element dummy (custom ops can't return tensors aliasing inputs).

This is what lets MLA participate in the PWCG out= plumbing like all other dynamic attention ops.

Two related tweaks in the same file:

• Workspace grows from 256 MB → 512 MB (DeepSeek's larger context buckets exhaust the smaller scratch and thop.attention warns/produces wrong output).
• auto_deploy::trtllm_mla_prepare_metadata is moved from _METADATA_PREP_OPS to _PERSISTENT_BUFFER_OPS in piecewise_utils.py. Its outputs already live in stable-address persistent buffers, not freshly allocated each call, so it should be classified that way.

e. Decode-only detection now considers extend requests

In DualModeCapturedGraph._is_decode_only:

# old
return num_prefill == 0
# new
return num_prefill == 0 and num_extend == 0

Extend requests are chunked-prefill continuation chunks. They need the prefill/mixed code path (cached-context attention), not the decode-only CUDA graph path. Without this fix, an "extend" batch would silently be treated as decode-only and corrupt output.

f. Capture-time buffer lifetime fix (ADPiecewiseRunner.finalize_capture)

Subtle but important. Dynamic-op outputs are allocated inside the preceding static runner's CUDA graph capture. Other downstream static segments capture in the same bucket and consume those buffers as inputs.

Previous code immediately weak-ref'd those dynamic-out buffers, which let CUDA's shared graph pool reuse the same addresses for later graph outputs during the same bucket capture — silently corrupting things. The fix keeps strong refs through the whole bucket capture, then finalize_capture(nt) downgrades them to weak refs once that bucket is fully captured.

The same commit also adds tail-zeroing in _copy_to_static_buffers (zero out padded rows of static-input buffers when runtime length < bucket length).

g. Tests added

• test_captured_graph.py (+225 lines), test_trtllm_mla_op.py (+264), test_ad_trtllm_serve.py (+83 lines, with a tiny-llama tokenizer/config helper for hermetic serve tests).
• The accuracy test for DeepSeek-R1-0528 PWCG is changed to skip MMLU (only run GSM8K) — see commits 2 and 3 for why.

---

2. fe686d83a4 — "[None][test] re-enable MMLU for DeepSeek-R1-0528 PWCG accuracy test"

Trivial diff — three lines in tests/integration/defs/accuracy/test_llm_api_autodeploy.py:

# old (from commit 1)
[
    # MMLU,
    GSM8K,
],
# new
[MMLU, GSM8K],

The commit message says PWCG + trtllm_mla now passes MMLU (~82.7) and GSM8K (~94) on DeepSeek-R1-0528.

One ordering note worth flagging: at the moment of this commit the actual fix that made MMLU pass (commit 3, ab5a33b8b2) wasn't committed yet — it landed 19 minutes later. So this re-enablement commit is technically "ahead" of the source change it depends on. Functionally fine after both land; just a quirk in the history.

---

3. ab5a33b8b2 — "[None][fix] AutoDeploy MLA: pass real past_kv to context FMHA for cache reuse"

Single-file fix in tensorrt_llm/_torch/auto_deploy/custom_ops/mla/trtllm_mla.py. This is the change that actually unlocked MMLU accuracy.

What was wrong

MLA's prefill code path (_handle_prefill_thop) was deliberately lying to the attention kernel:

# OLD: force "fresh prefill" semantics
thop.attention(
    ...,
    sequence_length        = context_lengths[:pf],             # new tokens only
    host_past_kv_lengths   = planner.past_kv_zeroed_host[:pf], # all zeros
    ...
)

The old comment said this was a workaround for the kernel's cached-KV path "producing garbage output (output magnitude blew up by ~10×num_decode)." Setting past_kv to 0 made the kernel treat every prefill as fresh, which worked as long as no prefill ever had cached prefix tokens.

But that assumption breaks the moment AD's scheduler does:

• Block reuse — e.g., MMLU has a shared 5-shot prompt prefix. The first request fills it into the KV cache; subsequent requests match that prefix and start with begin_compute > 0.
• Chunked prefill — long prompts split into multiple chunks, each chunk after the first sees cached prior chunks.

In both cases, host_past_kv_lengths is non-zero, but the old code zeroed it before passing to the kernel. Result: the kernel skipped the cached prefix entirely, attended only over new tokens, and produced a wrong output for the last prefill token — which is exactly the token whose logits drive the first generated token. Hence the catastrophic accuracy drop on MMLU specifically (long shared prefix, generation depends on that one logit).

GSM8K is less affected because the shared prefix is shorter and the answer extraction is more tolerant.

The fix

• Pass real host_past_kv_lengths[:pf] and real sequence_length[:pf] (which equals past + new) into thop.attention.
• ctx_total_kv_lens_host[0] now sums past + new tokens (so the kernel sizes its cached-context FMHA scratch correctly), not just new.
• The unused past_kv_zeroed_host planner buffer is deleted.
• Function signatures updated to accept sequence_length and host_past_kv_lengths; the caller in _mla_with_cache_impl passes them through.

The original "garbage output" the workaround was hiding evidently no longer reproduces — likely fixed in the underlying kernel between when the workaround was introduced and now — so simply trusting the real metadata works.

---

TL;DR

• 311c02 turns PWCG on for DeepSeek-R1 and adds the missing infrastructure for it: token-count-aware MoE, eager trailing partition with shape rewrite, MLA out= buffer, decode-only detection fix, capture-buffer lifetime fix.
• fe686d8 re-enables MMLU in the accuracy test (a one-liner anticipating commit 3).
• ab5a33b is the actual root-cause fix for MMLU: stop zeroing past_kv in the MLA prefill call, so cached-prefix / chunked-prefill scenarios produce correct attention output.

Why Did PWCG Expose These Bugs? Do They Affect Non-PWCG?

Per-fix classification of the changes in commits 311c02ddf4 and ab5a33b8b2, by whether PWCG introduced the bug or just exposed a pre-existing one, and whether non-PWCG runs are affected.

---

Pre-existing bugs that PWCG only exposed

These are real correctness issues in eager / non-PWCG mode too. PWCG didn't cause them, it just turned on workloads that finally hit them.

MLA prefill past_kv=0 workaround (ab5a33b)

This is a non-PWCG bug, full stop. The MLA prefill op was producing wrong output any time host_past_key_value_lengths[i] > 0 — i.e., whenever block reuse or chunked prefill kicked in — regardless of compile backend. Eager users with cache reuse have always been silently getting wrong logits on the last prefill token.

Why was it only caught now?

• Earlier accuracy gating used GSM8K (8-shot, short shared prefix, generation-tolerant) which barely exercised block reuse.
• Re-enabling MMLU (5-shot, long shared prefix, single-token answer extraction) was the workload that exposed the prefix-reuse path hard — and that re-enablement happened as part of the PWCG bring-up.
• Coincidentally, PWCG also bumps max_num_tokens to 15360 and turns on chunked prefill, which exercises the same code path from the other angle.

So the bug travels with PWCG only by circumstance, not by mechanism.

Workspace 256 MB → 512 MB

Same flavor. thop.attention's workspace is per-planner and used in prefill regardless of compile backend. The crash only happened because the DeepSeek PWCG config raised max_num_tokens to 15360 and added bucket sizes up to 8192. Any non-PWCG run with the same context sizes would have hit the same warning/garbage output.

---

True PWCG-only issues

These bugs literally cannot exist in non-PWCG mode because they're consequences of the PWCG execution model itself.

MoE token-count awareness

The whole problem ("padded rows still get routed to experts") only exists because PWCG introduces padding via bucketing. In non-PWCG runs the input tensor has exactly real-token rows; MoE has nothing extra to ignore.

Trailing eager shape rewrite (TrailingEagerStaticWrapper)

The shape mismatch — real-token tensors from dynamic ops alongside bucket-sized integer scalars baked into FX — is created by the dynamic/static split. No split, no mismatch.

MLA out= buffer

A stable output address is only required when the next op is replaying a captured CUDA graph. In non-PWCG (or any non-graph-capture) mode the old torch.zeros(...) was perfectly fine. The change is purely additive — when out=None the original allocation path still runs.

Capture-time dynamic_out_bufs lifetime fix (finalize_capture)

This bug only exists during the bucket capture loop, which only exists in PWCG.

Tail-zeroing in _copy_to_static_buffers

Padding only exists under PWCG.

trtllm_mla_prepare_metadata classification move

The _METADATA_PREP_OPS / _PERSISTENT_BUFFER_OPS categories are only consulted by the PWCG graph splitter. Misclassification is invisible without PWCG.

---

Mode-shared but graph-capture-specific

Decode-only detection considering num_extend

Lives in DualModeCapturedGraph._is_decode_only. DualMode is the dispatcher that picks between piecewise (prefill/mixed) and a decode-only CUDA graph. The bug — treating an extend batch as decode-only and skipping cached-context attention — would manifest in any captured-decode-graph mode, not strictly PWCG. In pure eager (no graph capture at all) extend just runs the normal model path and nothing breaks.

So: not PWCG-only, but also not present in pure eager.

---

Why did they all surface together?

Enabling PWCG for DeepSeek-R1 was a multi-axis change rolled out in one go:

1. New execution model — bucketing, splitting, capture/replay → exposes invariants nobody had to think about before (token-count awareness, stable addresses, capture pool lifetimes, shape pollution into the trailing partition).
2. Bigger workload knobs — max_num_tokens=15360, chunked prefill on, large context buckets → exercises code paths (workspace size, begin_compute > 0 prefills) that the previous tighter configs didn't reach.
3. Stricter accuracy gate — re-running MMLU with its 5-shot shared prefix is what actually finds the latent block-reuse bug.

The MLA prefill bug is the noteworthy one for non-PWCG users: if anyone was running DeepSeek-R1 (or any MLA model) in eager AutoDeploy with KV-cache block reuse or chunked prefill, they were getting silently degraded accuracy on the first generated token. That fix benefits every backend, not just PWCG.

---

Summary table

Fix	PWCG-only?	Affects non-PWCG?
MLA past_kv=0 workaround removal	No	Yes — silent wrong-logits on cache reuse / chunked prefill in any backend
MLA workspace 256→512 MB	No	Yes, for any large-context prefill
MoE token-count awareness	Yes	No
Trailing eager shape rewrite	Yes	No
MLA out= buffer	Yes	No (additive; old path preserved)
finalize_capture lifetime fix	Yes	No
Tail-zeroing static input buffers	Yes	No
trtllm_mla_prepare_metadata reclassification	Yes	No
Decode-only num_extend == 0 check	Graph-capture-only	Any captured-decode-graph mode, not pure eager

@coderabbitai summary

Description

Test Coverage

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.