[None][perf] Autotuner: distributed tuning for mhc + single-(runner,tactic) shortcut by hyukn · Pull Request #14458 · NVIDIA/TensorRT-LLM

hyukn · 2026-05-22T09:42:39Z

Summary

Three refinements for JIT-driven custom ops (DeepGEMM, Triton, CuteDSL) and multi-tactic MHC ops in the autotuner:

TuningConfig.exclude_from_cache flag (default False). For JIT-driven ops, the tuning call itself triggers JIT compilation as a side effect. Persisting the chosen tactic across process startups skips the tuning path on the next run, so the JIT never warms up and the first inference request pays the compile cost on the critical path. The flag tells the autotuner to:
- force a re-tune in tuning mode (drives the JIT warmup as a side effect),
- keep the in-process cache working,
- drop the op from disk persistence in _partition_cache_by_strategy.
Applied to fp8SwapABGemmRunner.tuning_config.
Single-(runner, tactic) shortcut in _profile_runners. When an op exposes exactly one candidate pair across all runners, there is nothing to compare against. The warmup + profile-repeat loop in _profile_single_kernel is pure overhead for JIT-driven kernels whose only purpose for the round-trip is firing the JIT compile. The shortcut fires one forward() call (no timing, no warmup loop) and records the pair directly. do_preparation still fires.

Total-pair count is computed BEFORE distributed-strategy splitting so PARALLEL — which can leave a rank with one tactic locally while the full set has several — still goes through the timed path so the merge picks a real winner.

For fp8SwapABGemmRunner this reduces per-shape-bucket warmup from ~5+ forward calls to 1.
Switch MHC ops to DistributedTuningStrategy.PARALLEL. MhcPreMappingRunner and MhcFusedHcRunner both expose multi-tactic spaces. Fanning their tactics out across TP ranks (each rank times a disjoint subset, then merge picks the global best) collapses the MhcPreMapping tuning span from ~5–6 min to ~2 min and MhcFusedHc from ~2 min to a few seconds with no change in the chosen tactic.

End-to-end measurements (cold autotune, no autotuner/DG cache, dsv4-325152 container, Pro 16K recipe, dataset 1024-1-20000-ratio-08):

config	autotune block 1	E2E throughput
TP=4 EP=4, baseline (PR 14329 already in container)	14m 35s	64.1 req/s
TP=4 EP=4, + this PR	5m 50s (-60%)	63.6 req/s (within ±1% noise band, n=3)
TEP8, this PR vs INDEPENDENT (n=1 each)	2m 41s vs 6m 19s (-57%)	99.01 vs 99.17 req/s (-0.17%)

Tactic-choice agreement vs INDEPENDENT reference at TEP8 (per-shape comparison, 627 unique shapes):

96% of shapes pick the same tactic
17 diff-tactic shapes (mhc_pre_mapping + mhc_fused_hc) all have A/B min_time ratio in 0.99–1.03 (max single-shape regression 2.8% at sub-10 µs)
0 shapes where PARALLEL is >5% slower than INDEPENDENT

🤖 Generated with Claude Code

hyukn · 2026-05-25T03:32:35Z

/bot run --disable-fail-fast --add-multi-gpu-test

tensorrt-cicd · 2026-05-25T03:37:53Z

PR_Github #50135 [ run ] triggered by Bot. Commit: 391d28f Link to invocation

tensorrt-cicd · 2026-05-25T06:18:07Z

PR_Github #50135 [ run ] completed with state SUCCESS. Commit: 391d28f
/LLM/main/L0_MergeRequest_PR pipeline #39687 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

hyukn · 2026-05-25T08:16:05Z

/bot run --disable-fail-fast --add-multi-gpu-test

tensorrt-cicd · 2026-05-25T08:22:34Z

PR_Github #50180 [ run ] triggered by Bot. Commit: 391d28f Link to invocation

tensorrt-cicd · 2026-05-25T10:08:22Z

PR_Github #50180 [ run ] completed with state SUCCESS. Commit: 391d28f
/LLM/main/L0_MergeRequest_PR pipeline #39723 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

hyukn · 2026-05-26T01:25:18Z

/bot run --disable-fail-fast --add-multi-gpu-test

tensorrt-cicd · 2026-05-26T01:31:22Z

PR_Github #50248 [ run ] triggered by Bot. Commit: 8e9770a Link to invocation

tensorrt-cicd · 2026-05-26T05:20:59Z

PR_Github #50248 [ run ] completed with state SUCCESS. Commit: 8e9770a
/LLM/main/L0_MergeRequest_PR pipeline #39781 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

hyukn · 2026-05-26T09:02:29Z

/bot run --disable-fail-fast --add-multi-gpu-test

tensorrt-cicd · 2026-05-26T09:08:06Z

PR_Github #50332 [ run ] triggered by Bot. Commit: 8e9770a Link to invocation

tensorrt-cicd · 2026-05-26T09:51:45Z

PR_Github #50332 [ run ] completed with state SUCCESS. Commit: 8e9770a
/LLM/main/L0_MergeRequest_PR pipeline #39859 completed with status: 'SUCCESS'

CI Report

Link to invocation

hyukn · 2026-05-27T06:36:46Z

/bot run --disable-fail-fast --add-multi-gpu-test

tensorrt-cicd · 2026-05-27T06:43:22Z

PR_Github #50483 [ run ] triggered by Bot. Commit: 83241b9 Link to invocation

…actic) shortcut Two related refinements for JIT-driven custom ops (DeepGEMM, Triton, CuteDSL) in the autotuner. 1. New TuningConfig.exclude_from_cache flag, default False. For JIT-driven ops (e.g. fp8_swap_ab_gemm with DeepGEMM) the tuning call itself is what triggers JIT compilation as a side effect. Persisting the chosen tactic across process startups skips the tuning path on the next run, so the JIT never warms up and the first inference request pays the compile cost on the critical path. When exclude_from_cache=True: - choose_one forces a re-tune in tuning mode (drives the JIT warmup as a side effect) - the in-process cache still works for subsequent calls within the same process - _partition_cache_by_strategy drops the op from disk persistence Applied to fp8SwapABGemmRunner.tuning_config (the reported case). 2. Single-(runner, tactic) shortcut in _profile_runners. When a custom_op exposes exactly one candidate (runner, tactic) pair across all runners, there is nothing to compare against. The warmup + profile-repeat loop in _profile_single_kernel is pure overhead -- typical for JIT-driven kernels whose only purpose for the tuning round-trip is to fire the JIT compile as a side effect. The shortcut fires exactly one forward() call (no timing, no warmup loop) and records the (runner, tactic) pair directly in the profiling cache. do_preparation hooks (e.g. AllReduceRunner buffer pre-allocation) still fire. The total-pair count is computed BEFORE distributed-strategy splitting so the PARALLEL strategy -- which can leave a rank with one tactic locally even though the full set has several -- still goes through the timed profile path so the merge picks a real winner. For fp8SwapABGemmRunner this reduces per-shape-bucket warmup from ~5+ forward calls to 1. UT coverage (tests/unittest/_torch/misc/test_autotuner.py): - test_single_tactic_skips_profile_kernel - test_multi_tactic_still_invokes_profile_kernel - test_single_tactic_do_preparation_still_fires - test_single_tactic_records_cache_entry - test_single_tactic_across_multiple_runners_takes_normal_path Signed-off-by: Yukun He <23156053+hyukn@users.noreply.github.com>

MhcPreMappingRunner.tuning_config and MhcFusedHcRunner.tuning_config both expose multi-tactic tuning spaces. Switching their distributed_tuning_strategy from the default INDEPENDENT to PARALLEL fans the tactic list out across TP ranks (each rank times a disjoint subset, _merge_cache_data picks the global best). With dsv4 ranks at TP=4, this collapses the MhcPreMapping span from ~5-6 min to ~2 min and MhcFusedHc from ~2 min to a few seconds, with no change in the chosen tactic. Signed-off-by: Yukun He <23156053+hyukn@users.noreply.github.com>

hyukn · 2026-05-27T06:49:33Z

/bot run --disable-fail-fast --add-multi-gpu-test

tensorrt-cicd · 2026-05-27T07:10:19Z

PR_Github #50490 [ run ] triggered by Bot. Commit: 66c81a0 Link to invocation

tensorrt-cicd · 2026-05-27T07:14:21Z

PR_Github #50483 [ run ] completed with state ABORTED. Commit: 83241b9

Link to invocation

tensorrt-cicd · 2026-05-27T10:33:59Z

PR_Github #50490 [ run ] completed with state SUCCESS. Commit: 66c81a0
/LLM/main/L0_MergeRequest_PR pipeline #40003 completed with status: 'SUCCESS'

CI Report

Link to invocation

PR NVIDIA#14458 enabled `distributed_tuning_strategy=DistributedTuningStrategy.PARALLEL` on both MhcPreMappingRunner and MhcFusedHcRunner. The MhcFusedHc setting triggers a non-deterministic CUDA Illegal Memory Access on TP=8 disagg GEN workers during the post-autotune max-batch CUDA-graph capture (observed on the 1024-1024 disagg_ctx1_gen1_tep8_batch256 recipe and the 8192-1024 disagg_ctx1_gen1_dep8_batch512 recipe; ~67% crash rate across 3 reruns of the same config). Mechanism: MhcFusedHcRunner has a per-instance LRU workspace cache (`_FusedHcWorkspaceCache`) keyed on `(B, num_k_splits, tile_m, device)` — all derived from the chosen tactic. Under PARALLEL, each rank measures only `idx % tp_size == tp_rank` of the tactics, so each rank's LRU ends up populated with workspace entries for a different tactic subset. After `_merge_cache_data` writes the global-best tactic to the autotuner cache, ranks that didn't locally measure that tactic hit a workspace-cache miss and call `torch.empty(...)` fresh — inside the post-autotune CUDA-graph-capture region for the max-batch warmup step. Fresh allocations are incompatible with capture; the next collective kernel (MLA attention) dereferences a stale/uncaptured pointer → IMA on the cache-miss rank (rank 6 at TP=8 in the observed cases). MhcPreMappingRunner does NOT have this problem: its forward() uses fresh `torch.empty(...)` every call, keyed only on `(M, n, hidden_size)`, independent of the chosen tactic. So per-rank tactic-subset profiling under PARALLEL produces identical allocation histories across ranks. Per-op cost data on this disagg recipe shows MhcFusedHc PARALLEL adds ~no autotune saving (GEN: 4m 42s with both PARALLEL vs 4m 46s with only MhcPre PARALLEL — within JIT noise), so dropping it is a clean revert with no perf cost. MhcPre PARALLEL still buys ~27% CTX autotune reduction on the same recipe. A/B/C/D variant matrix at TP=8 disagg (3 nodes, max_batch=256): - MhcPre=ON, MhcFused=ON (A): IMA in 2/3 reruns - MhcPre=OFF, MhcFused=OFF (B): clean - MhcPre=OFF, MhcFused=ON (C): IMA - MhcPre=ON, MhcFused=OFF (D): clean ← this PR's configuration Re-enabling MhcFusedHc PARALLEL is a follow-up after the workspace LRU is reworked (e.g. by dropping tactic-dependent keying, or by harmonizing the LRU across ranks post-merge). Signed-off-by: Yukun He <23156053+hyukn@users.noreply.github.com>

…actic) shortcut (NVIDIA#14458) Signed-off-by: Yukun He <23156053+hyukn@users.noreply.github.com>

…actic) shortcut (NVIDIA#14458) Signed-off-by: Yukun He <23156053+hyukn@users.noreply.github.com> (cherry picked from commit 5d59d34)

hyukn requested a review from a team as a code owner May 22, 2026 09:42

hyukn requested review from yizhang-nv and removed request for a team May 22, 2026 09:42

github-actions Bot assigned hyukn May 22, 2026

hyukn mentioned this pull request May 22, 2026

[None][fix] Autotuner cache: merge instead of replace, atomic save #14441

Merged

3 tasks done

hyukn added the deepseek-v4 label May 25, 2026

hyukn requested a review from a team as a code owner May 25, 2026 03:24

hyukn requested review from mingyangHao and xxi-nv and removed request for a team May 25, 2026 03:24

hyukn force-pushed the feat/autotune_jit_op_prune branch from 391d28f to 8e9770a Compare May 26, 2026 01:24

hyukn requested a review from lfr-0531 May 27, 2026 01:38

yizhang-nv approved these changes May 27, 2026

View reviewed changes

hyukn force-pushed the feat/autotune_jit_op_prune branch 3 times, most recently from bbe1b25 to 4e942f3 Compare May 27, 2026 06:35

hyukn force-pushed the feat/autotune_jit_op_prune branch from 4e942f3 to 83241b9 Compare May 27, 2026 06:39

hyukn added 2 commits May 27, 2026 06:44

hyukn force-pushed the feat/autotune_jit_op_prune branch from 83241b9 to 66c81a0 Compare May 27, 2026 06:44

hyukn changed the title ~~[None][perf] Autotuner: opt-out for JIT-driver ops + single-(runner,tactic) shortcut~~ [None][perf] Autotuner: distributed tuning for mhc + single-(runner,tactic) shortcut May 28, 2026

hyukn merged commit 5d59d34 into NVIDIA:feat/deepseek_v4 May 28, 2026
10 checks passed

hyukn mentioned this pull request May 28, 2026

[None][fix] mhc: disable PARALLEL distributed tuning for MhcFusedHc #14675

Merged

lfr-0531 pushed a commit to lfr-0531/TensorRT-LLM that referenced this pull request Jun 10, 2026

[None][perf] Autotuner: distributed tuning for mhc + single-(runner,t…

732e277

…actic) shortcut (NVIDIA#14458) Signed-off-by: Yukun He <23156053+hyukn@users.noreply.github.com>

lfr-0531 pushed a commit to lfr-0531/TensorRT-LLM that referenced this pull request Jun 11, 2026

[None][perf] Autotuner: distributed tuning for mhc + single-(runner,t…

e251cb8

…actic) shortcut (NVIDIA#14458) Signed-off-by: Yukun He <23156053+hyukn@users.noreply.github.com>

lfr-0531 pushed a commit to lfr-0531/TensorRT-LLM that referenced this pull request Jun 12, 2026

[None][perf] Autotuner: distributed tuning for mhc + single-(runner,t…

3f854b1

…actic) shortcut (NVIDIA#14458) Signed-off-by: Yukun He <23156053+hyukn@users.noreply.github.com>

Conversation

hyukn commented May 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Uh oh!

hyukn commented May 25, 2026

Uh oh!

tensorrt-cicd commented May 25, 2026

Uh oh!

tensorrt-cicd commented May 25, 2026

Uh oh!

hyukn commented May 25, 2026

Uh oh!

tensorrt-cicd commented May 25, 2026

Uh oh!

tensorrt-cicd commented May 25, 2026

Uh oh!

hyukn commented May 26, 2026

Uh oh!

tensorrt-cicd commented May 26, 2026

Uh oh!

tensorrt-cicd commented May 26, 2026

Uh oh!

hyukn commented May 26, 2026

Uh oh!

tensorrt-cicd commented May 26, 2026

Uh oh!

tensorrt-cicd commented May 26, 2026

Uh oh!

hyukn commented May 27, 2026

Uh oh!

tensorrt-cicd commented May 27, 2026

Uh oh!

hyukn commented May 27, 2026

Uh oh!

tensorrt-cicd commented May 27, 2026

Uh oh!

tensorrt-cicd commented May 27, 2026

Uh oh!

tensorrt-cicd commented May 27, 2026

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

hyukn commented May 22, 2026 •

edited

Loading