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

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

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

{flashinfer_python-0.2.14 → flashinfer_python-0.2.14.post1}/csrc/batch_attention.cu

@@ -111,7 +111,8 @@ void BatchPagedAttentionRun(at::Tensor float_workspace_buffer, at::Tensor int_wo
       DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE,
       AttentionVariant, PersistentParams, [&] {
         PersistentParams params[2];
-
+        IdType* len_kv_chunk =
+            GetPtrFromBaseOffset<IdType>(int_buffer_ptr, plan_info.len_kv_chunk_offset);
         for (int i = 0; i < 2; i++) {
           params[i].q = static_cast<DTypeQ*>(q.data_ptr());
           params[i].k = static_cast<DTypeKV*>(k_cache.data_ptr());
@@ -138,8 +139,7 @@ void BatchPagedAttentionRun(at::Tensor float_workspace_buffer, at::Tensor int_wo
               GetPtrFromBaseOffset<IdType>(int_buffer_ptr, plan_info.tasks[i].kv_head_idx_offset);
           params[i].work_indptr =
               GetPtrFromBaseOffset<IdType>(int_buffer_ptr, plan_info.tasks[i].work_indptr_offset);
-          params[i].len_kv_chunk =
-              GetPtrFromBaseOffset<IdType>(int_buffer_ptr, plan_info.tasks[i].len_kv_chunk_offset);
+          params[i].len_kv_chunk = len_kv_chunk + i;
 
           params[i].final_o = static_cast<DTypeO*>(o.data_ptr());
           params[i].final_lse =

{flashinfer_python-0.2.14 → flashinfer_python-0.2.14.post1}/csrc/fused_moe/cutlass_backend/cutlass_fused_moe_kernels.cuh

@@ -1431,13 +1431,14 @@ __global__ void expandInputRowsKernel(
   static_assert(BlockScalingType == TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::NONE ||
                     !PRE_QUANT_AWQ,
                 "AWQ and Block Scaling are mutually exclusive");
-#ifdef ENABLE_FP4
   constexpr bool is_mxfp8 =
       std::is_same_v<ExpandedActivationsType, __nv_fp8_e4m3> &&
       BlockScalingType == TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::MXFPX &&
       !PRE_QUANT_AWQ;
   constexpr bool is_mxfp8_input = is_mxfp8 && std::is_same_v<InputActivationsType, __nv_fp8_e4m3>;
   constexpr bool need_mxfp8_quant = is_mxfp8 && !is_mxfp8_input;
+
+#ifdef ENABLE_FP4
   constexpr bool is_nvfp4 =
       std::is_same_v<ExpandedActivationsType, __nv_fp4_e2m1> &&
       BlockScalingType == TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::NVFP4 &&
@@ -1445,9 +1446,6 @@ __global__ void expandInputRowsKernel(
   constexpr bool is_nvfp4_input = is_nvfp4 && std::is_same_v<InputActivationsType, __nv_fp4_e2m1>;
   constexpr bool need_nvfp4_quant = is_nvfp4 && !is_nvfp4_input;
 #else
-  constexpr bool is_mxfp8 = false;
-  constexpr bool is_mxfp8_input = false;
-  constexpr bool need_mxfp8_quant = false;
   constexpr bool is_nvfp4 = false;
   constexpr bool is_nvfp4_input = false;
   constexpr bool need_nvfp4_quant = false;
@@ -1755,7 +1753,6 @@ __global__ void finalizeMoeRoutingKernel(
     ScaleBiasType const* bias, float const* scales, int const* unpermuted_row_to_permuted_row,
     int const* token_selected_experts, int64_t const orig_cols, int64_t const experts_per_token,
     int const num_experts_per_node, int const start_expert_id) {
-  assert(orig_cols % 4 == 0);
   int64_t const original_row = blockIdx.x;
   int64_t const num_rows = gridDim.x;
   auto const offset = original_row * orig_cols;
@@ -1765,6 +1762,8 @@ __global__ void finalizeMoeRoutingKernel(
   constexpr int64_t FINALIZE_ELEM_PER_THREAD =
       128 / std::min(sizeof_bits<OutputType>::value, sizeof_bits<GemmOutputType>::value);
 
+  assert(orig_cols % FINALIZE_ELEM_PER_THREAD == 0);
+
   int64_t const start_offset = threadIdx.x;
   int64_t const stride = FINALIZE_THREADS_PER_BLOCK;
   int64_t const num_elems_in_col = orig_cols / FINALIZE_ELEM_PER_THREAD;
@@ -1795,6 +1794,11 @@ __global__ void finalizeMoeRoutingKernel(
       int64_t const expanded_original_row = original_row + k_idx * num_rows;
       int64_t const expanded_permuted_row = unpermuted_row_to_permuted_row[expanded_original_row];
 
+      int64_t expanded_rows = num_rows * experts_per_token;
+      if (expanded_permuted_row < 0 || expanded_permuted_row >= expanded_rows) {
+        continue;
+      }
+
       float const row_scale = (SCALE_MODE == ScaleMode::NO_SCALE) ? 1.f : scales[k_offset];
 
       auto const* expanded_permuted_rows_row_ptr =
@@ -1828,8 +1832,6 @@ __global__ void finalizeMoeRoutingNoFillingKernel(
     int const* permuted_row_to_unpermuted_row, int const* token_selected_experts,
     int64_t const* expert_first_token_offset, int64_t const num_rows, int64_t const orig_cols,
     int64_t const experts_per_token, int const num_experts_per_node, int const start_expert_id) {
-  assert(orig_cols % 4 == 0);
-
 #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
   asm volatile("griddepcontrol.wait;");
 #endif
@@ -1864,6 +1866,8 @@ __global__ void finalizeMoeRoutingNoFillingKernel(
     constexpr int64_t FINALIZE_ELEM_PER_THREAD =
         128 / std::min(sizeof_bits<OutputType>::value, sizeof_bits<GemmOutputType>::value);
 
+    assert(orig_cols % FINALIZE_ELEM_PER_THREAD == 0);
+
     int64_t const start_offset = threadIdx.x;
     int64_t const stride = FINALIZE_THREADS_PER_BLOCK;
     int64_t const num_elems_in_col = orig_cols / FINALIZE_ELEM_PER_THREAD;
@@ -1889,6 +1893,11 @@ __global__ void finalizeMoeRoutingNoFillingKernel(
 
         int64_t const expanded_permuted_row_from_k_idx =
             unpermuted_row_to_permuted_row[source_row + k_idx * num_rows];
+        int64_t valid_tokens = expert_first_token_offset[num_experts_per_node];
+        if (expanded_permuted_row_from_k_idx < 0 ||
+            expanded_permuted_row_from_k_idx >= valid_tokens) {
+          continue;
+        }
 
         float const row_scale = (SCALE_MODE == ScaleMode::NO_SCALE) ? 1.f : scales[k_offset];
 
@@ -4109,9 +4118,13 @@ CutlassMoeFCRunner<T, WeightType, OutputType, InputType, BackBoneType, Enable>::
 
       TmaWarpSpecializedGroupedGemmInput::MXFPXElementSF weight_block_scale_value_int{};
 #ifdef ENABLE_FP8
+#if (__CUDACC_VER_MAJOR__ * 10000 + __CUDACC_VER_MINOR__ * 100 >= 120800)
       __nv_fp8_e8m0 tmp;
       tmp.__x = __nv_cvt_float_to_e8m0(1.0f, __NV_SATFINITE, cudaRoundPosInf);
       std::memcpy(&weight_block_scale_value_int, &tmp, sizeof(tmp));
+#else
+      TLLM_CHECK_WITH_INFO(false, "WFP4AFP8 is not supported on CUDA ");
+#endif
 #endif
 
       auto act_sf_rows = std::min(expanded_num_rows, num_rows * num_experts_per_node);

{flashinfer_python-0.2.14 → flashinfer_python-0.2.14.post1}/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/cutlass_type_conversion.h

@@ -18,7 +18,9 @@
 
 #include <cuda_bf16.h>
 #include <cuda_fp16.h>
+#if (__CUDACC_VER_MAJOR__ * 10000 + __CUDACC_VER_MINOR__ * 100 >= 120800)
 #include <cuda_fp4.h>
+#endif
 #include <cuda_fp8.h>
 
 #include "cutlass/bfloat16.h"
@@ -91,11 +93,13 @@ struct TllmToCutlassTypeAdapter<__nv_fp8_e5m2> {
 #endif
 
 #if defined(ENABLE_FP4)
+#if (__CUDACC_VER_MAJOR__ * 10000 + __CUDACC_VER_MINOR__ * 100 >= 120800)
 template <>
 struct TllmToCutlassTypeAdapter<__nv_fp4_e2m1> {
   using type = cutlass::float_e2m1_t;
 };
 #endif
+#endif
 
 ///////////////////////////////////////////////////////////////////////////////////////////////////
 // Cutlass to Tllm
@@ -130,11 +134,13 @@ struct CutlassToTllmTypeAdapter<cutlass::float_e5m2_t> {
 #endif
 
 #if defined(ENABLE_FP4)
+#if (__CUDACC_VER_MAJOR__ * 10000 + __CUDACC_VER_MINOR__ * 100 >= 120800)
 template <>
 struct CutlassToTllmTypeAdapter<cutlass::float_e2m1_t> {
   using type = __nv_fp4_e2m1;
 };
 #endif
+#endif
 
 ///////////////////////////////////////////////////////////////////////////////////////////////////
 

{flashinfer_python-0.2.14 → flashinfer_python-0.2.14.post1}/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/include/moe_gemm_kernels.h

@@ -32,8 +32,10 @@
 #include "tensorrt_llm/cutlass_extensions/include/cutlass_extensions/gemm_configs.h"
 
 #ifdef ENABLE_FP4
+#if (__CUDACC_VER_MAJOR__ * 10000 + __CUDACC_VER_MINOR__ * 100 >= 120800)
 #include <cuda_fp4.h>
 #endif
+#endif
 
 namespace tensorrt_llm::kernels::cutlass_kernels {
 template <class T>
@@ -262,6 +264,7 @@ class MoeGemmRunner {
  public:
   MoeGemmRunner();
 
+#if defined(ENABLE_FP4)
 #if defined(ENABLE_BF16)
   static constexpr bool use_wfp4a16 = std::is_same_v<WeightType, __nv_fp4_e2m1> &&
                                       (std::is_same_v<T, half> || std::is_same_v<T, __nv_bfloat16>);
@@ -269,6 +272,9 @@ class MoeGemmRunner {
   static constexpr bool use_wfp4a16 =
       std::is_same_v<WeightType, __nv_fp4_e2m1> && std::is_same_v<T, half>;
 #endif
+#else
+  static constexpr bool use_wfp4a16 = false;
+#endif
 
 #if defined(ENABLE_FP8)
   static constexpr bool use_fp8 =

{flashinfer_python-0.2.14 → flashinfer_python-0.2.14.post1}/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/include/moe_kernels.h

@@ -527,6 +527,7 @@ class CutlassMoeFCRunner : public CutlassMoeFCRunnerInterface {
   using ScaleBiasType = BackBoneType;
   using Self = CutlassMoeFCRunner<T, WeightType, OutputType, InputType, BackBoneType>;
 
+#if defined(ENABLE_FP4)
 #if defined(ENABLE_BF16)
   static constexpr bool use_wfp4a16 = std::is_same_v<WeightType, __nv_fp4_e2m1> &&
                                       (std::is_same_v<T, half> || std::is_same_v<T, __nv_bfloat16>);
@@ -534,6 +535,9 @@ class CutlassMoeFCRunner : public CutlassMoeFCRunnerInterface {
   static constexpr bool use_wfp4a16 =
       std::is_same_v<WeightType, __nv_fp4_e2m1> && std::is_same_v<T, half>;
 #endif
+#else
+  static constexpr bool use_wfp4a16 = false;
+#endif
 
 #if defined(ENABLE_FP8)
   static constexpr bool use_fp8 =

{flashinfer_python-0.2.14 → flashinfer_python-0.2.14.post1}/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_kernels_bf16_fp4.cu

@@ -17,7 +17,7 @@
 #include "tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_template_dispatch.h"
 
 namespace tensorrt_llm::kernels::cutlass_kernels {
-#ifdef ENABLE_BF16
+#if defined(ENABLE_BF16) && defined(ENABLE_FP4)
 template class MoeGemmRunner<__nv_bfloat16, __nv_fp4_e2m1, __nv_bfloat16>;
 #endif
 }  // namespace tensorrt_llm::kernels::cutlass_kernels

{flashinfer_python-0.2.14 → flashinfer_python-0.2.14.post1}/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_kernels_fp16_fp4.cu

@@ -17,5 +17,7 @@
 #include "tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_template_dispatch.h"
 
 namespace tensorrt_llm::kernels::cutlass_kernels {
+#if defined(ENABLE_FP4)
 template class MoeGemmRunner<half, __nv_fp4_e2m1, half>;
-}
+#endif
+}  // namespace tensorrt_llm::kernels::cutlass_kernels

{flashinfer_python-0.2.14 → flashinfer_python-0.2.14.post1}/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_template_dispatch.h

@@ -96,7 +96,9 @@ struct genericMoeGemmKernelLauncher {
 
     static_assert(cutlass::platform::is_same<T, WeightType>::value ||
                   cutlass::platform::is_same<WeightType, uint8_t>::value ||
+#if defined(ENABLE_FP4)
                   cutlass::platform::is_same<WeightType, __nv_fp4_e2m1>::value ||
+#endif
                   cutlass::platform::is_same<WeightType, cutlass::uint4b_t>::value);
 
     static_assert(arch::kMinComputeCapability < 90,
@@ -675,13 +677,18 @@ void MoeGemmRunner<T, WeightType, OutputType, ScaleBiasType>::dispatchToArch(
                        "Hopper configuration provided for non-Hopper architecture");
 
   if (sm_ >= 75 && sm_ < 80) {
+#ifdef ENABLE_FP4
     if constexpr (!std::is_same_v<WeightType, __nv_fp4_e2m1>) {
       dispatchMoeGemmToCutlass<T, WeightType, ScaleBiasType, cutlass::arch::Sm75, EpilogueTag>(
           inputs, multi_processor_count_);
     } else {
       TLLM_THROW("FP4 data type is not supported on SM < 90");
     }
+#else
+    TLLM_THROW("FP4 data type is not supported on SM < 90");
+#endif
   } else if (sm_ >= 80 && sm_ < 90) {
+#ifdef ENABLE_FP4
     if constexpr (!std::is_same_v<WeightType, __nv_fp4_e2m1>) {
       if constexpr (use_fp8 || use_w4afp8) {
 #if defined(ENABLE_FP8)
@@ -700,6 +707,9 @@ void MoeGemmRunner<T, WeightType, OutputType, ScaleBiasType>::dispatchToArch(
     } else {
       TLLM_THROW("FP4 data type is not supported on SM < 90");
     }
+#else
+    TLLM_THROW("FP4 data type is not supported on SM < 90");
+#endif
   } else if (sm_ >= 90) {
     // For SM120+ FP8 MoE, redirect to SM89 (Ada) FP8 kernel implementations.
     if constexpr (use_fp8) {

{flashinfer_python-0.2.14 → flashinfer_python-0.2.14.post1}/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_template_dispatch_tma_ws.h

@@ -113,6 +113,7 @@ void dispatchMoeGemmSelectBiasTmaWarpSpecialized(TmaWarpSpecializedGroupedGemmIn
   }
 #endif
   else {
+#ifdef ENABLE_FP4
     auto getFunc = [&]() {
       if constexpr (std::is_same_v<T, __nv_fp8_e4m3> && std::is_same_v<WeightType, __nv_fp4_e2m1>) {
         TLLM_CHECK_WITH_INFO(hopper_input.fpX_block_scaling_type ==
@@ -131,6 +132,9 @@ void dispatchMoeGemmSelectBiasTmaWarpSpecialized(TmaWarpSpecializedGroupedGemmIn
       }
     };
     getFunc()(hopper_input, num_experts, multi_processor_count, stream, occupancy, workspace_size);
+#else
+    TLLM_THROW("FP4 data type is not supported on this architecture and CUDA version");
+#endif
   }
 }
 

{flashinfer_python-0.2.14 → flashinfer_python-0.2.14.post1}/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_template_dispatch_tma_ws_mixed_dtype.h

@@ -156,10 +156,16 @@ void sm90_dispatch_moe_mixed_dtype_gemm_to_cutlass(
   // We also only instantiate configs here where threadblockShapeM == warpShapeM since those usually
   // perform the best for mixed type gemms.
 
+#if defined(ENABLE_FP4)
   constexpr int Ntile = (std::is_same_v<WeightType, __nv_fp4_e2m1>) ? 64 : 128;
   constexpr int Ktile =
       (std::is_same_v<WeightType, __nv_fp4_e2m1>) ? 128 : 128 * PackedScalesNum / sizeof(T);
   TLLM_CHECK(sizeof(T) == (std::is_same_v<WeightType, __nv_fp4_e2m1>) ? 2 : 1);
+#else
+  constexpr int Ntile = 128;
+  constexpr int Ktile = 128 * PackedScalesNum / sizeof(T);
+  TLLM_CHECK(sizeof(T) == 2);
+#endif
 
   using _Ntile = Int<Ntile>;
   using _Ktile = Int<Ktile>;
@@ -240,7 +246,11 @@ void sm90_dispatch_moe_mixed_dtype_gemm_to_cutlass(
 template <typename T, typename WeightType, typename OutputType>
 size_t calcMaxWorkspaceSizeTmaWarpSpecializedMixedInput(int num_experts, int sm_count_) {
   size_t count = 0;
+#ifdef ENABLE_FP4
   constexpr int Ktile = (std::is_same_v<WeightType, __nv_fp4_e2m1>) ? 256 : 512;
+#else
+  constexpr int Ktile = 512;
+#endif
   using _Ktile = Int<Ktile>;
 
 #ifdef COMPILE_HOPPER_TMA_GROUPED_GEMMS

{flashinfer_python-0.2.14 → flashinfer_python-0.2.14.post1}/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_tma_warp_specialized_traits.h

@@ -68,8 +68,13 @@ constexpr bool isValidHopperMOESpecialisation() {
   return (cutlass::platform::is_same<T, WeightType>::value ||
           (cutlass::platform::is_same<cutlass::uint4b_t, WeightType>::value &&
            cutlass::platform::is_same<T, __nv_fp8_e4m3>::value) ||
+#ifdef ENABLE_FP4
           (cutlass::platform::is_same<__nv_fp4_e2m1, WeightType>::value &&
-           !cutlass::platform::is_same<T, __nv_fp8_e4m3>::value))
+           !cutlass::platform::is_same<T, __nv_fp8_e4m3>::value)
+#else
+          false
+#endif
+              )
 
 #ifdef ENABLE_FP4
          && !cutlass::platform::is_same<T, __nv_fp4_e2m1>::value

{flashinfer_python-0.2.14 → flashinfer_python-0.2.14.post1}/csrc/nv_internal/tensorrt_llm/thop/fp4Op.cpp

@@ -23,7 +23,9 @@
 #include <ATen/Tensor.h>
 #include <ATen/cuda/EmptyTensor.h>
 #include <cuda_fp16.h>
+#if (__CUDACC_VER_MAJOR__ * 10000 + __CUDACC_VER_MINOR__ * 100 >= 120800)
 #include <cuda_fp4.h>
+#endif
 #include <cuda_fp8.h>
 
 #include <cstdint>
@@ -367,14 +369,15 @@ at::Tensor mxfp4_dequantize_host(at::Tensor weight, at::Tensor scale, int64_t gr
   TORCH_CHECK(weight.size(1) * 2 == scale.size(1) * group_size,
               "weight and scale must have the same number of columns");
 
-  uint8_t* weight_packed_ptr = weight.data_ptr<uint8_t>();
-  __nv_fp8_e8m0* scale_ptr = reinterpret_cast<__nv_fp8_e8m0*>(scale.data_ptr<uint8_t>());
-
   int const n = weight.size(0);
   int const k = weight.size(1) * 2;
 
   at::Tensor dequant_weight =
       at::empty({n, k}, at::dtype(at::ScalarType::Float).device(at::kCPU).requires_grad(false));
+
+#if (__CUDACC_VER_MAJOR__ * 10000 + __CUDACC_VER_MINOR__ * 100 >= 120800)
+  uint8_t* weight_packed_ptr = weight.data_ptr<uint8_t>();
+  __nv_fp8_e8m0* scale_ptr = reinterpret_cast<__nv_fp8_e8m0*>(scale.data_ptr<uint8_t>());
   float* dequant_weight_ptr = dequant_weight.data_ptr<float>();
 
   float fp4_lut[] = {0.0, 0.5,  1.0,  1.5,  2.0,  3.0,  4.0,  6.0,
@@ -397,6 +400,10 @@ at::Tensor mxfp4_dequantize_host(at::Tensor weight, at::Tensor scale, int64_t gr
     dequant_weight_ptr[2 * packed_idx] = weight_low * scale_;
     dequant_weight_ptr[2 * packed_idx + 1] = weight_high * scale_;
   }
+#else
+  // NOTE(Zihao): __nv_fp8_e8m0 is not supported for CUDA < 12.8
+  TLLM_THROW("mxfp4_dequantize_host is not supported for CUDA < 12.8");
+#endif
 
   return dequant_weight;
 }

{flashinfer_python-0.2.14 → flashinfer_python-0.2.14.post1}/csrc/pytorch_extension_utils.h

@@ -33,8 +33,10 @@
 #endif
 
 #if defined(FLASHINFER_ENABLE_FP4_E2M1)
+#if (__CUDACC_VER_MAJOR__ * 10000 + __CUDACC_VER_MINOR__ * 100 >= 120800)
 #include <cuda_fp4.h>
 #endif
+#endif
 
 #ifndef FLASHINFER_EXT_MODULE_INITED
 #define FLASHINFER_EXT_MODULE_INITED

flashinfer_python-0.2.14.post1/flashinfer/_build_meta.py

@@ -0,0 +1 @@
+__version__ = '0.2.14.post1'

{flashinfer_python-0.2.14 → flashinfer_python-0.2.14.post1}/flashinfer/aot.py

@@ -11,7 +11,10 @@ from torch.utils.cpp_extension import _get_cuda_arch_flags
 
 from .activation import act_func_def_str, gen_act_and_mul_module
 from .cascade import gen_cascade_module
-from .fp4_quantization import gen_fp4_quantization_module
+from .fp4_quantization import (
+    gen_fp4_quantization_sm100_module,
+    gen_fp4_quantization_sm90_module,
+)
 from .fused_moe import (
     gen_cutlass_fused_moe_sm100_module,
     gen_cutlass_fused_moe_sm90_module,
@@ -332,11 +335,12 @@ def gen_all_modules(
 
     if add_moe:
         jit_specs.append(gen_gemm_module())
-        jit_specs.append(gen_fp4_quantization_module())
         if has_sm90:
             jit_specs.append(gen_gemm_sm90_module())
+            jit_specs.append(gen_fp4_quantization_sm90_module())
             jit_specs.append(gen_cutlass_fused_moe_sm90_module())
         if has_sm100:
+            jit_specs.append(gen_fp4_quantization_sm100_module())
             jit_specs.append(gen_cutlass_fused_moe_sm100_module())
             jit_specs.append(gen_gemm_sm100_module())
 

{flashinfer_python-0.2.14 → flashinfer_python-0.2.14.post1}/flashinfer/attention.py

@@ -45,7 +45,7 @@ class BatchAttention:
         self._kv_layout = kv_layout
 
         self.float_workspace_buffer = torch.empty(
-            256 * 1024 * 1024,
+            384 * 1024 * 1024,
             dtype=torch.uint8,
             device=torch.device(device),
         )
@@ -115,7 +115,6 @@ class BatchAttention:
         # No addtional buf allocated for CUDA graph tensor
         # Allocate outside FlashInfer
         self._kv_indices = kv_indices
-
         self._plan_info = self.module.plan(
             self.float_workspace_buffer,
             self.int_workspace_buffer,
@@ -154,8 +153,9 @@ class BatchAttention:
             out = torch.empty_like(q)
         if lse is None:
             # lse shape: [batch_size, num_qo_heads]
-            lse = torch.empty(q.shape[0], q.shape[1], device=q.device)
-
+            lse = torch.empty(
+                q.shape[0], q.shape[1], device=q.device, dtype=torch.float32
+            )
         head_dim_qk = q.shape[2]
         if self._sm_scale is None:
             self._sm_scale = 1.0 / math.sqrt(head_dim_qk)

{flashinfer_python-0.2.14 → flashinfer_python-0.2.14.post1}/flashinfer/fp4_quantization.py

@@ -17,13 +17,14 @@ limitations under the License.
 import functools
 from enum import Enum
 from types import SimpleNamespace
-from typing import Optional, Tuple
+from typing import List, Optional, Tuple
 
 import torch
 
 from .jit import JitSpec
 from .jit import env as jit_env
-from .jit import gen_jit_spec, sm100a_nvcc_flags
+from .jit import gen_jit_spec, sm100a_nvcc_flags, sm90a_nvcc_flags
+from .jit.cpp_ext import is_cuda_version_at_least
 from .utils import (
     device_support_pdl,
     get_shuffle_matrix_a_row_indices,
@@ -62,9 +63,17 @@ def _pad_scale_factors(
         ).contiguous()
 
 
-def gen_fp4_quantization_module() -> JitSpec:
+def gen_fp4_quantization_sm100_module() -> JitSpec:
+    return gen_fp4_quantization_module(sm100a_nvcc_flags, "100")
+
+
+def gen_fp4_quantization_sm90_module() -> JitSpec:
+    return gen_fp4_quantization_module(sm90a_nvcc_flags, "90")
+
+
+def gen_fp4_quantization_module(nvcc_flags: List[str], device_arch: str) -> JitSpec:
     return gen_jit_spec(
-        "fp4_quantization",
+        f"fp4_quantization_{device_arch}",
         [
             jit_env.FLASHINFER_CSRC_DIR
             / "nv_internal/tensorrt_llm/thop/fp4Quantize.cpp",
@@ -75,16 +84,16 @@ def gen_fp4_quantization_module() -> JitSpec:
             jit_env.FLASHINFER_CSRC_DIR / "nv_internal/cpp/common/stringUtils.cpp",
             jit_env.FLASHINFER_CSRC_DIR / "nv_internal/cpp/common/tllmException.cpp",
         ],
-        extra_cuda_cflags=sm100a_nvcc_flags
+        extra_cuda_cflags=nvcc_flags
         + [
             "-DENABLE_BF16",
             "-DENABLE_FP8",
-            "-DENABLE_FP4",
+            "-DENABLE_FP4" if is_cuda_version_at_least("12.8") else "",
         ],
         extra_cflags=[
             "-DENABLE_BF16",
             "-DENABLE_FP8",
-            "-DENABLE_FP4",
+            "-DENABLE_FP4" if is_cuda_version_at_least("12.8") else "",
         ],
         extra_include_paths=[
             jit_env.FLASHINFER_CSRC_DIR / "nv_internal",
@@ -94,8 +103,13 @@ def gen_fp4_quantization_module() -> JitSpec:
 
 
 @functools.cache
-def get_fp4_quantization_module():
-    module = gen_fp4_quantization_module().build_and_load()
+def get_fp4_quantization_module(backend: str = "100"):
+    if backend == "100":
+        module = gen_fp4_quantization_sm100_module().build_and_load()
+    elif backend == "90":
+        module = gen_fp4_quantization_sm90_module().build_and_load()
+    else:
+        raise ValueError(f"Invalid backend: {backend}")
 
     @register_custom_op(
         "flashinfer::fp4_quantize_sm100",
@@ -310,7 +324,7 @@ def fp4_quantize(
     assert input.shape[-1] % sf_vec_size == 0
     if enable_pdl is None:
         enable_pdl = device_support_pdl(input.device)
-    x_q, sf = get_fp4_quantization_module().fp4_quantize_sm100(
+    x_q, sf = get_fp4_quantization_module("100").fp4_quantize_sm100(
         input,
         global_scale,
         sf_vec_size,
@@ -346,7 +360,11 @@ def block_scale_interleave(unswizzled_sf: torch.Tensor) -> torch.Tensor:
     assert unswizzled_sf.dtype == torch.uint8, (
         f"Input dtype must be uint8, got {unswizzled_sf.dtype}"
     )
-    return get_fp4_quantization_module().block_scale_interleave_sm100(
+
+    major, minor = torch.cuda.get_device_capability()
+    device_arch = f"{major * 10 + minor}"
+
+    return get_fp4_quantization_module(device_arch).block_scale_interleave_sm100(
         unswizzled_sf,
     )
 
@@ -380,7 +398,11 @@ def e2m1_and_ufp8sf_scale_to_float(
         torch.Tensor: Dequantized float tensor of shape [M, K] with dtype float32.
 
     """
-    return get_fp4_quantization_module().e2m1_and_ufp8sf_scale_to_float_sm100(
+    major, minor = torch.cuda.get_device_capability()
+    device_arch = f"{major * 10 + minor}"
+    return get_fp4_quantization_module(
+        device_arch
+    ).e2m1_and_ufp8sf_scale_to_float_sm100(
         e2m1_tensor,
         ufp8_scale_tensor,
         global_scale_tensor,
@@ -547,7 +569,9 @@ def mxfp4_dequantize_host(
     Returns:
         torch.Tensor: Dequantized tensor of shape [M, K] with dtype float.
     """
-    return get_fp4_quantization_module().mxfp4_dequantize_host(
+    major, minor = torch.cuda.get_device_capability()
+    device_arch = f"{major * 10 + minor}"
+    return get_fp4_quantization_module(device_arch).mxfp4_dequantize_host(
         weight,
         scale,
         group_size,

{flashinfer_python-0.2.14 → flashinfer_python-0.2.14.post1}/flashinfer/fused_moe/core.py

@@ -32,6 +32,7 @@ from ..autotuner import (
 from ..jit import JitSpec
 from ..jit import env as jit_env
 from ..jit import gen_jit_spec, setup_cubin_loader, sm100a_nvcc_flags, sm90a_nvcc_flags
+from ..jit.cpp_ext import is_cuda_version_at_least
 from ..jit.cubin_loader import get_cubin
 from ..jit.cutlass_gemm.generate_kernels import generate_gemm_operations
 from ..utils import (
@@ -273,7 +274,7 @@ def gen_cutlass_fused_moe_sm90_module(use_fast_build: bool = False) -> JitSpec:
         "-DCOMPILE_HOPPER_TMA_GROUPED_GEMMS",
         "-DENABLE_BF16",
         "-DENABLE_FP8",
-        "-DENABLE_FP4",
+        "-DENABLE_FP4" if is_cuda_version_at_least("12.8") else "",
         "-DUSING_OSS_CUTLASS_MOE_GEMM",
     ]
     return gen_cutlass_fused_moe_module(nvcc_flags, "90", use_fast_build)

{flashinfer_python-0.2.14 → flashinfer_python-0.2.14.post1}/flashinfer/jit/cpp_ext.py

@@ -23,7 +23,7 @@ from . import env as jit_env
 
 
 @functools.cache
-def _get_cuda_version() -> Version:
+def get_cuda_version() -> Version:
     if CUDA_HOME is None:
         nvcc = "nvcc"
     else:
@@ -37,6 +37,10 @@ def _get_cuda_version() -> Version:
     return Version(matches[0])
 
 
+def is_cuda_version_at_least(version_str: str) -> bool:
+    return get_cuda_version() >= Version(version_str)
+
+
 def _get_glibcxx_abi_build_flags() -> List[str]:
     glibcxx_abi_cflags = [
         "-D_GLIBCXX_USE_CXX11_ABI=" + str(int(torch._C._GLIBCXX_USE_CXX11_ABI))
@@ -97,7 +101,7 @@ def generate_ninja_build_for_op(
         "--compiler-options=-fPIC",
         "--expt-relaxed-constexpr",
     ]
-    cuda_version = _get_cuda_version()
+    cuda_version = get_cuda_version()
     # enable -static-global-template-stub when cuda version >= 12.8
     if cuda_version >= Version("12.8"):
         cuda_cflags += [

{flashinfer_python-0.2.14 → flashinfer_python-0.2.14.post1}/flashinfer/jit/cutlass_gemm/generate_kernels.py

@@ -2,7 +2,22 @@ import enum
 import os
 from itertools import chain, product
 
-from .cutlass_library import *
+from .cutlass_library import (
+    enum_auto,
+    DataTypeNames,
+    DataTypeSize,
+    DataType,
+    DataTypeTag,
+    GemmKind,
+    GemmKindNames,
+    KernelScheduleType,
+    KernelScheduleTag,
+    KernelScheduleSuffixes,
+    EpilogueScheduleType,
+    EpilogueScheduleTag,
+    EpilogueScheduleSuffixes,
+)
+from ..cpp_ext import is_cuda_version_at_least
 
 
 ################################################################################
@@ -586,10 +601,20 @@ def generate_sm90_mixed_type_grouped_gemm_operations(is_arch_enabled):
         (DataType.e4m3, DataType.u4, DataType.f16, DataType.f16, DataType.f16),
         (DataType.e4m3, DataType.u4, DataType.bf16, DataType.bf16, DataType.bf16),
     ]
-    supported_dtypes_fp4 = [
-        (DataType.f16, DataType.e2m1, DataType.ue8m0, DataType.f16, DataType.f16),
-        (DataType.bf16, DataType.e2m1, DataType.ue8m0, DataType.bf16, DataType.bf16),
-    ]
+
+    if is_cuda_version_at_least("12.8"):
+        supported_dtypes_fp4 = [
+            (DataType.f16, DataType.e2m1, DataType.ue8m0, DataType.f16, DataType.f16),
+            (
+                DataType.bf16,
+                DataType.e2m1,
+                DataType.ue8m0,
+                DataType.bf16,
+                DataType.bf16,
+            ),
+        ]
+    else:
+        supported_dtypes_fp4 = []
 
     quant_ops = [TrtLlm_QuantOp.finegrained_scale_only]