Triton

https://triton-lang.org/main/
Admin
Triton is a Python-based language and compiler designed for parallel programming, enabling the...
Tokens:909,965
Snippets:4,645
Trust Score:8.1
Update:1 week ago
Show doc for...
Tokens:
Raw
### Setup and naive matmul kernel (partial)

Source: https://triton-lang.org/main/_downloads/b5537c431cd73b8ee3c32cdf90c58104/09-persistent-matmul.ipynb

Initialization and setup for the persistent matmul tutorial: imports, device detection (CUDA/HIP), cuBLAS/hipBLAS handles, capability checks (TMA, warp specialization, CLC), autotune configuration generation, and the start of the naive matmul kernel with launch metadata.

```python
import argparse
import itertools

import torch
import triton
import triton.language as tl
import triton.profiler as proton
from triton.tools.tensor_descriptor import TensorDescriptor
from contextlib import contextmanager

from typing import Optional


def is_cuda():
    return triton.runtime.driver.active.get_current_target().backend == "cuda"


def is_hip():
    return triton.runtime.driver.active.get_current_target().backend == "hip"


if is_cuda():
    from triton._C.libtriton import nvidia
    device_workspace = torch.empty(32 * 1024 * 1024, device="cuda", dtype=torch.uint8)
    device_blas = nvidia.cublas.CublasLt(device_workspace)
elif is_hip():
    from triton._C.libtriton import amd
    device_workspace = torch.empty(32 * 1024 * 1024, device="cuda", dtype=torch.uint8)
    device_blas = amd.hipblas.HipblasLt(device_workspace)
else:
    device_blas = None


def device_blas_name():
    return 'cuBLAS' if is_cuda() else 'hipBLAS'


def supports_tma():
    return is_cuda() and torch.cuda.get_device_capability()[0] >= 9


def is_hopper():
    return torch.cuda.get_device_capability()[0] == 9


def supports_ws():
    return is_cuda() and torch.cuda.get_device_capability()[0] >= 9


def supports_clc():
    return is_cuda() and torch.cuda.get_device_capability()[0] >= 10


def _matmul_launch_metadata(grid, kernel, args):
    ret = {}
    M, N, K, WS = args["M"], args["N"], args["K"], args.get("WARP_SPECIALIZE", False)
    ws_str = "_ws" if WS else ""
    ret["name"] = f"{kernel.name}{ws_str} [M={M}, N={N}, K={K}]"
    if "c_ptr" in args:
        bytes_per_elem = args["c_ptr"].element_size()
    else:
        bytes_per_elem = 1 if args["FP8_OUTPUT"] else 2
    ret[f"flops{bytes_per_elem * 8}"] = 2. * M * N * K
    ret["bytes"] = bytes_per_elem * (M * K + N * K + M * N)
    return ret


HAS_TENSOR_DESC = supports_tma() and hasattr(tl, "make_tensor_descriptor")
HAS_HOST_TENSOR_DESC = supports_tma() and hasattr(triton.tools.tensor_descriptor, "TensorDescriptor")
HAS_WARP_SPECIALIZE = supports_ws() and HAS_TENSOR_DESC
HAS_TMA_CLC = supports_clc() and HAS_HOST_TENSOR_DESC


def matmul_get_configs(pre_hook=None):
    return [
        triton.Config({'BLOCK_SIZE_M': BM, 'BLOCK_SIZE_N': BN, "BLOCK_SIZE_K": BK, "GROUP_SIZE_M": 8}, num_stages=s,
                      num_warps=w, pre_hook=pre_hook)
        for BM in [128]
        for BN in [128, 256]
        for BK in [64, 128]
        for s in ([2, 3, 4])
        for w in [4, 8]
    ]


@triton.autotune(
    configs=matmul_get_configs(),
    key=["M", "N", "K"],
)
@triton.jit(launch_metadata=_matmul_launch_metadata)
def matmul_kernel(a_ptr, b_ptr, c_ptr,  #
                  M, N, K,  #
                  stride_am, stride_ak,  #
                  stride_bk, stride_bn,  #
                  stride_cm, stride_cn,  #
                  BLOCK_SIZE_M: tl.constexpr,  #
                  BLOCK_SIZE_N: tl.constexpr,  #
                  BLOCK_SIZE_K: tl.constexpr,  #
                  GROUP_SIZE_M: tl.constexpr,  #
                  ):
    pid = tl.program_id(axis=0)
    num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
    num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
    num_pid_in_group = GROUP_SIZE_M * num_pid_n
    group_id = pid // num_pid_in_group
    first_pid_m = group_id * GROUP_SIZE_M
    group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
    pid_m = first_pid_m + (pid % group_size_m)
    pid_n = (pid % num_pid_in_group) // group_size_m

    start_m = pid_m * BLOCK_SIZE_M
    start_n = pid_n * BLOCK_SIZE_N

    offs_am = start_m + 
```

--------------------------------

### Install tutorial dependencies

Source: https://triton-lang.org/main/_sources/getting-started/tutorials/index.rst.txt

Navigate to the Triton repository and install the tutorial dependencies using the provided requirements file.

```bash
cd triton
pip install -r python/tutorials/requirements.txt
```

--------------------------------

### Main entry point for padded example

Source: https://triton-lang.org/main/getting-started/tutorials/gluon/conv-im2col.html

Runs the padded-region example when the script is executed directly, then prints a separator line.

```python
if __name__ == "__main__":
    run_tma_im2col_padded()
    print("\n" + "=" * 50 + "\n")
```

--------------------------------

### matmul_load_partition setup

Source: https://triton-lang.org/main/_sources/getting-started/tutorials/gluon/multicta.md.txt

Starts the load partition by extracting the block K dimension and the total K, then creates a counter for concurrent loads and obtains the CLC scheduler consumer via get_clc_consumer().

```python
@gluon.jit
def matmul_load_partition(p):
    block_k: gl.constexpr = p.a_desc.block_shape[1]
    K = p.a_desc.shape[1]

    concurrent_loads: gl.constexpr = p.load_ready_bars.shape[0]
    state = Counter.create(1, concurrent_loads)
    scheduler = p.get_clc_consumer()
```

--------------------------------

### Setup and imports for tcgen05_mma_scaled tutorial

Source: https://triton-lang.org/main/_sources/getting-started/tutorials/gluon/tcgen05-mma-scaled.md.txt

Imports and setup for the tutorial, including a check for Blackwell GPU capability and importing utilities from previous tutorials.

```python
import itertools
import importlib
import pytest
import triton
import torch
import triton.experimental.gluon as gluon
import triton.experimental.gluon.language as gl
from dataclasses import replace
from triton.tools.mxfp import MXFP4Tensor, MXScaleTensor
from triton.experimental.gluon.nvidia.hopper import TensorDescriptor
from triton.experimental.gluon.language.nvidia.blackwell import (
    TensorMemoryLayout,
    TensorMemoryScalesLayout,
    allocate_tensor_memory,
    fence_async_shared,
    tensor_memory_descriptor,
    tcgen05_copy,
    tcgen05_commit,
    tcgen05_mma_scaled,
    mbarrier,
    tma,
)


def is_blackwell():
    target = triton.runtime.driver.active.get_current_target()
    return target.backend == "cuda" and torch.cuda.get_device_capability()[0] == 10


if __name__ == "__main__" and not is_blackwell():
    raise RuntimeError("This tutorial requires a Blackwell NVIDIA GPU")

# Re-use utilities from the previous tutorials.
t7 = importlib.import_module("07-persistence")
t8 = importlib.import_module("08-warp-specialization")
```

--------------------------------

### Main guard to run offset example

Source: https://triton-lang.org/main/_sources/getting-started/tutorials/gluon/conv-im2col.md.txt

Entry point that runs the offset example and prints a separator line after it. Useful when executing this file directly.

```python
if __name__ == "__main__":
    run_tma_im2col_offset()
    print("\n" + "=" * 50 + "\n")
```

--------------------------------

### async_load_im2col starting position

Source: https://triton-lang.org/main/getting-started/tutorials/gluon/conv-im2col.html

Shows the starting position formula for async_load_im2col, where the first pixel is at (batch_idx, start_h + h_offset, start_w + w_offset, channel_start).

```text
Starting position (the first pixel): (batch_idx, start_h + h_offset, start_w + w_offset, channel_start)

```

--------------------------------

### Run padded TMA im2col example as main

Source: https://triton-lang.org/main/_sources/getting-started/tutorials/gluon/conv-im2col.md.txt

Runs the padded-region example as the main script and prints a separator line after the output.

```python

if __name__ == "__main__":
    run_tma_im2col_padded()
    print("\n" + "=" * 50 + "\n")
```

--------------------------------

### Program output for padded example

Source: https://triton-lang.org/main/getting-started/tutorials/gluon/conv-im2col.html

Program output from running the padded example: prints the title and the loaded pixel values, followed by a separator line.

```text
Example 2: Loading from padded region

Loaded data (pixel_id, value):
  pixel  0: 0
  pixel  1: 0
  pixel  2: 0
  pixel  3: 0
  pixel  4: 0
  pixel  5: 1
  pixel  6: 2
  pixel  7: 3
  pixel  8: 0
  pixel  9: 5
  pixel 10: 6
  pixel 11: 7
  pixel 12: 0
  pixel 13: 9
  pixel 14: 10
  pixel 15: 11

==================================================
```

--------------------------------

### Attention Forward Gluon Example (01-attention-forward.py)

Source: https://triton-lang.org/main/getting-started/examples/gluon/attention-forward.html

Full example file for attention forward pass using Triton Gluon. Imports required modules and defines layout utilities and data abstractions. The code is from python/examples/gluon/01-attention-forward.py.

```python
import copy
import math
import torch
import triton
import pytest
import itertools
from dataclasses import dataclass, fields, replace

from triton.experimental import gluon
from triton.experimental.gluon import language as gl
from triton.experimental.gluon.language.nvidia import blackwell as bw
from triton.experimental.gluon.nvidia.blackwell import TensorDescriptor
from triton.experimental.gluon.language.nvidia.blackwell import (
    TensorMemoryLayout,
    allocate_tensor_memory,
    fence_async_shared,
    tensor_memory_descriptor,
    tensor_memory_descriptor_type,
    tma,
    mbarrier,
    tcgen05_mma,
    tcgen05_mma_barrier_count,
    tcgen05_commit,
)

# ===-----------------------------------------------------------------------===#
# Layout Utilities
# ===-----------------------------------------------------------------------===#


@gluon.constexpr_function
def get_mma_instr_shape(shape, element_ty):
    m = 128 if shape[0] >= 128 else 64
    n = 256 if shape[1] >= 256 else shape[1]
    k = 256 // element_ty.primitive_bitwidth
    return (m, n, k)


@gluon.constexpr_function
def get_split_dim(cga_layout, dim):
    return 1 << sum(b[dim] != 0 for b in cga_layout)


def get_mma_operand_cga_layout(layout, op_idx):
    assert op_idx in (0, 1)
    if not layout:
        return layout

    # 2CTA performs an outer product so bases are [1, 0] and [0, 1].
    assert layout[0] == (1, 0)
    first = (1, 0) if op_idx == 0 else (0, 1)

    # Broadcast along K (the reduction dimension).
    def broadcast(b):
        return (b[0], 0) if op_idx == 0 else (0, 2 * b[1])

    return (first, *map(broadcast, layout[1:]))


# ===-----------------------------------------------------------------------===#
# Data Abstractions
# ===-----------------------------------------------------------------------===#


@gluon.aggregate
class BarrierCounter:
    index: gl.tensor
    phase: gl.tensor
    num_barriers: gl.constexpr

    @gluon.must_use_result
    @gluon.jit
    def increment(self):
        if self.num_barriers == 1:
            return BarrierCounter(gl.to_tensor(0), self.phase ^ 1, self.num_barriers)
        next_index = self.index + 1
        rollover = next_index == self.num_barriers
        index = gl.where(rollover, 0, next_index)
        phase = gl.where(rollover, self.phase ^ 1, self.phase)
        return BarrierCounter(index, phase, self.num_barriers)


def Channel(T, alloc_fn):

    @gluon.aggregate
    class ChannelType:
        mem: T
        ready_bars: gl.shared_memory_descriptor
        empty_bars: gl.shared_memory_descriptor
        num_buffers: gl.constexpr

        @gluon.jit
        def alloc(shape: gl.constexpr, dtype: gl.constexpr, layout: gl.constexpr, num_buffers: gl.constexpr,
                  producer_two_ctas: gl.constexpr = False, consumer_two_ctas: gl.constexpr = False):
            mem = alloc_fn(dtype, [num_buffers] + shape, layout)
            ready_bars = mbarrier.allocate_mbarrier(batch=num_buffers, two_ctas=consumer_two_ctas)
            empty_bars = mbarrier.allocate_mbarrier(batch=num_buffers, two_ctas=producer_two_ctas)
            return ChannelType(mem, ready_bars, empty_bars, num_buffers)

        @gluon.jit
        def init(self, num_producers: gl.constexpr = 1, num_consumers: gl.constexpr = 1):
            for i in gl.static_range(self.num_buffers):
                mbarrier.init(self.ready_bars.index(i), count=num_producers)
                mbarrier.init(self.empty_bars.index(i), count=num_consumers)

        @gluon.jit
        def prime(self, num_consumers: gl.constexpr = 1):
            for i in gl.static_range(self.num_buffers):
                mbarrier.arrive(self.empty_bars.index(i), count=num_consumers)

        @gluon.jit
        def acquire_producer(self, counter):
            index, phase = counter.index, counter.phase
            mem = self.mem.index(index)
            ready_bar = self.ready_bars.index(index)
            empty_bar = self.empty_bars.index(index)

            mbarrier.wait(empty_bar, phase)
            return mem, ready_bar

        @gluon.jit
        def acquire_consumer(self, counter):
            index, phase = counter.index, counter.phase
            mem = self.mem.index(index)
            ready_bar = self.ready_bars.index(index)
            empty_bar = self.empty_bars.index(index)

            mbarrier.wait(ready_bar, phase)
            return mem, empty_bar

        @gluon.jit
        def create_counter(self):
            return BarrierCounter(gl.to_tensor(0), gl.to_tensor(0), self.num_buffers)

        @gluon.jit
        def create_producer(self):
            return Producer(self, self.create_counter())

        @gluon.jit
        def create_consumer(self):
            return Consumer(self, self.create_counter())
```

--------------------------------

### proton.record example

Source: https://triton-lang.org/main/_sources/dialects/ProtonOps.md.txt

Example usage of the `proton.record` operation, showing how to start and end a named event scope. Scope names cannot be reused within the same function.

```mlir
proton.record start "name0"
...
proton.record end "name0"
```

--------------------------------

### Examples of amdg.update_tensor_descriptor

Source: https://triton-lang.org/main/_sources/dialects/TritonAMDGPUOps.md.txt

Three usage examples: K-loop interior bumping tile position, prologue setting position and predicate, and peel epilogue installing real OOB extent.

```mlir
// K-loop interior: bump tile position only
%d1 = amdg.update_tensor_descriptor %d, add_offsets=[%c0, %BLOCK_K]
        : !tt.tensordesc<64x64xf16, #shared>
```

```mlir
// Prologue: position at first tile + set the predicate
%d1 = amdg.update_tensor_descriptor %d,
        add_offsets=[%pid_m_off, %c0],
        pred=%p
        : !tt.tensordesc<64x64xf16, #shared>
```

```mlir
// Peel epilogue: install real OOB extent for the partial last tile
%d1 = amdg.update_tensor_descriptor %d, set_bounds=[%m_remain, %k_remain]
        : !tt.tensordesc<64x64xf16, #shared>
```

--------------------------------

### Async Fence for Load Hazard

Source: https://triton-lang.org/main/_sources/getting-started/tutorials/gluon/tma.md.txt

Example of using fence_async_shared to prevent an async_load from starting while a shared memory load is in progress.

```python
value = smem.load()
fence_async_shared()
tma.async_load(desc, [0, 0], bar, smem)
```

--------------------------------

### Example: Using make_tensor_descriptor in a Triton kernel

Source: https://triton-lang.org/main/python-api/generated/triton.language.make_tensor_descriptor.html

This example defines a Triton kernel that uses tl.make_tensor_descriptor to create a descriptor for a 2D tensor, then loads and stores blocks using the descriptor. It also shows the required global memory allocation setup with triton.set_allocator and the kernel launch.

```python
@triton.jit
def inplace_abs(in_out_ptr, M, N, M_BLOCK: tl.constexpr, N_BLOCK: tl.constexpr):
    desc = tl.make_tensor_descriptor(
        in_out_ptr,
        shape=[M, N],
        strides=[N, 1],
        block_shape=[M_BLOCK, N_BLOCK],
    )

    moffset = tl.program_id(0) * M_BLOCK
    noffset = tl.program_id(1) * N_BLOCK

    value = desc.load([moffset, noffset])
    desc.store([moffset, noffset], tl.abs(value))

# TMA descriptors require a global memory allocation
def alloc_fn(size: int, alignment: int, stream: Optional[int]):
    return torch.empty(size, device="cuda", dtype=torch.int8)

triton.set_allocator(alloc_fn)

M, N = 256, 256
x = torch.randn(M, N, device="cuda")
M_BLOCK, N_BLOCK = 32, 32
grid = (M / M_BLOCK, N / N_BLOCK)
inplace_abs[grid](x, M, N, M_BLOCK, N_BLOCK)
```

--------------------------------

### Get top 4 elements from a 1D tensor using tl.topk

Source: https://triton-lang.org/main/python-api/generated/triton.language.topk.html

Example usage of triton.language.topk to retrieve the k largest elements from a 1D tensor. The example creates a tensor with values 0 to 15 and returns the top 4 elements, which are [15, 14, 13, 12].

```python
# Get top 4 elements from a 1D tensor
x = tl.arange(0, 16)
top4 = tl.topk(x, 4)  # Returns [15, 14, 13, 12]

```

--------------------------------

### Setup and imports for warp specialization tutorial

Source: https://triton-lang.org/main/_sources/getting-started/tutorials/gluon/warp-specialization.md.txt

Imports required libraries and defines helper functions to check GPU capabilities. Raises an error if the GPU is not Hopper or newer.

```python
import pytest
import torch
import triton
import importlib
from functools import partial
from triton.experimental import gluon
from triton.experimental.gluon import language as gl

from triton.experimental.gluon.nvidia.hopper import TensorDescriptor
from triton.experimental.gluon.language.nvidia.hopper import tma, mbarrier, fence_async_shared
from triton.experimental.gluon.language.nvidia.blackwell import (
    TensorMemoryLayout,
    tensor_memory_descriptor,
    allocate_tensor_memory,
    tcgen05_mma,
    tcgen05_commit,
)

if torch.cuda.is_available():
    from triton._C.libtriton import nvidia
    cublas_workspace = torch.empty(32 * 1024 * 1024, device="cuda", dtype=torch.uint8)
    cublas = nvidia.cublas.CublasLt(cublas_workspace)
else:
    cublas = None

# Re-use utilities from the previous tutorial.
t3 = importlib.import_module("03-async-copy")
t4 = importlib.import_module("04-tma")
t7 = importlib.import_module("07-persistence")


def is_hopper_or_newer():
    target = triton.runtime.driver.active.get_current_target()
    return target.backend == "cuda" and torch.cuda.get_device_capability()[0] >= 9


def is_blackwell():
    target = triton.runtime.driver.active.get_current_target()
    return target.backend == "cuda" and torch.cuda.get_device_capability()[0] == 10


if __name__ == "__main__" and not is_hopper_or_newer():
    raise RuntimeError("This tutorial requires Hopper or newer NVIDIA GPU")
```

--------------------------------

### Run TMA im2col multi-batch example and show expected output

Source: https://triton-lang.org/main/_sources/getting-started/tutorials/gluon/conv-im2col.md.txt

Runs the multi-batch im2col example and prints the expected output showing how loading 16 pixels wraps from batch 0 into batch 1 when starting at (N=0, H=1, W=3). The output confirms the row-major traversal across batch boundaries.

```python
if __name__ == "__main__":
    run_tma_im2col_multi_batch()
    print("\n" + "=" * 50 + "\n")
```

```text
Example 4: Loading across multiple batches
Input shape: [2, 4, 4, 32] (2 batches of 4x4 images)
Starting at: batch=0, H=1, W=3 (pixel 8 in input)

Loaded data (pixel_id, value):
  pixel  0: value  8
  pixel  1: value  9
  pixel  2: value 10
  pixel  3: value 11
  pixel  4: value 12
  pixel  5: value 13
  pixel  6: value 14
  pixel  7: value 15
  pixel  8: value 16
  pixel  9: value 17
  pixel 10: value 18
  pixel 11: value 19
  pixel 12: value 20
  pixel 13: value 21
  pixel 14: value 22
  pixel 15: value 23

==================================================

```

--------------------------------

### Set up device and configure Triton imports

Source: https://triton-lang.org/main/_downloads/3176accb6c7288b0e45f41d94eebacb9/06-fused-attention.ipynb

Initial imports and device capability helper functions. These are the foundational setup for the Triton implementation, determining whether the kernel runs on HIP, CUDA, Hopper, or Blackwell hardware. The helpers are used later to configure kernel specialization and descriptor options.

```python
import pytest
import torch
import os
```

```python
import triton
import triton.language as tl
from triton.tools.tensor_descriptor import TensorDescriptor

DEVICE = triton.runtime.driver.active.get_active_torch_device()
```

```python
def is_hip():
    return triton.runtime.driver.active.get_current_target().backend == "hip"
```

```python
def is_cuda():
    return triton.runtime.driver.active.get_current_target().backend == "cuda"
```

```python
def supports_host_descriptor():
    return is_cuda() and torch.cuda.get_device_capability()[0] >= 9
```

```python
def is_blackwell():
    return is_cuda() and torch.cuda.get_device_capability()[0] == 10
```

```python
def is_hopper():
    return is_cuda() and torch.cuda.get_device_capability()[0] == 9
```

--------------------------------

### PartitionedSharedLayout.__init__

Source: https://triton-lang.org/main/_sources/gluon/api/generated/triton.experimental.gluon.language.amd.cdna5.PartitionedSharedLayout.rst.txt

Initializes a new instance of the PartitionedSharedLayout class.

```APIDOC
## PartitionedSharedLayout.__init__

### Description
Initializes a new instance of the PartitionedSharedLayout class.

### Method
__init__

### Parameters
No parameters are documented in the source.

### Returns
No return value is documented in the source.
```

--------------------------------

### Imports and setup for tcgen05_copy tutorial

Source: https://triton-lang.org/main/getting-started/tutorials/gluon/tcgen05-copy.html

Import block and setup for the tutorial. Defines is_blackwell() guard and raises RuntimeError if not running on Blackwell. Imports required gluon and language modules, plus imports prior tutorial modules for reuse.

```python
import itertools
import importlib
import pytest
import triton
import torch
import triton.experimental.gluon as gluon
import triton.experimental.gluon.language as gl
from triton.experimental.gluon.nvidia.hopper import TensorDescriptor
from triton.experimental.gluon.language.nvidia.blackwell import (
    TensorMemoryLayout,
    tensor_memory_descriptor,
    allocate_tensor_memory,
    fence_async_shared,
    tcgen05_copy,
    tcgen05_commit,
    tcgen05_mma,
    mbarrier,
    tma,
)


def is_blackwell():
    target = triton.runtime.driver.active.get_current_target()
    return target.backend == "cuda" and torch.cuda.get_device_capability()[0] == 10


if __name__ == "__main__" and not is_blackwell():
    raise RuntimeError("This tutorial requires a Blackwell NVIDIA GPU")

# Re-use utilities from the previous tutorials.
t7 = importlib.import_module("07-persistence")
t8 = importlib.import_module("08-warp-specialization")

```

--------------------------------

### Conv fprop: host-side entry point in Python

Source: https://triton-lang.org/main/getting-started/examples/gluon/conv-fprop.html

The page covers the host-side setup for a Triton convolution forward pass on Hopper/Blackwell, starting from autotuning configuration and input validation through TMA descriptor creation and grid computation. The final launch function is provided as an incomplete stub.

```python
conv2d_fprop_autotuned_kernel = triton.autotune(
    configs=conv2d_fprop_get_configs(pre_hook=conv2d_fprop_tma_set_block_size_hook),
    key=["out_h", "out_w", "stride_h", "stride_w"],
)(conv2d_fprop_kernel)
```

```python
def _prepare_conv_fprop_inputs(input_tensor, weight_tensor, stride, padding):
    N, H, W, Ci = input_tensor.shape
    Co, R, S, Ci_w = weight_tensor.shape
    assert Ci == Ci_w, "Input and weight channel dimensions must match"
    if input_tensor.dtype != TORCH_GEMM_DTYPE or weight_tensor.dtype != TORCH_GEMM_DTYPE:
        raise ValueError(
            f"conv2d_fprop expects bfloat16 input/weight tensors, got {input_tensor.dtype} and {weight_tensor.dtype}")
    stride_h, stride_w = normalize_2d(stride, "stride")
    pad_h, pad_w = normalize_2d(padding, "padding")
    if stride_h <= 0 or stride_w <= 0:
        raise ValueError(f"stride must be positive, got {(stride_h, stride_w)}")
    if pad_h < 0 or pad_w < 0:
        raise ValueError(f"padding must be non-negative, got {(pad_h, pad_w)}")

    # The Hopper/Blackwell TMA path requires outer strides to be 16-byte aligned.
    # For NHWC/OHWI bf16 tensors this means padding the channel dimension for
    # narrow inputs such as RGB (Ci=3).
    input_tensor, weight_tensor = maybe_pad_ci_for_tma(input_tensor, weight_tensor)
    N, H, W, Ci = input_tensor.shape
    Co, R, S, Ci_w = weight_tensor.shape

    out_h = (H + 2 * pad_h - R) // stride_h + 1
    out_w = (W + 2 * pad_w - S) // stride_w + 1
    if out_h <= 0 or out_w <= 0:
        raise ValueError("Invalid convolution geometry: computed output size "
                         f"({out_h}, {out_w}) from H={H}, W={W}, R={R}, S={S}, "
                         f"stride={(stride_h, stride_w)}, padding={(pad_h, pad_w)}.")

    output = torch.empty((N, out_h, out_w, Co), device=input_tensor.device, dtype=TORCH_GEMM_DTYPE)
    return input_tensor, weight_tensor, output, N, H, W, Ci, Co, R, S, out_h, out_w, stride_h, stride_w, pad_h, pad_w
```

```python
def _make_conv_fprop_descriptors(input_tensor, weight_tensor, out_h, out_w, stride_h, stride_w, pad_h, pad_w,
                                 input_block_shape, weight_block_shape):
    # TMA im2col descriptor for input: [N, H, W, Ci] in NHWC
    #
    # The pixel_box defines the access boundary per batch:
    #   Lower = pixel_box_lower_corner + offsets
    #   Upper = [H, W] + pixel_box_upper_corner + offsets
    # With element_strides = [1, stride_h, stride_w, 1], TMA steps by the
    # per-dimension convolution stride between output pixels. The window must
    # contain exactly out_h * out_w pixels per batch:
    #   pixels_h = floor((window_h - 1) / stride_h) + 1 = out_h
    #   pixels_w = floor((window_w - 1) / stride_w) + 1 = out_w
    #   => window_h = (out_h - 1) * stride_h + 1
    #   => window_w = (out_w - 1) * stride_w + 1
    _, H, W, _ = input_tensor.shape
    upper_h = (out_h - 1) * stride_h + 1 - H - pad_h
    upper_w = (out_w - 1) * stride_w + 1 - W - pad_w

    input_layout = gl.NVMMASharedLayout.get_default_for(input_block_shape, GL_GEMM_DTYPE)
    in_desc = TensorDescriptorIm2Col(
        base=input_tensor,
        shape=list(input_tensor.shape),
        strides=list(input_tensor.stride()),
        block_shape=input_block_shape,
        layout=input_layout,
        padding="zero",
        element_strides=[1, stride_h, stride_w, 1],
        pixel_box_lower_corner=[-pad_h, -pad_w],
        pixel_box_upper_corner=[upper_h, upper_w],
    )

    # TMA tiled descriptor for weight: (Co, R*S*Ci) = (N_GEMM, K_GEMM)
    Co, R, S, Ci = weight_tensor.shape
    weight_reshaped = weight_tensor.reshape(Co, R * S * Ci)
    weight_layout = gl.NVMMASharedLayout.get_default_for(weight_block_shape, GL_GEMM_DTYPE)
    weight_desc = TensorDescriptor.from_tensor(weight_reshaped, weight_block_shape, weight_layout)
    return in_desc, weight_desc
```

```python
def _make_grid(num_sms, M_GEMM, N_GEMM):

    def grid(meta):
        num_tiles = triton.cdiv(M_GEMM, meta["BLOCK_M"]) * triton.cdiv(N_GEMM, meta["BLOCK_N"])
        return (min(num_sms, num_tiles), )

    return grid
```

```python
def _launch_conv(
    kernel,
    grid,
    *,
    in_desc,
    weight_desc,
    output,
    N,
    H,
    W,
    Ci,
    Co,
    R,
    S,
    out_h,
    out_w,
    stride_h,
    stride_w,
    pad_h,
    pad_w,
    kernel_meta=None,
):
    if kernel_meta is None:
        kernel_meta = {}
```

--------------------------------

### Benchmark configuration setup for fused attention

Source: https://triton-lang.org/main/getting-started/tutorials/06-fused-attention.html

Defines benchmark configurations for fused attention, varying head dimension, mode (fwd/bwd), causality, and warp specialization. Enables warp specialization only for forward mode on Hopper (non-causal) or Blackwell. Includes FP8 support if available and FlashAttention-2 if installed.

```python
TORCH_HAS_FP8 = hasattr(torch, 'float8_e5m2')
BATCH, N_HEADS = 4, 32
# vary seq length for fixed head and batch=4
configs = []
for HEAD_DIM in [64, 128]:
    for mode in ["fwd", "bwd"]:
        for causal in [True, False]:
            # Enable warpspec for causal fwd on Hopper
            enable_ws = mode == "fwd" and (is_blackwell() or (is_hopper() and not causal))
            for warp_specialize in [False, True] if enable_ws else [False]:
                configs.append(
                    triton.testing.Benchmark(
                        x_names=["N_CTX"],
                        x_vals=[2**i for i in range(10, 15)],
                        line_arg="provider",
                        line_vals=["triton-fp16"] + (["triton-fp8"] if TORCH_HAS_FP8 else []) +
                        (["flash"] if HAS_FLASH else []),
                        line_names=["Triton [FP16]"] + (["Triton [FP8]"] if TORCH_HAS_FP8 else []) +
                        (["Flash-2"] if HAS_FLASH else []),
                        styles=[("red", "-"), ("blue", "-"), ("green", "-")],
                        ylabel="TFLOPS",
                        plot_name=
                        f"fused-attention-batch{BATCH}-head{N_HEADS}-d{HEAD_DIM}-{mode}-causal={causal}-warp_specialize={warp_specialize}",
                        args={
                            "H": N_HEADS,
                            "BATCH": BATCH,
                            "HEAD_DIM": HEAD_DIM,
                            "mode": mode,
                            "causal": causal,
                            "warp_specialize": warp_specialize,
                        },
                    ))
```

--------------------------------

### Setup and imports for TMA gather/scatter tutorial

Source: https://triton-lang.org/main/_sources/getting-started/tutorials/gluon/tma-gather-scatter.md.txt

Imports required modules and checks for Blackwell GPU. Raises an error if not on Blackwell.

```python
import sys
import pytest
import torch
import triton
import importlib
import triton.experimental.gluon as gluon
import triton.experimental.gluon.language as gl
from triton._C.libtriton import ir, gluon_ir

from triton.experimental.gluon.nvidia.hopper import TensorDescriptor
from triton.experimental.gluon.language.nvidia.blackwell import (tma, mbarrier, fence_async_shared)


def is_blackwell():
    target = triton.runtime.driver.active.get_current_target()
    return target.backend == "cuda" and torch.cuda.get_device_capability()[0] == 10


if __name__ == "__main__" and not is_blackwell():
    raise RuntimeError("This tutorial requires a Blackwell NVIDIA GPU")

# Re-use utilities from the previous tutorials.
t7 = importlib.import_module("07-persistence")
```

--------------------------------

### Example benchmark output

Source: https://triton-lang.org/main/getting-started/tutorials/gluon/wgmma.html

Sample output showing the performance metrics for various WGMMA configuration parameters.

```text
BLOCK_M BLOCK_N BLOCK_K num_warps instr_shape_n occupancy time (ms) tflops/s
    128     256     256         8           256         1      5.34   412.14
    256     128     256         8           128         1      5.67   387.74
     64     256     128         4           256         2      4.64   474.03
     64     128     256         4           128         2      6.18   355.60
    128     128     128         4           128         2      4.98   441.88
    128     128     128         8           128         2      5.79   380.08
```

--------------------------------

### Start of mma_scaled_tcgen05_copy_kernel with tcgen05_copy

Source: https://triton-lang.org/main/_sources/getting-started/tutorials/gluon/tcgen05-mma-scaled.md.txt

This is the beginning of a Gluon JIT kernel that uses tcgen05_copy to load scale factors into tensor memory. It reuses the unchanged setup from the previous mma_scaled_packed_block_kernel, allocating shared memory for A and B, tensor memory for scales and accumulator, and defining constexpr values for FP4 handling.

```python
@gluon.jit
def mma_scaled_tcgen05_copy_kernel(a_desc, b_desc, c_desc, a_scale_desc, b_scale_desc, VEC_SIZE: gl.constexpr):
    # ======= Begin unchanged code from `mma_scaled_packed_block_kernel` =======
    A_IS_FP4: gl.constexpr = a_desc.dtype == gl.uint8
    B_IS_FP4: gl.constexpr = b_desc.dtype == gl.uint8
    A_ELEM_PER_BYTE: gl.constexpr = 2 if A_IS_FP4 else 1
    B_ELEM_PER_BYTE: gl.constexpr = 2 if B_IS_FP4 else 1
    BLOCK_M: gl.constexpr = c_desc.block_type.shape[0]
    BLOCK_N: gl.constexpr = c_desc.block_type.shape[1]
    BLOCK_K: gl.constexpr = a_desc.block_type.shape[1] * A_ELEM_PER_BYTE
    K = a_desc.shape[1] * A_ELEM_PER_BYTE

    a_smem = gl.allocate_shared_memory(a_desc.dtype, a_desc.block_type.shape, a_desc.layout)
    b_smem = gl.allocate_shared_memory(b_desc.dtype, b_desc.block_type.shape, b_desc.layout)

    scale_layout: gl.constexpr = TensorMemoryScalesLayout()
    a_scale_tmem = allocate_tensor_memory(a_scale_desc.dtype, [BLOCK_M, BLOCK_K // VEC_SIZE], scale_layout)
    b_scale_tmem = allocate_tensor_memory(b_scale_desc.dtype, [BLOCK_N, BLOCK_K // VEC_SIZE], scale_layout)
    tmem_layout: gl.constexpr = TensorMemoryLayout([BLOCK_M, BLOCK_N], col_stride=1)
    acc_tmem = allocate_tensor_memory(gl.float32, [BLOCK_M, BLOCK_N], tmem_layout)
    use_acc = False

```

--------------------------------

### Run matmul benchmark and profile from command line

Source: https://triton-lang.org/main/_downloads/b5537c431cd73b8ee3c32cdf90c58104/09-persistent-matmul.ipynb

Main entry point that parses CLI arguments, validates small and large matmul cases, runs benchmarks with Triton profiling, and displays the profile.

```python
if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("-K", type=int, required=False, default=512)
    parser.add_argument("--K_range", type=int, nargs=2)
    parser.add_argument("--K_step", type=int, default=512)
    parser.add_argument("--prec", type=str, choices=["fp8", "fp16"], default="fp16")
    args = parser.parse_args()

    if args.prec == 'fp8' and (not hasattr(torch, "float8_e4m3fn") or not is_cuda()):
        print("This example requires CUDA/HIP with fp8 support.")
    else:
        dtype = torch.float8_e4m3fn if args.prec == 'fp8' else torch.float16

        if args.K and args.K_range is None:
            args.K_range = [args.K, args.K]
            args.K_step = 1  # doesn't matter as long as it's not 0

        torch.manual_seed(0)

        validate(32, 32, 32, dtype)
        validate(8192, 8192, args.K_range[0], dtype)

        proton.start("matmul", hook="triton")
        proton.deactivate()
        for K in range(args.K_range[0], args.K_range[1] + 1, args.K_step):
            bench(K, dtype)
        proton.finalize()
        show_profile(args.prec, "matmul")
```

--------------------------------

### triton.language.arange(start, end, _semantic=None)

Source: https://triton-lang.org/main/python-api/generated/triton.language.arange.html

Returns contiguous values within the half-open interval [start, end). The difference end - start must be less than or equal to TRITON_MAX_TENSOR_NUMEL = 1048576. Both start and end must be powers of two, with end greater than start.

```APIDOC
## triton.language.arange(start, end, _semantic=None)

### Description
Returns contiguous values within the half-open interval `[start, end)`. The difference `end - start` must be less than or equal to `TRITON_MAX_TENSOR_NUMEL = 1048576`.

### Parameters
- **start** (int32) - Required - Start of the interval. Must be a power of two.
- **end** (int32) - Required - End of the interval. Must be a power of two greater than `start`.

### Constraints
- `end - start` must be less than or equal to `TRITON_MAX_TENSOR_NUMEL = 1048576`.
- `start` must be a power of two.
- `end` must be a power of two greater than `start`.
```

--------------------------------

### Initialize softmax tile and scheduler loop

Source: https://triton-lang.org/main/_sources/getting-started/examples/gluon/attention-forward.md.txt

Sets up a softmax tile by creating consumers/producers, acquiring the correlation barrier, and iterating over programs from the scheduler. It initializes the running maximum and sum with packed FMA2-friendly layouts.

```python
@gluon.jit
def _softmax_tile(tile_id: gl.constexpr, config, M, desc_o, STAGE: gl.constexpr,  #
                  s_chnl, corr_chnl, exp_turnstile, use_tmem_red: gl.constexpr):
    qk_slice_dim1: gl.constexpr = gl.SliceLayout(1, config.qk_layout)
    sum_layout: gl.constexpr = _get_split_n_layout(config.qk_layout)

    s_consumer = s_chnl.create_consumer()
    corr_producer = corr_chnl.create_producer()
    _, corr_bar, corr_producer = corr_producer.acquire()

    scheduler = ProgramScheduler.create(config)
    for pid in range(scheduler.start_pid, scheduler.num_tiles, config.NUM_SMS):
        prog = scheduler.get_program(pid)

        offs_m = prog.start_m * config.BLOCK_M
        offs_m += gl.arange(tile_id * config.SPLIT_M, (1 + tile_id) * config.SPLIT_M)

        m_i = gl.full([config.SPLIT_M], -float("inf"), gl.float32, qk_slice_dim1)
        # Keep the two row sums adjacent so their updates use packed FMA2.
        l_i = gl.full([config.SPLIT_M], 0.0, gl.float32, gl.SliceLayout(1, sum_layout))
        l_i = gl.join(l_i, l_i)
```
Tokens:
Raw