[TRTLLM-14904][fix] Scope flashinfer autotuner WAR to the MLA decode config builder - #17256
Merged
brnguyen2 merged 1 commit intoAug 4, 2026
Conversation
…config builder Port of the reworked workaround from the main-branch PR (NVIDIA#17165, review by BowenFu): the class-wide DynamicTensorSpec.__eq__ replacement mutates a third-party class process-wide and its widened equality can cross-contaminate AutoTuner._override_config_cache entries whose configs differ only in tensor_initializers. Replace it with a workaround scoped to the pathological path: memoize flashinfer.mla._core._build_mla_decode_tuning_config, the one function that rebuilds a TuningConfig with fresh initializer closures on every decode call, keyed on the scalars that fully determine its output. Repeated calls reuse the same TuningConfig object, so autotuner cache hits resolve through flashinfer's own unmodified hash/eq — no equality semantics change anywhere. Every other TuningConfig construction site in flashinfer 0.6.15 already reuses its config across calls, so this one builder covers the entire pathology. Also folds in the two review fixes from NVIDIA#17165: drop the exc_info kwarg (tensorrt_llm's logger.debug(*msg) takes no kwargs, so the except handler would itself raise at import time in the incompatible-flashinfer case the guard exists for — latent in the previous revision here too), and Signature.bind()-check the builder so a positional-only refactor degrades to a no-op instead of failing the wrapper's keyword call. Guardrails as before: applied only when a behavioral probe shows the hash/eq inconsistency and the builder signature is keyword-callable with the expected parameters; any flashinfer refactor degrades the workaround to a no-op. Drop once a fixed flashinfer is picked up. Validation: serving benchmark at high concurrency on a large MoE model with MLA attention recovers output throughput to the pre-regression baseline band, identical to the __eq__-based workaround it replaces (~6500 output tok/s at concurrency 1024, low concurrency par). CPU microbenchmark of the autotuner cache probe: ~2 us per call with cache size 1, versus ~4-10 ms per call with an 8000-entry collision chain. Unit-level sanity: same-key calls return the same config object, each keyed scalar forces a rebuild, DynamicTensorSpec.__eq__ untouched, install idempotent, signature drift and in-handler exceptions no-op. The workaround block is byte-identical to the NVIDIA#17165 head. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #63772 [ ] completed with state |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Port of the reworked TRTLLM-14904 workaround from #17165 (main) to feat/kimi_k3, which currently carries the earlier
DynamicTensorSpec.__eq__-replacement variant (landed via #17164).Review on #17165 (@BowenFu) flagged that the class-wide
__eq__replacement mutates a third-party class process-wide (every flashinfer consumer in the interpreter) and that its widened equality can cross-contaminateAutoTuner._override_config_cacheentries whose configs differ only intensor_initializers. This PR replaces it with the scoped variant: memoizeflashinfer.mla._core._build_mla_decode_tuning_config— the one function that rebuilds aTuningConfigwith fresh initializer closures per decode call, and the only per-callTuningConfigconstruction site in flashinfer 0.6.15 — keyed on the scalars that fully determine its output. Repeated calls reuse the same config object, so autotuner cache hits resolve through flashinfer's own unmodified hash/eq, fixing the collision-chain growth at its root with no equality-semantics change anywhere.Also folds in the two CodeRabbit fixes from #17165: drop the
exc_infokwarg (tensorrt_llm'slogger.debug(*msg)takes no kwargs, so the except handler would itself raise at import time exactly when flashinfer is incompatible — latent in the__eq__variant on this branch too), and aSignature.bind()check so a positional-only builder refactor degrades to a no-op instead of failing the wrapper's keyword call at the first decode.Guardrails: the workaround installs only when a behavioral probe shows the hash/eq inconsistency and the builder is keyword-callable with the expected parameters; any flashinfer refactor degrades it to a no-op. Drop once a fixed flashinfer release is picked up (upstream issue being filed, proposing the builder be cached upstream as flashinfer's sparse-MLA path already does).
The workaround block is byte-identical to the #17165 head (3c6987c).
Test Coverage
__eq__variant this replaces (~6500 output tok/s at high concurrency); low-concurrency throughput unchanged.TuningConfigobject; each keyed scalar forces a rebuild;DynamicTensorSpec.__eq__untouched; install idempotent; signature drift (renamed or positional-only) and in-handler exceptions degrade to a no-op.PR Checklist