flashinfer-python 0.7.0rc3__tar.gz → 0.7.0rc4__tar.gz

{flashinfer_python-0.7.0rc3/flashinfer_python.egg-info → flashinfer_python-0.7.0rc4}/PKG-INFO

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

{flashinfer_python-0.7.0rc3 → flashinfer_python-0.7.0rc4}/csrc/cub_topk.cu

@@ -392,17 +392,22 @@ cudaError_t CUBBatchedTopKVarLenTransform(const DType* input, int64_t row_stride
   // synthesized by a counting iterator; d_values_out then receives the top-k source indices.
   auto d_values_in = cuda::make_constant_iterator(cuda::make_counting_iterator(int32_t{0}));
 
+  // Bound only CUB's selection K by the physical width. The caller's output
+  // iterators retain the requested top_k stride, and the existing tail fill pads
+  // to top_k. CUB still reads each logical length on device on every graph replay.
+  const int64_t selection_k = top_k < max_len ? top_k : max_len;
+
   if (row_starts != nullptr) {
     auto d_keys_in =
         cuda::make_transform_iterator(cuda::make_counting_iterator(int64_t{0}),
                                       CUBMakeRowIn<DType>{input, row_stride, row_starts});
     return CUBBatchedTopKDispatch(d_keys_in, d_keys_out, d_values_in, d_values_out, lengths,
-                                  maybe_workspace_buffer, num_rows, max_len, top_k, tie_break,
+                                  maybe_workspace_buffer, num_rows, max_len, selection_k, tie_break,
                                   query_bytes_out, stream);
   }
   auto d_keys_in = cuda::make_strided_iterator(cuda::make_counting_iterator(input), row_stride);
   return CUBBatchedTopKDispatch(d_keys_in, d_keys_out, d_values_in, d_values_out, lengths,
-                                maybe_workspace_buffer, num_rows, max_len, top_k, tie_break,
+                                maybe_workspace_buffer, num_rows, max_len, selection_k, tie_break,
                                 query_bytes_out, stream);
 }
 
@@ -439,7 +444,8 @@ cudaError_t CUBBatchedTopK(const DType* input, int64_t row_stride, DType* output
 }
 
 // Input-side validation shared by every CUB entry and workspace-size query.
-void CheckCUBTopKInput(const TensorView& input, int64_t top_k, int64_t tie_break) {
+void CheckCUBTopKInput(const TensorView& input, int64_t top_k, int64_t tie_break,
+                       bool allow_oversized_k = false) {
   // Rows only need to be individually contiguous: the row pitch is threaded through as
   // input.stride(0), so strided views (e.g. scores[:, :cur_len] of a wider buffer) work
   // without a .contiguous() copy.
@@ -451,16 +457,19 @@ void CheckCUBTopKInput(const TensorView& input, int64_t top_k, int64_t tie_break
       << ", expected 0 (none), 1 (prefer small indices), or 2 (prefer large indices)";
 
   const int64_t max_len = input.size(1);
-  TVM_FFI_ICHECK(top_k > 0 && top_k <= max_len)
-      << "cub_topk requires 0 < top_k <= d, got top_k=" << top_k << ", d=" << max_len;
-  TVM_FFI_ICHECK(max_len <= CUB_TOPK_MAX_LEN)
-      << "cub_topk supports d <= " << CUB_TOPK_MAX_LEN << ", got d=" << max_len;
+  TVM_FFI_ICHECK(top_k > 0) << "cub_topk requires top_k > 0, got top_k=" << top_k;
+  if (!allow_oversized_k) {
+    TVM_FFI_ICHECK(top_k <= max_len)
+        << "cub_topk requires top_k <= d, got top_k=" << top_k << ", d=" << max_len;
+  }
+  TVM_FFI_ICHECK(max_len > 0 && max_len <= CUB_TOPK_MAX_LEN)
+      << "cub_topk supports 0 < d <= " << CUB_TOPK_MAX_LEN << ", got d=" << max_len;
 }
 
 // Validation shared by the transform entries and their workspace-size queries.
 void CheckCUBTopKArgs(const TensorView& input, const TensorView& lengths, int64_t top_k,
                       int64_t tie_break) {
-  CheckCUBTopKInput(input, top_k, tie_break);
+  CheckCUBTopKInput(input, top_k, tie_break, /*allow_oversized_k=*/true);
   CHECK_INPUT(lengths);
   CHECK_DIM(1, lengths);  // lengths: (batch_size,)
   CHECK_INPUT_TYPE(lengths, dl_int32);
@@ -491,13 +500,7 @@ void cub_topk_page_table_transform(
                               maybe_output_raw_indices, top_k, page_size);
   // CUB-specific constraints; everything shared with the radix launcher lives in
   // CheckPageTableTransformArgs.
-  TVM_FFI_ICHECK(tie_break >= 0 && tie_break <= 2)
-      << "Invalid tie_break mode " << tie_break
-      << ", expected 0 (none), 1 (prefer small indices), or 2 (prefer large indices)";
-  TVM_FFI_ICHECK(top_k <= input.size(1))
-      << "cub_topk requires top_k <= d, got top_k=" << top_k << ", d=" << input.size(1);
-  TVM_FFI_ICHECK(input.size(1) <= CUB_TOPK_MAX_LEN)
-      << "cub_topk supports d <= " << CUB_TOPK_MAX_LEN << ", got d=" << input.size(1);
+  CheckCUBTopKInput(input, top_k, tie_break, /*allow_oversized_k=*/true);
 
   const auto* lengths_ptr = static_cast<const int32_t*>(lengths.data_ptr());
   const auto* row_to_batch_ptr =

{flashinfer_python-0.7.0rc3 → flashinfer_python-0.7.0rc4}/flashinfer/__main__.py

@@ -16,9 +16,11 @@ limitations under the License.
 
 import copy
 from email.parser import BytesParser
+from importlib.util import find_spec
 import os
 from pathlib import Path
 import re
+from shutil import which
 import subprocess
 import sys
 import tempfile
@@ -321,6 +323,21 @@ def _build_jit_cache_index_url(cuda_index_label: str, nightly: bool) -> str:
     return f"{base_url}/{cuda_index_label}"
 
 
+def _get_pip_install_cmd() -> list[str]:
+    uv = which("uv")
+    if uv is not None:
+        try:
+            config = (Path(sys.prefix) / "pyvenv.cfg").read_text(encoding="utf-8")
+        except (OSError, UnicodeError):
+            config = ""
+        created_by_uv = any(
+            line.partition("=")[0].strip() == "uv" for line in config.splitlines()
+        )
+        if created_by_uv or find_spec("pip") is None:
+            return [uv, "pip", "install", "--python", sys.executable]
+    return [sys.executable, "-m", "pip", "install"]
+
+
 def _build_pip_install_cmd(
     requirements: str | list[str],
     index_url: str,
@@ -329,13 +346,7 @@ def _build_pip_install_cmd(
 ) -> list[str]:
     if isinstance(requirements, str):
         requirements = [requirements]
-    cmd = [
-        sys.executable,
-        "-m",
-        "pip",
-        "install",
-        "--upgrade",
-    ]
+    cmd = [*_get_pip_install_cmd(), "--upgrade"]
     if no_deps:
         cmd.append("--no-deps")
     if nightly:

flashinfer_python-0.7.0rc4/flashinfer/_build_meta.py

@@ -0,0 +1,3 @@
+"""Build metadata for flashinfer package."""
+__version__ = "0.7.0rc4"
+__git_commit__ = "61c47a3abe92c061140d95a108683f30035979bc"

{flashinfer_python-0.7.0rc3 → flashinfer_python-0.7.0rc4}/flashinfer/comm/mnnvl.py

@@ -600,8 +600,8 @@ def is_mnnvl_fabric_supported(device_idx: int) -> bool:
         pynvml.nvmlDeviceGetGpuFabricInfoV(handle, ctypes.byref(fabric_info))
         return (
             fabric_info.state >= pynvml.NVML_GPU_FABRIC_STATE_COMPLETED
-            and fabric_info.clusterUuid
-            and fabric_info.clusterUuid[0] != 0
+            # A binary UUID may start with zero; only an all-zero UUID is absent.
+            and any(fabric_info.clusterUuid)
         )
     finally:
         pynvml.nvmlShutdown()

{flashinfer_python-0.7.0rc3 → flashinfer_python-0.7.0rc4}/flashinfer/comm/trtllm_moe_alltoall.py

@@ -463,6 +463,7 @@ def moe_a2a_dispatch(
     eplb_local_stats: Optional[torch.Tensor] = None,
     enable_rank_mask: bool = False,
     active_rank_mask: Optional[torch.Tensor] = None,
+    recv_view_cache: Optional[dict] = None,
 ):
     r"""Dispatch tokens and payloads to their target expert ranks.
 
@@ -506,6 +507,13 @@ def moe_a2a_dispatch(
         participates in this collective; tokens routed to a masked-off rank are dropped
         instead of hanging the collective.  Requires ``enable_rank_mask=True``; the local
         ``ep_rank``'s own bit must always be set.
+    recv_view_cache : dict, optional
+        Opaque cache of receive views.  Pass an empty dictionary to enable caching;
+        ``None`` (default) disables it and preserves the previous behaviour exactly.
+        The cache retains a reference to ``workspace`` and is cleared when a different
+        workspace tensor is passed.  Keep the backing allocation alive and do not change
+        the workspace or the cached views' shape, strides, or storage in place.  Clear
+        the dictionary to release its references; otherwise leave its contents unchanged.
 
     Returns
     -------
@@ -543,20 +551,48 @@ def moe_a2a_dispatch(
         active_rank_mask,
     )
 
+    # Bind the cache to this workspace once per dispatch, so the hot path below needs no
+    # pointer or device queries.  A cache carried over from a different (possibly freed
+    # and reallocated) workspace must not be reused.
+    if (
+        recv_view_cache is not None
+        and recv_view_cache.get("_workspace") is not workspace
+    ):
+        recv_view_cache.clear()
+        recv_view_cache["_workspace"] = workspace
+
     output_payloads = []
     for input_payload, offset, size in zip(
         input_payloads, recv_offsets, recv_sizes, strict=True
     ):
-        # This uses absolute offsets in the workspace, so skip indexing into the workspace
-        output_payloads.append(
-            moe_a2a_wrap_payload_tensor_in_workspace(
+        # This uses absolute offsets in the workspace, so skip indexing into the workspace.
+        # The view is a pure function of the key below, and it is a view *onto* the
+        # persistent workspace, so a reused view still observes the bytes this call
+        # writes -- only the Python-side reconstruction (four tensor-creating ops per
+        # payload) is skipped.
+        # ``int()`` on the offsets is deliberate and is a deviation from the upstream
+        # patch: they arrive from a TVM-FFI call, and anything int-like that hashes by
+        # identity would miss on every lookup -- silently disabling the cache while
+        # still growing the dict.  Coercing makes the key value-based either way.
+        key = (
+            ep_size,
+            runtime_max_tokens_per_rank,
+            int(offset),
+            int(size),
+            input_payload.dtype,
+        )
+        payload_view = None if recv_view_cache is None else recv_view_cache.get(key)
+        if payload_view is None:
+            payload_view = moe_a2a_wrap_payload_tensor_in_workspace(
                 workspace,
                 [ep_size, runtime_max_tokens_per_rank],
                 offset,
                 offset + size,
                 input_payload.dtype,
             )
-        )
+            if recv_view_cache is not None:
+                recv_view_cache[key] = payload_view
+        output_payloads.append(payload_view)
 
     eplb_gathered_stats = None
     if eplb_gathered_stats_offset >= 0:
@@ -1036,6 +1072,10 @@ class MoeAlltoAll:
         self.workspace = self._WORKSPACE["workspace"]
         self.metainfo = self._WORKSPACE["metainfo"]
         self._state = _A2AState()
+        # Keep receive views across dispatch calls.  Instance-owned rather than
+        # module-level: this object keeps ``self.workspace`` alive, so the cached views
+        # cannot alias a freed-and-reallocated workspace.
+        self._recv_view_cache: dict = {"_workspace": self.workspace}
 
     @property
     def eplb_gathered_stats(self) -> Optional[torch.Tensor]:
@@ -1220,6 +1260,7 @@ class MoeAlltoAll:
             eplb_local_stats=eplb_local_stats,
             enable_rank_mask=self.enable_rank_mask,
             active_rank_mask=active_rank_mask,
+            recv_view_cache=self._recv_view_cache,
         )
 
         # Update state

{flashinfer_python-0.7.0rc3 → flashinfer_python-0.7.0rc4}/flashinfer/kda_kernels/kda_decode_wy_output_only.py

@@ -2326,6 +2326,9 @@ def kda_wy_output_only(
     _total_ctas = HV * B
     _needed = math.ceil(_total_ctas / _num_sms)
     mbp = max(1, min(_needed + 1, 8))
+    if cc == (10, 7):
+        # mbp>1 makes the DSL auto-derive an SMEM carveout that faults (#4957).
+        mbp = 1
     _mbp_env = _os.environ.get("FLASHINFER_KDA_OO_MBP")
     if _mbp_env:
         mbp = int(_mbp_env)
@@ -2549,7 +2552,11 @@ def kda_recoverssm_verify(
     HV = H = num_heads
     t_disc = 4 if spec_query_len <= 4 else (8 if spec_query_len <= 8 else 16)
     _num_sms = torch.cuda.get_device_properties(device).multi_processor_count
+    cc = torch.cuda.get_device_capability(device)
     mbp = max(1, min(math.ceil(HV * batch / _num_sms) + 1, 8))
+    if cc == (10, 7):
+        # mbp>1 makes the DSL auto-derive an SMEM carveout that faults (#4957).
+        mbp = 1
     # Gate modes mirror recurrent_kda: precomputed log-space gate, the
     # Kimi K3 lower-bound sigmoid gate, or the Kimi-Linear softplus gate.
     if use_gate_in_kernel:
@@ -2559,7 +2566,6 @@ def kda_recoverssm_verify(
     lb = float(lower_bound) if lower_bound is not None else 0.0
     q_scale = float(scale) if scale is not None else float(key_dim**-0.5)
 
-    cc = torch.cuda.get_device_capability(device)
     cache_key = (
         "dropin",
         str(device),

{flashinfer_python-0.7.0rc3 → flashinfer_python-0.7.0rc4}/flashinfer/topk.py

@@ -1159,8 +1159,10 @@ def top_k_page_table_transform(
     - This is specifically designed for sparse attention's second stage.
     - ``input`` may have padding between rows, but its last dimension must be
       contiguous.
-    - If ``lengths[i] <= k``, raw indices are ``0..lengths[i]-1`` and remaining
-      positions are set to -1.
+    - If ``lengths[i] <= k``, all valid indices ``0..lengths[i]-1`` are
+      selected, including when ``k > max_len``. Outputs retain shape
+      ``(num_rows, k)``; unused positions are ``-1`` in both the mapped output
+      and optional raw indices. Padding is not translated through the page table.
 
     Examples
     --------
@@ -1369,8 +1371,9 @@ def top_k_ragged_transform(
     ----
     - This is specifically designed for sparse attention's second stage with
       ragged KV cache layout.
-    - If lengths[i] <= k, the output contains [offsets[i], offsets[i]+1, ..., offsets[i]+lengths[i]-1]
-      with remaining positions set to -1.
+    - If ``lengths[i] <= k``, all valid indices receive ``offsets[i]``,
+      including when ``k > max_len``. Output retains shape ``(num_rows, k)``;
+      unused positions remain ``-1`` without adding the offset.
 
     Examples
     --------

{flashinfer_python-0.7.0rc3 → flashinfer_python-0.7.0rc4/flashinfer_python.egg-info}/PKG-INFO

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

flashinfer_python-0.7.0rc3/flashinfer/_build_meta.py

@@ -1,3 +0,0 @@
-"""Build metadata for flashinfer package."""
-__version__ = "0.7.0rc3"
-__git_commit__ = "03296170f944590a611605675c4db5b61e8b7a98"