bugfix: move cum_m calculation inside kernels (#1060)

Author: yzh119

csrc/group_gemm_groupwise_sm100.cu

@@ -51,9 +51,8 @@ using namespace flashinfer;
 void CutlassGroupGemmGroupwiseScaledSM100(at::Tensor int_workspace_buffer,
                                           at::Tensor float_workspace_buffer, at::Tensor A,
                                           at::Tensor B, at::Tensor SFA, at::Tensor SFB,
-                                          at::Tensor C, at::Tensor m_indptr, int64_t cum_m,
-                                          int64_t n, int64_t k, int64_t scale_granularity_m,
-                                          int64_t scale_granularity_n,
+                                          at::Tensor C, at::Tensor m_indptr, int64_t n, int64_t k,
+                                          int64_t scale_granularity_m, int64_t scale_granularity_n,
                                           int64_t scale_granularity_k) {
   const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device());
   auto stream = at::cuda::getCurrentCUDAStream();
@@ -72,8 +71,8 @@ void CutlassGroupGemmGroupwiseScaledSM100(at::Tensor int_workspace_buffer,
               float_workspace_buffer.element_size() * float_workspace_buffer.size(0),
               static_cast<cutlass_t_in*>(A.data_ptr()), static_cast<cutlass_t_in*>(B.data_ptr()),
               static_cast<float*>(SFA.data_ptr()), static_cast<float*>(SFB.data_ptr()),
-              static_cast<cutlass_t_out*>(C.data_ptr()), static_cast<int*>(m_indptr.data_ptr()),
-              cum_m, n, k, batch_size, stream);
+              static_cast<cutlass_t_out*>(C.data_ptr()), static_cast<int*>(m_indptr.data_ptr()), n,
+              k, batch_size, stream);
           return true;
         });
   });

csrc/group_gemm_sm100_pybind.cu

@@ -18,9 +18,9 @@
 void CutlassGroupGemmGroupwiseScaledSM100(at::Tensor int_workspace_buffer,
                                           at::Tensor float_workspace_buffer, at::Tensor A,
                                           at::Tensor B, at::Tensor SFA, at::Tensor SFB,
-                                          at::Tensor C, at::Tensor m_indptr, int64_t cum_m,
-                                          int64_t n, int64_t k, int64_t scale_granularity_m,
-                                          int64_t scale_granularity_n, int64_t scale_granularity_k);
+                                          at::Tensor C, at::Tensor m_indptr, int64_t n, int64_t k,
+                                          int64_t scale_granularity_m, int64_t scale_granularity_n,
+                                          int64_t scale_granularity_k);
 
 TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) {
   m.def("group_gemm_fp8_nt_groupwise", CutlassGroupGemmGroupwiseScaledSM100);

flashinfer/gemm.py

@@ -871,7 +871,6 @@ def group_gemm_fp8_nt_groupwise(
             b_scale,
             out,
             m_indptr,
-            m_indptr[-1],
             n,
             k,
             *scale_granularity_mnk,
@@ -914,7 +913,6 @@ def group_gemm_fp8_nt_groupwise(
         b_scale,
         padded_out,
         padded_m_indptr,
-        padded_m_indptr[-1],
         n,
         k,
         *scale_granularity_mnk,

include/flashinfer/gemm/group_gemm_groupwise_sm100.cuh

@@ -33,19 +33,19 @@ template <typename ScaleConfig, typename DTypeIn, typename DTypeSF, typename DTy
           typename ProblemShape, typename StrideA, typename StrideB, typename StrideC,
           typename LayoutSFA, typename LayoutSFB>
 __global__ void compute_sm100_cutlass_group_gemm_args(
-    DTypeIn* A, DTypeIn* B, DTypeSF* SFA, DTypeSF* SFB, DTypeOut* C, int* m_indptr, int cum_m,
-    int n, int k, int batch_size, int scale_granularity_m, int scale_granularity_n,
-    int scale_granularity_k, ProblemShape* problem_sizes, const DTypeIn** A_ptr,
-    const DTypeIn** B_ptr, const DTypeSF** SFA_ptr, const DTypeSF** SFB_ptr, const DTypeOut** C_ptr,
-    DTypeOut** D_ptr, StrideA* stride_A, StrideB* stride_B, StrideC* stride_C,
-    LayoutSFA* layout_SFA, LayoutSFB* layout_SFB) {
+    DTypeIn* A, DTypeIn* B, DTypeSF* SFA, DTypeSF* SFB, DTypeOut* C, int* m_indptr, int n, int k,
+    int batch_size, int scale_granularity_m, int scale_granularity_n, int scale_granularity_k,
+    ProblemShape* problem_sizes, const DTypeIn** A_ptr, const DTypeIn** B_ptr,
+    const DTypeSF** SFA_ptr, const DTypeSF** SFB_ptr, const DTypeOut** C_ptr, DTypeOut** D_ptr,
+    StrideA* stride_A, StrideB* stride_B, StrideC* stride_C, LayoutSFA* layout_SFA,
+    LayoutSFB* layout_SFB) {
   int i = blockIdx.x;
   int m = m_indptr[i + 1] - m_indptr[i];
   problem_sizes[i] = ProblemShape(m, n, k);
   stride_A[i] = cutlass::make_cute_packed_stride(StrideA{}, {m, k, 1});
   stride_B[i] = cutlass::make_cute_packed_stride(StrideB{}, {n, k, 1});
   stride_C[i] = cutlass::make_cute_packed_stride(StrideC{}, {m, n, 1});
-  layout_SFA[i] = ScaleConfig::tile_atom_to_shape_SFA(make_shape(cum_m, n, k, 1));
+  layout_SFA[i] = ScaleConfig::tile_atom_to_shape_SFA(make_shape(m_indptr[batch_size], n, k, 1));
   layout_SFB[i] = ScaleConfig::tile_atom_to_shape_SFB(make_shape(m, n, k, 1));
   A_ptr[i] = A + m_indptr[i] * k;
   B_ptr[i] = B + i * k * n;
@@ -61,8 +61,8 @@ cudaError_t CutlassGroupwiseScaledGroupGEMMSM100(void* int_buffer, size_t int_bu
                                                  void* float_buffer,
                                                  size_t float_buffer_size_in_bytes, DTypeIn* A,
                                                  DTypeIn* B, float* SFA, float* SFB, DTypeOut* C,
-                                                 int* m_indptr, int cum_m, int n, int k,
-                                                 int batch_size, cudaStream_t stream) {
+                                                 int* m_indptr, int n, int k, int batch_size,
+                                                 cudaStream_t stream) {
   using ProblemShape = cutlass::gemm::GroupProblemShape<Shape<int, int, int>>;  // <M,N,K> per group
 
   using ElementA = DTypeIn;                   // Element type for A matrix operand
@@ -159,7 +159,7 @@ cudaError_t CutlassGroupwiseScaledGroupGEMMSM100(void* int_buffer, size_t int_bu
                                                        "sm100_groupwise_group_gemm_layout_SFB");
 
   compute_sm100_cutlass_group_gemm_args<ScaleConfig><<<batch_size, 1, 0, stream>>>(
-      A, B, SFA, SFB, C, m_indptr, cum_m, n, k, batch_size, ScaleGranularityM, ScaleGranularityN,
+      A, B, SFA, SFB, C, m_indptr, n, k, batch_size, ScaleGranularityM, ScaleGranularityN,
       ScaleGranularityK, problem_sizes, A_ptr, B_ptr, SFA_ptr, SFB_ptr, C_ptr, D_ptr, stride_A,
       stride_B, stride_C, layout_SFA, layout_SFB);