IMPORTANT: To view this page as Markdown, append `.md` to the URL (e.g. /max/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. /max/get-started.md).

Mojo struct

BlockwiseFP8TileWriter

struct BlockwiseFP8TileWriter[c_type: DType, c_smem_dim0: Int, c_smem_dim1: Int, accum_type: DType, accum_num_stages: Int, accum_num_elements: Int, /, *, block_tile_shape: IndexList[Int(3)], mma_shape: IndexList[Int(3)], is_lower_frag_required: Bool, cta_group: Int, num_output_stages: Int, num_output_warps: Int, c_swizzle: TensorMapSwizzle]

Write register accumulators to GMEM via SMEM and TMA.

Parameters​

  • ​c_type (DType): Element type of the C output tensor.
  • ​c_smem_dim0 (Int): M dimension of the C shared-memory tile.
  • ​c_smem_dim1 (Int): N dimension of the C shared-memory tile.
  • ​accum_type (DType): Element type of the accumulator registers.
  • ​accum_num_stages (Int): Number of accumulator pipeline stages to drain.
  • ​accum_num_elements (Int): Number of elements per accumulator fragment set.
  • ​block_tile_shape (IndexList[Int(3)]): Block tile shape as (BM, BN, BK).
  • ​mma_shape (IndexList[Int(3)]): MMA instruction shape as (MMA_M, MMA_N, MMA_K).
  • ​is_lower_frag_required (Bool): Whether the lower register fragment is populated.
  • ​cta_group (Int): Number of CTAs cooperating per output tile.
  • ​num_output_stages (Int): Number of SMEM buffer stages for the output epilogue.
  • ​num_output_warps (Int): Number of warps participating in the output epilogue.
  • ​c_swizzle (TensorMapSwizzle): TMA swizzle pattern applied to the C shared-memory tile.

Implemented traits​

AnyType, ImplicitlyDeletable

comptime members​

bits​

comptime bits = 256

BM​

comptime BM = block_tile_shape[Int(0)]

BN​

comptime BN = block_tile_shape[Int(1)]

c_smem_layout​

comptime c_smem_layout = row_major[c_smem_dim0, c_smem_dim1]()

CTileArray​

comptime CTileArray = SMemTileArray2DRowMajor[c_type, c_smem_dim0, c_smem_dim1, num_output_stages]

data_paths​

comptime data_paths = 16

epc​

comptime epc = EpilogueConfig.create(MMA_M=mma_shape[Int(0)], MMA_N=mma_shape[Int(1)], stageN=Int((mul (accum_num_elements // (Int(128) // _resolve_warp_size())), 8)), cta_group=cta_group, transpose_c=False, BM=block_tile_shape[Int(0)], BN=block_tile_shape[Int(1)])

fragment_size​

comptime fragment_size = (Int(128) // _resolve_warp_size())

fragments_per_stage​

comptime fragments_per_stage = ((Int(128) // _resolve_warp_size()) * (accum_num_elements // (Int(128) // _resolve_warp_size())))

MMA_M​

comptime MMA_M = mma_shape[Int(0)]

MMA_N​

comptime MMA_N = mma_shape[Int(1)]

num_elements​

comptime num_elements = accum_num_elements

num_elements_per_load​

comptime num_elements_per_load = 8

num_stages​

comptime num_stages = accum_num_stages

repeats​

comptime repeats = (BlockwiseFP8TileWriter[c_type, c_smem_dim0, c_smem_dim1, accum_type, accum_num_stages, accum_num_elements, block_tile_shape=block_tile_shape, mma_shape=mma_shape, is_lower_frag_required=is_lower_frag_required, cta_group=cta_group, num_output_stages=num_output_stages, num_output_warps=num_output_warps, c_swizzle=c_swizzle].num_elements // (Int(128) // _resolve_warp_size()))

SMEMWriter​

comptime SMEMWriter = TMEMToSMemWriter[c_type, accum_type, c_smem_dim0, c_smem_dim1, BlockwiseFP8TileWriter[c_type, c_smem_dim0, c_smem_dim1, accum_type, accum_num_stages, accum_num_elements, block_tile_shape=block_tile_shape, mma_shape=mma_shape, is_lower_frag_required=is_lower_frag_required, cta_group=cta_group, num_output_stages=num_output_stages, num_output_warps=num_output_warps, c_swizzle=c_swizzle].epc, num_output_warps, c_swizzle]

stageN​

comptime stageN = ((accum_num_elements // (Int(128) // _resolve_warp_size())) * Int(8))

Methods​

write​

static def write[c_rank: Int, c_tile_shape: IndexList[c_rank], c_desc_shape: IndexList[c_rank], cluster_size: Int](accum: BlockwiseFP8Accumulator[accum_type, accum_num_stages, accum_num_elements, is_lower_frag_required, block_tile_shape, mma_shape, cluster_size], c_tiles: SMemTileArray2DRowMajor[c_type, c_smem_dim0, c_smem_dim1, num_output_stages], c_tma_op: TMATensorTile[c_type, c_rank, c_tile_shape, c_desc_shape], c_coord: Tuple[Int, Int])

Write accumulated register tiles to GMEM via double-buffered SMEM.

Parameters:

  • ​c_rank (Int): Rank of the C output tensor.
  • ​c_tile_shape (IndexList[c_rank]): Tile shape of the C output tensor.
  • ​c_desc_shape (IndexList[c_rank]): Descriptor shape of the C output tensor.
  • ​cluster_size (Int): Size of the threadblock cluster for the matmul.

Args:

write_absolute_with_bounds_check​

static def write_absolute_with_bounds_check[c_tensor_layout: TensorLayout, cluster_size: Int](accum: BlockwiseFP8Accumulator[accum_type, accum_num_stages, accum_num_elements, is_lower_frag_required, block_tile_shape, mma_shape, cluster_size], c_tiles: SMemTileArray2DRowMajor[c_type, c_smem_dim0, c_smem_dim1, num_output_stages], m_abs: UInt32, n_abs: UInt32, m_end: UInt32, expert_scale: Float32, c_tensor: TileTensor[c_type, c_tensor_layout, MutAnyOrigin])

Write accumulated register tiles to GMEM with bounds checking.

Parameters:

  • ​c_tensor_layout (TensorLayout): Layout of the C output tensor in GMEM used for bounds-checked element stores.
  • ​cluster_size (Int): Number of CTAs in the threadblock cluster for the matmul.

Args:

Edit this page

Was this page helpful?

Models

Cookbook

Puzzles

Blog

Community

Contact