flashinfer-python 0.6.15__tar.gz → 0.6.15.post1__tar.gz

{flashinfer_python-0.6.15/flashinfer_python.egg-info → flashinfer_python-0.6.15.post1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: flashinfer-python
-Version: 0.6.15
+Version: 0.6.15.post1
 Summary: FlashInfer: Kernel Library for LLM Serving
 Author: FlashInfer team
 License-Expression: Apache-2.0

{flashinfer_python-0.6.15 → flashinfer_python-0.6.15.post1}/csrc/trtllm_fused_moe_kernel_launcher.cu

@@ -20,7 +20,10 @@
 #include <cstring>
 #include <iomanip>
 #include <iostream>
+#include <map>
+#include <memory>
 #include <set>
+#include <tuple>
 #include <unordered_map>
 #include <vector>
 
@@ -514,7 +517,7 @@ class FusedMoeLauncher {
   Tensor workspace_fc2;
   Tensor output;
   int64_t moe_tactic{-1};
-  std::unique_ptr<tensorrt_llm::kernels::trtllmgen_moe::MoE::Runner> moe_runner;
+  std::shared_ptr<tensorrt_llm::kernels::trtllmgen_moe::MoE::Runner> moe_runner;
 
   void prepare_moe_common(int64_t& moe_tactic) {
     using RunnerType = tensorrt_llm::kernels::trtllmgen_moe::MoE::Runner;
@@ -530,18 +533,53 @@ class FusedMoeLauncher {
     // gemm1 bias, use the weights-only Runner constructor to match the original kernel
     // path and numerics. DSFp8 + biasMn routes through the unified constructor below
     // (which accepts gemm1_bias_type).
-    if (this->mDtypeAct == btg::Dtype::E4m3 && this->mDtypeWeights == btg::Dtype::E4m3 &&
-        args->mUseDeepSeekFp8 && args->gemm1_bias_type == batchedGemm::gemm::BiasType::None) {
-      moe_runner = std::make_unique<RunnerType>(this->mDtypeWeights, args->mUseDeepSeekFp8,
-                                                (int32_t)tile_tokens_dim, this->use_shuffled_weight,
-                                                this->weight_layout, usePerTokenScalingGemm1,
-                                                usePerTokenScalingGemm2, false, false);
-    } else {
-      moe_runner = std::make_unique<RunnerType>(
-          this->mDtypeAct, this->mDtypeWeights, args->mUseDeepSeekFp8, (int32_t)tile_tokens_dim,
-          this->activation_type, this->use_shuffled_weight, this->weight_layout,
-          args->gemm1_bias_type, usePerTokenScalingGemm1, usePerTokenScalingGemm2);
+    bool const useWeightsOnlyConstructor =
+        this->mDtypeAct == btg::Dtype::E4m3 && this->mDtypeWeights == btg::Dtype::E4m3 &&
+        args->mUseDeepSeekFp8 && args->gemm1_bias_type == batchedGemm::gemm::BiasType::None;
+    bool constexpr usePerChannelScalingGemm1 = false;
+    bool constexpr usePerChannelScalingGemm2 = false;
+
+    // A Runner contains only constructor-derived kernel metadata and config indices. Reuse it on
+    // the same host thread instead of rebuilding and filtering the global config table per call.
+    using RunnerCacheKey =
+        std::tuple<int64_t, int64_t, bool, int32_t, int64_t, bool, int64_t, int64_t, bool, bool,
+                   bool, bool, bool, int32_t, int32_t, int32_t>;
+    static thread_local std::map<RunnerCacheKey, std::shared_ptr<RunnerType>> runnerCache;
+    RunnerCacheKey const runnerKey{static_cast<int64_t>(this->mDtypeAct),
+                                   static_cast<int64_t>(this->mDtypeWeights),
+                                   args->mUseDeepSeekFp8,
+                                   static_cast<int32_t>(tile_tokens_dim),
+                                   static_cast<int64_t>(this->activation_type),
+                                   this->use_shuffled_weight,
+                                   static_cast<int64_t>(this->weight_layout),
+                                   static_cast<int64_t>(args->gemm1_bias_type),
+                                   usePerTokenScalingGemm1,
+                                   usePerTokenScalingGemm2,
+                                   usePerChannelScalingGemm1,
+                                   usePerChannelScalingGemm2,
+                                   useWeightsOnlyConstructor,
+                                   std::get<0>(device_version),
+                                   std::get<1>(device_version),
+                                   hidden_states.device().device_id};
+
+    auto runnerIt = runnerCache.find(runnerKey);
+    if (runnerIt == runnerCache.end()) {
+      std::shared_ptr<RunnerType> runner;
+      if (useWeightsOnlyConstructor) {
+        runner = std::make_shared<RunnerType>(
+            this->mDtypeWeights, args->mUseDeepSeekFp8, static_cast<int32_t>(tile_tokens_dim),
+            this->use_shuffled_weight, this->weight_layout, usePerTokenScalingGemm1,
+            usePerTokenScalingGemm2, usePerChannelScalingGemm1, usePerChannelScalingGemm2);
+      } else {
+        runner = std::make_shared<RunnerType>(
+            this->mDtypeAct, this->mDtypeWeights, args->mUseDeepSeekFp8,
+            static_cast<int32_t>(tile_tokens_dim), this->activation_type, this->use_shuffled_weight,
+            this->weight_layout, args->gemm1_bias_type, usePerTokenScalingGemm1,
+            usePerTokenScalingGemm2, usePerChannelScalingGemm1, usePerChannelScalingGemm2);
+      }
+      runnerIt = runnerCache.emplace(runnerKey, std::move(runner)).first;
     }
+    moe_runner = runnerIt->second;
 
     int32_t const effectiveTopK = args->top_k + args->num_fused_shared_experts;
     int32_t const effectiveLocalExperts = args->local_num_experts + args->num_fused_shared_experts;
@@ -551,13 +589,10 @@ class FusedMoeLauncher {
                                                           args->intermediate_size,
                                                           effectiveLocalExperts, args->num_tokens);
     }
-    auto valid_cfgs =
-        moe_runner->getValidConfigIndices(effectiveTopK, args->hidden_size, args->intermediate_size,
-                                          effectiveLocalExperts, args->num_tokens);
-    auto valid_it = std::find(valid_cfgs.begin(), valid_cfgs.end(), moe_tactic);
-    FLASHINFER_CHECK(valid_it != valid_cfgs.end(), "Invalid MoE tactic ", moe_tactic,
-                     " for tile_N=", tile_tokens_dim, ". Number of valid tactics for this tile is ",
-                     valid_cfgs.size(),
+    FLASHINFER_CHECK(moe_runner->isValidConfigIndex(moe_tactic, effectiveTopK, args->hidden_size,
+                                                    args->intermediate_size, effectiveLocalExperts,
+                                                    args->num_tokens),
+                     "Invalid MoE tactic ", moe_tactic, " for tile_N=", tile_tokens_dim,
                      ". This often indicates a stale or mismatched autotuner cache entry.");
     this->moe_tactic = moe_tactic;
 
@@ -596,6 +631,7 @@ class FusedMoeLauncher {
   virtual Array<Tensor> run(int64_t moe_tactic, bool enable_pdl = true,
                             bool use_routing_scales_on_input = false,
                             bool use_deep_seek_fp8 = false, bool return_activation_output = false) {
+    ffi::CUDADeviceGuard device_guard(hidden_states.device().device_id);
     check_routing();
     prepare_routing();
 
@@ -1429,6 +1465,7 @@ class Fp8BlockScaleLauncher : public FusedMoeLauncher {
   Array<Tensor> run(int64_t moe_tactic, bool enable_pdl = true,
                     bool use_routing_scales_on_input = false, bool use_deep_seek_fp8 = false,
                     bool return_activation_output = false) override {
+    ffi::CUDADeviceGuard device_guard(hidden_states.device().device_id);
     check_routing();
     prepare_routing();
 
@@ -1988,6 +2025,7 @@ class FP4BlockScaleLauncher : public FusedMoeLauncher {
   Array<Tensor> run(int64_t moe_tactic, bool enable_pdl = true,
                     bool use_routing_scales_on_input = false, bool use_deep_seek_fp8 = false,
                     bool return_activation_output = false) override {
+    ffi::CUDADeviceGuard device_guard(hidden_states.device().device_id);
     TVM_FFI_ICHECK(!return_activation_output)
         << "return_activation_output is not supported for FP4 block-scale MoE";
     check_routing();

{flashinfer_python-0.6.15 → flashinfer_python-0.6.15.post1}/csrc/trtllm_fused_moe_runner.cu

@@ -773,6 +773,21 @@ std::vector<int64_t> Runner::getValidConfigIndices(int32_t topK, int32_t hiddenS
   return validIndices;
 }
 
+bool Runner::isValidConfigIndex(int64_t configIndex, int32_t topK, int32_t hiddenSize,
+                                int32_t intermediateSize, int32_t numLocalExperts,
+                                int32_t numTokens) const {
+  if (configIndex < 0 || configIndex >= static_cast<int64_t>(mPassingConfigs.size())) {
+    return false;
+  }
+
+  auto const& config = mPassingConfigs[configIndex];
+  return mPermuteGemm1.isValidConfigIndex(static_cast<int32_t>(config.gemm1Config), topK,
+                                          hiddenSize, intermediateSize, numLocalExperts,
+                                          numTokens) &&
+         mGemm2.isValidConfigIndex(static_cast<int32_t>(config.gemm2Config), topK, hiddenSize,
+                                   intermediateSize, numLocalExperts, numTokens);
+}
+
 int64_t Runner::getDefaultValidConfigIndex(int32_t topK, int32_t hiddenSize,
                                            int32_t intermediateSize, int32_t numLocalExperts,
                                            int32_t numTokens) const {

flashinfer_python-0.6.15.post1/flashinfer/_build_meta.py

@@ -0,0 +1,3 @@
+"""Build metadata for flashinfer package."""
+__version__ = "0.6.15.post1"
+__git_version__ = "de358e829443d2cab2261b4c2b403444f6925ab6"

{flashinfer_python-0.6.15 → flashinfer_python-0.6.15.post1}/flashinfer/autotuner/autotuner.py

@@ -174,6 +174,13 @@ def make_bucket_mapper(
 
 _METADATA_KEY = "_metadata"
 
+# Metadata values a writer records when it cannot determine a field.
+# ("None" covers files written by versions that stringified a null version
+# report.)  A cache carrying these values is ignored on load, but — unlike a
+# genuine environment mismatch — may be overwritten on save, so a poisoned
+# file heals on the next tuned run instead of forcing a re-tune forever.
+_INDETERMINATE_META_VALUES = (None, "unknown", "None")
+
 
 def _get_cublas_version() -> str:
     """Return the cuBLAS version as ``major.minor.patch``.
@@ -256,6 +263,32 @@ def _get_cublas_version() -> str:
     return "unknown"
 
 
+def _get_cudnn_backend_version() -> str:
+    """Return the cuDNN backend version as a string, or ``"unknown"``.
+
+    ``torch.backends.cudnn.version()`` returns ``None`` (or raises) when
+    torch has not loaded cuDNN in this process — e.g. before torch's lazy
+    cuDNN initialization, or in a CPU-only context.  Fall back to the
+    cudnn-frontend's ``backend_version()``, which queries libcudnn directly
+    without torch's lazy init (the same gate ``gemm_base`` uses).
+    """
+    try:
+        version = torch.backends.cudnn.version()
+        if version:
+            return str(version)
+    except Exception:
+        pass
+    try:
+        import cudnn as _cudnn_frontend
+
+        version = _cudnn_frontend.backend_version()
+        if version:
+            return str(version)
+    except Exception:
+        pass
+    return "unknown"
+
+
 def _collect_metadata() -> dict[str, str]:
     """Collect environment metadata that can affect tactic-to-kernel mappings.
 
@@ -283,10 +316,7 @@ def _collect_metadata() -> dict[str, str]:
     meta["flashinfer_version"] = _flashinfer_version
     meta["cuda_version"] = getattr(torch.version, "cuda", None) or "unknown"
     meta["cublas_version"] = _get_cublas_version()
-    try:
-        meta["cudnn_version"] = str(torch.backends.cudnn.version())
-    except Exception:
-        meta["cudnn_version"] = "unknown"
+    meta["cudnn_version"] = _get_cudnn_backend_version()
     try:
         # cudnn-frontend is an optional dependency; failing to import is
         # not fatal -- we just record "unknown" and let the metadata
@@ -305,6 +335,45 @@ def _collect_metadata() -> dict[str, str]:
     return meta
 
 
+def _classify_metadata_mismatches(
+    saved_meta: dict[str, Any], current_meta: dict[str, str]
+) -> tuple[dict[str, tuple[Any, str]], dict[str, tuple[Any, str]]]:
+    """Split saved-vs-current metadata mismatches by severity.
+
+    Returns ``(hard, soft)``, each mapping ``key -> (saved, current)``:
+
+      * ``hard`` — the writer recorded a definite value that differs from
+        the current environment: a genuine environment conflict.  Neither
+        reading the file's configs nor overwriting the file is safe.
+      * ``soft`` — the writer recorded an indeterminate value (or the key
+        is missing, for files from older flashinfer versions): the configs
+        cannot be trusted, but replacing the file with freshly tuned
+        configs and the current environment's metadata is safe.
+
+    A saved value of ``"*"`` is a wildcard and never mismatches.  Keys in
+    ``saved_meta`` that ``current_meta`` does not track are ignored.  A
+    corrupt (non-dict) ``saved_meta`` is treated as all-keys-missing: the
+    file is unattributable, i.e. every mismatch is soft.
+    """
+    if not isinstance(saved_meta, dict):
+        saved_meta = {}
+    hard: dict[str, tuple[Any, str]] = {}
+    soft: dict[str, tuple[Any, str]] = {}
+    for key, current in current_meta.items():
+        saved = saved_meta.get(key)
+        if saved in (current, "*"):
+            continue
+        target = soft if saved in _INDETERMINATE_META_VALUES else hard
+        target[key] = (saved, current)
+    return hard, soft
+
+
+def _format_meta_mismatches(mismatches: dict[str, tuple[Any, str]]) -> str:
+    return ", ".join(
+        f"{k}: saved={old} vs current={new}" for k, (old, new) in mismatches.items()
+    )
+
+
 def get_config_path(is_module: bool):
     """Return the module name or file path for bundled per-GPU tuning configs."""
     dev_name = torch.cuda.get_device_name(0).replace(" ", "_")
@@ -707,17 +776,16 @@ def autotune(
             "pass None (or omit) to inherit the current buckets"
         )
 
-    # Load configs from cache file on entry (if it exists).
-    # cache_valid is False when the file exists but has a metadata mismatch;
-    # in that case we skip saving on exit to avoid overwriting configs from
-    # a different environment.
-    cache_valid = True
+    # Load configs from cache file on entry (if it exists).  A file with
+    # mismatched metadata is ignored here; whether it may be overwritten on
+    # exit is decided by save_configs at save time (it re-reads the file,
+    # so a heal or replacement by another process is re-evaluated).
     if cache is not None:
         with tuner._lock:
             tuner._file_configs.clear()
             tuner._logged_file_hits.clear()
         if os.path.isfile(cache):
-            cache_valid = tuner.load_configs(cache)
+            tuner.load_configs(cache)
 
     # Push skip_ops onto per-thread stack.  Each entry is the cumulative
     # union so that _effective_skip_ops is an O(1) read from the top.
@@ -778,10 +846,12 @@ def autotune(
         if autotune_enabled:
             logger.info("[Autotuner]: Autotuning process ends")
 
-        # Save configs on exit when tuning with a cache path,
-        # but only if new profiling results were added this session
-        # and the cache file was valid (no environment mismatch).
-        if cache is not None and cache_valid and tune_mode and tuner._dirty:
+        # Save configs on exit when tuning with a cache path, but only if
+        # new profiling results were added this session.  save_configs
+        # itself refuses to overwrite a file that belongs to a genuinely
+        # different environment, and replaces one whose recorded
+        # environment is indeterminate.
+        if cache is not None and tune_mode and tuner._dirty:
             tuner.save_configs(cache)
 
 
@@ -1005,6 +1075,10 @@ class AutoTuner:
         self._file_configs: dict[str, tuple[str, Any]] = {}
         # Track which file config keys have been logged (to avoid per-call spam)
         self._logged_file_hits: set[tuple[str, str]] = set()
+        # Cache-metadata mismatch severities ("hard"/"soft") already logged at
+        # WARNING; repeats for further cache files are logged at DEBUG (a
+        # deployment can touch dozens of cache files per process).
+        self._warned_cache_mismatch: set[str] = set()
         # Track which (custom_op, profile-shape signature) pairs have already
         # produced an out-of-range cache-miss warning, so we warn at most
         # once per unique missing shape.
@@ -1884,6 +1958,21 @@ class AutoTuner:
         results taking priority for overlapping keys). This ensures the saved
         file is always a complete, self-contained config.
 
+        If a file already exists at ``path``, its ``_metadata`` decides how
+        the save proceeds:
+
+            * matches the current environment (or the file predates
+              metadata) — on-disk entries are merged with this process's
+              entries and the original "created by" record is preserved;
+            * indeterminate (``"unknown"`` values, or keys missing from
+              files written by older flashinfer versions) — the on-disk
+              entries were tuned in an environment the writer could not
+              identify, so the file is **replaced** with this process's
+              configs stamped with the current metadata (healing the file);
+            * definite conflict — the file belongs to a different
+              environment sharing this path; the save is skipped with a
+              warning so that environment's configs survive.
+
         Note:
             This is called automatically on exit from
             ``with autotune(True, cache=path):``. Direct calls are only needed
@@ -1929,16 +2018,51 @@ class AutoTuner:
 
         # Re-read the file from disk and merge to reduce lost updates when
         # multiple processes save to the same path.  Entries from this
-        # process take priority over on-disk entries.
+        # process take priority over on-disk entries.  The decision is made
+        # here, at save time, so that a file healed or replaced by another
+        # process since our load_configs() call is re-evaluated against its
+        # current on-disk metadata.
         abs_path = os.path.abspath(path)
         original_metadata = None
         try:
             with open(abs_path, "r") as f:
                 disk_configs = json.load(f)
-            # Preserve the original _metadata from disk (the "created by" record).
-            original_metadata = disk_configs.pop(_METADATA_KEY, None)
-            disk_configs.update(configs)
-            configs = disk_configs
+            if not isinstance(disk_configs, dict):
+                disk_configs = {}
+            disk_meta = disk_configs.pop(_METADATA_KEY, None)
+            hard, soft = (
+                _classify_metadata_mismatches(disk_meta, current_meta)
+                if disk_meta is not None
+                else ({}, {})
+            )
+            if hard:
+                # Genuine environment conflict: leave the other
+                # environment's file untouched (mirrors load_configs).
+                message = (
+                    f"[Autotuner]: Not saving configs to {path}: the file "
+                    f"was created in a different environment "
+                    f"({_format_meta_mismatches(hard)}). Use a different "
+                    f"cache path for the current environment, or delete "
+                    f"the file to re-prime it."
+                )
+                with self._lock:
+                    if "hard" in self._warned_cache_mismatch:
+                        logger.debug(message)
+                    else:
+                        self._warned_cache_mismatch.add("hard")
+                        logger.warning(message)
+                return
+            if not soft:
+                # Same environment: merge on-disk entries and preserve the
+                # original "created by" record.
+                disk_configs.update(configs)
+                configs = disk_configs
+                original_metadata = disk_meta
+            # else: the on-disk entries were tuned in an environment the
+            # writer could not identify ("unknown") — drop them and replace
+            # the file, stamping the current metadata.  This heals caches
+            # poisoned by e.g. cudnn_version="unknown", which would
+            # otherwise be ignored by every future load_configs() forever.
         except (FileNotFoundError, json.JSONDecodeError):
             pass  # file doesn't exist yet or is being replaced -- proceed with what we have
 
@@ -2000,10 +2124,21 @@ class AutoTuner:
               ``policy=ALL`` plan_index ordering independent of backend)
             * ``gpu`` (different SM may expose different engines)
 
-        Caches saved by older flashinfer versions that predate the
-        ``cudnn_frontend_version`` metadata field will fail this check
-        (since saved metadata does not contain the field) and need to
-        be regenerated.  See :func:`_collect_metadata` for the exact list.
+        A skipped cache is handled in one of two ways by the next
+        :meth:`save_configs` on the same path:
+
+            * the saved metadata records a **definite** conflicting value —
+              a genuine environment conflict; the file is never
+              overwritten.  Use a different cache path, or delete the file
+              to re-prime it for this environment.
+            * the saved metadata is **indeterminate** — ``"unknown"``
+              values (the writer could not detect e.g. its cuDNN version),
+              or keys missing entirely from files written by older
+              flashinfer versions.  The file is replaced with freshly
+              tuned configs and the current environment's metadata, so the
+              cache heals instead of being re-tuned forever.
+
+        See :func:`_collect_metadata` for the exact field list.
 
         Note:
             This is called automatically on entry to
@@ -2041,24 +2176,35 @@ class AutoTuner:
         # entirely to avoid silently using invalid or suboptimal tactics.
         if saved_meta is not None:
             current_meta = _collect_metadata()
-            mismatches = {
-                k: (saved_meta.get(k), current_meta.get(k))
-                for k in current_meta
-                if saved_meta.get(k) not in (current_meta.get(k), "*")
-            }
-            if mismatches:
-                details = ", ".join(
-                    f"{k}: saved={old} vs current={new}"
-                    for k, (old, new) in mismatches.items()
-                )
-                logger.warning(
-                    f"[Autotuner]: Cache file {path} was created in a different "
-                    f"environment ({details}). Ignoring cached configs. "
-                    f"Results will not be saved to this file to avoid "
-                    f"overwriting configs from a different environment. "
-                    f"Use a different cache path to save configs for the "
-                    f"current environment."
+            hard, soft = _classify_metadata_mismatches(saved_meta, current_meta)
+            if hard or soft:
+                if hard:
+                    severity = "hard"
+                    consequence = (
+                        "Results will not be saved to this file to avoid "
+                        "overwriting configs from a different environment. "
+                        "Use a different cache path for the current "
+                        "environment, or delete the file to re-prime it."
+                    )
+                else:
+                    severity = "soft"
+                    consequence = (
+                        "The file's recorded environment is indeterminate, "
+                        "so the next save will replace it with freshly "
+                        "tuned configs for the current environment."
+                    )
+                message = (
+                    f"[Autotuner]: Cache file {path} was created in a "
+                    f"different environment "
+                    f"({_format_meta_mismatches({**hard, **soft})}). "
+                    f"Ignoring cached configs. {consequence}"
                 )
+                with self._lock:
+                    if severity in self._warned_cache_mismatch:
+                        logger.debug(message)
+                    else:
+                        self._warned_cache_mismatch.add(severity)
+                        logger.warning(message)
                 return False
 
         with self._lock:

{flashinfer_python-0.6.15 → flashinfer_python-0.6.15.post1}/flashinfer/mla/_core.py

@@ -24,7 +24,13 @@ from typing import List, Literal, Optional, Sequence, Tuple, Union, cast, overlo
 import torch
 
 from ..api_logging import flashinfer_api
-from ..autotuner import AutoTuner, TunableRunner
+from flashinfer.autotuner import (
+    AutoTuner,
+    TunableRunner,
+    TuningConfig,
+    make_bucket_mapper,
+    DynamicTensorSpec,
+)
 from ..trace.templates.attention import (
     mla_paged_decode_trace,
     trtllm_batch_decode_mla_trace_dispatch,
@@ -2008,7 +2014,7 @@ def _cute_dsl_max_supported_batch(
 
 def _compute_mla_decode_buckets(
     workspace_buffer: torch.Tensor,
-    runner_names: List[str],
+    runner_names: Sequence[str],
     q_len: int,
     num_heads: int,
     kv_lora_rank: int,
@@ -2142,44 +2148,27 @@ def _cute_dsl_incompatibility_reason(
     return None
 
 
-def _build_mla_decode_tuning_config(
-    kv_cache: torch.Tensor,
-    block_tables: torch.Tensor,
-    workspace_buffer: torch.Tensor,
-    runner_names: List[str],
-    q_len: int,
-    num_heads: int,
-    kv_lora_rank: int,
-    max_seq_len: int,
-    device: torch.device,
-):
-    """Per-call TuningConfig with bucket-capped batch sweep and closure-based initializers.
+@functools.cache
+def _mla_decode_tuning_config(
+    buckets: tuple[int, ...],
+    num_pages: int,
+    profile_seq_len: int,
+) -> TuningConfig:
+    """One TuningConfig (and one pair of initializer closures) per key.
+
+    Memoized because ``AutoTuner._find_nearest_profile`` lru-caches on
+    ``(shapes, tuning_config)``: a fresh config per dispatcher call shares
+    its hash with all previous ones (``DynamicTensorSpec.__hash__`` skips
+    ``tensor_initializers``) but never compares equal (closures compare by
+    identity), so would result in a leak.
 
     The DynamicTensorSpec sweeps batch dim across all four ``inputs`` tensors
     (query, block_tables, seq_lens, out). ``block_tables`` is initialized via
     ``random_(0, num_pages)`` which wraps mod kv_cache size — safe for autotune
     profiling because MLA decode reads kv_cache and never writes it, so aliased
     page reads give correct timing measurements. ``seq_lens`` is filled
-    homogeneously with ``min(max_seq_len, provisioned_max_seq_len)``.
+    homogeneously with ``profile_seq_len``.
     """
-    from ..autotuner import DynamicTensorSpec, TuningConfig, make_bucket_mapper
-
-    # kv_cache may be 3D [num_pages, page_size, D] or 4D
-    # [num_pages, 1, page_size, D] after `_check_trtllm_gen_mla_shape` —
-    # page_size is the second-to-last dim in both layouts.
-    page_size = kv_cache.shape[-2]
-    provisioned_max_seq_len = block_tables.shape[-1] * page_size
-    profile_seq_len = min(max_seq_len, provisioned_max_seq_len)
-    num_pages = kv_cache.shape[0]
-
-    buckets = _compute_mla_decode_buckets(
-        workspace_buffer,
-        runner_names,
-        q_len,
-        num_heads,
-        kv_lora_rank,
-        device,
-    )
 
     def init_block_tables(shapes, dtype, device):
         tensor = torch.empty(shapes, dtype=dtype, device=device)
@@ -2198,7 +2187,7 @@ def _build_mla_decode_tuning_config(
                 dim_idx=(0, 0, 0, 0),
                 gen_tuning_buckets=buckets,
                 map_to_tuning_buckets=make_bucket_mapper(buckets, round_map=False),
-                tensor_initializers=[None, init_block_tables, init_seq_lens, None],
+                tensor_initializers=(None, init_block_tables, init_seq_lens, None),
             ),
         ),
         use_cuda_graph=True,
@@ -2206,6 +2195,44 @@ def _build_mla_decode_tuning_config(
     )
 
 
+def _build_mla_decode_tuning_config(
+    kv_cache: torch.Tensor,
+    block_tables: torch.Tensor,
+    workspace_buffer: torch.Tensor,
+    runner_names: Sequence[str],
+    q_len: int,
+    num_heads: int,
+    kv_lora_rank: int,
+    max_seq_len: int,
+    device: torch.device,
+) -> TuningConfig:
+    """Reduce call args to the memoization key of ``_mla_decode_tuning_config``.
+
+    Key stability: num_pages and profile_seq_len are fixed per KV-cache
+    allocation. A caller that varies max_seq_len or block-table width per
+    call gets one cached config per distinct value (bounded, unlike the
+    per-call configs this replaces).
+    """
+    # kv_cache may be 3D [num_pages, page_size, D] or 4D
+    # [num_pages, 1, page_size, D] after `_check_trtllm_gen_mla_shape` —
+    # page_size is the second-to-last dim in both layouts.
+    page_size = kv_cache.shape[-2]
+    provisioned_max_seq_len = block_tables.shape[-1] * page_size
+    profile_seq_len = min(max_seq_len, provisioned_max_seq_len)
+    num_pages = kv_cache.shape[0]
+
+    buckets = _compute_mla_decode_buckets(
+        workspace_buffer,
+        runner_names,
+        q_len,
+        num_heads,
+        kv_lora_rank,
+        device,
+    )
+
+    return _mla_decode_tuning_config(buckets, num_pages, profile_seq_len)
+
+
 def _resolve_trtllm_gen_multi_ctas_kv_counter_buffer(
     counter_buffer: Optional[torch.Tensor],
     batch_size: int,

{flashinfer_python-0.6.15 → flashinfer_python-0.6.15.post1/flashinfer_python.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: flashinfer-python
-Version: 0.6.15
+Version: 0.6.15.post1
 Summary: FlashInfer: Kernel Library for LLM Serving
 Author: FlashInfer team
 License-Expression: Apache-2.0

{flashinfer_python-0.6.15 → flashinfer_python-0.6.15.post1}/include/flashinfer/trtllm/fused_moe/runner.h

@@ -443,6 +443,10 @@ class Runner {
                                                            int32_t numLocalExperts,
                                                            int32_t numTokens) const;
 
+  [[nodiscard]] bool isValidConfigIndex(int64_t configIndex, int32_t topK, int32_t hiddenSize,
+                                        int32_t intermediateSize, int32_t numLocalExperts,
+                                        int32_t numTokens) const;
+
   [[nodiscard]] int64_t getDefaultValidConfigIndex(int32_t topK, int32_t hiddenSize,
                                                    int32_t intermediateSize,
                                                    int32_t numLocalExperts,

flashinfer_python-0.6.15/flashinfer/_build_meta.py

@@ -1,3 +0,0 @@
-"""Build metadata for flashinfer package."""
-__version__ = "0.6.15"
-__git_version__ = "8eccd0c1352165302840c0e19066bc42d36dbd7a"