Skip to content

[TRTLLM-14904][fix] Work around flashinfer 0.6.15 autotuner cache-key hash/eq inconsistency - #17164

Merged
brnguyen2 merged 1 commit into
NVIDIA:feat/kimi_k3from
brnguyen2:fix/TRTLLM-14904-flashinfer-autotuner-key
Aug 1, 2026
Merged

[TRTLLM-14904][fix] Work around flashinfer 0.6.15 autotuner cache-key hash/eq inconsistency#17164
brnguyen2 merged 1 commit into
NVIDIA:feat/kimi_k3from
brnguyen2:fix/TRTLLM-14904-flashinfer-autotuner-key

Conversation

@brnguyen2

Copy link
Copy Markdown
Collaborator

Description

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. 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__ (ignore tensor_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

  • Serving benchmark at high concurrency on a large MoE model with MLA attention: output throughput recovered ~15% (back to par with the pre-regression baseline band); low-concurrency throughput unchanged.
  • CPU microbenchmark of the autotuner cache probe: ~2 µs per call with cache size 1 post-patch, versus ~10 ms per call with an 8000-entry collision chain unpatched.

PR Checklist

  • PR title and description are self-explanatory
  • Commit message includes ticket ID and is signed off (DCO)
  • Change is guarded and degrades safely if flashinfer internals change

… 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>
@brnguyen2
brnguyen2 requested a review from a team as a code owner August 1, 2026 20:05
@brnguyen2
brnguyen2 merged commit da96d02 into NVIDIA:feat/kimi_k3 Aug 1, 2026
8 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant