[TRTLLM-14904][fix] Work around flashinfer 0.6.15 autotuner cache-key hash/eq inconsistency - #17165
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>
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
WalkthroughChangesFlashInfer MLA decode cache
Estimated code review effort: 3 (Moderate) | ~20 minutes Possibly related PRs
Suggested reviewers: Sequence Diagram(s)sequenceDiagram
participant MLARequest
participant CachedBuilder
participant OriginalBuilder
MLARequest->>CachedBuilder: Provide scalar geometry and device inputs
CachedBuilder->>CachedBuilder: Check bounded cache
CachedBuilder->>OriginalBuilder: Build configuration on cache miss
OriginalBuilder-->>CachedBuilder: Return tuning configuration
CachedBuilder-->>MLARequest: Return cached or new configuration
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
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
|
/bot run |
|
PR_Github #63286 [ run ] triggered by Bot. Commit: |
|
PR_Github #63286 [ run ] completed with state
|
…lause Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
|
/bot run |
|
PR_Github #63311 [ run ] triggered by Bot. Commit: |
|
PR_Github #63311 [ run ] completed with state
|
|
/bot run |
|
PR_Github #63331 [ run ] triggered by Bot. Commit: |
|
Held rather than approved. One finding on the replacement
Scope. The probe establishes that the upstream bug is present, not that the replacement is correct, and nothing here tests the replacement |
|
PR_Github #63331 [ run ] completed with state
|
…config builder Review flagged two problems with the DynamicTensorSpec.__eq__ replacement: it mutates a third-party class process-wide (every flashinfer consumer in the interpreter, not just the MLA decode path the WAR targets), and the widened equality lets AutoTuner's _override_config_cache return an override built from a different config's tensor_initializers when two configs differ only in their initializer closures. 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 (kv-cache paging geometry, workspace size, runner set, head/rank dims, max_seq_len, device). Repeated calls now reuse the same TuningConfig object, so autotuner cache hits resolve through flashinfer's own unmodified hash/eq — no equality semantics change anywhere, and the collision-chain growth is fixed at its root (per-call closure churn). Every other TuningConfig construction site in flashinfer 0.6.15 already reuses its config across calls (module-level constants, instance-level configs, functools.cache-wrapped builders), so memoizing this one builder covers the entire pathology. Guardrails as before: applied only when a behavioral probe shows the hash/eq inconsistency and the builder still has the expected signature; 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 previous workaround (~6500 output tok/s at concurrency 1024); low-concurrency throughput unchanged. 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 unpatched. Also verified: same-key calls return the same config object, each keyed scalar forces a rebuild, DynamicTensorSpec.__eq__ is untouched, install is idempotent, and signature drift no-ops. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
|
@BowenFu Both points were right, and the The workaround no longer touches On scope: this is bounded to the MLA decode path by construction. I checked the other Re-validated on the new revision: the high-concurrency serving benchmark recovers to the same throughput as the previous revision (par with the pre-regression baseline band), low concurrency unchanged. CPU microbenchmark of the cache probe: ~2 µs per call with cache size 1, versus ~4-10 ms per call with an 8000-entry collision chain unpatched. Also verified at unit level: same-key calls return the same object, each keyed scalar forces a rebuild, The upstream flashinfer issue will propose the same fix (cache the builder, as the sparse-MLA path already does). |
|
/bot run |
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 197-199: The logger.debug call with the "Skipping flashinfer MLA
decode tuning-config cache workaround." message is passing an unsupported
exc_info keyword argument. Remove the exc_info=True parameter from this
logger.debug call since tensorrt_llm.logger.debug only accepts positional
message arguments. Keep the debug message content unchanged.
- Around line 129-141: Update the signature guard around orig_build before
installing the wrapper to validate keyword-call compatibility, not just
parameter names and order. After confirming the expected parameters in
inspect.signature(orig_build), use Signature.bind with representative keyword
arguments for every wrapped call parameter; return without installing the
wrapper if binding raises TypeError. Keep the existing guard behavior for
refactored signatures.
🪄 Autofix
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: 5ed68bcb-ffae-41d2-88ce-395dd4f3d17d
📒 Files selected for processing (1)
tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py
|
PR_Github #63759 [ run ] triggered by Bot. Commit: |
…p unsupported exc_info Two CodeRabbit findings on the memoized-builder revision: - tensorrt_llm's logger.debug(*msg) takes no kwargs, so the except handler's exc_info=True would itself raise TypeError at import time — breaking the attention backend in exactly the incompatible-flashinfer case the guard exists for. Drop it. - The signature guard compared parameter names/order only; a refactor keeping the names but making them positional-only would pass the guard and then fail the wrapper's keyword call at the first decode. Add a Signature.bind() check so that case degrades to a no-op too. No hot-path change; CPU sanity extended to cover both paths (strict kwargs-rejecting logger stub, positional-only drift, exception path). Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
|
/bot run |
|
PR_Github #63770 [ run ] triggered by Bot. Commit: |
|
PR_Github #63759 [ run ] completed with state |
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.trtllm_batch_decode_with_kv_cache_mlabuilds a freshTuningConfigper call (via_build_mla_decode_tuning_config) 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:
flashinfer.mla._core._build_mla_decode_tuning_configis memoized on the scalars that fully determine its output (kv-cache paging geometry, workspace size, runner set, head/rank dims, max_seq_len, device). Repeated decode calls then reuse the sameTuningConfigobject, so autotuner cache hits resolve through flashinfer's own unmodified hash/eq — the collision-chain growth is fixed at its root (per-call closure churn) with no equality-semantics change anywhere. The scope is exactly the MLA decode path: every otherTuningConfigconstruction site in flashinfer 0.6.15 already reuses its config across calls (module-level constants, instance-level configs,functools.cache-wrapped builders), so this one builder covers the entire pathology.A behavioral probe applies the workaround only when the hash/eq inconsistency is actually present and the builder still has the expected signature; any failure degrades to not patching, so a future flashinfer release with different internals is unaffected.
This replaces the earlier revision of this PR, which patched
DynamicTensorSpec.__eq__class-wide. Review (see comments) flagged that the class-wide equality change affects every flashinfer consumer in the process and can cross-contaminateAutoTuner._override_config_cacheentries whose configs differ only intensor_initializers; the memoization approach has neither problem. The__eq__variant that landed on feat/kimi_k3 via #17164 will be aligned in a follow-up.The bug is upstream in flashinfer; an upstream issue is being filed (proposed upstream fix: cache the builder, as flashinfer's sparse-MLA path already does), and this workaround should be dropped once a fixed release is picked up.
Test Coverage
__eq__-patch revision of this PR); low-concurrency throughput unchanged.TuningConfigobject; each keyed scalar (paging geometry, workspace size, runners, q_len, heads, rank, max_seq_len) forces a rebuild;DynamicTensorSpec.__eq__is untouched; install is idempotent; builder signature drift degrades to a no-op.PR Checklist
Dev Engineer Review
flashinfer.mla._core._build_mla_decode_tuning_configusing the scalar inputs that determine eachTuningConfig.DynamicTensorSpechash and equality behavior.Signature.bind()checks before installation.exc_infologging arguments from compatibility and exception paths.QA Engineer Review
No test changes.