flashinfer-python 0.2.9rc1__tar.gz → 0.2.9rc2__tar.gz

{flashinfer_python-0.2.9rc1/flashinfer_python.egg-info → flashinfer_python-0.2.9rc2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: flashinfer-python
-Version: 0.2.9rc1
+Version: 0.2.9rc2
 Summary: FlashInfer: Kernel Library for LLM Serving
 Author: FlashInfer team
 Project-URL: Homepage, https://github.com/flashinfer-ai/flashinfer
@@ -16,7 +16,7 @@ Requires-Dist: pynvml
 Requires-Dist: einops
 Requires-Dist: nvidia-nvshmem-cu12
 Requires-Dist: nvidia-cudnn-cu12
-Requires-Dist: nvidia-cudnn-frontend
+Requires-Dist: nvidia-cudnn-frontend; platform_machine == "x86_64" or platform_machine == "AMD64"
 Dynamic: license-file
 Dynamic: requires-dist
 

{flashinfer_python-0.2.9rc1 → flashinfer_python-0.2.9rc2}/csrc/trtllm_batched_gemm_runner.cu

@@ -98,7 +98,8 @@ TrtllmGenBatchedGemmRunner::TrtllmGenBatchedGemmRunner(
         (!doesRouteImplUseNoRoute(options.mRouteImpl)) == mOptions.routeAct &&
         options.mFusedAct == mOptions.fusedAct && options.mIsStaticBatch == mOptions.staticBatch &&
         tileSize == mOptions.tileSize &&
-        options.mUseShuffledMatrixA == mOptions.useShuffledMatrixA) {
+        options.mUseShuffledMatrixA == mOptions.useShuffledMatrixA &&
+        options.mLayoutA == mOptions.weightLayout) {
       if (mOptions.transposeMmaOutput && options.mEpilogueTileM == mOptions.epilogueTileM) {
         mPassingConfigIndices.push_back(i);
       }

{flashinfer_python-0.2.9rc1 → flashinfer_python-0.2.9rc2}/csrc/trtllm_fused_moe_kernel_launcher.cu

@@ -46,7 +46,7 @@ namespace flashinfer {
 using tensorrt_llm::kernels::trtllmGenFp8BlockScaleMoe::Routing::RoutingMethodType;
 
 at::Tensor trtllm_fp8_per_tensor_scale_moe_launcher(
-    at::Tensor const& routing_logits, at::Tensor const& routing_bias,
+    at::Tensor const& routing_logits, at::optional<at::Tensor> routing_bias,
     at::Tensor const& hidden_states, at::Tensor const& gemm1_weights,
     at::Tensor const& output1_scales_scalar, at::Tensor const& output1_scales_gate_scalar,
     at::Tensor const& gemm2_weights, at::Tensor const& output2_scales_scalar,
@@ -77,10 +77,13 @@ at::Tensor trtllm_fp8_per_tensor_scale_moe_launcher(
   }
   TORCH_CHECK(routing_logits.dim() == 2, "routing_logits must be 2D.");
   TORCH_CHECK(routing_logits.sizes()[1] == num_experts, "routing_logits has incorrect shape.");
-  TORCH_CHECK(routing_bias.scalar_type() == at::ScalarType::BFloat16,
-              "routing_bias must be bfloat16.");
-  TORCH_CHECK(routing_bias.dim() == 1, "routing_bias must be 1D.");
-  TORCH_CHECK(routing_bias.sizes()[0] == num_experts, "routing_bias has incorrect shape.");
+  if (routing_bias.has_value()) {
+    TORCH_CHECK(routing_bias.value().scalar_type() == at::ScalarType::BFloat16,
+                "routing_bias must be bfloat16.");
+    TORCH_CHECK(routing_bias.value().dim() == 1, "routing_bias must be 1D.");
+    TORCH_CHECK(routing_bias.value().sizes()[0] == num_experts,
+                "routing_bias has incorrect shape.");
+  }
 
   if (n_group <= 0 || topk_group <= 0) {
     TORCH_CHECK(top_k == 1, "Current routing kernel (no groups) only supports top_k=1.");
@@ -114,7 +117,9 @@ at::Tensor trtllm_fp8_per_tensor_scale_moe_launcher(
   }
 
   args.routing_logits = routing_logits.data_ptr();
-  args.routing_bias = routing_bias.data_ptr();
+  auto const routing_bias_dtype =
+      routing_bias.has_value() ? routing_bias.value().scalar_type() : at::ScalarType::BFloat16;
+  args.routing_bias = routing_bias.has_value() ? routing_bias.value().data_ptr() : nullptr;
   args.hidden_states = hidden_states.data_ptr();
   args.gemm1_weights = gemm1_weights.data_ptr();
   args.output1_scales_scalar = output1_scales_scalar.data_ptr<float>();
@@ -145,9 +150,8 @@ at::Tensor trtllm_fp8_per_tensor_scale_moe_launcher(
       {args.num_tokens * args.top_k}, at::ScalarType::Int, routing_logits.device(), std::nullopt);
   at::Tensor permuted_idx_to_token_idx = at::detail::empty_cuda(
       {max_num_padded_tokens}, at::ScalarType::Int, routing_logits.device(), std::nullopt);
-  at::Tensor expert_weights =
-      at::detail::empty_cuda({args.num_tokens, args.top_k}, at::ScalarType::BFloat16,
-                             routing_logits.device(), std::nullopt);
+  at::Tensor expert_weights = at::detail::empty_cuda(
+      {args.num_tokens, args.top_k}, routing_bias_dtype, routing_logits.device(), std::nullopt);
   at::Tensor expert_indexes = at::detail::empty_cuda(
       {args.num_tokens, args.top_k}, at::ScalarType::Int, routing_logits.device(), std::nullopt);
   at::Tensor expert_count_histogram = at::detail::empty_cuda(
@@ -184,7 +188,7 @@ at::Tensor trtllm_fp8_per_tensor_scale_moe_launcher(
 
   tensorrt_llm::kernels::trtllmGenFp8BlockScaleMoe::Routing::Runner routing_runner(tile_tokens_dim);
   auto const& stream = at::cuda::getCurrentCUDAStream(routing_logits.get_device());
-  routing_runner.run(routing_logits.data_ptr(), routing_bias.data_ptr(), args.num_tokens,
+  routing_runner.run(routing_logits.data_ptr(), args.routing_bias, args.num_tokens,
                      args.num_experts, args.top_k, args.n_group, args.topk_group,
                      args.local_expert_offset, args.local_num_experts, args.routed_scaling_factor,
                      expert_indexes.data_ptr<int>(), expert_count_histogram.data_ptr<int>(),
@@ -287,7 +291,7 @@ at::Tensor trtllm_fp8_per_tensor_scale_moe_launcher(
 }
 
 at::Tensor trtllm_fp8_per_tensor_scale_moe(
-    at::Tensor routing_logits, at::Tensor routing_bias, at::Tensor hidden_states,
+    at::Tensor routing_logits, at::optional<at::Tensor> routing_bias, at::Tensor hidden_states,
     at::Tensor gemm1_weights, at::Tensor output1_scales_scalar,
     at::Tensor output1_scales_gate_scalar, at::Tensor gemm2_weights,
     at::Tensor output2_scales_scalar, int64_t num_experts, int64_t top_k, int64_t n_group,
@@ -308,7 +312,7 @@ at::Tensor trtllm_fp8_per_tensor_scale_moe(
 }
 
 at::Tensor trtllm_fp8_block_scale_moe_launcher(
-    at::Tensor const& routing_logits, at::Tensor const& routing_bias,
+    at::Tensor const& routing_logits, at::optional<at::Tensor> routing_bias,
     at::Tensor const& hidden_states, at::Tensor const& hidden_states_scale,
     at::Tensor const& gemm1_weights, at::Tensor const& gemm1_weights_scale,
     at::Tensor const& gemm2_weights, at::Tensor const& gemm2_weights_scale,
@@ -338,11 +342,14 @@ at::Tensor trtllm_fp8_block_scale_moe_launcher(
               "routing_logits and hidden_states must have the same number of tokens.");
   TORCH_CHECK(routing_logits.sizes()[1] == num_experts,
               "routing_logits dim1 must match num_experts.");
-  TORCH_CHECK(routing_bias.scalar_type() == at::ScalarType::BFloat16 ||
-                  routing_bias.scalar_type() == at::ScalarType::Float,
-              "routing_bias must be bfloat16 or float.");
-  TORCH_CHECK(routing_bias.dim() == 1, "routing_bias must be 1D.");
-  TORCH_CHECK(routing_bias.sizes()[0] == num_experts, "routing_bias has incorrect shape.");
+  if (routing_bias.has_value()) {
+    TORCH_CHECK(routing_bias.value().scalar_type() == at::ScalarType::BFloat16 ||
+                    routing_bias.value().scalar_type() == at::ScalarType::Float,
+                "routing_bias must be bfloat16 or float.");
+    TORCH_CHECK(routing_bias.value().dim() == 1, "routing_bias must be 1D.");
+    TORCH_CHECK(routing_bias.value().sizes()[0] == num_experts,
+                "routing_bias has incorrect shape.");
+  }
 
   if (n_group <= 0 || topk_group <= 0) {
     TORCH_CHECK(top_k == 1, "Current routing kernel (no groups) only supports top_k=1.");
@@ -375,11 +382,13 @@ at::Tensor trtllm_fp8_block_scale_moe_launcher(
     TORCH_CHECK(false, "Unsupported input dtype for MoE: ", dtype);
   }
 
-  args.mDtypeExpW = routing_bias.scalar_type() == at::ScalarType::BFloat16
+  auto const routing_bias_dtype =
+      routing_bias.has_value() ? routing_bias.value().scalar_type() : at::ScalarType::BFloat16;
+  args.mDtypeExpW = routing_bias_dtype == at::ScalarType::BFloat16
                         ? batchedGemm::trtllm::gen::Dtype::Bfloat16
                         : batchedGemm::trtllm::gen::Dtype::Fp32;
   args.routing_logits = routing_logits.data_ptr<float>();
-  args.routing_bias = routing_bias.data_ptr();
+  args.routing_bias = routing_bias.has_value() ? routing_bias.value().data_ptr() : nullptr;
   args.hidden_states = hidden_states.data_ptr();
   args.hidden_states_scale = hidden_states_scale.data_ptr<float>();
   args.gemm1_weights = gemm1_weights.data_ptr();
@@ -410,9 +419,8 @@ at::Tensor trtllm_fp8_block_scale_moe_launcher(
       {args.num_tokens * args.top_k}, at::ScalarType::Int, routing_logits.device(), std::nullopt);
   at::Tensor permuted_idx_to_token_idx = at::detail::empty_cuda(
       {max_num_padded_tokens}, at::ScalarType::Int, routing_logits.device(), std::nullopt);
-  at::Tensor expert_weights =
-      at::detail::empty_cuda({args.num_tokens, args.top_k}, routing_bias.scalar_type(),
-                             routing_logits.device(), std::nullopt);
+  at::Tensor expert_weights = at::detail::empty_cuda(
+      {args.num_tokens, args.top_k}, routing_bias_dtype, routing_logits.device(), std::nullopt);
   at::Tensor expert_indexes = at::detail::empty_cuda(
       {args.num_tokens, args.top_k}, at::ScalarType::Int, routing_logits.device(), std::nullopt);
   int64_t const size_of_expert_count_histogram = std::max(num_experts * 2, int64_t(256 * 2));
@@ -451,7 +459,7 @@ at::Tensor trtllm_fp8_block_scale_moe_launcher(
   tensorrt_llm::kernels::trtllmGenFp8BlockScaleMoe::Routing::Runner routing_runner(tile_tokens_dim);
   auto const& stream = at::cuda::getCurrentCUDAStream(routing_logits.get_device());
   routing_runner.run(
-      routing_logits.data_ptr<float>(), routing_bias.data_ptr(), args.num_tokens, args.num_experts,
+      routing_logits.data_ptr<float>(), args.routing_bias, args.num_tokens, args.num_experts,
       args.top_k, args.n_group, args.topk_group, args.local_expert_offset, args.local_num_experts,
       args.routed_scaling_factor, expert_indexes.data_ptr<int>(),
       expert_count_histogram.data_ptr<int>(), total_num_padded_tokens.data_ptr<int>(),
@@ -474,12 +482,26 @@ at::Tensor trtllm_fp8_block_scale_moe_launcher(
               "hidden_states_scale dim1 must match num_tokens.");
   TORCH_CHECK(gemm1_weights.scalar_type() == at::ScalarType::Float8_e4m3fn,
               "gemm1_weights must be fp8.");
-  TORCH_CHECK(gemm1_weights.dim() == 3, "gemm1_weights must be 3D.");
-  TORCH_CHECK(gemm1_weights.sizes()[1] % 2 == 0, "the second dimension of weights must be even.");
-  TORCH_CHECK(intermediate_size == gemm1_weights.sizes()[1] / 2,
-              "intermediate_size has incorrect shape.");
-  TORCH_CHECK(gemm1_weights.sizes()[2] == hidden_states.sizes()[1],
-              "the third dimension of weights must be equal to hidden_size.");
+
+  TORCH_CHECK(gemm1_weights.dim() == 3 || gemm1_weights.dim() == 4,
+              "gemm1_weights must be 3D or 4D.");
+  {
+    int64_t Mn = 0, K = 0;
+    if (gemm1_weights.dim() == 3) {
+      // MajorK [num_experts, M, K]
+      Mn = gemm1_weights.sizes()[1];
+      K = gemm1_weights.sizes()[2];
+    } else if (gemm1_weights.dim() == 4) {
+      // BlockMajorK [num_experts, K/block_k, M, block_k]
+      Mn = gemm1_weights.sizes()[2];
+      int64_t block_k = gemm1_weights.sizes()[3];
+      K = gemm1_weights.sizes()[1] * block_k;
+    }
+    TORCH_CHECK(Mn % 2 == 0, "the second dimension of weights must be even.");
+    TORCH_CHECK(intermediate_size == Mn / 2, "intermediate_size has incorrect shape.");
+    TORCH_CHECK(K == hidden_states.sizes()[1],
+                "the third dimension of weights must be equal to hidden_size.");
+  }
   TORCH_CHECK(gemm1_weights_scale.scalar_type() == at::ScalarType::Float,
               "gemm1_weights_scale must be float.");
   TORCH_CHECK(gemm1_weights_scale.dim() == 3, "gemm1_weights_scale must be 3D.");
@@ -494,9 +516,22 @@ at::Tensor trtllm_fp8_block_scale_moe_launcher(
               "gemm1_weights_scale has incorrect shape.");
   TORCH_CHECK(gemm2_weights.scalar_type() == at::ScalarType::Float8_e4m3fn,
               "gemm2_weights must be fp8.");
-  TORCH_CHECK(gemm2_weights.dim() == 3, "gemm2_weights must be 3D.");
-  TORCH_CHECK(gemm2_weights.sizes()[2] == intermediate_size,
-              "the third dimension of weights must be equal to intermediate_size.");
+
+  TORCH_CHECK(gemm2_weights.dim() == 3 || gemm2_weights.dim() == 4,
+              "gemm2_weights must be 3D or 4D.");
+  {
+    int64_t K = 0;
+    if (gemm2_weights.dim() == 3) {
+      // MajorK [num_experts, M, K]
+      K = gemm2_weights.sizes()[2];
+    } else if (gemm2_weights.dim() == 4) {
+      // BlockMajorK [num_experts, K/block_k, M, block_k]
+      int64_t block_k = gemm2_weights.sizes()[3];
+      K = gemm2_weights.sizes()[1] * block_k;
+    }
+    TORCH_CHECK(K == intermediate_size,
+                "the third dimension of weights must be equal to intermediate_size.");
+  }
   TORCH_CHECK(gemm2_weights_scale.scalar_type() == at::ScalarType::Float,
               "gemm2_weights_scale must be float.");
   TORCH_CHECK(gemm2_weights_scale.dim() == 3, "gemm2_weights_scale must be 3D.");
@@ -554,13 +589,14 @@ at::Tensor trtllm_fp8_block_scale_moe_launcher(
 }
 
 at::Tensor trtllm_fp8_block_scale_moe(
-    at::Tensor const& routing_logits, at::Tensor const& routing_bias,
+    at::Tensor const& routing_logits, at::optional<at::Tensor> routing_bias,
     at::Tensor const& hidden_states, at::Tensor const& hidden_states_scale,
     at::Tensor const& gemm1_weights, at::Tensor const& gemm1_weights_scale,
     at::Tensor const& gemm2_weights, at::Tensor const& gemm2_weights_scale, int64_t num_experts,
     int64_t top_k, int64_t n_group, int64_t topk_group, int64_t intermediate_size,
     int64_t local_expert_offset, int64_t local_num_experts, double routed_scaling_factor,
-    int64_t tile_tokens_dim, int64_t routing_method_type, bool use_shuffled_weight) {
+    int64_t tile_tokens_dim, int64_t routing_method_type, bool use_shuffled_weight,
+    int64_t weight_layout) {
   auto dtype = hidden_states.dtype();
   if (dtype == at::ScalarType::Half || dtype == at::ScalarType::BFloat16 ||
       dtype == at::ScalarType::Float8_e4m3fn) {
@@ -570,9 +606,13 @@ at::Tensor trtllm_fp8_block_scale_moe(
         batchedGemm::trtllm::gen::Dtype::E4m3};  // FP8 runner so hard-coded
     bool mUseDeepSeekFp8{true};                  // Always true for BlockScaleMoe
 
+    TORCH_CHECK(0 <= weight_layout && weight_layout <= 2,
+                "the value of weight_layout is not recognized");
+
     // Properly initialize the runner using make_unique like in the original code
-    auto mRunner = std::make_unique<RunnerType>(mDtypeElt, mUseDeepSeekFp8, tile_tokens_dim,
-                                                use_shuffled_weight);
+    auto mRunner = std::make_unique<RunnerType>(
+        mDtypeElt, mUseDeepSeekFp8, tile_tokens_dim, use_shuffled_weight,
+        static_cast<batchedGemm::gemm::MatrixLayout>(weight_layout));
 
     // Always use fallback config (equivalent to moeConfigIndex == -1 case from original code)
     auto const num_tokens = hidden_states.sizes()[0];
@@ -921,7 +961,8 @@ std::vector<at::Tensor> trtllm_fp4_block_scale_moe(
 
   // Properly initialize the runner using make_unique like in the original code
   auto mRunner = std::make_unique<RunnerType>(mDtypeElt, mUseDeepSeekFp8, tile_tokens_dim,
-                                              /*useShuffledMatrixA*/ true);
+                                              /*useShuffledMatrixA*/ true,
+                                              batchedGemm::gemm::MatrixLayout::MajorK);
 
   auto const num_tokens = hidden_states.sizes()[0];
 

{flashinfer_python-0.2.9rc1 → flashinfer_python-0.2.9rc2}/csrc/trtllm_fused_moe_runner.cu

@@ -174,10 +174,9 @@ void Runner::run(void* routingLogits, void* routingBias, int32_t numTokens, int3
 
 namespace PermuteGemm1 {
 
-tensorrt_llm::kernels::TrtllmGenBatchedGemmRunnerOptions getOptions(btg::Dtype dtypeElt,
-                                                                    int32_t tileTokensDim,
-                                                                    bool useDeepSeekFp8,
-                                                                    bool useShuffledMatrixA) {
+tensorrt_llm::kernels::TrtllmGenBatchedGemmRunnerOptions getOptions(
+    btg::Dtype dtypeElt, int32_t tileTokensDim, bool useDeepSeekFp8, bool useShuffledMatrixA,
+    batchedGemm::gemm::MatrixLayout weightLayout) {
   tensorrt_llm::kernels::TrtllmGenBatchedGemmRunnerOptions options = {
       .eltType = dtypeElt,
       .outputType = dtypeElt,
@@ -188,15 +187,17 @@ tensorrt_llm::kernels::TrtllmGenBatchedGemmRunnerOptions getOptions(btg::Dtype d
       .transposeMmaOutput = true,
       .tileSize = tileTokensDim,
       .epilogueTileM = useDeepSeekFp8 ? 64 : 128,
-      .useShuffledMatrixA = useShuffledMatrixA};
+      .useShuffledMatrixA = useShuffledMatrixA,
+      .weightLayout = weightLayout};
   return options;
 }
 
-Runner::Runner(btg::Dtype dtypeElt, bool useDeepSeekFp8, int tileTokensDim, bool useShuffledMatrixA)
+Runner::Runner(btg::Dtype dtypeElt, bool useDeepSeekFp8, int tileTokensDim, bool useShuffledMatrixA,
+               batchedGemm::gemm::MatrixLayout weightLayout)
     : mDtypeElt(dtypeElt),
       mTileTokensDim(tileTokensDim),
-      mRunner(tensorrt_llm::kernels::TrtllmGenBatchedGemmRunner(
-          getOptions(mDtypeElt, mTileTokensDim, useDeepSeekFp8, useShuffledMatrixA))) {}
+      mRunner(tensorrt_llm::kernels::TrtllmGenBatchedGemmRunner(getOptions(
+          mDtypeElt, mTileTokensDim, useDeepSeekFp8, useShuffledMatrixA, weightLayout))) {}
 
 void Runner::run(void* hiddenState, void* hiddenStateScale, void* weights, void* weightsScale,
                  void* expertWeights, float* outputScalesScalar, float* outputScalesGateScalar,
@@ -253,11 +254,9 @@ std::vector<int64_t> Runner::getPassingConfigIndices() const {
 }  // namespace PermuteGemm1
 
 namespace Gemm2 {
-tensorrt_llm::kernels::TrtllmGenBatchedGemmRunnerOptions getOptions(btg::Dtype dtypeElt,
-                                                                    btg::Dtype dtypeOut,
-                                                                    int32_t tileTokensDim,
-                                                                    bool useDeepSeekFp8,
-                                                                    bool useShuffledMatrixA) {
+tensorrt_llm::kernels::TrtllmGenBatchedGemmRunnerOptions getOptions(
+    btg::Dtype dtypeElt, btg::Dtype dtypeOut, int32_t tileTokensDim, bool useDeepSeekFp8,
+    bool useShuffledMatrixA, batchedGemm::gemm::MatrixLayout weightLayout) {
   tensorrt_llm::kernels::TrtllmGenBatchedGemmRunnerOptions options = {
       .eltType = dtypeElt,
       .outputType = dtypeOut,
@@ -268,17 +267,19 @@ tensorrt_llm::kernels::TrtllmGenBatchedGemmRunnerOptions getOptions(btg::Dtype d
       .transposeMmaOutput = true,
       .tileSize = tileTokensDim,
       .epilogueTileM = useDeepSeekFp8 ? 64 : 128,
-      .useShuffledMatrixA = useShuffledMatrixA};
+      .useShuffledMatrixA = useShuffledMatrixA,
+      .weightLayout = weightLayout};
   return options;
 }
 
 Runner::Runner(btg::Dtype dtypeElt, btg::Dtype outputDtype, bool useDeepSeekFp8, int tileTokensDim,
-               bool useShuffledMatrixA)
+               bool useShuffledMatrixA, batchedGemm::gemm::MatrixLayout weightLayout)
     : mDtypeElt(dtypeElt),
       mOutputDtype(outputDtype),
       mTileTokensDim(tileTokensDim),
-      mRunner(tensorrt_llm::kernels::TrtllmGenBatchedGemmRunner(getOptions(
-          mDtypeElt, mOutputDtype, mTileTokensDim, useDeepSeekFp8, useShuffledMatrixA))) {}
+      mRunner(tensorrt_llm::kernels::TrtllmGenBatchedGemmRunner(
+          getOptions(mDtypeElt, mOutputDtype, mTileTokensDim, useDeepSeekFp8, useShuffledMatrixA,
+                     weightLayout))) {}
 
 void Runner::run(void* permutedHiddenState, void* permutedHiddenStateScale, void* weights,
                  void* weightsScale, float* outputScalesScalar, void* output, void* outputScale,
@@ -336,11 +337,11 @@ std::vector<int64_t> Runner::getPassingConfigIndices() const {
 
 namespace MoE {
 Runner::Runner(btg::Dtype dtypeElt, bool useDeepSeekFp8, int32_t tileTokensDim,
-               bool useShuffledMatrixA)
-    : mPermuteGemm1(
-          PermuteGemm1::Runner(dtypeElt, useDeepSeekFp8, tileTokensDim, useShuffledMatrixA)),
+               bool useShuffledMatrixA, batchedGemm::gemm::MatrixLayout weightLayout)
+    : mPermuteGemm1(PermuteGemm1::Runner(dtypeElt, useDeepSeekFp8, tileTokensDim,
+                                         useShuffledMatrixA, weightLayout)),
       mGemm2(Gemm2::Runner(dtypeElt, btg::Dtype::Bfloat16, useDeepSeekFp8, tileTokensDim,
-                           useShuffledMatrixA)) {
+                           useShuffledMatrixA, weightLayout)) {
   auto const& gemm1PassingIndices = mPermuteGemm1.getPassingConfigIndices();
   auto const& gemm2PassingIndices = mGemm2.getPassingConfigIndices();
 

{flashinfer_python-0.2.9rc1 → flashinfer_python-0.2.9rc2}/csrc/trtllm_mnnvl_allreduce.cu

@@ -37,6 +37,8 @@ void trtllm_mnnvl_all_reduce(at::Tensor& in, int64_t multicast_buffer_ptr, int64
     int64_t token_dim = in.size(1);
 
     // Validate input parameters
+    TORCH_CHECK(token_dim % (sizeof(float2) / sizeof(c_type)) == 0,
+                "token_dim must be divisible by ", sizeof(float2) / sizeof(c_type));
     TORCH_CHECK(nranks >= 2 && nranks <= 64, "nranks must be between 2 and 64, got ", nranks);
     TORCH_CHECK(rank >= 0 && rank < nranks, "rank must be between 0 and nranks-1, got ", rank);
     TORCH_CHECK(out.has_value() || !wait_for_results,

{flashinfer_python-0.2.9rc1 → flashinfer_python-0.2.9rc2}/flashinfer/__init__.py

@@ -117,4 +117,7 @@ from .sampling import top_k_top_p_sampling_from_probs as top_k_top_p_sampling_fr
 from .sampling import top_p_renorm_probs as top_p_renorm_probs
 from .sampling import top_p_sampling_from_probs as top_p_sampling_from_probs
 from .sparse import BlockSparseAttentionWrapper as BlockSparseAttentionWrapper
+from .sparse import (
+    VariableBlockSparseAttentionWrapper as VariableBlockSparseAttentionWrapper,
+)
 from .utils import next_positive_power_of_2 as next_positive_power_of_2

flashinfer_python-0.2.9rc2/flashinfer/_build_meta.py

@@ -0,0 +1 @@
+__version__ = '0.2.9rc2'

{flashinfer_python-0.2.9rc1 → flashinfer_python-0.2.9rc2}/flashinfer/comm/mnnvl.py

@@ -18,7 +18,7 @@ import logging
 import platform
 import sys
 from dataclasses import dataclass
-from typing import List
+from typing import List, Optional
 
 import pynvml
 import torch
@@ -110,6 +110,26 @@ def test_cuda_memory_access(ptr: int, size: int, device_id: int) -> bool:
         return False
 
 
+def alloc_and_copy_to_cuda(host_ptr_array: List[int]) -> Optional[int]:
+    """
+    A helper function that allocates memory on cuda and copies the data from the host to the device.
+    """
+    if not host_ptr_array:
+        return None
+
+    ArrayType = ctypes.c_uint64 * len(host_ptr_array)
+    c_array = ArrayType(*host_ptr_array)
+    size_in_bytes = ctypes.sizeof(c_array)
+
+    device_ptr: cuda.CUdeviceptr = checkCudaErrors(cuda.cuMemAlloc(size_in_bytes))
+    checkCudaErrors(
+        cuda.cuMemcpyHtoD(device_ptr, ctypes.addressof(c_array), size_in_bytes)
+    )
+    # c_array should be freed by GC
+
+    return device_ptr
+
+
 class MpiComm:
     _comm: MPI.Intracomm = MPI.COMM_WORLD
 
@@ -423,7 +443,9 @@ class McastDeviceMemory:
         # CUDA memory handles and pointers
         self.mc_ptr = 0  # CUdeviceptr mMcPtr
         self.uc_ptrs: List[int] = []  # std::vector<CUdeviceptr> mUcPtrs
-        self.signal_pads_dev: List[int] = []  # std::vector<CUdeviceptr> mSignalPadsDev
+        self.signal_pads: List[int] = []  # mSignalPads
+        self.signal_pads_dev = 0  # std::vector<CUdeviceptr> mSignalPadsDev
+        self.uc_ptrs_dev = 0
         self.mc_handle = 0  # CUmemGenericAllocationHandle mMcHandle
         self.uc_handles: List[int] = (
             []
@@ -475,14 +497,18 @@ class McastDeviceMemory:
             raise NotImplementedError("Single-node NVLS allocation not implemented yet")
 
         # Initialize signal pads
-        self.signal_pads_dev = [0] * self.group_size
+        self.signal_pads = [0] * self.group_size
         for i in range(self.group_size):
-            self.signal_pads_dev[i] = self.uc_ptrs[i] + self.signal_pad_offset
+            self.signal_pads[i] = self.uc_ptrs[i] + self.signal_pad_offset
             if i == self.group_rank:
                 checkCudaErrors(
-                    cuda.cuMemsetD8(self.signal_pads_dev[i], 0, self.SIGNAL_PAD_SIZE)
+                    cuda.cuMemsetD8(self.signal_pads[i], 0, self.SIGNAL_PAD_SIZE)
                 )
 
+        # Create device pointers
+        self.signal_pads_dev = alloc_and_copy_to_cuda(self.signal_pads)
+        self.uc_ptrs_dev = alloc_and_copy_to_cuda(self.uc_ptrs)
+
     def __del__(self):
         """Destructor - cleanup allocated memory"""
 
@@ -505,6 +531,12 @@ class McastDeviceMemory:
             print(f"Destructor: CUDA context invalid, skipping cleanup: {e}")
             return
 
+        # Free device pointers
+        if self.signal_pads_dev:
+            checkCudaErrors(cuda.cuMemFree(self.signal_pads_dev))
+        if self.uc_ptrs_dev:
+            checkCudaErrors(cuda.cuMemFree(self.uc_ptrs_dev))
+
         # Unmap UC regions and release their handles
         if hasattr(self, "uc_handles") and self.uc_handles:
             for rank in range(self.group_size):
@@ -541,14 +573,22 @@ class McastDeviceMemory:
             except Exception as e:
                 print(f"Destructor: Failed to release MC handle: {e}")
 
-    def get_signal_pad_ptrs_dev(self) -> List[int]:
+    def get_signal_pad_ptrs_host(self) -> List[int]:
         """Get the raw array of signal pad pointers to all ranks (including self)"""
-        return self.signal_pads_dev
+        return self.signal_pads
 
-    def get_buffer_ptrs_dev(self) -> List[int]:
+    def get_buffer_ptrs_host(self) -> List[int]:
         """Get the raw array of unicast pointers to all ranks (including self)"""
         return self.uc_ptrs
 
+    def get_signal_pad_ptrs_dev(self) -> int:
+        """Get the raw array of signal pad pointers to all ranks (including self)"""
+        return self.signal_pads_dev
+
+    def get_buffer_ptrs_dev(self) -> int:
+        """Get the raw array of unicast pointers to all ranks (including self)"""
+        return self.uc_ptrs_dev
+
     def get_unicast_ptr(self, rank: int) -> int:
         """Get the raw unicast pointer to a given rank"""
         if rank >= len(self.uc_ptrs):
@@ -755,16 +795,8 @@ class McastDeviceMemory:
             )
         )
 
-    def get_multicast_ptr_as_int64(self) -> int:
-        """Get multicast pointer as int64 (legacy compatibility)"""
-        return self.get_multicast_ptr()
-
-    def get_buffer_ptrs_dev_as_int64(self) -> int:
-        """Get buffer pointers device as int64 (returning first UC pointer for now) (legacy compatibility)"""
-        return self.uc_ptrs[0] if self.uc_ptrs else 0
-
     def lamport_initialize(self, rank: int, dtype: torch.dtype):
-        if dtype == torch.bfloat16:
+        if dtype == torch.bfloat16 or dtype == torch.float16:
             neg_zero = 0x8000
             dsize = 2
             memset_func = cuda.cuMemsetD16
@@ -838,33 +870,6 @@ class McastGPUBuffer:
         """Get the raw multicast pointer"""
         return self.mcast_device_memory.get_multicast_ptr()
 
-    def get_multicast_ptr_as_int64(self) -> int:
-        """Get the multicast pointer as int64"""
-        return self.get_multicast_ptr()
-
-    def get_buffer_ptrs_dev(self) -> List[int]:
+    def get_buffer_ptrs_dev(self) -> int:
         """Get the buffer pointers device array"""
         return self.mcast_device_memory.get_buffer_ptrs_dev()
-
-    def get_buffer_ptrs_dev_as_int64(self) -> int:
-        """Get the buffer pointers device as int64 (returning first UC pointer)"""
-        ptrs = self.get_buffer_ptrs_dev()
-        assert ptrs is not None
-        return ptrs[0] if ptrs else 0
-
-    def get_buffer_ptrs_dev_as_ctypes_ptr(self) -> int:
-        """
-        Get buffer pointers as ctypes array pointer (equivalent to C++ void**).
-        Returns the address of a ctypes array that can be cast to int64_t and back to void**.
-
-        This matches the C++ pattern:
-        reinterpret_cast<int64_t>(reinterpret_cast<void**>(mUcPtrs.data()))
-        """
-        # Create ctypes array of void pointers
-        ArrayType = ctypes.c_void_p * len(self.mcast_device_memory.uc_ptrs)
-        self._buffer_ptrs_array = ArrayType(
-            *self.mcast_device_memory.uc_ptrs
-        )  # Keep reference to prevent GC
-
-        # Return the address of this array (equivalent to .data() in C++)
-        return ctypes.cast(self._buffer_ptrs_array, ctypes.c_void_p).value

{flashinfer_python-0.2.9rc1 → flashinfer_python-0.2.9rc2}/flashinfer/comm/trtllm_mnnvl_ar.py

@@ -184,9 +184,9 @@ def get_allreduce_mnnvl_workspace(
     mpi_barrier()
 
     # This is a buffer to maintain the state of this allreduce Op
-    # [Buffer_ptr, Clear_ptr, Buffer_size, atomic access counter]
+    # [Buffer_ptr, Clear_ptr, Buffer_size, num_tokens_prev, atomic access counter]
     buffer_flags = torch.tensor(
-        [0, 2, max_num_elements, 0],
+        [0, 2, max_num_elements, 0, 0],
         dtype=torch.uint32,
         device=torch.device("cuda", mapping.local_rank),
     )

{flashinfer_python-0.2.9rc1 → flashinfer_python-0.2.9rc2}/flashinfer/decode.py

@@ -1798,13 +1798,13 @@ class TrtllmGenDecodeModule:
         max_seq_len: int,
         bmm1_scale: float,  # todo(Yingyi): add dynamic scale tensor later
         bmm2_scale: float,
-        sm_count: int,
         window_left: int = -1,
         out: Optional[torch.Tensor] = None,
     ) -> torch.Tensor:
         if out is None:
             out = torch.empty_like(query)
-        sm_count = get_device_sm_count(query.device)
+        if self._sm_count is None:
+            self._sm_count = get_device_sm_count(query.device)
         self._op.trtllm_paged_attention_decode(
             out,
             query.unsqueeze(
@@ -1818,7 +1818,7 @@ class TrtllmGenDecodeModule:
             bmm1_scale,
             bmm2_scale,
             window_left,
-            sm_count,
+            self._sm_count,
         )
         return out
 
@@ -1826,6 +1826,7 @@ class TrtllmGenDecodeModule:
         pass
 
     def __init__(self):
+        self._sm_count: Optional[int] = None
         self._mod = trtllm_gen_fmha_module()
         self._op = self._mod.build_and_load()
         from flashinfer.jit.cubin_loader import (