Why cutile's rmsnorm required full register

[MARD1NO]

I just test 16384x512 bf16 dtype rmsnorm, Nsight compute report says rms_norm_kernel_static_persistent use 255 registers, and I wonder how to control register num to increase occupancy.

[MARD1NO]

and I observe a set max reg instruction in SASS:

[xjmxyt]

Hi, you can increase occupancy to increase occupancy, see examples in https://github.com/NVIDIA/TileGym/blob/main/src/tilegym/ops/cutile/matmul.py#L101

[MARD1NO]

Hi, you can increase occupancy to increase occupancy, see examples in https://github.com/NVIDIA/TileGym/blob/main/src/tilegym/ops/cutile/matmul.py#L101

And it seems trigger a deadlock when I use occupancy=16, and set x load allow_tma=False:

@ct.kernel(occupancy=16)
def _rms_norm_kernel_static_persistent(
    X,  # Input tensor
    Y,  # Output tensor
    W,  # Weight tensor
    Rstd,  # rstd output (for backward)
    TILE_SIZE_M: ct.Constant[int],  # rows per tile
    TILE_SIZE_N: ct.Constant[int],  # columns per tile
    EPS: ct.Constant[float],  # Epsilon value
    OFFSET: ct.Constant[float],  # Offset value
):
    """
    CuTile static persistent RMSNorm kernel that uses a persistent approach,
    where NUM_SMS tile blocks are launched and each tile block processes multiple output tiles
    for better efficiency.

    Formula: y = norm(x) * (offset + w)
    For Llama: offset=0.0, For Gemma3: offset=1.0
    """
    bid = ct.bid(0)

    # Infer tensor dimensions from input shape
    M = X.shape[0]  # Number of rows
    N = X.shape[1]  # Number of columns

    upper_bound = (M + TILE_SIZE_M - 1) // TILE_SIZE_M

    # Load weight vector once (shared across all tiles processed by this program).
    # padding_mode=ZERO keeps out-of-range columns (when N is not a power of two
    # and TILE_SIZE_N = next_power_of_2(N) > N) from pulling in uninitialized
    # memory; combined with the same mode on X, OOB contributions to y are zero
    # and never reach the stored output columns [0, N).
    w = ct.load(W, index=(0,), shape=(TILE_SIZE_N,), padding_mode=ct.PaddingMode.ZERO)
    w = ct.astype(w, ct.float32)

    # Static persistent loop: each  processes multiple tiles
    num_tile_blocks = ct.num_blocks(0)
    for current_bid in range(bid, upper_bound, num_tile_blocks):
        # Load input tile.
        # padding_mode=ZERO is required when N is not a power of two so that
        # the out-of-range columns [N, TILE_SIZE_N) contribute 0 to the
        # sum-of-squares reduction below; otherwise uninitialized memory
        # inflates the variance and compresses the normalized output.
        x = ct.load(
            X,
            index=(current_bid, 0),
            shape=(TILE_SIZE_M, TILE_SIZE_N),
            padding_mode=ct.PaddingMode.ZERO,
            latency=10,  # +2% perf from this hint
            allow_tma=False, 
        )
        x = ct.astype(x, ct.float32)

        x_squared = x * x

        x2_sum = ct.sum(x_squared, axis=1, keepdims=True)  # Shape: [TILE_SIZE_M, 1]

        N_f32 = ct.full((TILE_SIZE_M, 1), N * 1.0, dtype=ct.float32)
        variance = x2_sum / N_f32

        eps_tensor = ct.full((TILE_SIZE_M, 1), EPS, dtype=ct.float32)
        variance_eps = variance + eps_tensor
        rsqrt_var = ct.rsqrt(variance_eps)

        # Store rstd for backward pass
        ct.store(Rstd, index=(current_bid,), tile=ct.reshape(rsqrt_var, (TILE_SIZE_M,)), allow_tma=False)

        x_normalized = x * rsqrt_var

        # Broadcast weight to match input shape
        w_broadcasted = ct.reshape(w, (1, TILE_SIZE_N))

        # Apply offset to weight: (offset + w)
        offset_tensor = ct.full((1, TILE_SIZE_N), OFFSET, dtype=ct.float32)
        w_with_offset = offset_tensor + w_broadcasted

        # Apply linear transformation: y = x_normalized * (offset + w)
        y = x_normalized * w_with_offset
        y = ct.astype(y, X.dtype)
        ct.store(
            Y,
            index=(current_bid, 0),
            tile=y,
            allow_tma=False,  # +30% perf
            latency=3,  # +3% perf from this hint
        )

My cuda-tile version is 1.5.0, with cuda13.2 + torch2.12

[gengjiayao]

Hi @MARD1NO, thanks for reporting this! The occupancy=16 + allow_tma=False hang is a confirmed compiler bug, fixed in an upcoming CUDA release — until then keep TMA on or use occupancy <= 8. Hope this helps!

[MARD1NO]

Hi @MARD1NO, thanks for reporting this! The occupancy=16 + allow_tma=False hang is a confirmed compiler bug, fixed in an upcoming CUDA release — until then keep TMA on or use occupancy <= 8. Hope this helps!

Thanks！