tilelang.contrib.cutedsl.gemm_v2_sp¶

Sparse (2:4 structured) WGMMA helpers for the CuTeDSL backend.

CuTeDSL has no sparse WGMMA primitive, so the wgmma.mma_async.sp family is emitted as inline PTX here. The operand layout mirrors the cute::SM90::GMMA::SPARSE::GMMA_*_SS/_RS functors in cute/arch/mma_sm90_gmma_sparse.hpp: the accumulator registers are the tied D/C operands, e is the per-instruction 32-bit metadata register taken from the E fragment, and the sparse selector is an immediate.

Functions¶

wgmma_sp_ss(A_dtype, B_dtype, C_dtype, M, N, K, tnspA, ...)

Sparse WGMMA shared-shared: A and B come from shared memory descriptors.

wgmma_sp_rs(A_dtype, B_dtype, C_dtype, M, N, K, tnspB, ...)

Sparse WGMMA register-shared: A from registers, B from a descriptor.

Module Contents¶

tilelang.contrib.cutedsl.gemm_v2_sp.wgmma_sp_ss(A_dtype, B_dtype, C_dtype, M, N, K, tnspA, tnspB, scaleA, scaleB, spsel, desc_a, desc_b, C_ptr, scale_out, E_ptr, e_offset)¶

Sparse WGMMA shared-shared: A and B come from shared memory descriptors.

Parameters:

  • A_dtype (str)

  • B_dtype (str)

  • C_dtype (str)

  • M (cutlass.cutlass_dsl.Constexpr[int])

  • N (cutlass.cutlass_dsl.Constexpr[int])

  • K (cutlass.cutlass_dsl.Constexpr[int])

  • tnspA (cutlass.cutlass_dsl.Constexpr[bool])

  • tnspB (cutlass.cutlass_dsl.Constexpr[bool])

  • scaleA (cutlass.cutlass_dsl.Constexpr[int])

  • scaleB (cutlass.cutlass_dsl.Constexpr[int])

  • spsel (cutlass.cutlass_dsl.Constexpr[int])

  • desc_a (tilelang.contrib.cutedsl.gemm_v2.GmmaDescriptor)

  • desc_b (tilelang.contrib.cutedsl.gemm_v2.GmmaDescriptor)

  • C_ptr (cutlass.cute.Pointer)

  • scale_out (int)

  • E_ptr (cutlass.cute.Pointer)

tilelang.contrib.cutedsl.gemm_v2_sp.wgmma_sp_rs(A_dtype, B_dtype, C_dtype, M, N, K, tnspB, scaleA, scaleB, spsel, A_ptr, desc_b, C_ptr, scale_out, E_ptr, e_offset)¶

Sparse WGMMA register-shared: A from registers, B from a descriptor.

Sparse register A is always K-major, so it takes no trans-a operand.

Parameters:

  • A_dtype (str)

  • B_dtype (str)

  • C_dtype (str)

  • M (cutlass.cutlass_dsl.Constexpr[int])

  • N (cutlass.cutlass_dsl.Constexpr[int])

  • K (cutlass.cutlass_dsl.Constexpr[int])

  • tnspB (cutlass.cutlass_dsl.Constexpr[bool])

  • scaleA (cutlass.cutlass_dsl.Constexpr[int])

  • scaleB (cutlass.cutlass_dsl.Constexpr[int])

  • spsel (cutlass.cutlass_dsl.Constexpr[int])

  • A_ptr (cutlass.cute.Pointer)

  • desc_b (tilelang.contrib.cutedsl.gemm_v2.GmmaDescriptor)

  • C_ptr (cutlass.cute.Pointer)

  • scale_out (int)

  • E_ptr (cutlass.cute.Pointer)