/tilegym-improve-cutile-kernel-perf

cuTile API Reference

Contents

Quick Lookup: Most Common Mistakes Import & Decorator Indexing Memory Operations Tensor Creation Reductions Scan Operations Matrix Operations Type & Shape Operations Slicing & Extraction Math Functions Comparison Operations Bitwise Operations Atomic Operations Debug & Utility Functions Host Functions Data Types Enums: PaddingMode, RoundingMode, MemoryOrder, MemoryScope Launch Pattern Kernel Compilation Hints Critical Rules (The 18 Rules)

For patterns, debug tables, and conversion reference: See cutile-patterns-reference.md

Quick Lookup: Most Common Mistakes

Import & Decorator

import cuda.tile as ct  # NOT import cutile as ct!

@ct.kernel
def kernel(X, Y, BLOCK: ct.Constant[int]):
    ...

ConstInt = ct.Constant[int]  # Type alias for cleaner signatures

Indexing

Persistent scheduling pattern (kernel processes multiple blocks):

@ct.kernel
def persistent_kernel(X, Y, BLOCK: ConstInt):
    num_blks = ct.num_blocks(0)       # total blocks in grid
    for bid in range(ct.bid(0), total_tiles, num_blks):
        x = ct.load(X, index=(bid,), shape=(BLOCK,))
        ct.store(Y, index=(bid,), tile=x)

Memory Operations

⚠️ TMA-FIRST STRATEGY

ALWAYS try TMA (ct.load/ct.store) FIRST! TMA is 2-4x faster than gather/scatter due to hardware acceleration.

TMA Load/Store (PREFERRED - Block-aligned)

Parameters:

• order — 'C' (default, no permutation), 'F' (reversed axes), or tuple of ints for custom axis permutation
• padding_mode — What value to use for out-of-bounds reads (see PaddingMode)
• latency — Hint for DRAM traffic intensity, int 1 (low) to 10 (high), or None (auto)
• allow_tma — If False, disables TMA for this load/store. Default None (TMA allowed)
• memory_order — Memory ordering for non-TMA load/store. Default MemoryOrder.WEAK (see MemoryOrder)
• memory_scope — Memory scope for non-TMA load/store. Default MemoryScope.NONE (see MemoryScope)

⚠️ CRITICAL: index and shape must have the SAME number of dimensions as the source tensor!

⚠️ CRITICAL: index is BLOCK INDEX (which block), NOT element offset!

# CORRECT: index=(bid,) means "load block number `bid`"
bid = ct.bid(0)
x = ct.load(X, index=(bid,), shape=(BLOCK,))  # Loads elements [bid*BLOCK : (bid+1)*BLOCK]

# WRONG: Do NOT multiply bid by BLOCK_SIZE
# x = ct.load(X, index=(bid * BLOCK,), shape=(BLOCK,))  # WRONG!

# Example: Loading 2D tile from 4D tensor [batch, head, seq, dim]
# CORRECT: index and shape both have 4 elements, then reshape
q = ct.load(
    Q, index=(batch_idx, head_idx, bid_x, 0), shape=(1, 1, TILE_M, TILE_D)
).reshape((TILE_M, TILE_D))

# WRONG: mismatched dimensions
# q = ct.load(Q, index=(batch_idx, head_idx, bid_x, 0), shape=(TILE_M, TILE_D))  # ERROR!

# Load with transpose
tile = ct.load(array2d, (0, 0), shape=(4, 2), order='F')

# Load a single element as 0d tile
tile = ct.load(array3d, (0, 0, 0), shape=())

Gather/Scatter (FALLBACK - Arbitrary offset)

Use ONLY when TMA truly fails (truly sparse random access). Most "paged" or "ragged" patterns CAN use TMA - see patterns below!

gather parameters:

• indices — Tuple of integer tiles (length = array rank), or single tile for 1D arrays
• mask — Boolean tile; where False, returns padding_value instead of loading
• padding_value — Value for masked/OOB elements (default: 0)
• check_bounds — If True (default), OOB indices return padding_value. If False, OOB is undefined behavior
• latency — DRAM traffic hint (1-10), or None (auto)

scatter parameters:

• indices — Same as gather
• value — Tile or scalar to store
• mask — Boolean tile; where False, no store occurs
• check_bounds — If True (default), OOB indices are skipped. If False, OOB is undefined behavior
• latency — DRAM traffic hint (1-10), or None (auto)

Note: When both mask and check_bounds=True are provided, the effective mask is the logical AND of both.

TMA with Runtime Index (ct.gather().item() Pattern) - CRITICAL!

⚠️ TMA works with RUNTIME indices! For paged attention or indirect access:

# ⚠️ WRONG (78x slower!): Using gather for all loads
page_id_tile = ct.gather(block_tables, (idx,))
k_indices = compute_flat_indices(page_id_tile, ...)
k_tile = ct.gather(k_cache.view(-1), k_indices)  # NO TMA!

# ✅ CORRECT: Extract scalar with .item(), then use ct.load(allow_tma=True)
page_id = ct.gather(block_tables, (idx,), padding_value=0).item()
k_tile = ct.load(k_cache, index=(page_id, ...), shape=(...), allow_tma=True)

Tensor Creation

⚠️ shape must be compile-time constants (literals or ct.Constant params), NOT X.shape.

Reductions

axis: None (reduce all), int, or tuple[int, ...].

ct.reduce example:

# Custom sum via reduce
result = ct.reduce(x, axis=0, func=lambda a, b: a + b, identity=0)

# Multi-tile reduce (x is a tuple of tiles)
# func takes 2N args and returns N combined tiles

Scan Operations

ct.scan example:

# Custom cumsum via scan
result = ct.scan(x, axis=0, func=lambda a, b: a + b, identity=0)

Matrix Operations

ct.mma signature: def mma(x, y, /, acc) -> Tile

acc is a positional parameter (not keyword-only). Both forms work:

acc = ct.mma(a, b, acc)       # positional — OK
acc = ct.mma(a, b, acc=acc)   # keyword — also OK

Supported mma dtypes:

⚠️ ct.mma does NOT auto-cast f32→tf32. You must manually cast:

a_tf32 = ct.astype(a, ct.tfloat32)
b_tf32 = ct.astype(b, ct.tfloat32)
acc = ct.mma(a_tf32, b_tf32, acc)

Block-Scaled MMA

⚠️ Note: mma_scaled is defined in _stub.py but is not yet exported from cuda.tile.__init__. The datatypes float8_e8m0fnu and float4_e2m1fn required by this API are also not yet exported. Confirm with the cuTile team before using.

ct.mma_scaled(x, x_scale, y, y_scale, /, acc) — block-scaled matrix multiply-accumulate for microscaling (MX) formats.

Computes: result[i,j] = sum_k (x[i,k] * x_scale[i,k/V]) * (y[k,j] * y_scale[k/V,j]) + acc[i,j]

tx = ct.full((16, 32), 1, dtype=ct.float8_e4m3fn)
sx = ct.full((16, 1), 1, dtype=ct.float8_e8m0fnu)   # scale shape: [M, K_s]
ty = ct.full((32, 16), 1, dtype=ct.float8_e4m3fn)
sy = ct.full((1, 16), 1, dtype=ct.float8_e8m0fnu)    # scale shape: [K_s, N]
acc = ct.full((16, 16), 0, dtype=ct.float32)
result = ct.mma_scaled(tx, sx, ty, sy, acc)

Type & Shape Operations

Tile properties: tile.dtype, tile.shape, tile.ndim Tile methods: tile.item() (reshape to 0D scalar), tile.reshape(shape), tile.permute(axes), tile.transpose(axis0, axis1), tile.astype(dtype), tile.extract(index, shape)

Array properties: array.dtype, array.shape, array.strides, array.ndim Array methods: array.slice(axis, start, stop) — creates a view with restricted range along one axis

Slicing & Extraction

# ct.extract: Extract a sub-tile from a larger tile
a_reshaped = ct.reshape(a_interleaved, (TILE_M, TILE_N, 2))

# Extract first slice along dim 2
gelu_part = ct.reshape(
    ct.extract(a_reshaped, index=(0, 0, 0), shape=(TILE_M, TILE_N, 1)),
    (TILE_M, TILE_N)
)
# Extract second slice along dim 2
linear_part = ct.reshape(
    ct.extract(a_reshaped, index=(0, 0, 1), shape=(TILE_M, TILE_N, 1)),
    (TILE_M, TILE_N)
)

# array.slice: Create a view of an array with restricted range
segment = A.slice(axis=1, start=offset, stop=offset + length)
tile = ct.load(segment, (0, 0), shape=(TILE_M, TILE_N))

Math Functions

Unary Math

flush_to_zero (bool): If True, flushes subnormal inputs/results to sign-preserving zero. Default False.

rounding_mode (RoundingMode): Controls rounding behavior for float ops. See RoundingMode enum.

Binary Math

Recommended: Use Python +, -, *, /, //, %, ** operators for all arithmetic on both tiles and scalars. Use ct.add/ct.mul/ct.sub/ct.truediv only when you need flush_to_zero= or rounding_mode= parameters (e.g., ct.truediv(x, y, rounding_mode=RoundingMode.APPROX)). The ct.* forms may also promote int32 to float — another reason to prefer Python operators for general use.

Conditional

Missing Functions (Must Implement Manually)

Comparison Operations

All comparisons return boolean tiles and support broadcasting + dtype promotion.

Bitwise Operations

Atomic Operations

All atomic operations follow the same index convention as ct.gather/ct.scatter.

Common parameters:

• memory_order — MemoryOrder.RELAXED, .ACQUIRE, .RELEASE, .ACQ_REL (default)
• memory_scope — MemoryScope.BLOCK, .DEVICE (default), .SYS
• check_bounds — If True (default), OOB indices are skipped

Debug & Utility Functions

# printf example (C-style format strings)
ct.printf("value: %d", tile)
ct.printf("two tiles: %d, %f", tile_a, tile_b)

# print example (Python-style, supports f-strings)
ct.print(f"tile={tile}")
ct.print(f"x={tile:.5f}", end='')
ct.print("tile:", tile, sep='=')

# static_eval example — select tile based on compile-time condition
x_or_y = ct.static_eval(x if N % 2 == 0 else y)

# static_assert example
ct.static_assert(x.dtype == y.dtype, f"Expected {x} and {y} to have same dtype.")

# static_iter example — compile-time unrolled loop
for i in ct.static_iter(range(4)):
    ...

Host Functions

# Prefer Python arithmetic on host (simpler, no ct import needed)
grid = ((N + BLOCK - 1) // BLOCK, 1, 1)

# ct.cdiv also valid on host, but Python arithmetic is preferred
# grid = (ct.cdiv(N, BLOCK), 1, 1)

# ct.cdiv in kernel code (operates on tiles)
num_iters = ct.cdiv(K, BLOCK_K)

Power-of-2 Utility

def next_power_of_2(x: int) -> int:
    """Round up to nearest power of 2 (required for tile shapes)"""
    return 1 << (x - 1).bit_length()

Data Types

ct.float16, ct.float32, ct.float64, ct.bfloat16
ct.tfloat32
ct.float8_e4m3fn, ct.float8_e5m2
ct.float8_e8m0fnu                        # 8-bit exponent-only (scale factor for mma_scaled) ⚠️ Not yet exported from cuda.tile.__init__
ct.float4_e2m1fn                         # 4-bit MX format (for mma_scaled) ⚠️ Not yet exported from cuda.tile.__init__
ct.int8, ct.int16, ct.int32, ct.int64
ct.uint8, ct.uint16, ct.uint32, ct.uint64
ct.bool_

Enums

PaddingMode (for ct.load)

RoundingMode (for math ops)

MemoryOrder (for load/store and atomics)

MemoryScope (for load/store and atomics)

ByTarget (for kernel hints)

from cuda.tile import ByTarget

# Different values per GPU architecture
@ct.kernel(num_ctas=ByTarget(sm_100=8, sm_120=4, default=2))
def kernel_fn(x):
    ...

Launch Pattern

# Grid can be 1-tuple, 2-tuple, or 3-tuple
grid = ((N + BLOCK - 1) // BLOCK,)    # 1D grid — OK
grid = (grid_m, grid_n)               # 2D grid — OK
grid = (grid_m, grid_n, 1)            # 3D grid — OK

ct.launch(torch.cuda.current_stream(), grid, kernel, (x, y, BLOCK, n))

ct.launch signature: launch(stream, grid, kernel, kernel_args)

• stream — CUDA stream (e.g., torch.cuda.current_stream())
• grid — Tuple of 1, 2, or 3 ints
• kernel — Function decorated with @ct.kernel
• kernel_args — Tuple of arguments to pass to the kernel

Kernel Compilation Hints

ct.kernel accepts optional hints that affect compilation and scheduling:

@ct.kernel(num_ctas=2, occupancy=4)
def kernel(X, Y, BLOCK: ct.Constant[int]):
    ...

# Or with ByTarget for architecture-specific values:
@ct.kernel(num_ctas=ct.ByTarget(sm_100=2), occupancy=ct.ByTarget(sm_100=4))
def kernel(X, Y, BLOCK: ct.Constant[int]):
    ...

Note: occupancy CAN be passed directly to @ct.kernel, but for production code with autotuning, passing it via hints_fn in autotune_launch is the recommended approach:

# Direct (simple cases):
@ct.kernel(occupancy=4)
def kernel(...): ...

# Via autotune (production):
ct_experimental.autotune_launch(
    stream, grid_fn=..., kernel=kernel, args_fn=...,
    hints_fn=lambda cfg: {"num_ctas": cfg.num_ctas, "occupancy": cfg.occupancy},
    search_space=configs,
)

Critical Rules (The 18 Rules)

Rule 1: Import Statement

import cuda.tile as ct  # NOT import cutile as ct!

Rule 2: Index = Block Index, NOT Element Offset

# cuTile uses block index for TMA, or computed indices for gather
x = ct.load(X, index=(bid,), shape=(BLOCK,))
# OR
indices = bid * BLOCK + ct.arange(BLOCK, dtype=ct.int32)
x = ct.gather(X, indices, check_bounds=True)

Rule 3: Python Operators for Index Math

# WRONG — ct.add/ct.mul promote int32 to float
indices = ct.add(ct.mul(bid, BLOCK), ct.arange(BLOCK, dtype=ct.int32))

# CORRECT — use Python +, *, /
indices = bid * BLOCK + ct.arange(BLOCK, dtype=ct.int32)

Rule 4: ct.mma — acc is Positional

# Both forms are correct:
acc = ct.mma(a, b, acc)       # positional — OK
acc = ct.mma(a, b, acc=acc)   # keyword — also OK

Rule 5: No None in ct.launch()

# WRONG
ct.launch(stream, grid, kernel, (x, None, n))

# CORRECT
dummy = torch.zeros(1, device=x.device)
ct.launch(stream, grid, kernel, (x, dummy, n))

Rule 6: Prefer Python Arithmetic on Host; Use ct.cdiv() in Kernel

# Host — prefer Python arithmetic:
grid = ((N + BLOCK - 1) // BLOCK, 1, 1)  # preferred
# grid = (ct.cdiv(N, BLOCK), 1, 1)       # also valid, but Python is simpler

# Kernel — ct.cdiv() operates on tiles:
num_iters = ct.cdiv(K, BLOCK_K)

Rule 7: ct.astype(), Not .to() or .cast()

# WRONG
y = x.to(ct.float32)

# CORRECT — function form
y = ct.astype(x, ct.float32)
# CORRECT — method form (preferred for chaining)
y = x.astype(ct.float32)
# CORRECT — chained on load
tile = ct.load(X, index=(bid,), shape=(BLOCK,)).astype(ct.float32)

Rule 8: Helper Functions - No @ct.kernel

# WRONG
@ct.kernel
def helper(x): return ct.exp(x)

# CORRECT - plain function
def helper(x): return ct.exp(x)

@ct.kernel
def main_kernel(X, Y, N: ConstInt):
    y = helper(x)

Rule 9: Pre-define Variables Before Branches

# WRONG — Variable only defined in one branch
if condition:
    result = ct.zeros((M,), dtype=ct.float32)
    result = ct.load(X, ...)
else:
    # result is undefined here!
    pass
output = result  # ERROR: result may not exist

# CORRECT — Pre-define ALL variables used across branches
result = ct.zeros((M,), dtype=ct.float32)  # Pre-define before branch
if condition:
    result = ct.load(X, ...)
else:
    result = ct.zeros((M,), dtype=ct.float32)
output = result  # OK: always defined

Rule 10: No break/continue in Loops

# WRONG
for i in range(N):
    if condition: break

# CORRECT - use conditionals
for i in range(N):
    if not condition:
        # loop body

Rule 11: Grid Must Be Tuple (1, 2, or 3 elements)

# WRONG
grid = N // BLOCK          # bare int
grid = [N // BLOCK, 1, 1]  # list

# CORRECT — tuple of 1, 2, or 3 ints
grid = ((N + BLOCK - 1) // BLOCK,) # 1-tuple
grid = (grid_m, grid_n)            # 2-tuple
grid = (grid_m, grid_n, 1)         # 3-tuple

Rule 12: ct.arange Starts at 0

# ct.arange(N) produces [0, 1, ..., N-1] — always starts at 0, no start param
offs = ct.arange(BLOCK, dtype=ct.int32)

Rule 13: NHWC Tensors - Use tensor.stride()

# WRONG: Assumes NCHW layout
offset = n * C * H * W + c * H * W + h * W + w  # WRONG for NHWC!

# CORRECT: Use actual strides from tensor
stride_n, stride_c, stride_h, stride_w = tensor.stride()
offset = n * stride_n + c * stride_c + h * stride_h + w * stride_w

# CRITICAL: tensor.view(-1) MAY REORDER DATA for non-contiguous!
# WRONG
flat = nhwc_tensor.view(-1)  # May silently reorder!

# CORRECT - Use torch.as_strided()
flat = torch.as_strided(tensor, (tensor.numel(),), (1,), storage_offset=tensor.storage_offset())

Rule 14: Block > Dim Masking - Apply ct.where AFTER gather

# When BLOCK_SIZE > actual dimension size
# WRONG - No mask applied
offsets = ct.arange(BLOCK_C, dtype=ct.int32)
data = ct.gather(input, base + offsets)
sum_val = ct.sum(data, axis=0)  # WRONG: includes padding!

# CORRECT - Use gather's mask parameter
offsets = ct.arange(BLOCK_C, dtype=ct.int32)
mask = offsets < actual_C
data = ct.gather(input, base + offsets, mask=mask, padding_value=0)
sum_val = ct.sum(data, axis=0)  # Correct!

# Alternative - Mask AFTER gather with ct.where
data = ct.gather(input, base + offsets)
data = ct.where(mask, data, ct.zeros((BLOCK_C,), dtype=data.dtype))
sum_val = ct.sum(data, axis=0)  # Correct!

# CRITICAL: Divide by actual_size, NOT BLOCK_SIZE
mean = sum_val / actual_C  # Correct
mean = sum_val / BLOCK_C   # WRONG!

Rule 15: Masked Scatter — Use mask= or Out-of-Bounds Offsets

# ct.scatter now supports mask= parameter!

# PREFERRED: Use scatter's mask parameter directly
offsets = ct.arange(BLOCK, dtype=ct.int32)
mask = offsets < actual_size
ct.scatter(Y, offsets, data, mask=mask)  # Masked elements are skipped

# ALTERNATIVE: Out-of-bounds offsets (ct.scatter skips OOB indices when check_bounds=True)
ARRAY_SIZE = Y.numel()  # Pass as kernel arg
oob_offset = ct.full((BLOCK,), ARRAY_SIZE, dtype=ct.int32)
offsets_masked = ct.where(mask, offsets, oob_offset)
ct.scatter(Y, offsets_masked, data)  # OOB positions skipped!

Rule 16: Constant Types — No Strings

# ct.Constant works with int, float, bool — but NOT str
# WRONG
@ct.kernel
def kernel(X, MODE: ct.Constant[str]):  # ERROR: str not supported!
    if MODE == "relu":
        ...

# CORRECT — Use integer enum
RELU = 0
GELU = 1
@ct.kernel
def kernel(X, MODE: ct.Constant[int]):
    if MODE == RELU:
        ...

# float and bool constants are also fine:
@ct.kernel
def kernel(X, SCALE: ct.Constant[float], USE_BIAS: ct.Constant[bool]):
    ...

Rule 17: Shape Args to ct.full/ct.zeros/ct.ones Must Be Static

# ct.full / ct.zeros / ct.ones shape arguments must be compile-time constants.
# WRONG — X.shape is dynamic, cannot be used as shape arg to ct.full
@ct.kernel
def kernel(X, N: ct.Constant[int]):
    result = ct.full(X.shape, 0.0, dtype=ct.float32)  # ERROR!

# CORRECT — Use compile-time constant
@ct.kernel
def kernel(X, N: ct.Constant[int], BLOCK: ct.Constant[int]):
    result = ct.full((BLOCK,), 0.0, dtype=ct.float32)  # OK: BLOCK is constexpr

# NOTE: X.shape IS fine for arithmetic, loop bounds, and comparisons:
@ct.kernel
def kernel(X, BLOCK: ct.Constant[int]):
    mask = ct.arange(BLOCK, dtype=ct.int32) < X.shape[0]  # OK!
    num_iters = ct.cdiv(X.shape[0], BLOCK)                 # OK!

Rule 18: No Dead Code

# cuTile compiles ALL parameters. Unused params waste registers and may cause errors.
# WRONG
@ct.kernel
def kernel(X, Y, Z, UNUSED: ct.Constant[int]):  # UNUSED wastes a register
    x = ct.load(X, ...)
    ct.store(Y, ...)
    # Z and UNUSED are never used!

# CORRECT — Remove unused parameters
@ct.kernel
def kernel(X, Y):
    x = ct.load(X, ...)
    ct.store(Y, ...)