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

tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py

@@ -67,6 +67,146 @@
     )
 
 
+def _install_flashinfer_mla_decode_tuning_config_cache() -> None:
+    """Work around an autotuner cache pathology on flashinfer 0.6.15's MLA decode path.
+
+    ``flashinfer.mla._core._build_mla_decode_tuning_config`` builds a fresh
+    ``TuningConfig`` — with fresh initializer closures — on every
+    ``trtllm_batch_decode_with_kv_cache_mla`` call. ``DynamicTensorSpec``'s
+    custom ``__hash__`` skips ``tensor_initializers``, but its
+    dataclass-generated ``__eq__`` compares them by value, so each call's
+    config is hash-equal but eq-unequal to all previous ones: the
+    ``lru_cache`` on ``AutoTuner._find_nearest_profile`` accumulates
+    colliding keys up to the cache cap, and every autotuner cache probe on
+    the MLA decode path then walks the entire collision chain in Python
+    ``__eq__`` (tens of milliseconds per attention call, per layer).
+
+    The workaround memoizes the config builder on the scalars that fully
+    determine its output (kv-cache paging geometry, workspace size, runner
+    set, head/rank dims, max_seq_len, device), so repeated calls reuse the
+    same ``TuningConfig`` object. Cache hits then resolve through
+    flashinfer's own unmodified hash/eq — no equality semantics change
+    anywhere, and the scope is exactly the MLA decode path that exhibits the
+    problem.
+
+    MLA decode is the only path that needs this: every other ``TuningConfig``
+    construction site in flashinfer 0.6.15 already reuses its config across
+    calls (module-level constants in ``gemm``, instance-level configs in the
+    cute-dsl MoE tuner, ``functools.cache``-wrapped builders in
+    ``mla/_sparse_mla_sm120.py``), so those specs hit the autotuner caches by
+    object identity and never accumulate collision chains. Memoizing the one
+    per-call builder therefore covers the whole pathology; extending it to
+    other ops would add nothing.
+
+    The bug is upstream in flashinfer; the patch is applied only when a
+    behavioral probe shows the hash/eq inconsistency and the builder still
+    has the expected signature, so a fixed or refactored flashinfer release
+    degrades this to a no-op.
+    """
+    try:
+        import inspect
+
+        from flashinfer.autotuner.autotuner import DynamicTensorSpec
+        from flashinfer.mla import _core as _flashinfer_mla_core
+
+        def _probe_spec() -> DynamicTensorSpec:
+            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],
+            )
+
+        a, b = _probe_spec(), _probe_spec()
+        if not (hash(a) == hash(b) and a != b):
+            return  # flashinfer without the inconsistency; nothing to do
+
+        orig_build = _flashinfer_mla_core._build_mla_decode_tuning_config
+        if getattr(orig_build, "_trtllm_mla_tuning_config_cache", False):
+            return  # already installed
+
+        expected_params = (
+            "kv_cache",
+            "block_tables",
+            "workspace_buffer",
+            "runner_names",
+            "q_len",
+            "num_heads",
+            "kv_lora_rank",
+            "max_seq_len",
+            "device",
+        )
+        signature = inspect.signature(orig_build)
+        if tuple(signature.parameters) != expected_params:
+            return  # builder was refactored; memo key may no longer be complete
+        try:
+            signature.bind(**dict.fromkeys(expected_params))
+        except TypeError:
+            return  # parameters went positional-only; the keyword call below would fail
+
+        cache: dict = {}
+
+        def _cached_build(
+            kv_cache,
+            block_tables,
+            workspace_buffer,
+            runner_names,
+            q_len,
+            num_heads,
+            kv_lora_rank,
+            max_seq_len,
+            device,
+        ):
+            # Everything the builder (and _compute_mla_decode_buckets) reads
+            # from its tensor arguments reduces to these scalars.
+            key = (
+                kv_cache.shape[0],  # num_pages, captured by init_block_tables
+                kv_cache.shape[-2],  # page_size, feeds profile_seq_len
+                block_tables.shape[-1],  # pages per sequence, feeds profile_seq_len
+                workspace_buffer.numel() * workspace_buffer.element_size(),  # cute-dsl bucket cap
+                tuple(runner_names),
+                q_len,
+                num_heads,
+                kv_lora_rank,
+                max_seq_len,
+                device,
+            )
+            config = cache.get(key)
+            if config is None:
+                if len(cache) >= 256:  # backstop; a serving process sees a handful of keys
+                    cache.clear()
+                config = orig_build(
+                    kv_cache=kv_cache,
+                    block_tables=block_tables,
+                    workspace_buffer=workspace_buffer,
+                    runner_names=runner_names,
+                    q_len=q_len,
+                    num_heads=num_heads,
+                    kv_lora_rank=kv_lora_rank,
+                    max_seq_len=max_seq_len,
+                    device=device,
+                )
+                cache[key] = config
+            return config
+
+        _cached_build._trtllm_mla_tuning_config_cache = True
+        _flashinfer_mla_core._build_mla_decode_tuning_config = _cached_build
+        logger.debug(
+            "Memoized flashinfer _build_mla_decode_tuning_config "
+            "(autotuner cache collision-chain workaround)."
+        )
+    except (ImportError, AttributeError, TypeError, ValueError):
+        # A future flashinfer refactor (moved module, changed fields) must not
+        # break attention; it just loses the workaround.
+        # tensorrt_llm's logger.debug(*msg) takes no exc_info/kwargs.
+        logger.debug("Skipping flashinfer MLA decode tuning-config cache workaround.")
+
+
+if IS_FLASHINFER_AVAILABLE:
+    _install_flashinfer_mla_decode_tuning_config_cache()
+
+
 _MULTI_CTAS_KV_COUNTER_ALIGNMENT = 8