/tilegym-improve-cutile-kernel-perf

cuTile Patterns Quick-Reference Card

Quick-lookup tables, unique patterns, and debug reference for cuTile kernels.

For core API (functions, types, 18 rules): See cutile-api-reference.md For advanced conversion patterns (NHWC, masking, TMA decisions, ragged tensors): See advanced-patterns.md

Contents

• Unique Patterns
• Quick Debug Reference Table
• Appendix: Block vs Tile Terminology

Unique Patterns

Scalar Extraction from Tensor

Load a single element as a scalar tile for use in multi-dim indexing:

# Load single element, reshape to scalar
idx_tile = ct.load(input_ids, index=(row,), shape=(1,))
scalar_idx = ct.reshape(idx_tile, ())  # (1,) → ()

# Use scalar in multi-dim gather
embedding = ct.gather(weight_2d, (scalar_idx, col_offsets))

Scalar Load (0D Tile)

# Load single element as 0D tile (scalar)
scalar_val = ct.load(X, index=(0,), shape=())       # 1D array
scalar_val = ct.load(X, index=(0, 0, 0), shape=())  # 3D array
# Note: index tuple length must match source array rank

Batched MMA (3D Tiles)

ct.mma supports 2D and 3D tiles natively. For batched matmul, load 3D tiles and call ct.mma directly — no reshape needed:

@ct.kernel
def matmul_batched(A, B, C, B_DIM: ConstInt, M: ConstInt, N: ConstInt, K: ConstInt,
                   BLOCK_B: ConstInt, BLOCK_M: ConstInt, BLOCK_N: ConstInt):
    bid_b, bid_m, bid_n = ct.bid(0), ct.bid(1), ct.bid(2)

    # Load 3D tiles: batch × rows/cols × contraction dim
    a_tile = ct.load(A, index=(bid_b, bid_m, 0), shape=(BLOCK_B, BLOCK_M, K))
    b_tile = ct.load(B, index=(bid_b, 0, bid_n), shape=(BLOCK_B, K, BLOCK_N))

    # mma supports 3D directly — batch dims are broadcast
    acc = ct.zeros((BLOCK_B, BLOCK_M, BLOCK_N), dtype=ct.float32)
    acc = ct.mma(a_tile, b_tile, acc=acc)

    ct.store(C, index=(bid_b, bid_m, bid_n), tile=acc)

For true 4D tensors (e.g. shape (B, H, M, K)), reshape to 3D before ct.mma:

@ct.kernel
def matmul_4d(A, B, C, BATCH: ConstInt, HEADS: ConstInt, M: ConstInt, N: ConstInt, K: ConstInt,
              BLOCK_M: ConstInt, BLOCK_N: ConstInt):
    bid_bh, bid_m, bid_n = ct.bid(0), ct.bid(1), ct.bid(2)

    # Load 4D tiles (batch and head merged into one grid dim)
    # bid_bh indexes the flattened (BATCH * HEADS) dimension
    b_idx = bid_bh // HEADS
    h_idx = bid_bh % HEADS

    a_tile = ct.load(A, index=(b_idx, h_idx, bid_m, 0),
                     shape=(1, 1, BLOCK_M, K))         # 4D: (1, 1, BLOCK_M, K)
    b_tile = ct.load(B, index=(b_idx, h_idx, 0, bid_n),
                     shape=(1, 1, K, BLOCK_N))          # 4D: (1, 1, K, BLOCK_N)

    # Reshape 4D → 2D for mma
    a_2d = ct.reshape(a_tile, (BLOCK_M, K))             # (BLOCK_M, K)
    b_2d = ct.reshape(b_tile, (K, BLOCK_N))             # (K, BLOCK_N)

    acc = ct.zeros((BLOCK_M, BLOCK_N), dtype=ct.float32)
    acc = ct.mma(a_2d, b_2d, acc=acc)

    # Reshape back to 4D for store
    result = ct.reshape(acc, (1, 1, BLOCK_M, BLOCK_N))
    ct.store(C, index=(b_idx, h_idx, bid_m, bid_n), tile=result)

Multi-dimensional Index with Reshape (4D → 2D)

@ct.kernel
def attention_pattern(Q, K, V, Out,
                      batch_idx: ConstInt, head_idx: ConstInt,
                      TILE_M: ConstInt, TILE_N: ConstInt, TILE_D: ConstInt):
    bid_m = ct.bid(0)

    # Load 4D slice, reshape to 2D for computation
    q = ct.load(Q, index=(batch_idx, head_idx, bid_m, 0),
                shape=(1, 1, TILE_M, TILE_D)).reshape((TILE_M, TILE_D))

    # ... compute attention ...

    # Store back: reshape to 4D
    ct.store(Out, index=(batch_idx, head_idx, bid_m, 0),
             tile=result.reshape((1, 1, TILE_M, TILE_D)))

Cross-Reference: Advanced Patterns

For detailed coverage of these patterns, see the corresponding documents linked from the SKILL.md Reference Documents table:

Quick Debug Reference Table

Appendix: Block vs Tile Terminology

TileGym uses mixed terminology:

Convention in TileGym cuTile code:

• Prefer TILE_M, TILE_N, TILE_K over BLOCK_M, BLOCK_N, BLOCK_K
• Both are accepted in kernel_configs dicts
• ct.bid(0) returns the tile index, despite "block" in the name