IMPORTANT: To view this page as Markdown, append `.md` to the URL (e.g. /get-started.md). For the complete documentation index, see llms.txt.
Skip to main content
For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).

Mojo struct

B200BlockScaledMatmulSmem

struct B200BlockScaledMatmulSmem[a_type: DType, b_type: DType, c_type: DType, sfa_dtype: DType, sfb_dtype: DType, transpose_b: Bool, *, config: BlockScaledMatmulConfig[a_type, b_type, c_type, sfa_dtype, sfb_dtype, transpose_b]]

Owns the shared-memory buffers and mbarrier pools for one kernel block.

Allocates the A, B, C, SFA, and SFB SMEM tiles plus the TMA/MMA, accumulator, CLC, SFB, and cross-CTA SFB-readiness mbarrier arrays used to synchronize the producer/consumer warps of the small-BN block-scaled matmul kernel.

Parameters​

  • ​a_type (DType): Element dtype of the A operand matrix.
  • ​b_type (DType): Element dtype of the B operand matrix.
  • ​c_type (DType): Element dtype of the C output matrix.
  • ​sfa_dtype (DType): Element dtype of the A scale factors.
  • ​sfb_dtype (DType): Element dtype of the B scale factors.
  • ​transpose_b (Bool): Whether B is stored in K-major (transposed) layout.
  • ​config (BlockScaledMatmulConfig[a_type, b_type, c_type, sfa_dtype, sfb_dtype, transpose_b]): Kernel configuration controlling tile shapes, pipeline stages, cluster geometry, and swizzle modes.

Fields​

  • ​a_smem (Array[Scalar[a_type], Int((mul config.block_tile_shape[Int(0)], config.block_tile_shape[Int(2)], config.num_pipeline_stages))]):
  • ​b_smem (Array[Scalar[b_type], Int((mul config.block_tile_shape[Int(1)], config.block_tile_shape[Int(2)], config.num_pipeline_stages))]):
  • ​c_smem (Array[Scalar[c_type], Int((mul config.output_tile_shape[Int(0)], config.output_tile_shape[Int(1)], config.num_output_stages))]):
  • ​sfa_smem (Array[Scalar[sfa_dtype], Int((mul (config.block_tile_shape[Int(0)] // Int((mul (load_from_mem Tuple(Int(32), Int(4)).__getitem_param__[Int(0)]()), (load_from_mem Tuple(Int(32), Int(4)).__getitem_param__[Int(1)]())))), (load_from_mem Tuple(Int(32), Int(4)).__getitem_param__[Int(0)]()), (load_from_mem Tuple(Int(32), Int(4)).__getitem_param__[Int(1)]()), config.num_pipeline_stages, config.num_sf_k_tiles, 4))]):
  • ​sfb_smem (Array[Scalar[sfb_dtype], Int((mul (align_up(config.mma_shape[Int(1)], Int((mul (load_from_mem Tuple(Int(32), Int(4)).__getitem_param__[Int(0)]()), (load_from_mem Tuple(Int(32), Int(4)).__getitem_param__[Int(1)]())))) // Int((mul (load_from_mem Tuple(Int(32), Int(4)).__getitem_param__[Int(0)]()), (load_from_mem Tuple(Int(32), Int(4)).__getitem_param__[Int(1)]())))), (load_from_mem Tuple(Int(32), Int(4)).__getitem_param__[Int(0)]()), (load_from_mem Tuple(Int(32), Int(4)).__getitem_param__[Int(1)]()), config.num_pipeline_stages, config.num_sf_k_tiles, 4))]):
  • ​tma_mma_mbars (Array[SharedMemBarrier, ((config // config) * Int(2))]):
  • ​accum_mbars (Array[SharedMemBarrier, (config * Int(2))]):
  • ​clc_mbars_full (Array[SharedMemBarrier, config.num_clc_pipeline_stages]):
  • ​clc_mbars_empty (Array[SharedMemBarrier, config.num_clc_pipeline_stages]):
  • ​clc_throttle_mbars (Array[SharedMemBarrier, (config * Int(2))]):
  • ​clc_response (Array[UInt128, config.num_clc_pipeline_stages]):
  • ​sfb_mbars (Array[SharedMemBarrier, ((config // config) * Int(2))]):
  • ​sfb_ready_mbars (Array[SharedMemBarrier, (config // config) if (eq config.cta_group, 2) else Int(0)]):
  • ​tmem_dealloc_mbar (Array[SharedMemBarrier, Int(1)]):
  • ​tmem_addr (Array[UInt32, Int(1)]):

Implemented traits​

AnyType, Deinitable, Movable

comptime members​

a_smem_size​

comptime a_smem_size = (Int((mul config.block_tile_shape[Int(0)], config.block_tile_shape[Int(2)])) * config)

AScalesType​

comptime AScalesType = Scalar[sfa_dtype]

AType​

comptime AType = Scalar[a_type]

b_smem_size​

comptime b_smem_size = (Int((mul config.block_tile_shape[Int(1)], config.block_tile_shape[Int(2)])) * config)

BK​

comptime BK = config.block_tile_shape[Int(2)]

BM​

comptime BM = config.block_tile_shape[Int(0)]

BN​

comptime BN = config.block_tile_shape[Int(1)]

BScalesType​

comptime BScalesType = Scalar[sfb_dtype]

BType​

comptime BType = Scalar[b_type]

c_smem_size​

comptime c_smem_size = (Int((mul config.output_tile_shape[Int(0)], config.output_tile_shape[Int(1)])) * config)

CType​

comptime CType = Scalar[c_type]

MMA_K​

comptime MMA_K = config.mma_shape[Int(2)]

MMA_M​

comptime MMA_M = config.mma_shape[Int(0)]

MMA_N​

comptime MMA_N = config.mma_shape[Int(1)]

num_group_pipeline_stages​

comptime num_group_pipeline_stages = (config // config)

OutputM​

comptime OutputM = config.output_tile_shape[Int(0)]

OutputN​

comptime OutputN = config.output_tile_shape[Int(1)]

sfa_smem_size​

comptime sfa_smem_size = (Int((mul (config.block_tile_shape[Int(0)] // Int((mul (load_from_mem Tuple(Int(32), Int(4)).__getitem_param__[Int(0)]()), (load_from_mem Tuple(Int(32), Int(4)).__getitem_param__[Int(1)]())))), (load_from_mem Tuple(Int(32), Int(4)).__getitem_param__[Int(0)]()), (load_from_mem Tuple(Int(32), Int(4)).__getitem_param__[Int(1)]()), config.num_sf_k_tiles, 4)) * config)

sfb_smem_size​

comptime sfb_smem_size = (Int((mul (align_up(config.mma_shape[Int(1)], Int((mul (load_from_mem Tuple(Int(32), Int(4)).__getitem_param__[Int(0)]()), (load_from_mem Tuple(Int(32), Int(4)).__getitem_param__[Int(1)]())))) // Int((mul (load_from_mem Tuple(Int(32), Int(4)).__getitem_param__[Int(0)]()), (load_from_mem Tuple(Int(32), Int(4)).__getitem_param__[Int(1)]())))), (load_from_mem Tuple(Int(32), Int(4)).__getitem_param__[Int(0)]()), (load_from_mem Tuple(Int(32), Int(4)).__getitem_param__[Int(1)]()), config.num_sf_k_tiles, 4)) * config)

Models

Cookbook

Puzzles

Blog

Community

Contact