flashinfer-python 0.6.0__tar.gz → 0.6.0rc1__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: flashinfer-python
-Version: 0.6.0
+Version: 0.6.0rc1
 Summary: FlashInfer: Kernel Library for LLM Serving
 Author: FlashInfer team
 License-Expression: Apache-2.0
@@ -14,7 +14,7 @@ Requires-Dist: einops
 Requires-Dist: ninja
 Requires-Dist: numpy
 Requires-Dist: nvidia-cudnn-frontend>=1.13.0
-Requires-Dist: nvidia-cutlass-dsl>=4.3.4
+Requires-Dist: nvidia-cutlass-dsl>=4.3.2
 Requires-Dist: nvidia-ml-py
 Requires-Dist: packaging>=24.2
 Requires-Dist: requests
@@ -142,27 +142,20 @@ pip install -U --pre flashinfer-cubin --index-url https://flashinfer.ai/whl/nigh
 pip install -U --pre flashinfer-jit-cache --index-url https://flashinfer.ai/whl/nightly/cu129
 ```
 
-### CLI Tools
+### Verify Installation
 
-FlashInfer provides several CLI commands for configuration, module management, and development:
+After installation, verify that FlashInfer is correctly installed and configured:
 
 ```bash
-# Verify installation and view configuration
 flashinfer show-config
-
-# List and inspect modules
-flashinfer list-modules
-flashinfer module-status
-
-# Manage artifacts and cache
-flashinfer download-cubin
-flashinfer clear-cache
-
-# For developers: generate compile_commands.json for IDE integration
-flashinfer export-compile-commands [output_path]
 ```
 
-For complete documentation, see the [CLI reference](https://docs.flashinfer.ai/cli.html).
+This command displays:
+
+- FlashInfer version and installed packages (flashinfer-python, flashinfer-cubin, flashinfer-jit-cache)
+- PyTorch and CUDA version information
+- Environment variables and artifact paths
+- Downloaded cubin status and module compilation status
 
 ### Trying it out
 

{flashinfer_python-0.6.0 → flashinfer_python-0.6.0rc1}/README.md

@@ -117,27 +117,20 @@ pip install -U --pre flashinfer-cubin --index-url https://flashinfer.ai/whl/nigh
 pip install -U --pre flashinfer-jit-cache --index-url https://flashinfer.ai/whl/nightly/cu129
 ```
 
-### CLI Tools
+### Verify Installation
 
-FlashInfer provides several CLI commands for configuration, module management, and development:
+After installation, verify that FlashInfer is correctly installed and configured:
 
 ```bash
-# Verify installation and view configuration
 flashinfer show-config
-
-# List and inspect modules
-flashinfer list-modules
-flashinfer module-status
-
-# Manage artifacts and cache
-flashinfer download-cubin
-flashinfer clear-cache
-
-# For developers: generate compile_commands.json for IDE integration
-flashinfer export-compile-commands [output_path]
 ```
 
-For complete documentation, see the [CLI reference](https://docs.flashinfer.ai/cli.html).
+This command displays:
+
+- FlashInfer version and installed packages (flashinfer-python, flashinfer-cubin, flashinfer-jit-cache)
+- PyTorch and CUDA version information
+- Environment variables and artifact paths
+- Downloaded cubin status and module compilation status
 
 ### Trying it out
 

{flashinfer_python-0.6.0 → flashinfer_python-0.6.0rc1}/csrc/flashinfer_norm_binding.cu

@@ -17,15 +17,9 @@
 
 void rmsnorm(TensorView out, TensorView input, TensorView weight, double eps, bool enable_pdl);
 
-void rmsnorm_quant(TensorView out, TensorView input, TensorView weight, double scale, double eps,
-                   bool enable_pdl);
-
 void fused_add_rmsnorm(TensorView input, TensorView residual, TensorView weight, double eps,
                        bool enable_pdl);
 
-void fused_add_rmsnorm_quant(TensorView output, TensorView input, TensorView residual,
-                             TensorView weight, double scale, double eps, bool enable_pdl);
-
 void gemma_rmsnorm(TensorView out, TensorView input, TensorView weight, double eps,
                    bool enable_pdl);
 
@@ -35,9 +29,7 @@ void gemma_fused_add_rmsnorm(TensorView input, TensorView residual, TensorView w
 void layernorm(Tensor out, Tensor input, Tensor gamma, Tensor beta, double eps);
 
 TVM_FFI_DLL_EXPORT_TYPED_FUNC(rmsnorm, rmsnorm);
-TVM_FFI_DLL_EXPORT_TYPED_FUNC(rmsnorm_quant, rmsnorm_quant);
 TVM_FFI_DLL_EXPORT_TYPED_FUNC(fused_add_rmsnorm, fused_add_rmsnorm);
-TVM_FFI_DLL_EXPORT_TYPED_FUNC(fused_add_rmsnorm_quant, fused_add_rmsnorm_quant);
 TVM_FFI_DLL_EXPORT_TYPED_FUNC(gemma_rmsnorm, gemma_rmsnorm);
 TVM_FFI_DLL_EXPORT_TYPED_FUNC(gemma_fused_add_rmsnorm, gemma_fused_add_rmsnorm);
 TVM_FFI_DLL_EXPORT_TYPED_FUNC(layernorm, layernorm);

flashinfer_python-0.6.0rc1/csrc/flashinfer_topk_binding.cu

@@ -0,0 +1,24 @@
+/*
+ * Copyright (c) 2024 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"
+
+using tvm::ffi::Optional;
+
+void radix_topk(TensorView input, TensorView output_indices, TensorView output_values,
+                Optional<TensorView> maybe_row_states_buffer, int64_t top_k);
+
+// Radix-based Top-K selection
+TVM_FFI_DLL_EXPORT_TYPED_FUNC(radix_topk, radix_topk);

{flashinfer_python-0.6.0 → flashinfer_python-0.6.0rc1}/csrc/norm.cu

@@ -77,43 +77,6 @@ void rmsnorm(TensorView output, TensorView input, TensorView weight, double eps,
   }
 }
 
-void rmsnorm_quant(TensorView output, TensorView input, TensorView weight, double scale, double eps,
-                   bool enable_pdl) {
-  CHECK_LAST_DIM_CONTIGUOUS_INPUT(input);
-  CHECK_LAST_DIM_CONTIGUOUS_INPUT(output);
-  CHECK_LAST_DIM_CONTIGUOUS_INPUT(weight);
-  CHECK_DEVICE(input, weight);
-  CHECK_DIM(1, weight);  // weight: (hidden_size)
-
-  auto input_ndim = input.ndim();
-  if (input_ndim == 2) {
-    // Normal RMSNorm: [batch_size, hidden_size]
-    // Use CTA parallelization for better parallelism
-    CHECK_DIM(2, output);
-    TVM_FFI_ICHECK_EQ(input.size(1), weight.size(0));
-    unsigned int batch_size = input.size(0);
-    unsigned int hidden_size = input.size(1);
-    TVM_FFI_ICHECK_EQ(output.size(0), batch_size);
-    TVM_FFI_ICHECK_EQ(output.size(1), hidden_size);
-    ffi::CUDADeviceGuard device_guard(input.device().device_id);
-    const cudaStream_t stream = get_stream(input.device());
-
-    DISPATCH_DLPACK_DTYPE_TO_CTYPE_FP16(input.dtype(), c_type, [&] {
-      return DISPATCH_DLPACK_DTYPE_TO_CTYPE_FP8(output.dtype(), o_type, [&] {
-        cudaError_t status = norm::RMSNormQuant(
-            static_cast<c_type*>(input.data_ptr()), static_cast<c_type*>(weight.data_ptr()),
-            static_cast<o_type*>(output.data_ptr()), batch_size, hidden_size, input.stride(0),
-            output.stride(0), static_cast<float>(scale), eps, enable_pdl, stream);
-        TVM_FFI_ICHECK(status == cudaSuccess)
-            << "RMSNormQuant failed with error code " << cudaGetErrorString(status);
-        return true;
-      });
-    });
-  } else {
-    TVM_FFI_ICHECK(false) << "Unsupported input dimension: " << input_ndim;
-  }
-}
-
 void fused_add_rmsnorm(TensorView input, TensorView residual, TensorView weight, double eps,
                        bool enable_pdl) {
   CHECK_LAST_DIM_CONTIGUOUS_INPUT(input);
@@ -144,42 +107,6 @@ void fused_add_rmsnorm(TensorView input, TensorView residual, TensorView weight,
   });
 }
 
-void fused_add_rmsnorm_quant(TensorView output, TensorView input, TensorView residual,
-                             TensorView weight, double scale, double eps, bool enable_pdl) {
-  CHECK_LAST_DIM_CONTIGUOUS_INPUT(input);
-  CHECK_LAST_DIM_CONTIGUOUS_INPUT(residual);
-  CHECK_LAST_DIM_CONTIGUOUS_INPUT(weight);
-  CHECK_LAST_DIM_CONTIGUOUS_INPUT(output);
-  CHECK_DEVICE(input, residual);
-  CHECK_DEVICE(input, weight);
-  CHECK_DEVICE(input, output);
-  CHECK_DIM(2, input);     // input: (batch_size, hidden_size)
-  CHECK_DIM(2, residual);  // residual: (batch_size, hidden_size)
-  CHECK_DIM(1, weight);    // weight: (hidden_size)
-  CHECK_DIM(2, output);
-  unsigned int batch_size = input.size(0);
-  unsigned int hidden_size = input.size(1);
-  TVM_FFI_ICHECK_EQ(residual.size(0), batch_size);
-  TVM_FFI_ICHECK_EQ(residual.size(1), hidden_size);
-  TVM_FFI_ICHECK_EQ(weight.size(0), hidden_size);
-  ffi::CUDADeviceGuard device_guard(input.device().device_id);
-  const cudaStream_t stream = get_stream(input.device());
-
-  DISPATCH_DLPACK_DTYPE_TO_CTYPE_FP16(input.dtype(), c_type, [&] {
-    return DISPATCH_DLPACK_DTYPE_TO_CTYPE_FP8(output.dtype(), o_type, [&] {
-      cudaError_t status = norm::FusedAddRMSNormQuant(
-          static_cast<c_type*>(input.data_ptr()), static_cast<c_type*>(residual.data_ptr()),
-          static_cast<c_type*>(weight.data_ptr()), static_cast<o_type*>(output.data_ptr()),
-          batch_size, hidden_size, input.stride(0), residual.stride(0), output.stride(0), scale,
-          eps, enable_pdl, stream);
-
-      TVM_FFI_ICHECK(status == cudaSuccess)
-          << "FusedAddRMSNormQuant failed with error code " << cudaGetErrorString(status);
-      return true;
-    });
-  });
-}
-
 void gemma_rmsnorm(TensorView output, TensorView input, TensorView weight, double eps,
                    bool enable_pdl) {
   CHECK_LAST_DIM_CONTIGUOUS_INPUT(input);

{flashinfer_python-0.6.0 → flashinfer_python-0.6.0rc1}/csrc/nv_internal/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.cu

@@ -294,97 +294,87 @@ __global__ void moeA2ADispatchKernel(
   int thread_idx = ThreadingPolicy::offset();
   int local_token_idx = ThreadingPolicy::token_idx();
 
-  if (local_num_tokens == 0) {
-    // Special case: If local_num_tokens == 0,
-    // we need to keep the threads where local_token_idx == 0 alive to participate in the
-    // synchronization. Other threads should return.
-    if (local_token_idx > 0) return;
+  if (local_token_idx >= local_num_tokens) {
+    return;
+  }
+
+  // Prepare per-policy shared-memory tiles for this token
+  extern __shared__ int smem[];
+  int* smem_topk_target_ranks;
+  int* smem_topk_send_indices;
+  int warps_per_block = blockDim.x / warpSize;
+  if constexpr (std::is_same<ThreadingPolicy, WarpPolicy>::value) {
+    int lane_id = threadIdx.x / warpSize;
+    smem_topk_target_ranks = smem + lane_id * TOP_K;
+    smem_topk_send_indices = smem + warps_per_block * TOP_K + lane_id * TOP_K;
   } else {
-    // Threads that do not have a token to process should return.
-    if (local_token_idx >= local_num_tokens) return;
-
-    // Prepare per-policy shared-memory tiles for this token
-    extern __shared__ int smem[];
-    int* smem_topk_target_ranks;
-    int* smem_topk_send_indices;
-    int warps_per_block = blockDim.x / warpSize;
-    if constexpr (std::is_same<ThreadingPolicy, WarpPolicy>::value) {
-      int lane_id = threadIdx.x / warpSize;
-      smem_topk_target_ranks = smem + lane_id * TOP_K;
-      smem_topk_send_indices = smem + warps_per_block * TOP_K + lane_id * TOP_K;
-    } else {
-      smem_topk_target_ranks = smem;
-      smem_topk_send_indices = smem + TOP_K;
-    }
+    smem_topk_target_ranks = smem;
+    smem_topk_send_indices = smem + TOP_K;
+  }
 
-    uint64_t already_copied = 0;
-    for (int k = 0; k < TOP_K; k++) {
-      int expert_id = token_selected_experts[local_token_idx * TOP_K + k];
-      // Use contiguous partitioning to determine target rank
-      int target_rank = compute_target_rank_id(expert_id, num_experts_per_rank);
-
-      if (already_copied & (1ULL << target_rank)) {
-        if (thread_idx == 0) {
-          ptrs.topk_target_ranks[local_token_idx * TOP_K + k] = -1;
-          ptrs.topk_send_indices[local_token_idx * TOP_K + k] = -1;
-          // Mirror to shared memory immediately
-          smem_topk_target_ranks[k] = -1;
-          smem_topk_send_indices[k] = -1;
-        }
-        continue;
-      }
+  uint64_t already_copied = 0;
+  for (int k = 0; k < TOP_K; k++) {
+    int expert_id = token_selected_experts[local_token_idx * TOP_K + k];
+    // Use contiguous partitioning to determine target rank
+    int target_rank = compute_target_rank_id(expert_id, num_experts_per_rank);
 
-      // Only one thread per warp should increment the counter
-      int dst_token_idx;
+    if (already_copied & (1ULL << target_rank)) {
       if (thread_idx == 0) {
-        dst_token_idx = atomicAdd(&ptrs.send_counters[target_rank], 1);
-
-        ptrs.topk_target_ranks[local_token_idx * TOP_K + k] = target_rank;
-        ptrs.topk_send_indices[local_token_idx * TOP_K + k] = dst_token_idx;
+        ptrs.topk_target_ranks[local_token_idx * TOP_K + k] = -1;
+        ptrs.topk_send_indices[local_token_idx * TOP_K + k] = -1;
         // Mirror to shared memory immediately
-        smem_topk_target_ranks[k] = target_rank;
-        smem_topk_send_indices[k] = dst_token_idx;
+        smem_topk_target_ranks[k] = -1;
+        smem_topk_send_indices[k] = -1;
       }
-      already_copied |= 1ULL << target_rank;
+      continue;
     }
-    // Sync before dispatching data
-    ThreadingPolicy::sync();
 
-    // Read staged routing once into registers per thread
-    int topk_target_ranks[TOP_K];
-    int topk_send_indices[TOP_K];
-#pragma unroll
-    for (int k = 0; k < TOP_K; ++k) {
-      topk_target_ranks[k] = smem_topk_target_ranks[k];
-      topk_send_indices[k] = smem_topk_send_indices[k];
+    // Only one thread per warp should increment the counter
+    int dst_token_idx;
+    if (thread_idx == 0) {
+      dst_token_idx = atomicAdd(&ptrs.send_counters[target_rank], 1);
+
+      ptrs.topk_target_ranks[local_token_idx * TOP_K + k] = target_rank;
+      ptrs.topk_send_indices[local_token_idx * TOP_K + k] = dst_token_idx;
+      // Mirror to shared memory immediately
+      smem_topk_target_ranks[k] = target_rank;
+      smem_topk_send_indices[k] = dst_token_idx;
     }
+    already_copied |= 1ULL << target_rank;
+  }
+  // Sync before dispatching data
+  ThreadingPolicy::sync();
 
-    // Perform a single source load and TOP_K fanout per payload
-    for (int payload_idx = 0; payload_idx < num_payloads; payload_idx++) {
-      uint8_t const* src_data = static_cast<uint8_t const*>(ptrs.src_data_ptrs[payload_idx]);
-      int bytes_per_token = ptrs.payload_bytes_per_token[payload_idx];
-      uint8_t const* src_ptr = src_data + local_token_idx * bytes_per_token;
+  // Read staged routing once into registers per thread
+  int topk_target_ranks[TOP_K];
+  int topk_send_indices[TOP_K];
+#pragma unroll
+  for (int k = 0; k < TOP_K; ++k) {
+    topk_target_ranks[k] = smem_topk_target_ranks[k];
+    topk_send_indices[k] = smem_topk_send_indices[k];
+  }
 
-      vectorized_dispatch<TOP_K, ThreadingPolicy>(src_ptr, bytes_per_token, rank_id,
-                                                  max_tokens_per_rank, payload_idx, ptrs,
-                                                  topk_target_ranks, topk_send_indices);
-    }
+  // Perform a single source load and TOP_K fanout per payload
+  for (int payload_idx = 0; payload_idx < num_payloads; payload_idx++) {
+    uint8_t const* src_data = static_cast<uint8_t const*>(ptrs.src_data_ptrs[payload_idx]);
+    int bytes_per_token = ptrs.payload_bytes_per_token[payload_idx];
+    uint8_t const* src_ptr = src_data + local_token_idx * bytes_per_token;
 
-    ThreadingPolicy::sync();
+    vectorized_dispatch<TOP_K, ThreadingPolicy>(src_ptr, bytes_per_token, rank_id,
+                                                max_tokens_per_rank, payload_idx, ptrs,
+                                                topk_target_ranks, topk_send_indices);
   }
 
+  ThreadingPolicy::sync();
+
   bool is_first_warp = threadIdx.x / warpSize == 0;
   if (is_first_warp) {
     int lane_id = threadIdx.x % warpSize;
 
     bool is_last_token = false;
     if (lane_id == 0) {
-      if (local_num_tokens != 0) {
-        int cnt = atomicAdd(ptrs.local_token_counter, 1);
-        is_last_token = cnt + 1 == local_num_tokens;
-      } else {
-        is_last_token = true;
-      }
+      int cnt = atomicAdd(ptrs.local_token_counter, 1);
+      is_last_token = cnt + 1 == local_num_tokens;
     }
     is_last_token = __shfl_sync(0xffffffff, is_last_token, 0);
 
@@ -456,7 +446,7 @@ void moe_a2a_dispatch_launch(MoeA2ADispatchParams const& params) {
   // Validate parameters
   TLLM_CHECK(params.top_k > 0 && params.top_k <= kMaxTopK);
   TLLM_CHECK(params.ep_size > 0 && params.ep_size <= kMaxRanks);
-  TLLM_CHECK(params.local_num_tokens >= 0);
+  TLLM_CHECK(params.local_num_tokens > 0);
   TLLM_CHECK(params.num_payloads > 0 && params.num_payloads <= kMaxPayloads);
 
   // Prepare kernel pointers struct
@@ -496,11 +486,6 @@ void moe_a2a_dispatch_launch(MoeA2ADispatchParams const& params) {
   // Configure kernel launch
   if (params.one_block_per_token) {
     int grid_size = params.local_num_tokens;
-    // If local_num_tokens is 0, we still need to launch a minimal kernel to participate in the
-    // synchronization.
-    if (grid_size == 0) {
-      grid_size = 1;
-    }
     int shared_bytes = 2 * params.top_k * (int)sizeof(int);
     SWITCH_TOP_K(params.top_k, TOP_K,
                  moeA2ADispatchKernel<BlockPolicy, TOP_K>
@@ -510,11 +495,6 @@ void moe_a2a_dispatch_launch(MoeA2ADispatchParams const& params) {
                      params.ep_size, params.num_experts_per_rank))
   } else {
     int grid_size = ceilDiv(params.local_num_tokens, kWarpsPerBlock);
-    // If local_num_tokens is 0, we still need to launch a minimal kernel to participate in the
-    // synchronization.
-    if (grid_size == 0) {
-      grid_size = 1;
-    }
     int shared_bytes = 2 * kWarpsPerBlock * params.top_k * (int)sizeof(int);
     SWITCH_TOP_K(params.top_k, TOP_K,
                  moeA2ADispatchKernel<WarpPolicy, TOP_K>
@@ -700,14 +680,8 @@ __global__ void moeA2ACombineKernel(
   int local_token_idx = ThreadingPolicy::token_idx();
   int const size_per_token = elements_per_token * sizeof(T);
 
-  if (local_num_tokens == 0) {
-    // Special case: If local_num_tokens == 0,
-    // we need to keep the threads where local_token_idx == 0 alive to participate in the
-    // synchronization. Other threads should return.
-    if (local_token_idx > 0) return;
-  } else {
-    // Threads that do not have a token to process should return.
-    if (local_token_idx >= local_num_tokens) return;
+  if (local_token_idx >= local_num_tokens) {
+    return;
   }
 
 #if !DISABLE_SYNC_FOR_PROFILING
@@ -763,8 +737,6 @@ __global__ void moeA2ACombineKernel(
   __syncthreads();
 #endif
 
-  if (local_num_tokens == 0) return;
-
   // Get output location for this token (using src_data_ptrs[0] as output)
   T* token_output = static_cast<T*>(ptrs.src_data_ptrs[0]) + local_token_idx * elements_per_token;
 
@@ -821,7 +793,7 @@ void moe_a2a_combine_launch(MoeA2ACombineParams const& params) {
   // Validate parameters
   TLLM_CHECK(params.top_k > 0 && params.top_k <= kMaxTopK);
   TLLM_CHECK(params.ep_size > 0 && params.ep_size <= kMaxRanks);
-  TLLM_CHECK(params.local_num_tokens >= 0);
+  TLLM_CHECK(params.local_num_tokens > 0);
   TLLM_CHECK(params.elements_per_token > 0);
 
   // Configure kernel launch
@@ -829,14 +801,6 @@ void moe_a2a_combine_launch(MoeA2ACombineParams const& params) {
   int const kWarpsPerBlock = kBlockSize / 32;  // warpSize
   int grid_size_warp = ceilDiv(params.local_num_tokens, kWarpsPerBlock);
   int grid_size_block = params.local_num_tokens;
-  // If local_num_tokens is 0, we still need to launch a minimal kernel to participate in the
-  // synchronization.
-  if (grid_size_warp == 0) {
-    grid_size_warp = 1;
-  }
-  if (grid_size_block == 0) {
-    grid_size_block = 1;
-  }
 
   // Prepare kernel pointers struct for combine
   CombineKernelPointers kernel_ptrs = {};  // Zero-initialize

{flashinfer_python-0.6.0 → flashinfer_python-0.6.0rc1}/csrc/rope.cu

@@ -547,77 +547,71 @@ void rope_quantize_append_paged_kv_cache(
 
   DISPATCH_DLPACK_DTYPE_TO_CTYPE_FP16(q_rope_in.dtype(), c_type, [&] {
     return DISPATCH_DLPACK_DTYPE_TO_CTYPE_FP8(q_rope_out.dtype(), c_quant_type, [&] {
-      return DISPATCH_DLPACK_IDTYPE_TO_CTYPE(pos_ids.dtype(), c_idtype, [&] {
-        cudaError_t status;
-
-        if (is_mla) {
-          // MLA: Construct paged_kv_mla_t struct
-          auto ckv_strides = ckv_cache.strides();
-          auto kpe_strides = kpe_cache.strides();
-
-          paged_kv_mla_t<c_quant_type, int32_t> paged_kv_mla(
-              page_size, no_rope_dim, rope_dim, batch_size,
-              static_cast<c_quant_type*>(ckv_cache.data_ptr()), ckv_strides.data(),
-              static_cast<c_quant_type*>(kpe_cache.data_ptr()), kpe_strides.data(),
-              static_cast<int32_t*>(kv_indices.data_ptr()),
-              static_cast<int32_t*>(kv_indptr.data_ptr()),
-              nullptr  // last_page_len not needed for this kernel
-          );
-
-          status = RopeQuantizeAppendPagedMLACache(
-              static_cast<c_type*>(q_rope_in.data_ptr()),
-              static_cast<c_type*>(k_rope_in.data_ptr()),
-              static_cast<c_type*>(q_nope_in.data_ptr()),
-              static_cast<c_type*>(k_nope_in.data_ptr()),
-              static_cast<c_quant_type*>(q_rope_out.data_ptr()),
-              static_cast<c_quant_type*>(q_nope_out.data_ptr()), paged_kv_mla,
-              static_cast<int32_t*>(batch_indices.data_ptr()),
-              static_cast<int32_t*>(positions.data_ptr()),
-              static_cast<float*>(cos_sin_cache.data_ptr()),
-              static_cast<c_idtype*>(pos_ids.data_ptr()), nnz, num_qo_heads, rope_dim, no_rope_dim,
-              q_rope_in_stride_n, q_rope_in_stride_h, q_nope_in_stride_n, q_nope_in_stride_h,
-              q_rope_out_stride_n, q_rope_out_stride_h, q_nope_out_stride_n, q_nope_out_stride_h,
-              k_rope_in_stride, k_nope_in_stride, quant_scale_q, quant_scale_kv, interleave,
-              enable_pdl, stream);
-
-        } else {
-          // GQA/MHA: Construct paged_kv_t struct
-          auto k_strides = k_cache.strides();
-          auto v_strides = v_cache.strides();
-          uint32_t head_dim = rope_dim + no_rope_dim;
-
-          paged_kv_t<c_quant_type, int32_t> paged_kv(
-              num_kv_heads, page_size, head_dim, batch_size, kv_layout,
-              static_cast<c_quant_type*>(k_cache.data_ptr()),
-              static_cast<c_quant_type*>(v_cache.data_ptr()), k_strides.data(),
-              static_cast<int32_t*>(kv_indices.data_ptr()),
-              static_cast<int32_t*>(kv_indptr.data_ptr()),
-              nullptr  // last_page_len not needed for this kernel
-          );
-
-          status = RopeQuantizeAppendPagedKVCache(
-              static_cast<c_type*>(q_rope_in.data_ptr()),
-              static_cast<c_type*>(k_rope_in.data_ptr()),
-              static_cast<c_type*>(q_nope_in.data_ptr()),
-              static_cast<c_type*>(k_nope_in.data_ptr()), static_cast<c_type*>(v_in.data_ptr()),
-              static_cast<c_quant_type*>(q_rope_out.data_ptr()),
-              static_cast<c_quant_type*>(q_nope_out.data_ptr()), paged_kv,
-              static_cast<int32_t*>(batch_indices.data_ptr()),
-              static_cast<int32_t*>(positions.data_ptr()),
-              static_cast<float*>(cos_sin_cache.data_ptr()),
-              static_cast<c_idtype*>(pos_ids.data_ptr()), nnz, num_qo_heads, num_kv_heads, rope_dim,
-              no_rope_dim, q_rope_in_stride_n, q_rope_in_stride_h, q_nope_in_stride_n,
-              q_nope_in_stride_h, q_rope_out_stride_n, q_rope_out_stride_h, q_nope_out_stride_n,
-              q_nope_out_stride_h, k_rope_in_stride, k_rope_in_stride_h, k_nope_in_stride,
-              k_nope_in_stride_h, v_in_stride, v_in_stride_h, quant_scale_q, quant_scale_kv,
-              interleave, enable_pdl, stream);
-        }
+      cudaError_t status;
+
+      if (is_mla) {
+        // MLA: Construct paged_kv_mla_t struct
+        auto ckv_strides = ckv_cache.strides();
+        auto kpe_strides = kpe_cache.strides();
+
+        paged_kv_mla_t<c_quant_type, int32_t> paged_kv_mla(
+            page_size, no_rope_dim, rope_dim, batch_size,
+            static_cast<c_quant_type*>(ckv_cache.data_ptr()), ckv_strides.data(),
+            static_cast<c_quant_type*>(kpe_cache.data_ptr()), kpe_strides.data(),
+            static_cast<int32_t*>(kv_indices.data_ptr()),
+            static_cast<int32_t*>(kv_indptr.data_ptr()),
+            nullptr  // last_page_len not needed for this kernel
+        );
+
+        status = RopeQuantizeAppendPagedMLACache(
+            static_cast<c_type*>(q_rope_in.data_ptr()), static_cast<c_type*>(k_rope_in.data_ptr()),
+            static_cast<c_type*>(q_nope_in.data_ptr()), static_cast<c_type*>(k_nope_in.data_ptr()),
+            static_cast<c_quant_type*>(q_rope_out.data_ptr()),
+            static_cast<c_quant_type*>(q_nope_out.data_ptr()), paged_kv_mla,
+            static_cast<int32_t*>(batch_indices.data_ptr()),
+            static_cast<int32_t*>(positions.data_ptr()),
+            static_cast<float*>(cos_sin_cache.data_ptr()),
+            static_cast<int32_t*>(pos_ids.data_ptr()), nnz, num_qo_heads, rope_dim, no_rope_dim,
+            q_rope_in_stride_n, q_rope_in_stride_h, q_nope_in_stride_n, q_nope_in_stride_h,
+            q_rope_out_stride_n, q_rope_out_stride_h, q_nope_out_stride_n, q_nope_out_stride_h,
+            k_rope_in_stride, k_nope_in_stride, quant_scale_q, quant_scale_kv, interleave,
+            enable_pdl, stream);
+
+      } else {
+        // GQA/MHA: Construct paged_kv_t struct
+        auto k_strides = k_cache.strides();
+        auto v_strides = v_cache.strides();
+        uint32_t head_dim = rope_dim + no_rope_dim;
+
+        paged_kv_t<c_quant_type, int32_t> paged_kv(
+            num_kv_heads, page_size, head_dim, batch_size, kv_layout,
+            static_cast<c_quant_type*>(k_cache.data_ptr()),
+            static_cast<c_quant_type*>(v_cache.data_ptr()), k_strides.data(),
+            static_cast<int32_t*>(kv_indices.data_ptr()),
+            static_cast<int32_t*>(kv_indptr.data_ptr()),
+            nullptr  // last_page_len not needed for this kernel
+        );
+
+        status = RopeQuantizeAppendPagedKVCache(
+            static_cast<c_type*>(q_rope_in.data_ptr()), static_cast<c_type*>(k_rope_in.data_ptr()),
+            static_cast<c_type*>(q_nope_in.data_ptr()), static_cast<c_type*>(k_nope_in.data_ptr()),
+            static_cast<c_type*>(v_in.data_ptr()),
+            static_cast<c_quant_type*>(q_rope_out.data_ptr()),
+            static_cast<c_quant_type*>(q_nope_out.data_ptr()), paged_kv,
+            static_cast<int32_t*>(batch_indices.data_ptr()),
+            static_cast<int32_t*>(positions.data_ptr()),
+            static_cast<float*>(cos_sin_cache.data_ptr()),
+            static_cast<int32_t*>(pos_ids.data_ptr()), nnz, num_qo_heads, num_kv_heads, rope_dim,
+            no_rope_dim, q_rope_in_stride_n, q_rope_in_stride_h, q_nope_in_stride_n,
+            q_nope_in_stride_h, q_rope_out_stride_n, q_rope_out_stride_h, q_nope_out_stride_n,
+            q_nope_out_stride_h, k_rope_in_stride, k_rope_in_stride_h, k_nope_in_stride,
+            k_nope_in_stride_h, v_in_stride, v_in_stride_h, quant_scale_q, quant_scale_kv,
+            interleave, enable_pdl, stream);
+      }
 
-        TVM_FFI_ICHECK(status == cudaSuccess)
-            << "RopeQuantizeAppendPagedKVCache failed with error code "
-            << cudaGetErrorString(status);
-        return true;
-      });
+      TVM_FFI_ICHECK(status == cudaSuccess)
+          << "RopeQuantizeAppendPagedKVCache failed with error code " << cudaGetErrorString(status);
+      return true;
     });
   });
 }