[Wanli-Jiang]

@coderabbitai summary

Description

Seven independent changes found while serving Qwen3.8-Flash-Next-FP8 with the modeling_v2 target (disaggregated and aggregated, GB300). Each is one commit with its own tests and can be reviewed, or split out, on its own; none depends on another or on the model. Six are fixes (a crash in error handling, a dropped request hint, Triton kernels in speculative decoding, causal_conv1d and QSA that recompiled inside served steps, AutoTuner tactics that depended on the first tuned call); one is a small opt-in performance change (an early PDL release in the MNNVL one-shot all-reduce).

Stop passing exc_info to the TensorRT-LLM logger

The TensorRT-LLM Logger methods take only positional message parts
(Logger.error(*msg)), so a call with exc_info=True raises TypeError
inside the except block that was meant to log the error. In the serve
iteration-stats collector this ended the collector loop on the first
stats RPC timeout: a disaggregated context-server log shows an RPC
timeout in fetch_stats_wait_async followed by "TypeError: Logger.error()
got an unexpected keyword argument 'exc_info'", and no iteration stats
after it.

Format the traceback into the message instead, at the four call sites
on main: the serve stats collector, the two perplexity handlers in
executor/result.py and the RPC client response reader. The new test
checks that no tensorrt_llm module passes exc_info to it (modules using
a stdlib logging.getLogger logger are skipped).

serve: carry the sub-agent parent hint on /v1/completions

The disaggregated router signs the sub-agent parent id
(conversation_params.subagent_affinity_id) into the request headers for
both chat and completions requests, but only the /v1/chat/completions
handler turned it into SchedulingParams. /v1/completions dropped it, so
a context or generation worker serving completions never saw the parent
id. A completions replay of a sub-agent trace showed no effect from
parent-based cache handling until the hint was carried.

Build SchedulingParams in the completions handler too and pass it to
generate_async. _get_scheduling_params now accepts a CompletionRequest
(it has no agent_hierarchy field, so that stays None).

Only requests that carry an authenticated parent header change. For
those, the ADP router's conversation affinity on /v1/completions now
uses the parent id, as it already does on chat.

The new test round-trips a context-only completion request through the
wire format and checks that the worker recovers the parent id.

speculative: compile the fused accept and step-0 kernels once across batch splits

fused_accept._strict_accept_kernel and _step0_epilogue_kernel
specialized on the per-step request counts (== 1, divisible by 16), the
accept tile size followed the batch, and the accept kernel and
spec_tail_kernels._strict_accept_kernel specialized on the alignment of
the token buffers, whose start follows the batch layout. Each new batch
split therefore compiled and loaded a new variant inside served
requests, which shows up as latency spikes and voided two benchmark
screens that audit compilation in their timed window. Mark the counts
do_not_specialize, the token pointers do_not_specialize_on_alignment,
and fix the accept tile at 64 rows (the kernel loops over row blocks,
rows are independent). This extends what #21424 did for
spec_tail_kernels. A sweep of 1,704 calls (0 to 17 context and 0 to 128
generation requests, draft length 2 and 3, sliced token inputs at
several offsets, the three forced-acceptance modes) compiled 12
variants of the fused kernels instead of 312 and 2 of the spec tail
kernel instead of 160, with bitwise identical outputs.

Tests: the accept and step-0 kernels keep one compiled variant across
batch splits and token-view offsets; the spec tail compile-count test
also takes the target tokens at an unaligned offset. The existing
fused-vs-unfused tests cover the results.

Without this change, a 3600 s AgentX measurement of a one-model MTP-Eagle deployment on main (Qwen3.8-Flash-Next-FP8, tep4, MTP3, concurrency 1) compiled a new _step0_epilogue_kernel variant inside its timed window.

causal_conv1d Triton: no per-slot-count or per-batch specialization

_causal_conv1d_fwd_kernel and _causal_conv1d_update_kernel took
num_cache_lines as a tl.constexpr, so every conv-state pool size (it
follows the memory left at startup) compiled its own variant, although
the value only bounds a mask. The update kernel also let Triton
specialize batch (a value of 1 gets its own cache entry with the same
machine code) and the 16-byte alignment of the slot, snapshot-row,
accepted-count and circular-buffer index vectors, which are read one
element per program and are often views into a larger vector on mixed
steps. A new request count or slot offset therefore compiled a new
variant inside served requests.

Make num_cache_lines a runtime argument kept out of the cache key, and
mark batch do_not_specialize and the index pointers
do_not_specialize_on_alignment. The generated code is the same, so the
outputs are bitwise unchanged. This applies to every Mamba / GDN model
that uses these kernels (Qwen3-Next, Qwen3.5, Nemotron-H).

A latency benchmark of a GDN model failed its timed-compilation audit
twice on _causal_conv1d_update_kernel: the TTIR of the new variants
differed from the nearest cached one only in the slot-count constant of
the mask (699 vs 696). After num_cache_lines became a runtime argument,
the remaining new variants differed only in the batch == 1
specialization. With both fixes the next run passed the audit on the
first try.

The new test runs the update kernel over request counts 1 to 5, index
vector offsets 1 to 3 elements and slot counts 16 to 32, checks that no
new variant is compiled, and compares each launch with the same launch
on aligned copies of the index vectors.

Without this change, a 3600 s AgentX measurement on main (Qwen3.8-Flash-Next-FP8, tep4, concurrency 1) compiled a new _causal_conv1d_update_kernel variant inside its timed window.

QSA Triton: no per-batch or per-offset specialization of the metadata kernels

Several stock QSA Triton kernels let Triton specialize on per-step row
and request counts (== 1, divisible by 16) or on the 16-byte alignment
of row-metadata pointers. Both change from step to step: the counts
follow the batch, and the metadata vectors are sliced at request or
chunk boundaries (for example the chunked index-scores loop in
module.py). Each first-seen combination compiled and loaded a new
variant inside served requests. Mark the counts do_not_specialize and
the metadata pointers do_not_specialize_on_alignment in the unscale,
decode and generation token mapping, prefill compress, block-index
expansion, paged index scores (row-wise and query tile) and MTP top-k
reuse kernels. The integer arithmetic is unchanged, so outputs are the
same.

Served benchmark runs compiled a new _qsa_generation_token_mapping_kernel
variant inside their timed windows because of the count and
pointer-alignment specializations: a high-throughput point at
concurrency 192 (twice, 700 to 870 s into the window) and a latency
point at concurrency 1. In a Triton-cache sweep over the runtime ranges
of those deployments (metadata view offsets 0 to 3), the decode token
mapping compiled 6 new cubins in the timed phase before the change and
none after. In an isolated microbenchmark the generation token mapping
without the alignment specialization was 1.36x to 1.39x slower at 64
and 192 requests (4.65 to 6.47 us).

Tests: the decode and generation token mapping, block-index expansion,
paged index scores and block-table unscale keep one compiled kernel per
launch class across row counts and metadata view offsets, with the
reference (mapping) or aligned-copy (others) results.

trtllm-gen MoE AutoTuner: draw dummy routing for every profiled bucket

With precomputed top-k (attention DP), the trtllm-gen MoE runners tune
on dummy routing from prepare_dummy_topk_and_hook. The AutoTuner builds
each bucket's inputs from the tuned call's tensors: float tensors are
redrawn at the bucket's shape, but an int [num_tokens, top_k] tensor is
filled by repeating the tuned call's own rows. Its shape therefore
already matches the bucket when the pre-hook sees it, and the pre-hook
only redrew the dummy top-k when the row count differed. A bucket larger
than the first tuned call was then profiled with that call's few routing
rows repeated, every row on the same few experts, where the widest token
tile wins. The chosen tactics depended on which call happened to be
tuned first.

Redraw the dummy top-k in the pre-hook for every profiled bucket,
whatever its shape.

Seen on an FP8 block-scale MoE (512 experts, 256-expert local window,
top-10): when the first tuned call was an 8192-row warm-up, every
decode bucket chose token tile 8. When it was a small call, the buckets
chose the widest tiles: [32, 8] for 32 to 216 rows, [64, 2] for 288 to
480 rows and [128, 0] for 528 and 576 rows. Measured with real routing,
[64, 2] is 6% to 26% slower than (32, 8) at 288 to 480 rows, and
[128, 0] takes 385 us vs 195 to 201 us at 528 and 576 rows, about
+9 ms per step at batch 176 and 192. The same hook serves every
trtllm-gen MoE runner (FP8 block scale, NVFP4, MXFP4, W4A8), so tactics
can change for those models: they no longer depend on the order of the
first tuned call.

Tests: the pre-hook gives every bucket its own dummy routing (random
and balanced distributions); both cases failed on the unfixed source
("64-row bucket profiled with 8 distinct routing rows"). With the
fix, the same model's decode tactics were [8, 8] up to 480 rows and
[16, 8] at 528 and 576 rows whatever the tuning order, and a
batch-192 decode step was 9% faster than with the order-dependent
tactics.

MNNVL one-shot all-reduce: optional early PDL release

Add trigger_completion_at_end (default True, the current behavior) to
trtllm::mnnvl_fusion_allreduce, the same argument the IPC all-reduce op
already takes. With False, the one-shot kernel calls
cudaTriggerProgrammaticLaunchCompletion right after its own grid
dependency wait instead of after the Lamport reduction. A dependent reads
the output only after its own griddepcontrol.wait, which waits for the
whole all-reduce grid, so values are unchanged; only the dependent's
prologue (for example a weight prefetch before its wait) overlaps the
exchange. The two-shot path ignores the argument.

A caller passes it per call, so no process environment variable is
needed.

Evidence (Qwen3.8 FP8, GB300, tep4 MTP-2, batch 1, 98k context; every
decode block's TP sum runs on the one-shot kernel; same-node step
microbenchmark, p10 over 3 requests): 3.218 ms/step with the early
release, 3.287 ms with only it turned off (about -2.1 %). The existing
test_mnnvl_allreduce.py suite passed (201) with the early release forced
on for every one-shot call by a process switch, the form this argument
replaces.

test_mnnvl_oneshot_early_pdl_release chains one-shot calls (sum, sum of
the sum, residual RMSNorm), each a PDL dependent reading the previous
output, and checks them bit for bit against the default trigger.

Test Coverage

• tests/unittest/utils/test_logger.py: a source scan that no logger.<level>(..., exc_info=...) call is left in tensorrt_llm.
• tests/unittest/disaggregated/test_disagg_internal_auth.py::test_completion_worker_affinity_survives_wire_roundtrip: a /v1/completions request carrying the signed parent hint keeps it through the wire format, and the worker recovers the parent id.
• tests/unittest/_torch/speculative/test_fused_accept.py::test_accept_kernels_compile_once_across_batch_splits and test_spec_tail_kernels.py: one compiled accept / step-0 kernel across batch splits and token-view offsets; the existing fused-vs-unfused tests cover the results.
• tests/unittest/_torch/modules/mamba/test_gdn_kernel_optimizations.py::test_conv1d_update_compiles_once_across_request_counts_and_slots: one compiled update kernel across request counts, index-vector offsets and slot counts, with the results of aligned copies.
• tests/unittest/_torch/attention/sparse/qsa/test_qsa_sparse.py: decode and generation token mapping, block-index expansion, paged index scores and block-table unscale keep one compiled kernel per launch class across row counts and metadata offsets.
• tests/unittest/_torch/misc/test_autotuner.py::test_trtllm_gen_moe_dummy_topk_is_redrawn_at_every_bucket (random and balanced): fails before the change, passes after.
• tests/unittest/_torch/multi_gpu/test_mnnvl_allreduce.py::test_mnnvl_oneshot_early_pdl_release: chained one-shot calls with trigger_completion_at_end=False match the default bit for bit.
• On GB300 (4 GPUs), native code rebuilt, at 22d14c8 (the same fixes rebased on main a6c28b4, before the QSA and test trimming):
  • tests/unittest/utils/test_logger.py and tests/unittest/disaggregated/test_disagg_internal_auth.py: 102 passed.
  • tests/unittest/_torch/speculative/test_fused_accept.py, test_spec_tail_kernels.py, hw_agnostic/: 1586 passed; the 33 failures need model checkpoints the test machine does not have (HFValidationError on the model path), and those rerun on main a6c28b4 fail the same way.
  • tests/unittest/_torch/modules/mamba: 1498 passed, 2 skipped.
  • tests/unittest/_torch/attention/sparse/qsa: 266 passed.
  • tests/unittest/_torch/moe/test_moe_backend.py -k TRTLLM: 80 passed, 16 skipped.
  • tests/unittest/_torch/multi_gpu/test_mnnvl_allreduce.py: 204 passed.
• At this head (bce2ebc) on GB300, native code rebuilt: test_logger.py and test_disagg_internal_auth.py 101 passed; test_fused_accept.py and test_spec_tail_kernels.py 43 passed; test_gdn_kernel_optimizations.py 103 passed; test_qsa_sparse.py 104 passed; test_autotuner.py 89 passed; test_mnnvl_allreduce.py -k early_pdl 2 passed.

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)

• If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.

• 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.