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]]
Defines the shared memory layout for the B200 block-scaled matmul kernel, including A/B/C tiles, scale factor tiles, and pipeline barriers for TMA-MMA, accumulator, CLC, and TMEM deallocation.
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]): - tmem_dealloc_mbar (
Array[SharedMemBarrier, Int(1)]): - tmem_addr (
Array[UInt32, Int(1)]):
Implemented traits
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)
output_m
comptime output_m = config.output_tile_shape[Int(0)]
output_n
comptime output_n = 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)