[TRTLLM-10288][perf] Reduce AutoTuner host overhead in inference hot path - #13116
[TRTLLM-10288][perf] Reduce AutoTuner host overhead in inference hot path#13116hyukn wants to merge 2 commits into
Conversation
hyukn
changed the title
|
/bot run --disable-fail-fast |
📝 WalkthroughWalkthroughThe changes introduce performance optimizations in the autotuner and custom operations modules. The autotuner adds a memoized cache key prefix property and implements a fast path for non-tuning mode inference to skip unnecessary capture/replay branching. The custom ops module applies singleton pattern caching to avoid redundant instantiations of GEMM and quantization runners. Changes
Sequence Diagram(s)sequenceDiagram
participant Client
participant AutoTuner
participant Cache as AutoTunerProfilingCache
participant Runner as TunableRunner
Client->>AutoTuner: choose_one() [inference, not tuning]
activate AutoTuner
alt Fast Path [is_tuning_mode=false, no active capture]
AutoTuner->>Cache: get_cache_key(runner)
activate Cache
Cache->>Runner: _cache_key_prefix
activate Runner
Runner-->>Cache: memoized (class_name, unique_id)
deactivate Runner
Cache-->>AutoTuner: cache_key
deactivate Cache
AutoTuner->>AutoTuner: lookup cache by key
alt Cache Hit
AutoTuner-->>Client: (best_runner, best_tactic)
else Cache Miss
AutoTuner->>AutoTuner: log fallback warning
AutoTuner-->>Client: (best_runner, best_tactic)
end
else Slower Path [tuning mode or capture active]
AutoTuner->>AutoTuner: capture/replay logic
AutoTuner-->>Client: (best_runner, best_tactic)
end
deactivate AutoTuner
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes 🚥 Pre-merge checks | ✅ 3✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (2)
tensorrt_llm/_torch/autotuner.py (1)
256-269: Annotate the new cached property.The new
_cache_key_prefixproperty should declare its return type, e.g.-> tuple[str, str].As per coding guidelines "Always annotate Python function return types; use None if the function does not return anything."Suggested tweak
`@property` - def _cache_key_prefix(self): + def _cache_key_prefix(self) -> tuple[str, str]:🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tensorrt_llm/_torch/autotuner.py` around lines 256 - 269, Annotate the new cached property _cache_key_prefix with an explicit return type; change its signature to declare it returns a tuple of two strings (e.g., -> tuple[str, str]) so static type checkers and linters know the expected type, leaving the implementation using self.__cache_key_prefix and str(self.unique_id()) unchanged.tensorrt_llm/_torch/custom_ops/torch_custom_ops.py (1)
2237-2242: Annotate the new helper’s return type.
_get_quantize_e4m3_runner()is newly added and should declare-> QuantizeE4M3PerTensorRunnerto match the repo’s Python typing rules.As per coding guidelines "Always annotate Python function return types; use None if the function does not return anything."Suggested tweak
-def _get_quantize_e4m3_runner(): +def _get_quantize_e4m3_runner() -> QuantizeE4M3PerTensorRunner:🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tensorrt_llm/_torch/custom_ops/torch_custom_ops.py` around lines 2237 - 2242, Add a return type annotation to the singleton accessor: change the signature of _get_quantize_e4m3_runner so it declares "-> QuantizeE4M3PerTensorRunner". Update the function definition for _get_quantize_e4m3_runner() to include this return type (use a forward-reference string if QuantizeE4M3PerTensorRunner is not yet defined in the file) so it conforms to the repo's typing rules.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@tensorrt_llm/_torch/autotuner.py`:
- Around line 913-926: The fast path currently builds input_shapes by iterating
inputs and calling t.shape, which can raise AttributeError when callers pass
None placeholders (see choose_one callers); replace that logic by calling
self._get_input_sizes(inputs) to compute input_sizes safely and use that tuple
in the profiling_cache.search_cache call and log message (keep the rest of the
flow intact: guard on self.is_tuning_mode and self._active_capture, call
profiling_cache.search_cache with custom_op, runners, the result of
self._get_input_sizes(inputs), tuning_config, apply_map_to_tuning_buckets=True,
then use best_runner_id/best_tactic as before).
---
Nitpick comments:
In `@tensorrt_llm/_torch/autotuner.py`:
- Around line 256-269: Annotate the new cached property _cache_key_prefix with
an explicit return type; change its signature to declare it returns a tuple of
two strings (e.g., -> tuple[str, str]) so static type checkers and linters know
the expected type, leaving the implementation using self.__cache_key_prefix and
str(self.unique_id()) unchanged.
In `@tensorrt_llm/_torch/custom_ops/torch_custom_ops.py`:
- Around line 2237-2242: Add a return type annotation to the singleton accessor:
change the signature of _get_quantize_e4m3_runner so it declares "->
QuantizeE4M3PerTensorRunner". Update the function definition for
_get_quantize_e4m3_runner() to include this return type (use a forward-reference
string if QuantizeE4M3PerTensorRunner is not yet defined in the file) so it
conforms to the repo's typing rules.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: 2a4397f1-8afe-42e4-a9e3-dba4d7126e53
📒 Files selected for processing (2)
tensorrt_llm/_torch/autotuner.pytensorrt_llm/_torch/custom_ops/torch_custom_ops.py
|
PR_Github #43715 [ run ] triggered by Bot. Commit: |
|
PR_Github #43715 [ run ] completed with state
|
0a69051 to
1e313ca
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #45275 [ run ] triggered by Bot. Commit: |
|
/bot run --disable-fail-fast |
|
PR_Github #45348 [ run ] triggered by Bot. Commit: |
|
PR_Github #45348 [ run ] completed with state
|
19f4cce to
d57fd9a
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #45708 [ run ] triggered by Bot. Commit: |
|
/bot run --disable-fail-fast |
|
PR_Github #45708 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
1 similar comment
|
/bot run --disable-fail-fast |
d57fd9a to
40593c6
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #46355 [ run ] triggered by Bot. Commit: |
|
PR_Github #46355 [ run ] completed with state |
40593c6 to
7e9ca65
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #46880 [ run ] triggered by Bot. Commit: |
7e9ca65 to
e609b32
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #46948 [ run ] triggered by Bot. Commit: |
|
PR_Github #46948 [ run ] completed with state |
|
/bot run --disable-fail-fast |
|
PR_Github #46971 [ run ] triggered by Bot. Commit: |
|
PR_Github #46971 [ run ] completed with state |
e609b32 to
f741c77
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #47064 [ run ] triggered by Bot. Commit: |
| return (runners[best_runner_id], best_tactic) | ||
|
|
||
| # Cache miss — resolve via search_cache, populate cache | ||
| input_shapes = tuple(self._get_input_sizes(inputs)) |
There was a problem hiding this comment.
Duplicate code for the tuning mode and non tuning mode.
There was a problem hiding this comment.
The duplication is intentional — the two paths serve different purposes and have different lifetimes:
-
Non-tuning path (the new cache):
search_cachehere is only hit once per unique shape (on_choose_one_cachemiss). After warmup, this code is effectively dead — all subsequent calls hit the fastdict.get()above and return immediately. Storing the result in_choose_one_cacheis what makes this a one-time cost. -
Tuning path (existing code):
search_cachehere runs every iteration during active profiling and feeds into the "should we re-profile?" decision logic below it. It cannot use_choose_one_cachebecause the cache is cleared/invalid during tuning.
I considered extracting a shared helper, but it would need to return different things for each path (the non-tuning path just needs (runner_id, tactic) to cache; the tuning path also needs is_cache_hit and min_time for the profiling decision). The 4-line search_cache call is straightforward enough that a shared helper would add indirection without meaningful deduplication.
Happy to add a brief comment in the code clarifying why both exist if that would help readability.
|
PR_Github #47064 [ run ] completed with state
|
…path Phase 1 — WAR optimizations: - Cache runner key prefix (_cache_key_prefix) to avoid recomputing class name + unique_id string on every call - Fast path in choose_one() for non-tuning mode - Singleton runner caches for nvfp4_gemm and quantize_e4m3_per_tensor Phase 2 — dispatch() + bind_tactic(): - Fused resolve+execute via AutoTuner.dispatch() with _dispatch_cache - _make_fast_key() for lightweight keying on dynamic dimensions - TunableRunner.bind_tactic() returns callable bound to a tactic - Converted 21 callsites to dispatch(); 10 remain as choose_one() where tactic value is needed Profiling (B200, nvfp4_gemm, CUTLASS, N=1000): dispatch() cache hit: +1.08us overhead (saves 3.24us vs old path) Signed-off-by: Yukun He <23156053+hyukn@users.noreply.github.com> Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
f741c77 to
e831ea6
Compare
Move the fast-path cache from dispatch() into choose_one() so all callers benefit. The cache key now includes runner_ids (each runner's _cache_key_prefix) to prevent collisions when different runner configurations share the same custom_op name and bucketed dimensions (e.g. int8 vs int4 weight_only_quant_gemm runners). Cache stores (runner_id, tactic) instead of runner instances, resolving via the caller's runners list on hit. Simplify dispatch() to a thin wrapper over choose_one(). Delete the separate _dispatch_cache. Signed-off-by: Yukun He <23156053+hyukn@users.noreply.github.com> Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
e831ea6 to
a3c0cb0
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #48216 [ run ] triggered by Bot. Commit: |
|
PR_Github #48216 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #48315 [ run ] triggered by Bot. Commit: |
|
PR_Github #48315 [ run ] completed with state
|
|
The optimization will affect some of the tuning process and may hurt the overall inference perf. |
3 participants
Description
Reduces AutoTuner host overhead in the inference hot path. Profiling on B200 showed the
choose_one()→search_cache()→get_cache_key()path adds significant Python overhead on every custom op call. This PR eliminates the bulk of it.Changes
Phase 1: WAR Optimizations
get_cache_key()(avoids recomputing__class__.__name__+str(unique_id())per call)nvfp4_gemmandquantize_e4m3_per_tensorPhase 2:
_choose_one_cachefast path inchoose_one()_choose_one_cache: maps(custom_op, runner_ids, *bucketed_dims)→(runner_id, tactic). On cache hit: one_make_fast_key()+ onedict.get(), then resolve runner from caller's list._make_fast_key(): lightweight key using only dynamic dimensions + runner identity (_cache_key_prefix).dispatch(): simplified to thin wrapper overchoose_one()— delegates all caching to_choose_one_cache.TunableRunner.bind_tactic(): returns callable pre-bound to a specific tactic. Subclasses pre-resolve sub-runners.dispatch()acrosstorch_custom_ops.py(12) andcute_dsl_custom_ops.py(9).Cache collision fix (runner identity in key)
runner_ids(tuple of each runner's_cache_key_prefix) so that different runner configurations for the samecustom_op+ dimensions don't collide (e.g. int8 vs int4weight_only_quant_gemmrunners).(runner_id, tactic)instead of runner instances — resolution always uses caller'srunnerslist.Profiling Results
Clean wall-clock measurements on B200, nvfp4_gemm (M=128, K=4096, N=4096, CUTLASS, 500 iters, no nsys/NVTX overhead):
Overhead breakdown
The `_choose_one_cache` reduces Python overhead from 6.2us to 1.3us by replacing per-call runner construction + `choose_one()` tactic search with `_make_fast_key()` + `dict.get()`.
The remaining 6.5us is the `@torch.library.custom_op` framework dispatch cost (DispatchKeySet traversal, auto-functionalization, schema validation). This is addressable separately via `fast_custom_op` (PR #13149), which reduces it to 2.3us.
Per-step impact (40 tunable ops)
Correctness Verification
Cache collision fix
Single-GPU collision repro (B200, `weight_only_quant_gemm`, same dims 4096x7168x2112):
Unit tests
Files Modified
PR Checklist