Autows Authoring

AutoWS Kernel Authoring Guide

AutoWS is Meta's compiler-driven warp specialization for Triton kernels. Instead of manually writing producer/consumer partitions with TLX primitives (barriers, tlx.async_tasks, local_alloc), you annotate tl.range() with warp_specialize=True and the compiler automatically partitions ops, inserts barriers, allocates SMEM/TMEM buffers, and handles multi-buffering.

Minimal recipe:

1. Set TRITON_USE_META_WS=1 (or triton.knobs.nvidia.use_meta_ws = True)
2. Use tl.range(..., warp_specialize=True) on your loop
3. Pass num_warps >= 4 at launch

Related skills: autows-testing (run tests), ir-debugging (IR dumps), autows-docs (compiler internals).

Key Authoring Rules

1. Place warp_specialize=True on the outermost/persistent loop. For persistent kernels, annotate the tile loop (tl.range(start_pid, num_tiles, NUM_SMS, warp_specialize=True)), not the inner K-reduction loop. The inner loop uses plain range(). For non-persistent kernels, annotate the main compute loop.

2. Use TMA loads. AutoWS partitions loads into a producer warp group and compute into consumer warp groups. This partitioning is most effective with TMA descriptor loads (a_desc.load(...)) rather than pointer-based tl.load(). TMA enables async bulk copies that the producer can issue independently while consumers compute. All reference kernels use TensorDescriptor + desc.load()/desc.store().

3. Memory allocation and partition scheduling kwargs are Meta-only. The tl.range() kwargs for controlling memory allocation (smem_alloc_algo, tmem_alloc_algo, smem_budget, smem_circular_reuse) and partition scheduling (merge_epilogue, merge_correction, merge_epilogue_to_computation, separate_epilogue_store) are consumed exclusively by Meta's WS passes. They do not exist in OSS Triton's tl.range() — see OSS Fallback for how to gate them.

4. Compute the post-loop epilogue from the genuine MMA accumulator — never a synthetic matmul. The partition scheduler keeps a post-loop reduction/gate/epilogue only if it sits in a real MMA's backward slice. If the reduction is fed by a synthetic matmul inserted only to create an MMA (e.g. (x·x)@e0 to get a row-sum), the post-loop region belongs to no genuine MMA and the scheduler elides the entire epilogue: the kernel still emits ttg.warp_specialize and passes a structural WS check, but the tl.sum / normalize / tt.store ops are all dropped and the output is unwritten garbage (a silent correctness failure, not a compile error). Reduce or gate from the real A@B accumulator instead. Confirm the WS TTGIR has a dedicated "epilogue" Meta partition and that the store ops survive.

Enabling AutoWS

Environment variable (recommended for running kernels from the command line):

TRITON_USE_META_WS=1 python my_kernel.py

Also add TRITON_USE_META_WS=1 to the kernel script's module docstring so users know it's required:

"""
My AutoWS GEMM kernel.

Usage:
    TRITON_USE_META_WS=1 python my_gemm.py
"""

Programmatic (recommended for correctness tests only):

with triton.knobs.nvidia.scope():
    triton.knobs.nvidia.use_meta_ws = True
    # ... launch kernel ...

Use triton.knobs.nvidia.scope() in correctness tests so the knob is automatically restored after the test, preventing state leakage between test cases. Do NOT use scope() for actual runtime/benchmark scripts — use the TRITON_USE_META_WS=1 env var instead.

TRITON_USE_META_WS is a cache-invalidating env var (listed in include/triton/Tools/Sys/GetEnv.hpp), meaning changing it forces recompilation.

Additional Environment Variables

Source: python/triton/knobs.py lines 502-534

tl.range() Kwargs Reference

Defined in python/triton/language/core.py (tl.range.__init__).

Core

Memory Allocation (Meta WS only)

These kwargs control how the compiler allocates SMEM/TMEM buffers. They are consumed by WSMemoryPlanner.cpp via loop attributes.

Partition Scheduling (Meta WS only)

These kwargs control how PartitionSchedulingMeta assigns ops to partitions. They override pass-level defaults (all false) via per-loop attributes, read at PartitionSchedulingMeta.cpp lines 2614-2629.

JIT-Level Options

Passed at kernel launch time. Defined in CUDAOptions at third_party/nvidia/backend/compiler.py lines 145-180.

Register Budget (minRegAutoWS / maxRegAutoWS)

These control how the 64K hardware register file is divided across warp-specialized partitions. Each partition runs a subset of warps and can have a different register cap (emitted as PTX setmaxnreg instructions).

Partition types:

• Non-tensor partitions — partitions without MMA/dot ops (typically the TMA load producer). These get minRegAutoWS registers.
• Tensor partitions — partitions with MMA/dot ops (computation, correction, reduction). These get maxRegAutoWS registers (if set) or split the remainder evenly (if not set).
• Default partition — runs outside the WS region. Gets leftover registers.

When maxRegAutoWS is NOT set (default): Non-tensor partitions get minRegAutoWS (24). Tensor partitions AND the default partition all receive sentinel value -1, meaning "split the remaining register pool evenly after deducting the fixed allocations":

regsPerThread = (totalRegs - fixedRegs) / leftoverThreads

Computed in AllocateWarpGroups.cpp lines 244-280.

When maxRegAutoWS IS set: Non-tensor partitions get minRegAutoWS. Tensor partitions get maxRegAutoWS. The default partition gets ALL leftover registers. Computed in OptimizePartitionWarps.cpp lines 295-314.

Both values must be divisible by 8 (compiler.py:140-142).

Checking Register Allocations

To see which partitions map to which register budget and verify the actual allocation:

1. IR inspection: Set MLIR_ENABLE_DUMP=1. Look for the ttg.warp_specialize op which carries:

   • requestedRegisters = array<i32: ...> — what OptimizePartitionWarps requested
   • actualRegisters = array<i32: ...> — what AllocateWarpGroups computed
   • Array order: [default_partition, partition_0, partition_1, ...]
   • A -1 in requestedRegisters means "split evenly"
   • Example: requestedRegisters = array<i32: 24, -1, 24> means load partition gets 24, computation splits leftovers, epilogue store gets 24

2. PTXAS log: TRITON_DUMP_PTXAS_LOG=1 prints ptxas verbose output showing register usage.

3. PTX inspection: kernel.asm['ptx'] — search for setmaxnreg.inc.sync.aligned and setmaxnreg.dec.sync.aligned to see register reallocation at partition boundaries.

Kernel Patterns

GEMM (K-loop WS)

Warp-specialize the inner K-reduction loop. Based on matmul_kernel_tma_ws in python/test/unit/language/test_tutorial09_warp_specialization.py lines 34-94.

"""
GEMM with AutoWS (K-loop warp specialization).

Usage:
    TRITON_USE_META_WS=1 python my_gemm.py
"""

@triton.jit
def matmul_kernel_ws(a_desc, b_desc, c_desc, M, N, K,
                     BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr,
                     BLOCK_SIZE_K: tl.constexpr,
                     DATA_PARTITION_FACTOR: tl.constexpr,
                     SMEM_ALLOC_ALGO: tl.constexpr,
                     SEPARATE_EPILOGUE_STORE: tl.constexpr, ...):
    # ... pid computation ...
    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
    for k in tl.range(
            k_tiles,
            warp_specialize=True,
            data_partition_factor=DATA_PARTITION_FACTOR,
            smem_alloc_algo=SMEM_ALLOC_ALGO,
            separate_epilogue_store=SEPARATE_EPILOGUE_STORE,
    ):
        offs_k = k * BLOCK_SIZE_K
        a = a_desc.load([offs_am, offs_k])
        b = b_desc.load([offs_bn, offs_k])
        accumulator = tl.dot(a, b.T, accumulator)

    c_desc.store([offs_cm, offs_cn], accumulator.to(dtype))

Persistent GEMM (tile-loop WS)

Warp-specialize the outer persistent loop. Inner K-loop uses plain range(). Based on matmul_kernel_tma_persistent_ws in same file, lines 102-167.

"""
Persistent GEMM with AutoWS (tile-loop warp specialization).

Usage:
    TRITON_USE_META_WS=1 python my_persistent_gemm.py
"""

@triton.jit
def matmul_persistent_ws(a_desc, b_desc, c_desc, M, N, K,
                         BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr,
                         BLOCK_SIZE_K: tl.constexpr, NUM_SMS: tl.constexpr,
                         DATA_PARTITION_FACTOR: tl.constexpr,
                         SMEM_ALLOC_ALGO: tl.constexpr,
                         SEPARATE_EPILOGUE_STORE: tl.constexpr, ...):
    start_pid = tl.program_id(axis=0)
    num_tiles = tl.cdiv(M, BLOCK_SIZE_M) * tl.cdiv(N, BLOCK_SIZE_N)
    k_tiles = tl.cdiv(K, BLOCK_SIZE_K)

    for tile_id in tl.range(
            start_pid, num_tiles, NUM_SMS,
            warp_specialize=True,                           # on the OUTER loop
            data_partition_factor=DATA_PARTITION_FACTOR,
            smem_alloc_algo=SMEM_ALLOC_ALGO,
            separate_epilogue_store=SEPARATE_EPILOGUE_STORE,
    ):
        # ... compute pid_m, pid_n from tile_id ...
        accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
        for ki in range(k_tiles):                           # plain range()
            a = a_desc.load([offs_am, ki * BLOCK_SIZE_K])
            b = b_desc.load([offs_bn, ki * BLOCK_SIZE_K])
            accumulator = tl.dot(a, b.T, accumulator)

        c_desc.store([offs_am, offs_bn], accumulator.to(dtype))

Flash Attention (inner-loop WS)

Warp-specialize the KV-iteration loop with partition scheduling hints. Based on python/tutorials/fused-attention-ws-device-tma-hopper-or-blackwell.py lines 150-155.

for start_n in tl.range(
        lo, hi, BLOCK_N,
        warp_specialize=warp_specialize,
        merge_epilogue=True,
        merge_correction=True,
        smem_alloc_algo=1,
        data_partition_factor=DP_FACTOR,
):
    # TMA descriptor loads for K and V
    k = k_desc.load([start_n, offs_k])
    v = v_desc.load([start_n, offs_d])
    # QK dot product
    qk = tl.dot(q, k)
    # ... softmax ...
    # PV dot product
    acc = tl.dot(p.to(dtype), v, acc)

Test & Launch Boilerplate

Correctness tests should use triton.knobs.nvidia.use_meta_ws = True (not the env var) inside a scope() block to avoid state leakage between test cases.

Based on test_tutorial09_matmul_tma_warp_specialize in python/test/unit/language/test_tutorial09_warp_specialization.py lines 416-492.

import torch
import triton
from triton.tools.tensor_descriptor import TensorDescriptor

def test_my_kernel():
    with triton.knobs.nvidia.scope():
        triton.knobs.nvidia.use_meta_ws = True

        M, N, K = 8192, 8192, 1024
        BLOCK_M, BLOCK_N, BLOCK_K = 128, 128, 64
        dtype = torch.float16

        torch.manual_seed(42)
        A = torch.randn((M, K), dtype=dtype, device="cuda")
        B = torch.randn((N, K), dtype=dtype, device="cuda")
        C = torch.empty((M, N), dtype=dtype, device="cuda")

        # TMA requires a custom allocator
        def alloc_fn(size, align, stream):
            return torch.empty(size, dtype=torch.int8, device="cuda")
        triton.set_allocator(alloc_fn)

        a_desc = TensorDescriptor(A, A.shape, A.stride(), [BLOCK_M, BLOCK_K])
        b_desc = TensorDescriptor(B, B.shape, B.stride(), [BLOCK_N, BLOCK_K])
        c_desc = TensorDescriptor(C, C.shape, C.stride(), [BLOCK_M, BLOCK_N])

        grid = lambda META: (triton.cdiv(M, META["BLOCK_SIZE_M"])
                           * triton.cdiv(N, META["BLOCK_SIZE_N"]),)

        # Launch — capture handle to inspect IR
        kernel = matmul_kernel_ws[grid](
            a_desc, b_desc, c_desc, M, N, K,
            BLOCK_SIZE_M=BLOCK_M, BLOCK_SIZE_N=BLOCK_N, BLOCK_SIZE_K=BLOCK_K,
            DATA_PARTITION_FACTOR=1, SMEM_ALLOC_ALGO=0,
            SEPARATE_EPILOGUE_STORE=False,
            num_stages=3, num_warps=4,
        )

        # 1. Verify WS was applied
        ttgir = kernel.asm["ttgir"]
        assert "ttg.warp_specialize" in ttgir

        # 2. Verify correct HW instructions (pick one per arch)
        # Blackwell: assert "ttng.tc_gen5_mma" in ttgir
        # Hopper:    assert "ttng.warp_group_dot" in ttgir
        assert "ttng.async_tma_copy_global_to_local" in ttgir

        # 3. Compare against reference
        ref = torch.matmul(A.float(), B.T.float()).to(dtype)
        torch.testing.assert_close(ref, C, atol=0.03, rtol=0.03)

Verifying AutoWS is Working

1. IR check: kernel.asm["ttgir"] must contain "ttg.warp_specialize" — confirms WS was applied.

2. MMA check: Look for "ttng.tc_gen5_mma" (Blackwell) or "ttng.warp_group_dot" (Hopper) — confirms tensor core usage.

3. TMA check: Look for "ttng.async_tma_copy_global_to_local" — confirms async TMA copies in the producer partition.

4. Partition check: Set MLIR_ENABLE_DUMP=1 and look for ttg.partition = array<i32: N> attributes in the IR after PartitionSchedulingMeta. This shows how ops were assigned to partitions.

5. Full IR dump: Set TRITON_KERNEL_DUMP=<kernel_name> + TRITON_ALWAYS_COMPILE=1 to dump IR at each compilation stage. See ir-debugging skill for details.

6. AutoWS vs TLX comparison:

   TRITON_USE_META_WS=1 python python/tutorials/test_hopper_fwd_autows_vs_tlx.py

ttg.warp_specialize is necessary, not sufficient. Its presence or count alone does not prove a correct WS kernel — an elided epilogue (Key Authoring Rule 5) leaves the op in the IR while the stores are gone. Prefer the named Meta partition list as the reliable signal: a "load" / "gemm" / "computation" set, plus a dedicated "epilogue" partition whenever the kernel has a post-loop epilogue.

2-CTA (Multi-CTA) with AutoWS

2-CTA allows two CTAs in a cluster to cooperatively execute a single MMA, doubling the N dimension of the output tile. This is a Blackwell-only, Meta WS-only feature.