tilelang.cuda.intrinsics.macro.mma_sm120_macro_generator

Attributes

lift

Classes

BlockScaleMmaConfig	Static SM120 warp-level block-scale MMA configuration.
SM120BlockScaleTile	Validated tile geometry for the SM120 packed-scale register pipeline.
TensorCoreIntrinEmitterSM120	Warp-level MMA emitter, with optional SM120 block-scale mode.

Module Contents

tilelang.cuda.intrinsics.macro.mma_sm120_macro_generator.lift

class tilelang.cuda.intrinsics.macro.mma_sm120_macro_generator.BlockScaleMmaConfig

Static SM120 warp-level block-scale MMA configuration.

kind: str

mma_prefix: str

atom_k: int

scale_vec_size: int

sf_vec_size: int

scale_type: str

a_dtype_abbrv: str

b_dtype_abbrv: str

accum_dtype: str = 'float32'

active_sfa_threads: int = 16

active_sfb_threads: int = 8

class tilelang.cuda.intrinsics.macro.mma_sm120_macro_generator.SM120BlockScaleTile

Validated tile geometry for the SM120 packed-scale register pipeline.

tile_m: int

tile_n: int

tile_k: int

block_row_warps: int

block_col_warps: int

warp_rows: int

warp_cols: int

warp_row_tiles: int

warp_col_tiles: int

kblocks: int

micro_size_m: int

micro_size_n: int

micro_size_k: int

sf_layout: str

sfa_words: int

sfb_words: int

warp_issues: int

warpgroup_issues: int

classmethod from_emitter(emitter, *, sf_layout, kblocks=None)

Parameters:

• emitter (TensorCoreIntrinEmitterSM120)

• sf_layout (str)

• kblocks (int | None)

Return type:

SM120BlockScaleTile

validate()

Return type:

None

compact_selector_scale_rows(lane, warp_m, warp_n)

Return SFA/SFB semantic rows loaded by the current compact TV package.

Parameters:

• lane (int)

• warp_m (int)

• warp_n (int)

Return type:

tuple[tuple[int, int], tuple[int, int]]

compact_selector_scale_word_offsets(lane, warp_m, warp_n, kblock)

Parameters:

• lane (int)

• warp_m (int)

• warp_n (int)

• kblock (int)

Return type:

tuple[tuple[int, int], tuple[int, int]]

static sfa_selector_source_lane(lane, scale_a_thread_id)

Parameters:

• lane (int)

• scale_a_thread_id (int)

Return type:

int

static sfb_selector_source_lane(lane, scale_b_thread_id)

Parameters:

• lane (int)

• scale_b_thread_id (int)

Return type:

int

compact_selector_effective_rows(lane, warp_m, warp_n, issue)

Return semantic SFA/SFB rows consumed by one compact-selector issue.

Parameters:

• lane (int)

• warp_m (int)

• warp_n (int)

• issue (tuple[int, int, int, int, int, int, int])

Return type:

tuple[int, int]

package_pingpong_lifecycle()

Return the current copy/gemm package lifecycle from the CUDA helper.

Each tuple is (op, register_package, kblock). Scale and A/B packages share the same register package id in the current implementation.

Return type:

tuple[tuple[str, int, int], Ellipsis]

omma_sf_issue_schedule_per_warp()

Return the current per-warp OMMA.SF issue schedule.

Each tuple is (mma_i, mma_j, n8_half, sfa_word, sfb_word, scale_a_thread_id, scale_b_thread_id). sfa_word and sfb_word are indices inside the current compact-selector scale package, not source-memory offsets.

Return type:

tuple[tuple[int, int, int, int, int, int, int], Ellipsis]

class tilelang.cuda.intrinsics.macro.mma_sm120_macro_generator.TensorCoreIntrinEmitterSM120(a_dtype='float16', b_dtype='float16', accum_dtype='float16', a_transposed=False, b_transposed=False, block_row_warps=2, block_col_warps=2, warp_row_tiles=8, warp_col_tiles=8, chunk=16, reduce_k=1, num_elems_per_byte=1, is_m_first=False, thread_var=None, is_blockscaled=False, kind='mxf4nvf4', scale_vec_size=4, stype='ue4m3')

Bases: tilelang.cuda.intrinsics.macro.mma_macro_generator.TensorCoreIntrinEmitter

Warp-level MMA emitter, with optional SM120 block-scale mode.

The block-scale mode keeps scale-factor storage explicit, matching TileLang’s TCGEN05 block-scaled style while targeting warp-level mma.sync.

Parameters:

• a_dtype (str)

• b_dtype (str)

• accum_dtype (str)

• a_transposed (bool)

• b_transposed (bool)

• block_row_warps (int)

• block_col_warps (int)

• warp_row_tiles (int)

• warp_col_tiles (int)

• chunk (int)

• reduce_k (int)

• num_elems_per_byte (int)

• is_m_first (bool | None)

• thread_var (tvm.tirx.Var | None)

• is_blockscaled (bool)

• kind (str)

• scale_vec_size (int)

• stype (str)

is_blockscaled = False

ldmatrix_a(A_local_buf, A_shared_buf, ki, rk=0)

Parameters:

• A_local_buf (tvm.tirx.Buffer)

• A_shared_buf (tvm.tirx.Buffer | tvm.tirx.BufferRegion)

• ki (tvm.tirx.PrimExpr)

• rk (tvm.tirx.PrimExpr | None)

ldmatrix_b(B_local_buf, B_shared_buf, ki, rk=0)

Parameters:

• B_local_buf (tvm.tirx.Buffer)

• B_shared_buf (tvm.tirx.Buffer | tvm.tirx.BufferRegion)

• ki (tvm.tirx.PrimExpr)

• rk (tvm.tirx.PrimExpr | None)

mma(A_local_buf, B_local_buf, C_local_buf, k_inner=0, *, SFA_buf=None, SFB_buf=None, k_start=0, sf_a_granularity_k=None, sf_b_granularity_k=None, sf_layout='rowmajor')

Parameters:

• k_inner (tvm.tirx.PrimExpr | None)

• k_start (tvm.tirx.PrimExpr)

• sf_a_granularity_k (int | None)

• sf_b_granularity_k (int | None)

• sf_layout (str)

ldscale(SFA_local_buf, SFB_local_buf, SFB_rep_local_buf, SFA_buf, SFB_buf, ki=0, k_start=0, sf_a_granularity_k=None, sf_b_granularity_k=None, sf_layout='rowmajor')

Parameters:

• ki (tvm.tirx.PrimExpr)

• k_start (tvm.tirx.PrimExpr)

• sf_a_granularity_k (int | None)

• sf_b_granularity_k (int | None)

• sf_layout (str)

ldscale_fragment(SFA_fragment_buf, SFB_fragment_buf, SFB_rep_fragment_buf, SFA_buf, SFB_buf, ki=0, k_start=0, sf_a_granularity_k=None, sf_b_granularity_k=None, sf_layout='rowmajor')

Load SM120 block-scale fragments into local registers.

This is currently a thin wrapper over the existing scale-word load. The separate name gives the SM120 MMA lowering a stable hook for a CUTLASS-like scale-fragment copy path.

Parameters:

• ki (tvm.tirx.PrimExpr)

• k_start (tvm.tirx.PrimExpr)

• sf_a_granularity_k (int | None)

• sf_b_granularity_k (int | None)

• sf_layout (str)

mma_blockscaled_fulltile(A_shared_buf, B_shared_buf, C_local_buf, SFA_buf, SFB_buf, sf_layout='rowmajor')

Emit an SM120 full-tile block-scaled MMA register micro-pipeline.

Parameters:

• A_shared_buf (tvm.tirx.Buffer | tvm.tirx.BufferRegion)

• B_shared_buf (tvm.tirx.Buffer | tvm.tirx.BufferRegion)

• C_local_buf (tvm.tirx.Buffer)

• SFA_buf (tvm.tirx.Buffer | tvm.tirx.BufferRegion)

• SFB_buf (tvm.tirx.Buffer | tvm.tirx.BufferRegion)

• sf_layout (str)

mma_full_b_atom_with_scale_fragments(A_local_buf, B_local_buf, C_local_buf, SFA_fragment_buf, SFB_fragment_buf, SFB_rep_fragment_buf, inst_m_idx, inst_n_idx)

Issue one SM120 block-scaled MMA atom from a full B fragment tile.

Parameters:

• inst_m_idx (tvm.tirx.PrimExpr | int)

• inst_n_idx (tvm.tirx.PrimExpr | int)

mma_full_b_atom_with_prefetched_scales(A_local_buf, B_local_buf, C_local_buf, SFA_local_buf, SFB_local_buf, SFB_rep_local_buf, inst_m_idx, inst_n_idx)

Parameters:

• inst_m_idx (tvm.tirx.PrimExpr | int)

• inst_n_idx (tvm.tirx.PrimExpr | int)