[QST] Does splitK cause computation errors on shapes with very large k?

[umiswing]

Hello! I'm using splitK in kGemmSplitKParallel mode with examples/36. Some shapes have a very large k (e.g. 5x1e4). I set the spltk_slices to max(min(256, k/128), 1). I find the kernels with splitK give an error output. I have some questions:

1. Should I use splitK on shapes with very large k?
2. Do I need to call cudaDeviceSynchronize() after the splitK kernel?
3. How can I use streamK in examples/36?

hardware: a100
cuda version: cuda 11.7

[jackkosaian]

Hi, @umiswing.

A few questions:

1. What sort of error do you see? Can you list the error message? Or is there no error message, but the outputs of matrix multiplication are incorrect?
2. What GEMM sizes M and N are you using alongside K of 5x1e4 that results in the error you mention?
3. Have you added a reduction kernel after the split K GEMM? This is needed when using kGemmSplitKParallel with split_k_slices > 1. If you launch the reduction kernel in the same stream as the split-K GEMM kernel, you should not need to call cudaDeviceSynchronize() in between the split-K kernel and the reduction kernel.

[umiswing]

Hi, @jackkosaian

What sort of error do you see? Can you list the error message? Or is there no error message, but the outputs of matrix multiplication are incorrect?

Sorry for my poor english. There is no error message, the outputs of matrix multiplication are incorrect.

What GEMM sizes M and N are you using alongside K of 5x1e4 that results in the error you mention?

M and N ranges from [16, 32, 64, 128], just common channels in convolution.

Have you added a reduction kernel after the split K GEMM? This is needed when using kGemmSplitKParallel with split_k_slices > 1. If you launch the reduction kernel in the same stream as the split-K GEMM kernel, you should not need to call cudaDeviceSynchronize() in between the split-K kernel and the reduction kernel.

No, I have not added a reduction kernel. But does https://github.com/NVIDIA/cutlass/blob/master/include/cutlass/gemm/kernel/gemm_universal.h do the reduction for me? If not, how could I add a reduction kernel? Are there any codes I can refer to?

Finally, I am a bit confused about how to set the arguments when in kGemmSplitKParallel mode, does the following code launch the kernel correctly?

  template <typename Config>                                                                                         
  void launchKernel(float const alpha,                                                                               
                    float const beta,                                                                                
                    const GPUContext& dev_ctx,  // GPUContext is a custom class in my system.                                                                     
                    const float* const a,                                                                            
                    const float* const b,                                                                            
                    const float* const c,                                                                            
                    float* const d,                                                                                  
                    const int m,                                                                                     
                    const int n,                                                                                     
                    const int k,                                                                                     
                    const int32_t* a_indices,                                                                        
                    const int32_t* b_indices,                                                                        
                    const int32_t* c_d_indices) {                                                                    
    cutlass::gemm::GemmCoord problem_size_real({m, n, k});                                                           
    using Gemm = typename Config::Gemm;                                                                              
    int split_k_slices = std::max(std::min(256, k / 128), 1);                                                        
    typename Gemm::Arguments arguments{                                                                              
        Config::Mode,                                                                                                
        problem_size_real,                                                                                           
        split_k_slices,                                                                                              
        {alpha                                           
         beta},                                           
        a,                                                              
        b,                                                              
        c,                                                              
        d,                                                                    
        typename Gemm::Base::LayoutA().capacity(problem_size_real.mk()),                                             
        typename Gemm::Base::LayoutB().capacity(problem_size_real.kn()),                                             
        typename Gemm::Base::LayoutC().capacity(problem_size_real.mn()),                                             
        typename Gemm::Base::LayoutC().capacity(problem_size_real.mn()),                                             
        std::is_same<typename Gemm::Base::LayoutA,                                 
                     cutlass::layout::RowMajor>::value                                                               
            ? problem_size_real.k()                                                                                  
            : problem_size_real.m(),                                                                                 
        std::is_same<typename Gemm::Base::LayoutB,                                                                   
                     cutlass::layout::RowMajor>::value                                                               
            ? problem_size_real.n()                                                                                  
            : problem_size_real.k(),                                                                                 
        std::is_same<typename Gemm::Base::LayoutC,                                                                   
                     cutlass::layout::RowMajor>::value                                                               
            ? problem_size_real.n()                                                                                  
            : problem_size_real.m(),                                                                                 
        std::is_same<typename Gemm::Base::LayoutC,                                                                   
                     cutlass::layout::RowMajor>::value                                                               
            ? problem_size_real.n()                                                                                  
            : problem_size_real.m(),                                                                                 
        a_indices,                                                                                                   
        b_indices,                                                                                                   
        c_d_indices};                                                                                                
    size_t workspace_size = Gemm::get_workspace_size(arguments);                                                     
    cutlass::device_memory::allocation<uint8_t> workspace(workspace_size);                                           
    Gemm gemm_op;                                                                                                    
    cutlass::Status status = gemm_op.can_implement(arguments);                                                       
    GATHER_GEMM_SCATTER_CHECK(status);                                                                               
    status = gemm_op.initialize(arguments, workspace.get());                                                         
    GATHER_GEMM_SCATTER_CHECK(status);                                                                               
    gemm_op(dev_ctx.stream());                                                                                       
  }
template<cutlass::gemm::GemmUniversalMode Mode_ =                                                                    
             cutlass::gemm::GemmUniversalMode::kGemm>                                                                
struct cutlass_tensorop_s1688gemm_64x128_16x4_tn_align4 {                                                            
  using Gemm =                                                                                                       
    cutlass::gemm::device::GemmUniversal<                                                                            
      float,                                                                                                         
      cutlass::layout::ColumnMajor,                                                                                  
      float,                                                                                                         
      cutlass::layout::RowMajor,                                                                                     
      float,                                                                                                         
      cutlass::layout::RowMajor,                                                                                     
      float,                                                                                                         
      cutlass::arch::OpClassTensorOp,                                                                                
      cutlass::arch::Sm80,                                                                                           
      cutlass::gemm::GemmShape<64, 128, 16>,                                                                         
      cutlass::gemm::GemmShape<32, 64, 16>,                                                                          
      cutlass::gemm::GemmShape<16, 8, 8>,                                                                            
                                                                                                                     
    cutlass::epilogue::thread::LinearCombination<                                                                    
      float,                                                                                                         
      4,                                                                                                             
      float,                                                                                                         
      float                                                                                                          
    >                                                                                                                
,                                                                                                                    
      cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<8>,                                                 
      4,                                                                                                             
      4,                                                                                                             
      4,                                                                                                             
      cutlass::arch::OpMultiplyAddFastF32,                                                                           
      cutlass::ComplexTransform::kNone,                                                                              
      cutlass::ComplexTransform::kNone,
      true, // gather a                                                                                              
      true, // gather b                                                                                              
      false // scatter d                                                                                             
    >;                                                                                                               
  static const cutlass::gemm::GemmUniversalMode Mode = Mode_;                                                        
}; 
/* some code stuffs*/
launchKernel<cutlass_tensorop_s1688gemm_64x128_16x4_tn_align4<cutlass::gemm::GemmUniversalMode::kGemmSplitKParallel>>(/*some args*/);                                                                                    

By the way, is it possible to add streamK to examples/36?

[jackkosaian]

Hi, @umiswing. No worries. Thank you for clarifying!

GemmUniversal will only perform a reduction for you if you use GEMM mode kGemmSplitKSerial. For kGemmSplitKSerial, each threadblock writes out the outputs for the k slice that it worked over, and a second reduction kernel is needed to aggregate these partial outputs. You can see an example of this being used in example 30 (albeit, for a convolution, rather than a GEMM) or as part of GemmSplitKParallel here.

By the way, is it possible to add streamK to examples/36?

Yes, it should be possible to use stream K alongside gather/scatter GEMM. To enable this, you'll need to do some plumbing in kernel::GemmUniversalStreamK to add the gather/scatter pointer arguments and propagate them through iterators etc.

[umiswing]

Thanks for your reply @jackkosaian . Could stream K still be faster than split K alongside gather/scatter GEMM, with the shapes [m, n, k] where m and n rangs from [16, 32, 64, 128] and k ranges from 1e3~1e5?

[jackkosaian]

It's certainly possible that stream K with gather/scatter could be faster than split K with gather/scatter. In particular, stream K should help in allowing one to reduce the number of tile sizes that must be profiled to find the most-performant configuration of the kernel.

Perhaps a good first estimate would be to benchmark stream K vs. split K for a regular GEMM with the problem sizes you list. You could do this by using example 47.

[umiswing]

I will update this issue to share some code stuffs after finishing the code combining gather/scatter with stream K/splitK. Thank you! @jackkosaian

[umiswing]

Hi @jackkosaian ! Is anything wrong with the reduction in the following code? The outputs are still incorrect. I set the arguments according to example 47 and device::GemmSplitKParallel.
And according to kernel::GemmUniversal, the gather/scatter fusion kernel writes to ptr_D directly, not the workspace. Actually, it just allocates an empty workspace and illegal memory accessed occurs when ref_workspace point to the workspace. So I point ref_workspace to d in the following code.
Note that input a is in col_major, while b, c and d are in row_major. Also c and d actually point to the same matrix.

  template <typename Config>                                                                                         
  void launchKernel(float const alpha,                                                                               
                    float const beta,                                                                                
                    const GPUContext& dev_ctx,  // GPUContext is a custom class in my system.                                                                     
                    const float* const a,                                                                            
                    const float* const b,                                                                            
                    const float* const c,                                                                            
                    float* const d,                                                                                  
                    const int m,                                                                                     
                    const int n,                                                                                     
                    const int k,                                                                                     
                    const int32_t* a_indices,                                                                        
                    const int32_t* b_indices,                                                                        
                    const int32_t* c_d_indices) {                                                                    
    cutlass::gemm::GemmCoord problem_size_real({m, n, k});                                                           
    using Gemm = typename Config::Gemm;                                                                              
    int split_k_slices = std::max(std::min(256, k / 128), 1);                                                        
    typename Gemm::Arguments arguments{                                                                              
        Config::Mode,                                                                                                
        problem_size_real,                                                                                           
        split_k_slices,                                                                                              
        {alpha                                           
         beta},                                           
        a,                                                              
        b,                                                              
        c,                                                              
        d,                                                                    
        typename Gemm::Base::LayoutA().capacity(problem_size_real.mk()),                                             
        typename Gemm::Base::LayoutB().capacity(problem_size_real.kn()),                                             
        typename Gemm::Base::LayoutC().capacity(problem_size_real.mn()),                                             
        typename Gemm::Base::LayoutC().capacity(problem_size_real.mn()),                                             
        std::is_same<typename Gemm::Base::LayoutA,                                 
                     cutlass::layout::RowMajor>::value                                                               
            ? problem_size_real.k()                                                                                  
            : problem_size_real.m(),                                                                                 
        std::is_same<typename Gemm::Base::LayoutB,                                                                   
                     cutlass::layout::RowMajor>::value                                                               
            ? problem_size_real.n()                                                                                  
            : problem_size_real.k(),                                                                                 
        std::is_same<typename Gemm::Base::LayoutC,                                                                   
                     cutlass::layout::RowMajor>::value                                                               
            ? problem_size_real.n()                                                                                  
            : problem_size_real.m(),                                                                                 
        std::is_same<typename Gemm::Base::LayoutC,                                                                   
                     cutlass::layout::RowMajor>::value                                                               
            ? problem_size_real.n()                                                                                  
            : problem_size_real.m(),                                                                                 
        a_indices,                                                                                                   
        b_indices,                                                                                                   
        c_d_indices};                                                                                                
    size_t workspace_size = Gemm::get_workspace_size(arguments);                                                     
    cutlass::device_memory::allocation<uint8_t> workspace(workspace_size);                                           
    Gemm gemm_op;                                                                                                    
    cutlass::Status status = gemm_op.can_implement(arguments);                                                       
    GATHER_GEMM_SCATTER_CHECK(status);                                                                               
    status = gemm_op.initialize(arguments, workspace.get());                                                         
    GATHER_GEMM_SCATTER_CHECK(status);                                                                               
    gemm_op(dev_ctx.stream());
   if (Config::Mode ==                                                                                                                                      
        cutlass::gemm::GemmUniversalMode::kGemmSplitKParallel) {                                                                                             
      using ReductionOp = cutlass::reduction::thread::ReduceAdd<                                                                                             
          typename Gemm::ElementAccumulator,                                                                                                                 
          typename Gemm::EpilogueOutputOp::ElementAccumulator,                                                                                               
          Gemm::EpilogueOutputOp::kCount>;                                                                                                                   
                                                                                                                                                             
      using ReductionKernel = cutlass::reduction::kernel::ReduceSplitK<                                                                                      
          cutlass::MatrixShape<4, 32 * Gemm::EpilogueOutputOp::kCount>,                                                                                      
          typename Gemm::EpilogueOutputOp,                                                                                                                   
          ReductionOp>;                                                                                                                                      
      using ReductionDevice =                                                                                                                                
          typename cutlass::reduction::device::ReduceSplitK<ReductionKernel>;                                                                                
      ReductionDevice reduction_op;                                                                                                                          
      cutlass::TensorRef<typename Gemm::ElementAccumulator,                                                                                                  
                         cutlass::layout::RowMajor>                                                                                                          
          ref_workspace(reinterpret_cast<typename Gemm::ElementAccumulator*>(                                                                                
                            d),                                                                                                                
                        problem_size_real.n());                                                                                                              
      cutlass::TensorRef<typename Gemm::Base::ElementC,                                                                                                      
                         typename Gemm::Base::LayoutC>                                                                                                       
          ref_c(reinterpret_cast<typename Gemm::Base::ElementC* const>(d),                                                                                   
                problem_size_real.n());                                                                                                                      
      cutlass::TensorRef<typename Gemm::Base::ElementC,                                                                                                      
                         typename Gemm::Base::LayoutC>                                                                                                       
          ref_d(reinterpret_cast<typename Gemm::Base::ElementC* const>(d),                                                                                   
                problem_size_real.n());                                                                                                                      
      typename ReductionDevice::Arguments reduction_args(                                                                                                    
          problem_size_real.mn(),                                                                                                                            
          split_k_slices,                                                                                                                                    
          static_cast<size_t>(problem_size_real.m() * problem_size_real.n()),                                                                                
          ref_workspace,                                                                                                                                     
          ref_d,                                                                                                                                             
          ref_c,                                                                                                                                             
          {alpha, beta});                                                                                 
      status = reduction_op.initialize(reduction_args);                                                                                                      
      status = reduction_op(dev_ctx.stream());                                                                                                               
    }                                                                                                                                                        
  }

template<cutlass::gemm::GemmUniversalMode Mode_ =                                                                    
             cutlass::gemm::GemmUniversalMode::kGemm>                                                                
struct cutlass_tensorop_s1688gemm_64x128_16x4_tn_align4 {                                                            
  using Gemm =                                                                                                       
    cutlass::gemm::device::GemmUniversal<                                                                            
      float,                                                                                                         
      cutlass::layout::ColumnMajor,                                                                                  
      float,                                                                                                         
      cutlass::layout::RowMajor,                                                                                     
      float,                                                                                                         
      cutlass::layout::RowMajor,                                                                                     
      float,                                                                                                         
      cutlass::arch::OpClassTensorOp,                                                                                
      cutlass::arch::Sm80,                                                                                           
      cutlass::gemm::GemmShape<64, 128, 16>,                                                                         
      cutlass::gemm::GemmShape<32, 64, 16>,                                                                          
      cutlass::gemm::GemmShape<16, 8, 8>,                                                                            
                                                                                                                     
    cutlass::epilogue::thread::LinearCombination<                                                                    
      float,                                                                                                         
      4,                                                                                                             
      float,                                                                                                         
      float                                                                                                          
    >                                                                                                                
,                                                                                                                    
      cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<8>,                                                 
      4,                                                                                                             
      4,                                                                                                             
      4,                                                                                                             
      cutlass::arch::OpMultiplyAddFastF32,                                                                           
      cutlass::ComplexTransform::kNone,                                                                              
      cutlass::ComplexTransform::kNone,
      true, // gather a                                                                                              
      true, // gather b                                                                                              
      false // scatter d                                                                                             
    >;                                                                                                               
  static const cutlass::gemm::GemmUniversalMode Mode = Mode_;                                                        
}; 
/* some code stuffs*/
launchKernel<cutlass_tensorop_s1688gemm_64x128_16x4_tn_align4<cutlass::gemm::GemmUniversalMode::kGemmSplitKParallel>>(/*some args*/);

[jackkosaian]

Actually, it just allocates an empty workspace and illegal memory accessed occurs when ref_workspace point to the workspace

Do you mean that workspace_size gets returned as 0?

Assuming that workspace_size is not 0, could you try setting arguments.ptr_D = workspace.get(); before cutlass::Status status = gemm_op.can_implement(arguments); when you run your GEMM?

Then, you will want to set ref_workspace to have workspace.get() as its pointer, rather than d.

[umiswing]

Do you mean that workspace_size gets returned as 0?

Yes, workspace_size gets returned as 0 in size_t workspace_size = Gemm::get_workspace_size(arguments);. Setting arguments.ptr_D = workspace.get(); would lead to illegal memory access.

[jackkosaian]

Could you double check that split_k_slices or batch_size_D being passed into your GEMM are not being set to 0 by accident, and that themode is GemmUniversalMode::kGemmSplitKParallel?

Provided that mode is set to GemmUniversalMode::kGemmSplitKParallel, the workspace size should be calculated as follows (see here):

      workspace_bytes =
        sizeof(ElementC) *
        size_t(batch_stride_D) *
        size_t(grid_tiled_shape.k());

where grid_tiled_shape.k() == split_k_slices.

[hwu36]

check this gemm example for how to use serial and parallel split k useage. https://github.com/NVIDIA/cutlass/tree/master/examples/23_ampere_gemm_operand_reduction_fusion . Ignore the kernel fusion part.