[TRTLLM-14904][fix] Work around flashinfer 0.6.15 autotuner cache-key hash/eq inconsistency - #17164
Merged
brnguyen2 merged 1 commit intoAug 1, 2026
Conversation
… hash/eq inconsistency flashinfer 0.6.15's DynamicTensorSpec defines a custom __hash__ that skips tensor_initializers (and hashes callables by identity) while keeping the dataclass-generated __eq__, which compares all fields by value. Callers such as trtllm_batch_decode_with_kv_cache_mla build a fresh TuningConfig per call with fresh initializer closures, so the lru_cache on AutoTuner._find_nearest_profile accumulates hash-equal but eq-unequal keys up to its 16384 cap. Every autotuner cache probe on the MLA decode path then walks the whole collision chain in Python __eq__ — a 17-19 ms host stall per eager MLA generation call, which cost ~15% output throughput in a serving benchmark at high concurrency on a large MoE model (mixed prefill+decode iterations run eagerly; decode-only iterations replay CUDA graphs and are unaffected). 0.6.14 built a fresh bucket mapper per call, so the keys hashed differently and missed in O(1). Install a guarded TRT-LLM-side workaround where the fmha flashinfer backend imports flashinfer: replace DynamicTensorSpec.__eq__ with one consistent with its __hash__ (ignore tensor_initializers, compare callables by identity). A behavioral probe applies the patch only when the inconsistency is present, and any failure degrades to not patching, so a future flashinfer release with different internals is unaffected. The bug is upstream in flashinfer; this workaround should be dropped once a fixed version is picked up. Validation: the serving benchmark recovered from ~5200 to ~6500 output tok/s at concurrency 1024 (par with the pre-regression baseline band), with low-concurrency throughput unchanged. A CPU microbenchmark of the cache probe stays at ~2 us per call and cache size 1, versus ~10 ms per call with an 8000-entry collision chain unpatched. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
Merged
3 tasks
3 tasks
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
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
flashinfer 0.6.15's
DynamicTensorSpecdefines a custom__hash__that skipstensor_initializers(and hashes callables by identity) while keeping the dataclass-generated__eq__, which compares all fields by value. Callers such astrtllm_batch_decode_with_kv_cache_mlabuild a freshTuningConfigper call with fresh initializer closures, so thelru_cacheonAutoTuner._find_nearest_profileaccumulates hash-equal but eq-unequal keys up to its 16384 cap. Every autotuner cache probe on the MLA decode path then walks the whole collision chain in Python__eq__— a 17–19 ms host stall per eager MLA generation call. Mixed prefill+decode iterations run eagerly and pay this on every call; decode-only iterations replay CUDA graphs and are unaffected, which made it present as a prefill-side regression. flashinfer 0.6.14 built a fresh bucket mapper per call, so keys hashed differently and cache probes missed in O(1); the 0.6.14→0.6.15 bump (#16530) exposed the inconsistency.This PR installs a guarded TRT-LLM-side workaround where the fmha flashinfer backend imports flashinfer:
DynamicTensorSpec.__eq__is replaced with one consistent with its existing__hash__(ignoretensor_initializers, compare callables by identity). A behavioral probe applies the patch only when the inconsistency is actually present, and any failure degrades to not patching, so a future flashinfer release with different internals is unaffected.The bug is upstream in flashinfer; an upstream issue is being filed, and this workaround should be dropped once a fixed release is picked up.
Test Coverage
PR Checklist