[TRTLLM-14904][fix] Work around flashinfer 0.6.15 autotuner cache-key hash/eq inconsistency - #17165
Open
brnguyen2 wants to merge 1 commit into
Open
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>
Contributor
WalkthroughChangesFlashInfer compatibility
Estimated code review effort: 2 (Simple) | ~10 minutes Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py`:
- Around line 134-137: Update the exception handler surrounding the flashinfer
DynamicTensorSpec.__eq__ workaround probe and replacement to catch only expected
compatibility errors, such as ImportError, AttributeError, and TypeError,
instead of Exception. Preserve the existing debug logging and fallback behavior
while allowing unexpected implementation defects to propagate.
- Around line 98-105: Add precise type annotations to the nested function
_probe_spec, including its DynamicTensorSpec return type. Also annotate self,
other, and the return value of _hash_consistent_eq, preserving the existing
comparison behavior; apply these changes at both referenced locations in
flashinfer_trtllm_gen.py (lines 98-105 and 111-129).
🪄 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: Enterprise
Run ID: 9ba2af7e-8beb-46f1-aa0a-22306b2f779f
📒 Files selected for processing (1)
tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py
Comment on lines
+98
to
+105
| def _probe_spec(): | ||
| return DynamicTensorSpec( | ||
| input_idx=(0,), | ||
| dim_idx=(0,), | ||
| gen_tuning_buckets=(1,), | ||
| map_to_tuning_buckets=_probe_spec, # any stable callable | ||
| tensor_initializers=[lambda shapes, dtype, device: None], | ||
| ) |
Contributor
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Add annotations to the nested functions.
tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py#L98-L105: Add a precise return annotation to_probe_spec.tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py#L111-L129: Add precise annotations forself,other, and the return value of_hash_consistent_eq.
As per coding guidelines, “Annotate every function.”
📍 Affects 1 file
tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py#L98-L105(this comment)tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py#L111-L129
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py` around
lines 98 - 105, Add precise type annotations to the nested function _probe_spec,
including its DynamicTensorSpec return type. Also annotate self, other, and the
return value of _hash_consistent_eq, preserving the existing comparison
behavior; apply these changes at both referenced locations in
flashinfer_trtllm_gen.py (lines 98-105 and 111-129).
Source: Coding guidelines
Comment on lines
+134
to
+137
| except Exception: | ||
| # A future flashinfer refactor (moved class, changed fields) must not | ||
| # break attention; it just loses the workaround. | ||
| logger.debug("Skipping flashinfer DynamicTensorSpec.__eq__ workaround.", exc_info=True) |
Contributor
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Catch only expected compatibility errors.
Line 134 hides defects in the probe and replacement implementation. Catch only expected import and API-shape errors, such as ImportError, AttributeError, and TypeError. This keeps the compatibility fallback while exposing unexpected defects.
As per coding guidelines, “Catch specific exceptions instead of using broad or bare exception handling such as except:.”
🧰 Tools
🪛 Ruff (0.16.0)
[warning] 134-134: Do not catch blind exception: Exception
(BLE001)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py` around
lines 134 - 137, Update the exception handler surrounding the flashinfer
DynamicTensorSpec.__eq__ workaround probe and replacement to catch only expected
compatibility errors, such as ImportError, AttributeError, and TypeError,
instead of Exception. Preserve the existing debug logging and fallback behavior
while allowing unexpected implementation defects to propagate.
Sources: Coding guidelines, Linters/SAST tools
brnguyen2
marked this pull request as ready for review
August 1, 2026 21:10
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
1 participant
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.
Same change as #17164 (feat/kimi_k3, where the regression was measured), cherry-picked onto main, which carries the same flashinfer 0.6.15 pin.
Test Coverage
PR Checklist
Dev Engineer Review
DynamicTensorSpechash/equality inconsistency.tensor_initializers, consistent with__hash__.QA Engineer Review
No test changes.