rocm-kernels

by huggingface

Provides guidance for writing and benchmarking optimized Triton kernels for AMD GPUs (MI355X, R9700) on ROCm, targeting HuggingFace diffusers (LTX-Video, SD3,… npx skills add https://github.com/huggingface/kernels --skill rocm-kernels Download ZIPGitHub703

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.

Target Model: LTX-Video

Architecture Overview

Component Class Has Weight Count Kernel transformer_blocks.*.norm1 RMSNorm No (elementwise_affine=False) 56 RMSNorm transformer_blocks.*.norm2 RMSNorm No 56 RMSNorm transformer_blocks.*.attn1.norm_q torch.nn.RMSNorm Yes 28 RMSNorm transformer_blocks.*.attn1.norm_k torch.nn.RMSNorm Yes 28 RMSNorm transformer_blocks.*.ff FeedForward - 28 GELU (not GEGLU!) Rotary position encoding LTXVideoRotaryPosEmbed - 1 RoPE 3D

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

Target Kernels

Kernel Use Case Input Layout Key Challenge RMSNorm Normalization [..., hidden_size] Weight may be None; 168 instances RoPE 3D Video position encoding [batch, t*h*w, heads, head_dim] 3D → temporal + spatial decomposition GEGLU Gated activation (SD3/FLUX) [batch, seq, 2*hidden] → [batch, seq, hidden] Gate/value split AdaLN Conditioned normalization (DiT) norm(x) * weight * (1+scale) + shift Fused norm + condition

Supported Hardware

GPU Architecture Wave Size LDS/CU Mem BW Key Feature Verified MI355X CDNA3+ (gfx950) Wave64 160 KB 8 TB/s 32 XCDs, XCD Swizzle for GEMM Yes R9700 RDNA4 (gfx1201) Wave32 64 KB ~608 GB/s 256B cacheline, inference-focused Yes

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 **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 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 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 **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 See 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

• RMSNorm weight may be None — LTX-Video uses elementwise_affine=False

• Diffusers RMSNorm != torch.nn.RMSNorm — Use type(module).__name__ not isinstance()

• LTX-Video uses GELU, not GEGLU — Don't patch GEGLU for LTX-Video

• Inject BEFORE CPU offloading — inject_kernels() then enable_model_cpu_offload()

Performance Expectations

Micro-benchmark Results (MI355X, BF16)

Kernel Avg Speedup Best Config Speedup Status RMSNorm 1.71x 2.44x ([4×4096×3072]) PASS RoPE 3D 1.21x 1.52x ([2×4096×16×128]) PASS GEGLU 1.43x 2.13x ([4×4096×8192]) PASS AdaLN 2.22x 2.77x ([4×4096×3072]) PASS

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

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

Mode Time (s) Per Step (s) Peak Mem (GB) Speedup baseline 1.20 0.040 18.58 1.00x triton 0.98 0.033 18.58 1.22x torch.compile 0.78 0.026 18.58 1.54x

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)

Kernel Avg Speedup Best Config Speedup Status RMSNorm 2.90x 3.97x ([1×8192×2048]) PASS RoPE 3D 2.09x 2.38x ([1×1024×16×64]) PASS GEGLU 1.69x 1.93x ([2×1024×8192]) PASS AdaLN 3.00x 3.67x ([4×4096×3072]) PASS

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)

Mode Time (s) Per Step (s) Peak Mem (GB) Speedup baseline (mean of 3) 6.91 0.231 18.58 1.00x triton (mean of 3) 6.10 0.203 18.58 1.13x torch.compile (single run) 5.05 0.168 18.58 1.37x

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

Test Result Transformers injection (TinyLlama 1.1B) PASS — 45 RMSNorm patched, 99.9 tokens/s HuggingFace Kernels Hub integration PASS — Hub kernel loads and runs on ROCm Local Triton vs Hub kernel (small shape) Local 5.92x vs Hub 1.27x (lower launch overhead) Local Triton vs Hub kernel (large shape) Local 3.59x vs Hub 3.57x (comparable) num_warps sweep (2/4/8/16/32) Default heuristic (4/8/16) is near-optimal; nw=32 always worst rocprof kernel fusion analysis Triton fuses 4 PyTorch kernels (pow+mean+rsqrt+mul) into 1

CUDA Reference (H100, for comparison)

Shape Custom (ms) PyTorch (ms) Speedup [1×1024×2048] 0.019 0.065 3.37x [2×4096×3072] 0.087 0.208 2.41x

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

Optimization Targets

Kernel MI355X R9700 Target Priority RMSNorm 1.71x 2.90x >3x (R9700) P0 — MI355X bandwidth util (44%→60%+) AdaLN 2.22x 3.00x >3.5x (R9700) P1 — already strong on both GEGLU 1.43x 1.69x >2x P1 — tanh overhead RoPE 3D 1.21x 2.09x >2.5x (R9700) P2 — small head_dim launch overhead

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