Rocm Kernels

ROCm Triton Kernels for Diffusers & Transformers

This skill provides patterns and guidance for developing optimized Triton kernels targeting AMD GPUs (MI355X, R9700) on ROCm, for use with HuggingFace diffusers (LTX-Video, SD3, FLUX) and transformers libraries.

Quick Start

Diffusers (LTX-Video)

Inject optimized kernels into LTX-Video pipeline:

import os
os.environ['TRITON_HIP_USE_BLOCK_PINGPONG'] = '1'
os.environ['TRITON_HIP_USE_ASYNC_COPY'] = '1'

from diffusers import LTXPipeline
pipe = LTXPipeline.from_pretrained("Lightricks/LTX-Video", torch_dtype=torch.bfloat16)
pipe.to("cuda")  # ROCm uses same API via HIP
inject_optimized_kernels(pipe)  # BEFORE CPU offloading
pipe.enable_model_cpu_offload()

For a minimal integration example (~150 lines):

python scripts/ltx_kernel_injection_example.py

Isolated Kernel Micro-benchmarks

# All 4 kernels: correctness + performance + bandwidth
python scripts/benchmark_kernels.py

# Single kernel
python scripts/benchmark_kernels.py --kernel rmsnorm
python scripts/benchmark_kernels.py --kernel rope
python scripts/benchmark_kernels.py --kernel geglu
python scripts/benchmark_kernels.py --kernel adaln

End-to-End Pipeline Benchmark

# Compare baseline vs Triton vs torch.compile
python scripts/benchmark_e2e.py --mode all

# Quick test
python scripts/benchmark_e2e.py --mode triton --num-frames 9 --steps 5

# Save results for comparison
python scripts/benchmark_e2e.py --mode all --output-json results.json

Target Model: LTX-Video

Architecture Overview

Total RMSNorm modules: 168 (56 with weights, 112 without)

Target Kernels

Supported Hardware

See MI355X guide | R9700 guide

When This Skill Applies

Use this skill when:

• Writing Triton kernels for RMSNorm, RoPE, GEGLU, AdaLN on AMD GPUs
• Integrating custom kernels with diffusers pipelines (LTX-Video, SD3, FLUX)
• Benchmarking kernel performance against PyTorch baseline on ROCm
• Optimizing existing kernels for MI355X or R9700 architecture
• Debugging ROCm/HIP-specific kernel issues

Critical ROCm Constraints

Things That DON'T Work on AMD

# FORBIDDEN - CUDA only, NOT available on ROCm
tl.libdevice.tanh(x)          # Use manual formula below
tl.libdevice.log1p(x)         # Use: tl.log(1.0 + x)
tl.math.tanh(x)               # Also NOT available on ROCm Triton

# Manual tanh (ONLY reliable method on ROCm):
e2x = tl.exp(2.0 * x)
tanh_x = (e2x - 1.0) / (e2x + 1.0)

# FORBIDDEN - Triton limitations on ROCm
break / continue               # Use: tl.where()
min(a, b) / max(a, b)          # Use: tl.minimum(a, b) / tl.maximum(a, b)

Mandatory Environment Variables

import os
os.environ['TRITON_HIP_USE_BLOCK_PINGPONG'] = '1'
os.environ['TRITON_HIP_USE_ASYNC_COPY'] = '1'

Core Kernel Implementations

1. RMSNorm (Core Optimization Target)

Row-wise reduction pattern. 168 instances in LTX-Video, ~5% of total compute.

CRITICAL: Do NOT autotune BLOCK_D. Autotune may pick BLOCK_D < D, causing partial row processing and wrong results. Always compute BLOCK_D = triton.next_power_of_2(D) in the Python wrapper.

@triton.jit
def rmsnorm_kernel(
    x_ptr, weight_ptr, out_ptr,
    stride_x, D,
    eps: tl.constexpr,
    HAS_WEIGHT: tl.constexpr,
    BLOCK_D: tl.constexpr,
):
    row = tl.program_id(0)
    offs = tl.arange(0, BLOCK_D)
    mask = offs < D
    x = tl.load(x_ptr + row * stride_x + offs, mask=mask, other=0.0).to(tl.float32)

    variance = tl.sum(x * x, axis=0) / D
    rms_inv = tl.rsqrt(variance + eps)

    if HAS_WEIGHT:
        w = tl.load(weight_ptr + offs, mask=mask, other=1.0).to(tl.float32)
        out = x * rms_inv * w
    else:
        out = x * rms_inv

    tl.store(out_ptr + row * stride_x + offs, out.to(x.dtype), mask=mask)


def triton_rmsnorm(x, weight=None, eps=1e-6):
    x_2d = x.contiguous().view(-1, x.shape[-1])
    out = torch.empty_like(x_2d)
    M, D = x_2d.shape
    has_weight = weight is not None
    if not has_weight:
        weight = torch.empty(0, device=x.device)

    BLOCK_D = triton.next_power_of_2(D)
    num_warps = 4 if BLOCK_D <= 1024 else (8 if BLOCK_D <= 4096 else 16)
    rmsnorm_kernel[(M,)](
        x_2d, weight, out, x_2d.stride(0), D, eps, has_weight,
        BLOCK_D=BLOCK_D, num_warps=num_warps, num_stages=2,
    )
    return out.view_as(x)

LTX-Video pitfall: Weight may be None!

has_weight = hasattr(module, 'weight') and module.weight is not None

2. RoPE 3D (Video Position Encoding)

Element-wise pattern. LTX-Video splits head_dim into temporal + spatial components.

CRITICAL: cos/sin have shape `[seq_len, head_dim]`. When grid flattens batch dimension (batch * seq_len), use pid_s % seq_len to index cos/sin, otherwise batch > 1 causes OOB GPU crash.

@triton.jit
def rope_3d_kernel(
    qk_ptr, cos_ptr, sin_ptr, out_ptr,
    seq_len, num_heads, head_dim,
    stride_s, stride_h, stride_d,
    BLOCK_HD: tl.constexpr,
):
    pid_s = tl.program_id(0)  # batch * seq_len
    pid_h = tl.program_id(1)  # head index
    half_dim = head_dim // 2
    offs = tl.arange(0, BLOCK_HD)
    mask = offs < half_dim

    base = pid_s * stride_s + pid_h * stride_h
    x0 = tl.load(qk_ptr + base + offs, mask=mask, other=0.0).to(tl.float32)
    x1 = tl.load(qk_ptr + base + half_dim + offs, mask=mask, other=0.0).to(tl.float32)

    seq_idx = pid_s % seq_len  # wrap for batch > 1
    cos_val = tl.load(cos_ptr + seq_idx * head_dim + offs, mask=mask, other=1.0).to(tl.float32)
    sin_val = tl.load(sin_ptr + seq_idx * head_dim + offs, mask=mask, other=0.0).to(tl.float32)

    out0 = x0 * cos_val - x1 * sin_val
    out1 = x0 * sin_val + x1 * cos_val

    tl.store(out_ptr + base + offs, out0.to(x0.dtype), mask=mask)
    tl.store(out_ptr + base + half_dim + offs, out1.to(x0.dtype), mask=mask)


def triton_rope_3d(qk, cos, sin):
    qk = qk.contiguous()
    out = torch.empty_like(qk)
    batch, seq_len, num_heads, head_dim = qk.shape
    qk_flat = qk.view(batch * seq_len, num_heads, head_dim)
    out_flat = out.view(batch * seq_len, num_heads, head_dim)
    BLOCK_HD = triton.next_power_of_2(head_dim // 2)
    num_warps = 4 if BLOCK_HD <= 64 else 8
    rope_3d_kernel[(batch * seq_len, num_heads)](
        qk_flat, cos, sin, out_flat,
        seq_len, num_heads, head_dim,
        qk_flat.stride(0), qk_flat.stride(1), qk_flat.stride(2),
        BLOCK_HD=BLOCK_HD, num_warps=num_warps, num_stages=2,
    )
    return out

3. GEGLU (For SD3/FLUX, NOT LTX-Video)

Element-wise gated activation. Input [batch, seq, 2*hidden] → Output [batch, seq, hidden].

Same BLOCK_SIZE rule: compute dynamically, do NOT autotune.

@triton.jit
def geglu_kernel(
    input_ptr, output_ptr,
    stride_in, stride_out, hidden_size,
    BLOCK_H: tl.constexpr,
):
    row = tl.program_id(0)
    offs = tl.arange(0, BLOCK_H)
    mask = offs < hidden_size

    gate = tl.load(input_ptr + row * stride_in + offs, mask=mask, other=0.0).to(tl.float32)
    value = tl.load(input_ptr + row * stride_in + hidden_size + offs, mask=mask, other=0.0).to(tl.float32)

    # GELU approx — manual tanh (tl.math.tanh NOT available on ROCm)
    k = 0.7978845608028654  # sqrt(2/pi)
    tanh_arg = k * (gate + 0.044715 * gate * gate * gate)
    e2x = tl.exp(2.0 * tanh_arg)
    tanh_val = (e2x - 1.0) / (e2x + 1.0)
    gate_gelu = 0.5 * gate * (1.0 + tanh_val)
    result = gate_gelu * value

    tl.store(output_ptr + row * stride_out + offs, result.to(gate.dtype), mask=mask)


def triton_geglu(x):
    x = x.contiguous()
    *batch_dims, double_h = x.shape
    hidden_size = double_h // 2
    x_2d = x.view(-1, double_h)
    M = x_2d.shape[0]
    out = torch.empty(M, hidden_size, device=x.device, dtype=x.dtype)
    BLOCK_H = triton.next_power_of_2(hidden_size)
    num_warps = 4 if BLOCK_H <= 1024 else (8 if BLOCK_H <= 4096 else 16)
    geglu_kernel[(M,)](
        x_2d, out, x_2d.stride(0), out.stride(0), hidden_size,
        BLOCK_H=BLOCK_H, num_warps=num_warps, num_stages=2,
    )
    return out.view(*batch_dims, hidden_size)

Warning: LTX-Video uses GELU, NOT GEGLU. GEGLU is for SD3/FLUX.

4. AdaLN (Adaptive Layer Normalization for DiT)

Fused normalization + conditioning: norm(x) * weight * (1 + scale) + shift

Same BLOCK_D rule: compute dynamically.

@triton.jit
def adaln_kernel(
    x_ptr, weight_ptr, scale_ptr, shift_ptr, out_ptr,
    stride_x, stride_cond, D,
    eps: tl.constexpr,
    BLOCK_D: tl.constexpr,
):
    row = tl.program_id(0)
    offs = tl.arange(0, BLOCK_D)
    mask = offs < D
    x = tl.load(x_ptr + row * stride_x + offs, mask=mask, other=0.0).to(tl.float32)

    variance = tl.sum(x * x, axis=0) / D
    rms_inv = tl.rsqrt(variance + eps)
    x_norm = x * rms_inv

    w = tl.load(weight_ptr + offs, mask=mask, other=1.0).to(tl.float32)
    scale = tl.load(scale_ptr + row * stride_cond + offs, mask=mask, other=0.0).to(tl.float32)
    shift = tl.load(shift_ptr + row * stride_cond + offs, mask=mask, other=0.0).to(tl.float32)

    out = x_norm * w * (1.0 + scale) + shift
    tl.store(out_ptr + row * stride_x + offs, out.to(x.dtype), mask=mask)


def triton_adaln(x, weight, scale, shift, eps=1e-6):
    x_flat = x.contiguous().view(-1, x.shape[-1])
    scale_flat = scale.contiguous().view(-1, x.shape[-1])
    shift_flat = shift.contiguous().view(-1, x.shape[-1])
    out = torch.empty_like(x_flat)
    M, D = x_flat.shape
    BLOCK_D = triton.next_power_of_2(D)
    num_warps = 4 if BLOCK_D <= 1024 else (8 if BLOCK_D <= 4096 else 16)
    adaln_kernel[(M,)](
        x_flat, weight, scale_flat, shift_flat, out,
        x_flat.stride(0), scale_flat.stride(0), D, eps,
        BLOCK_D=BLOCK_D, num_warps=num_warps, num_stages=2,
    )
    return out.view_as(x)

Diffusers Integration

See [diffusers-integration.md](references/diffusers-integration.md) for the complete guide.

Minimal Integration Pattern

def patch_rmsnorm_modules(model):
    """Patch all RMSNorm modules to use custom Triton kernel."""
    for name, module in model.named_modules():
        if type(module).__name__ == 'RMSNorm':
            eps = getattr(module, 'eps', 1e-6)
            has_weight = hasattr(module, 'weight') and module.weight is not None
            if has_weight:
                def make_forward(mod, epsilon):
                    def forward(x):
                        return triton_rmsnorm(x, mod.weight, eps=epsilon)
                    return forward
                module.forward = make_forward(module, eps)
            else:
                def make_forward(epsilon):
                    def forward(x):
                        w = torch.ones(x.shape[-1], device=x.device, dtype=x.dtype)
                        return triton_rmsnorm(x, w, eps=epsilon)
                    return forward
                module.forward = make_forward(eps)

pipe = LTXPipeline.from_pretrained("Lightricks/LTX-Video", torch_dtype=torch.bfloat16)
pipe.to("cuda")
patch_rmsnorm_modules(pipe.transformer)
pipe.enable_model_cpu_offload()

Diffusers Critical Pitfalls

1. RMSNorm weight may be None — LTX-Video uses elementwise_affine=False 2. Diffusers RMSNorm != torch.nn.RMSNorm — Use type(module).__name__ not isinstance() 3. LTX-Video uses GELU, not GEGLU — Don't patch GEGLU for LTX-Video 4. Inject BEFORE CPU offloading — inject_kernels() then enable_model_cpu_offload()

Performance Expectations

Micro-benchmark Results (MI355X, BF16)

RMSNorm bandwidth utilization: 3554 GB/s (MI355X theoretical: 8 TB/s, ~44%).

End-to-End LTX-Video (MI355X, 25 frames, 30 steps)

Key finding: MI355X Triton E2E speedup (22%) is significantly higher than H100 CUDA reference (6%), because MI355X's default PyTorch RMSNorm path has more room for optimization.

Micro-benchmark Results (R9700, BF16)

RMSNorm bandwidth utilization: 483 GB/s (R9700 theoretical: ~608 GB/s, ~79%).

R9700 speedups are higher than MI355X because PyTorch's default RDNA4 backend is less mature, leaving more room for Triton optimization. The bandwidth utilization (79%) is also significantly better than MI355X (44%).

End-to-End LTX-Video (R9700, 25 frames, 30 steps)

Reviewer-facing benchmark files for this comparison live in examples/ltx-video-benchmark/, including:

• Summary table with gen_time_s, time_per_step_s, peak_memory_gb, and speedup
• Consolidated JSON results in examples/ltx-video-benchmark/benchmark_results.json
• OpenCode run outputs in examples/ltx-video-benchmark/trace/opencode_live/results.json
• OpenCode parsed trace in examples/ltx-video-benchmark/trace/opencode_live/opencode_trace_result.json

R9700 Additional Validation

CUDA Reference (H100, for comparison)

H100 E2E: ~6% (RMSNorm is ~5% of total compute).

Optimization Targets

Common Issues on ROCm

See troubleshooting.md for all common issues.

Performance Profiling

rocprof --stats python your_kernel.py
rocprofv3 -i metrics.txt python your_kernel.py
rocm-bandwidth-test
rocminfo | grep -E "Name|Compute Unit|Wavefront"

See Also

Benchmark & Test Scripts

• benchmark_kernels.py - Micro-benchmark all 4 kernels (correctness + perf + bandwidth)
• benchmark_e2e.py - End-to-end LTX-Video pipeline benchmark (baseline vs Triton vs compile)
• sweep_num_warps.py - num_warps sweep for R9700 Wave32 optimization
• ltx_kernel_injection_example.py - Minimal diffusers injection example
• transformers_injection_example.py - Minimal transformers injection example
• huggingface_kernels_example.py - HuggingFace Kernels Hub integration example

Integration Guides

• diffusers-integration.md - LTX-Video pipeline integration
• transformers-integration.md - LLaMA/Mistral/Qwen integration
• huggingface-kernels-integration.md - HuggingFace Kernels Hub (get_kernel)
• kernel-templates.md - Complete Triton kernel templates (incl. GEMM with XCD Swizzle)

GPU Optimization Guides

• mi355x-optimization-guide.md - MI355X (gfx950) deep dive
• r9700-optimization-guide.md - R9700 (RDNA4) deep dive

Reference

• troubleshooting.md - Common issues and solutions
• kernelbench-classification.md - KernelBench operator taxonomy
• skill-evaluation-methodology.md - How to evaluate and improve skills
• kernel-agent-knowledge-base.md - Knowledge from kernel-agent project

External Resources

• Triton Documentation
• ROCm Documentation
• HuggingFace Kernels Hub
• LTX-Video on HuggingFace
• HuggingFace Diffusers

Changelog

v0.3 (2026-04-10)

Added

• Consolidated benchmark package: examples/ltx-video-benchmark/benchmark_results.json now provides a single JSON artifact with formal R9700 benchmark runs and summary statistics.
• Live harness benchmark trace:
• examples/ltx-video-benchmark/trace/opencode_live/opencode_trace_result.json
• Supporting run output in examples/ltx-video-benchmark/trace/opencode_live/results.json
• Unified examples documentation: examples/ltx-video-benchmark/README.md is now the single reviewer-facing README for the examples package.
• Dependency lock-in entrypoint: scripts/requirements.txt is used as the canonical install file for benchmark/integration examples.

Changed

• E2E benchmark CLI/output semantics (scripts/benchmark_e2e.py):
• Standardized output flag to --output-dir
• Updated docs and usage examples to match current examples layout
• Structured outputs under examples/ltx-video-benchmark/ and trace/*_live/
• Examples content structure:
• Merged previous top-level examples README content into examples/ltx-video-benchmark/README.md
• Updated trace/result paths in docs to current live locations
• Script requirement notes:
• Updated benchmark/example script docstrings to install dependencies from scripts/requirements.txt

Fixed

• Removed ambiguity between "harness replay" and "live benchmark execution" by recording explicit live_benchmark: true trace outputs with executed command/config.
• Aligned reviewer-facing materials to use one benchmark documentation entrypoint and one consolidated formal benchmark JSON.

Removed

• Legacy split/duplicate examples documentation and outdated trace path references in the examples package.
• Codex-specific benchmark trace artifacts and non-essential video artifacts from the reviewer package.

v0.2 (2026-03-12)

Added

• Transformers integration: references/transformers-integration.md — LLaMA/Mistral/Qwen RMSNorm patching, Flash Attention 2, epsilon handling differences
• Transformers injection script: scripts/transformers_injection_example.py — minimal runnable example (~150 lines)
• HuggingFace Kernels Hub integration: references/huggingface-kernels-integration.md — get_kernel, has_kernel, publishing, ROCm compatibility notes
• HuggingFace Kernels example script: scripts/huggingface_kernels_example.py — Hub loading, benchmarking, model integration with fallback
• GEMM template with XCD swizzle: Template 5 in kernel-templates.md — full GEMM kernel with XCD swizzle for MI355X, L2 cache grouping, autotune configs, Python API, and benchmark
• CHANGELOG.md: Version tracking for skill iterations

Fixed

• Broken cross-references: "Template 2" for GEMM → corrected to "Template 5" in troubleshooting.md, kernelbench-classification.md, and skill-evaluation-methodology.md
• R9700 Memory Bandwidth: filled in ~608 GB/s (was TBD) in SKILL.md

Updated

• SKILL.md See Also section: added new integration guides, scripts, and Hub links
• SKILL.md argument-hint: added gemm, transformers, huggingface-kernels, get_kernel
• manifest.txt: added all new files

v0.1 (2026-03-10)

Added

• Initial skill with SKILL.md, 4 kernel templates (RMSNorm, RoPE 3D, GEGLU, AdaLN)
• MI355X and R9700 GPU optimization guides
• Diffusers integration guide (LTX-Video)
• Troubleshooting guide (14 ROCm-specific issues)
• Benchmark scripts: micro-benchmark (benchmark_kernels.py) and E2E (benchmark_e2e.py)
• LTX-Video injection example (ltx_kernel_injection_example.py)
• KernelBench classification and evaluation methodology docs
• Kernel-agent knowledge base

ROCm Kernels Examples: LTX-Video Benchmark

This is the consolidated reviewer-facing README for ROCm kernels benchmark examples.

It answers:

Could we also see some numbers with and without these kernels?

Dependencies

python -m pip install -r skills/rocm-kernels/scripts/requirements.txt

What is included

• A formal benchmark table comparing baseline and triton
• A supplemental compile run for reference
• A single consolidated JSON file: benchmark_results.json
• A single live harness trace from OpenCode

Benchmark configuration

• Model: Lightricks/LTX-Video
• Device: AMD Radeon Graphics
• ROCm: 7.1.52802-26aae437f6
• Config: 25 frames, 480x704, 30 steps, warmup=1
• Repeats: 3 for baseline and triton
• Same prompt and seed across modes

Commands

Run from skills/rocm-kernels:

# Regenerate benchmark outputs
python scripts/benchmark_e2e.py --mode all --num-frames 25 --steps 30 --warmup 1 --output-dir examples/ltx-video-benchmark

Summary table (formal benchmark)

Notes:

• speedup is computed against baseline mean.
• compile is included as a supplemental reference and is not used as the main claim for the custom-kernel comparison.

Output files

• benchmark_results.json
• trace/opencode_live/results.json
• trace/opencode_live/opencode_trace_result.json

Coding Harness Trace

This package includes a single coding-harness trace file:

• trace/opencode_live/opencode_trace_result.json

OpenCode trace was generated with model opencode/minimax-m2.5-free.

PR-ready response

We added baseline vs Triton numbers for LTX-Video on ROCm with consolidated benchmark output and a single live harness trace. Using the same prompt, seed, resolution, frame count, and inference steps, Triton improves end-to-end latency from 6.91s to 6.10s on average (1.13x speedup) in the formal benchmark package.

Diffusers Pipeline Integration Guide (ROCm)

Integrating custom Triton kernels into HuggingFace diffusers pipelines on AMD GPUs.

Dependencies

Use the existing dependency file:

python -m pip install -r skills/rocm-kernels/scripts/requirements.txt

Required (for this guide)

• torch
• triton
• diffusers
• transformers
• accelerate

Optional

• Video export: imageio-ffmpeg (pip) or system ffmpeg
• Profiling/trace: rocprof, rocprofv3

Tested Environment

Clean Install + Smoke Test

python -m venv .venv-rocm-kernels
source .venv-rocm-kernels/bin/activate
python -m pip install --upgrade pip
python -m pip install -r skills/rocm-kernels/scripts/requirements.txt

python - <<'PY'
import torch
from diffusers import LTXPipeline

assert torch.cuda.is_available(), "ROCm device is not available"
pipe = LTXPipeline.from_pretrained("Lightricks/LTX-Video", torch_dtype=torch.bfloat16)
pipe.to("cuda")
_ = pipe(
    prompt="A quick smoke test on ROCm",
    num_frames=9,
    height=480,
    width=704,
    num_inference_steps=1,
    guidance_scale=7.5,
)
print("Smoke test passed.")
PY

Overview

This guide covers injecting optimized Triton kernels (RMSNorm, RoPE 3D, GEGLU, AdaLN) into diffusers pipelines running on ROCm. The patterns are analogous to the CUDA kernel integration but use Triton instead of CUDA C.

LTX-Video Architecture

Module Inventory

from diffusers import LTXPipeline
pipe = LTXPipeline.from_pretrained("Lightricks/LTX-Video", torch_dtype=torch.bfloat16)

# Analyze RMSNorm modules
for name, module in pipe.transformer.named_modules():
    if 'Norm' in type(module).__name__:
        has_weight = hasattr(module, 'weight') and module.weight is not None
        print(f"{name}: {type(module).__name__} (has_weight={has_weight})")

Kernel Applicability in LTX-Video

Integration Pattern

Step 1: Triton RMSNorm Wrapper

import os
os.environ['TRITON_HIP_USE_BLOCK_PINGPONG'] = '1'
os.environ['TRITON_HIP_USE_ASYNC_COPY'] = '1'

import torch
import triton
import triton.language as tl


@triton.jit
def _rmsnorm_kernel(
    x_ptr, weight_ptr, out_ptr,
    stride_x_row, D,
    eps: tl.constexpr,
    HAS_WEIGHT: tl.constexpr,
    BLOCK_D: tl.constexpr,
):
    row = tl.program_id(0)
    offs = tl.arange(0, BLOCK_D)
    mask = offs < D

    x = tl.load(x_ptr + row * stride_x_row + offs, mask=mask, other=0.0).to(tl.float32)
    sq_sum = tl.sum(x * x, axis=0)
    rms_inv = tl.rsqrt(sq_sum / D + eps)

    if HAS_WEIGHT:
        w = tl.load(weight_ptr + offs, mask=mask, other=1.0).to(tl.float32)
        out = x * rms_inv * w
    else:
        out = x * rms_inv

    tl.store(out_ptr + row * stride_x_row + offs, out.to(x.dtype), mask=mask)


def triton_rmsnorm(x: torch.Tensor, weight: torch.Tensor = None, eps: float = 1e-6):
    """Drop-in replacement for RMSNorm forward pass."""
    x_flat = x.contiguous().view(-1, x.shape[-1])
    out = torch.empty_like(x_flat)
    M, D = x_flat.shape
    has_weight = weight is not None

    if not has_weight:
        weight = torch.ones(D, device=x.device, dtype=x.dtype)

    # CRITICAL: BLOCK_D must be >= D. Never autotune BLOCK_D.
    BLOCK_D = triton.next_power_of_2(D)
    num_warps = 4 if BLOCK_D <= 1024 else (8 if BLOCK_D <= 4096 else 16)
    _rmsnorm_kernel[(M,)](
        x_flat, weight, out,
        x_flat.stride(0), D,
        eps, has_weight,
        BLOCK_D=BLOCK_D, num_warps=num_warps, num_stages=2,
    )
    return out.view_as(x)

Step 2: Module Patcher

def patch_rmsnorm_modules(model) -> int:
    """
    Patch all RMSNorm modules to use Triton kernel on ROCm.

    Handles both:
    - Modules WITH weight (elementwise_affine=True) — attention norm_q/norm_k
    - Modules WITHOUT weight (elementwise_affine=False) — transformer block norms
    """
    patched = 0
    for name, module in model.named_modules():
        # IMPORTANT: Use class name, NOT isinstance
        if type(module).__name__ == 'RMSNorm':
            eps = getattr(module, 'eps', 1e-6)
            has_weight = hasattr(module, 'weight') and module.weight is not None

            if has_weight:
                def make_forward(mod, epsilon):
                    def forward(x):
                        return triton_rmsnorm(x, mod.weight, eps=epsilon)
                    return forward
                module.forward = make_forward(module, eps)
            else:
                def make_forward_no_weight(epsilon):
                    def forward(x):
                        return triton_rmsnorm(x, None, eps=epsilon)
                    return forward
                module.forward = make_forward_no_weight(eps)

            patched += 1
    return patched

Step 3: Pipeline Injection

def inject_optimized_kernels(pipe) -> dict:
    """
    Inject Triton kernels into LTX-Video pipeline.

    Call AFTER pipe.to("cuda"), BEFORE pipe.enable_model_cpu_offload().
    """
    stats = {'rmsnorm_modules': 0}

    if not hasattr(pipe, 'transformer'):
        print("WARNING: Pipeline has no 'transformer' attribute!")
        return stats

    stats['rmsnorm_modules'] = patch_rmsnorm_modules(pipe.transformer)
    return stats

Step 4: Usage

import torch
from diffusers import LTXPipeline
from diffusers.utils import export_to_video

pipe = LTXPipeline.from_pretrained("Lightricks/LTX-Video", torch_dtype=torch.bfloat16)
pipe.to("cuda")  # ROCm via HIP

stats = inject_optimized_kernels(pipe)
print(f"RMSNorm modules patched: {stats['rmsnorm_modules']}")
# Expected: 168

pipe.enable_model_cpu_offload()  # AFTER injection

output = pipe(
    prompt="A cat sleeping in the sun",
    num_frames=25, height=480, width=704,
    num_inference_steps=30,
)
export_to_video(output.frames[0], "output.mp4", fps=24)

Model-Specific Notes

LTX-Video

• Uses GELU (activation_fn="gelu-approximate"), NOT GEGLU
• RMSNorm in blocks: elementwise_affine=False (no weight)
• RMSNorm in attention: elementwise_affine=True (has weight)
• RoPE: Computed by diffusers via LTXVideoRotaryPosEmbed

SD3 / FLUX

• Uses GEGLU in FeedForward blocks
• Different attention patterns
• May have different normalization conventions
• Verify architecture before applying LTX-Video patterns

ROCm-Specific Considerations

BF16 vs FP16

# MI355X supports BF16 — use it for diffusers
pipe = LTXPipeline.from_pretrained(..., torch_dtype=torch.bfloat16)

# R9700 (RDNA4) — check BF16 support, may need FP16
# torch_dtype=torch.float16

ROCm Memory Management

# ROCm uses same API as CUDA via HIP
pipe.to("cuda")                    # Works on ROCm
pipe.enable_model_cpu_offload()    # Works on ROCm
torch.cuda.empty_cache()           # Works on ROCm

Triton on ROCm vs CUDA C Kernels

Verification

# Check injection worked
for name, module in pipe.transformer.named_modules():
    if type(module).__name__ == 'RMSNorm':
        x = torch.randn(1, 10, 2048, device='cuda', dtype=torch.bfloat16)
        out = module(x)
        print(f"RMSNorm forward: {x.shape} -> {out.shape}")
        break

# Compare with PyTorch reference
def pytorch_rmsnorm(x, weight, eps=1e-6):
    rms = torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps)
    if weight is not None:
        return x * rms * weight
    return x * rms

# Verify correctness
torch.testing.assert_close(
    triton_rmsnorm(x, weight, eps=1e-6),
    pytorch_rmsnorm(x, weight, eps=1e-6),
    rtol=1e-2, atol=1e-3
)

Benchmark Evidence (R9700)

For reviewer-facing benchmark materials, see examples/ltx-video-benchmark/.

Reviewer-facing artifacts are organized under examples/ltx-video-benchmark/:

• benchmark_results.json
• trace/opencode_live/results.json
• trace/opencode_live/opencode_trace_result.json

Troubleshooting

HuggingFace Kernels Integration Guide (ROCm)

Complete guide for using and publishing kernels with the HuggingFace Kernels library (get_kernel) on ROCm.

Quick Start: See huggingface_kernels_example.py for a minimal working example.

Overview

The HuggingFace Kernels library enables dynamic loading of pre-compiled kernels from the Hugging Face Hub. This eliminates the need for local compilation and ensures compatibility across different Python, PyTorch, and CUDA/ROCm versions.

Key Benefits:

• No local compilation — download pre-built binaries
• Version management — load specific kernel versions
• Multi-version support — multiple versions coexist in one Python process
• Automatic compatibility — matches your PyTorch/ROCm configuration

ROCm Note: Not all Hub kernels have ROCm builds. Triton-based kernels (e.g., triton-layer-norm) are more likely to work on ROCm than CUDA C kernels. Always check with has_kernel() first.

Installation

pip install kernels torch numpy

Requirements:

• PyTorch >= 2.5 (ROCm build)
• ROCm-capable AMD GPU
• Python 3.8+

Core API

get_kernel

Download and load a kernel from the Hub:

from kernels import get_kernel

kernel = get_kernel("kernels-community/triton-layer-norm")

# With specific version
kernel = get_kernel("kernels-community/triton-layer-norm", version=1)

# With specific revision
kernel = get_kernel("kernels-community/flash-attn", revision="v2.0.0")

Parameters:

Returns: ModuleType — the imported kernel module

has_kernel

Check if a kernel build exists for your environment:

from kernels import has_kernel

if has_kernel("kernels-community/triton-layer-norm"):
    kernel = get_kernel("kernels-community/triton-layer-norm")
else:
    print("No compatible build for this ROCm/PyTorch version")

get_local_kernel

Load a kernel from a local path (useful during development):

from kernels import get_local_kernel

kernel = get_local_kernel("/path/to/my-kernel")

load_kernel & get_locked_kernel

For reproducible, offline-capable deployments using lockfiles:

from kernels import load_kernel, get_locked_kernel

kernel = load_kernel("lockfile.json")
kernel = get_locked_kernel("kernels-community/activation", lockfile="kernel.lock")

Usage Examples

1. RMSNorm Kernel from Hub

Note: The actual function name may vary by kernel version. Use dir(kernel) to inspect, and check for rms_norm_fn, rms_norm, or rmsnorm.

import torch
from kernels import get_kernel, has_kernel

repo_id = "kernels-community/triton-layer-norm"

if has_kernel(repo_id):
    layer_norm = get_kernel(repo_id)

    # Inspect available functions
    print([f for f in dir(layer_norm) if not f.startswith('_')])
    # e.g. ['layer_norm', 'layer_norm_fn', 'rms_norm_fn', ...]

    x = torch.randn(2, 1024, 2048, dtype=torch.bfloat16, device="cuda")
    weight = torch.ones(2048, dtype=torch.bfloat16, device="cuda")

    # Use the actual function name (rms_norm_fn in current version)
    out = layer_norm.rms_norm_fn(x, weight, eps=1e-6)
    print(f"Output shape: {out.shape}")
else:
    print("No ROCm-compatible build available")

2. Integration with Transformers Models

import torch
from kernels import get_kernel, has_kernel

repo_id = "kernels-community/triton-layer-norm"

if has_kernel(repo_id):
    rmsnorm_kernel = get_kernel(repo_id)

    def patch_rmsnorm_with_hub_kernel(model):
        """Patch model's RMSNorm to use Hub kernel."""
        patched = 0
        for name, module in model.named_modules():
            if 'RMSNorm' in type(module).__name__:
                eps = getattr(module, 'variance_epsilon', None) or getattr(module, 'eps', 1e-6)

                def make_forward(mod, epsilon):
                    def forward(hidden_states):
                        return rmsnorm_kernel.rms_norm(hidden_states, mod.weight, eps=epsilon)
                    return forward

                module.forward = make_forward(module, eps)
                patched += 1
        return patched

3. Integration with Diffusers Pipelines

import torch
from diffusers import LTXPipeline
from kernels import get_kernel, has_kernel

if has_kernel("kernels-community/triton-layer-norm"):
    rmsnorm_kernel = get_kernel("kernels-community/triton-layer-norm")

    def patch_rmsnorm(model):
        for name, module in model.named_modules():
            if type(module).__name__ == 'RMSNorm':
                eps = getattr(module, 'eps', 1e-6)
                has_weight = hasattr(module, 'weight') and module.weight is not None

                if has_weight:
                    def make_forward(mod, epsilon):
                        def forward(x):
                            return rmsnorm_kernel.rms_norm(x, mod.weight, eps=epsilon)
                        return forward
                    module.forward = make_forward(module, eps)

    pipe = LTXPipeline.from_pretrained("Lightricks/LTX-Video", torch_dtype=torch.bfloat16)
    pipe.to("cuda")
    patch_rmsnorm(pipe.transformer)

4. Benchmark Hub Kernel vs PyTorch

import time
import torch
from kernels import get_kernel

kernel = get_kernel("kernels-community/triton-layer-norm")

sizes = [(2, 1024, 2048), (4, 4096, 4096)]
for shape in sizes:
    x = torch.randn(shape, dtype=torch.bfloat16, device="cuda")
    w = torch.ones(shape[-1], dtype=torch.bfloat16, device="cuda")

    for _ in range(10):
        kernel.rms_norm(x, w, eps=1e-6)
        variance = x.pow(2).mean(-1, keepdim=True)
        _ = x * torch.rsqrt(variance + 1e-6) * w
    torch.cuda.synchronize()

    iters = 100
    start = time.perf_counter()
    for _ in range(iters):
        kernel.rms_norm(x, w, eps=1e-6)
    torch.cuda.synchronize()
    hub_ms = (time.perf_counter() - start) / iters * 1000

    start = time.perf_counter()
    for _ in range(iters):
        variance = x.pow(2).mean(-1, keepdim=True)
        _ = x * torch.rsqrt(variance + 1e-6) * w
    torch.cuda.synchronize()
    pt_ms = (time.perf_counter() - start) / iters * 1000

    print(f"Shape {shape}: Hub={hub_ms:.3f}ms, PyTorch={pt_ms:.3f}ms, Speedup={pt_ms/hub_ms:.2f}x")

ROCm-Specific Notes

Kernel Compatibility

Not all Hub kernels have ROCm builds:

Always check availability first:

from kernels import has_kernel

if has_kernel("kernels-community/triton-layer-norm"):
    print("ROCm build available")
else:
    print("No ROCm build — use local Triton kernel instead")

Fallback Strategy

When a Hub kernel is not available for ROCm, fall back to the local Triton implementation:

from kernels import has_kernel, get_kernel

def get_rmsnorm_function():
    """Get best available RMSNorm implementation."""
    if has_kernel("kernels-community/triton-layer-norm"):
        kernel = get_kernel("kernels-community/triton-layer-norm")
        return lambda x, w, eps: kernel.rms_norm(x, w, eps=eps)
    else:
        from your_local_kernels import triton_rmsnorm
        return triton_rmsnorm

Environment Check

import torch
print(f"PyTorch: {torch.__version__}")
print(f"HIP version: {getattr(torch.version, 'hip', 'N/A')}")
print(f"GPU: {torch.cuda.get_device_name()}")
print(f"GPU arch: {torch.cuda.get_device_capability()}")

Publishing Kernels to Hub

Triton Kernel Project Structure

For Triton-based kernels (best ROCm compatibility):

my-triton-kernel/
├── build.toml
├── kernel_src/
│   └── rmsnorm.py          # Triton kernel source
└── torch-ext/
    ├── torch_binding.cpp
    └── my_kernels/
        └── __init__.py

build.toml for Triton Kernels

[general]
name = "my_triton_kernels"
backends = ["cuda", "rocm"]   # Include ROCm backend

[torch]
src = ["torch-ext/torch_binding.cpp"]

[kernel.rmsnorm]
backend = "triton"
src = ["kernel_src/rmsnorm.py"]
depends = ["torch"]

Build and Publish

pip install kernel-builder
kernel-builder build

huggingface-cli repo create your-org/your-kernel --type model
huggingface-cli upload your-org/your-kernel ./dist

Others Load It

from kernels import get_kernel

rmsnorm = get_kernel("your-org/your-kernel")

Available Community Kernels

Popular kernels from kernels-community:

Browse all kernels: https://huggingface.co/kernels-community

Caching and Offline Usage

import os
os.environ["HF_HUB_OFFLINE"] = "1"

# Will only use cached kernels
kernel = get_kernel("kernels-community/triton-layer-norm")

Best Practices

1. Always check availability — has_kernel() before get_kernel() 2. Pin versions — get_kernel(repo, version=1) for reproducibility 3. Have a fallback — local Triton kernel when Hub build is unavailable 4. Use lockfiles in production — load_kernel("kernel.lock") 5. Test on your GPU — verify correctness after loading

See Also

• HuggingFace Kernels Documentation
• HuggingFace Kernels GitHub
• Kernel Builder Documentation
• Community Kernels
• Blog: Learn the Kernel Hub in 5 Minutes

Triton Kernel Templates for ROCm (LTX-Video Operators)

Copy-paste ready Triton kernel templates for RMSNorm, RoPE 3D, GEGLU, and AdaLN on AMD GPUs.

Required Header

Every kernel file MUST start with:

import os
os.environ['TRITON_HIP_USE_BLOCK_PINGPONG'] = '1'
os.environ['TRITON_HIP_USE_ASYNC_COPY'] = '1'

import torch
import torch.nn as nn
import triton
import triton.language as tl

Template 1: RMSNorm (Core Target)

Row-wise reduction. 168 instances in LTX-Video. Handles both with-weight and no-weight variants.

CRITICAL: Do NOT autotune BLOCK_D. Autotune may select BLOCK_D < D, causing partial row processing and completely wrong results. Always compute BLOCK_D = triton.next_power_of_2(D) dynamically.

Triton Kernel

@triton.jit
def rmsnorm_fwd_kernel(
    x_ptr, weight_ptr, out_ptr,
    stride_x, D,
    eps: tl.constexpr,
    HAS_WEIGHT: tl.constexpr,
    BLOCK_D: tl.constexpr,
):
    row = tl.program_id(0)
    col_offsets = tl.arange(0, BLOCK_D)
    mask = col_offsets < D

    x = tl.load(x_ptr + row * stride_x + col_offsets, mask=mask, other=0.0).to(tl.float32)

    variance = tl.sum(x * x, axis=0) / D
    rms_inv = tl.rsqrt(variance + eps)

    if HAS_WEIGHT:
        w = tl.load(weight_ptr + col_offsets, mask=mask, other=1.0).to(tl.float32)
        result = x * rms_inv * w
    else:
        result = x * rms_inv

    tl.store(out_ptr + row * stride_x + col_offsets, result.to(x.dtype), mask=mask)

Python API

def triton_rmsnorm(
    x: torch.Tensor,
    weight: torch.Tensor = None,
    eps: float = 1e-6,
) -> torch.Tensor:
    orig_shape = x.shape
    x_2d = x.contiguous().view(-1, x.shape[-1])
    out = torch.empty_like(x_2d)
    M, D = x_2d.shape

    has_weight = weight is not None
    if not has_weight:
        weight = torch.empty(0, device=x.device)

    BLOCK_D = triton.next_power_of_2(D)
    num_warps = 4 if BLOCK_D <= 1024 else (8 if BLOCK_D <= 4096 else 16)
    rmsnorm_fwd_kernel[(M,)](
        x_2d, weight, out,
        x_2d.stride(0), D, eps, has_weight,
        BLOCK_D=BLOCK_D, num_warps=num_warps, num_stages=2,
    )
    return out.view(orig_shape)

Benchmark

def benchmark_rmsnorm():
    configs = [
        (1, 1024, 2048),
        (2, 1024, 2048),
        (2, 4096, 3072),
    ]
    for batch, seq, hidden in configs:
        x = torch.randn(batch, seq, hidden, device='cuda', dtype=torch.float16)
        w = torch.ones(hidden, device='cuda', dtype=torch.float16)

        # Reference
        ref = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + 1e-6) * w

        # Custom
        out = triton_rmsnorm(x, w, eps=1e-6)

        # Verify
        torch.testing.assert_close(out, ref, rtol=1e-2, atol=1e-3)
        print(f"[{batch}x{seq}x{hidden}] ✓ Correct")

Template 2: RoPE 3D (Video Position Encoding)

Element-wise rotation. Splits head_dim into temporal + spatial (height + width) components.

*CRITICAL: cos/sin have shape `[seq_len, head_dim]`, NOT `[batchseq_len, ...].** When the grid flattens the batch dimension, use pid_s % seq_len` to index cos/sin, otherwise batch > 1 causes out-of-bounds GPU crash.

Triton Kernel

@triton.jit
def rope_3d_fwd_kernel(
    qk_ptr, cos_ptr, sin_ptr, out_ptr,
    seq_len, num_heads, head_dim,
    stride_s, stride_h, stride_d,
    BLOCK_HD: tl.constexpr,
):
    pid_s = tl.program_id(0)  # ranges [0, batch * seq_len)
    pid_h = tl.program_id(1)

    half_dim = head_dim // 2
    offs = tl.arange(0, BLOCK_HD)
    mask = offs < half_dim

    base = pid_s * stride_s + pid_h * stride_h
    x0 = tl.load(qk_ptr + base + offs, mask=mask, other=0.0).to(tl.float32)
    x1 = tl.load(qk_ptr + base + half_dim + offs, mask=mask, other=0.0).to(tl.float32)

    seq_idx = pid_s % seq_len  # wrap for batch > 1
    cos_val = tl.load(cos_ptr + seq_idx * head_dim + offs, mask=mask, other=1.0).to(tl.float32)
    sin_val = tl.load(sin_ptr + seq_idx * head_dim + offs, mask=mask, other=0.0).to(tl.float32)

    out0 = x0 * cos_val - x1 * sin_val
    out1 = x0 * sin_val + x1 * cos_val

    tl.store(out_ptr + base + offs, out0.to(x0.dtype), mask=mask)
    tl.store(out_ptr + base + half_dim + offs, out1.to(x0.dtype), mask=mask)

Python API

def triton_rope_3d(
    qk: torch.Tensor,
    cos: torch.Tensor,
    sin: torch.Tensor,
) -> torch.Tensor:
    """
    Apply 3D RoPE to Q or K tensor.

    Args:
        qk: [batch, seq_len, num_heads, head_dim]
        cos: [seq_len, head_dim]  — NOT batch-expanded!
        sin: [seq_len, head_dim]
    """
    qk = qk.contiguous()
    out = torch.empty_like(qk)
    batch, seq_len, num_heads, head_dim = qk.shape

    qk_flat = qk.view(batch * seq_len, num_heads, head_dim)
    out_flat = out.view(batch * seq_len, num_heads, head_dim)

    BLOCK_HD = triton.next_power_of_2(head_dim // 2)
    num_warps = 4 if BLOCK_HD <= 64 else 8

    rope_3d_fwd_kernel[(batch * seq_len, num_heads)](
        qk_flat, cos, sin, out_flat,
        seq_len, num_heads, head_dim,
        qk_flat.stride(0), qk_flat.stride(1), qk_flat.stride(2),
        BLOCK_HD=BLOCK_HD, num_warps=num_warps, num_stages=2,
    )
    return out

Template 3: GEGLU (For SD3/FLUX)

Gated activation: GELU(gate) * value. Input splits in half along last dim.

Note: LTX-Video uses GELU, NOT GEGLU. This template is for SD3/FLUX.

Triton Kernel

@triton.jit
def geglu_fwd_kernel(
    input_ptr, output_ptr,
    stride_in, stride_out, hidden_size,
    BLOCK_H: tl.constexpr,
):
    row = tl.program_id(0)
    offs = tl.arange(0, BLOCK_H)
    mask = offs < hidden_size

    gate = tl.load(input_ptr + row * stride_in + offs,
                   mask=mask, other=0.0).to(tl.float32)
    value = tl.load(input_ptr + row * stride_in + hidden_size + offs,
                    mask=mask, other=0.0).to(tl.float32)

    # Manual tanh — tl.math.tanh / tl.libdevice.tanh NOT available on ROCm
    SQRT_2_OVER_PI = 0.7978845608028654
    tanh_arg = SQRT_2_OVER_PI * (gate + 0.044715 * gate * gate * gate)
    e2x = tl.exp(2.0 * tanh_arg)
    tanh_val = (e2x - 1.0) / (e2x + 1.0)
    cdf = 0.5 * (1.0 + tanh_val)
    gelu_gate = gate * cdf

    result = gelu_gate * value
    tl.store(output_ptr + row * stride_out + offs, result.to(gate.dtype), mask=mask)

Python API

def triton_geglu(x: torch.Tensor) -> torch.Tensor:
    """
    GEGLU activation: GELU(x[..., :H]) * x[..., H:]

    Input: [..., 2*hidden_size] → Output: [..., hidden_size]
    """
    x = x.contiguous()
    *batch_dims, double_h = x.shape
    hidden_size = double_h // 2

    x_2d = x.view(-1, double_h)
    M = x_2d.shape[0]
    out = torch.empty(M, hidden_size, device=x.device, dtype=x.dtype)

    BLOCK_H = triton.next_power_of_2(hidden_size)
    num_warps = 4 if BLOCK_H <= 1024 else (8 if BLOCK_H <= 4096 else 16)
    geglu_fwd_kernel[(M,)](
        x_2d, out,
        x_2d.stride(0), out.stride(0), hidden_size,
        BLOCK_H=BLOCK_H, num_warps=num_warps, num_stages=2,
    )
    return out.view(*batch_dims, hidden_size)

Template 4: AdaLN (Adaptive Layer Normalization)

Fused RMSNorm + adaptive conditioning for DiT blocks. Formula: norm(x) * weight * (1 + scale) + shift

Triton Kernel

@triton.jit
def adaln_fwd_kernel(
    x_ptr, weight_ptr, scale_ptr, shift_ptr, out_ptr,
    stride_x, stride_cond, D,
    eps: tl.constexpr,
    BLOCK_D: tl.constexpr,
):
    row = tl.program_id(0)
    offs = tl.arange(0, BLOCK_D)
    mask = offs < D

    x = tl.load(x_ptr + row * stride_x + offs, mask=mask, other=0.0).to(tl.float32)

    variance = tl.sum(x * x, axis=0) / D
    rms_inv = tl.rsqrt(variance + eps)
    x_norm = x * rms_inv

    w = tl.load(weight_ptr + offs, mask=mask, other=1.0).to(tl.float32)
    scale = tl.load(scale_ptr + row * stride_cond + offs, mask=mask, other=0.0).to(tl.float32)
    shift = tl.load(shift_ptr + row * stride_cond + offs, mask=mask, other=0.0).to(tl.float32)

    out = x_norm * w * (1.0 + scale) + shift
    tl.store(out_ptr + row * stride_x + offs, out.to(x.dtype), mask=mask)

Python API

def triton_adaln(
    x: torch.Tensor,
    weight: torch.Tensor,
    scale: torch.Tensor,
    shift: torch.Tensor,
    eps: float = 1e-6,
) -> torch.Tensor:
    """
    Adaptive Layer Normalization for DiT blocks.

    Args:
        x: [batch, seq, hidden]
        weight: [hidden]
        scale: [batch, seq, hidden] or [batch, 1, hidden]
        shift: [batch, seq, hidden] or [batch, 1, hidden]
    """
    x_flat = x.contiguous().view(-1, x.shape[-1])
    scale_flat = scale.contiguous().view(-1, x.shape[-1])
    shift_flat = shift.contiguous().view(-1, x.shape[-1])
    out = torch.empty_like(x_flat)
    M, D = x_flat.shape

    BLOCK_D = triton.next_power_of_2(D)
    num_warps = 4 if BLOCK_D <= 1024 else (8 if BLOCK_D <= 4096 else 16)
    adaln_fwd_kernel[(M,)](
        x_flat, weight, scale_flat, shift_flat, out,
        x_flat.stride(0), scale_flat.stride(0), D, eps,
        BLOCK_D=BLOCK_D, num_warps=num_warps, num_stages=2,
    )
    return out.view_as(x)

Common Math Replacements for ROCm

Kernel-Specific Guidelines

RMSNorm

• Input: [..., hidden_size] — flatten to 2D [M, D]
• Epsilon default: 1e-6
• Weight may be None if elementwise_affine=False
• Always accumulate x*x sum in FP32
• BLOCK_D = `triton.next_power_of_2(D)` — compute in wrapper, NEVER autotune
• Autotuning BLOCK_D is dangerous: if BLOCK_D < D, only partial row is processed → wrong results

RoPE 3D

• 1D: [batch, seq, heads, head_dim] for text
• 3D: [batch, t*h*w, heads, head_dim] for video
• LTX-Video computes RoPE via LTXVideoRotaryPosEmbed — kernel replaces the apply step
• head_dim typically 64 or 128
• cos/sin shape is `[seq_len, head_dim]` — use pid_s % seq_len for batch > 1

GEGLU vs GELU

• GEGLU: Input [B, S, 2*H] → Output [B, S, H] — gate/value split
• GELU: Standard activation, no split
• LTX-Video uses GELU, NOT GEGLU
• GEGLU is for SD3/FLUX

AdaLN

• Formula: norm(x) * weight * (1 + scale) + shift
• Scale/shift come from timestep embedding MLP
• DiT computes 6 values per block: (scale1, shift1, gate1, scale2, shift2, gate2)
• Fusing norm + conditioning saves one memory round-trip

Template 5: GEMM with XCD Swizzle (MI355X)

Tiled matrix multiplication with XCD swizzle for MI355X (32 XCDs). Mandatory for any GEMM-like operation on MI355X — without it, work clusters on a few chiplets, wasting 90%+ of the GPU.

See mi355x-optimization-guide.md for architecture details.

When to use XCD swizzle: GEMM, batched GEMM, attention (Q@K, score@V). NOT needed for elementwise, reduction, or normalization kernels.

Triton Kernel

NUM_XCDS = 32  # MI355X has 32 XCDs

@triton.autotune(
    configs=[
        triton.Config({'BLOCK_M': 256, 'BLOCK_N': 256, 'BLOCK_K': 64, 'GROUP_M': 8},
                      num_stages=2, num_warps=8),
        triton.Config({'BLOCK_M': 256, 'BLOCK_N': 128, 'BLOCK_K': 64, 'GROUP_M': 8},
                      num_stages=2, num_warps=8),
        triton.Config({'BLOCK_M': 128, 'BLOCK_N': 256, 'BLOCK_K': 64, 'GROUP_M': 8},
                      num_stages=2, num_warps=8),
        triton.Config({'BLOCK_M': 128, 'BLOCK_N': 128, 'BLOCK_K': 64, 'GROUP_M': 8},
                      num_stages=3, num_warps=8),
        triton.Config({'BLOCK_M': 128, 'BLOCK_N': 128, 'BLOCK_K': 32, 'GROUP_M': 8},
                      num_stages=4, num_warps=4),
        triton.Config({'BLOCK_M': 64, 'BLOCK_N': 256, 'BLOCK_K': 64, 'GROUP_M': 8},
                      num_stages=3, num_warps=8),
    ],
    key=['M', 'N', 'K'],
)
@triton.jit
def gemm_xcd_swizzle_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,
    GROUP_M: tl.constexpr,
):
    pid = tl.program_id(0)
    num_pid_m = tl.cdiv(M, BLOCK_M)
    num_pid_n = tl.cdiv(N, BLOCK_N)
    num_pids = num_pid_m * num_pid_n

    # --- XCD Swizzle: distribute blocks across 32 chiplets ---
    pids_per_xcd = (num_pids + NUM_XCDS - 1) // NUM_XCDS
    xcd_id = pid % NUM_XCDS
    local_pid = pid // NUM_XCDS
    if local_pid < pids_per_xcd:
        remapped_pid = xcd_id * pids_per_xcd + local_pid
        if remapped_pid < num_pids:
            pid = remapped_pid

    # --- L2 Cache Grouping (after XCD swizzle) ---
    num_pid_in_group = GROUP_M * num_pid_n
    group_id = pid // num_pid_in_group
    first_pid_m = group_id * GROUP_M
    group_size_m = tl.minimum(num_pid_m - first_pid_m, GROUP_M)
    pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m)
    pid_n = (pid % num_pid_in_group) // group_size_m

    # --- Compute GEMM tile ---
    offs_am = (pid_m * BLOCK_M + tl.arange(0, BLOCK_M)) % M
    offs_bn = (pid_n * BLOCK_N + tl.arange(0, BLOCK_N)) % N
    offs_k = tl.arange(0, BLOCK_K)

    a_ptrs = a_ptr + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak)
    b_ptrs = b_ptr + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn)

    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
    for k in range(0, tl.cdiv(K, BLOCK_K)):
        a = tl.load(a_ptrs, mask=offs_k[None, :] < K - k * BLOCK_K, other=0.0)
        b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_K, other=0.0)
        acc += tl.dot(a, b)
        a_ptrs += BLOCK_K * stride_ak
        b_ptrs += BLOCK_K * stride_bk

    # --- Store result ---
    offs_cm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
    offs_cn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
    c_ptrs = c_ptr + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :]
    mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N)
    tl.store(c_ptrs, acc.to(tl.float16), mask=mask)

Python API

def triton_gemm(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
    """
    Matrix multiplication C = A @ B with XCD swizzle for MI355X.

    Args:
        a: [M, K] input matrix
        b: [K, N] input matrix
    Returns:
        c: [M, N] output matrix
    """
    assert a.shape[1] == b.shape[0], "Inner dimensions must match"
    assert a.is_contiguous() and b.is_contiguous()
    M, K = a.shape
    K, N = b.shape
    c = torch.empty((M, N), device=a.device, dtype=a.dtype)

    grid = lambda META: (triton.cdiv(M, META['BLOCK_M']) * triton.cdiv(N, META['BLOCK_N']),)
    gemm_xcd_swizzle_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),
    )
    return c

Benchmark

def benchmark_gemm():
    configs = [(4096, 4096, 4096), (8192, 8192, 4096), (2048, 8192, 2048)]
    for M, N, K in configs:
        a = torch.randn(M, K, device='cuda', dtype=torch.float16)
        b = torch.randn(K, N, device='cuda', dtype=torch.float16)

        ref = torch.mm(a, b)
        out = triton_gemm(a, b)
        torch.testing.assert_close(out, ref, rtol=1e-2, atol=1e-1)

        # Benchmark
        for _ in range(10):
            triton_gemm(a, b)
            torch.mm(a, b)
        torch.cuda.synchronize()

        import time
        iters = 50
        start = time.perf_counter()
        for _ in range(iters):
            triton_gemm(a, b)
        torch.cuda.synchronize()
        custom_ms = (time.perf_counter() - start) / iters * 1000

        start = time.perf_counter()
        for _ in range(iters):
            torch.mm(a, b)
        torch.cuda.synchronize()
        torch_ms = (time.perf_counter() - start) / iters * 1000

        print(f"[{M}x{N}x{K}] Custom: {custom_ms:.2f}ms, Torch: {torch_ms:.2f}ms, "
              f"Speedup: {torch_ms/custom_ms:.2f}x")

GEMM-Specific Guidelines

• XCD Swizzle is MANDATORY on MI355X for any GEMM — without it, expect 0.3-0.5x
• L2 Cache Grouping (GROUP_M=8-16): Improves L2 hit rate after XCD swizzle
• MFMA: Use matrix_instr_nonkdim=16 for MI355X matrix cores
• FP32 accumulation: Always accumulate in FP32, cast at store
• LDS budget: Check BLOCK_M * BLOCK_K + BLOCK_K * BLOCK_N dtype num_stages < 160 KB
• Autotune: GEMM benefits heavily from autotuning — always include 4+ configs
• R9700: Does NOT have XCDs — remove the XCD swizzle section for RDNA4

KernelBench Operator Classification & Skill Mapping

This document classifies KernelBench operators into categories and maps each to the appropriate kernel skill/pattern.

Classification Taxonomy

Level 1: Basic Operators (53 operators)

Category A: GEMM / Matrix Multiplication (18 operators)

Key Pattern: Template 5 (GEMM with XCD Swizzle) Critical Optimization: XCD swizzle + L2 cache grouping + MFMA 16x16

Category B: Elementwise / Activation Functions (14 operators)

Key Pattern: Template 1 (Elementwise) Critical Optimization: Large BLOCK_SIZE (4096-16384), FP32 compute

Category C: Normalization (8 operators)

Key Pattern: Template 3 (Row-wise Reduction) Critical Optimization: FP32 accumulation, proper reduction

Category D: Pooling (6 operators)

Key Challenge: 3D grid mapping with Triton's program_id limits

Category E: Reduction (7 operators)

Key Pattern: Template 5 (Dimension Reduction) Key Challenge: Argmax/Argmin require manual implementation

Level 2: Fused Operators (20+ operators)

Combine multiple operations into single kernels.

Key Pattern: Template 6 (Fused GEMM + Activation)

Level 3-4: Network Models / Transformers

Full models requiring multiple kernel types. Decompose into Level 1 operators.

Level 6-7: Advanced / Expert

Category → Skill Mapping

Performance Expectations by Category

Based on kernel-agent test results:

Recommended Skill Development Order

1. Phase 1 (Quick wins): Elementwise activations, Sum/Mean reduction 2. Phase 2 (Core): GEMM with XCD swizzle, LayerNorm/RMSNorm 3. Phase 3 (Advanced): Softmax, Pooling, Attention 4. Phase 4 (Expert): Fused operators, BatchNorm, Quantized GEMM

MI355X (gfx950) Optimization Guide

Deep dive into MI355X-specific optimizations for Triton kernels on ROCm.

MI355X CDNA3+ Architecture

Key Specifications

Critical Architecture Differences from MI300X

1. 32 XCDs vs 8: XCD swizzle must use NUM_XCDS=32 2. 8 CUs per XCD vs 28: Finer-grained chiplet distribution 3. 160 KB LDS vs 64 KB: 2.5x larger local memory per CU 4. 8 TB/s vs 5.3 TB/s: 50% more memory bandwidth 5. OCP FP8 vs AMD FP8: Different FP8 format

XCD Swizzle (MANDATORY for GEMM)

MI355X has 32 XCDs. Without proper swizzle, GEMM blocks cluster on a few XCDs, wasting 90%+ of the GPU.

When to Use XCD Swizzle

XCD Swizzle Implementation

NUM_XCDS = 32

@triton.jit
def gemm_with_xcd_swizzle(...):
    pid = tl.program_id(0)
    num_pid_m = tl.cdiv(M, BLOCK_M)
    num_pid_n = tl.cdiv(N, BLOCK_N)
    num_pids = num_pid_m * num_pid_n

    # Step 1: XCD Swizzle
    pids_per_xcd = (num_pids + NUM_XCDS - 1) // NUM_XCDS
    xcd_id = pid % NUM_XCDS
    local_pid = pid // NUM_XCDS
    if local_pid < pids_per_xcd:
        remapped_pid = xcd_id * pids_per_xcd + local_pid
        if remapped_pid < num_pids:
            pid = remapped_pid

    # Step 2: L2 Cache Grouping (after XCD swizzle)
    num_pid_in_group = GROUP_M * num_pid_n
    group_id = pid // num_pid_in_group
    first_pid_m = group_id * GROUP_M
    group_size_m = tl.minimum(num_pid_m - first_pid_m, GROUP_M)
    pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m)
    pid_n = (pid % num_pid_in_group) // group_size_m

Performance Impact

MFMA Instructions

Use 16x16 MFMA for optimal matrix core utilization:

# Launch kernel with MFMA hint
kernel[grid](..., matrix_instr_nonkdim=16)

LDS Optimization

MI355X has 160 KB LDS per CU—2.5x more than MI300X.

LDS Budget Calculation

LDS usage = BLOCK_M × BLOCK_K × dtype_size + BLOCK_K × BLOCK_N × dtype_size
           × num_stages

Example (BLOCK_M=256, BLOCK_N=256, BLOCK_K=64, FP16, num_stages=2):
  = (256×64×2 + 64×256×2) × 2 = 131,072 bytes = 128 KB < 160 KB ✓

Same config on MI300X (64 KB LDS):
  128 KB > 64 KB ✗ → Need num_stages=1 or smaller blocks

Stage Configuration

Autotune Configurations

Elementwise Operations

@triton.autotune(
    configs=[
        triton.Config({'BLOCK_SIZE': 1024}, num_warps=4, num_stages=2),
        triton.Config({'BLOCK_SIZE': 2048}, num_warps=8, num_stages=2),
        triton.Config({'BLOCK_SIZE': 4096}, num_warps=8, num_stages=2),
        triton.Config({'BLOCK_SIZE': 4096}, num_warps=16, num_stages=2),
        triton.Config({'BLOCK_SIZE': 8192}, num_warps=16, num_stages=2),
        triton.Config({'BLOCK_SIZE': 16384}, num_warps=16, num_stages=2),
    ],
    key=['n_elements'],
)

GEMM Operations

@triton.autotune(
    configs=[
        triton.Config({'BLOCK_M': 256, 'BLOCK_N': 256, 'BLOCK_K': 64, 'GROUP_M': 8},
                      num_stages=2, num_warps=8),
        triton.Config({'BLOCK_M': 256, 'BLOCK_N': 128, 'BLOCK_K': 64, 'GROUP_M': 8},
                      num_stages=2, num_warps=8),
        triton.Config({'BLOCK_M': 128, 'BLOCK_N': 256, 'BLOCK_K': 64, 'GROUP_M': 8},
                      num_stages=2, num_warps=8),
        triton.Config({'BLOCK_M': 128, 'BLOCK_N': 128, 'BLOCK_K': 64, 'GROUP_M': 8},
                      num_stages=3, num_warps=8),
        triton.Config({'BLOCK_M': 128, 'BLOCK_N': 128, 'BLOCK_K': 32, 'GROUP_M': 8},
                      num_stages=4, num_warps=4),
        triton.Config({'BLOCK_M': 64, 'BLOCK_N': 256, 'BLOCK_K': 64, 'GROUP_M': 8},
                      num_stages=3, num_warps=8),
    ],
    key=['M', 'N', 'K'],
)

Problem-Specific Block Sizes

Precision and Numerical Stability

FP32 Accumulation (Required)

# Always accumulate in FP32
acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
for k in range(...):
    acc += tl.dot(a, b)
# Cast at store
c = acc.to(tl.float16)

Math Operations

# Cast to FP32 for transcendental functions
x_f32 = x.to(tl.float32)
result = tl.exp(x_f32)      # ✓
result = tl.log(x_f32)      # ✓
result = tl.sqrt(x_f32)     # ✓
result = 1.0 / x_f32        # ✓ (division in FP32)

# tanh workaround (tl.tanh not supported on AMD)
e2x = tl.exp(2.0 * x_f32)
tanh_x = (e2x - 1.0) / (e2x + 1.0)

Performance Profiling

# Basic kernel profiling
rocprof --stats python your_kernel.py

# Detailed metrics
rocprofv3 -i metrics.txt python your_kernel.py

# Key metrics to watch:
# - L2 cache hit rate (target >70%)
# - VGPR usage (128+ may limit occupancy)
# - LDS usage (max 160 KB on MI355X)
# - Memory bandwidth utilization (target 40-60% of 8 TB/s)

Environment Variables

import os
# Block ping-pong for better latency hiding
os.environ['TRITON_HIP_USE_BLOCK_PINGPONG'] = '1'
# Async memory copies
os.environ['TRITON_HIP_USE_ASYNC_COPY'] = '1'

Best Practices Summary

1. XCD Swizzle: Always for GEMM, never for elementwise 2. MFMA: Use matrix_instr_nonkdim=16 3. LDS: Leverage 160 KB, but check with num_stages 4. num_stages: 2-3 (safe), up to 4 if LDS permits 5. num_warps: 8 is default, autotune 4-16 6. BLOCK_SIZE: Larger than MI300X (1024-16384 for 1D) 7. GROUP_M: 8 or 16 for L2 cache grouping 8. FP32 acc: Always accumulate in FP32 9. Env vars: Set BLOCK_PINGPONG and ASYNC_COPY 10. Profile: Use rocprof to validate optimizations

R9700 (RDNA4, gfx1201) Optimization Guide

Deep dive into R9700-specific optimizations for Triton kernels on ROCm.

R9700 RDNA4 Architecture

Key Specifications

Critical RDNA4 vs CDNA3+ Differences

1. Wave32 vs Wave64: Warp size is 32, same as NVIDIA 2. No XCD Swizzle: Single die, no chiplet distribution needed 3. Limited Matrix Cores: No FP8 MFMA support 4. Smaller LDS: 64 KB vs 160 KB 5. Smaller L2 Cache: 8 MB vs 256 MB 6. 256B Cacheline: Stricter memory alignment requirements 7. Consumer GPU: Optimized for inference, not training

Wave32 Implications

num_warps Mapping

On RDNA4, num_warps still means "number of wavefronts per block":

• 1 warp = 32 threads (Wave32)
• Max 32 waves per CU
• num_warps range: 2-8 (smaller than CDNA)

# CDNA (MI355X): 1 warp = 64 threads
# num_warps=8 → 512 threads/block

# RDNA4 (R9700): 1 warp = 32 threads
# num_warps=8 → 256 threads/block
# Use higher num_warps if needed for same thread count

Reduction Code

Warp-level reductions use different offsets:

# CDNA (Wave64): offsets = 32, 16, 8, 4, 2, 1
# RDNA4 (Wave32): offsets = 16, 8, 4, 2, 1

# In Triton this is handled automatically by tl.sum(), tl.max(), etc.
# No manual shuffle code needed in Triton

Memory Hierarchy

256B Cacheline Alignment

R9700 uses 256-byte cachelines (vs 128B on RDNA3). Misaligned accesses are penalized more.

# Ensure contiguous memory access
x = x.contiguous()

# In kernel: sequential access pattern
offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
x = tl.load(x_ptr + offsets, mask=mask)  # Coalesced

L2 Cache Strategy

With only 8 MB L2, cache reuse is limited:

# For GEMM: use smaller tiles to fit in L2
# BLOCK_M=64, BLOCK_N=64, BLOCK_K=32
# Tile = 64×32×2 + 32×64×2 = 8 KB per stage
# With 2 stages: 16 KB fits in L2 per block

LDS (64 KB) Budget

Max LDS per CU = 64 KB

GEMM example (BLOCK_M=64, BLOCK_N=128, BLOCK_K=32, FP16, num_stages=2):
  = (64×32×2 + 32×128×2) × 2 = 24,576 bytes = 24 KB ✓

GEMM example (BLOCK_M=128, BLOCK_N=128, BLOCK_K=64, FP16, num_stages=2):
  = (128×64×2 + 64×128×2) × 2 = 65,536 bytes = 64 KB → Borderline!

Autotune Configurations

Elementwise Operations

@triton.autotune(
    configs=[
        triton.Config({'BLOCK_SIZE': 256}, num_warps=2, num_stages=2),
        triton.Config({'BLOCK_SIZE': 512}, num_warps=4, num_stages=2),
        triton.Config({'BLOCK_SIZE': 1024}, num_warps=4, num_stages=2),
        triton.Config({'BLOCK_SIZE': 1024}, num_warps=8, num_stages=2),
        triton.Config({'BLOCK_SIZE': 2048}, num_warps=8, num_stages=2),
    ],
    key=['n_elements'],
)

GEMM Operations

@triton.autotune(
    configs=[
        triton.Config({'BLOCK_M': 64, 'BLOCK_N': 64, 'BLOCK_K': 32, 'GROUP_M': 8},
                      num_stages=2, num_warps=4),
        triton.Config({'BLOCK_M': 64, 'BLOCK_N': 128, 'BLOCK_K': 32, 'GROUP_M': 8},
                      num_stages=2, num_warps=4),
        triton.Config({'BLOCK_M': 128, 'BLOCK_N': 64, 'BLOCK_K': 32, 'GROUP_M': 8},
                      num_stages=2, num_warps=4),
        triton.Config({'BLOCK_M': 128, 'BLOCK_N': 128, 'BLOCK_K': 32, 'GROUP_M': 8},
                      num_stages=2, num_warps=8),
    ],
    key=['M', 'N', 'K'],
)

Grid Sizing

With 64 CUs:

# Aim for multiples of 64 blocks
grid = (triton.cdiv(N, BLOCK_SIZE),)
# For GEMM: grid = (cdiv(M, BLOCK_M) * cdiv(N, BLOCK_N),)

Precision Considerations

• FP16 Matrix TFLOPS = 191 (2x FP32 vector)
• FP16 Vector TFLOPS = 95.7 (2x FP32 vector)
• No FP8 MFMA: Cannot use FP8 matrix operations
• INT8 Matrix TOPS = 383 (quantized inference)
• Use FP16 for compute, FP32 for accumulation

Best Practices Summary

1. Wave32 awareness: Use num_warps=2-8 2. No XCD Swizzle: Not needed on single-die 3. Smaller blocks: 64-128 for GEMM tiles 4. 256B alignment: Ensure contiguous memory access 5. LDS budget: Max 64 KB, keep num_stages=2 6. Grid sizing: Multiples of 64 CUs 7. FP16 preferred: Best throughput, no FP8 MFMA 8. L3 cache: 64 MB can help with model weights 9. Inference focus: Best suited for inference workloads 10. Cacheline: 256B alignment is stricter than MI355X

Transformers Library Integration Guide (ROCm / Triton)

Complete guide for integrating custom Triton kernels into HuggingFace transformers models on AMD GPUs.

Quick Start: See transformers_injection_example.py for a minimal working example (~150 lines).

Overview

The HuggingFace transformers library has different architecture patterns than diffusers. Understanding these patterns is critical for successful kernel integration with models like LLaMA, Mistral, Qwen, and other LLMs on ROCm.

Key difference from diffusers: All transformers RMSNorm modules have weights (elementwise_affine=True). No need to handle the weight-less variant.

Model Architecture Analysis

from transformers import AutoModelForCausalLM, AutoConfig
import torch

config = AutoConfig.from_pretrained("Qwen/Qwen3-8B")
print(f"Hidden size: {config.hidden_size}")       # 4096
print(f"Num layers: {config.num_hidden_layers}")   # 32
print(f"Num heads: {config.num_attention_heads}")   # 32
print(f"RMS norm eps: {config.rms_norm_eps}")       # 1e-6

model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-8B",
    torch_dtype=torch.bfloat16,
    device_map="cuda"  # ROCm uses same API via HIP
)

for name, module in model.named_modules():
    class_name = type(module).__name__
    if 'Norm' in class_name:
        has_weight = hasattr(module, 'weight') and module.weight is not None
        print(f"{name}: {class_name} (has_weight={has_weight})")

Common Transformers Architectures

LLaMA / Llama-2 / Llama-3

Mistral / Mixtral

Qwen / Qwen2 / Qwen3

Kernel Applicability

Integration Pattern

Step 1: Set ROCm Environment Variables

import os
os.environ['TRITON_HIP_USE_BLOCK_PINGPONG'] = '1'
os.environ['TRITON_HIP_USE_ASYNC_COPY'] = '1'

Step 2: Define the Triton RMSNorm Kernel

import torch
import triton
import triton.language as tl

@triton.jit
def rmsnorm_fwd_kernel(
    x_ptr, weight_ptr, out_ptr,
    stride_x, D,
    eps: tl.constexpr,
    BLOCK_D: tl.constexpr,
):
    row = tl.program_id(0)
    offs = tl.arange(0, BLOCK_D)
    mask = offs < D
    x = tl.load(x_ptr + row * stride_x + offs, mask=mask, other=0.0).to(tl.float32)

    variance = tl.sum(x * x, axis=0) / D
    rms_inv = tl.rsqrt(variance + eps)

    w = tl.load(weight_ptr + offs, mask=mask, other=1.0).to(tl.float32)
    out = x * rms_inv * w

    tl.store(out_ptr + row * stride_x + offs, out.to(x.dtype), mask=mask)


def triton_rmsnorm(x, weight, eps=1e-6):
    x_2d = x.contiguous().view(-1, x.shape[-1])
    out = torch.empty_like(x_2d)
    M, D = x_2d.shape
    BLOCK_D = triton.next_power_of_2(D)
    num_warps = 4 if BLOCK_D <= 1024 else (8 if BLOCK_D <= 4096 else 16)
    rmsnorm_fwd_kernel[(M,)](
        x_2d, weight, out, x_2d.stride(0), D, eps,
        BLOCK_D=BLOCK_D, num_warps=num_warps, num_stages=2,
    )
    return out.view_as(x)

Step 3: Create RMSNorm Patcher

def patch_rmsnorm_modules(model) -> int:
    """
    Patch all RMSNorm modules to use Triton kernel on ROCm.

    Works with LlamaRMSNorm, MistralRMSNorm, Qwen2RMSNorm, etc.
    """
    patched_count = 0

    for name, module in model.named_modules():
        class_name = type(module).__name__

        if 'RMSNorm' in class_name:
            # LLaMA uses 'variance_epsilon', others use 'eps'
            eps = getattr(module, 'variance_epsilon', None)
            if eps is None:
                eps = getattr(module, 'eps', 1e-6)

            has_weight = hasattr(module, 'weight') and module.weight is not None

            if has_weight:
                def make_patched_forward(mod, epsilon):
                    def patched_forward(hidden_states):
                        return triton_rmsnorm(hidden_states, mod.weight, eps=epsilon)
                    return patched_forward
                module.forward = make_patched_forward(module, eps)
                patched_count += 1
            else:
                print(f"WARNING: {name} has no weight, skipping")

    return patched_count

Step 4: Use in Script

from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-8B",
    torch_dtype=torch.bfloat16,
    device_map="cuda"  # ROCm uses same device API via HIP
)
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-8B")

count = patch_rmsnorm_modules(model)
print(f"Patched {count} RMSNorm modules")
# Expected: 65 modules (32 layers * 2 + 1 final)

inputs = tokenizer("The capital of France is", return_tensors="pt").to("cuda")
outputs = model.generate(**inputs, max_new_tokens=20)
print(tokenizer.decode(outputs[0]))

Key Differences from Diffusers

1. RMSNorm Always Has Weight

Unlike diffusers (where some RMSNorm modules have elementwise_affine=False), transformers RMSNorm modules always have weights. The HAS_WEIGHT branch is always true, so you can simplify the kernel to always load weights.

2. Different Epsilon Attribute Names

# LLaMA uses 'variance_epsilon'
eps = getattr(module, 'variance_epsilon', 1e-6)

# Some models use 'eps'
eps = getattr(module, 'eps', 1e-6)

# Safe pattern
eps = getattr(module, 'variance_epsilon', None) or getattr(module, 'eps', 1e-6)

3. No Attention Processor Pattern

Diffusers uses set_processor() for attention modules. Transformers does not:

# Transformers: Use Flash Attention 2 instead of custom processors
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.bfloat16,
    attn_implementation="flash_attention_2"
)

4. Device Map vs Manual Move

# Transformers — use device_map for large models
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    device_map="auto"  # Handles multi-GPU automatically
)

# Diffusers — manual move then CPU offload
pipe = DiffusionPipeline.from_pretrained(model_id)
pipe.to("cuda")
pipe.enable_model_cpu_offload()

ROCm-Specific Considerations

1. ROCm Environment Setup

import os
os.environ['TRITON_HIP_USE_BLOCK_PINGPONG'] = '1'
os.environ['TRITON_HIP_USE_ASYNC_COPY'] = '1'

2. No tl.libdevice / tl.math.tanh

If you extend beyond RMSNorm (e.g., custom SiLU activation), remember tanh is not available:

# Manual tanh for ROCm
e2x = tl.exp(2.0 * x)
tanh_x = (e2x - 1.0) / (e2x + 1.0)

3. Verify HIP Backend

import torch
print(f"HIP version: {torch.version.hip}")  # Should show ROCm version
print(f"GPU: {torch.cuda.get_device_name()}")

4. torch.compile on ROCm

Custom Triton kernels and torch.compile can coexist on ROCm since Triton is already the compilation backend. However, test thoroughly as behavior may differ from eager mode.

Model-Specific Integration

LLaMA Models

from transformers import LlamaForCausalLM

model = LlamaForCausalLM.from_pretrained(
    "meta-llama/Llama-2-7b-hf",
    torch_dtype=torch.bfloat16,
    device_map="cuda"
)

count = patch_rmsnorm_modules(model)
print(f"Patched {count} LlamaRMSNorm modules")
# Expected: 65 modules (32 layers * 2 + 1 final)

Qwen3-8B

model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-8B",
    torch_dtype=torch.bfloat16,
    device_map="cuda"
)

count = patch_rmsnorm_modules(model)
print(f"Patched {count} Qwen2RMSNorm modules")
# Expected: 65 modules (32 layers * 2 + 1 final)

Verification

Verify Injection Worked

x = torch.randn(1, 10, model.config.hidden_size, device='cuda', dtype=torch.bfloat16)
for name, module in model.named_modules():
    if 'RMSNorm' in type(module).__name__:
        out = module(x)
        print(f"RMSNorm forward pass: {x.shape} -> {out.shape}")
        break

Run Generation Test

inputs = tokenizer("Hello, my name is", return_tensors="pt").to("cuda")
with torch.inference_mode():
    outputs = model.generate(**inputs, max_new_tokens=20)
print(tokenizer.decode(outputs[0]))

Profile on ROCm

rocprof --stats python your_script.py
rocprofv3 -i metrics.txt python your_script.py

Performance Optimization

Enable Flash Attention 2

model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.bfloat16,
    attn_implementation="flash_attention_2",
    device_map="cuda"
)

Combine with Custom Kernels

model = AutoModelForCausalLM.from_pretrained(model_id, ...)
patch_rmsnorm_modules(model)  # Inject Triton RMSNorm
# Flash Attention 2 handles attention optimization

ROCm Triton Kernel Troubleshooting Guide

Common issues and solutions when developing Triton kernels for AMD GPUs.

Build / Import Issues

1. tl.libdevice Not Found

Error: AttributeError: module 'triton.language' has no attribute 'libdevice'

Cause: tl.libdevice is CUDA-only (NVIDIA's libdevice library).

Fix: Replace with manual implementations:

# WRONG (CUDA only)
tl.libdevice.tanh(x)
tl.libdevice.log1p(x)

# CORRECT (ROCm compatible)
e2x = tl.exp(2.0 * x); tanh_x = (e2x - 1.0) / (e2x + 1.0)
log1p_x = tl.log(1.0 + x)

2. tl.tanh / tl.math.tanh Not Available

Error: AttributeError: module 'triton.language.math' has no attribute 'tanh'

Cause: Neither tl.tanh, tl.math.tanh, nor tl.libdevice.tanh exist on ROCm Triton. This is the most common GEGLU compilation failure.

Fix — manual tanh (ONLY reliable method):

x_f32 = x.to(tl.float32)
e2x = tl.exp(2.0 * x_f32)
tanh_x = (e2x - 1.0) / (e2x + 1.0)

Runtime Errors

3. HIP Runtime Error: Invalid Argument

Error: hipErrorInvalidValue or HIP Error: invalid argument

Common causes:

• Grid/block size exceeds hardware limits
• Mismatched tensor shapes
• LDS overflow

Fix:

# Check grid size
grid = (triton.cdiv(N, BLOCK_SIZE),)
assert grid[0] > 0, f"Grid size must be > 0, got {grid[0]}"

# Ensure contiguous tensors
x = x.contiguous()

# Reduce num_stages to avoid LDS overflow
# num_stages=2 is safest

4. HIP Out of Memory (LDS)

Error: AMDGPU_KERNEL_ERROR_OUT_OF_MEMORY or LDS size exceeds limit

Cause: Kernel uses more LDS than available (64 KB on R9700, 160 KB on MI355X).

Fix:

# Reduce num_stages
num_stages=2  # instead of 3 or 4

# Reduce block sizes
BLOCK_M=64, BLOCK_N=64, BLOCK_K=32  # smaller tiles

5. Kernel Timeout

Error: Kernel hangs or times out.

Common cause: Grid and Program ID mismatch.

# WRONG: 1D grid but 2D program_id
grid = (triton.cdiv(M, BLOCK_M) * triton.cdiv(N, BLOCK_N),)
pid_m = tl.program_id(0)  # OK
pid_n = tl.program_id(1)  # ERROR: axis 1 doesn't exist in 1D grid

# CORRECT: Compute 2D indices from 1D grid
grid = (triton.cdiv(M, BLOCK_M) * triton.cdiv(N, BLOCK_N),)
pid = tl.program_id(0)
pid_m = pid // triton.cdiv(N, BLOCK_N)
pid_n = pid % triton.cdiv(N, BLOCK_N)

Correctness Issues

6. Autotuning BLOCK_D Causes Wrong Results

Symptom: RMSNorm/AdaLN/GEGLU correctness fails with large max_abs errors (4-8+). Kernel runs fast but produces garbage.

Cause: @triton.autotune with BLOCK_D configs (e.g., 512, 1024, 2048, 4096) may select a BLOCK_D < D (hidden dimension). Since tl.arange(0, BLOCK_D) only covers BLOCK_D elements, the kernel processes a partial row, computing wrong variance and writing incomplete output.

Fix: Never autotune BLOCK_D for row-reduction kernels. Compute it dynamically:

# WRONG — autotune may pick BLOCK_D=512 when D=2048
@triton.autotune(configs=[
    triton.Config({'BLOCK_D': 512}, num_warps=4),
    triton.Config({'BLOCK_D': 1024}, num_warps=8),
], key=['D'])

# CORRECT — compute in Python wrapper
BLOCK_D = triton.next_power_of_2(D)
num_warps = 4 if BLOCK_D <= 1024 else (8 if BLOCK_D <= 4096 else 16)
kernel[(M,)](..., BLOCK_D=BLOCK_D, num_warps=num_warps, num_stages=2)

7. RoPE cos/sin Out-of-Bounds GPU Crash (batch > 1)

Symptom: Memory access fault by GPU node crash. Only happens when batch_size > 1.

Cause: cos/sin tensors have shape [seq_len, head_dim], but when the grid is (batch * seq_len, num_heads), pid_s ranges [0, batch * seq_len). For pid_s >= seq_len, cos_ptr + pid_s * head_dim is out of bounds.

Fix: Use modular indexing for cos/sin:

# WRONG — crashes when pid_s >= seq_len
cos_val = tl.load(cos_ptr + pid_s * head_dim + offs, ...)

# CORRECT — wrap position index for batch dimension
seq_idx = pid_s % seq_len
cos_val = tl.load(cos_ptr + seq_idx * head_dim + offs, ...)

8. FP16/BF16 Precision Loss

Symptom: Results differ from PyTorch reference by more than tolerance.

Fix: Always accumulate in FP32:

# WRONG: Accumulate in FP16
acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float16)

# CORRECT: Accumulate in FP32, cast at store
acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
# ... computation ...
result = acc.to(tl.float16)
tl.store(out_ptr + ..., result, mask=mask)

Tolerance guidelines:

• BF16 (7-bit mantissa): atol=0.1, rtol=1e-2
• FP16 (10-bit mantissa): atol=0.01, rtol=1e-3

9. Mask Errors

Error: ValueError: Mask argument cannot be block type

Fix: Ensure mask dimensions match pointer dimensions:

# 1D kernel
mask = offsets < n_elements
x = tl.load(x_ptr + offsets, mask=mask, other=0.0)

# 2D kernel
mask = (offs_m[:, None] < M) & (offs_n[None, :] < N)
x = tl.load(ptr + offs_m[:, None] * stride + offs_n[None, :], mask=mask, other=0.0)

10. Python min/max Inside Kernel

Error: TypeError or incorrect results.

Fix:

# WRONG: Python builtins
result = min(a, b)
result = max(a, b)

# CORRECT: Triton functions
result = tl.minimum(a, b)
result = tl.maximum(a, b)

Performance Issues

11. GEMM Extremely Slow (0.3-0.5x)

Cause: Missing XCD swizzle on MI355X.

Fix: Add XCD swizzle pattern (see Template 5: GEMM with XCD Swizzle in kernel-templates.md).

12. Elementwise Kernel Slow

Common causes: 1. BLOCK_SIZE too small → not utilizing bandwidth 2. Internal loops → should process full block 3. Missing autotune → not finding optimal config

Fix:

# Use large BLOCK_SIZE for elementwise
@triton.autotune(
    configs=[
        triton.Config({'BLOCK_SIZE': 2048}, num_warps=8, num_stages=2),
        triton.Config({'BLOCK_SIZE': 4096}, num_warps=8, num_stages=2),
        triton.Config({'BLOCK_SIZE': 8192}, num_warps=16, num_stages=2),
    ],
    key=['n_elements'],
)

13. Missing @triton.autotune (for elementwise)

Symptom: Kernel runs but performance is poor.

Fix: EVERY kernel must have autotune with 4+ configs. Fixed block sizes are almost never optimal.

14. tl.store() Keyword Argument Error

Error: TypeError: store() got an unexpected keyword argument

Fix: Check Triton version API. Use positional arguments if needed:

# Check your Triton version
# tl.store(ptr, value, mask=mask)  # Most versions
# tl.store(ptr, value, mask)       # Some older versions

15. eps: tl.constexpr Causes Recompilation Crash

Error: AttributeError("'NoneType' object has no attribute 'type'") during Triton compilation

Cause: When eps is declared as tl.constexpr, the kernel is compiled separately for each unique eps value. If the kernel first compiles with eps=1e-6 and later is called with eps=1e-8 (e.g., from nn.RMSNorm.eps), the recompilation on ROCm Triton can crash.

Fix: Remove tl.constexpr from eps and pass it as a regular runtime parameter:

# WRONG — triggers recompilation for each eps value, may crash on ROCm
@triton.jit
def rmsnorm_kernel(x_ptr, ..., eps: tl.constexpr, BLOCK_D: tl.constexpr):
    ...

# CORRECT — eps is a regular runtime float, no recompilation
@triton.jit
def rmsnorm_kernel(x_ptr, ..., eps, BLOCK_D: tl.constexpr):
    ...

# Also ensure eps is a plain float in the wrapper
rmsnorm_kernel[(M,)](..., float(eps), BLOCK_D=BLOCK_D, ...)

Note: Only BLOCK_D, HAS_WEIGHT, and other values that change kernel structure should be tl.constexpr. Parameters like eps that only affect numerical values should be regular parameters.

Debugging Tips

Check GPU Architecture

rocminfo | grep "Name"
# Should show gfx950 (MI355X) or gfx1201 (R9700)

Verify ROCm Triton Installation

import triton
print(triton.__version__)
import torch
print(torch.version.hip)  # Should show ROCm version
print(torch.cuda.get_device_properties(0))

Profile Kernel

# Basic profiling
rocprof --stats python your_kernel.py

# Detailed kernel metrics
rocprofv3 -i metrics.txt python your_kernel.py

Test Kernel Correctness

# Compare with PyTorch reference
ref_output = reference_model(inputs)
custom_output = custom_model(inputs)

torch.testing.assert_close(
    custom_output, ref_output,
    rtol=1e-2, atol=1e-3  # FP16 tolerance
)