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.

Requirements

1. ctas_per_cga=(2, 1, 1) — pass this at kernel launch time (not cluster_dims or num_ctas). ctas_per_cga is the correct way to enable 2-CTA because it bypasses PlanCTA's unreliable CTASplitNum encoding by forcing num_ctas=1 internally, then using Transform2CTALoads + Insert2CTASync for B-operand splitting and cross-CTA synchronization.

2. two_ctas=True on tl.dot() — tells the compiler this dot product should use 2-CTA MMA. The compiler automatically splits the B load across CTAs and inserts cross-CTA barriers.

3. TRITON_USE_META_WS=1 — 2-CTA with WS is only supported by Meta's WS pipeline. Upstream OAI WS does not support cluster_dims >= 2.

4. num_stages=1 — current 2-CTA implementations use 1 pipeline stage.

5. Grid M dimension must be >= 2 — the grid must launch at least 2 CTAs in the cluster dimension: grid = (max(triton.cdiv(M, BLOCK_M), 2), ...).

Kernel Pattern

"""
2-CTA GEMM with AutoWS.

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

@triton.jit
def matmul_2cta_ws_kernel(a_ptr, b_ptr, c_ptr, M, N, K,
                          stride_am, stride_ak, stride_bk, stride_bn,
                          stride_cm, stride_cn,
                          BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr,
                          BLOCK_K: tl.constexpr):
    pid_m = tl.program_id(0)
    pid_n = tl.program_id(1)
    offs_am = pid_m * BLOCK_M
    offs_bn = pid_n * BLOCK_N

    a_desc = tl.make_tensor_descriptor(
        a_ptr, shape=[M, K], strides=[stride_am, stride_ak],
        block_shape=[BLOCK_M, BLOCK_K])
    b_desc = tl.make_tensor_descriptor(
        b_ptr, shape=[K, N], strides=[stride_bk, stride_bn],
        block_shape=[BLOCK_K, BLOCK_N])
    c_desc = tl.make_tensor_descriptor(
        c_ptr, shape=[M, N], strides=[stride_cm, stride_cn],
        block_shape=[BLOCK_M, BLOCK_N])

    accumulator = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
    k_tiles = tl.cdiv(K, BLOCK_K)

    for k in tl.range(0, k_tiles, warp_specialize=True):
        offs_k = k * BLOCK_K
        a = a_desc.load([offs_am, offs_k])
        b = b_desc.load([offs_k, offs_bn])
        accumulator = tl.dot(a, b, accumulator, two_ctas=True)  # <-- 2-CTA MMA

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

Launch Pattern

grid = (max(triton.cdiv(M, BLOCK_M), 2), triton.cdiv(N, BLOCK_N))

matmul_2cta_ws_kernel[grid](
    a, b, c, M, N, K,
    a.stride(0), a.stride(1), b.stride(0), b.stride(1),
    c.stride(0), c.stride(1),
    BLOCK_M=128, BLOCK_N=128, BLOCK_K=64,
    num_stages=1,
    ctas_per_cga=(2, 1, 1),     # <-- enables 2-CTA cluster
)

What the compiler does

When ctas_per_cga is set with two_ctas=True on tl.dot():

1. Transform2CTALoads splits the B-operand load: each CTA loads half of BLOCK_N ([BLOCK_K, BLOCK_N/2]), offset by cluster_cta_id * BLOCK_N/2
2. Insert2CTASync inserts cross-CTA barriers before the 2-CTA MMA using the "arrive remote, wait local" pattern via mapa instructions
3. Both CTAs issue the MMA cooperatively

Reference files

• 2-CTA AutoWS test: third_party/tlx/tutorials/blackwell-triton-addmm-2cta_test.py
• TLX 2-CTA GEMM (manual WS): third_party/tlx/tutorials/blackwell_gemm_2cta.py
• Design doc: docs/design/2cta-autoWS-sync.md
• Transform2CTALoads: third_party/nvidia/hopper/lib/Transforms/Transform2CTALoads.cpp
• Insert2CTASync: third_party/nvidia/hopper/lib/Transforms/Insert2CTASync.cpp

Restrictions (When WS Bails Out)

If any of these conditions are violated, the compiler silently strips WS annotations and the kernel runs without specialization.

• Minimum 4 warps — num_warps >= 4 required (WarpSpecialization.cpp:148-153)
• No else blocks — scf.if with non-trivial else blocks not supported (WarpSpecialization.cpp:156-170)
• Max 16 total warps — if estimated warp budget exceeds 16, WS is stripped (PartitionSchedulingMeta.cpp:2664-2698)
• No distance > 1 loop-carried deps (ScheduleLoops.cpp:40-41)
• No outer loops for pipelining eligibility (ScheduleLoops.cpp:42-43)
• No barriers, asserts, or prints in the WS loop body (ScheduleLoops.cpp:44-47)
• Register alignment — minRegAutoWS and maxRegAutoWS must be divisible by 8 (compiler.py:140-142)
• 2-CTA + upstream WS not supported — only Meta's WS supports cluster_dims >= 2 (compiler.py:703-706)
• flatten=True skips WS — the kernel runs but WS is not applied
• data_partition_factor != 1 requires sufficient BLOCK_SIZE_M (256 for Blackwell dp=2, 128 for Hopper dp=2)
• Pointer-typed tensors should not be cross-partition

OSS Triton Fallback

The basic tl.range(warp_specialize=True) syntax works with both Meta WS and upstream OAI WS. The difference is entirely which compiler passes run, controlled by TRITON_USE_META_WS:

• Blackwell: upstream uses add_warp_specialize; Meta uses add_partition_scheduling_meta + add_hopper_warpspec
• Hopper: upstream uses add_hopper_warpspec only (internal doTaskPartition); Meta runs the full pipeline

Meta WS-specific features do NOT work with OSS Triton. The following kwargs and options are consumed only by Meta's compiler passes. They do not exist in OSS Triton's tl.range() signature — passing them will cause errors, not silent no-ops. If the kernel must run on both, use completely separate code paths gated by tl.constexpr:

• Partition scheduling kwargs: merge_epilogue, merge_correction, merge_epilogue_to_computation, separate_epilogue_store
• Memory allocation kwargs: smem_alloc_algo, tmem_alloc_algo, smem_budget, smem_circular_reuse
• Register budget options: minRegAutoWS, maxRegAutoWS, pingpongAutoWS
• 2-CTA / multi-CTA: multi_cta=True and cluster_dims >= 2 with WS
• early_tma_store_lowering

Dual-mode kernel pattern

Because these kwargs do not exist in OSS Triton's tl.range(), you cannot conditionally pass them (e.g., smem_alloc_algo=1 if X else None). You must use completely separate tl.range() calls:

@triton.jit
def _kernel_body(a_desc, b_desc, c_desc, ...):
    # ... shared load/dot/store logic ...

@triton.jit
def my_kernel(..., USE_META_WS: tl.constexpr):
    if USE_META_WS:
        for tile_id in tl.range(
                start_pid, num_tiles, NUM_SMS,
                warp_specialize=True,
                separate_epilogue_store=True,
                smem_alloc_algo=1,
                merge_epilogue=True,
        ):
            _kernel_body(...)
    else:
        for tile_id in tl.range(
                start_pid, num_tiles, NUM_SMS,
                warp_specialize=True,
        ):
            _kernel_body(...)

Exhaustive Reference

All tl.range() kwargs (from python/triton/language/core.py)

This is the complete list of every kwarg accepted by tl.range(). Items marked (Meta only) do not exist in OSS Triton and require separate code paths.

All AutoWS-relevant environment variables (from python/triton/knobs.py)

All AutoWS-relevant JIT-level options (from CUDAOptions)

Reference Files

• tl.range definition: python/triton/language/core.py (lines 3454-3484)
• Knobs: python/triton/knobs.py (line 516)
• Backend options: third_party/nvidia/backend/compiler.py (lines 145-180)
• Compiler pipeline: third_party/nvidia/backend/compiler.py (lines 659-716)
• GEMM test: python/test/unit/language/test_tutorial09_warp_specialization.py
• AddMM test: python/test/unit/language/test_autows_addmm.py
• FA test: third_party/tlx/tutorials/testing/test_correctness_autows.py
• FA kernel: third_party/tlx/tutorials/fused_attention_ws_device_tma.py
• FA kernel (DP): third_party/tlx/tutorials/fused_attention_ws_device_tma_dp.py
• AutoWS vs TLX: python/tutorials/test_hopper_fwd_autows_vs_tlx.py
• LIT tests: test/Hopper/WarpSpecialization/