flashinfer-python 0.6.2__tar.gz → 0.6.3__tar.gz

{flashinfer_python-0.6.2 → flashinfer_python-0.6.3}/PKG-INFO

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

flashinfer_python-0.6.3/csrc/flashinfer_mamba_binding.cu

@@ -0,0 +1,53 @@
+/*
+ * Copyright (c) 2025 by FlashInfer team.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include "tvm_ffi_utils.h"
+
+// Declare the function(s) implemented in selective_state_update.cu
+using tvm::ffi::Optional;
+
+namespace flashinfer::mamba {
+
+void selective_state_update(
+    TensorView state,              // (batch, dim, dstate) or (batch, nheads, dim, dstate)
+    TensorView x,                  // (batch, dim) or (batch, nheads, dim) for single-token
+                                   // or (batch, T, nheads, dim) for multi-token
+    TensorView dt,                 // (batch, dim) or (batch, nheads, dim) for single-token
+                                   // or (batch, T, nheads, dim) for multi-token
+    TensorView A,                  // (dim, dstate) or (nheads, dim, dstate)
+    TensorView B,                  // (batch, dstate) or (batch, ngroups, dstate) for single-token
+                                   // or (batch, T, ngroups, dstate) for multi-token
+    TensorView C,                  // (batch, dstate) or (batch, ngroups, dstate) for single-token
+                                   // or (batch, T, ngroups, dstate) for multi-token
+    TensorView D,                  // (dim,) or (nheads, dim)
+    Optional<TensorView> z,        // (batch, dim) or (batch, nheads, dim) for single-token
+                                   // or (batch, T, nheads, dim) for multi-token
+    Optional<TensorView> dt_bias,  // (dim,) or (nheads, dim)
+    bool dt_softplus,
+    Optional<TensorView> state_batch_indices,  // (batch,)
+    int64_t pad_slot_id,
+    TensorView output,  // same as x
+    bool disable_state_update,
+    Optional<TensorView> intermediate_states_buffer,  // (batch, cache_steps, nheads, dim, dstate)
+    Optional<TensorView> intermediate_state_indices,  // (batch,)
+    int64_t cache_steps);
+
+}  // namespace flashinfer::mamba
+
+// Export the function(s) via TVM-FFI
+// This enables cross-language bindings (not just PyTorch)
+TVM_FFI_DLL_EXPORT_TYPED_FUNC(selective_state_update, flashinfer::mamba::selective_state_update);
+
+// Add more mamba operations here as they are implemented

flashinfer_python-0.6.3/csrc/flat_prefill_kernel_delta_rule_sm90_extern.inc

@@ -0,0 +1,64 @@
+/*
+ * Copyright (c) 2025 by FlashInfer team.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+// Extern template declarations to prevent implicit instantiation in the dispatcher.
+// Explicit instantiations are in separate generated files for parallel compilation.
+
+#pragma once
+
+#include <cuda_bf16.h>
+#include <cuda_fp16.h>
+#include "cutlass/arch/arch.h"
+
+namespace flat {
+
+// clang-format off
+
+#define FOR_EACH_BOOL_4(MACRO, ...) \
+    MACRO(false, false, false, false, __VA_ARGS__) \
+    MACRO(false, false, false, true, __VA_ARGS__) \
+    MACRO(false, false, true, false, __VA_ARGS__) \
+    MACRO(false, false, true, true, __VA_ARGS__) \
+    MACRO(false, true, false, false, __VA_ARGS__) \
+    MACRO(false, true, false, true, __VA_ARGS__) \
+    MACRO(false, true, true, false, __VA_ARGS__) \
+    MACRO(false, true, true, true, __VA_ARGS__) \
+    MACRO(true, false, false, false, __VA_ARGS__) \
+    MACRO(true, false, false, true, __VA_ARGS__) \
+    MACRO(true, false, true, false, __VA_ARGS__) \
+    MACRO(true, false, true, true, __VA_ARGS__) \
+    MACRO(true, true, false, false, __VA_ARGS__) \
+    MACRO(true, true, false, true, __VA_ARGS__) \
+    MACRO(true, true, true, false, __VA_ARGS__) \
+    MACRO(true, true, true, true, __VA_ARGS__)
+
+#define DECLARE_TEMPLATE_INSTANCE(is_gva, needs_beta, needs_alpha, init_state, ctype) \
+extern template void launch_delta_rule_prefill_kernel_gbai<is_gva, needs_beta, needs_alpha, init_state, cutlass::arch::Sm90, ctype, ctype, float>( \
+    cudaStream_t, ctype*, float*, ctype const*, ctype const*, ctype const*, \
+    float const*, float const*, float const*, int64_t const*, uint8_t*, int32_t, int32_t, \
+    int32_t, int32_t, int32_t, int32_t, int64_t, float, int32_t);
+
+// Extern template declarations for half
+FOR_EACH_BOOL_4(DECLARE_TEMPLATE_INSTANCE, half)
+
+// Extern template declarations for nv_bfloat16
+FOR_EACH_BOOL_4(DECLARE_TEMPLATE_INSTANCE, nv_bfloat16)
+
+#undef DECLARE_TEMPLATE_INSTANCE
+#undef FOR_EACH_BOOL_4
+
+// clang-format on
+
+}  // namespace flat

{flashinfer_python-0.6.2 → flashinfer_python-0.6.3}/csrc/fp4_gemm_cutlass.jinja

@@ -26,6 +26,7 @@ INSTANTIATE_FP4_GEMM_KERNEL_LAUNCHER({{ type }}, {{ cta_m }}, {{ cta_n }}, {{ ct
 INSTANTIATE_FP4_GEMM_KERNEL_LAUNCHER({{ type }}, {{ cta_m }}, {{ cta_n }}, {{ cta_k }}, 2, 4, 1, _2SM)
 INSTANTIATE_FP4_GEMM_KERNEL_LAUNCHER({{ type }}, {{ cta_m }}, {{ cta_n }}, {{ cta_k }}, 4, 2, 1, _2SM)
 INSTANTIATE_FP4_GEMM_KERNEL_LAUNCHER({{ type }}, {{ cta_m }}, {{ cta_n }}, {{ cta_k }}, 4, 4, 1, _2SM)
+INSTANTIATE_FP4_GEMM_KERNEL_LAUNCHER({{ type }}, {{ cta_m }}, {{ cta_n }}, {{ cta_k }}, 4, 1, 1, _2SM)
 
 }  // namespace gemm
 }  // namespace flashinfer

flashinfer_python-0.6.3/csrc/fp4_gemm_cutlass_sm103.cu

@@ -0,0 +1,193 @@
+/*
+ * Copyright (c) 2020-2023, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include <cuda_fp16.h>
+
+#include <cstddef>
+#include <cstdint>
+#include <functional>
+#include <type_traits>
+#include <vector>
+
+#include "flashinfer/gemm/cutlass_gemm_configs.h"
+#include "flashinfer/gemm/fp4_gemm_cutlass.h"
+#include "flashinfer/gemm/fp4_gemm_cutlass_template_sm103.h"
+#include "tvm_ffi_utils.h"
+
+using flashinfer::gemm::ClusterShape;
+using flashinfer::gemm::CutlassFp4GemmRunner;
+using flashinfer::gemm::CutlassFp4GemmRunnerInterface;
+using flashinfer::gemm::CutlassGemmConfig;
+using flashinfer::gemm::CutlassTileConfigSM100;
+using flashinfer::gemm::EpilogueScheduleType;
+using flashinfer::gemm::FP4GemmType;
+using flashinfer::gemm::MainloopScheduleType;
+
+namespace flashinfer {
+namespace gemm {
+template class CutlassFp4GemmRunner<__nv_bfloat16, FP4GemmType::W4A4_NVFP4_NVFP4>;
+template class CutlassFp4GemmRunner<half, FP4GemmType::W4A4_NVFP4_NVFP4>;
+}  // namespace gemm
+}  // namespace flashinfer
+
+namespace torch_ext {
+
+namespace {
+
+CutlassGemmConfig getFp4GemmConfig(int64_t m, int64_t n, int64_t k, int64_t tactic) {
+  auto getCutlassFp4GemmConfigs = []() {
+    CutlassFp4GemmRunner<__nv_bfloat16, FP4GemmType::W4A4_NVFP4_NVFP4> gemmRunner;
+    return gemmRunner.getConfigs();
+  };
+  static std::vector<CutlassGemmConfig> globalConfigs = getCutlassFp4GemmConfigs();
+  TVM_FFI_ICHECK(tactic >= 0 && tactic < globalConfigs.size())
+      << "tactic must be between 0 and " << globalConfigs.size();
+  return globalConfigs[tactic];
+}
+
+template <typename T>
+void runGemm(TensorView out, TensorView mat1, TensorView mat2, TensorView mat1Scale,
+             TensorView mat2Scale, TensorView globalScale, int64_t m, int64_t n, int64_t k,
+             int64_t batch_count, CutlassGemmConfig const& gemmConfig,
+             TensorView workspace_buffer) {
+  CutlassFp4GemmRunner<T, FP4GemmType::W4A4_NVFP4_NVFP4> gemmRunner;
+
+  int64_t const required_workspace_size = gemmRunner.getWorkspaceSize(m, n, k, batch_count);
+  int64_t const provided_workspace_size =
+      workspace_buffer.numel() * get_element_size(workspace_buffer);
+
+  auto runKernel = [&](void* workspace) {
+    gemmRunner.gemm(out.data_ptr(), mat1.data_ptr(), mat2.data_ptr(), mat1Scale.data_ptr(),
+                    mat2Scale.data_ptr(), static_cast<float*>(globalScale.data_ptr()), m, n, k,
+                    batch_count, gemmConfig, reinterpret_cast<char*>(workspace),
+                    required_workspace_size, get_stream(mat1.device()));
+  };
+
+  if (provided_workspace_size < required_workspace_size) {
+    Tensor new_workspace =
+        alloc_tensor({required_workspace_size}, DLDataType{kDLInt, 8, 1}, mat1.device());
+    runKernel(new_workspace.data_ptr());
+  } else {
+    runKernel(workspace_buffer.data_ptr());
+  }
+}
+
+constexpr auto FLOAT4_E2M1X2 = dl_uint8;  // uint8_t
+constexpr auto SF_DTYPE = dl_uint8;       // uint8_t
+
+// mat1: [B, M, K / 2], FLOAT4_E2M1X2 or [B, M, K], FLOAT8_E4M3FN
+// mat2: [B, N, K / 2], FLOAT4_E2M1X2
+// out: [B, M, N], fp16/bf16/fp32
+// mat1Scale: ceil(M / 128) * 128 * ceil(K / sfVecSize / 4) * 4, SF_DTYPE (UE4M3 or UE8M0)
+// mat2Scale: ceil(N / 128) * 128 * ceil(K / sfVecSize / 4) * 4, SF_DTYPE (UE4M3 or UE8M0)
+// globalScale: [1], 1 / (((448 * 6) / mat1.abs().max()) * ((448 * 6) / mat2.abs().max()))
+// B = 1 for GEMM op as a special case
+void fp4_bmm_impl(TensorView mat1, TensorView mat2, TensorView mat1Scale, TensorView mat2Scale,
+                  TensorView globalScale, TensorView out, TensorView workspace_buffer,
+                  int64_t tactic) {
+  CHECK_INPUT_AND_TYPE(mat1, FLOAT4_E2M1X2);
+  CHECK_INPUT_AND_TYPE(mat2, FLOAT4_E2M1X2);
+
+  int mat2_k_scale = 1;
+
+  CHECK_INPUT_AND_TYPE(mat1Scale, SF_DTYPE);
+  CHECK_INPUT_AND_TYPE(mat2Scale, SF_DTYPE);
+
+  CHECK_INPUT_AND_TYPE(globalScale, dl_float32);
+
+  int64_t m, n, k, b;
+  if (mat1.ndim() == 2) {
+    TVM_FFI_ICHECK_EQ(mat2.ndim(), 2) << "mat2 must be a matrix";
+    TVM_FFI_ICHECK_EQ(mat1.size(1), mat2.size(1) * mat2_k_scale)
+        << "mat1 and mat2 shapes cannot be multiplied (" << mat1.size(0) << "x" << mat1.size(1)
+        << " and " << mat2.size(0) << "x" << mat2.size(1) << ")";
+    m = mat1.size(0);
+    n = mat2.size(0);
+    k = mat2.size(1) * 2;
+    b = 1;
+  } else if (mat1.ndim() == 3) {
+    TVM_FFI_ICHECK_EQ(mat2.ndim(), 3) << "mat2 must be a batch of matrices";
+    TVM_FFI_ICHECK_EQ(mat1.size(0), mat2.size(0)) << "mat1 and mat2 must have the same batch size ("
+                                                  << mat1.size(0) << " and " << mat2.size(0) << ")";
+    TVM_FFI_ICHECK_EQ(mat1.size(2), mat2.size(2) * mat2_k_scale)
+        << "mat1 and mat2 shapes cannot be multiplied (" << mat1.size(1) << "x" << mat1.size(2)
+        << " and " << mat2.size(1) << "x" << mat2.size(2) << ")";
+    m = mat1.size(1);
+    n = mat2.size(1);
+    k = mat2.size(2) * 2;
+    b = mat1.size(0);
+  } else {
+    TVM_FFI_LOG_AND_THROW(NotImplementedError) << "mat1 must be a matrix or a batch of matrices";
+  }
+
+  // No heuristic for now, we rely on the autotuner to select the best tactic.
+  if (tactic == -1) {
+    tactic = 0;
+  }
+  auto config = getFp4GemmConfig(m, n, k, tactic);
+
+  constexpr int alignment = 32;
+  TVM_FFI_ICHECK_EQ(k % alignment, 0)
+      << "Expected k to be divisible by " << alignment << ", but got mat1 shape: (" << mat1.size(0)
+      << "x" << mat1.size(1) << "), k: " << k << ".";
+  TVM_FFI_ICHECK_EQ(n % alignment, 0)
+      << "Expected n to be divisible by " << alignment << ", but got mat2 shape: (" << mat2.size(0)
+      << "x" << mat2.size(1) << ").";
+
+  // Validate out dimensions
+  std::vector<int64_t> out_shape =
+      mat1.ndim() == 2 ? std::vector<int64_t>{m, n} : std::vector<int64_t>{b, m, n};
+  TVM_FFI_ICHECK_EQ(out.ndim(), out_shape.size())
+      << "out must have " << out_shape.size() << " dimensions, but got " << out.ndim();
+  for (int i = 0; i < out_shape.size(); ++i) {
+    TVM_FFI_ICHECK_EQ(out.size(i), out_shape[i])
+        << "out shape mismatch at dimension " << i << ": expected " << out_shape[i] << ", got "
+        << out.size(i);
+  }
+
+  switch (encode_dlpack_dtype(out.dtype())) {
+    case float16_code:
+      runGemm<half>(out, mat1, mat2, mat1Scale, mat2Scale, globalScale, m, n, k, b, config,
+                    workspace_buffer);
+      break;
+    case bfloat16_code:
+      runGemm<__nv_bfloat16>(out, mat1, mat2, mat1Scale, mat2Scale, globalScale, m, n, k, b, config,
+                             workspace_buffer);
+      break;
+    default:
+      TVM_FFI_ICHECK(false) << "out_dtype must be one of fp16/bf16.";
+  }
+}
+
+}  // namespace
+
+void fp4_gemm(TensorView mat1, TensorView mat2, TensorView mat1Scale, TensorView mat2Scale,
+              TensorView globalScale, TensorView out, TensorView workspace_buffer, int64_t tactic) {
+  fp4_bmm_impl(mat1, mat2, mat1Scale, mat2Scale, globalScale, out, workspace_buffer, tactic);
+}
+
+int64_t fp4_gemm_tactic_num() {
+  auto getCutlassConfigs = []() {
+    CutlassFp4GemmRunner<__nv_bfloat16, FP4GemmType::W4A4_NVFP4_NVFP4> gemmRunner;
+    return gemmRunner.getConfigs();
+  };
+  static int64_t totalTactics = getCutlassConfigs().size();
+  return totalTactics;
+}
+
+}  // namespace torch_ext
+
+TVM_FFI_DLL_EXPORT_TYPED_FUNC(fp4_gemm, torch_ext::fp4_gemm);
+TVM_FFI_DLL_EXPORT_TYPED_FUNC(fp4_gemm_tactic_num, torch_ext::fp4_gemm_tactic_num);

flashinfer_python-0.6.3/csrc/fp4_gemm_cutlass_sm103.jinja

@@ -0,0 +1,32 @@
+/*
+ * Copyright (c) 2020-2023, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "flashinfer/gemm/fp4_gemm_cutlass_template_sm103.h"
+
+namespace flashinfer {
+namespace gemm {
+INSTANTIATE_FP4_ULTRA_GEMM_KERNEL_LAUNCHER({{ type }}, {{ cta_m }}, {{ cta_n }}, {{ cta_k }}, 1, 1, 1, _1SM_sm103)
+INSTANTIATE_FP4_ULTRA_GEMM_KERNEL_LAUNCHER({{ type }}, {{ cta_m }}, {{ cta_n }}, {{ cta_k }}, 1, 2, 1, _1SM_sm103)
+INSTANTIATE_FP4_ULTRA_GEMM_KERNEL_LAUNCHER({{ type }}, {{ cta_m }}, {{ cta_n }}, {{ cta_k }}, 1, 4, 1, _1SM_sm103)
+INSTANTIATE_FP4_ULTRA_GEMM_KERNEL_LAUNCHER({{ type }}, {{ cta_m }}, {{ cta_n }}, {{ cta_k }}, 2, 1, 1, _2SM_sm103)
+INSTANTIATE_FP4_ULTRA_GEMM_KERNEL_LAUNCHER({{ type }}, {{ cta_m }}, {{ cta_n }}, {{ cta_k }}, 2, 2, 1, _2SM_sm103)
+INSTANTIATE_FP4_ULTRA_GEMM_KERNEL_LAUNCHER({{ type }}, {{ cta_m }}, {{ cta_n }}, {{ cta_k }}, 2, 4, 1, _2SM_sm103)
+INSTANTIATE_FP4_ULTRA_GEMM_KERNEL_LAUNCHER({{ type }}, {{ cta_m }}, {{ cta_n }}, {{ cta_k }}, 4, 2, 1, _2SM_sm103)
+INSTANTIATE_FP4_ULTRA_GEMM_KERNEL_LAUNCHER({{ type }}, {{ cta_m }}, {{ cta_n }}, {{ cta_k }}, 4, 4, 1, _2SM_sm103)
+INSTANTIATE_FP4_ULTRA_GEMM_KERNEL_LAUNCHER({{ type }}, {{ cta_m }}, {{ cta_n }}, {{ cta_k }}, 4, 1, 1, _2SM_sm103)
+
+}  // namespace gemm
+}  // namespace flashinfer

{flashinfer_python-0.6.2 → flashinfer_python-0.6.3}/csrc/gdn_prefill_launcher.cu

@@ -25,7 +25,7 @@
 #include <iostream>
 #include <sstream>
 
-#include "flat/prefill/prefill_kernel.hpp"
+#include "flashinfer/flat/prefill/prefill_kernel.hpp"
 
 using tvm::ffi::Optional;
 using tvm::ffi::TensorView;
@@ -35,10 +35,10 @@ namespace flashinfer {
 
 void gdn_prefill_launcher(void* output, void* output_state, void* q, void* k, void* v,
                           void* input_state, void* alpha, void* beta, int64_t* cu_seqlens,
-                          int64_t num_seqs, int64_t num_q_heads, int64_t num_k_heads,
-                          int64_t num_v_heads, int64_t num_o_heads, int64_t head_size,
-                          int64_t packed_seq, float scale, int64_t sm_count, DLDataType dtype,
-                          cudaStream_t stream) {
+                          uint8_t* workspace_buffer, int64_t num_seqs, int64_t num_q_heads,
+                          int64_t num_k_heads, int64_t num_v_heads, int64_t num_o_heads,
+                          int64_t head_size, int64_t packed_seq, float scale, int64_t sm_count,
+                          DLDataType dtype, cudaStream_t stream) {
   DISPATCH_DLPACK_DTYPE_TO_CTYPE_FP16(dtype, DType, [&] {
     int dev_id;
     cudaGetDevice(&dev_id);
@@ -51,8 +51,8 @@ void gdn_prefill_launcher(void* output, void* output_state, void* q, void* k, vo
           stream, static_cast<DType*>(output), static_cast<float*>(output_state),
           static_cast<DType const*>(q), static_cast<DType const*>(k), static_cast<DType const*>(v),
           static_cast<float const*>(input_state), static_cast<float const*>(alpha),
-          static_cast<float const*>(beta), cu_seqlens, num_seqs, num_q_heads, num_k_heads,
-          num_v_heads, num_o_heads, head_size, packed_seq, scale, sm_count);
+          static_cast<float const*>(beta), cu_seqlens, workspace_buffer, num_seqs, num_q_heads,
+          num_k_heads, num_v_heads, num_o_heads, head_size, packed_seq, scale, sm_count);
       return true;
     } else {
       std::ostringstream err_msg;
@@ -70,7 +70,8 @@ void gdn_prefill_launcher(void* output, void* output_state, void* q, void* k, vo
 
 void gdn_prefill(TensorView output, TensorView output_state, TensorView q, TensorView k,
                  TensorView v, TensorView cu_seqlens, Optional<TensorView> input_state,
-                 Optional<TensorView> alpha, Optional<TensorView> beta, double scale) {
+                 Optional<TensorView> alpha, Optional<TensorView> beta, double scale,
+                 TensorView workspace_buffer) {
   int64_t num_seqs = cu_seqlens.size(0) - 1;
   int64_t packed_seq = q.size(0);
   int64_t head_size = q.size(2);
@@ -109,6 +110,7 @@ void gdn_prefill(TensorView output, TensorView output_state, TensorView q, Tenso
   CHECK_INPUT(k);
   CHECK_INPUT(v);
   CHECK_INPUT(cu_seqlens);
+  CHECK_INPUT(workspace_buffer);
 
   TVM_FFI_ICHECK(output.dtype() == dl_float16 || output.dtype() == dl_bfloat16);
   TVM_FFI_ICHECK_EQ(output_state.dtype(), dl_float32);
@@ -116,6 +118,7 @@ void gdn_prefill(TensorView output, TensorView output_state, TensorView q, Tenso
   TVM_FFI_ICHECK_EQ(output.dtype(), k.dtype());
   TVM_FFI_ICHECK_EQ(output.dtype(), v.dtype());
   TVM_FFI_ICHECK_EQ(cu_seqlens.dtype(), dl_int64);
+  TVM_FFI_ICHECK_EQ(workspace_buffer.dtype(), dl_uint8);
 
   TVM_FFI_ICHECK_EQ(packed_seq, k.size(0));
   TVM_FFI_ICHECK_EQ(packed_seq, v.size(0));
@@ -164,7 +167,8 @@ void gdn_prefill(TensorView output, TensorView output_state, TensorView q, Tenso
 
   gdn_prefill_launcher(output.data_ptr(), output_state.data_ptr(), q.data_ptr(), k.data_ptr(),
                        v.data_ptr(), input_state_ptr, alpha_ptr, beta_ptr,
-                       static_cast<int64_t*>(cu_seqlens.data_ptr()), num_seqs, num_q_heads,
+                       static_cast<int64_t*>(cu_seqlens.data_ptr()),
+                       static_cast<uint8_t*>(workspace_buffer.data_ptr()), num_seqs, num_q_heads,
                        num_k_heads, num_v_heads, num_o_heads, head_size, packed_seq,
                        static_cast<float>(scale), sm_count, q.dtype(), stream);
 }

flashinfer_python-0.6.3/csrc/gdn_prefill_sm90_kernel_inst.jinja

@@ -0,0 +1,37 @@
+/*
+ * Copyright (c) 2025 by FlashInfer team.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+// Auto-generated file for separate compilation of GDN prefill kernel variants.
+// Template parameters: dtype={{ dtype }}, is_gva={{ is_gva }}, needs_beta={{ needs_beta }},
+//                      needs_alpha={{ needs_alpha }}, init_state={{ init_state }}
+
+// CUDA type definitions for half and nv_bfloat16
+#include <cuda_bf16.h>
+#include <cuda_fp16.h>
+
+// Include the header which defines the function template
+// The header includes all necessary CUTLASS type definitions
+#include "flashinfer/flat/prefill/prefill_kernel_delta_rule_sm90.cuh"
+
+namespace flat {
+
+// Explicit template instantiation for launch_delta_rule_prefill_kernel_gbai
+// Parameter types must exactly match the extern template declaration in prefill_kernel_delta_rule_sm90_extern.inc
+template void launch_delta_rule_prefill_kernel_gbai<{{ is_gva }}, {{ needs_beta }}, {{ needs_alpha }}, {{ init_state }}, cutlass::arch::Sm90, {{ dtype }}, {{ dtype }}, float>(
+    cudaStream_t, {{ dtype }}*, float*, {{ dtype }} const*, {{ dtype }} const*, {{ dtype }} const*,
+    float const*, float const*, float const*, int64_t const*, uint8_t*, int32_t, int32_t,
+    int32_t, int32_t, int32_t, int32_t, int64_t, float, int32_t);
+
+}  // namespace flat

{flashinfer_python-0.6.2 → flashinfer_python-0.6.3}/csrc/nv_internal/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.cu

@@ -399,7 +399,11 @@ __global__ void moeA2ADispatchKernel(
 #if !DISABLE_SYNC_FOR_PROFILING
       uint32_t expected_value = *ptrs.flag_val;
 
+#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
       asm volatile("fence.release.sys;");
+#else
+      __threadfence_system();
+#endif
 #pragma unroll 1  // No unroll as one iter is typically enough
       for (int target_rank = lane_id; target_rank < ep_size; target_rank += warpSize) {
         uint32_t* flag_addr = &ptrs.completion_flags[target_rank][rank_id];

{flashinfer_python-0.6.2 → flashinfer_python-0.6.3}/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_template_dispatch.h

@@ -807,6 +807,24 @@ void MoeGemmRunner<T, WeightType, OutputType, ScaleBiasType>::dispatchToArch(
             hopper_inputs.isValid(),
             "Calling TMA warp specialized configuration with invalid hopper config");
 
+        // For NoSmem epilogue schedule, output N must be 256-bit aligned.
+        // No need to check isGatedActivation here - inputs.n is already set correctly:
+        //   fc1_out_size = is_gated_activation ? inter_size * 2 : inter_size
+        // So gated activations naturally have better alignment (2x the size).
+        // Also, FINALIZE fusion is not supported with NO_SMEM.
+        // This check is here so the autotuner can catch invalid tactics during profiling.
+        if (inputs.gemm_config.epilogue_schedule ==
+            cutlass_extensions::EpilogueScheduleType::NO_SMEM) {
+          TLLM_CHECK_WITH_INFO(
+              inputs.n % (256 / cutlass::sizeof_bits<OutputType>::value) == 0,
+              "Output N %ld does not meet minimum alignment requirements for NO_SMEM epilogue %d",
+              (long)inputs.n, (int)(256 / cutlass::sizeof_bits<OutputType>::value));
+          TLLM_CHECK_WITH_INFO(
+              inputs.gemm_config.epilogue_fusion_type !=
+                  cutlass_extensions::CutlassGemmConfig::EpilogueFusionType::FINALIZE,
+              "NO_SMEM epilogue schedule is not supported with FINALIZE fusion");
+        }
+
         // Select the appropriate fusion function
         auto select_function = [&]() {
           switch (hopper_inputs.fusion) {

{flashinfer_python-0.6.2 → flashinfer_python-0.6.3}/csrc/nv_internal/tensorrt_llm/kernels/quantization_utils.cuh

@@ -496,7 +496,10 @@ __device__ uint64_t cvt_warp_fp16_to_mxfp8(PackedVec<Type, CVT_ELTS_PER_THREAD>&
   SFValue = static_cast<float>(tmpSFVal);
   fp8SFVal = tmpSFVal.__x;
   // Get the output scale (reciprocal of the SFValue).
-  float outputScale = vecMax != 0.f ? reciprocal_approximate_ftz(SFValue) : 0.0f;
+  // Note: Check SFValue != 0 (not vecMax != 0) because E8M0 conversion can underflow
+  // very small vecMax values to zero. Using vecMax != 0 would cause division by zero
+  // (reciprocal of 0 = infinity), leading to NaN when multiplied with denormal inputs.
+  float outputScale = SFValue != 0.f ? reciprocal_approximate_ftz(SFValue) : 0.0f;
 
   if (SFout) {
     // Write the SF to global memory (STG.8).

{flashinfer_python-0.6.2/csrc/flat/prefill → flashinfer_python-0.6.3/csrc}/prefill_kernel_delta_rule_sm90.cu

@@ -15,7 +15,11 @@
  */
 #include <cuda_bf16.h>
 
-#include "prefill_kernel_delta_rule_sm90.cuh"
+#include "flashinfer/flat/prefill/prefill_kernel_delta_rule_sm90.cuh"
+
+// Extern template declarations prevent implicit instantiation here.
+// Explicit instantiations are in separate generated files for parallel compilation.
+#include "flat_prefill_kernel_delta_rule_sm90_extern.inc"
 
 namespace flat {
 
@@ -27,19 +31,20 @@ void launch_delta_rule_prefill_kernel(cudaStream_t stream, TO* output, TState* o
                                       TQKV const* q, TQKV const* k, TQKV const* v,
                                       TState const* input_state, float const* alpha,
                                       float const* beta, int64_t const* cu_seqlens,
-                                      int32_t num_seqs, int32_t num_q_heads, int32_t num_k_heads,
-                                      int32_t num_v_heads, int32_t num_o_heads, int32_t head_size,
-                                      int64_t total_seqlen, float scale, int32_t sm_count) {
+                                      uint8_t* workspace_buffer, int32_t num_seqs,
+                                      int32_t num_q_heads, int32_t num_k_heads, int32_t num_v_heads,
+                                      int32_t num_o_heads, int32_t head_size, int64_t total_seqlen,
+                                      float scale, int32_t sm_count) {
   bool is_gva = num_v_heads > num_q_heads;
   bool needs_beta = beta != nullptr;
   bool needs_alpha = alpha != nullptr;
   bool init_state = input_state != nullptr;
 
-#define LAUNCH(is_gva, needs_beta, needs_alpha, init_state)                                    \
-  launch_delta_rule_prefill_kernel_gbai<is_gva, needs_beta, needs_alpha, init_state, ArchTag>( \
-      stream, output, output_state, q, k, v, input_state, alpha, beta, cu_seqlens, num_seqs,   \
-      num_q_heads, num_k_heads, num_v_heads, num_o_heads, head_size, total_seqlen, scale,      \
-      sm_count);
+#define LAUNCH(is_gva, needs_beta, needs_alpha, init_state)                                      \
+  launch_delta_rule_prefill_kernel_gbai<is_gva, needs_beta, needs_alpha, init_state, ArchTag>(   \
+      stream, output, output_state, q, k, v, input_state, alpha, beta, cu_seqlens,               \
+      workspace_buffer, num_seqs, num_q_heads, num_k_heads, num_v_heads, num_o_heads, head_size, \
+      total_seqlen, scale, sm_count);
 
   if (init_state) {
     if (is_gva && needs_beta && needs_alpha) {
@@ -86,18 +91,21 @@ void launch_delta_rule_prefill_kernel(cudaStream_t stream, TO* output, TState* o
 #undef LAUNCH
 }
 
+// Explicit instantiations for the outer dispatch function only.
+// The inner launch_delta_rule_prefill_kernel_gbai instantiations are in separate files.
 template void launch_delta_rule_prefill_kernel<cutlass::arch::Sm90, half, half, float>(
     cudaStream_t stream, half* output, float* state, half const* q, half const* k, half const* v,
     float const* input_state, float const* alpha, float const* beta, int64_t const* cu_seqlens,
-    int32_t num_seqs, int32_t num_q_heads, int32_t num_k_heads, int32_t num_v_heads,
-    int32_t num_o_heads, int32_t head_size, int64_t total_seqlen, float scale, int32_t sm_count);
+    uint8_t* workspace_buffer, int32_t num_seqs, int32_t num_q_heads, int32_t num_k_heads,
+    int32_t num_v_heads, int32_t num_o_heads, int32_t head_size, int64_t total_seqlen, float scale,
+    int32_t sm_count);
 
 template void
 launch_delta_rule_prefill_kernel<cutlass::arch::Sm90, nv_bfloat16, nv_bfloat16, float>(
     cudaStream_t stream, nv_bfloat16* output, float* state, nv_bfloat16 const* q,
     nv_bfloat16 const* k, nv_bfloat16 const* v, float const* input_state, float const* alpha,
-    float const* beta, int64_t const* cu_seqlens, int32_t num_seqs, int32_t num_q_heads,
-    int32_t num_k_heads, int32_t num_v_heads, int32_t num_o_heads, int32_t head_size,
-    int64_t total_seqlen, float scale, int32_t sm_count);
+    float const* beta, int64_t const* cu_seqlens, uint8_t* workspace_buffer, int32_t num_seqs,
+    int32_t num_q_heads, int32_t num_k_heads, int32_t num_v_heads, int32_t num_o_heads,
+    int32_t head_size, int64_t total_seqlen, float scale, int32_t sm_count);
 
 }  // namespace flat