flashinfer-python 0.6.7.post1__tar.gz → 0.6.7.post2__tar.gz

{flashinfer_python-0.6.7.post1 → flashinfer_python-0.6.7.post2}/PKG-INFO

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

{flashinfer_python-0.6.7.post1 → flashinfer_python-0.6.7.post2}/csrc/trtllm_fmha_kernel_launcher.cu

@@ -87,7 +87,8 @@ void trtllm_paged_attention_launcher(
     int64_t o_sf_vec_size, int64_t o_sf_start_index, int64_t window_left, int64_t sum_seq_q,
     int64_t sparse_mla_top_k, float skip_softmax_threshold_scale_factor, bool skips_softmax,
     bool uses_shared_paged_kv_idx, int64_t sm_count, bool enable_pdl, int64_t workspace_size,
-    cudaStream_t stream) {
+    int64_t k_sf_stride_heads, int64_t k_sf_stride_batch, int64_t v_sf_stride_heads,
+    int64_t v_sf_stride_batch, cudaStream_t stream) {
   if (num_qo_heads % num_kv_heads != 0) {
     std::ostringstream err_msg;
     err_msg << "num_qo_heads must be a multiple of num_kv_heads, got num_kv_heads: " << num_kv_heads
@@ -126,6 +127,10 @@ void trtllm_paged_attention_launcher(
   runner_params.vStrideKeysValues = kv_stride_keys_values;
   runner_params.vStrideHeads = kv_stride_heads;
   runner_params.vStrideBatch = kv_stride_batch;
+  runner_params.kSfStrideHeads = k_sf_stride_heads;
+  runner_params.kSfStrideBatch = k_sf_stride_batch;
+  runner_params.vSfStrideHeads = v_sf_stride_heads;
+  runner_params.vSfStrideBatch = v_sf_stride_batch;
   runner_params.mNumPagesInMemPool = num_pages_in_mem_pool;
   runner_params.stream = stream;
   // the scaleSoftmaxLog2Ptr and outputScalePtr have higher priority than the scaleSoftmaxLog2 and
@@ -299,6 +304,19 @@ void trtllm_paged_attention_decode(
   const void* v_block_scales_ptr =
       value_block_scales.has_value() ? value_block_scales.value().data_ptr() : nullptr;
 
+  // Read actual scale factor strides from the scale tensors (HND layout: [pages, heads, N, D/16]).
+  // These are passed separately to the kernel instead of being derived from KV data strides.
+  int k_sf_stride_heads = 0, k_sf_stride_batch = 0;
+  int v_sf_stride_heads = 0, v_sf_stride_batch = 0;
+  if (key_block_scales.has_value()) {
+    k_sf_stride_heads = key_block_scales.value().stride(-3);
+    k_sf_stride_batch = key_block_scales.value().stride(0);
+  }
+  if (value_block_scales.has_value()) {
+    v_sf_stride_heads = value_block_scales.value().stride(-3);
+    v_sf_stride_batch = value_block_scales.value().stride(0);
+  }
+
   const auto stream = get_stream(query.device());
   void* output_sf_ptr =
       out_scale_factor.has_value() ? out_scale_factor.value().data_ptr() : nullptr;
@@ -345,7 +363,8 @@ void trtllm_paged_attention_decode(
       max_num_blocks_per_seq, bmm1_scale_value, bmm2_scale_value, bmm1_scale_log2_ptr,
       bmm2_scale_ptr, o_sf_scale, o_sf_vec_size, o_sf_start_index, window_left, sum_seq_q,
       sparse_mla_top_k, skip_softmax_threshold_scale_factor_value, skips_softmax,
-      uses_shared_paged_kv_idx_value, sm_count, enable_pdl, workspace_size, stream);
+      uses_shared_paged_kv_idx_value, sm_count, enable_pdl, workspace_size, k_sf_stride_heads,
+      k_sf_stride_batch, v_sf_stride_heads, v_sf_stride_batch, stream);
 }
 
 void trtllm_paged_attention_context(
@@ -407,6 +426,18 @@ void trtllm_paged_attention_context(
   const void* v_block_scales_ptr =
       value_block_scales.has_value() ? value_block_scales.value().data_ptr() : nullptr;
 
+  // Read actual scale factor strides from the scale tensors (HND layout: [pages, heads, N, D/16]).
+  int k_sf_stride_heads = 0, k_sf_stride_batch = 0;
+  int v_sf_stride_heads = 0, v_sf_stride_batch = 0;
+  if (key_block_scales.has_value()) {
+    k_sf_stride_heads = key_block_scales.value().stride(-3);
+    k_sf_stride_batch = key_block_scales.value().stride(0);
+  }
+  if (value_block_scales.has_value()) {
+    v_sf_stride_heads = value_block_scales.value().stride(-3);
+    v_sf_stride_batch = value_block_scales.value().stride(0);
+  }
+
   const auto stream = get_stream(query.device());
   void* output_sf_ptr =
       out_scale_factor.has_value() ? out_scale_factor.value().data_ptr() : nullptr;
@@ -455,7 +486,8 @@ void trtllm_paged_attention_context(
       kv_stride_heads, kv_stride_batch, max_num_blocks_per_seq, bmm1_scale_value, bmm2_scale_value,
       bmm1_scale_log2_ptr, bmm2_scale_ptr, o_sf_scale, o_sf_vec_size, o_sf_start_index, window_left,
       sum_seq_q, /*sparse_mla_top_k=*/0, skip_softmax_threshold_scale_factor_value, skips_softmax,
-      uses_shared_paged_kv_idx_value, sm_count, enable_pdl, workspace_size, stream);
+      uses_shared_paged_kv_idx_value, sm_count, enable_pdl, workspace_size, k_sf_stride_heads,
+      k_sf_stride_batch, v_sf_stride_heads, v_sf_stride_batch, stream);
 }
 
 void trtllm_ragged_attention_launcher(

flashinfer_python-0.6.7.post2/flashinfer/_build_meta.py

@@ -0,0 +1,3 @@
+"""Build metadata for flashinfer package."""
+__version__ = "0.6.7.post2"
+__git_version__ = "49e2c92c35f7dec3d1e69731486fc5da2a17637f"

{flashinfer_python-0.6.7.post1 → flashinfer_python-0.6.7.post2}/flashinfer/artifacts.py

@@ -135,7 +135,7 @@ class ArtifactPath:
     When compiling new cubins for backend directories, update the corresponding path.
     """
 
-    TRTLLM_GEN_FMHA: str = "e7afc4134bb53eaab63fb85163d5943fb190621c/fmha/trtllm-gen/"
+    TRTLLM_GEN_FMHA: str = "55bba55929d4093682e32d817bd11ffb0441c749/fmha/trtllm-gen/"
     TRTLLM_GEN_BMM: str = (
         "b55211623be7f5697c5262ffd8361fc06c147bc9/batched_gemm-b3c1646-c111d7c/"
     )
@@ -155,7 +155,7 @@ class CheckSumHash:
     """
 
     TRTLLM_GEN_FMHA: str = (
-        "5bd87798e560a63e883902fc5468146ffff0d3551bf337d2f81bd02893e9dc39"
+        "f2c0aad1e74391c4267a2f9a20ec819358b59e04588385cffb452ed341500b99"
     )
     TRTLLM_GEN_BMM: str = (
         "0af823880730c4f0b3832d2208fab035946694b83444410b9309db5613d60195"

{flashinfer_python-0.6.7.post1 → flashinfer_python-0.6.7.post2}/flashinfer/decode.py

@@ -2245,6 +2245,11 @@ def trtllm_batch_decode_with_kv_cache(
         or [num_pages, page_size, num_kv_heads, head_dim] if :attr:`kv_layout` is ``NHD``.
         The first tensor is the key cache, and the second tensor is the value cache.
 
+        **Contiguity requirements (trtllm-gen backend):**
+
+        - The ``head_dim`` (last dim) **must** have stride 1. This is a TMA hardware constraint
+        - The head and batch/page dims can have arbitrary strides.
+
     workspace_buffer : torch.Tensor. Must be initialized to 0 for its first use.
         workspace
 
@@ -2290,6 +2295,10 @@ def trtllm_batch_decode_with_kv_cache(
     kv_layout : str = "HND"
         The layout of the input k/v tensors, could be either ``NHD`` or ``HND``.
         Defaults to ``HND``.
+        For the trtllm-gen backend with NVFP4 KV cache, using ``NHD`` will trigger an
+        automatic transpose and ``.contiguous()`` copy of both the KV data and block scale
+        tensors to convert them to HND layout. This incurs extra memory allocation and
+        data copy overhead. Use ``HND`` for better performance.
 
     enable_pdl : Optional[bool] = None
         Whether to enable Programmatic Dependent Launch (PDL). See https://docs.nvidia.com/cuda/cuda-c-programming-guide/#programmatic-dependent-launch-and-synchronization
@@ -2317,6 +2326,20 @@ def trtllm_batch_decode_with_kv_cache(
         Only supported by trtllm-gen backend. Must be provided together with ``max_q_len``.
         When None, all requests use uniform query length specified by ``q_len_per_req``.
 
+    kv_block_scales : Optional[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]] = None
+        Per-block scale factors for NVFP4 KV cache. Either a tuple of (k_scales, v_scales) or
+        a single tensor with shape ``[num_pages, 2, ...]`` that will be unbound along dim=1.
+        Each scale tensor has shape ``[num_pages, num_kv_heads, page_size, head_dim // 16]``
+        in HND layout, with dtype ``torch.float8_e4m3fn``.
+
+        **Contiguity requirements (trtllm-gen backend):**
+
+        - The last two dims (``page_size``, ``head_dim // 16``) **must** be contiguous
+          (i.e., ``stride[-1] == 1`` and ``stride[-2] == head_dim // 16``). This is because
+          the kernel reshapes them into ``(16, page_size * head_dim / 16 / 16)`` to satisfy
+          TMA's 16-byte box width minimum.
+        - The head and batch/page dims can have arbitrary strides.
+
     skip_softmax_threshold_scale_factor: Optional[float] = None
         threshold scale factor for skipping softmax operations.
         Providing a value for this parameter enables skip-softmax sparsity as described in: https://arxiv.org/abs/2512.12087

{flashinfer_python-0.6.7.post1 → flashinfer_python-0.6.7.post2}/flashinfer/prefill.py

@@ -3555,7 +3555,7 @@ def trtllm_ragged_attention_deepseek(
     window_left: int,
     cum_seq_lens_q: torch.Tensor,
     cum_seq_lens_kv: torch.Tensor,
-    enable_pdl: bool | None,
+    enable_pdl: Optional[bool],
     is_causal: bool,
     return_lse: bool,
     attention_sinks: Optional[torch.Tensor] = None,
@@ -3748,6 +3748,11 @@ def trtllm_batch_context_with_kv_cache(
         If kv_cache is a tuple of two tensors, it should be a tuple of two tensors with shape [num_pages, num_kv_heads, page_size, head_dim] if :attr:`kv_layout` is "HND",
         or [num_pages, page_size, num_kv_heads, head_dim] if :attr:`kv_layout` is "NHD".
         The first tensor is the key cache, the second tensor is the value cache.
+
+        **Contiguity requirements (trtllm-gen backend):**
+
+        - The ``head_dim`` (last dim) **must** have stride 1. This is a TMA hardware constraint
+        - The head and batch/page dims can have arbitrary strides.
     workspace_buffer : torch.Tensor. Must be initialized to 0 for its first use.
         workspace
     block_tables : torch.Tensor
@@ -3789,11 +3794,25 @@ def trtllm_batch_context_with_kv_cache(
         Defaults to ``None``, which means it will be enabled if the device supports PDL.
     kv_layout : str = "HND"
         Layout of kv-cache, can be "HND" or "NHD", default is "HND".
+        For the trtllm-gen backend with NVFP4 KV cache, using ``NHD`` will trigger an
+        automatic transpose and ``.contiguous()`` copy of both the KV data and block scale
+        tensors to convert them to HND layout. This incurs extra memory allocation and
+        data copy overhead. Use ``HND`` for better performance.
     sinks : Optional[List[torch.Tensor]] = None
         additional value per head in the denominator of the softmax.
     kv_block_scales : Optional[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]] = None
         Per-block scale factors for NVFP4 KV cache. Either a tuple of (k_scales, v_scales) or
-        a single tensor with shape [num_pages, 2, ...] that will be unbound along dim=1.
+        a single tensor with shape ``[num_pages, 2, ...]`` that will be unbound along dim=1.
+        Each scale tensor has shape ``[num_pages, num_kv_heads, page_size, head_dim // 16]``
+        in HND layout, with dtype ``torch.float8_e4m3fn``.
+
+        **Contiguity requirements (trtllm-gen backend):**
+
+        - The last two dims (``page_size``, ``head_dim // 16``) **must** be contiguous
+          (i.e., ``stride[-1] == 1`` and ``stride[-2] == head_dim // 16``). This is because
+          the kernel reshapes them into ``(16, page_size * head_dim / 16 / 16)`` to satisfy
+          TMA's 16-byte box width minimum.
+        - The head and batch/page dims can have arbitrary strides.
     skip_softmax_threshold_scale_factor: Optional[float] = None
         threshold scale factor for skipping softmax operations.
         Providing a value for this parameter enables skip-softmax sparsity as described in: https://arxiv.org/abs/2512.12087

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

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

{flashinfer_python-0.6.7.post1 → flashinfer_python-0.6.7.post2}/include/flashinfer/trtllm/fmha/fmhaRunnerParams.h

@@ -251,6 +251,15 @@ struct TllmGenFmhaRunnerParams {
   // The stride between different batches for V.
   int vStrideBatch;
 
+  // The stride between different heads for K scaling factors.
+  int kSfStrideHeads;
+  // The stride between different batches for K scaling factors.
+  int kSfStrideBatch;
+  // The stride between different heads for V scaling factors.
+  int vSfStrideHeads;
+  // The stride between different batches for V scaling factors.
+  int vSfStrideBatch;
+
   // Head dimension for Q and K.
   int mHeadDimQk;
   // Head dimension for V.

{flashinfer_python-0.6.7.post1 → flashinfer_python-0.6.7.post2}/include/flashinfer/trtllm/fmha/kernelParams.h

@@ -407,10 +407,17 @@ struct KernelParams {
     return std::make_tuple(strideKeysVals, strideHeads, strideBatch);
   }
 
-  // Create the TMA shape/stride for K.
+  // Create the TMA shape/stride for K/V data tensors.
+  //
+  // Layout requirement (HND): [num_pages, num_kv_heads, page_size, head_dim]
+  //   - head_dim (last dim) MUST have stride 1. This is a TMA hardware constraint:
+  //     cuTensorMapEncodeTiled does not accept a stride for dim 0 and implicitly assumes 1.
+  //   - Other dimensions (heads, batch/pages) can have arbitrary strides; the actual
+  //     strides are read from the tensor and passed to the TMA descriptor.
   template <class FmhaOptions>
   static auto makeTmaShapeStrideKv(FmhaOptions const& options, KernelParams const& params,
-                                   Data_type dtypeKv, bool isK, bool storeTransformedKvInTmem) {
+                                   Data_type dtypeKv, bool isK, bool storeTransformedKvInTmem,
+                                   int reshapeFactor) {
     // The shape elements.
     auto [numKeys, numHeadsQPerKv, batchSize] = makeShapeKv(options, params);
     // The stride elements.
@@ -429,31 +436,47 @@ struct KernelParams {
     // dimension (it's stride is always 0). For V, the headDim dimension is already the first
     // dimension so no swapping is needed.
 
-    // Therefore, the resulting TMA layout is 4D: (headDim, numKeys, numHeadsKv, batchSize):(1,
-    // strideKeys, strideHeads, strideBatch)
+    // The TMA layout is 4D:
+    // Shape:  (headDim, numKeys, numHeadsKv, batchSize)
+    // Stride: (1, strideKeys, strideHeads, strideBatch)
+    //
+    // We reshape by a factor r so the headDim is 128B+, reducing the number of TMA read requests:
+    // Shape:  (headDim * r, numKeys / r, numHeadsKv, batchSize)
+    // Stride: (1, strideKeys * r, strideHeads, strideBatch)
 
     // Note that for FP4 KV input, elements are stored as uint8_t, each packs 2 FP4 elements.
     // The column index and strides needs to divide by 2.
     auto const colIdxDivisor = dtypeKv == DATA_TYPE_E2M1 ? 2 : 1;
     auto shape = std::vector<uint64_t>{
-        storeTransformedKvInTmem ? headDim : static_cast<uint64_t>(headDim / colIdxDivisor),
-        static_cast<uint64_t>(numKeys), static_cast<uint64_t>(options.mNumHeadsKv),
+        (storeTransformedKvInTmem ? headDim : static_cast<uint64_t>(headDim / colIdxDivisor)) *
+            reshapeFactor,
+        static_cast<uint64_t>(numKeys / reshapeFactor), static_cast<uint64_t>(options.mNumHeadsKv),
         static_cast<uint64_t>(batchSize)};
-    auto stride = std::vector<uint64_t>{1, static_cast<uint64_t>(strideKeys / colIdxDivisor),
-                                        static_cast<uint64_t>(strideHeads / colIdxDivisor),
-                                        static_cast<uint64_t>(strideBatch / colIdxDivisor)};
+    auto stride =
+        std::vector<uint64_t>{1, static_cast<uint64_t>(strideKeys / colIdxDivisor * reshapeFactor),
+                              static_cast<uint64_t>(strideHeads / colIdxDivisor),
+                              static_cast<uint64_t>(strideBatch / colIdxDivisor)};
 
     return std::make_tuple(shape, stride);
   }
 
-  // Create the TMA shape/stride for KV scaling factors.
+  // Create the TMA shape/stride for KV scaling factors (block scales for NVFP4 KV cache).
+  //
+  // Layout requirement (HND): [num_pages, num_kv_heads, page_size, head_dim // 16]
+  //   - The last two dims (page_size, head_dim // 16) MUST be contiguous (stride[-1] = 1,
+  //     stride[-2] = head_dim // 16). This is because we reshape them into
+  //     (16, page_size * head_dim / 16 / 16) with hardcoded stride[1] = 16 to satisfy TMA's
+  //     16-byte box width requirement. Each scale factor is 1 byte (FP8), and head_dim // 16
+  //     can be < 16 (e.g., 8 for head_dim=128), so we must merge with page_size to reach 16.
+  //   - The head and batch/page strides are read from the actual scale tensors (kSfStrideHeads,
+  //     kSfStrideBatch) and can differ from the KV data strides.
+  //   - cuTensorMapEncodeTiled requires all non-dim0 strides to be multiples of 16 bytes, so
+  //     sfStrideHeads and sfStrideBatch must each be a multiple of 16.
   template <class FmhaOptions>
   static auto makeTmaShapeStrideKvSf(FmhaOptions const& options, KernelParams const& params,
-                                     bool isK) {
+                                     bool isK, int reshapeFactor) {
     // The shape elements.
     auto [numKeys, numHeadsQPerKv, batchSize] = makeShapeKv(options, params);
-    // The stride elements.
-    auto [strideKeys, strideHeads, strideBatch] = makeStrideKv(options, isK);
 
     // The headDim.
     // Note that contiguousKv or pagedKv will pad K and V to maxHeadDimKv.
@@ -464,20 +487,25 @@ struct KernelParams {
     // The number of elements per SF.
     int32_t NumEltsPerSf = 16;
 
+    // Use actual scale factor strides instead of deriving from KV strides.
+    int32_t sfStrideHeads = isK ? options.kSfStrideHeads : options.vSfStrideHeads;
+    int32_t sfStrideBatch = isK ? options.kSfStrideBatch : options.vSfStrideBatch;
+
     // The KV shape is: (headDim, numKeys, numHeadsKv, batchSize)
-    // Therefore, the KV SF shape should be (headDim / NumEltsPerSf, numKeys, numHeadsKv,
-    // batchSize). Considering the TMA requires box width to be multiple of 16B, without changing
-    // the underlying layout, we reshape into (16, numKeys * headDim / NumEltsPerSf / 16,
-    // numHeadsKv, batchSize)
+    // The KV SF shape is: (headDim / NumEltsPerSf, numKeys, numHeadsKv, batchSize).
+    // Considering the TMA requires box width to be multiple of 16B, and ideally >= 128B, we reshape
+    // with a factor into: (headDim / NumEltsPerSf * r, numKeys / r, numHeadsKv, batchSize)
 
     // Note that it only works for pagedKv layout.
     FLASHINFER_CHECK(isPagedKv(options.mQkvLayout), "The qkvLayout is not supported.");
 
     auto shape = std::vector<uint64_t>{
-        16, static_cast<uint64_t>(numKeys * headDim / NumEltsPerSf / 16),
-        static_cast<uint64_t>(options.mNumHeadsKv), static_cast<uint64_t>(batchSize)};
-    auto stride = std::vector<uint64_t>{1, 16, static_cast<uint64_t>(strideHeads / NumEltsPerSf),
-                                        static_cast<uint64_t>(strideBatch / NumEltsPerSf)};
+        static_cast<uint64_t>(headDim / NumEltsPerSf * reshapeFactor),
+        static_cast<uint64_t>(numKeys / reshapeFactor), static_cast<uint64_t>(options.mNumHeadsKv),
+        static_cast<uint64_t>(batchSize)};
+    auto stride = std::vector<uint64_t>{
+        1, static_cast<uint64_t>(headDim / NumEltsPerSf * reshapeFactor),
+        static_cast<uint64_t>(sfStrideHeads), static_cast<uint64_t>(sfStrideBatch)};
 
     return std::make_tuple(shape, stride);
   }
@@ -676,20 +704,39 @@ struct KernelParams {
     bool const storeTransformedKvInTmem{kernelMeta.mDataTypeKv == DATA_TYPE_E2M1 &&
                                         kernelMeta.mDataTypeQ == DATA_TYPE_E4M3 &&
                                         maxHeadDimKv >= 128 && isSwapsMmaAb};
-    // Shape/stride for gmem tensor Kv.
-    auto [shapeK, strideK] = makeTmaShapeStrideKv(options, params, kernelMeta.mDataTypeKv,
-                                                  /*isK*/ true, storeTransformedKvInTmem);
-    auto [shapeV, strideV] = makeTmaShapeStrideKv(options, params, kernelMeta.mDataTypeKv,
-                                                  /*isK*/ false, storeTransformedKvInTmem);
     // Whether swizzle is needed for K/V.
     bool const swizzleKv{storeTransformedKvInTmem || !transformsKv};
+    // Whether we can reshape the TMA box for K/V to widen it to 128B.
+    bool const canReshapeTmaKv{isPagedKv(options.mQkvLayout) &&
+                               options.mHeadDimQk == options.mHeadDimV && !swizzleKv};
+    // The reshape factor for K/V TMA box: aim for 128B box width.
+    //   - 128 / maxHeadDimKv: keeps first-dim tile <= 128 elts (CU_TENSOR_MAP_SWIZZLE_128B limit).
+    //   - 128 / (maxHeadDimKv * bytesPerElt): factor needed to reach 128B box width.
+    //   - numKeysPerTile: can't reshape more rows than the tile covers.
+    int32_t const reshapeFactorKv{
+        canReshapeTmaKv
+            ? std::max(
+                  1,
+                  std::min(
+                      {128 / maxHeadDimKv,
+                       128 / (maxHeadDimKv *
+                              static_cast<int32_t>(get_size_in_bits(kernelMeta.mDataTypeKv)) / 8),
+                       numKeysPerTile}))
+            : 1};
+    // Shape/stride for gmem tensor Kv.
+    auto [shapeK, strideK] =
+        makeTmaShapeStrideKv(options, params, kernelMeta.mDataTypeKv,
+                             /*isK*/ true, storeTransformedKvInTmem, reshapeFactorKv);
+    auto [shapeV, strideV] =
+        makeTmaShapeStrideKv(options, params, kernelMeta.mDataTypeKv,
+                             /*isK*/ false, storeTransformedKvInTmem, reshapeFactorKv);
     // Note that for FP4 KV input, elements are stored as uint8_t, each packs 2 FP4 elements.
     auto const numEltsDivisor =
         kernelMeta.mDataTypeKv == DATA_TYPE_E2M1 && !storeTransformedKvInTmem ? 2 : 1;
     // The tileShapes for K/V.
     std::vector<uint32_t> tileShapeKv(shapeK.size(), 1);
-    tileShapeKv[0] = numEltsInClampedHeadDimKv / numEltsDivisor;
-    tileShapeKv[1] = numKeysPerTile;
+    tileShapeKv[0] = numEltsInClampedHeadDimKv / numEltsDivisor * reshapeFactorKv;
+    tileShapeKv[1] = numKeysPerTile / reshapeFactorKv;
 
     // If sparse MLA is enabled, the shape and stride for K need to be updated for 2D layout
     // (numTokensKvInPagedKv, headDimQk).
@@ -712,18 +759,22 @@ struct KernelParams {
     if (kernelMeta.mDataTypeKv == DATA_TYPE_E2M1) {
       // The number of elements per SF.
       int32_t NumEltsPerSf = 16;
+      // The reshape factor for K/V SF: aim for box width 128B, limit to numKeysPerTile.
+      int32_t const reshapeFactorKvSf =
+          std::min(128 / (maxHeadDimKv / NumEltsPerSf), numKeysPerTile);
       // Compute the shape and stride for SF tensor.
       // FIXME: assume K and V uses the same shape.
-      auto [shapeKvSf, strideKvSf] = makeTmaShapeStrideKvSf(options, params, /*isK*/ true);
+      auto [shapeKvSf, strideKvSf] =
+          makeTmaShapeStrideKvSf(options, params, /*isK*/ true, reshapeFactorKvSf);
 
       // The tileShapes for K/V.
       std::vector<uint32_t> tileShapeKvSf(shapeKvSf.size(), 1);
-      tileShapeKvSf[0] = 16;
-      tileShapeKvSf[1] = numKeysPerTile * maxHeadDimKv / NumEltsPerSf / 16;
+      tileShapeKvSf[0] = maxHeadDimKv / NumEltsPerSf * reshapeFactorKvSf;
+      tileShapeKvSf[1] = numKeysPerTile / reshapeFactorKvSf;
 
-      // The tile box is reshaped from (headDim / NumEltsPerSf, tileSizeKv) into (16, tileSizeKv *
-      // headDim / NumEltsPerSf / 16). See makeTmaShapeStrideKvSf for details. Build tma descriptor
-      // for K SF.
+      // The tile box is reshaped from (headDim / NumEltsPerSf, tileSizeKv) into
+      // (headDim / NumEltsPerSf * reshapeFactorKvSf, tileSizeKv / reshapeFactorKvSf).
+      // Build tma descriptor for K SF.
       params.tmaKSf_ = buildNdTmaDescriptor(options, DATA_TYPE_E4M3, shapeKvSf, strideKvSf,
                                             tileShapeKvSf, const_cast<void*>(options.kSfBasePtr),
                                             /*swizzled = */ false);

flashinfer_python-0.6.7.post1/flashinfer/_build_meta.py

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