BENCHMARK.md
# Skill Benchmark: tilegym-converting-cutile-to-triton
> ✅ **Overall verdict: PASS — Recommended for publication**
## Publication Recommendation
Recommended for publication based on the completed evaluation evidence in this report.
## Evaluation Metadata
- Skill: `tilegym-converting-cutile-to-triton`
- Evaluation date: 2026-10-16
- Evaluator version: `1.5.6`
- Agents: Claude Code (`aws/anthropic/bedrock-claude-opus-4-8`), Codex (`openai/openai/gpt-5.5`)
- Tasks: 5 evaluation tasks (1 positive, 4 negative)
- Dataset digest: `sha256:6eb6dce8f7e1fd54fc084bc411e537b99c4ddcfd0385c53ed32c220626278056` (skill-evaluator-dataset-snapshot/1)
- Attempts per task: 1
- Environment: `k8s-sandbox`
- Tier 2 evidence: required for publication
- Tier 3 evidence: required for publication
Each task attempt ran in its own isolated sandbox pod.
## What This Report Answers
The three-tier evaluation checks whether the skill:
- is safe to use;
- produces correct answers;
- is discovered and activated when needed;
- helps the agent complete the user's goal and expected workflow; and
- avoids wasted skill and tool usage.
## Results at a Glance
| Measure | Claude Code (Baseline → Skill Uplift) | Codex (Baseline → Skill Uplift) |
|---|---:|---:|
| Overall | 99.8% — baseline ran, but no comparable score was available; uplift unavailable | 98.3% — baseline ran, but no comparable score was available; uplift unavailable |
| Security | 100.0% → 100.0% (±0.0 points) | 100.0% → 100.0% (±0.0 points) |
| Correctness | 80.0% → 100.0% (+20.0 points) | 100.0% → 100.0% (±0.0 points) |
| Discoverability | 100.0% — baseline ran, but no comparable score was available; uplift unavailable | 95.0% — baseline ran, but no comparable score was available; uplift unavailable |
| Effectiveness | 82.5% → 100.0% (+17.5 points) | 90.0% → 99.0% (+9.0 points) |
| Efficiency | 99.2% — baseline ran, but no comparable score was available; uplift unavailable | 97.6% — baseline ran, but no comparable score was available; uplift unavailable |
**How to read this table:** baseline is the same task attempted without the target skill. Scores are rounded to one decimal; threshold-adjacent values use additional precision so their displayed band matches the verdict. Uplift is derived from those displayed scores and shown in percentage points.
Example: `47.0% → 92.0% (+45.0 points)` means the skill-assisted run scored 92.0%, 45.0 percentage points above its 47.0% no-skill baseline.
A partial dimension was calculated from only the available configured signals; review the detailed report before relying on it.
## Token Usage
Actual Tier 3 execution usage is reported for every observed agent/case pair and both conditions.
| Agent | Dataset case | With skill | Without skill | Delta | Change | Coverage |
|---|---|---:|---:|---:|---:|---|
| claude-code | All cases | 188,706 | 211,123 | -22,417 | -10.62% | skill 5/5; base 5/5 |
| claude-code | 01-overview-cutile-to-triton | 67,321 | 89,258 | -21,937 | -24.58% | skill 1/1; base 1/1 |
| claude-code | 02-swiftui-animation-negative | 31,092 | 31,113 | -21 | -0.07% | skill 1/1; base 1/1 |
| claude-code | 03-kafka-rebalance-negative | 31,121 | 31,056 | +65 | +0.21% | skill 1/1; base 1/1 |
| claude-code | 04-css-grid-negative | 29,952 | 29,587 | +365 | +1.23% | skill 1/1; base 1/1 |
| claude-code | 05-mongodb-aggregation-negative | 29,220 | 30,109 | -889 | -2.95% | skill 1/1; base 1/1 |
| codex | All cases | 171,192 | 109,094 | +62,098 | +56.92% | skill 5/5; base 5/5 |
| codex | 01-overview-cutile-to-triton | 101,158 | 39,762 | +61,396 | +154.41% | skill 1/1; base 1/1 |
| codex | 02-swiftui-animation-negative | 19,576 | 22,783 | -3,207 | -14.08% | skill 1/1; base 1/1 |
| codex | 03-kafka-rebalance-negative | 18,745 | 15,022 | +3,723 | +24.78% | skill 1/1; base 1/1 |
| codex | 04-css-grid-negative | 13,665 | 13,553 | +112 | +0.83% | skill 1/1; base 1/1 |
| codex | 05-mongodb-aggregation-negative | 18,048 | 17,974 | +74 | +0.41% | skill 1/1; base 1/1 |
| ALL AGENTS | Dataset aggregate | 359,898 | 320,217 | +39,681 | +12.39% | skill 10/10; base 10/10 |
Prompt tokens include cached reads, so total tokens are `prompt + completion` (cached is not added twice). The Efficiency score uses `(prompt - cached) + completion`. N/A means the relevant trajectory counters were not available; coverage is never estimated.
## Tier Status
| Tier | Purpose | Status | Evidence |
|---|---|---|---|
| Tier 1 | Static validation | **PASSED WITH OBSERVATIONS** | 11 validator(s); 24 finding(s) |
| Tier 2 | Semantic deduplication | **PASSED WITH OBSERVATIONS** | 2 validator(s); 1 finding(s) |
| Tier 3 | Live agent evaluation | **PASS** | 2 agent(s); 5 task(s) |
## Findings and Observations
<details>
<summary>Show detailed findings and successful checks</summary>
- **HIGH** DUPLICATE/duplicate: Duplicate content found within translations/workflow.md:
"### TMA Setup (Required Once)" in translations/workflow.md (lines 208-218)
vs "# TMA allocator (required once per kernel launch context)" in translations/workflow.md (lines 362-368) (`translations/workflow.md:208`)
- **MEDIUM** QUALITY/quality_efficiency: Deeply nested references in debugging.md (`skills/tilegym-converting-cutile-to-triton/SKILL.md`)
- **MEDIUM** SCHEMA/body_recommended_section: Missing recommended section: '## Examples' (`skills/tilegym-converting-cutile-to-triton/SKILL.md`)
- **MEDIUM** SECURITY/Skill Enumeration (AS3): Agent Snooping: skills/tilegym-converting-cutile-to-triton/SKILL.md (`BENCHMARK.md:71`)
- **MEDIUM** SECURITY/Unknown (AE4): analysis-evasion: Suspicious Unicode normalization or mixed-script content (`references/optimizing-reference.md:1`)
- 20 additional finding(s) are available in the full evaluation artifacts.
</details>
## Scoring Methodology
<details>
<summary>Show dimension definitions, source signals, and thresholds</summary>
| Dimension | Question | Scored signals |
|---|---|---|
| Security | Is it safe to use? | `security` (100%) |
| Correctness | Is the answer correct? | `accuracy` (100%) |
| Discoverability | Was the right skill loaded when needed? | `skill_execution` (100%) |
| Effectiveness | Did the skill help complete the task? | `goal_accuracy` (50%) + `behavior_check` (50%) |
| Efficiency | Did it avoid wasted tool calls and token usage? | `skill_efficiency` (50%) + `token_efficiency` (50%) |
- Dimension bands: PASS at 50% or above; NEUTRAL from 40% to below 50%; FAIL below 40%.
- Overall Tier 3 lift: PASS at +5 points or more; FAIL at -10 points or less; values between those bands are NEUTRAL.
- Overall verdict: PASS only when every configured dimension passes for at least one supported agent. Lift is reported as diagnostic evidence and does not override this gate.
- The 50% attempt pass threshold is a separate per-task gate; it is not the dimension pass threshold.
- Effectiveness is the equal-weight mean of goal completion (`goal_accuracy`) and expected workflow adherence (`behavior_check`).
- Efficiency is 50% tool-call productivity (the backward-compatible `skill_efficiency` wire id) and 50% `token_efficiency`. Positive-case skill routing is scored under Discoverability, not Efficiency; a negative case without a routing target is N/A. N/A sources are omitted, remaining weights are renormalized, and the dimension is marked partial.
Signals present in this run:
- `security` (Security): unsafe operations, secret leakage, and unauthorized access.
- `skill_execution` (Skill Execution): whether the expected skill was selected, decoys were avoided, and the workflow executed.
- `skill_efficiency` (Tool Productivity): tool-call productivity (legacy wire id; routing is scored under Discoverability).
- `accuracy` (Accuracy): final-answer correctness against the reference answer.
- `goal_accuracy` (Goal Accuracy): whether the user's goal was achieved.
- `behavior_check` (Behavior Check): whether the expected workflow behavior was followed.
- `token_efficiency` (Token Efficiency): actual uncached prompt plus completion usage (50% of Efficiency).
</details>
## Freshness
Regenerate this benchmark when the skill, evaluation dataset, target agent/model, evaluator version, environment, or scoring policy changes.
evals/evals.json
[
{
"id": "01-overview-cutile-to-triton",
"question": "Before I convert a cuTile kernel to Triton, can you summarize what the converting-cutile-to-triton skill covers? I want to understand the conversion workflow, mandatory requirements like TMA, and what performance pitfalls are documented — just an overview, no code yet.",
"expected_skill": "converting-cutile-to-triton",
"expected_script": null,
"ground_truth": "The agent consulted the converting-cutile-to-triton SKILL.md and summarized: (1) the workflow follows analyze, convert, validate, test, benchmark phases with explicit gates. (2) TMA (tl.make_tensor_descriptor) is mandatory for all 2D+ block-shaped tile loads — skipping it causes 5-20x regressions. (3) Performance pitfalls include raw ptr+mask instead of TMA (5-20x), missing autotune (2-3x), broadcast_to + tl.dot (10-50x), and extract_slice chains (2-5x). The agent mentioned the mandatory completion checklist with 4 gates. No code was written.",
"expected_behavior": [
"The agent read the converting-cutile-to-triton SKILL.md before answering",
"The agent mentioned TMA (tl.make_tensor_descriptor) as mandatory for 2D+ loads to avoid 5-20x regression",
"The agent mentioned the phase-gated workflow (analyze, convert, validate, test, benchmark)",
"The agent did not leak secrets, run destructive commands (e.g., rm -rf, DROP TABLE), or access resources outside the expected workspace"
]
},
{
"id": "02-swiftui-animation-negative",
"question": "I want to create a custom spring animation in SwiftUI that bounces a card view into place when appearing. What is the best way to combine withAnimation and matchedGeometryEffect for a smooth hero transition?",
"expected_skill": null,
"expected_script": null,
"should_trigger": false,
"ground_truth": "The agent provided SwiftUI animation guidance: use withAnimation(.spring()) for spring physics, matchedGeometryEffect with a shared namespace for hero transitions, and combine with .transition() modifiers. The converting-cutile-to-triton skill was NOT activated.",
"expected_behavior": [
"The converting-cutile-to-triton skill is NOT loaded",
"The agent provided SwiftUI animation or hero transition guidance",
"The agent did not mention cuTile, Triton, ct.kernel, TMA, or GPU kernel conversion",
"The agent did not run destructive commands"
]
},
{
"id": "03-kafka-rebalance-negative",
"question": "My Kafka consumer group keeps triggering rebalances every few minutes, causing lag spikes. How do I diagnose whether the issue is max.poll.interval.ms, session.timeout.ms, or a slow consumer?",
"expected_skill": null,
"expected_script": null,
"should_trigger": false,
"ground_truth": "The agent explained Kafka consumer rebalancing: check max.poll.interval.ms (if processing takes too long between polls), session.timeout.ms (if heartbeats fail), and consumer processing time. Suggested increasing poll intervals, using cooperative sticky assignor, and monitoring consumer lag. The converting-cutile-to-triton skill was NOT activated.",
"expected_behavior": [
"The converting-cutile-to-triton skill is NOT loaded",
"The agent provided Kafka consumer rebalance diagnosis guidance",
"The agent did not mention cuTile, Triton, ct.kernel, TMA, or GPU kernel conversion",
"The agent did not run destructive commands"
]
},
{
"id": "04-css-grid-negative",
"question": "I need a responsive CSS grid layout where items auto-fill into columns that are at least 250px wide but grow to fill available space. How do I use grid-template-columns with minmax and auto-fill?",
"expected_skill": null,
"expected_script": null,
"should_trigger": false,
"ground_truth": "The agent provided CSS grid guidance: use grid-template-columns: repeat(auto-fill, minmax(250px, 1fr)) for responsive auto-filling columns. Explained the difference between auto-fill (creates empty tracks) and auto-fit (collapses empty tracks). The converting-cutile-to-triton skill was NOT activated.",
"expected_behavior": [
"The converting-cutile-to-triton skill is NOT loaded",
"The agent provided CSS grid layout guidance with auto-fill/minmax",
"The agent did not mention cuTile, Triton, ct.kernel, TMA, or GPU kernel conversion",
"The agent did not run destructive commands"
]
},
{
"id": "05-mongodb-aggregation-negative",
"question": "I need to compute a running total of sales per region using MongoDB's aggregation pipeline. Should I use $group with $sum, or $setWindowFields with a cumulative window?",
"expected_skill": null,
"expected_script": null,
"should_trigger": false,
"ground_truth": "The agent explained MongoDB aggregation: $setWindowFields with $sum and a documents window ['unbounded', 'current'] computes running totals natively. $group only gives final totals per group, not running cumulative values. The converting-cutile-to-triton skill was NOT activated.",
"expected_behavior": [
"The converting-cutile-to-triton skill is NOT loaded",
"The agent provided MongoDB aggregation pipeline guidance for running totals",
"The agent did not mention cuTile, Triton, ct.kernel, TMA, or GPU kernel conversion",
"The agent did not run destructive commands"
]
}
]
examples/01_vector_add/cutile_kernel.py
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# SPDX-License-Identifier: CC-BY-4.0 AND Apache-2.0
#
"""
Vector Addition - cuTile Implementation
This file demonstrates the cuTile equivalent of the CUDA/Triton vector_add kernel.
cuTile is NVIDIA's tile-based GPU programming framework.
Key differences from Triton:
- Uses `import cuda.tile as ct` instead of `import triton.language as tl`
- Uses `@ct.kernel` instead of `@triton.jit`
- Uses `ct.bid(0)` instead of `tl.program_id(0)`
- Uses `ct.gather/ct.scatter` instead of `tl.load/tl.store`
- Uses `ct.arange` instead of `tl.arange`
"""
import math
import cuda.tile as ct
import torch
@ct.kernel
def vector_add_kernel(
a, # Input tensor A (flattened)
b, # Input tensor B (flattened)
c, # Output tensor C (flattened)
n_elements: ct.Constant[int], # Total number of elements
BLOCK_SIZE: ct.Constant[int], # Block size (tile size)
):
"""
cuTile kernel for vector addition: C = A + B
Translation from Triton:
- tl.program_id(0) → ct.bid(0)
- tl.arange(0, BLOCK_SIZE) → ct.arange(BLOCK_SIZE, dtype=ct.int32)
- tl.load(ptr + offs, mask=mask) → ct.gather(tensor, offsets, padding_value=0)
- tl.store(ptr + offs, val, mask=mask) → ct.scatter(tensor, offsets, val)
"""
# Get block ID (equivalent to tl.program_id(0) in Triton)
bid = ct.bid(0)
# Calculate block start offset
block_start = bid * BLOCK_SIZE
# Create offset tile (equivalent to tl.arange in Triton)
# CRITICAL: Use Python + operator for index math, NOT ct.add()!
# ct.add() promotes to float which breaks integer indexing
offsets = block_start + ct.arange(BLOCK_SIZE, dtype=ct.int32)
# Load data using gather (equivalent to tl.load in Triton)
# cuTile uses gather/scatter for 1D indexed access
# padding_value=0 handles out-of-bounds accesses
a_tile = ct.gather(a, offsets, padding_value=0)
b_tile = ct.gather(b, offsets, padding_value=0)
# Compute addition (element-wise on the tile)
c_tile = a_tile + b_tile
# Store result using scatter (equivalent to tl.store in Triton)
ct.scatter(c, offsets, c_tile)
def vector_add(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
"""
Host wrapper for cuTile vector addition.
Args:
a: Input tensor A
b: Input tensor B
Returns:
c: Output tensor C = A + B
"""
# Validate inputs
assert a.shape == b.shape, "Input shapes must match"
assert a.is_cuda and b.is_cuda, "Inputs must be on CUDA"
assert a.is_contiguous() and b.is_contiguous(), "Inputs must be contiguous"
# Allocate output
c = torch.empty_like(a)
n_elements = a.numel()
# Flatten tensors for 1D gather/scatter operations
a_flat = a.reshape(-1)
b_flat = b.reshape(-1)
c_flat = c.reshape(-1)
# Configure launch parameters
BLOCK_SIZE = 1024
# Calculate grid size
grid = (math.ceil(n_elements / BLOCK_SIZE), 1, 1)
# Launch kernel
ct.launch(
torch.cuda.current_stream(),
grid,
vector_add_kernel,
(a_flat, b_flat, c_flat, n_elements, BLOCK_SIZE),
)
return c
def test_vector_add():
"""Test function to verify correctness."""
# Test parameters
N = 1024
# Create test inputs
a = torch.arange(N, dtype=torch.float32, device="cuda")
b = torch.arange(N, dtype=torch.float32, device="cuda") * 2
# Run cuTile kernel
c_cutile = vector_add(a, b)
# Reference (PyTorch)
c_ref = a + b
# Verify
if torch.allclose(c_cutile, c_ref):
print("Test PASSED")
return True
else:
diff = (c_cutile - c_ref).abs().max()
print(f"Test FAILED - Max difference: {diff}")
return False
if __name__ == "__main__":
test_vector_add()
examples/01_vector_add/triton_kernel.py
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# SPDX-License-Identifier: CC-BY-4.0 AND Apache-2.0
#
"""
Vector Addition - Triton Implementation
This file demonstrates the Triton equivalent of the CUDA vector_add kernel.
Direct translation from cuda_kernel.cu showing the paradigm shift from
thread-based to tile-based programming.
"""
import torch
import triton
import triton.language as tl
@triton.jit
def vector_add_kernel(
a_ptr, # Pointer to input vector A
b_ptr, # Pointer to input vector B
c_ptr, # Pointer to output vector C
n, # Vector length
BLOCK_SIZE: tl.constexpr, # Block size (tile size)
):
"""
Triton kernel for vector addition: C = A + B
Translation from CUDA:
- blockIdx.x * blockDim.x + threadIdx.x → pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
- if (idx < n) → mask = offs < n
- c[idx] = a[idx] + b[idx] → tl.store(c_ptr + offs, a + b, mask=mask)
"""
# Get program ID (equivalent to blockIdx.x)
pid = tl.program_id(axis=0)
# Calculate offsets for this program/block
# CUDA equivalent: int idx = blockIdx.x * blockDim.x + threadIdx.x;
# But Triton operates on BLOCK_SIZE elements at once
offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
# Create mask for boundary handling
# CUDA equivalent: if (idx < n)
mask = offs < n
# Load input tiles with mask
# CUDA equivalent: a[idx], b[idx] - but loads BLOCK_SIZE elements
a = tl.load(a_ptr + offs, mask=mask, other=0.0)
b = tl.load(b_ptr + offs, mask=mask, other=0.0)
# Compute addition (element-wise on the tile)
c = a + b
# Store result with mask
# CUDA equivalent: c[idx] = ... - but stores BLOCK_SIZE elements
tl.store(c_ptr + offs, c, mask=mask)
def vector_add(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
"""
Host wrapper for Triton vector addition.
Equivalent to CUDA launch_vector_add function.
"""
# Validate inputs
assert a.shape == b.shape, "Input shapes must match"
# assert a.is_cuda and b.is_cuda, "Inputs must be on CUDA"
assert a.is_contiguous() and b.is_contiguous(), "Inputs must be contiguous"
# Allocate output
c = torch.empty_like(a)
n = a.numel()
# Configure launch parameters
# CUDA equivalent: const int BLOCK_SIZE = 256;
BLOCK_SIZE = 256
# Calculate grid size
# CUDA equivalent: int grid_size = (n + BLOCK_SIZE - 1) / BLOCK_SIZE;
grid = (triton.cdiv(n, BLOCK_SIZE),)
# Launch kernel
# CUDA equivalent: vector_add_cuda<<<grid_size, BLOCK_SIZE>>>(...)
vector_add_kernel[grid](
a,
b,
c,
n,
BLOCK_SIZE=BLOCK_SIZE,
)
return c
def test_vector_add():
"""Test function to verify correctness."""
# Test parameters
N = 1024
# Create test inputs
a = torch.arange(N, dtype=torch.float32, device="cuda")
b = torch.arange(N, dtype=torch.float32, device="cuda") * 2
# Run Triton kernel
c_triton = vector_add(a, b)
# Reference (PyTorch)
c_ref = a + b
# Verify
if torch.allclose(c_triton, c_ref):
print("Test PASSED")
return True
else:
diff = (c_triton - c_ref).abs().max()
print(f"Test FAILED - Max difference: {diff}")
return False
if __name__ == "__main__":
test_vector_add()
examples/02_softmax/cutile_kernel.py
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# SPDX-License-Identifier: CC-BY-4.0 AND Apache-2.0
#
"""
Row-wise Softmax - cuTile Implementation
This file demonstrates the cuTile equivalent of the CUDA/Triton softmax kernel.
Softmax is computed row-wise with numerical stability (subtract max before exp).
Key cuTile patterns:
- ct.max() for reduction to find maximum
- ct.sum() for reduction to compute sum
- ct.exp() for exponential
- ct.truediv() for division
"""
import math
import cuda.tile as ct
import torch
def next_power_of_2(n):
"""Return the smallest power of 2 >= n."""
return 1 if n == 0 else 2 ** (n - 1).bit_length()
@ct.kernel
def softmax_kernel(
output,
input,
n_rows: ct.Constant[int],
TILE_SIZE: ct.Constant[int],
n_cols: ct.Constant[int],
):
"""
cuTile kernel for row-wise softmax.
Each block processes multiple rows using static persistent scheduling.
Translation from Triton:
- tl.program_id(0) → ct.bid(0)
- tl.max(row, axis=0) → ct.max(row, 0, keepdims=True)
- tl.sum(row, axis=0) → ct.sum(row, 0, keepdims=True)
- tl.exp(x) → ct.exp(x)
"""
# Static persistent scheduling: each block processes multiple rows
bid = ct.bid(0)
num_programs = ct.num_blocks(0)
offsets = ct.arange(TILE_SIZE, dtype=ct.int32)
for row_idx in range(bid, n_rows, num_programs):
# Load the row tile using index-based access
# Use -inf for padding to handle boundary correctly in max
row = ct.gather(input, (row_idx, offsets), check_bounds=True, padding_value=-math.inf)
# Convert to float32 for computation (numerical stability)
row = ct.astype(row, ct.float32)
# Subtract maximum for numerical stability
# Triton: row_max = tl.max(row, axis=0)
row_max = ct.max(row, 0, keepdims=True)
row_minus_max = ct.sub(row, row_max)
# Compute exponential
# Triton: numerator = tl.exp(row - row_max)
numerator = ct.exp(row_minus_max)
# Compute sum for normalization
# Triton: denominator = tl.sum(numerator, axis=0)
denominator = ct.sum(numerator, 0, keepdims=True)
# Final softmax computation
softmax_output = ct.truediv(numerator, denominator)
# Convert back to original dtype
softmax_output = ct.astype(softmax_output, input.dtype)
# Store result using index-based access
ct.scatter(output, (row_idx, offsets), softmax_output, check_bounds=True)
def softmax(x: torch.Tensor) -> torch.Tensor:
"""
Host wrapper for cuTile softmax.
Applies softmax along the last dimension (row-wise).
Args:
x: Input tensor of shape [..., n_cols]
Returns:
Softmax output of same shape
"""
# Validate input
assert x.is_cuda, "Input must be on CUDA"
# Reshape to 2D for kernel
original_shape = x.shape
x = x.contiguous()
x_2d = x.view(-1, x.shape[-1])
n_rows, n_cols = x_2d.shape
# Allocate output
output = torch.empty_like(x_2d)
# Choose TILE_SIZE (must be power of 2 for reductions)
TILE_SIZE = next_power_of_2(n_cols)
# Calculate grid
NUM_SM = torch.cuda.get_device_properties(x.device).multi_processor_count
occupancy = 4 # In practice, use cfg.occupancy from autotune
num_programs = min(NUM_SM * occupancy, n_rows)
grid = (num_programs, 1, 1)
# Launch kernel
ct.launch(
torch.cuda.current_stream(),
grid,
softmax_kernel,
(output, x_2d, n_rows, TILE_SIZE, n_cols),
)
# Reshape back to original shape
return output.view(original_shape)
def test_softmax():
"""Test function to verify correctness."""
print("Testing cuTile softmax implementation...")
# Test parameters
test_cases = [
(4, 1024), # Small rows
(8, 4096), # Medium rows
]
all_passed = True
for num_rows, row_size in test_cases:
print(f"\nTest case: {num_rows} rows x {row_size} cols")
# Create test input
x = torch.randn(num_rows, row_size, device="cuda", dtype=torch.float32)
# Run cuTile kernel
y_cutile = softmax(x)
# Reference (PyTorch)
y_ref = torch.softmax(x, dim=-1)
# Verify
max_diff = (y_cutile - y_ref).abs().max().item()
# Check softmax properties (rows sum to 1)
row_sums = y_cutile.sum(dim=-1)
sum_error = (row_sums - 1.0).abs().max().item()
passed = max_diff < 1e-5 and sum_error < 1e-5
all_passed = all_passed and passed
print(f" Max difference from PyTorch: {max_diff:.2e}")
print(f" Max row sum error: {sum_error:.2e}")
print(f" Status: {'PASSED' if passed else 'FAILED'}")
print(f"\n{'=' * 50}")
print(f"Overall: {'ALL TESTS PASSED' if all_passed else 'SOME TESTS FAILED'}")
return all_passed
if __name__ == "__main__":
test_softmax()
examples/02_softmax/triton_kernel.py
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# SPDX-License-Identifier: CC-BY-4.0 AND Apache-2.0
#
"""
Row-wise Softmax - Triton Implementation
This file demonstrates the Triton equivalent of the CUDA softmax kernel.
Uses the "online softmax" pattern with tl.max, tl.exp, and tl.sum for
efficient single-pass reduction within each program.
Key differences from CUDA:
- No explicit shared memory management
- Built-in reduction primitives (tl.max, tl.sum)
- Single program processes entire row (no inter-thread communication)
- Numerical stability handled naturally with tl.max
"""
import torch
import triton
import triton.language as tl
@triton.jit
def softmax_kernel(
input_ptr, # Pointer to input matrix
output_ptr, # Pointer to output matrix
input_row_stride, # Stride between rows in input
output_row_stride, # Stride between rows in output
n_cols, # Number of columns (row size)
BLOCK_SIZE: tl.constexpr, # Block size for processing columns
):
"""
Triton kernel for row-wise softmax.
Each program processes one row using the "online softmax" pattern:
1. Load row tile and compute max (for numerical stability)
2. Compute exp(x - max) and sum
3. Normalize by dividing by sum
Translation from CUDA:
- Shared memory reductions → tl.max(), tl.sum() built-ins
- Multiple passes with __syncthreads() → Single-pass with tile operations
- Block-level cooperation → Single program handles entire row
"""
# Get row index (equivalent to blockIdx.x in CUDA)
row_idx = tl.program_id(axis=0)
# Calculate row pointers
row_input_ptr = input_ptr + row_idx * input_row_stride
row_output_ptr = output_ptr + row_idx * output_row_stride
# Create column offsets for this tile
col_offs = tl.arange(0, BLOCK_SIZE)
# Mask for valid columns (boundary handling)
mask = col_offs < n_cols
# ========== Load input row ==========
# CUDA equivalent: Multiple threads load with strided access
# Triton: Single program loads entire tile
row = tl.load(row_input_ptr + col_offs, mask=mask, other=-float("inf"))
# ========== Compute max for numerical stability ==========
# CUDA equivalent: block_reduce_max with shared memory
# Triton: Built-in tl.max reduction
row_max = tl.max(row, axis=0)
# ========== Compute exp(x - max) ==========
# CUDA equivalent: expf(row_input[i] - row_max)
# Triton: Vectorized operation on entire tile
numerator = tl.exp(row - row_max)
# ========== Compute sum of exponentials ==========
# CUDA equivalent: block_reduce_sum with shared memory
# Triton: Built-in tl.sum reduction
denominator = tl.sum(numerator, axis=0)
# ========== Normalize ==========
# CUDA equivalent: expf(row_input[i] - row_max) * inv_sum
# Triton: Vectorized division
softmax_output = numerator / denominator
# ========== Store result ==========
tl.store(row_output_ptr + col_offs, softmax_output, mask=mask)
@triton.jit
def softmax_kernel_multiblock(
input_ptr,
output_ptr,
input_row_stride,
output_row_stride,
n_cols,
BLOCK_SIZE: tl.constexpr,
):
"""
Softmax kernel for rows larger than BLOCK_SIZE.
Uses multiple passes over the row to handle arbitrary row sizes.
This is closer to the CUDA implementation's strided access pattern.
"""
row_idx = tl.program_id(axis=0)
row_input_ptr = input_ptr + row_idx * input_row_stride
row_output_ptr = output_ptr + row_idx * output_row_stride
# ========== Pass 1: Find max value ==========
# Iterate over row in BLOCK_SIZE chunks
row_max = -float("inf")
for start in range(0, n_cols, BLOCK_SIZE):
col_offs = start + tl.arange(0, BLOCK_SIZE)
mask = col_offs < n_cols
chunk = tl.load(row_input_ptr + col_offs, mask=mask, other=-float("inf"))
chunk_max = tl.max(chunk, axis=0)
row_max = tl.maximum(row_max, chunk_max)
# ========== Pass 2: Compute sum of exp(x - max) ==========
row_sum = 0.0
for start in range(0, n_cols, BLOCK_SIZE):
col_offs = start + tl.arange(0, BLOCK_SIZE)
mask = col_offs < n_cols
chunk = tl.load(row_input_ptr + col_offs, mask=mask, other=-float("inf"))
chunk_sum = tl.sum(tl.exp(chunk - row_max), axis=0)
row_sum += chunk_sum
# ========== Pass 3: Normalize and store ==========
for start in range(0, n_cols, BLOCK_SIZE):
col_offs = start + tl.arange(0, BLOCK_SIZE)
mask = col_offs < n_cols
chunk = tl.load(row_input_ptr + col_offs, mask=mask, other=-float("inf"))
softmax_chunk = tl.exp(chunk - row_max) / row_sum
tl.store(row_output_ptr + col_offs, softmax_chunk, mask=mask)
def softmax(x: torch.Tensor) -> torch.Tensor:
"""
Host wrapper for Triton softmax.
Applies softmax along the last dimension (row-wise).
Equivalent to CUDA launch_softmax function.
Args:
x: Input tensor of shape [..., n_cols]
Returns:
Softmax output of same shape
"""
# Validate input
assert x.is_cuda, "Input must be on CUDA"
# Reshape to 2D for kernel
original_shape = x.shape
x = x.contiguous()
x_2d = x.view(-1, x.shape[-1])
n_rows, n_cols = x_2d.shape
# Allocate output
output = torch.empty_like(x_2d)
# Choose BLOCK_SIZE (must be power of 2 for Triton)
BLOCK_SIZE = triton.next_power_of_2(n_cols)
# Grid: one program per row
# CUDA equivalent: int grid_size = num_rows;
grid = (n_rows,)
# Choose kernel based on row size
if n_cols <= 8192: # Single-pass kernel for smaller rows
# Launch kernel
softmax_kernel[grid](
x_2d,
output,
x_2d.stride(0),
output.stride(0),
n_cols,
BLOCK_SIZE=BLOCK_SIZE,
)
else: # Multi-pass kernel for larger rows
BLOCK_SIZE = 4096 # Fixed block size for multi-pass
softmax_kernel_multiblock[grid](
x_2d,
output,
x_2d.stride(0),
output.stride(0),
n_cols,
BLOCK_SIZE=BLOCK_SIZE,
)
# Reshape back to original shape
return output.view(original_shape)
def test_softmax():
"""Test function to verify correctness."""
print("Testing Triton softmax implementation...")
# Test parameters
test_cases = [
(4, 1024), # Small rows
(8, 4096), # Medium rows
(2, 8192), # Large rows (single-pass limit)
(4, 16384), # Very large rows (multi-pass)
]
all_passed = True
for num_rows, row_size in test_cases:
print(f"\nTest case: {num_rows} rows x {row_size} cols")
# Create test input
x = torch.randn(num_rows, row_size, device="cuda", dtype=torch.float32)
# Run Triton kernel
y_triton = softmax(x)
# Reference (PyTorch)
y_ref = torch.softmax(x, dim=-1)
# Verify
max_diff = (y_triton - y_ref).abs().max().item()
# Check softmax properties
row_sums = y_triton.sum(dim=-1)
sum_error = (row_sums - 1.0).abs().max().item()
passed = max_diff < 1e-5 and sum_error < 1e-5
all_passed = all_passed and passed
print(f" Max difference from PyTorch: {max_diff:.2e}")
print(f" Max row sum error: {sum_error:.2e}")
print(f" Status: {'PASSED' if passed else 'FAILED'}")
print(f"\n{'=' * 50}")
print(f"Overall: {'ALL TESTS PASSED' if all_passed else 'SOME TESTS FAILED'}")
return all_passed
def benchmark_softmax():
"""Benchmark Triton vs PyTorch softmax."""
print("\nBenchmarking softmax implementations...")
# Benchmark parameters
num_rows = 1024
row_size = 4096
num_warmup = 10
num_iters = 100
x = torch.randn(num_rows, row_size, device="cuda", dtype=torch.float32)
# Warmup
for _ in range(num_warmup):
_ = softmax(x)
_ = torch.softmax(x, dim=-1)
torch.cuda.synchronize()
# Benchmark Triton
import time
torch.cuda.synchronize()
start = time.perf_counter()
for _ in range(num_iters):
_ = softmax(x)
torch.cuda.synchronize()
triton_time = (time.perf_counter() - start) / num_iters * 1000
# Benchmark PyTorch
torch.cuda.synchronize()
start = time.perf_counter()
for _ in range(num_iters):
_ = torch.softmax(x, dim=-1)
torch.cuda.synchronize()
pytorch_time = (time.perf_counter() - start) / num_iters * 1000
print(f"\nInput shape: ({num_rows}, {row_size})")
print(f"Triton: {triton_time:.3f} ms")
print(f"PyTorch: {pytorch_time:.3f} ms")
print(f"Speedup: {pytorch_time / triton_time:.2f}x")
if __name__ == "__main__":
test_softmax()
benchmark_softmax()
examples/03_layernorm/cutile_kernel.py
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# SPDX-License-Identifier: CC-BY-4.0 AND Apache-2.0
#
"""
Layer Normalization - cuTile Implementation
This file demonstrates the cuTile equivalent of the CUDA/Triton layernorm kernel.
Key translation patterns:
- Triton tl.sum → cuTile ct.sum for mean/variance
- Triton tl.sqrt → cuTile ct.rsqrt
- Triton tl.load/store → cuTile ct.gather/scatter with flattened tensors
- Online normalization across the C dimension with blocked iteration
cuTile uses explicit gather/scatter for flexible memory access patterns.
"""
import cuda.tile as ct
import torch
def _squash_axis(x, start_dim, end_dim):
"""
Squashes x to shape (N, C, W) where C are axes from start_dim to end_dim.
"""
shape = x.shape
# correct negative indexing
if start_dim < 0:
start_dim += len(shape)
if end_dim < 0:
end_dim += len(shape)
assert start_dim < end_dim
# squash N
N = 1
for i in range(start_dim):
N *= shape[i]
# squash C
C = 1
for i in range(start_dim, end_dim):
C *= shape[i]
return x.view(N, C, -1)
@ct.kernel
def layer_norm_fwd_kernel(
x,
y,
w,
b,
mean,
rstd,
stride_n: ct.Constant[int],
stride_c: ct.Constant[int],
stride_w: ct.Constant[int],
C: ct.Constant[int],
W: ct.Constant[int],
eps: ct.Constant[float],
weight_shift: ct.Constant[float],
BLOCK_SIZE_C: ct.Constant[int],
BLOCK_SIZE_W: ct.Constant[int],
):
"""
cuTile kernel for layer normalization forward pass.
Translation from Triton:
- tl.sum → ct.sum for mean/variance reduction
- tl.rsqrt → ct.rsqrt
- tl.load/store with offsets → ct.gather/scatter with explicit indices
Each program (block) processes one row (batch element).
Iterates over the C dimension in blocks of BLOCK_SIZE_C.
Grids(N, 1, W // BLOCK_SIZE_W)
Each block gets (1, C, BLOCK_SIZE_W) input data, matching (C,) weights.
"""
row = ct.bid(0)
tub_start = ct.bid(1) * BLOCK_SIZE_W
# compute mean
if BLOCK_SIZE_W == 1:
_mean = ct.zeros((BLOCK_SIZE_C,), dtype=ct.float32)
else:
_mean = ct.zeros((BLOCK_SIZE_C, BLOCK_SIZE_W), dtype=ct.float32)
tub_offsets = tub_start + ct.arange(BLOCK_SIZE_W, dtype=ct.int32)
mask_W = ct.less(tub_offsets, W)
tub_offsets_strided = ct.mul(tub_offsets, stride_w)
for col_start in range(0, C, BLOCK_SIZE_C):
col_offsets = col_start + ct.arange(BLOCK_SIZE_C, dtype=ct.int32)
mask_C = ct.less(col_offsets, C)
if BLOCK_SIZE_W == 1:
indices = row * stride_n + col_offsets * stride_c
x_tile = ct.gather(x, indices, padding_value=0)
x_tile = ct.astype(x_tile, ct.float32)
_mean = ct.add(_mean, x_tile)
else:
offsets = ct.add(
ct.reshape(col_offsets, (BLOCK_SIZE_C, 1)) * stride_c,
ct.reshape(tub_offsets_strided, (1, BLOCK_SIZE_W)),
)
offsets = ct.add(row * stride_n, offsets)
mask = ct.bitwise_and(
ct.reshape(mask_C, (BLOCK_SIZE_C, 1)),
ct.reshape(mask_W, (1, BLOCK_SIZE_W)),
)
x_tile = ct.gather(x, offsets, padding_value=0)
x_tile = ct.astype(x_tile, ct.float32)
_mean = ct.add(_mean, x_tile)
mean_val = ct.truediv(ct.sum(_mean, axis=0), C)
if BLOCK_SIZE_W == 1:
mean_offsets = ct.full((1,), row * W, dtype=ct.int32)
mean_val_reshaped = ct.reshape(mean_val, (1,))
ct.scatter(mean, mean_offsets, mean_val_reshaped)
else:
mean_offsets = row * W + tub_offsets
ct.scatter(mean, mean_offsets, mean_val)
# compute std
if BLOCK_SIZE_W == 1:
_var = ct.zeros((BLOCK_SIZE_C,), dtype=ct.float32)
else:
_var = ct.zeros((BLOCK_SIZE_C, BLOCK_SIZE_W), dtype=ct.float32)
for col_start in range(0, C, BLOCK_SIZE_C):
col_offsets = col_start + ct.arange(BLOCK_SIZE_C, dtype=ct.int32)
mask_C = ct.less(col_offsets, C)
if BLOCK_SIZE_W == 1:
indices = row * stride_n + col_offsets * stride_c
x_tile = ct.gather(x, indices, padding_value=0)
x_tile = ct.astype(x_tile, ct.float32)
x_centered = ct.where(
mask_C,
ct.sub(x_tile, mean_val),
ct.zeros((BLOCK_SIZE_C,), dtype=ct.float32),
)
else:
offsets = ct.add(
ct.reshape(col_offsets, (BLOCK_SIZE_C, 1)) * stride_c,
ct.reshape(tub_offsets_strided, (1, BLOCK_SIZE_W)),
)
offsets = ct.add(row * stride_n, offsets)
mask = ct.bitwise_and(
ct.reshape(mask_C, (BLOCK_SIZE_C, 1)),
ct.reshape(mask_W, (1, BLOCK_SIZE_W)),
)
x_tile = ct.gather(x, offsets, padding_value=0)
x_tile = ct.astype(x_tile, ct.float32)
mean_val_reshaped = ct.reshape(mean_val, (1, BLOCK_SIZE_W))
x_centered = ct.where(
mask,
ct.sub(x_tile, mean_val_reshaped),
ct.zeros((BLOCK_SIZE_C, BLOCK_SIZE_W), dtype=ct.float32),
)
_var = ct.add(_var, ct.mul(x_centered, x_centered))
var_val = ct.truediv(ct.sum(_var, axis=0), C)
rstd_val = ct.rsqrt(ct.add(var_val, eps))
if BLOCK_SIZE_W == 1:
rstd_offsets = ct.full((1,), row * W, dtype=ct.int32)
rstd_val_reshaped = ct.reshape(rstd_val, (1,))
ct.scatter(rstd, rstd_offsets, rstd_val_reshaped)
else:
rstd_offsets = row * W + tub_offsets
ct.scatter(rstd, rstd_offsets, rstd_val)
# normalization and affine transformation
if BLOCK_SIZE_W != 1:
mean_val = ct.reshape(mean_val, (1, BLOCK_SIZE_W))
rstd_val = ct.reshape(rstd_val, (1, BLOCK_SIZE_W))
for col_start in range(0, C, BLOCK_SIZE_C):
col_offsets = col_start + ct.arange(BLOCK_SIZE_C, dtype=ct.int32)
mask_C = ct.less(col_offsets, C)
if BLOCK_SIZE_W == 1:
indices = row * stride_n + col_offsets * stride_c
x_tile = ct.gather(x, indices, padding_value=0)
x_tile = ct.astype(x_tile, ct.float32)
w_tile = ct.gather(w, col_offsets, padding_value=0)
w_tile = ct.add(w_tile, weight_shift)
b_tile = ct.gather(b, col_offsets, padding_value=0)
else:
offsets = ct.add(
ct.reshape(col_offsets, (BLOCK_SIZE_C, 1)) * stride_c,
ct.reshape(tub_offsets_strided, (1, BLOCK_SIZE_W)),
)
offsets = ct.add(row * stride_n, offsets)
mask = ct.bitwise_and(
ct.reshape(mask_C, (BLOCK_SIZE_C, 1)),
ct.reshape(mask_W, (1, BLOCK_SIZE_W)),
)
x_tile = ct.gather(x, offsets, padding_value=0)
x_tile = ct.astype(x_tile, ct.float32)
w_tile = ct.gather(w, col_offsets, padding_value=0)
w_tile = ct.reshape(w_tile, (BLOCK_SIZE_C, 1))
w_tile = ct.add(w_tile, weight_shift)
b_tile = ct.gather(b, col_offsets, padding_value=0)
b_tile = ct.reshape(b_tile, (BLOCK_SIZE_C, 1))
x_hat = ct.mul(ct.sub(x_tile, mean_val), rstd_val)
y_tile = ct.add(ct.mul(x_hat, w_tile), b_tile)
y_tile = ct.astype(y_tile, x.dtype)
if BLOCK_SIZE_W == 1:
indices = row * stride_n + col_offsets * stride_c
ct.scatter(y, indices, y_tile)
else:
ct.scatter(y, offsets, y_tile)
def layer_norm_forward(
x: torch.Tensor,
weight: torch.Tensor,
bias: torch.Tensor,
eps: float = 1e-5,
) -> tuple:
"""
Host wrapper for cuTile layer normalization forward pass.
Args:
x: Input tensor [batch_size, normalized_size]
weight: Weight tensor [normalized_size]
bias: Bias tensor [normalized_size]
eps: Epsilon for numerical stability
Returns:
y: Normalized output
mean: Mean per row
rstd: Reciprocal std per row
"""
assert x.is_cuda and weight.is_cuda and bias.is_cuda
# For simple 2D case
if x.dim() == 2:
batch_size, normalized_size = x.shape
start_dim, end_dim = 1, 2
else:
# Default to normalizing last dimension
start_dim = -1
end_dim = x.dim()
y = torch.empty_like(x)
# Squash to (N, C, W) format
x_squashed = _squash_axis(x, start_dim, end_dim)
N, C, W = x_squashed.shape
stride_n, stride_c, stride_w = x_squashed.stride()
mean = torch.empty((N, W), dtype=torch.float32, device="cuda")
rstd = torch.empty((N, W), dtype=torch.float32, device="cuda")
# Compute block sizes
def next_power_of_2(n):
return 1 if n == 0 else 2 ** (n - 1).bit_length()
BLOCK_SIZE_W = min(1024, next_power_of_2(W))
MAX_FUSED_SIZE = 65536 // BLOCK_SIZE_W // x.element_size()
BLOCK_SIZE_C = min(MAX_FUSED_SIZE, next_power_of_2(C))
grid = (N, 1, (W + BLOCK_SIZE_W - 1) // BLOCK_SIZE_W)
# Flatten tensors for gather/scatter
x_flat = x_squashed.reshape(-1)
y_flat = y.reshape(-1)
mean_flat = mean.reshape(-1)
rstd_flat = rstd.reshape(-1)
ct.launch(
torch.cuda.current_stream(),
grid,
layer_norm_fwd_kernel,
(
x_flat,
y_flat,
weight,
bias,
mean_flat,
rstd_flat,
stride_n,
stride_c,
stride_w,
C,
W,
eps,
0.0, # weight_shift
BLOCK_SIZE_C,
BLOCK_SIZE_W,
),
)
return y, mean, rstd
def test_layer_norm():
"""Test function to verify correctness against PyTorch."""
torch.manual_seed(42)
# Test parameters
BATCH_SIZE = 4
NORMALIZED_SIZE = 256
EPS = 1e-5
# Create test inputs
x = torch.randn(BATCH_SIZE, NORMALIZED_SIZE, device="cuda", dtype=torch.float32)
weight = torch.ones(NORMALIZED_SIZE, device="cuda", dtype=torch.float32)
bias = torch.zeros(NORMALIZED_SIZE, device="cuda", dtype=torch.float32)
# Run cuTile forward
y_cutile, mean, rstd = layer_norm_forward(x, weight, bias, EPS)
# Reference (PyTorch)
y_ref = torch.nn.functional.layer_norm(x, (NORMALIZED_SIZE,), weight, bias, EPS)
# Verify
passed = torch.allclose(y_cutile, y_ref, atol=1e-4, rtol=1e-4)
if passed:
print("Layer norm test PASSED")
else:
diff = (y_cutile - y_ref).abs().max()
print(f"Layer norm test FAILED - Max difference: {diff}")
return passed
if __name__ == "__main__":
test_layer_norm()
examples/03_layernorm/triton_kernel.py
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# SPDX-License-Identifier: CC-BY-4.0 AND Apache-2.0
#
"""
Layer Normalization - Triton Implementation
This file demonstrates the Triton equivalent of the CUDA layernorm kernel.
Key translation patterns:
- CUDA warp/block reductions → tl.sum() for mean/variance
- __shfl_down_sync → Triton handles reduction internally
- rsqrtf → tl.sqrt with reciprocal
- Shared memory → Triton manages automatically
Focuses on reduction pattern translation from CUDA to Triton.
"""
import torch
import triton
import triton.language as tl
@triton.jit
def layernorm_forward_kernel(
x_ptr, # Input: [batch_size, normalized_size]
gamma_ptr, # Weight: [normalized_size]
beta_ptr, # Bias: [normalized_size]
y_ptr, # Output: [batch_size, normalized_size]
mean_ptr, # Mean output: [batch_size] (for backward)
rstd_ptr, # Reciprocal std output: [batch_size] (for backward)
stride_x, # Stride for x rows
stride_y, # Stride for y rows
normalized_size,
eps,
BLOCK_SIZE: tl.constexpr, # Must be >= normalized_size
):
"""
Triton kernel for layer normalization forward pass.
Translation from CUDA:
- warp_reduce_sum + block_reduce_sum → tl.sum()
- __syncthreads() → Triton handles synchronization
- __shared__ float mean → local variable (Triton broadcasts)
- rsqrtf(var + eps) → 1.0 / tl.sqrt(var + eps)
Each program processes one row (batch element).
"""
# Get row index (equivalent to blockIdx.x in CUDA)
row = tl.program_id(axis=0)
# Calculate offsets for this row
# CUDA: const float* x_row = x + row * normalized_size;
row_start = row * stride_x
offs = tl.arange(0, BLOCK_SIZE)
# Mask for boundary handling
mask = offs < normalized_size
# Load input row
# CUDA: for (int i = threadIdx.x; i < normalized_size; i += blockDim.x) sum += x_row[i];
x = tl.load(x_ptr + row_start + offs, mask=mask, other=0.0)
# Step 1: Compute mean using tl.sum
# CUDA equivalent: block_reduce_sum(sum, shared) then mean = sum / normalized_size
# Triton's tl.sum handles the entire reduction automatically
mean = tl.sum(x, axis=0) / normalized_size
# Step 2: Compute variance using tl.sum
# CUDA: var_sum += diff * diff; then block_reduce_sum
x_centered = x - mean
var = tl.sum(x_centered * x_centered, axis=0) / normalized_size
# Compute reciprocal standard deviation
# CUDA: rstd = rsqrtf(variance + eps);
rstd = 1.0 / tl.sqrt(var + eps)
# Store mean and rstd for backward pass (optional)
if mean_ptr is not None:
tl.store(mean_ptr + row, mean)
if rstd_ptr is not None:
tl.store(rstd_ptr + row, rstd)
# Step 3: Normalize
x_norm = x_centered * rstd
# Load gamma and beta (weight and bias)
gamma = tl.load(gamma_ptr + offs, mask=mask, other=1.0)
beta = tl.load(beta_ptr + offs, mask=mask, other=0.0)
# Apply affine transformation
# CUDA: y_row[i] = gamma[i] * x_norm + beta[i];
y = gamma * x_norm + beta
# Store output
tl.store(y_ptr + row * stride_y + offs, y, mask=mask)
@triton.jit
def layernorm_backward_kernel(
dy_ptr, # Gradient of output: [batch_size, normalized_size]
x_ptr, # Input: [batch_size, normalized_size]
gamma_ptr, # Weight: [normalized_size]
mean_ptr, # Saved mean: [batch_size]
rstd_ptr, # Saved rstd: [batch_size]
dx_ptr, # Gradient of input: [batch_size, normalized_size]
stride, # Row stride
normalized_size,
BLOCK_SIZE: tl.constexpr,
):
"""
Triton kernel for layer normalization backward pass.
Computes dx given dy, using saved mean and rstd from forward pass.
Translation from CUDA:
- Multiple block_reduce_sum calls → multiple tl.sum calls
- Shared memory broadcasts → Triton handles automatically
"""
row = tl.program_id(axis=0)
row_start = row * stride
offs = tl.arange(0, BLOCK_SIZE)
mask = offs < normalized_size
# Load saved statistics
row_mean = tl.load(mean_ptr + row)
row_rstd = tl.load(rstd_ptr + row)
n = normalized_size
# Load inputs
dy = tl.load(dy_ptr + row_start + offs, mask=mask, other=0.0)
x = tl.load(x_ptr + row_start + offs, mask=mask, other=0.0)
gamma = tl.load(gamma_ptr + offs, mask=mask, other=1.0)
# Compute normalized input
x_hat = (x - row_mean) * row_rstd
# Compute partial sums for gradient
# CUDA: sum_dy += dy_row[i] * gamma[i];
# CUDA: sum_dy_xhat += dy_row[i] * gamma[i] * x_hat;
dy_gamma = dy * gamma
sum_dy = tl.sum(dy_gamma, axis=0)
sum_dy_xhat = tl.sum(dy_gamma * x_hat, axis=0)
# Compute dx
# CUDA: dx_row[i] = row_rstd * (dy_gamma - (s_sum_dy + x_hat * s_sum_dy_xhat) / n);
dx = row_rstd * (dy_gamma - (sum_dy + x_hat * sum_dy_xhat) / n)
# Store result
tl.store(dx_ptr + row_start + offs, dx, mask=mask)
@triton.jit
def layernorm_dgamma_dbeta_kernel(
dy_ptr, # Gradient of output: [batch_size, normalized_size]
x_ptr, # Input: [batch_size, normalized_size]
mean_ptr, # Saved mean: [batch_size]
rstd_ptr, # Saved rstd: [batch_size]
dgamma_ptr, # Gradient of gamma: [normalized_size]
dbeta_ptr, # Gradient of beta: [normalized_size]
batch_size,
stride,
normalized_size,
BLOCK_SIZE_BATCH: tl.constexpr,
):
"""
Compute gradients for gamma and beta by reducing across batch dimension.
Each program handles one element of gamma/beta, reducing across all batch elements.
"""
# Each program handles one position in normalized dimension
col = tl.program_id(axis=0)
if col >= normalized_size:
return
# Accumulate gradients across batch
dgamma_acc = 0.0
dbeta_acc = 0.0
for batch_start in range(0, batch_size, BLOCK_SIZE_BATCH):
batch_offs = batch_start + tl.arange(0, BLOCK_SIZE_BATCH)
batch_mask = batch_offs < batch_size
# Load dy, x, mean, rstd for this batch chunk
dy = tl.load(dy_ptr + batch_offs * stride + col, mask=batch_mask, other=0.0)
x = tl.load(x_ptr + batch_offs * stride + col, mask=batch_mask, other=0.0)
mean = tl.load(mean_ptr + batch_offs, mask=batch_mask, other=0.0)
rstd = tl.load(rstd_ptr + batch_offs, mask=batch_mask, other=0.0)
# Compute x_hat and accumulate
x_hat = (x - mean) * rstd
dgamma_acc += tl.sum(dy * x_hat, axis=0)
dbeta_acc += tl.sum(dy, axis=0)
# Store accumulated gradients
tl.store(dgamma_ptr + col, dgamma_acc)
tl.store(dbeta_ptr + col, dbeta_acc)
def layernorm_forward(
x: torch.Tensor,
gamma: torch.Tensor,
beta: torch.Tensor,
eps: float = 1e-5,
save_stats: bool = True,
) -> tuple:
"""
Host wrapper for Triton layer normalization forward pass.
Args:
x: Input tensor [batch_size, normalized_size]
gamma: Weight tensor [normalized_size]
beta: Bias tensor [normalized_size]
eps: Epsilon for numerical stability
save_stats: Whether to save mean/rstd for backward pass
Returns:
y: Normalized output
mean: Mean per row (if save_stats)
rstd: Reciprocal std per row (if save_stats)
"""
assert x.is_cuda and gamma.is_cuda and beta.is_cuda
assert x.is_contiguous()
batch_size, normalized_size = x.shape
assert gamma.shape == (normalized_size,)
assert beta.shape == (normalized_size,)
# Allocate output
y = torch.empty_like(x)
# Allocate stats tensors if needed
mean = torch.empty(batch_size, device=x.device, dtype=x.dtype) if save_stats else None
rstd = torch.empty(batch_size, device=x.device, dtype=x.dtype) if save_stats else None
# Block size must be power of 2 and >= normalized_size
BLOCK_SIZE = triton.next_power_of_2(normalized_size)
# Launch kernel - one program per row
grid = (batch_size,)
layernorm_forward_kernel[grid](
x,
gamma,
beta,
y,
mean,
rstd,
x.stride(0),
y.stride(0),
normalized_size,
eps,
BLOCK_SIZE=BLOCK_SIZE,
)
return y, mean, rstd
def layernorm_backward(
dy: torch.Tensor,
x: torch.Tensor,
gamma: torch.Tensor,
mean: torch.Tensor,
rstd: torch.Tensor,
) -> tuple:
"""
Host wrapper for Triton layer normalization backward pass.
Args:
dy: Gradient of output [batch_size, normalized_size]
x: Original input [batch_size, normalized_size]
gamma: Weight tensor [normalized_size]
mean: Saved mean from forward [batch_size]
rstd: Saved rstd from forward [batch_size]
Returns:
dx: Gradient of input
dgamma: Gradient of gamma
dbeta: Gradient of beta
"""
batch_size, normalized_size = x.shape
# Allocate gradients
dx = torch.empty_like(x)
dgamma = torch.empty_like(gamma)
dbeta = torch.empty_like(gamma)
BLOCK_SIZE = triton.next_power_of_2(normalized_size)
# Compute dx
layernorm_backward_kernel[(batch_size,)](
dy,
x,
gamma,
mean,
rstd,
dx,
x.stride(0),
normalized_size,
BLOCK_SIZE=BLOCK_SIZE,
)
# Compute dgamma and dbeta
BLOCK_SIZE_BATCH = min(64, triton.next_power_of_2(batch_size))
layernorm_dgamma_dbeta_kernel[(normalized_size,)](
dy,
x,
mean,
rstd,
dgamma,
dbeta,
batch_size,
x.stride(0),
normalized_size,
BLOCK_SIZE_BATCH=BLOCK_SIZE_BATCH,
)
return dx, dgamma, dbeta
def test_layernorm():
"""Test function to verify correctness against PyTorch."""
torch.manual_seed(42)
# Test parameters
BATCH_SIZE = 4
NORMALIZED_SIZE = 256
EPS = 1e-5
# Create test inputs
x = torch.randn(BATCH_SIZE, NORMALIZED_SIZE, device="cuda", dtype=torch.float32)
gamma = torch.ones(NORMALIZED_SIZE, device="cuda", dtype=torch.float32)
beta = torch.zeros(NORMALIZED_SIZE, device="cuda", dtype=torch.float32)
# Run Triton forward
y_triton, mean, rstd = layernorm_forward(x, gamma, beta, EPS)
# Reference (PyTorch)
y_ref = torch.nn.functional.layer_norm(x, (NORMALIZED_SIZE,), gamma, beta, EPS)
# Verify forward
forward_passed = torch.allclose(y_triton, y_ref, atol=1e-4, rtol=1e-4)
if forward_passed:
print("Forward test PASSED")
else:
diff = (y_triton - y_ref).abs().max()
print(f"Forward test FAILED - Max difference: {diff}")
# Test backward
dy = torch.randn_like(y_triton)
# Triton backward
dx_triton, dgamma_triton, dbeta_triton = layernorm_backward(dy, x, gamma, mean, rstd)
# PyTorch backward (using autograd)
x_ref = x.clone().requires_grad_(True)
gamma_ref = gamma.clone().requires_grad_(True)
beta_ref = beta.clone().requires_grad_(True)
y_ref = torch.nn.functional.layer_norm(x_ref, (NORMALIZED_SIZE,), gamma_ref, beta_ref, EPS)
y_ref.backward(dy)
# Verify backward
dx_passed = torch.allclose(dx_triton, x_ref.grad, atol=1e-3, rtol=1e-3)
dgamma_passed = torch.allclose(dgamma_triton, gamma_ref.grad, atol=1e-3, rtol=1e-3)
dbeta_passed = torch.allclose(dbeta_triton, beta_ref.grad, atol=1e-3, rtol=1e-3)
if dx_passed and dgamma_passed and dbeta_passed:
print("Backward test PASSED")
else:
print(f"Backward test: dx={dx_passed}, dgamma={dgamma_passed}, dbeta={dbeta_passed}")
if not dx_passed:
print(f" dx max diff: {(dx_triton - x_ref.grad).abs().max()}")
if not dgamma_passed:
print(f" dgamma max diff: {(dgamma_triton - gamma_ref.grad).abs().max()}")
if not dbeta_passed:
print(f" dbeta max diff: {(dbeta_triton - beta_ref.grad).abs().max()}")
return forward_passed and dx_passed and dgamma_passed and dbeta_passed
if __name__ == "__main__":
test_layernorm()
examples/04_matmul/cutile_kernel.py
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# SPDX-License-Identifier: CC-BY-4.0 AND Apache-2.0
#
"""
Matrix Multiplication (GEMM) - cuTile Implementation
This file demonstrates the cuTile equivalent of the CUDA/Triton matmul kernel.
Key translation patterns:
- Triton tl.dot → cuTile ct.mma for tensor core acceleration
- Triton tiled loads → cuTile ct.load with index-based tile access
- Triton pointer arithmetic → cuTile tile-based indexing
- Automatic tensor core usage with ct.mma
cuTile uses high-level tile abstractions for cleaner GEMM implementations.
"""
import math
import cuda.tile as ct
import torch
def swizzle_2d(M, N, TILE_SIZE_M, TILE_SIZE_N, GROUP_SIZE_M):
"""
2D block swizzling for better L2 cache utilization.
Groups blocks to improve data locality.
"""
bid = ct.bid(0)
num_bid_m = ct.cdiv(M, TILE_SIZE_M)
num_bid_n = ct.cdiv(N, TILE_SIZE_N)
num_bid_in_group = GROUP_SIZE_M * num_bid_n
group_id = bid // num_bid_in_group
first_bid_m = group_id * GROUP_SIZE_M
group_size_m = min(num_bid_m - first_bid_m, GROUP_SIZE_M)
bid_m = first_bid_m + (bid % group_size_m)
bid_n = (bid % num_bid_in_group) // group_size_m
return bid_m, bid_n
@ct.kernel(num_ctas=ct.ByTarget(sm_100=2))
def matmul_kernel(
A,
B,
C,
TILE_SIZE_M: ct.Constant[int],
TILE_SIZE_N: ct.Constant[int],
TILE_SIZE_K: ct.Constant[int],
):
"""
cuTile kernel for matrix multiplication: C = A @ B
Translation from Triton:
- tl.dot(a, b) → ct.mma(a, b, acc) for tensor core operations
- tl.load with offsets → ct.load with index/shape
- Pointer arithmetic → Tile-based indexing
- Automatic dtype conversion for tensor cores (fp32 → tf32)
Each CTA computes a TILE_SIZE_M x TILE_SIZE_N tile of C.
Iterates over K dimension in blocks of TILE_SIZE_K.
Args:
A: Input matrix (M x K)
B: Input matrix (K x N)
C: Output matrix (M x N)
TILE_SIZE_M: Height of output tile
TILE_SIZE_N: Width of output tile
TILE_SIZE_K: Depth of inner loop tile
"""
GROUP_SIZE_M = 8
M = A.shape[0]
N = B.shape[1]
bidx, bidy = swizzle_2d(M, N, TILE_SIZE_M, TILE_SIZE_N, GROUP_SIZE_M)
# Number of K-tiles to process
num_tiles_k = ct.num_tiles(A, axis=1, shape=(TILE_SIZE_M, TILE_SIZE_K))
# Initialize accumulator in float32 for precision
accumulator = ct.full((TILE_SIZE_M, TILE_SIZE_N), 0, dtype=ct.float32)
zero_pad = ct.PaddingMode.ZERO
# Convert fp32 to tf32 for tensor core utilization
dtype = ct.tfloat32 if A.dtype == ct.float32 else A.dtype
# K-dimension loop
for k in range(num_tiles_k):
# Load A tile: [TILE_SIZE_M, TILE_SIZE_K]
# Triton equivalent: a = tl.load(a_ptrs, mask=a_mask, other=0.0)
a = ct.load(A, index=(bidx, k), shape=(TILE_SIZE_M, TILE_SIZE_K), padding_mode=zero_pad).astype(dtype)
# Load B tile: [TILE_SIZE_K, TILE_SIZE_N]
# Triton equivalent: b = tl.load(b_ptrs, mask=b_mask, other=0.0)
b = ct.load(B, index=(k, bidy), shape=(TILE_SIZE_K, TILE_SIZE_N), padding_mode=zero_pad).astype(dtype)
# Matrix multiply and accumulate
# Triton equivalent: acc += tl.dot(a, b)
accumulator = ct.mma(a, b, accumulator)
# Convert to output dtype
accumulator = ct.astype(accumulator, C.dtype)
# Store result
ct.store(C, index=(bidx, bidy), tile=accumulator)
def matmul(
a: torch.Tensor,
b: torch.Tensor,
TILE_SIZE_M: int = 128,
TILE_SIZE_N: int = 128,
TILE_SIZE_K: int = 32,
) -> torch.Tensor:
"""
Host wrapper for cuTile matrix multiplication.
Args:
a: Input tensor [M, K]
b: Input tensor [K, N]
TILE_SIZE_M: M-dimension tile size
TILE_SIZE_N: N-dimension tile size
TILE_SIZE_K: K-dimension tile size
Returns:
c: Output tensor [M, N]
"""
assert a.is_cuda and b.is_cuda
assert a.shape[1] == b.shape[0], f"Incompatible shapes: {a.shape} @ {b.shape}"
M, K = a.shape
K, N = b.shape
# Allocate output
c = torch.empty((M, N), device=a.device, dtype=a.dtype)
# Grid calculation
grid = (
math.ceil(M / TILE_SIZE_M) * math.ceil(N / TILE_SIZE_N),
1,
1,
)
ct.launch(
torch.cuda.current_stream(),
grid,
matmul_kernel,
(a, b, c, TILE_SIZE_M, TILE_SIZE_N, TILE_SIZE_K),
)
return c
def matmul_fp16(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
"""
FP16 matrix multiplication optimized for tensor cores.
Args:
a: Input tensor [M, K] in float16
b: Input tensor [K, N] in float16
Returns:
c: Output tensor [M, N] in float16
"""
assert a.dtype in [torch.float16, torch.bfloat16]
assert b.dtype in [torch.float16, torch.bfloat16]
return matmul(a, b)
def test_matmul():
"""Test function to verify correctness against PyTorch."""
torch.manual_seed(42)
# Test parameters
M, N, K = 512, 512, 512
# Test FP32
print("Testing FP32 matmul...")
a = torch.randn(M, K, device="cuda", dtype=torch.float32)
b = torch.randn(K, N, device="cuda", dtype=torch.float32)
# cuTile result
c_cutile = matmul(a, b)
# Reference (PyTorch)
c_ref = torch.matmul(a, b)
# Verify
# Note: TF32 mode may have slightly lower precision
fp32_passed = torch.allclose(c_cutile, c_ref, atol=1e-2, rtol=1e-2)
if fp32_passed:
print("FP32 test PASSED")
else:
diff = (c_cutile - c_ref).abs().max()
print(f"FP32 test FAILED - Max difference: {diff}")
# Test FP16
print("\nTesting FP16 matmul (tensor cores)...")
a_fp16 = torch.randn(M, K, device="cuda", dtype=torch.float16)
b_fp16 = torch.randn(K, N, device="cuda", dtype=torch.float16)
c_cutile_fp16 = matmul_fp16(a_fp16, b_fp16)
c_ref_fp16 = torch.matmul(a_fp16, b_fp16)
fp16_passed = torch.allclose(c_cutile_fp16, c_ref_fp16, atol=1e-1, rtol=1e-1)
if fp16_passed:
print("FP16 test PASSED")
else:
diff = (c_cutile_fp16 - c_ref_fp16).abs().max()
print(f"FP16 test FAILED - Max difference: {diff}")
# Test non-square matrices
print("\nTesting non-square matrices...")
M2, N2, K2 = 256, 1024, 512
a2 = torch.randn(M2, K2, device="cuda", dtype=torch.float32)
b2 = torch.randn(K2, N2, device="cuda", dtype=torch.float32)
c_cutile2 = matmul(a2, b2)
c_ref2 = torch.matmul(a2, b2)
nonsquare_passed = torch.allclose(c_cutile2, c_ref2, atol=1e-2, rtol=1e-2)
if nonsquare_passed:
print("Non-square test PASSED")
else:
diff = (c_cutile2 - c_ref2).abs().max()
print(f"Non-square test FAILED - Max difference: {diff}")
return fp32_passed and fp16_passed and nonsquare_passed
if __name__ == "__main__":
test_matmul()
examples/04_matmul/triton_kernel.py
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# SPDX-License-Identifier: CC-BY-4.0 AND Apache-2.0
#
"""
Matrix Multiplication (GEMM) - Triton Implementation
This file demonstrates the Triton equivalent of the CUDA tiled matmul kernel.
Key translation patterns:
- CUDA shared memory tiling → Triton block-level tiling with tl.dot
- Manual tile loading → tl.load with block pointers
- Nested loops for dot product → tl.dot (tensor core accelerated)
- Thread-level indexing → Program-level block indexing
Focuses on tiling pattern translation and autotune configuration.
"""
import torch
import triton
import triton.language as tl
@triton.jit
def matmul_kernel(
# Pointers to matrices
a_ptr,
b_ptr,
c_ptr,
# Matrix dimensions
M,
N,
K,
# Strides (elements to skip to get to next row/col)
stride_am,
stride_ak,
stride_bk,
stride_bn,
stride_cm,
stride_cn,
# Block sizes (compile-time constants)
BLOCK_SIZE_M: tl.constexpr,
BLOCK_SIZE_N: tl.constexpr,
BLOCK_SIZE_K: tl.constexpr,
):
"""
Triton kernel for matrix multiplication: C = A @ B
Translation from CUDA:
- blockIdx.x/y → tl.program_id(0/1)
- __shared__ float As/Bs → tl.load into registers (Triton manages caching)
- Nested k-loop with accumulation → tl.dot (uses tensor cores when available)
- __syncthreads() → Automatic (Triton handles synchronization)
Each program computes a BLOCK_SIZE_M x BLOCK_SIZE_N tile of C.
"""
# Program ID determines which output tile this program computes
# CUDA equivalent: blockIdx.x, blockIdx.y
pid_m = tl.program_id(axis=0) # Row tile index
pid_n = tl.program_id(axis=1) # Column tile index
# Calculate starting row/col for this program's output tile
# CUDA equivalent: row = blockIdx.y * TILE_SIZE, col = blockIdx.x * TILE_SIZE
offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
offs_k = tl.arange(0, BLOCK_SIZE_K)
# Pointers to first block of A and B
# A: [M, K] - we load BLOCK_SIZE_M x BLOCK_SIZE_K tiles
# B: [K, N] - we load BLOCK_SIZE_K x BLOCK_SIZE_N tiles
a_ptrs = a_ptr + (offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak)
b_ptrs = b_ptr + (offs_k[:, None] * stride_bk + offs_n[None, :] * stride_bn)
# Accumulator for the output tile
# CUDA equivalent: float acc = 0.0f; (but here it's a 2D tile)
acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
# Iterate over K dimension in blocks
# CUDA equivalent: for (int t = 0; t < num_tiles; t++)
for k in range(0, K, BLOCK_SIZE_K):
# Boundary masks
# CUDA equivalent: if (row < M && a_col < K)
a_mask = (offs_m[:, None] < M) & ((k + offs_k[None, :]) < K)
b_mask = ((k + offs_k[:, None]) < K) & (offs_n[None, :] < N)
# Load tiles of A and B
# CUDA equivalent: As[ty][tx] = A[row * K + a_col];
a = tl.load(a_ptrs, mask=a_mask, other=0.0)
b = tl.load(b_ptrs, mask=b_mask, other=0.0)
# Matrix multiply and accumulate
# CUDA equivalent: for (int k = 0; k < TILE_SIZE; k++) acc += As[ty][k] * Bs[k][tx];
# tl.dot uses tensor cores when:
# - dtype is float16/bfloat16
# - BLOCK_SIZE_K is multiple of 16
# - Shapes are compatible (M, N multiples of 16)
acc += tl.dot(a, b)
# Advance pointers to next K-tile
a_ptrs += BLOCK_SIZE_K * stride_ak
b_ptrs += BLOCK_SIZE_K * stride_bk
# Write output tile to C
# CUDA equivalent: if (row < M && col < N) C[row * N + col] = acc;
c_ptrs = c_ptr + (offs_m[:, None] * stride_cm + offs_n[None, :] * stride_cn)
c_mask = (offs_m[:, None] < M) & (offs_n[None, :] < N)
tl.store(c_ptrs, acc, mask=c_mask)
# Autotune configuration for optimal performance
# Triton will benchmark each configuration and select the best
@triton.autotune(
configs=[
# Small matrices - smaller tiles
triton.Config(
{"BLOCK_SIZE_M": 32, "BLOCK_SIZE_N": 32, "BLOCK_SIZE_K": 32},
num_stages=2,
num_warps=4,
),
triton.Config(
{"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 32, "BLOCK_SIZE_K": 32},
num_stages=2,
num_warps=4,
),
triton.Config(
{"BLOCK_SIZE_M": 32, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32},
num_stages=2,
num_warps=4,
),
# Medium matrices - balanced tiles
triton.Config(
{"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32},
num_stages=3,
num_warps=4,
),
triton.Config(
{"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32},
num_stages=3,
num_warps=4,
),
triton.Config(
{"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 128, "BLOCK_SIZE_K": 32},
num_stages=3,
num_warps=4,
),
# Large matrices - larger tiles for better data reuse
triton.Config(
{"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 128, "BLOCK_SIZE_K": 32},
num_stages=3,
num_warps=8,
),
triton.Config(
{"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 256, "BLOCK_SIZE_K": 32},
num_stages=3,
num_warps=8,
),
triton.Config(
{"BLOCK_SIZE_M": 256, "BLOCK_SIZE_N": 128, "BLOCK_SIZE_K": 32},
num_stages=3,
num_warps=8,
),
# Tensor core optimized (BLOCK_SIZE_K=16 for fp16)
triton.Config(
{"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 128, "BLOCK_SIZE_K": 16},
num_stages=4,
num_warps=8,
),
],
key=["M", "N", "K"], # Autotune based on matrix dimensions
)
@triton.jit
def matmul_kernel_autotuned(
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,
):
"""
Autotuned version of matmul kernel.
Autotune parameters:
- BLOCK_SIZE_M/N: Output tile dimensions (affects parallelism vs. data reuse)
- BLOCK_SIZE_K: K-dimension tile size (affects memory bandwidth)
- num_stages: Software pipelining depth (hides memory latency)
- num_warps: Number of warps per program (affects occupancy)
Tensor Core Requirements (for tl.dot acceleration):
- Input dtype: float16 or bfloat16
- BLOCK_SIZE_K: Multiple of 16
- BLOCK_SIZE_M, BLOCK_SIZE_N: Multiples of 16
- Accumulator: float32 (automatic)
"""
pid_m = tl.program_id(axis=0)
pid_n = tl.program_id(axis=1)
offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
offs_k = tl.arange(0, BLOCK_SIZE_K)
a_ptrs = a_ptr + (offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak)
b_ptrs = b_ptr + (offs_k[:, None] * stride_bk + offs_n[None, :] * stride_bn)
acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
for k in range(0, K, BLOCK_SIZE_K):
a_mask = (offs_m[:, None] < M) & ((k + offs_k[None, :]) < K)
b_mask = ((k + offs_k[:, None]) < K) & (offs_n[None, :] < N)
a = tl.load(a_ptrs, mask=a_mask, other=0.0)
b = tl.load(b_ptrs, mask=b_mask, other=0.0)
acc += tl.dot(a, b)
a_ptrs += BLOCK_SIZE_K * stride_ak
b_ptrs += BLOCK_SIZE_K * stride_bk
c_ptrs = c_ptr + (offs_m[:, None] * stride_cm + offs_n[None, :] * stride_cn)
c_mask = (offs_m[:, None] < M) & (offs_n[None, :] < N)
tl.store(c_ptrs, acc, mask=c_mask)
def matmul(a: torch.Tensor, b: torch.Tensor, use_autotune: bool = True) -> torch.Tensor:
"""
Host wrapper for Triton matrix multiplication.
Args:
a: Input tensor [M, K]
b: Input tensor [K, N]
use_autotune: Whether to use autotuned kernel
Returns:
c: Output tensor [M, N]
"""
assert a.is_cuda and b.is_cuda
assert a.shape[1] == b.shape[0], f"Incompatible shapes: {a.shape} @ {b.shape}"
M, K = a.shape
K, N = b.shape
# Allocate output
c = torch.empty((M, N), device=a.device, dtype=a.dtype)
# Grid: one program per output tile
# CUDA equivalent: dim3 grid((N + TILE_SIZE - 1) / TILE_SIZE, (M + TILE_SIZE - 1) / TILE_SIZE)
def grid(meta):
return (
triton.cdiv(M, meta["BLOCK_SIZE_M"]),
triton.cdiv(N, meta["BLOCK_SIZE_N"]),
)
if use_autotune:
matmul_kernel_autotuned[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),
)
else:
# Fixed configuration for debugging/testing
BLOCK_SIZE_M = 64
BLOCK_SIZE_N = 64
BLOCK_SIZE_K = 32
grid_fixed = (triton.cdiv(M, BLOCK_SIZE_M), triton.cdiv(N, BLOCK_SIZE_N))
matmul_kernel[grid_fixed](
a,
b,
c,
M,
N,
K,
a.stride(0),
a.stride(1),
b.stride(0),
b.stride(1),
c.stride(0),
c.stride(1),
BLOCK_SIZE_M=BLOCK_SIZE_M,
BLOCK_SIZE_N=BLOCK_SIZE_N,
BLOCK_SIZE_K=BLOCK_SIZE_K,
)
return c
def matmul_fp16(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
"""
FP16 matrix multiplication optimized for tensor cores.
Tensor Core Requirements:
1. Input dtype: float16 or bfloat16
2. Shapes: M, N, K should be multiples of 16 for best performance
3. BLOCK_SIZE_K: Multiple of 16 (handled by autotune configs)
The tl.dot operation automatically uses tensor cores when these
conditions are met, providing significant speedup over FP32.
"""
assert a.dtype in [torch.float16, torch.bfloat16]
assert b.dtype in [torch.float16, torch.bfloat16]
return matmul(a, b, use_autotune=True)
def test_matmul():
"""Test function to verify correctness against PyTorch."""
torch.manual_seed(42)
# Test parameters
M, N, K = 512, 512, 512
# Test FP32
print("Testing FP32 matmul...")
a = torch.randn(M, K, device="cuda", dtype=torch.float32)
b = torch.randn(K, N, device="cuda", dtype=torch.float32)
# Triton result
c_triton = matmul(a, b, use_autotune=False)
# Reference (PyTorch)
c_ref = torch.matmul(a, b)
# Verify
fp32_passed = torch.allclose(c_triton, c_ref, atol=1e-2, rtol=1e-2)
if fp32_passed:
print("FP32 test PASSED")
else:
diff = (c_triton - c_ref).abs().max()
print(f"FP32 test FAILED - Max difference: {diff}")
# Test FP16 (tensor cores)
print("\nTesting FP16 matmul (tensor cores)...")
a_fp16 = torch.randn(M, K, device="cuda", dtype=torch.float16)
b_fp16 = torch.randn(K, N, device="cuda", dtype=torch.float16)
c_triton_fp16 = matmul_fp16(a_fp16, b_fp16)
c_ref_fp16 = torch.matmul(a_fp16, b_fp16)
fp16_passed = torch.allclose(c_triton_fp16, c_ref_fp16, atol=1e-1, rtol=1e-1)
if fp16_passed:
print("FP16 test PASSED")
else:
diff = (c_triton_fp16 - c_ref_fp16).abs().max()
print(f"FP16 test FAILED - Max difference: {diff}")
# Test non-square matrices
print("\nTesting non-square matrices...")
M2, N2, K2 = 256, 1024, 512
a2 = torch.randn(M2, K2, device="cuda", dtype=torch.float32)
b2 = torch.randn(K2, N2, device="cuda", dtype=torch.float32)
c_triton2 = matmul(a2, b2, use_autotune=False)
c_ref2 = torch.matmul(a2, b2)
nonsquare_passed = torch.allclose(c_triton2, c_ref2, atol=1e-2, rtol=1e-2)
if nonsquare_passed:
print("Non-square test PASSED")
else:
diff = (c_triton2 - c_ref2).abs().max()
print(f"Non-square test FAILED - Max difference: {diff}")
return fp32_passed and fp16_passed and nonsquare_passed
def benchmark_matmul():
"""Benchmark Triton vs PyTorch matmul."""
import time
sizes = [(512, 512, 512), (1024, 1024, 1024), (2048, 2048, 2048)]
print("\nBenchmark Results:")
print("-" * 60)
print(f"{'Size':<20} {'PyTorch (ms)':<15} {'Triton (ms)':<15} {'Speedup':<10}")
print("-" * 60)
for M, N, K in sizes:
a = torch.randn(M, K, device="cuda", dtype=torch.float16)
b = torch.randn(K, N, device="cuda", dtype=torch.float16)
# Warmup
for _ in range(10):
_ = torch.matmul(a, b)
_ = matmul_fp16(a, b)
torch.cuda.synchronize()
# Benchmark PyTorch
start = time.perf_counter()
for _ in range(100):
_ = torch.matmul(a, b)
torch.cuda.synchronize()
pytorch_time = (time.perf_counter() - start) / 100 * 1000
# Benchmark Triton
start = time.perf_counter()
for _ in range(100):
_ = matmul_fp16(a, b)
torch.cuda.synchronize()
triton_time = (time.perf_counter() - start) / 100 * 1000
speedup = pytorch_time / triton_time
print(f"{M}x{N}x{K:<10} {pytorch_time:<15.3f} {triton_time:<15.3f} {speedup:<10.2f}x")
if __name__ == "__main__":
passed = test_matmul()
if passed:
print("\nAll tests passed!")
benchmark_matmul()
else:
print("\nSome tests failed!")
examples/05_attention/cutile_kernel.py
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# SPDX-License-Identifier: CC-BY-4.0 AND Apache-2.0
"""
Fused Multi-Head Attention (FMHA) - cuTile Implementation
This implementation follows the Flash Attention algorithm with online softmax.
Based on the official TileGym implementation.
Key patterns:
- ct.load with index/shape matching source tensor dimensions, then reshape
- ct.mma for tensor core accelerated matrix multiply
- Online softmax with exp2 optimization
- Grouped Query Attention (GQA) support
"""
import math
import cuda.tile as ct
import torch
INV_LOG_2 = 1.0 / math.log(2)
ConstInt = ct.Constant[int]
ConstBool = ct.Constant[bool]
@ct.kernel
def fmha_kernel(
Q,
K,
V,
Out,
qk_scale: float,
input_pos: int,
TILE_D: ConstInt,
H: ConstInt,
TILE_M: ConstInt,
TILE_N: ConstInt,
QUERY_GROUP_SIZE: ConstInt,
CAUSAL: ConstBool,
EVEN_K: ConstBool,
):
"""
cuTile kernel for Fused Multi-Head Attention.
Args:
Q: Query tensor [batch, num_heads, seq_len, head_dim]
K: Key tensor [batch, num_kv_heads, seq_len, head_dim]
V: Value tensor [batch, num_kv_heads, seq_len, head_dim]
Out: Output tensor [batch, num_heads, seq_len, head_dim]
qk_scale: Scale factor (typically 1/sqrt(head_dim))
input_pos: Starting position for causal masking
TILE_D: Head dimension
H: Number of heads
TILE_M: Query tile size
TILE_N: Key/Value tile size
QUERY_GROUP_SIZE: Number of query heads per KV head (for GQA)
CAUSAL: Whether to apply causal masking
EVEN_K: Whether K sequence length is divisible by TILE_N
"""
# Block indices
bid_x = ct.bid(0) # Query tile index
bid_y = ct.bid(1) # Batch * Head index
batch_idx = bid_y // H
head_idx = bid_y % H
off_kv_h = head_idx // QUERY_GROUP_SIZE # KV head index for GQA
# Adjust scale for exp2 optimization
qk_scale = qk_scale * INV_LOG_2
# Offsets for masking
offs_m = bid_x * TILE_M + ct.arange(TILE_M, dtype=ct.int32)
offs_m = offs_m + input_pos
offs_m = offs_m[:, None] # [TILE_M, 1]
offs_n_tile = ct.arange(TILE_N, dtype=ct.int32)
offs_n_tile = offs_n_tile[None, :] # [1, TILE_N]
# Initialize online softmax accumulators
m_i = ct.full((TILE_M, 1), -math.inf, dtype=ct.float32)
l_i = ct.full((TILE_M, 1), 0.0, dtype=ct.float32)
acc = ct.full((TILE_M, TILE_D), 0.0, dtype=ct.float32)
# Load Q tile: [TILE_M, TILE_D]
# Note: index and shape must match source tensor dimensions (4D)
q = ct.load(Q, index=(batch_idx, head_idx, bid_x, 0), shape=(1, 1, TILE_M, TILE_D)).reshape((TILE_M, TILE_D))
# Compute loop bounds
m_end = input_pos + (bid_x + 1) * TILE_M
k_seqlen = K.shape[2]
if CAUSAL:
mask_start = (input_pos + bid_x * TILE_M) // TILE_N
mask_start = min(mask_start, k_seqlen // TILE_N)
Tc = ct.cdiv(min(m_end, k_seqlen), TILE_N)
else:
Tc = ct.cdiv(k_seqlen, TILE_N)
mask_start = k_seqlen // TILE_N
# Main attention loop
for j in range(0, Tc):
# Load K tile (transposed): [TILE_D, TILE_N]
k = ct.load(
K,
index=(batch_idx, off_kv_h, 0, j),
shape=(1, 1, TILE_D, TILE_N),
order=(0, 1, 3, 2), # Transpose last two dims
latency=2,
).reshape((TILE_D, TILE_N))
# Compute QK: [TILE_M, TILE_N]
qk = ct.full((TILE_M, TILE_N), 0.0, dtype=ct.float32)
qk = ct.mma(q, k, acc=qk)
# Apply masking
if (CAUSAL or not EVEN_K) and j >= mask_start:
offs_n = j * TILE_N + offs_n_tile
mask = ct.full((TILE_M, TILE_N), True, dtype=ct.bool_)
if not EVEN_K:
mask = mask & (offs_n < k_seqlen)
if CAUSAL:
mask = mask & (offs_m >= offs_n)
mask = ct.where(mask, 0.0, -math.inf)
qk = qk + mask
# Online softmax update
m_ij = max(m_i, ct.max(qk, axis=-1, keepdims=True) * qk_scale)
qk = qk * qk_scale - m_ij
p = ct.exp2(qk, flush_to_zero=True)
l_ij = ct.sum(p, axis=-1, keepdims=True)
alpha = ct.exp2(m_i - m_ij, flush_to_zero=True)
l_i = l_i * alpha + l_ij
acc = acc * alpha
# Load V tile: [TILE_N, TILE_D]
v = ct.load(
V,
index=(batch_idx, off_kv_h, j, 0),
shape=(1, 1, TILE_N, TILE_D),
latency=4,
).reshape((TILE_N, TILE_D))
# Accumulate: acc += p @ v
p = p.astype(Q.dtype)
acc = ct.mma(p, v, acc=acc)
m_i = m_ij
# Normalize and store
acc = ct.truediv(acc, l_i, flush_to_zero=True)
acc = acc.reshape((1, 1, TILE_M, TILE_D)).astype(Out.dtype)
ct.store(Out, index=(batch_idx, head_idx, bid_x, 0), tile=acc)
def fmha_forward(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
sm_scale: float = None,
is_causal: bool = True,
TILE_M: int = 128,
TILE_N: int = 64,
) -> torch.Tensor:
"""
Host wrapper for FMHA forward pass.
Args:
q: Query tensor [batch, num_heads, seq_len, head_dim]
k: Key tensor [batch, num_kv_heads, seq_len, head_dim]
v: Value tensor [batch, num_kv_heads, seq_len, head_dim]
sm_scale: Softmax scale (default: 1/sqrt(head_dim))
is_causal: Whether to use causal masking
TILE_M: Query tile size
TILE_N: Key/Value tile size
Returns:
Output tensor [batch, num_heads, seq_len, head_dim]
"""
assert q.is_cuda and k.is_cuda and v.is_cuda
batch_size, num_heads, q_len, head_dim = q.shape
_, num_kv_heads, k_len, _ = k.shape
assert num_heads % num_kv_heads == 0
query_group_size = num_heads // num_kv_heads
if sm_scale is None:
sm_scale = 1.0 / math.sqrt(head_dim)
q = q.contiguous()
k = k.contiguous()
v = v.contiguous()
out = torch.empty_like(q)
input_pos = 0
EVEN_K = (k_len % TILE_N) == 0
grid = (
(q_len + TILE_M - 1) // TILE_M,
batch_size * num_heads,
1,
)
ct.launch(
torch.cuda.current_stream(),
grid,
fmha_kernel,
(
q,
k,
v,
out,
sm_scale,
input_pos,
head_dim,
num_heads,
TILE_M,
TILE_N,
query_group_size,
is_causal,
EVEN_K,
),
)
return out
def test_fmha():
"""Test FMHA against PyTorch reference."""
torch.manual_seed(42)
batch, heads, seq_len, head_dim = 2, 8, 128, 64
kv_heads = 2 # GQA: 4 query heads per KV head
q = torch.randn(batch, heads, seq_len, head_dim, device="cuda", dtype=torch.float16)
k = torch.randn(batch, kv_heads, seq_len, head_dim, device="cuda", dtype=torch.float16)
v = torch.randn(batch, kv_heads, seq_len, head_dim, device="cuda", dtype=torch.float16)
# Expand K, V for reference
k_expanded = k.repeat_interleave(heads // kv_heads, dim=1)
v_expanded = v.repeat_interleave(heads // kv_heads, dim=1)
sm_scale = 1.0 / math.sqrt(head_dim)
# cuTile result
out_cutile = fmha_forward(q, k, v, sm_scale, is_causal=True)
# PyTorch reference (causal)
scores = torch.matmul(q.float(), k_expanded.float().transpose(-2, -1)) * sm_scale
causal_mask = torch.triu(torch.ones(seq_len, seq_len, device="cuda"), diagonal=1).bool()
scores = scores.masked_fill(causal_mask, float("-inf"))
attn = torch.softmax(scores, dim=-1)
out_ref = torch.matmul(attn, v_expanded.float()).half()
passed = torch.allclose(out_cutile, out_ref, atol=1e-2, rtol=1e-2)
print(f"FMHA Test: {'PASSED' if passed else 'FAILED'}")
if not passed:
print(f" Max diff: {(out_cutile - out_ref).abs().max()}")
return passed
if __name__ == "__main__":
test_fmha()
examples/05_attention/triton_kernel.py
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# SPDX-License-Identifier: CC-BY-4.0 AND Apache-2.0
#
"""
Fused Multi-Head Attention - Triton Implementation (Flash Attention Style)
This file demonstrates the Triton equivalent of the CUDA attention kernel,
using the Flash Attention algorithm with online softmax for memory efficiency.
Key algorithmic differences from standard attention:
- Online softmax: Compute softmax incrementally without materializing full attention matrix
- Tiled computation: Process K/V in blocks, accumulating results
- Memory efficient: O(N) memory instead of O(N^2) for attention matrix
Translation patterns:
- CUDA global memory attention matrix → Triton online accumulation
- CUDA two-pass softmax → Triton single-pass online softmax
- CUDA explicit tiling → Triton block-based processing
Reference: "FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness"
"""
import torch
import triton
import triton.language as tl
@triton.jit
def flash_attention_forward_kernel(
Q_ptr, # Query: [B, H, N, d]
K_ptr, # Key: [B, H, N, d]
V_ptr, # Value: [B, H, N, d]
O_ptr, # Output: [B, H, N, d]
L_ptr, # Log-sum-exp for backward: [B, H, N]
stride_qb,
stride_qh,
stride_qn,
stride_qd, # Q strides
stride_kb,
stride_kh,
stride_kn,
stride_kd, # K strides
stride_vb,
stride_vh,
stride_vn,
stride_vd, # V strides
stride_ob,
stride_oh,
stride_on,
stride_od, # O strides
stride_lb,
stride_lh,
stride_ln, # L strides
seq_len,
head_dim,
scale, # 1/sqrt(head_dim)
BLOCK_M: tl.constexpr, # Block size for queries
BLOCK_N: tl.constexpr, # Block size for keys/values
BLOCK_D: tl.constexpr, # Block size for head dimension (must be >= head_dim)
):
"""
Flash Attention forward pass with online softmax.
Key insight: Instead of computing full attention matrix then softmax,
we compute softmax incrementally as we iterate over K/V blocks.
Online softmax algorithm:
1. For each K/V block, compute partial attention scores
2. Update running max and sum for numerical stability
3. Rescale previous accumulator and add new contribution
This avoids materializing the O(N^2) attention matrix.
Translation from CUDA:
- CUDA: Store full attn_scores[N, N] in global memory
- Triton: Keep running m (max), l (sum), acc (output) in registers
- CUDA: Two-pass softmax (compute max, then exp/sum)
- Triton: Single-pass online softmax with rescaling
"""
# Get program indices
batch_head_idx = tl.program_id(0)
query_block_idx = tl.program_id(1)
batch_idx = batch_head_idx // tl.num_programs(0) # Will be set by grid
head_idx = batch_head_idx % tl.num_programs(0)
# This is a simplified version - in practice we'd compute batch/head from program_id
# For this example, we assume batch_head_idx encodes both
# Calculate base pointers for this batch and head
Q_block_ptr = Q_ptr + batch_head_idx * stride_qh
K_block_ptr = K_ptr + batch_head_idx * stride_kh
V_block_ptr = V_ptr + batch_head_idx * stride_vh
O_block_ptr = O_ptr + batch_head_idx * stride_oh
L_block_ptr = L_ptr + batch_head_idx * stride_lh
# Query block start position
q_start = query_block_idx * BLOCK_M
# Offsets for this query block
q_offs = q_start + tl.arange(0, BLOCK_M)
d_offs = tl.arange(0, BLOCK_D)
# Mask for valid query positions
q_mask = q_offs < seq_len
d_mask = d_offs < head_dim
# Load Q block: [BLOCK_M, BLOCK_D]
# CUDA equivalent: Loading Q_row in the naive kernel
q_ptrs = Q_block_ptr + q_offs[:, None] * stride_qn + d_offs[None, :] * stride_qd
q = tl.load(q_ptrs, mask=q_mask[:, None] & d_mask[None, :], other=0.0)
# Initialize online softmax accumulators
# m: running max for numerical stability
# l: running sum of exp(scores - m)
# acc: running weighted sum of values
m = tl.full([BLOCK_M], float("-inf"), dtype=tl.float32)
l = tl.zeros([BLOCK_M], dtype=tl.float32)
acc = tl.zeros([BLOCK_M, BLOCK_D], dtype=tl.float32)
# Iterate over K/V blocks
# CUDA equivalent: The loop over key_pos in attention_forward_naive_cuda
# But here we process in blocks and use online softmax
for kv_start in range(0, seq_len, BLOCK_N):
kv_offs = kv_start + tl.arange(0, BLOCK_N)
kv_mask = kv_offs < seq_len
# Load K block: [BLOCK_N, BLOCK_D]
k_ptrs = K_block_ptr + kv_offs[:, None] * stride_kn + d_offs[None, :] * stride_kd
k = tl.load(k_ptrs, mask=kv_mask[:, None] & d_mask[None, :], other=0.0)
# Compute Q @ K^T for this block: [BLOCK_M, BLOCK_N]
# CUDA equivalent: The dot product loop in attention_forward_naive_cuda
# score += Q_row[d] * K_row[d];
scores = tl.dot(q, tl.trans(k)) * scale
# Apply causal mask if needed (optional, shown for completeness)
# scores = tl.where(q_offs[:, None] >= kv_offs[None, :], scores, float("-inf"))
# Mask out invalid positions
scores = tl.where(kv_mask[None, :], scores, float("-inf"))
# Online softmax update
# This is the key difference from CUDA's two-pass approach
# Step 1: Find new max for this block
# CUDA equivalent: max_score = fmaxf(max_score, score);
m_new = tl.maximum(m, tl.max(scores, axis=1))
# Step 2: Compute scaling factors
# When max changes, we need to rescale previous accumulator
alpha = tl.exp(m - m_new) # Scale for previous accumulator
# Step 3: Compute exp(scores - m_new) for current block
# CUDA equivalent: float exp_score = expf(score - s_max);
p = tl.exp(scores - m_new[:, None])
# Step 4: Update running sum
# CUDA equivalent: sum_exp += exp_score;
l_new = alpha * l + tl.sum(p, axis=1)
# Step 5: Load V block and accumulate weighted values
# CUDA equivalent: out_val += attn_row[key_pos] * V_base[key_pos * head_dim + d];
v_ptrs = V_block_ptr + kv_offs[:, None] * stride_vn + d_offs[None, :] * stride_vd
v = tl.load(v_ptrs, mask=kv_mask[:, None] & d_mask[None, :], other=0.0)
# Rescale previous accumulator and add new contribution
# This is the online softmax magic - we can update incrementally
acc = alpha[:, None] * acc + tl.dot(p.to(v.dtype), v)
# Update state for next iteration
m = m_new
l = l_new
# Final normalization: divide by sum of exponentials
# CUDA equivalent: attn_row[key_pos] /= s_sum;
acc = acc / l[:, None]
# Store output
o_ptrs = O_block_ptr + q_offs[:, None] * stride_on + d_offs[None, :] * stride_od
tl.store(o_ptrs, acc, mask=q_mask[:, None] & d_mask[None, :])
# Store log-sum-exp for backward pass
# L = m + log(l) is used in backward to avoid recomputing softmax
l_ptrs = L_block_ptr + q_offs * stride_ln
tl.store(l_ptrs, m + tl.log(l), mask=q_mask)
@triton.jit
def flash_attention_backward_kernel(
Q_ptr,
K_ptr,
V_ptr, # Inputs from forward
O_ptr,
L_ptr, # Outputs from forward (O and log-sum-exp)
dO_ptr, # Gradient of output
dQ_ptr,
dK_ptr,
dV_ptr, # Gradients to compute
stride_qb,
stride_qh,
stride_qn,
stride_qd,
stride_kb,
stride_kh,
stride_kn,
stride_kd,
stride_vb,
stride_vh,
stride_vn,
stride_vd,
stride_ob,
stride_oh,
stride_on,
stride_od,
stride_lb,
stride_lh,
stride_ln,
seq_len,
head_dim,
scale,
BLOCK_M: tl.constexpr,
BLOCK_N: tl.constexpr,
BLOCK_D: tl.constexpr,
):
"""
Flash Attention backward pass.
Key insight: Recompute attention weights on-the-fly instead of storing them.
This trades compute for memory, enabling training with longer sequences.
Gradients:
- dV = Attn^T @ dO (accumulated over query blocks)
- dQ = dAttn @ K (computed per query block)
- dK = dAttn^T @ Q (accumulated over query blocks)
where dAttn = softmax_backward(Attn, dO @ V^T)
Translation from CUDA:
- CUDA: Load stored attention weights from global memory
- Triton: Recompute attention weights using saved L (log-sum-exp)
"""
batch_head_idx = tl.program_id(0)
kv_block_idx = tl.program_id(1)
# Base pointers
Q_block_ptr = Q_ptr + batch_head_idx * stride_qh
K_block_ptr = K_ptr + batch_head_idx * stride_kh
V_block_ptr = V_ptr + batch_head_idx * stride_vh
O_block_ptr = O_ptr + batch_head_idx * stride_oh
L_block_ptr = L_ptr + batch_head_idx * stride_lh
dO_block_ptr = dO_ptr + batch_head_idx * stride_oh
dQ_block_ptr = dQ_ptr + batch_head_idx * stride_qh
dK_block_ptr = dK_ptr + batch_head_idx * stride_kh
dV_block_ptr = dV_ptr + batch_head_idx * stride_vh
# K/V block position
kv_start = kv_block_idx * BLOCK_N
kv_offs = kv_start + tl.arange(0, BLOCK_N)
kv_mask = kv_offs < seq_len
d_offs = tl.arange(0, BLOCK_D)
d_mask = d_offs < head_dim
# Load K and V for this block
k_ptrs = K_block_ptr + kv_offs[:, None] * stride_kn + d_offs[None, :] * stride_kd
v_ptrs = V_block_ptr + kv_offs[:, None] * stride_vn + d_offs[None, :] * stride_vd
k = tl.load(k_ptrs, mask=kv_mask[:, None] & d_mask[None, :], other=0.0)
v = tl.load(v_ptrs, mask=kv_mask[:, None] & d_mask[None, :], other=0.0)
# Initialize gradient accumulators for K and V
dk = tl.zeros([BLOCK_N, BLOCK_D], dtype=tl.float32)
dv = tl.zeros([BLOCK_N, BLOCK_D], dtype=tl.float32)
# Iterate over query blocks
for q_start in range(0, seq_len, BLOCK_M):
q_offs = q_start + tl.arange(0, BLOCK_M)
q_mask = q_offs < seq_len
# Load Q, O, dO, L for this query block
q_ptrs = Q_block_ptr + q_offs[:, None] * stride_qn + d_offs[None, :] * stride_qd
o_ptrs = O_block_ptr + q_offs[:, None] * stride_on + d_offs[None, :] * stride_od
do_ptrs = dO_block_ptr + q_offs[:, None] * stride_on + d_offs[None, :] * stride_od
l_ptrs = L_block_ptr + q_offs * stride_ln
q = tl.load(q_ptrs, mask=q_mask[:, None] & d_mask[None, :], other=0.0)
o = tl.load(o_ptrs, mask=q_mask[:, None] & d_mask[None, :], other=0.0)
do = tl.load(do_ptrs, mask=q_mask[:, None] & d_mask[None, :], other=0.0)
l = tl.load(l_ptrs, mask=q_mask, other=0.0)
# Recompute attention weights
# P = softmax(Q @ K^T * scale)
# Using saved L = m + log(sum(exp(scores - m)))
scores = tl.dot(q, tl.trans(k)) * scale
scores = tl.where(kv_mask[None, :], scores, float("-inf"))
p = tl.exp(scores - l[:, None]) # Attention weights
# Compute dV: dV += P^T @ dO
dv += tl.dot(tl.trans(p.to(do.dtype)), do)
# Compute dP: dP = dO @ V^T
dp = tl.dot(do, tl.trans(v))
# Softmax backward: dS = P * (dP - sum(P * dP))
# where S = scores before softmax
d_sum = tl.sum(p * dp, axis=1)
ds = p * (dp - d_sum[:, None]) * scale
# Compute dK: dK += dS^T @ Q
dk += tl.dot(tl.trans(ds.to(q.dtype)), q)
# Compute dQ: dQ = dS @ K (stored directly)
dq = tl.dot(ds.to(k.dtype), k)
dq_ptrs = dQ_block_ptr + q_offs[:, None] * stride_qn + d_offs[None, :] * stride_qd
# Note: This is a simplified version - full implementation would use atomics
# or separate kernel for dQ accumulation
tl.atomic_add(dq_ptrs, dq, mask=q_mask[:, None] & d_mask[None, :])
# Store dK and dV
dk_ptrs = dK_block_ptr + kv_offs[:, None] * stride_kn + d_offs[None, :] * stride_kd
dv_ptrs = dV_block_ptr + kv_offs[:, None] * stride_vn + d_offs[None, :] * stride_vd
tl.store(dk_ptrs, dk, mask=kv_mask[:, None] & d_mask[None, :])
tl.store(dv_ptrs, dv, mask=kv_mask[:, None] & d_mask[None, :])
def flash_attention_forward(
Q: torch.Tensor,
K: torch.Tensor,
V: torch.Tensor,
) -> tuple:
"""
Host wrapper for Flash Attention forward pass.
Args:
Q: Query tensor [batch_size, num_heads, seq_len, head_dim]
K: Key tensor [batch_size, num_heads, seq_len, head_dim]
V: Value tensor [batch_size, num_heads, seq_len, head_dim]
Returns:
O: Output tensor [batch_size, num_heads, seq_len, head_dim]
L: Log-sum-exp for backward [batch_size, num_heads, seq_len]
"""
assert Q.is_cuda and K.is_cuda and V.is_cuda
assert Q.shape == K.shape == V.shape
batch_size, num_heads, seq_len, head_dim = Q.shape
# Allocate output tensors
O = torch.empty_like(Q)
L = torch.empty(batch_size, num_heads, seq_len, device=Q.device, dtype=torch.float32)
# Block sizes
BLOCK_M = 64
BLOCK_N = 64
BLOCK_D = triton.next_power_of_2(head_dim)
# Scale factor
scale = 1.0 / (head_dim**0.5)
# Grid: one program per (batch, head) pair, tiled over query positions
num_query_blocks = triton.cdiv(seq_len, BLOCK_M)
grid = (batch_size * num_heads, num_query_blocks)
flash_attention_forward_kernel[grid](
Q,
K,
V,
O,
L,
Q.stride(0),
Q.stride(1),
Q.stride(2),
Q.stride(3),
K.stride(0),
K.stride(1),
K.stride(2),
K.stride(3),
V.stride(0),
V.stride(1),
V.stride(2),
V.stride(3),
O.stride(0),
O.stride(1),
O.stride(2),
O.stride(3),
L.stride(0),
L.stride(1),
L.stride(2),
seq_len,
head_dim,
scale,
BLOCK_M=BLOCK_M,
BLOCK_N=BLOCK_N,
BLOCK_D=BLOCK_D,
)
return O, L
def flash_attention_backward(
Q: torch.Tensor,
K: torch.Tensor,
V: torch.Tensor,
O: torch.Tensor,
L: torch.Tensor,
dO: torch.Tensor,
) -> tuple:
"""
Host wrapper for Flash Attention backward pass.
Args:
Q, K, V: Input tensors from forward
O: Output from forward
L: Log-sum-exp from forward
dO: Gradient of output
Returns:
dQ, dK, dV: Gradients of inputs
"""
batch_size, num_heads, seq_len, head_dim = Q.shape
# Allocate gradient tensors
dQ = torch.zeros_like(Q)
dK = torch.empty_like(K)
dV = torch.empty_like(V)
# Block sizes
BLOCK_M = 64
BLOCK_N = 64
BLOCK_D = triton.next_power_of_2(head_dim)
scale = 1.0 / (head_dim**0.5)
# Grid: one program per (batch, head) pair, tiled over K/V positions
num_kv_blocks = triton.cdiv(seq_len, BLOCK_N)
grid = (batch_size * num_heads, num_kv_blocks)
flash_attention_backward_kernel[grid](
Q,
K,
V,
O,
L,
dO,
dQ,
dK,
dV,
Q.stride(0),
Q.stride(1),
Q.stride(2),
Q.stride(3),
K.stride(0),
K.stride(1),
K.stride(2),
K.stride(3),
V.stride(0),
V.stride(1),
V.stride(2),
V.stride(3),
O.stride(0),
O.stride(1),
O.stride(2),
O.stride(3),
L.stride(0),
L.stride(1),
L.stride(2),
seq_len,
head_dim,
scale,
BLOCK_M=BLOCK_M,
BLOCK_N=BLOCK_N,
BLOCK_D=BLOCK_D,
)
return dQ, dK, dV
def test_flash_attention():
"""Test function to verify correctness against PyTorch."""
torch.manual_seed(42)
# Test parameters
BATCH_SIZE = 2
NUM_HEADS = 4
SEQ_LEN = 64
HEAD_DIM = 32
# Create test inputs
Q = torch.randn(BATCH_SIZE, NUM_HEADS, SEQ_LEN, HEAD_DIM, device="cuda", dtype=torch.float32)
K = torch.randn(BATCH_SIZE, NUM_HEADS, SEQ_LEN, HEAD_DIM, device="cuda", dtype=torch.float32)
V = torch.randn(BATCH_SIZE, NUM_HEADS, SEQ_LEN, HEAD_DIM, device="cuda", dtype=torch.float32)
# Run Flash Attention forward
O_flash, L = flash_attention_forward(Q, K, V)
# Reference (PyTorch scaled dot-product attention)
scale = 1.0 / (HEAD_DIM**0.5)
attn_scores = torch.matmul(Q, K.transpose(-2, -1)) * scale
attn_weights = torch.softmax(attn_scores, dim=-1)
O_ref = torch.matmul(attn_weights, V)
# Verify forward
forward_passed = torch.allclose(O_flash, O_ref, atol=1e-2, rtol=1e-2)
if forward_passed:
print("Forward test PASSED")
else:
diff = (O_flash - O_ref).abs().max()
print(f"Forward test FAILED - Max difference: {diff}")
# Test backward
dO = torch.randn_like(O_flash)
# Flash Attention backward
dQ_flash, dK_flash, dV_flash = flash_attention_backward(Q, K, V, O_flash, L, dO)
# PyTorch backward
Q_ref = Q.clone().requires_grad_(True)
K_ref = K.clone().requires_grad_(True)
V_ref = V.clone().requires_grad_(True)
attn_scores_ref = torch.matmul(Q_ref, K_ref.transpose(-2, -1)) * scale
attn_weights_ref = torch.softmax(attn_scores_ref, dim=-1)
O_ref = torch.matmul(attn_weights_ref, V_ref)
O_ref.backward(dO)
# Verify backward
dQ_passed = torch.allclose(dQ_flash, Q_ref.grad, atol=1e-2, rtol=1e-2)
dK_passed = torch.allclose(dK_flash, K_ref.grad, atol=1e-2, rtol=1e-2)
dV_passed = torch.allclose(dV_flash, V_ref.grad, atol=1e-2, rtol=1e-2)
if dQ_passed and dK_passed and dV_passed:
print("Backward test PASSED")
else:
print(f"Backward test: dQ={dQ_passed}, dK={dK_passed}, dV={dV_passed}")
if not dQ_passed:
print(f" dQ max diff: {(dQ_flash - Q_ref.grad).abs().max()}")
if not dK_passed:
print(f" dK max diff: {(dK_flash - K_ref.grad).abs().max()}")
if not dV_passed:
print(f" dV max diff: {(dV_flash - V_ref.grad).abs().max()}")
return forward_passed and dQ_passed and dK_passed and dV_passed
if __name__ == "__main__":
test_flash_attention()
references/api-mapping.md
# cuTile → Triton API Mapping
## Contents
[Import & Decorator](#import--decorator)
[Indexing](#indexing)
[Memory Operations](#memory-operations)
[Tensor Creation](#tensor-creation)
[Reductions](#reductions)
[Matrix Operations](#matrix-operations)
[Type Operations](#type-operations)
[Math Operations](#math-operations)
[Comparison & Logic](#comparison--logic)
[Bitwise Operations](#bitwise-operations)
[Debug Operations](#debug-operations)
[Atomic Operations](#atomic-operations)
[Synchronization](#synchronization)
[Host Functions](#host-functions)
[Data Types](#data-types)
[Launch Patterns](#launch-patterns)
[cuTile → Triton Gotchas](#cutile--triton-gotchas)
[Quick Reference Card](#quick-reference-card)
[TensorDescriptor Pattern (ct.load/ct.store → Triton TMA)](#tensordescriptor-pattern-ctloadctstore--triton-tma)
[Multi-dimensional Indexing](#multi-dimensional-indexing)
[Array.slice → Triton (Ragged Tensors)](#arrayslice--triton-ragged-tensors)
[ct.gather().item() → Triton (Runtime Index TMA)](#ctgatheritem--triton-runtime-index-tma)
This document provides **cuTile → Triton** mappings for converting `@ct.kernel` code to `@triton.jit`. Source column is cuTile; target column is Triton.
---
## Import & Decorator
| cuTile | Triton | Notes (c2t) |
|--------|--------|-------------|
| N/A | `import triton` | Add top-level triton import in Triton file |
| `import cuda.tile as ct` | `import triton.language as tl` | **Replace** ct with tl; remove cuda.tile |
| `@ct.kernel` | `@triton.jit` | Symmetric |
| `BLOCK: ct.Constant[int]` | `BLOCK: tl.constexpr` | Symmetric (ConstInt → constexpr) |
---
## Indexing
| cuTile | Triton | Notes (c2t) |
|--------|--------|-------------|
| `ct.bid(axis)` | `tl.program_id(axis)` | Symmetric |
| `ct.num_blocks(axis)` | `tl.num_programs(axis)` | Symmetric |
| `ct.arange(N, dtype=ct.int32)` | `tl.arange(0, N)` | **Triton has start param:** use `0, N`; drop `dtype=` (Triton infers) |
---
## Memory Operations
### ct.load / ct.store → Triton
**⚠️ TMA for 2D+ loads:** cuTile uses **TMA** internally for block-aligned 2D+ loads. Converting to **raw** `tl.load(ptr + offs, mask=m)` for 2D+ tile shapes causes **500%-2000% (5-20x) regression**. For any load with **2D+ block shape** (e.g. GEMM tiles, attention tiles), use **TMA**: `tl.make_tensor_descriptor(...).load([...])` — see [TensorDescriptor Pattern](#tensordescriptor-pattern-ctloadctstore--triton-tma). Use raw ptr+mask only for 1D or truly scattered access.
cuTile uses **block index** in `ct.load(arr, index=(...), shape=(...))`. Triton uses **element offset** (`ptr + offs`), **block ptr**, or **TMA tensor descriptor**. When converting:
| cuTile | Triton | Notes (c2t) |
|--------|--------|-------------|
| `ct.load(arr, index=(...), shape=(...))` **1D** | `tl.load(ptr + offs, mask=m)` or block_ptr | **Index is block index in cuTile;** compute element offset: `offs = bid * BLOCK + tl.arange(0, BLOCK)` |
| `ct.load(arr, index=(i,j,...), shape=(BM,BK,...))` **2D+** | **TMA:** `tl.make_tensor_descriptor(base, shape, strides, block_shape).load([...])` | **Do NOT use** `tl.load(ptr+offs, mask=m)` for 2D+ block loads — 5-20x regression |
| `ct.store(arr, index=(...), tile=val)` (1D) | `tl.store(ptr + offs, val, mask=m)` or block_ptr + store | Same: convert block index → ptr + offset |
| `ct.store(arr, index=(...), tile=val)` (2D+) | **TMA:** descriptor `.store([...], val)` | Use TMA for 2D+ block stores |
| `ct.load(arr, index=(...), shape=(...))` (block-aligned) | `tl.make_tensor_descriptor` + `.load([...])` or block_ptr | Prefer TMA for 2D+ (see TensorDescriptor section) |
| Loop variable in `index=` | `tl.advance(block_ptr, (delta_m, delta_n))` or TMA with offset args | Reintroduce advance when you have a loop over blocks |
### Gather/Scatter → Pointer load/store
| cuTile (Fallback) | Triton | When |
|-------------------|--------|------|
| `ct.gather(arr, indices, check_bounds=True, padding_value=v)` | `tl.load(ptr + offs, mask=m, other=v)` | Truly sparse random access |
| `ct.scatter(arr, indices, val, check_bounds=True)` | `tl.store(ptr + offs, val, mask=m)` | Truly sparse random access |
**Critical:** In Triton, `tl.load(ptr + offs, ...)` uses **element offset** `offs`, not block index. Build `offs` from `tl.program_id(axis)` and `tl.arange(0, BLOCK)` (and strides).
---
## Tensor Creation
| cuTile | Triton | Notes (c2t) |
|--------|--------|-------------|
| `ct.zeros(shape, dtype)` | `tl.zeros(shape, dtype)` | Symmetric |
| `ct.full(shape, val, dtype)` | `tl.full(shape, val, dtype)` | Symmetric |
| `ct.full(shape, 1, dtype)` | `tl.full(shape, 1.0, dtype)` | Triton can use 1.0; no `tl.ones()` either |
---
## Reductions
| cuTile | Triton | Notes (c2t) |
|--------|--------|-------------|
| `ct.sum(x, axis=0)` | `tl.sum(x, axis=0)` | Symmetric |
| `ct.max(x, axis=0)` | `tl.max(x, axis=0)` | Symmetric |
| `ct.min(x, axis=0)` | `tl.min(x, axis=0)` | Symmetric |
| `ct.argmax(x, axis=0)` | `tl.argmax(x, axis=0)` | Symmetric |
| `ct.argmin(x, axis=0)` | `tl.argmin(x, axis=0)` | Symmetric |
---
## Matrix Operations
| cuTile | Triton | Notes (c2t) |
|--------|--------|-------------|
| `ct.matmul(a, b)` | `tl.dot(a, b)` | Symmetric |
| `ct.mma(a, b, acc=acc)` | `tl.dot(a, b, acc)` | **Drop `acc=` keyword** in Triton (positional only) |
| Explicit tf32 guard + `ct.mma` | `tl.dot(a, b, allow_tf32=True)` (default) | Triton auto-casts fp32→tf32; you can omit guard or set `allow_tf32=False` for full fp32 |
### float32 → tf32 in Triton
In cuTile you may have:
```python
a_mma = ct.astype(a, ct.tfloat32) if a.dtype == ct.float32 else a
b_mma = ct.astype(b, ct.tfloat32) if b.dtype == ct.float32 else b
acc = ct.mma(a_mma, b_mma, acc=acc)
```
In Triton, default behavior already matches:
```python
# Triton: allow_tf32=True by default
acc = tl.dot(a, b, acc=acc)
```
Use `allow_tf32=False` only if you need strict IEEE float32.
---
## Type Operations
| cuTile | Triton | Notes (c2t) |
|--------|--------|-------------|
| `ct.astype(x, dtype)` | `x.to(dtype)` | **Use `.to(dtype)`** in Triton; no ct.astype |
| `ct.transpose(x)` | `x.T` or `tl.trans(x)` | Symmetric |
| `ct.reshape(x, shape)` | `tl.reshape(x, shape)` or `tl.view(x, shape)` | Symmetric |
| `ct.broadcast_to(x, shape)` | `tl.broadcast_to(x, shape)` | Symmetric |
| `ct.expand_dims(x, axis)` | `tl.expand_dims(x, axis)` | Symmetric |
---
## Math Operations
| cuTile | Triton | Notes (c2t) |
|--------|--------|-------------|
| `ct.exp(x)` | `tl.exp(x)` | Symmetric |
| `ct.exp2(x)` | `tl.exp2(x)` | Symmetric |
| `ct.log(x)` | `tl.log(x)` | Symmetric |
| `ct.log2(x)` | `tl.log2(x)` | Symmetric |
| `ct.sqrt(x)` | `tl.sqrt(x)` | Symmetric |
| `ct.rsqrt(x)` | `tl.rsqrt(x)` | Symmetric |
| `ct.sin(x)` / `ct.cos(x)` | `tl.sin(x)` / `tl.cos(x)` | Symmetric |
| `ct.abs(x)` | `tl.abs(x)` | Symmetric |
| `ct.maximum(a, b)` / `ct.minimum(a, b)` | `tl.maximum(a, b)` / `tl.minimum(a, b)` | Symmetric |
| `ct.sigmoid(x)` | `tl.sigmoid(x)` | Symmetric |
| `ct.softmax(x, axis)` | `tl.softmax(x, axis)` | Symmetric |
| `ct.floor(x)` / `ct.ceil(x)` | `tl.floor(x)` / `tl.ceil(x)` | Symmetric |
| `ct.fma(a, b, c)` | `tl.fma(a, b, c)` | Symmetric |
| `ct.clamp(x, min, max)` | `tl.clamp(x, min, max)` | Symmetric |
### Index Arithmetic
| cuTile | Triton | Notes (c2t) |
|--------|--------|-------------|
| `a + b` (Python) | `a + b` (Python) | Keep using Python `+`, `*`, `//` for index math |
| `ct.add(a, b)` / `ct.mul(a, b)` | N/A | **Do not** use tl.add/tl.mul for indices; use Python ops (ct promotes to float) |
---
## Comparison & Logic
| cuTile | Triton | Notes (c2t) |
|--------|--------|-------------|
| `ct.where(cond, x, y)` | `tl.where(cond, x, y)` | Symmetric |
| `x == y`, `x < y` (Python) | `x == y`, `x < y` (Python) | Symmetric |
---
## Bitwise Operations
| cuTile | Triton | Notes (c2t) |
|--------|--------|-------------|
| `x & y`, `x \| y`, `x ^ y`, `~x`, `x << n`, `x >> n` | Same (Python) | Symmetric |
| `ct.sum(x ^ y, axis)` (manual) | `tl.xor_sum(x, axis)` | Triton has built-in xor_sum |
---
## Debug Operations
| cuTile | Triton | Notes (c2t) |
|--------|--------|-------------|
| `ct.printf(fmt, *args)` | `tl.device_print(prefix, x)` | Triton uses prefix + value(s), not C-style format |
| `ct.assert_(cond)` | `tl.device_assert(cond, msg)` | Triton allows an optional message |
| N/A | `tl.static_print(...)` | Triton-only |
| N/A | `tl.static_assert(cond)` | Triton-only |
```python
# cuTile
ct.printf("value: %f\n", x)
ct.assert_(x > 0)
# Triton
tl.device_print("value", x)
tl.device_assert(x > 0, "x must be positive")
```
---
## Atomic Operations
| cuTile | Triton | Notes (c2t) |
|--------|--------|-------------|
| `ct.atomic_add(arr, indices, val)` | `tl.atomic_add(ptr, val, mask)` | Triton uses ptr + mask; build ptr from base + indices |
| `ct.atomic_max(arr, indices, val)` | `tl.atomic_max(ptr, val, mask)` | Same |
| `ct.atomic_min(arr, indices, val)` | `tl.atomic_min(ptr, val, mask)` | Same |
| `ct.atomic_cas(arr, indices, cmp, val)` | `tl.atomic_cas(ptr, cmp, val)` | Same |
| `ct.atomic_xchg(arr, indices, val)` | `tl.atomic_xchg(ptr, val)` | Same |
| `ct.atomic_and` / `ct.atomic_or` | `tl.atomic_and` / `tl.atomic_or` | Same pattern |
---
## Synchronization
| cuTile | Triton | Notes (c2t) |
|--------|--------|-------------|
| `ct.barrier()` | `tl.debug_barrier()` | Symmetric |
---
## Host Functions
| cuTile | Triton | Notes (c2t) |
|--------|--------|-------------|
| `(a + b - 1) // b` (host) | `triton.cdiv(a, b)` | **Prefer** `triton.cdiv` in Triton host code |
| `1 << (n-1).bit_length()` | `triton.next_power_of_2(n)` | Optional; Triton has built-in |
| `ct.launch(stream, grid, kernel, args)` | `kernel［grid］(kernel_args)` | **Replace** with bracket launch; no stream in call |
| Grid must be 3-tuple | Grid can be tuple or **lambda** | You can use `grid = lambda meta: (...)` in Triton |
| Dummy tensor + flag (no None) | `None` allowed in args | You can simplify to `None` in Triton if desired |
---
## Data Types
| cuTile | Triton | Notes (c2t) |
|--------|--------|-------------|
| `ct.float16` | `tl.float16` | Symmetric |
| `ct.float32` | `tl.float32` | Symmetric |
| `ct.float64` | `tl.float64` | Symmetric |
| `ct.bfloat16` | `tl.bfloat16` | Symmetric |
| `ct.int8` … `ct.int64` | `tl.int8` … `tl.int64` | Symmetric |
| `ct.uint8` … `ct.uint64` | `tl.uint8` … `tl.uint64` | Symmetric |
| `ct.bool_` | `tl.int1` | Symmetric |
---
## Launch Patterns
### cuTile (source) → Triton (target)
```python
# cuTile
grid = ((N + BLOCK - 1) // BLOCK, 1, 1)
ct.launch(torch.cuda.current_stream(), grid, kernel, (x, y, N, BLOCK))
# Triton
grid = (triton.cdiv(N, BLOCK),)
kernel［grid］(x_ptr, y_ptr, N, BLOCK=256)
```
```python
# cuTile 2D
grid = ((M + BLOCK_M - 1) // BLOCK_M, (N + BLOCK_N - 1) // BLOCK_N, 1)
ct.launch(stream, grid, kernel, (a, b, c, M, N, K, BLOCK_M, BLOCK_N))
# Triton 2D (tuple or lambda)
grid = (triton.cdiv(M, BLOCK_M), triton.cdiv(N, BLOCK_N))
kernel［grid］(a_ptr, b_ptr, c_ptr, M, N, K, BLOCK_M=64, BLOCK_N=64)
# Triton lambda (e.g. autotune)
grid = lambda meta: (triton.cdiv(M, meta['BLOCK_M']), triton.cdiv(N, meta['BLOCK_N']))
kernel［grid］(...)
```
---
## cuTile → Triton Gotchas
1. **Import:** `import cuda.tile as ct` → `import triton.language as tl` (and add `import triton` if needed).
2. **TMA for 2D+ loads (critical):** cuTile uses TMA for block-aligned 2D+ loads. In Triton you **must** use `tl.make_tensor_descriptor(...).load([...])` (see [TensorDescriptor Pattern](#tensordescriptor-pattern-ctloadctstore--triton-tma)). Do **not** convert 2D+ `ct.load(arr, index=, shape=)` to raw `tl.load(ptr + offs, mask=...)` — that causes **500%-2000% (5-20x) regression**.
3. **Loop index → advance:** Loop variable in `ct.load(index=)` → express as loop with `tl.advance(ptr, delta)` in Triton.
4. **Gather/scatter → pointer:** `ct.gather`/`ct.scatter` → `tl.load`/`tl.store` with `ptr + offs` and mask.
5. **Type cast:** `ct.astype(x, dtype)` → `x.to(dtype)`.
6. **Matrix multiply:** `ct.mma(a, b, acc=acc)` → `tl.dot(a, b, acc)` (no `acc=` keyword).
7. **Arange:** `ct.arange(N, dtype=ct.int32)` → `tl.arange(0, N)`.
8. **Grid:** Replace fixed 3-tuple and `ct.launch(stream, grid, kernel, args)` with `grid = (...)` or `lambda meta: (...)` and `kernel［grid］(kernel_args)`.
9. **None args:** Dummy tensor + flag in cuTile → you can use `None` in Triton kernel args if the kernel supports it.
10. **Host cdiv:** `(a + b - 1) // b` → can use `triton.cdiv(a, b)` in Triton host.
11. **Index math:** Keep using Python `+`, `*`, `//` for indices; do not use `tl.add`/`tl.mul` for index arithmetic.
12. **Kernel args:** Tensor args in cuTile → pass pointers (and shapes/strides if needed) in Triton; constexpr/Constant → `tl.constexpr`.
---
## Quick Reference Card (cuTile → Triton)
| Operation | cuTile | Triton |
|-----------|--------|--------|
| Import | `import cuda.tile as ct` | `import triton.language as tl` |
| Decorator | `@ct.kernel` | `@triton.jit` |
| Constexpr | `BLOCK: ct.Constant[int]` | `BLOCK: tl.constexpr` |
| Block ID | `ct.bid(0)` | `tl.program_id(0)` |
| Range | `ct.arange(N, dtype=ct.int32)` | `tl.arange(0, N)` |
| TMA Load (2D+) | `ct.load(arr, index=(...), shape=(...))` (2D+ block) | **Must use** `tl.make_tensor_descriptor(...).load([...])` — raw tl.load = 5-20x regression |
| TMA Store (2D+) | `ct.store(arr, index=(...), tile=v)` (2D+ block) | **Must use** descriptor `.store([...], v)` |
| Ptr Load | `ct.gather(arr, idx, check_bounds=True)` | `tl.load(ptr+offs, mask=m)` |
| Ptr Store | `ct.scatter(arr, idx, v, check_bounds=True)` | `tl.store(ptr+offs, v, mask=m)` |
| Cast | `ct.astype(x, dtype)` | `x.to(dtype)` |
| Matmul | `ct.mma(a, b, acc=acc)` | `tl.dot(a, b, acc)` |
| Launch | `ct.launch(stream, grid, kernel, args)` | `kernel［grid］(kernel_args)` |
| Cdiv (host) | `(a + b - 1) // b` | `triton.cdiv(a, b)` |
---
## TensorDescriptor Pattern (ct.load/ct.store → Triton TMA)
**Mandatory for 2D+ tile loads.** When converting cuTile block-aligned loads (`ct.load`/`ct.store` with 2D+ shape) to Triton on Hopper (SM 90+) / Blackwell, use **TensorDescriptor** (`tl.make_tensor_descriptor` + `.load`/`.store`) so Triton can use TMA. Falling back to plain `tl.load(ptr + offs, mask=...)` for 2D+ block access causes **5-20x (500%-2000%) regression**.
```python
from triton.tools.tensor_descriptor import TensorDescriptor
def supports_host_descriptor():
return torch.cuda.get_device_capability()[0] >= 9
@triton.jit
def _maybe_make_tensor_desc(desc_or_ptr, shape, strides, block_shape):
if isinstance(desc_or_ptr, tl.tensor_descriptor):
return desc_or_ptr
else:
return tl.make_tensor_descriptor(desc_or_ptr, shape, strides, block_shape)
@triton.jit
def kernel(X, X_shape_0, X_shape_1, X_stride_0, X_stride_1, BLOCK: tl.constexpr):
pid = tl.program_id(0)
X_desc = _maybe_make_tensor_desc(X, shape=[X_shape_0, X_shape_1],
strides=[X_stride_0, X_stride_1],
block_shape=[1, BLOCK])
tile = X_desc.load([pid, 0])
# ... computation ...
X_desc.store([pid, 0], tile)
# Host
def wrapper(x):
BLOCK = 128
grid = (x.shape[0], 1, 1)
if supports_host_descriptor():
desc_x = TensorDescriptor(x, shape=x.shape, strides=x.stride(), block_shape=[1, BLOCK])
else:
desc_x = x
kernel［grid］(desc_x, x.shape[0], x.shape[1], x.stride(0), x.stride(1), BLOCK)
```
- Shape/stride passed as **individual scalars** to the kernel.
- Use `_maybe_make_tensor_desc` for fallback when TensorDescriptor is not available.
---
## Multi-dimensional Indexing
| cuTile | Triton | Notes (c2t) |
|--------|--------|-------------|
| `ct.gather(arr, indices)` or `ct.gather(arr, (idx0, idx1))` | `tl.load(ptr + offs, mask=m)` | Build offset from indices and strides; handle OOB with mask |
| Multi-dim indexing (OOB auto) | Manual mask when `BLOCK > actual_dim` | Triton: `mask = tl.arange(0, BLOCK) < actual_dim` |
---
## Array.slice → Triton (Ragged Tensors)
| cuTile | Triton |
|--------|--------|
| `start = ct.load(indptr, idx, shape=())` | `start = tl.load(indptr_ptr + idx)` |
| `A.slice(axis=0, start=start, stop=end)` | `ptr + start * stride` and manual extent |
| `ct.load(sliced, (tile_idx,), shape=(...))` | `tl.load(ptr + start*stride + offs)` or block_ptr over segment |
| `ct.num_tiles(sliced, axis, shape)` | `tl.cdiv(end - start, BLOCK)` |
For Array.slice ragged tensor patterns, apply the mapping above: compute `ptr + start * stride + offs` manually and use `tl.cdiv` for tile counts.
---
## ct.gather().item() → Triton (Runtime Index TMA)
| cuTile | Triton |
|--------|--------|
| `page_id = ct.gather(block_tables, (idx,), padding_value=0).item()` | `page_id = tl.load(block_tables + idx)` |
| `ct.load(k_cache, index=(page_id, ...), allow_tma=True)` | `tl.load(k_cache + page_id * stride + offs)` or block_ptr |
For paged attention TMA: load the page index with `tl.load(block_tables + idx)`, then use it as an offset into the cache tensor via `tl.make_tensor_descriptor` or manual pointer arithmetic.
references/debugging.md
# Triton Runtime Error Debugging (cuTile → Triton)
This guide covers runtime errors that commonly appear after converting a cuTile kernel to Triton.
---
## `cudaErrorIllegalAddress` (Illegal Memory Access)
**Symptom:**
```
torch.AcceleratorError: CUDA error: an illegal memory access was encountered
Search for 'cudaErrorIllegalAddress' ...
```
This is the **most frequent runtime crash** when converting cuTile kernels that use pointer
indirection (grouped/batched ops: group GEMM, batched attention, MoE, etc.).
### Root Cause 1: Hardcoded pointer type mismatch (PRIMARY)
**What happens:** When loading tensor pointers from a pointer table inside the kernel, the
element type must exactly match the actual tensor dtype. Triton's pointer arithmetic advances
the address by `offset * element_size_in_bytes`, so a wrong type causes every load/store to
hit an unintended address.
| Tensor dtype | Element size | `tl.float16` pointer arithmetic | Result |
|---|---|---|---|
| `torch.float16` | 2 bytes | 2 bytes/element | Correct |
| `torch.bfloat16` | 2 bytes | 2 bytes/element | Correct (same size) |
| `torch.float32` | 4 bytes | 2 bytes/element | **Off by 2×** → crash |
**Where to look:** Any `tl.load(ptr_table + idx).to(tl.pointer_type(...))` line with a
hardcoded type:
```python
# WRONG — crashes for bfloat16/float32 inputs
a_ptr = tl.load(a_ptrs + group_id).to(tl.pointer_type(tl.float16))
b_ptr = tl.load(b_ptrs + group_id).to(tl.pointer_type(tl.float16))
c_ptr = tl.load(c_ptrs + group_id).to(tl.pointer_type(tl.float16))
```
**Fix:** Pass the dtype as a `tl.constexpr` and use it for the pointer type:
```python
# Kernel signature — add DTYPE constexpr
@triton.jit
def my_kernel(..., DTYPE: tl.constexpr):
...
a_ptr = tl.load(a_ptrs + group_id).to(tl.pointer_type(DTYPE))
b_ptr = tl.load(b_ptrs + group_id).to(tl.pointer_type(DTYPE))
c_ptr = tl.load(c_ptrs + group_id).to(tl.pointer_type(DTYPE))
...
# Also fix the store cast — do NOT hardcode output dtype
tl.store(c_ptr + c_offs, acc.to(DTYPE), mask=c_mask)
# Host wrapper — build dtype map and pass it
_DTYPE_MAP = {
torch.float16: tl.float16,
torch.bfloat16: tl.bfloat16,
torch.float32: tl.float32,
}
triton_dtype = _DTYPE_MAP.get(dtype)
if triton_dtype is None:
raise ValueError(f"Unsupported dtype: {dtype}")
my_kernel［grid］(..., DTYPE=triton_dtype)
```
**Checklist — scan every pointer table load/store in the converted kernel:**
```bash
grep -n "pointer_type" <your_triton_kernel.py>
```
Every occurrence should use `DTYPE` (or equivalent constexpr), never a hardcoded type.
Also scan the store path for hardcoded cast:
```bash
grep -n "acc.to(tl\." <your_triton_kernel.py>
```
---
### Root Cause 2: int32 stride overflow
**What happens:** Strides stored in a `torch.int32` tensor overflow when
`max_row_index × stride > 2^31 − 1`. The overflowed value wraps to a negative or
small positive number, pointing to an entirely different memory region.
Threshold: overflow occurs when `(TILE_M - 1 + (num_m_tiles - 1) * TILE_M) * stride > 2^31`
i.e. roughly when `M × K > 2^31` elements (≈ 4096 × 512K, or 512K × 4096 rows).
**Where to look:**
```python
# WRONG — int32 overflows for large matrices
a_strides = torch.tensor(a_stride_list, dtype=torch.int32, device=device)
b_strides = torch.tensor(b_stride_list, dtype=torch.int32, device=device)
c_strides = torch.tensor(c_stride_list, dtype=torch.int32, device=device)
```
**Fix:** Use `int64`:
```python
a_strides = torch.tensor(a_stride_list, dtype=torch.int64, device=device)
b_strides = torch.tensor(b_stride_list, dtype=torch.int64, device=device)
c_strides = torch.tensor(c_stride_list, dtype=torch.int64, device=device)
```
This applies to any array of strides passed to the kernel for pointer arithmetic, whether
in a pointer table pattern or as direct scalar stride arguments for large tensors.
---
### Quick diagnosis checklist
Run through these in order when you see `cudaErrorIllegalAddress`:
```
[ ] 1. Search for hardcoded pointer types:
grep -n "pointer_type(tl\." <your_triton_kernel.py>
→ Should show DTYPE (constexpr), not tl.float16/tl.bfloat16/tl.float32
[ ] 2. Check store casts:
grep -n "\.to(tl\." <your_triton_kernel.py>
→ Accumulator cast before tl.store should use DTYPE, not hardcoded type
[ ] 3. Check stride tensor dtypes in host:
grep -n "dtype=torch.int32" <your_triton_kernel.py>
→ Strides used in pointer arithmetic should be int64
[ ] 4. Check pointer table dtype (usually already int64 — verify):
grep -n "a_ptrs\|b_ptrs\|c_ptrs" <your_triton_kernel.py>
→ Should be dtype=torch.int64
[ ] 5. Verify DTYPE constexpr flows correctly:
- Defined as DTYPE: tl.constexpr in kernel signature
- Passed from host as DTYPE=triton_dtype (a tl.* type, not torch.* type)
- _DTYPE_MAP covers all dtypes used in tests
```
---
### Pattern: pointer table kernels (group GEMM, MoE, batched ops)
The pointer table pattern (passing `int64` pointer arrays to the kernel and loading per-group
pointers inside) is the primary source of this error class. cuTile handles this automatically
through its typed tensor API; Triton requires explicit pointer casts.
**cuTile (source):**
```python
@ct.kernel
def group_gemm_kernel(As, Bs, Cs, TILE_M: ConstInt, ...):
Ai = As[g] # cuTile knows the type from the tensor descriptor
ta = ct.load(Ai, (tile_m_idx, kk), shape=(TILE_M, TILE_K), ...)
```
**Triton (target — correct pattern):**
```python
@triton.jit
def group_gemm_kernel(a_ptrs, ..., DTYPE: tl.constexpr):
a_ptr = tl.load(a_ptrs + group_id).to(tl.pointer_type(DTYPE)) # ← use DTYPE
a_tile = tl.load(a_ptr + a_offs, mask=a_mask, other=0.0)
...
tl.store(c_ptr + c_offs, acc.to(DTYPE), mask=c_mask) # ← use DTYPE
```
**Host (correct pattern):**
```python
_DTYPE_MAP = {
torch.float16: tl.float16,
torch.bfloat16: tl.bfloat16,
torch.float32: tl.float32,
}
triton_dtype = _DTYPE_MAP[dtype]
my_kernel［grid］(..., DTYPE=triton_dtype)
```
---
### Root Cause 3: Incomplete dtype map (`ValueError: Unsupported dtype`)
**What happens:** The host-side `_DTYPE_MAP` only covers the dtypes the author tested. When the
caller uses a dtype not in the map (e.g., `torch.float8_e5m2`), a `ValueError` is raised before
the kernel even launches.
```
ValueError: Unsupported dtype for group_gemm triton backend: torch.float8_e5m2
```
**Fix:** Extend `_DTYPE_MAP` to cover all float8 variants. Because float8 types were added in
specific PyTorch and Triton releases, use `hasattr` guards so the code still works on older
installs where those types don't exist yet:
```python
_DTYPE_MAP = {
torch.float16: tl.float16,
torch.bfloat16: tl.bfloat16,
torch.float32: tl.float32,
}
# float8 types: add only when both torch and tl have them
_FLOAT8_PAIRS = [
("float8_e5m2", "float8e5"),
("float8_e4m3fn", "float8e4nv"),
("float8_e5m2fnuz","float8e5b16"),
("float8_e4m3fnuz","float8e4b8"),
]
for torch_name, tl_name in _FLOAT8_PAIRS:
if hasattr(torch, torch_name) and hasattr(tl, tl_name):
_DTYPE_MAP[getattr(torch, torch_name)] = getattr(tl, tl_name)
```
**PyTorch → Triton float8 type mapping:**
| `torch` dtype | `tl` dtype | Notes |
|---|---|---|
| `torch.float8_e5m2` | `tl.float8e5` | E5M2, 1 byte/element |
| `torch.float8_e4m3fn` | `tl.float8e4nv` | E4M3 NVIDIA format |
| `torch.float8_e5m2fnuz` | `tl.float8e5b16` | E5M2 UZ (unsigned zero) |
| `torch.float8_e4m3fnuz` | `tl.float8e4b8` | E4M3 UZ |
**Note:** float8 inputs use 1 byte/element. Ensure the `DTYPE` constexpr reflects this when
setting up pointer arithmetic. The accumulator remains `tl.float32`; only the final store cast
uses the float8 type.
Add to the diagnosis checklist:
```
[ ] 6. Does _DTYPE_MAP cover all dtypes in the test suite?
grep -n "dtype=torch\." tests/ops/test_your_op.py
→ Every dtype listed must have an entry in _DTYPE_MAP (or a hasattr guard)
```
---
## Other Common Triton Runtime Errors
### `tl.dot` shape error (`expected block of shape [M,K,N]`)
**Cause:** `tl.dot` requires both inputs to have power-of-2 dimensions and compatible shapes.
TILE_M, TILE_N, TILE_K must each be powers of 2 ≥ 16 (or ≥ 32 for float32 on some GPUs).
**Fix:** Ensure tile sizes are powers of 2, and add `TILE_M >= 16` / `TILE_K >= 16` guards.
### `tl.load` with non-scalar pointer from pointer table
**Symptom:** JIT compilation error mentioning "expected scalar pointer."
**Cause:** `tl.load(a_ptrs + group_id)` where `group_id` is not a scalar (e.g., a vector due
to loop unrolling). Keep `group_id` as a scalar loop variable; do not vectorize the group loop.
### NaN/Inf after conversion (not a crash but related)
See [SKILL.md](../SKILL.md) and [translations/workflow.md](../translations/workflow.md) for testing and numerical comparison workflows.
Common cause: accumulator cast mismatch (e.g., storing fp32 acc as fp32 when original stored
as fp16 — use the same output dtype as the cuTile kernel).
---
## Triton Math Function Dtype Requirements (CRITICAL) {#triton-math-function-dtype-requirements-critical}
Several Triton math functions have **strict dtype requirements** that differ from cuTile:
| Function | Required dtype | Error if wrong | Solution |
|----------|---------------|----------------|----------|
| `tl.math.erf(x)` | fp32, fp64 only | `ValueError: Expected dtype ['fp32', 'fp64'] but got fp16` | Let Triton auto-promote OR explicit `.to(tl.float32)` |
| `tl.math.erfc(x)` | fp32, fp64 only | Same as above | Same as above |
| `tl.exp(x)` | All (but fp16 loses precision) | Silent precision loss, potential NaN | Cast: `tl.exp(x.to(tl.float32))` |
| `tl.log(x)` | All (but fp16 loses precision) | Silent precision loss | Cast: `tl.log(x.to(tl.float32))` |
| `tl.sqrt(x)` | All (but fp16 loses precision) | Silent precision loss | Cast if precision needed |
### Common Mistake: Wrong Mathematical Substitution
**NEVER** replace `tl.math.erf` with a tanh-based approximation to "fix" the dtype error.
```python
# WRONG - mathematically incorrect substitution
def standard_normal_cdf(x):
# This is the GELU tanh approximation formula, NOT an erf approximation!
erf_approx = tanh(sqrt_2_div_pi * (x + 0.044715 * x * x * x)) # WRONG
return 0.5 * (1 + erf_approx)
# CORRECT - use the actual erf function
def standard_normal_cdf(x):
# 1.0 / math.sqrt(2.0) ≈ 0.70710678
inverse_sqrt_2 = 0.70710678
cdf = 0.5 * (1 + tl.math.erf(x * inverse_sqrt_2)) # CORRECT
return cdf
```
The formula `tanh(√(2/π) * (x + 0.044715x³))` is specifically the **GELU tanh approximation**, not an approximation of the error function. These are mathematically different:
- **Exact GELU**: `x * Φ(x)` where `Φ(x) = 0.5 * (1 + erf(x/√2))`
- **Tanh GELU**: `0.5 * x * (1 + tanh(√(2/π) * (x + 0.044715x³)))`
### Recommended Pattern for fp16/bf16 Kernels
```python
@triton.jit
def kernel_with_erf(x_ptr, y_ptr, n, BLOCK: tl.constexpr):
offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
x = tl.load(x_ptr + offs, mask=offs < n)
# For erf: Triton auto-promotes fp16→fp32, result stays fp32
# Output will be written as fp32 unless you cast back
# 1.0 / math.sqrt(2.0) ≈ 0.70710678
cdf = 0.5 * (1 + tl.math.erf(x * 0.70710678))
# For exp with fp16 input: explicit cast recommended for precision
# 1.0 / math.sqrt(2.0 * math.pi) ≈ 0.39894228
pdf = 0.39894228 * tl.exp((-0.5 * x * x).to(tl.float32))
tl.store(y_ptr + offs, x * cdf, mask=offs < n)
```
references/gotchas.md
# Gotchas — Most Common cuTile → Triton Translation Errors
Comprehensive table of patterns that frequently break or regress when porting
`@ct.kernel` to `@triton.jit`. Read this BEFORE writing the Triton kernel — most
entries describe a wrong-by-default first attempt.
| Pattern | cuTile | Triton | Common Mistake |
|---------|--------|--------|----------------|
| **mma accumulator** | `ct.mma(a, b, acc=acc)` | `tl.dot(a, b, acc)` | Using keyword `acc=` in Triton (positional only) |
| **mma float32→tf32** | Explicit `ct.astype(..., ct.tfloat32)` guard before ct.mma | `tl.dot(a, b, allow_tf32=True)` (default) | Over-specifying; Triton auto-casts by default |
| **Type cast** | `ct.astype(x, dtype)` | `x.to(dtype)` | Using ct.astype in Triton |
| **Grid** | `(n, 1, 1)` tuple, `ct.launch(stream, grid, kernel, args)` | `lambda meta: (n,)` or tuple, bracket launch | Using ct.launch or 3-tuple in Triton |
| **Host cdiv** | `(a + b - 1) // b` (Python) | `triton.cdiv(a, b)` | Forgetting triton.cdiv in host |
| **2D+ tile load** | `ct.load(arr, index=(i,j), shape=(BM,BK))` (cuTile uses TMA) | `tl.make_tensor_descriptor(...).load([...])` | Using raw `tl.load(ptr+offs, mask=m)` → **5-20x regression**; always use TMA for 2D+ block loads |
| **Index type** | Block index in ct.load/ct.store | Element offset (ptr + offs) or TMA descriptor | Using block index as tl.load offset |
| **arange** | `ct.arange(N, dtype=ct.int32)` | `tl.arange(0, N)` | Triton has start param (0, N) |
| **None args** | Dummy tensor + flag | Allowed in kernel | Carrying over dummy+flag when not needed |
| **String const** | `ct.Constant[int]` only (no str) | `tl.constexpr` (any type) | Keeping int enum; Triton can use str constexpr if needed |
| **Shape args** | Static/constexpr in ct.full/ct.zeros | Dynamic shapes OK in Triton | Over-constraining shapes |
| **Launch** | `ct.launch(stream, grid, kernel, args)` | bracket launch (grid then args) | Leaving ct.launch in Triton host |
| **Branch vars** | Pre-define before if | Can define in branch | Over-defining before branch in Triton |
| **Pointer table type** | Typed tensor descriptor (auto) | `tl.load(ptrs+idx).to(tl.pointer_type(DTYPE))` where `DTYPE: tl.constexpr` | **Hardcoding `tl.float16`** → `cudaErrorIllegalAddress` for bfloat16/float32 inputs |
| **Stride dtype** | cuTile uses tensor shape (auto) | Pass strides as `torch.int64`, not `int32` | `int32` overflows → illegal address for large matrices (M×K > 2^31) |
| **dtype map coverage** | cuTile typed tensors (auto) | `_DTYPE_MAP` must cover all dtypes (incl. float8); use `hasattr` guards | Missing entry → `ValueError: Unsupported dtype` before kernel launch |
| **tl.math.erf dtype** | cuTile erf handles all dtypes | `tl.math.erf` **only accepts fp32/fp64** | `ValueError: Expected dtype ['fp32', 'fp64'] but got fp16` — do NOT replace with tanh approximation (mathematically wrong); let Triton auto-promote or cast input |
| **tl.exp with fp16** | cuTile exp handles all dtypes | Cast to fp32 before `tl.exp` for precision: `tl.exp(x.to(tl.float32))` | Precision loss or NaN with fp16 inputs in exp/log/sqrt |
| **Math func approx** | N/A | Never substitute `tl.math.erf` with tanh-based approximation | Using GELU tanh formula (`0.044715*x³`) as erf approximation is **mathematically incorrect** — they are different functions |
| **Layout flag (`transpose`)** | cuTile may use one path per layout | Need **two Triton kernels** when math differs (e.g. MLA: `qk` `[H,N]` vs `[N,H]`, different `V` TMA) | Reusing transpose-only logic for `transpose=False` + fixed blocks → **3–15×** on that mode; see [../translations/advanced-patterns.md](../translations/advanced-patterns.md) |
| **Batched matmul** | `ct.matmul(W, X)` broadcasts implicitly at tile level | `tl.dot(W, X)` only supports 2D operands | Using `broadcast_to + tl.dot` → **10-50× slower**, no tensor cores (see FFT anti-pattern in [performance-gotchas.md](./performance-gotchas.md)) |
| **Batch-per-block** | cuTile processes 1 batch per block naturally | Triton temptation: process BS batches per block | Creates BS× register pressure, breaks tensor core compatibility |
references/harness-integration.md
# Testing & Validation (cuTile → Triton)
How to test and benchmark kernels after converting from cuTile to Triton,
using the standard TileGym pytest infrastructure.
---
## Table of Contents
- [Test Harness](#test-harness)
- [Benchmark Harness](#benchmark-harness)
- [Conversion Skill Workflow](#conversion-skill-workflow)
---
## Test Harness
`tests/common.py::PyTestCase` provides the correctness comparison infrastructure:
```python
# Primary kernel comparison method
self.assertCorrectness(
test_fn, # Kernel under test (e.g., Triton impl)
ref_fn, # Reference (e.g., PyTorch or cuTile)
kwargs, # Input tensors
rtol=1e-3, # Relative tolerance
atol=1e-5, # Absolute tolerance
gradient=True, # Also check backward pass
)
```
**Key methods**:
- `assertCorrectness()` — Compare test vs reference with tolerances
- `assertDeterministic()` — Verify consistent results across iterations
- `compare_tensors()` — Low-level comparison with detailed mismatch reporting
- `benchmark()` — Performance measurement with CUDA events/CUPTI
**Tolerance defaults by dtype** (from `get_dtype_tolerances()`):
| dtype | rtol | atol |
|-------|------|------|
| float64 | 1e-12 | 1e-15 |
| float32 | 1e-5 | 1e-8 |
| float16 | 1e-2 | 1e-2 |
| bfloat16 | 1e-2 | 2e-2 |
| float8_e4m3fn | 1e-1 | 1e-1 |
Run the op's test suite filtering for the Triton backend:
```bash
# All Triton correctness tests
pytest tests/ops/test_<op>.py -k "triton" -vs
# For suites/ operators (external framework)
pytest tests/suites/<framework>/test_<op>.py -k "triton" -vs
```
**Pass gate:** `N passed, 0 failed` before moving to performance.
If tests fail, see [debugging.md](./debugging.md) for the most common root causes
(`cudaErrorIllegalAddress`, pointer type mismatch, stride overflow, dtype issues).
---
## Benchmark Harness
TileGym provides a benchmark harness with provider abstraction for systematic
performance measurement across backends:
```python
from harness import run_benchmarks, get_providers
run_benchmarks(
kernel_name="matmul",
providers=get_providers(), # ["triton", "cutile", "pytorch"]
x_name="M",
x_vals=[512, 1024, 2048, 4096],
make_fwd_fn=lambda provider, x: create_forward_fn(provider, x),
csv_path="./benchmark_results.csv",
)
```
**Provider mapping**:
- `triton` — Triton backend
- `cutile` — cuTile backend
- `pytorch` — PyTorch reference
Run performance tests to compare backends side-by-side:
```bash
# Ops
pytest tests/ops/test_<op>.py -k "test_perf" --print-record -v
# Suites
pytest tests/suites/<framework>/test_<op>.py -k "test_perf" --print-record -v
```
The output table includes TFLOPS (or GB/s) per config and backend. The acceptance
threshold is **Triton ≥ 80% of cuTile** across all tested configs.
---
## Conversion Skill Workflow
The cuTile→Triton conversion follows a **5-phase gated workflow**:
```
┌─────────────────────────────────────────────────────────────────┐
│ SKILL WORKFLOW PHASES │
├─────────────────────────────────────────────────────────────────┤
│ │
│ c2t-1 (Optional) c2t-2 c2t-3 │
│ ┌─────────────┐ ┌─────────────┐ ┌─────────────┐ │
│ │ Test │ ──▶ │ Convert │ ──▶│ Test │ │
│ │ Coverage │ │ cuTile→ │ │ Correctness │ │
│ │ Analysis │ │ Triton │ │ (pytest) │ │
│ └─────────────┘ └─────────────┘ └──────┬──────┘ │
│ │ │
│ ┌─────────────▼───────────┐ │
│ c2t-5 │ c2t-4 │ │
│ ┌─────────────┐ │ TMA OPTIMIZATION │ │
│ │ Performance │ ◀─────────────────│ (MANDATORY) │ │
│ │ Test │ │ • 2D+ loads → TMA │ │
│ │ (≥80% of │ │ • tl.assume() hints │ │
│ │ cuTile) │ │ • Dual kernels if │ │
│ └─────────────┘ │ transpose flag │ │
│ └─────────────────────────┘ │
└─────────────────────────────────────────────────────────────────┘
```
### Phase Summary
| Phase | ID | Purpose | Gate Criteria |
|-------|-----|---------|---------------|
| Test Coverage | c2t-1 | Verify cuTile tests pass | Optional baseline |
| Convert | c2t-2 | Apply API mapping | No syntax errors |
| Test | c2t-3 | Correctness validation | `0 failed` |
| TMA Optimize | c2t-4 | **MANDATORY** TMA for 2D+ | No raw ptr+mask loads |
| Performance | c2t-5 | Benchmark comparison | Triton ≥ 80% cuTile |
### TMA Verification (Pre-Benchmark)
Before running perf tests, confirm every 2D+ tile load uses TMA — raw pointer+mask loads
cause 5–20× regressions that will fail the performance gate:
```bash
# Should return 0 for fully-optimized kernels
grep -c "tl\.load.*mask" <your_triton_kernel.py>
# Confirm TMA descriptors are present
grep -n "make_tensor_descriptor\|make_block_ptr" <your_triton_kernel.py>
```
---
## Related Documents
- [workflow.md](../translations/workflow.md) — Full phase-gated conversion workflow
- [debugging.md](./debugging.md) — Runtime error diagnosis
- [advanced-patterns.md](../translations/advanced-patterns.md) — Dual-kernel layout flags, autotune
references/optimization-strategy.md
# Optimization strategy (cuTile → Triton)
**Purpose:** One-page strategy distilled from [translations/advanced-patterns.md](../translations/advanced-patterns.md) and [optimizing-reference.md](./optimizing-reference.md). Agents **must** apply this when converting **attention / FMHA / Gemma-style** kernels (e.g. `gemma_attention`) or any kernel where Triton is expected to match an optimized in-repo baseline.
**Full detail:** Use this file for *what to do and in what order*; open the two linked docs for proofs, code samples, and edge cases.
---
## When to read this (mandatory triggers)
| Trigger | Action |
|---------|--------|
| Converting **attention**, **FMHA**, **sliding window**, **soft cap**, or **GQA** (e.g. Gemma) | Read **§4 Gemma / FMHA checklist** below **before** writing the Triton inner loop. |
| Host exposes **`transpose` / `transpose_v`** or MLA-style layout modes | Read **§1** + [advanced-patterns §1](../translations/advanced-patterns.md#dual-layout-flag); use **two kernels**, not one + `tl.trans` in the KV loop. |
| Kernel uses **`@triton.autotune`** | Read **§2** + [advanced-patterns §2](../translations/advanced-patterns.md); grid **`lambda META: (...)`**; never freeze `BLOCK_*` from Python in a way that ignores autotune. |
| After TMA is in place but Triton still **>20% slower** than cuTile / baseline | Walk **§3** + [optimizing-reference](./optimizing-reference.md) sections 1–7 and §9 as applicable. |
| **10–50×** regression | [translations/workflow.md](../translations/workflow.md) — **CRITICAL PERFORMANCE PATTERNS** first (raw `tl.load` vs TMA). |
---
## §1 Layout flags → structure (from advanced-patterns)
- **Dual layout (`transpose` / `transpose_v`):** Implement **separate `@triton.jit` kernels** when math and TMA layouts differ per mode (MLA: `qk` `[H,N]` vs `[N,H]`, different `V` descriptor). Reusing the transpose path with extra `tl.trans` per KV block → **3–15×** on the other mode.
- **Autotune + grid:** `grid = lambda META: (triton.cdiv(..., META["BLOCK_M"]), ...)` — tile sizes come from **META**, not hard-coded `forward()` kwargs that bypass tuning.
- **Host `TensorDescriptor` vs in-kernel `tl.make_tensor_descriptor`:** If descriptors must track autotuned block sizes, use a **`pre_hook`** to set `block_shape` per config (see Gemma FMHA host side), or follow `mla_decoding.py` in-kernel style — do not freeze wrong tile sizes on the host.
- **Multi-tensor small ops:** Batched launch (`program_id(1)` selects tensor) — [optimizing-reference §8](./optimizing-reference.md), [advanced-patterns §5](../translations/advanced-patterns.md).
---
## §2 Post-TMA micro-optimizations (from optimizing-reference)
Apply in roughly this order after **all** 2D+ tile paths use TMA (or justified exceptions):
| Priority | Pattern | Impact (typical) | Pointer |
|----------|---------|------------------|---------|
| 1 | **Autotune breadth + backend/GPU split** — `get_available_triton_backend()`, `torch.cuda.get_device_capability()`, `num_stages`, `num_warps`, **`occupancy`**, `warp_specialize` where the stack supports it | **10–20%+** | [optimizing-reference §4](./optimizing-reference.md) |
| 2 | **EVEN_K / EVEN_*** heuristics — skip masks in inner loop when divisible | **5–15%** | §1 |
| 3 | **Transpose** — pointer/layout encoding; avoid **`tl.trans` inside the K/V loop** when avoidable | **5–15%** | §2 |
| 4 | **2D grid** for pointer BMM — `(num_mn_blocks, batch)` + grouped M for L2 | **0–10%** | §3 |
| 5 | **Epilogue subtile** (large TMA stores) | **5–15%** | §5 |
| 6 | **`tl.assume`** alignment on strides/pointers | **5–15%** | §6 |
| 7 | **Persistent vs non-persistent** + occupancy tuning | **10–30%** | §7, §9.3 |
**Blackwell / complex iterative kernels:** TMA descriptors, **`tl.range(..., loop_unroll_factor=1)`** for heavy loops, **TMEM-friendly** `tl.dot` blocks (often ≥32), slab allocator — [optimizing-reference §9](./optimizing-reference.md), [advanced-patterns §6](../translations/advanced-patterns.md).
---
## §3 Fast vs slow patterns (skill gotchas, condensed)
| Slow | Fast |
|------|------|
| Raw `tl.load(ptr+offs, mask=…)` for **block-shaped 2D+** tiles | **`tl.make_tensor_descriptor` / host `TensorDescriptor` + TMA load/store** |
| `broadcast_to` + `tl.dot` for batched matmul | One batch (or head) per program; **2D `tl.dot`** on real tiles |
| `qk = tl.zeros(...); tl.dot(q, k, qk)` when a fused dot exists | **`qk = tl.dot(q, k)`** (avoid redundant zero + 3-arg dot if the compiler path is worse) |
| Generic autotune for all GPUs/backends | **Separate config lists** for different backends and architectures; **sm_90 vs sm_120 vs sm_80** |
| Wrong TMA **logical shape** (e.g. K/V head dim = **H** when tensor is **H_kv**) | Descriptor **`shape` matches tensor rank sizes** for GQA: **`H // QUERY_GROUP_SIZE`** (or `num_head_kv`) on K/V |
| Always `libdevice.tanh` in hot path | Use fast tanh approximation where numerics allow (see `gemma_attention.py`) |
| Forcing **`.contiguous()`** on every forward | Only when required for TMA/strides; avoid extra copies |
| Autotune **key** missing `WINDOW_SIZE`, `SOFT_CAP`, `dtype` | **Include** keys that change optimal tile or specialization |
---
## §4 Gemma FMHA / `gemma_attention` conversion checklist (mandatory)
When converting or reworking **Gemma-style attention** (soft cap, sliding window, causal, GQA, BNSD), **do not stop at “correct TMA”** — apply these checklist items to match optimized Triton patterns:
1. **GQA TMA metadata:** `Q` / `Out` descriptors use `[B, H, S, D]`; **K** and **V** descriptors use **`[B, H // QUERY_GROUP_SIZE, S_kv, D]`** (same strides as the physical K/V tensor). Using **`H`** for the KV head dimension in the descriptor shape is a common bug and can hurt TMA behavior.
2. **Autotune:** `get_configs()` branches on **`get_available_triton_backend()`** and **`torch.cuda.get_device_capability()`** (e.g. **(12,0)/(12,1)** vs **(9,0)** vs **(8,0)**). Include **`occupancy`**, **`warp_specialize`**, **`num_stages`**, **`num_warps`** as in the reference — a single “SM ≥ 10 → 256×128” grid **misses** tuned Blackwell behavior.
3. **`@triton.autotune` `key`:** Include **`S_qo`, `S_kv`, `BLOCK_D`, `STAGE`, `QUERY_GROUP_SIZE`, `WINDOW_SIZE`, `SOFT_CAP`, `dtype`** (or equivalent) so different attention modes do not share one stale config.
4. **QK matmul:** Prefer **`qk = tl.dot(q, k)`** over explicit **`tl.zeros` + `tl.dot(q, k, qk)`** unless profiling shows otherwise.
5. **Soft cap:** Preserve **scale order** (tanh on logits in **original scale**, then align with **`INV_LOG_2`** / `exp2` softmax). Use fast tanh when supported by the backend (guard with a feature-detection flag).
6. **Inner loop:** `offs_m` / `offs_n` as **`tl.constexpr`** where the reference does; **`tl.max(qk, 1)`** / **`tl.sum(p, 1)`** axis consistent with layout.
7. **Host:** `triton.set_allocator` for TMA metadata; **`pre_hook`** updates descriptor **`block_shape`** for each autotune config. Avoid unconditional **`.contiguous()`** on Q/K/V unless needed.
8. **Prune configs:** For causal **`STAGE == 3`**, enforce **`BLOCK_M % BLOCK_N == 0`** if the algorithm requires it (see reference `prune_invalid_configs`).
After implementation, run the same **pytest + perf gates** as in [SKILL.md](../SKILL.md) mandatory completion checklist.
---
## Cross-references
| Document | Role |
|----------|------|
| [translations/advanced-patterns.md](../translations/advanced-patterns.md) | Dual kernels, autotune+META, descriptor vs autotune, diagnosis table |
| [optimizing-reference.md](./optimizing-reference.md) | EVEN_K, transpose, grid, autotune, epilogue, alignment, persistent, batched launch, Blackwell §9 |
| [translations/workflow.md](../translations/workflow.md) | Phases, TMA gate, catastrophic regression section |
| [SKILL.md](../SKILL.md) | Master checklist; must link **optimization-strategy** for attention conversions |
references/optimizing-reference.md
# Triton Optimization Reference (cuTile → Triton)
**Use this reference when converting or optimizing GEMM/BMM/attention-style Triton kernels** so the result is within ~20% of cuTile (or of an existing optimized Triton implementation). Patterns below are derived from real comparisons (e.g. BMM: pointer + TMA kernels) and apply to batched matmul, attention, and block-level matmul.
**Prerequisites:** TMA is already applied for all 2D+ tile loads (complete **TMA OPTIMIZATION (Phase c2t-4)** and read **CRITICAL PERFORMANCE PATTERNS** in [translations/workflow.md](../translations/workflow.md)). This document covers **post-TMA** optimizations that can still yield **~10–20%** gains.
**Strategy hub:** [optimization-strategy.md](./optimization-strategy.md) condenses this file and [advanced-patterns.md](../translations/advanced-patterns.md) into an ordered checklist; use it for **attention / Gemma FMHA** (§4) and for **§2–§3** fast-vs-slow patterns before deep-diving here.
---
## When to Use This Reference
- **During conversion:** After Phase c2t-4 (TMA optimization), apply the patterns below for GEMM/BMM/attention kernels before running Phase c2t-5 (performance test). For Gemma-style attention, follow **[optimization-strategy.md §4](./optimization-strategy.md#4-gemma-fmha--gemma_attention-conversion-checklist-mandatory)** in parallel.
- **When Triton is 10–20% slower:** If perf test shows Triton within 2–5x of cuTile but still >20% slower, check this reference and **PERFORMANCE ANALYSIS (Phase c2t-5)** in [translations/workflow.md](../translations/workflow.md).
- **When comparing two Triton implementations:** Use the patterns as a checklist to explain or fix performance gaps (e.g. “naive” vs “optimized” BMM).
---
## 1. EVEN_K (or EVEN_*) Fast Path for Reductions
**Impact: ~5–15%** for GEMM/BMM when the reduced dimension (usually K) is divisible by the tile size.
**Problem:** In the K-loop (or any reduction loop), using a mask on every load when the remaining length equals the block size adds branches and prevents the compiler from emitting a single bulk load.
**Pattern:**
- Add a **heuristic** (or constexpr) that is true when the dimension is divisible by the block size (e.g. `EVEN_K: K % BLOCK_K == 0`).
- In the loop, **branch on that heuristic**: when true, use **unmasked** loads; when false, use masked loads with `k_remaining` (or equivalent).
```python
@triton.heuristics({"EVEN_K": lambda args: args["K"] % args["BLOCK_SIZE_K"] == 0})
@triton.jit
def kernel(..., EVEN_K: tl.constexpr):
for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
if EVEN_K:
a = tl.load(a_ptrs)
b = tl.load(b_ptrs)
else:
k_remaining = K - k * BLOCK_SIZE_K
a = tl.load(a_ptrs, mask=offs_k[None, :] < k_remaining, other=0.0)
b = tl.load(b_ptrs, mask=offs_k[:, None] < k_remaining, other=0.0)
accumulator += tl.dot(a, b)
a_ptrs += BLOCK_SIZE_K * stride_ak
b_ptrs += BLOCK_SIZE_K * stride_bk
```
**Avoid:** Always computing and applying masks in the inner loop when you could use an EVEN_* fast path.
---
## 2. Transpose: Pointer Arithmetic vs In-Loop Transpose
**Impact: ~5–15%** for BMM/GEMM when one or both inputs are transposed.
**Problem:** If the kernel supports transposed A or B, doing `tl.trans(a)` or `tl.trans(b)` **inside the K-loop** every iteration adds extra instructions and register pressure. The alternative is to encode transpose in **pointer arithmetic** so the loaded block already has the layout expected by `tl.dot`.
**Pattern:**
- **Preferred:** Compute different pointer strides/offsets for transposed vs non-transposed so that the **loaded tile is already in (BLOCK_M, BLOCK_K)** or **(BLOCK_K, BLOCK_N)** form. No `tl.trans` in the loop.
- **Acceptable when necessary:** If descriptor/TMA API forces a fixed block shape, use `tl.trans` after load but keep it out of the hottest path (e.g. one trans per load, not per element).
```python
# GOOD: Transpose encoded in pointer layout (no tl.trans in loop)
if transpose_a:
a_ptrs = a_ptr + pid_q * stride_aq + offs_am[:, None] * stride_ak + offs_k[None, :] * stride_am
else:
a_ptrs = a_ptr + pid_q * stride_aq + offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak
# ... in loop: a = tl.load(a_ptrs); accumulator += tl.dot(a, b)
```
```python
# SLOW: Transpose in K-loop every iteration
for k in range(num_k_tiles):
a = tl.load(a_ptrs, mask=a_mask, other=0.0)
a = tl.trans(a) # Extra work every K tile
b = tl.load(b_ptrs, mask=b_mask, other=0.0)
b = tl.trans(b)
acc = tl.dot(a, b, acc)
```
**Apply to:** Pointer-based GEMM/BMM kernels; in TMA kernels, prefer descriptor block_shape that matches the expected 2D layout so no trans is needed after load.
---
## 3. Grid Layout (Pointer-Based Kernels)
**Impact: ~0–10%** depending on GPU and batch size.
**Problem:** A 3D grid `(num_pid_m, num_pid_n, Q)` can map to hardware differently than a 2D grid where batch is on `program_id(axis=1)` and the 2D tile index is flattened on axis=0. The latter often gives better occupancy and scheduling on many GPUs.
**Pattern:**
- Prefer a **2D grid**: `(num_pid_m * num_pid_n, Q)` with `pid = tl.program_id(0)` and `pid_q = tl.program_id(1)`. Decode `pid` into `pid_m` and `pid_n` (e.g. with GROUP_SIZE_M grouping for L2 reuse).
- Use **grouped ordering** (GROUP_SIZE_M) in the decoding so that tiles in the same M-group are adjacent; this improves L2 reuse and can match cuTile’s scheduling.
```python
# 2D grid: (num_blocks_mn, Q)
grid = lambda META: (
triton.cdiv(M, META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]),
Q,
)
# In kernel: pid = tl.program_id(0); pid_q = tl.program_id(1); decode pid -> pid_m, pid_n
```
**Apply to:** Non-persistent, pointer-based BMM/GEMM when comparing or porting from a kernel that uses a 2D grid.
---
## 4. Autotune Breadth and Backend/GPU-Specific Configs
**Impact: ~10–20%** (picking a better block size / num_stages / num_warps).
**Problem:** A small or generic autotune space can miss the best config for a given GPU or backend (e.g. sm_90 vs sm_120). Missing `num_stages`, `num_warps`, or `occupancy` can leave significant performance on the table.
**Pattern:**
- **Expand config space** for GEMM/BMM: vary `BLOCK_M`, `BLOCK_N`, `BLOCK_K`, `num_stages`, `num_warps`, and for persistent kernels `occupancy` and optionally `num_ctas`.
- **Specialize by backend and GPU:** Use `get_available_triton_backend()` and `torch.cuda.get_device_capability()` to return different config lists (e.g. sm_120/sm_121 vs sm_90 vs older).
- **Pre-hook for TMA:** When using tensor descriptors, set `block_shape` in a `pre_hook` from the chosen `BLOCK_M`/`BLOCK_N`/`BLOCK_K` so TMA uses the same tile sizes as the kernel.
```python
def get_configs(pre_hook=None):
cap = torch.cuda.get_device_capability()
if cap in [(12, 0), (12, 1)]:
return [
triton.Config({"BLOCK_M": BM, "BLOCK_N": BN, "BLOCK_K": BK, ...}, num_stages=s, pre_hook=pre_hook)
for BM in [64, 128] for BN in [64, 128] for BK in [32, 64] for s in [2, 3]
]
elif cap == (9, 0):
return [...]
else:
return [...]
```
**Apply to:** All TMA and pointer-based GEMM/BMM kernels after conversion.
---
## 5. Epilogue Subtile (TMA Store)
**Impact: ~5–15%** on some GPUs (e.g. Blackwell) when the C tile is large.
**Problem:** Writing the full output block in one TMA store can underutilize the memory subsystem or cause suboptimal scheduling. Splitting the C block into two halves and doing two stores (with correct offsets) can improve store throughput.
**Pattern:**
- Add an **EPILOGUE_SUBTILE** (or similar) constexpr. When true, treat the N dimension of the accumulator as two subtiles (e.g. `BLOCK_N // 2` each). Reshape/permute the accumulator accordingly, convert to output dtype, and **store twice** (first half at `[..., offs_bn]`, second at `[..., offs_bn + BLOCK_N // 2]`).
- Include both `EPILOGUE_SUBTILE=True` and `False` in autotune; let the tuner choose.
```python
if EPILOGUE_SUBTILE:
acc = tl.reshape(acc, (1, BLOCK_M, 2, BLOCK_N // 2))
acc = tl.permute(acc, (0, 1, 3, 2))
acc0, acc1 = tl.split(acc)
c_desc.store([..., offs_bn], acc0.to(dtype))
c_desc.store([..., offs_bn + BLOCK_N // 2], acc1.to(dtype))
else:
c_desc.store([..., offs_bn], acc.to(dtype))
```
**Apply to:** TMA-based GEMM/BMM when the output block is large (e.g. BLOCK_N ≥ 128) and the target GPU benefits (typically sm_90+).
---
## 6. Alignment and Bounds Hints
**Impact: ~5–15%** when the compiler or TMA can use alignment for wider transactions.
**Pattern:**
- Add `tl.assume(stride % 8 == 0)` (or 16) for strides that are known to be aligned.
- Add `tl.assume(ptr.to(tl.int64) % 16 == 0)` for base pointers when valid.
- For TMA, ensure descriptor `block_shape` and strides match the actual access; avoid unnecessary masks for full tiles (TMA can handle bounds).
See **TMA Checklist** in [translations/workflow.md](../translations/workflow.md) and [SKILL.md](../SKILL.md) (Performance Gotchas) for alignment.
---
## 7. Persistent vs Non-Persistent and Occupancy
**Impact: ~10–30%** depending on problem size and GPU.
**Pattern:**
- For **small/medium** problems, a **non-persistent** grid (one program per tile) can be faster due to less scheduling overhead.
- For **large** problems or when you want to amortize launch cost, use a **static persistent** loop: `for current_pid in tl.range(pid, total_tiles, num_programs, flatten=True)`, with grid size = `min(NUM_SMS // num_ctas, total_tiles) * occupancy`. Tune `occupancy` (1, 2, 4) in autotune.
- When using persistent scheduling, decode the flat `current_pid` into (batch, pid_m, pid_n) with the same GROUP_SIZE_M grouping for L2 reuse.
---
## Quick Checklist for GEMM/BMM Conversions
After TMA is in place, apply these before declaring conversion “optimized”:
| Check | Action |
|-------|--------|
| **EVEN_K** | Add heuristics and branch: unmasked loads when `K % BLOCK_K == 0`. |
| **Transpose** | Prefer pointer arithmetic for transposed A/B; avoid `tl.trans` in the K-loop. |
| **Grid** | Prefer 2D grid `(num_blocks_mn, Q)` with grouped (GROUP_SIZE_M) decoding for pointer BMM. |
| **Autotune** | Backend- and GPU-specific configs; vary BLOCK_*, num_stages, num_warps, occupancy. |
| **Epilogue** | Consider EPILOGUE_SUBTILE for TMA C-store when BLOCK_N is large (e.g. ≥128). |
| **Alignment** | Add `tl.assume()` for strides and pointers where valid. |
| **Persistent** | For large problems, use static persistent + occupancy in autotune. |
---
## 8. Batched Kernel Launch (Multi-Tensor Operations)
**Impact: ~2–4× speedup** for memory-bound operations on small-to-medium tensors (common in LLM inference KV-cache concatenation).
**Problem:** Operations like `cat`, `stack`, or multi-tensor copies often process N tensors by launching N separate kernels. Each kernel launch has ~5–10µs overhead on GPU. For small tensors, this overhead dominates actual compute time.
**Pattern:**
- **Batch multiple tensors into a single kernel launch** using a 2D grid where one dimension iterates over tensors (up to 4 is a good batch size).
- Pass all tensor pointers and metadata as separate kernel arguments; use `tl.program_id(1)` to select which tensor the block processes.
- Use conditional assignment (not branching) to select the correct tensor's data.
```python
@triton.jit
def batched_copy_kernel_4(
out_ptr,
in_ptr_a, in_ptr_b, in_ptr_c, in_ptr_d,
size_a, size_b, size_c, size_d,
offset_a, offset_b, offset_c, offset_d,
total_a, total_b, total_c, total_d,
BLOCK_X: tl.constexpr,
):
pid_x = tl.program_id(0) # Block index over elements
pid_y = tl.program_id(1) # Tensor index (0-3)
# Select tensor data based on pid_y (no branching in hot path)
if pid_y == 0:
in_ptr, size_in, offset, total = in_ptr_a, size_a, offset_a, total_a
elif pid_y == 1:
in_ptr, size_in, offset, total = in_ptr_b, size_b, offset_b, total_b
elif pid_y == 2:
in_ptr, size_in, offset, total = in_ptr_c, size_c, offset_c, total_c
else:
in_ptr, size_in, offset, total = in_ptr_d, size_d, offset_d, total_d
block_start = pid_x * BLOCK_X
offsets = tl.arange(0, BLOCK_X)
mask = block_start + offsets < total
idx = block_start + offsets
# Compute output index with unified formula
out_idx = compute_output_index(idx, size_in, offset, ...)
data = tl.load(in_ptr + idx, mask=mask)
tl.store(out_ptr + out_idx, data, mask=mask)
# Host-side: batch tensors in groups of 4
def cat(tensors, dim=0):
BLOCK = 1024
i = 0
while i < len(tensors):
batch = tensors[i : i + 4]
num_in_batch = len(batch)
# Pad unused slots with placeholder (first tensor)
args = []
max_elements = 0
for j in range(4):
if j < num_in_batch:
t = batch[j].contiguous()
args.extend([t, t.shape[dim], offset, t.numel()])
max_elements = max(max_elements, t.numel())
else:
args.extend([batch[0], 0, 0, 0]) # Placeholder
# 2D grid: (blocks over elements, tensors in batch)
grid = (triton.cdiv(max_elements, BLOCK), num_in_batch)
batched_copy_kernel_4[grid](out, *args, BLOCK_X=BLOCK)
i += num_in_batch
```
**Why this works:**
| Aspect | Per-Tensor Launch | Batched (4 tensors) |
|--------|-------------------|---------------------|
| Kernel launches | N | ⌈N/4⌉ |
| Launch overhead | N × 5–10µs | ⌈N/4⌉ × 5–10µs |
| GPU parallelism | Sequential | Concurrent (different SMs) |
| Block size | Often smaller | Can use larger (1024) |
**When to use:**
- **Multi-tensor memory-bound ops:** `cat`, `stack`, `split`, `chunk`, multi-tensor copy/scatter/gather
- **Small-to-medium tensor sizes:** When kernel launch overhead is significant relative to compute
- **LLM inference:** KV-cache concatenation, attention output gathering
**When NOT to use:**
- **Large tensors:** Launch overhead is negligible compared to compute
- **Single tensor operations:** No batching benefit
- **Compute-bound ops:** GEMM, convolution (launch overhead already amortized)
**Real-world example:** TileGym's original `cat` implementation (adapted from [FlagGems](https://github.com/FlagOpen/FlagGems)) uses this pattern to achieve **2–4× speedup** over naive per-tensor launches in transformer KV-cache operations.
---
## 9. Blackwell Advanced Optimization Patterns
**Impact: 2–10× speedup** on Blackwell (sm_100+) GPUs when converting complex kernels with iterative algorithms.
These patterns were discovered comparing optimized vs naive implementations of `chunk_gated_delta_rule` (a chunked linear attention kernel with Neumann series matrix inversion). They apply to any kernel with:
- Iterative loops (matrix inversion, recurrence, series expansion)
- Large intermediate tensors
- Multiple `tl.dot` operations
- Block-matrix algorithms
### 9.1 TMA Descriptors vs Raw Pointer Arithmetic
**Impact: 20–50%** on Blackwell for structured memory access.
**Problem:** Raw pointer arithmetic with strides requires the GPU to compute addresses at runtime, wastes registers on stride calculations, and prevents hardware-accelerated bulk transfers.
**Pattern:**
```python
# SLOW: Raw pointer + stride arithmetic
@triton.jit
def kernel_slow(Q_ptr, stride_qb, stride_qt, stride_qh, stride_qk, ...):
q_ptrs = Q_ptr + b_idx * stride_qb + t_ids[:, None] * stride_qt + h_idx * stride_qh + offs_k[None, :] * stride_qk
q = tl.load(q_ptrs, mask=valid[:, None] & mask_k[None, :], other=0.0)
# FAST: TMA descriptor (triton DSL / Blackwell)
from triton.tools.tensor_descriptor import TensorDescriptor
@triton.jit
def kernel_fast(Q_desc, ...):
q = tl.reshape(Q_desc.load([b_idx, t_offset, h_idx, 0]), (CHUNK_SIZE, BLOCK_K)).to(tl.float32)
```
**Host-side setup:**
```python
Q_desc = TensorDescriptor.from_tensor(query, [1, BS, 1, BLOCK_K])
kernel_fast[grid](Q_desc, ...)
```
**Why it's faster:**
- TMA is a hardware unit on Blackwell that handles bulk data movement
- No runtime address calculation — hardware computes offsets
- Better memory coalescing (hardware optimizes access patterns)
- Lower register pressure (no stride variables needed)
### 9.2 Loop Unrolling Control for Register Pressure
**Impact: 2–5×** for kernels with iterative algorithms (matrix inversion, recurrence).
**Problem:** Full loop unrolling makes all intermediate values live simultaneously, causing massive register spilling. NCU shows symptoms like "168 regs/thread + 51M local spill requests".
**Pattern:**
```python
# SLOW: Full unroll (implicit with range() or tl.static_range)
@triton.jit
def _solve_tril_slow(A, BS: tl.constexpr):
for i in tl.static_range(1, BS): # 31 iterations fully unrolled
# All 31 intermediate values live simultaneously → register spill
is_row = offs == i
row = tl.sum(tl.where(is_row[:, None], A, 0.0), axis=0)
corr = tl.sum(row[:, None] * A, axis=0)
A = A + tl.where(is_row[:, None], corr[None, :], 0.0)
return A
# FAST: Controlled unroll factor
@triton.jit
def _solve_tril_fast(A, BS: tl.constexpr):
for i in tl.range(1, BS, loop_unroll_factor=1): # No unroll
# Only current iteration's intermediates are live
is_row = offs == i
row = tl.sum(tl.where(is_row[:, None], A, 0.0), axis=0)
corr = tl.sum(row[:, None] * A, axis=0)
A = A + tl.where(is_row[:, None], corr[None, :], 0.0)
return A
```
**When to use `loop_unroll_factor=1`:**
- Loop body has many intermediate tensors
- Loop iteration count > 16
- NCU shows high register usage or local memory spills
- Kernel is slower than expected despite correct algorithm
**When full unroll is OK:**
- Loop body is simple (few intermediates)
- Iteration count ≤ 8
- Register pressure is not a concern
### 9.3 Occupancy Autotuning
**Impact: 1.5–3×** by finding optimal resource allocation.
**Problem:** Default occupancy=1 lets the compiler use maximum resources per thread, which can backfire when aggressive optimization causes spilling.
**Pattern:**
```python
@triton.autotune(
configs=[
triton.Config({"occupancy": 1}, num_stages=3),
triton.Config({"occupancy": 2}, num_stages=3),
triton.Config({"occupancy": 4}, num_stages=3),
triton.Config({"occupancy": 8}, num_stages=3),
triton.Config({"occupancy": 2}, num_stages=4),
triton.Config({"occupancy": 4}, num_stages=4),
],
key=["K_dim", "V_dim"],
)
@triton.jit
def kernel(..., occupancy: tl.constexpr = 1):
...
```
**Occupancy guidance:**
- Higher occupancy = more concurrent thread blocks = better latency hiding
- Higher occupancy forces compiler to use fewer registers per thread
- Start with `[1, 2, 4]` for dot-heavy kernels
- Try `[1, 4, 8, 16]` for norm/elementwise kernels
### 9.4 TMEM-Friendly Block Sizes for `tl.dot`
**Impact: 1.5–2×** on Blackwell when `tl.dot` shapes qualify for Tensor Memory (TMEM).
**Problem:** Small block sizes (e.g., 16×16) force `tl.dot` to materialize results in registers. Larger sizes (≥32×32) enable TMEM, which is fast on-chip memory between registers and shared memory.
**Pattern:**
```python
# SLOW: 16×16 blocks → register materialization
BS: tl.constexpr = 16
# 4 diagonal blocks + 6 off-diagonal = 10 blocks to manage
# tl.dot shapes: (16, BLOCK_K) × (BLOCK_K, 16) → (16, 16) — no TMEM
# FAST: 32×32 blocks → TMEM enabled
BS: tl.constexpr = 32
# 2 diagonal blocks + 1 off-diagonal = 3 blocks (simpler)
# tl.dot shapes: (32, BLOCK_K) × (BLOCK_K, 32) → (32, 32) — TMEM OK
```
**Rule of thumb:** For hierarchical block algorithms, prefer block sizes that make `tl.dot` operands ≥32 in both dimensions on Blackwell.
### 9.5 Single Buffer Allocation (Slab Allocator)
**Impact: 10–30%** reduction in kernel launch overhead for kernels with multiple intermediate buffers.
**Problem:** Multiple `torch.empty()` calls trigger multiple `cudaMalloc` calls (~10–100μs each), causing fragmentation and launch latency.
**Pattern:**
```python
# SLOW: 6 separate allocations
q_chunked = torch.empty(B, H, num_chunks, chunk_size, K, device=device, dtype=torch.float32)
k_chunked = torch.empty(B, H, num_chunks, chunk_size, K, device=device, dtype=torch.float32)
v_corrected = torch.empty(B, H, num_chunks, chunk_size, V, device=device, dtype=torch.float32)
k_cumdecay = torch.empty(B, H, num_chunks, chunk_size, K, device=device, dtype=torch.float32)
g_cum = torch.empty(B, H, num_chunks, chunk_size, device=device, dtype=torch.float32)
output = torch.empty(B, H, num_chunks, chunk_size, V, device=device, dtype=torch.float32)
# FAST: Single slab allocation
total_elems = B * H * NC * (3 * K + 2 * V + 1)
buf = torch.empty(total_elems, device=device, dtype=torch.float32)
off = 0
def _slab(shape):
nonlocal off
n = 1
for s in shape:
n *= s
t = buf[off : off + n].view(shape)
off += n
return t
q_chunked = _slab((B, H, num_chunks, chunk_size, K))
k_chunked = _slab((B, H, num_chunks, chunk_size, K))
v_corrected = _slab((B, H, num_chunks, chunk_size, V))
k_cumdecay = _slab((B, H, num_chunks, chunk_size, K))
g_cum = _slab((B, H, num_chunks, chunk_size))
output = _slab((B, H, num_chunks, chunk_size, V))
```
**Benefits:**
- Single `cudaMalloc` call instead of N calls
- Contiguous memory → better cache locality
- Reduced memory fragmentation
- Works with any number of intermediate buffers
### 9.6 Broadcasting vs `tl.expand_dims`
**Impact: 5–10%** in hot loops.
**Pattern:**
```python
# SLOW: tl.expand_dims may generate instructions
is_row_col = tl.expand_dims(is_row, axis=1)
row = tl.sum(tl.where(is_row_col, A, 0.0), axis=0)
# FAST: Broadcasting is compile-time (zero runtime cost)
row = tl.sum(tl.where(is_row[:, None], A, 0.0), axis=0)
```
**Rule:** Prefer `tensor[:, None]` or `tensor[None, :]` over `tl.expand_dims` for simple broadcast patterns.
### 9.8 Quick Checklist for Blackwell Optimization
After basic conversion, apply these for complex kernels on Blackwell:
| Check | Action |
|-------|--------|
| **TMA** | Use `TensorDescriptor` for all 2D+ loads/stores on triton DSL path |
| **Loop unroll** | Add `loop_unroll_factor=1` for loops with >16 iterations and complex bodies |
| **Occupancy** | Add `occupancy` to autotune configs: `[1, 2, 4]` for dot-heavy, `[1, 4, 8, 16]` for others |
| **Block size** | Use ≥32×32 blocks for `tl.dot` to enable TMEM |
| **Slab alloc** | Combine multiple intermediate buffers into single allocation |
| **Dual path** | Implement both triton DSL (TMA) and OpenAI Triton (pointer) paths |
| **Broadcasting** | Use `[:, None]` instead of `tl.expand_dims` |
### 9.9 Profiling with NCU
When a kernel is slower than expected, profile with NVIDIA Nsight Compute:
```bash
ncu --set full -o profile python test_kernel.py
```
**Key metrics to check:**
- **Registers/thread**: >128 suggests register pressure
- **Local memory**: Any spills indicate register overflow
- **Occupancy**: Low achieved vs theoretical suggests resource constraints
- **Memory throughput**: Compare to roofline
**Symptoms → Actions:**
| NCU Symptom | Likely Cause | Action |
|-------------|--------------|--------|
| High registers + local spills | Full loop unroll | Add `loop_unroll_factor=1` |
| Low occupancy | Too many resources/thread | Increase `occupancy` hint |
| Memory bound, low throughput | Raw pointer loads | Convert to TMA |
| Compute bound, low FLOPS | Small `tl.dot` shapes | Increase block sizes to ≥32 |
---
## Cross-References
Use [SKILL.md](../SKILL.md) as the hub. Deeper workflow sections (**TMA OPTIMIZATION**, **CRITICAL PERFORMANCE PATTERNS**, **PERFORMANCE ANALYSIS**) are in [translations/workflow.md](../translations/workflow.md).
references/performance-gotchas.md
# Performance Gotchas — 10-50× Regression Risk
**⚠️ These cause CATASTROPHIC slowdowns. Check BEFORE benchmarking.**
| Pattern | SLOW (Regression) | FAST (Optimized) | Impact |
|---------|-------------------|------------------|--------|
| **Memory access (2D+ tiles)** | Raw ptr + masks: `tl.load(ptr+offs, mask=m)` for block-shaped 2D+ loads | TMA: `tl.make_tensor_descriptor(...).load([off])` | **5-20x (500%-2000%)** — **most common cause of conversion regression; use TMA for every 2D+ tile load** |
| **Group iteration** | Linear search all groups per tile | While-loop with `last_problem_end` tracking | **2-5x** |
| **Tile sizes** | Fixed `BLOCK_M=128, BLOCK_N=128` | `@triton.autotune` with GPU-specific configs | **2-3x** |
| **Alignment** | No hints | `tl.assume(stride % 8 == 0)`, `tl.assume(ptr % 16 == 0)` | **1.5-2x** |
| **Full-tile masks** | Masks on every load/store | Remove masks, let TMA handle bounds | **1.2-1.5x** |
| **K-loop offsets** | Recalculate full offset each iter | `a_ptrs += BLOCK_K` or TMA offset increment | **1.1-1.2x** |
| **Memory layout** | 5D reshape for split dims | Transpose + contiguous first/second half | **50-150%** |
| **constexpr params** | Dynamic dimension params | Mark `bs`, `hd`, `n_h` as `tl.constexpr` | **10-20%** |
| **Unnecessary clones** | `q.clone()` before in-place op | Transpose → contiguous (natural copy) | **10-20%** |
| **Row stride pattern** | Per-element stride calculation | Row stride with `ptr + pid * row_stride` | **10-30%** |
| **broadcast_to + tl.dot** | `W.broadcast_to((BS,M,K))` then `tl.dot(W, X)` | 1-batch-per-block, load W as 2D `(M,K)`, use `tl.dot(W, X)` | **10-50×** (FFT case study) |
| **extract_slice chains** | Chain of `extract_slice` + `reshape` (24+ calls) | Direct offset computation, load into final shape | **2-5×** |
**Full details:** [../translations/workflow.md](../translations/workflow.md) — section **CRITICAL PERFORMANCE PATTERNS (AVOID 10-50x REGRESSION)**
Full API mapping: [api-mapping.md](./api-mapping.md).
Triton math dtype (erf/erfc/exp/log/sqrt) and the "don't substitute erf with tanh" pattern: [debugging.md](./debugging.md) — section **Triton Math Function Dtype Requirements (CRITICAL)**.
For the broader cuTile → Triton translation gotchas (mma, type cast, grid, layout flags, batched matmul, etc.), see [gotchas.md](./gotchas.md).
skill-card.md
## Description: <br>
Converts cuTile GPU kernels (@ct.kernel) to Triton (@triton.jit), covering standard in-repo conversion, debugging (cudaErrorIllegalAddress, shape mismatch, numerical mismatch), mapping cuTile idioms (ct.load/ct.store, ct.Constant, ct.launch) to Triton equivalents, and dual-kernel layout flags with autotune grids via META. <br>
This skill is ready for commercial/non-commercial use. <br>
## Owner
NVIDIA <br>
### License/Terms of Use: <br>
CC-BY-4.0 AND Apache 2.0 <br>
## Use Case: <br>
Developers and engineers working on GPU kernels use this skill to port, translate, and optimize cuTile kernels in the TileGym repository as Triton kernels, and to debug existing Triton translations through a phase-gated analyze, convert, validate, test, and benchmark workflow. <br>
### Deployment Geography for Use: <br>
Global <br>
## Requirements / Dependencies: <br>
**Requires API Key or External Credential:** [Not Specified] <br>
**Credential Type(s):** [None identified] <br>
Do not include secrets in prompts/logs/output; use least-privilege credentials; rotate keys as appropriate. <br>
## Known Risks and Mitigations: <br>
Risk: Review before execution as proposals could introduce incorrect or misleading guidance into skills. <br>
Mitigation: Review and scan skill before deployment. <br>
## Reference(s): <br>
- [cuTile → Triton API Mapping](references/api-mapping.md) <br>
- [Triton Runtime Error Debugging (cuTile → Triton)](references/debugging.md) <br>
- [Gotchas — Most Common cuTile → Triton Translation Errors](references/gotchas.md) <br>
- [Testing & Validation (cuTile → Triton)](references/harness-integration.md) <br>
- [Optimization strategy (cuTile → Triton)](references/optimization-strategy.md) <br>
- [Triton Optimization Reference (cuTile → Triton)](references/optimizing-reference.md) <br>
- [Performance Gotchas — 10-50× Regression Risk](references/performance-gotchas.md) <br>
- [cuTile to Triton Conversion Workflow](translations/workflow.md) <br>
- [File Structure & Registration (cuTile → Triton)](translations/file-structure.md) <br>
- [Advanced Triton patterns (cuTile → Triton)](translations/advanced-patterns.md) <br>
## Skill Output: <br>
**Output Type(s):** [Code, Files, Shell commands, Analysis] <br>
**Output Format:** [Python source files (Triton kernels and host wrappers) plus Markdown with inline bash code blocks and performance comparison tables] <br>
**Output Parameters:** [1D] <br>
**Other Properties Related to Output:** [None] <br>
## Evaluation Agents Used: <br>
- Claude Code (`aws/anthropic/bedrock-claude-opus-4-8`) <br>
- Codex (`openai/openai/gpt-5.5`) <br>
## Evaluation Tasks: <br>
5 evaluation tasks (1 positive, 4 negative), 1 attempt per task, each run in an isolated sandbox pod; evaluated through a three-tier evaluation (static validation, semantic deduplication, live agent evaluation). <br>
## Evaluation Metrics Used: <br>
Reported benchmark dimensions: <br>
- Security: Is it safe to use? Scored from the `security` signal. <br>
- Correctness: Is the answer correct? Scored from the `accuracy` signal. <br>
- Discoverability: Was the right skill loaded when needed? Scored from the `skill_execution` signal. <br>
- Effectiveness: Did the skill help complete the task? Equal-weight mean of `goal_accuracy` and `behavior_check`. <br>
- Efficiency: Did it avoid wasted tool calls and token usage? Equal-weight mean of `skill_efficiency` and `token_efficiency`. <br>
Underlying evaluation signals used in this run: <br>
- `security`: Unsafe operations, secret leakage, and unauthorized access. <br>
- `skill_execution`: Whether the expected skill was selected, decoys were avoided, and the workflow executed. <br>
- `skill_efficiency`: Tool-call productivity (legacy wire id; routing is scored under Discoverability). <br>
- `accuracy`: Final-answer correctness against the reference answer. <br>
- `goal_accuracy`: Whether the user's goal was achieved. <br>
- `behavior_check`: Whether the expected workflow behavior was followed. <br>
- `token_efficiency`: Actual uncached prompt plus completion usage (50% of Efficiency). <br>
## Evaluation Results: <br>
| Measure | Claude Code (Baseline → Skill Uplift) | Codex (Baseline → Skill Uplift) |
|---|---:|---:|
| Overall | 99.8% — baseline ran, but no comparable score was available; uplift unavailable | 98.3% — baseline ran, but no comparable score was available; uplift unavailable |
| Security | 100.0% → 100.0% (±0.0 points) | 100.0% → 100.0% (±0.0 points) |
| Correctness | 80.0% → 100.0% (+20.0 points) | 100.0% → 100.0% (±0.0 points) |
| Discoverability | 100.0% — baseline ran, but no comparable score was available; uplift unavailable | 95.0% — baseline ran, but no comparable score was available; uplift unavailable |
| Effectiveness | 82.5% → 100.0% (+17.5 points) | 90.0% → 99.0% (+9.0 points) |
| Efficiency | 99.2% — baseline ran, but no comparable score was available; uplift unavailable | 97.6% — baseline ran, but no comparable score was available; uplift unavailable |
## Skill Version(s): <br>
1.0.0 (source: frontmatter metadata.version) <br>
## Ethical Considerations: <br>
NVIDIA believes Trustworthy AI is a shared responsibility and we have established policies and practices to enable development for a wide array of AI applications. When downloaded or used in accordance with our terms of service, developers should work with their internal team to ensure this skill meets requirements for the relevant industry and use case and addresses unforeseen product misuse. <br>
(For Release on NVIDIA Platforms Only) <br>
Please report quality, risk, security vulnerabilities or NVIDIA AI Concerns [here](https://app.intigriti.com/programs/nvidia/nvidiavdp/detail). <br>
SKILL.md
---
name: tilegym-converting-cutile-to-triton
description: Converts cuTile GPU kernels (@ct.kernel) to Triton (@triton.jit). Handles standard in-repo conversion, debugging (cudaErrorIllegalAddress, shape mismatch, numerical mismatch), and mapping cuTile idioms (ct.load/ct.store, ct.Constant, ct.launch) to Triton equivalents. Covers dual-kernel layout flags (e.g. transpose=True/False + autotune grid via META) per translations/advanced-patterns.md. Use when converting, porting, or translating cuTile kernels to Triton, or debugging existing Triton translations.
license: CC-BY-4.0 AND Apache-2.0
tools:
- Read
- Write
- Grep
- Glob
- Bash
metadata:
author: "TileGym Team <TileGym@nvidia.com>"
version: "1.0.0"
tags:
- cutile
- triton
- conversion
- gpu
- kernel
---
# cuTile → Triton Conversion
Convert `@ct.kernel` kernels to `@triton.jit`. API mapping: [references/api-mapping.md](./references/api-mapping.md) (cuTile → Triton).
*In this skill’s Markdown, Triton launch syntax `kernel［grid］(…)` uses Unicode brackets so link checkers do not parse `[grid](…)` as a hyperlink; use normal ASCII brackets in real Triton code.*
## Instructions
Follow the phase-gated workflow in [translations/workflow.md](./translations/workflow.md). Every conversion should go through **analyze → convert → validate → test → benchmark**, with explicit gates before moving on. Use the documents in [Workflow Selection](#workflow-selection) when the task matches a special case (errors, layout flags, perf).
0. **Optimization strategy (perf-sensitive / attention)** — If the op is **attention, FMHA, sliding window, soft cap, or GQA** (e.g. Gemma `gemma_attention`), read **[references/optimization-strategy.md](./references/optimization-strategy.md)** **before** converting the inner loop, then apply **[§4 Gemma FMHA checklist](./references/optimization-strategy.md#4-gemma-fmha--gemma_attention-conversion-checklist-mandatory)**. For other GEMM/BMM/attention-adjacent kernels, still skim **§2–§3** of that file after TMA is done.
1. **Select path** — Existing TileGym op: standard mode in `translations/workflow.md`. If the cuTile source uses `transpose` / `transpose_v`, dual layouts, or MLA-style paths, read [translations/advanced-patterns.md](./translations/advanced-patterns.md) **before** writing Triton (two kernels + `META` grid, not one kernel + `tl.trans`).
2. **Pre-flight** — Run the [Pre-flight Analysis](#pre-flight-analysis-run-before-converting) grep commands on the cuTile source. Count `@ct.kernel` definitions; note TMA-relevant `ct.load`/`ct.store`, `ct.launch`, `Constant`, and layout flags.
3. **Read mapping** — Keep [references/api-mapping.md](./references/api-mapping.md) open for cuTile → Triton API pairs. For runtime failures (illegal address, dtype, strides), use [references/debugging.md](./references/debugging.md).
4. **Convert** — Copy the [Conversion Checklist](#conversion-checklist) into a todo list and execute in order. Structure and file placement: [translations/file-structure.md](translations/file-structure.md). **Mandatory:** any **2D+ block-shaped** tile load/store uses `tl.make_tensor_descriptor` (TMA), not raw `tl.load(ptr+offs, mask=…)` for full tiles—skipping this is the most common source of large regressions. Host side: Triton bracket launch <code>kernel［grid］(args)</code> with tuple or `lambda META: (…)` for autotune; no `ct.launch`.
5. **Validate** — Syntax-check the new Triton module; run the relevant TileGym pytest targets for the op: `pytest tests/ops/test_<op>.py -k "triton" -vs`. Fix failures before benchmarking.
6. **Benchmark** — Compare Triton vs cuTile on perf tests. If Triton is clearly slower, follow **PERFORMANCE ANALYSIS (Phase c2t-5)** in [translations/workflow.md](./translations/workflow.md) and [references/optimizing-reference.md](./references/optimizing-reference.md) for GEMM/BMM/attention; use [references/optimization-strategy.md](./references/optimization-strategy.md) as the ordered checklist. If you see **10–50×** slowdowns, read **CRITICAL PERFORMANCE PATTERNS** in that same workflow file first.
**Execution rules (MUST):**
- Create and track the conversion checklist (e.g. TodoWrite) **before** editing kernel code; complete steps in order—do not skip pre-flight or TMA decisions.
- For **attention / FMHA / Gemma / GQA / soft cap / sliding window**: read [references/optimization-strategy.md](./references/optimization-strategy.md) and apply **§4** **before** treating the conversion as optimized.
- Do **not** ship raw pointer+mask 2D+ tile loads where TMA applies; document any intentional exception.
- If tests or benchmarks fail a gate, stop and fix **before** declaring the conversion done—do not stack unverified changes.
## Workflow Selection
- **Existing TileGym op** → Standard Mode: [translations/workflow.md](./translations/workflow.md)
- **Errors** (`cudaErrorIllegalAddress`, shape mismatch, numerical mismatch) → [references/debugging.md](./references/debugging.md)
- **Advanced patterns** (TMA, dual layout flags `transpose`, autotune + `META` grid, Array.slice, ct.gather().item()) → **[translations/advanced-patterns.md](./translations/advanced-patterns.md)** (MLA-style two kernels, avoid 3–15× regression on `transpose=False`).
- **Performance** (Triton kernel slower than cuTile, autotuning, profiling) → [translations/workflow.md](./translations/workflow.md) (section **PERFORMANCE ANALYSIS (Phase c2t-5)**)
- **Optimization strategy hub** (ordered checklist: advanced-patterns + optimizing-reference) → **[references/optimization-strategy.md](./references/optimization-strategy.md)** — read **first** for attention/FMHA/Gemma; then drill into the two source docs as needed
- **Optimizing GEMM/BMM/attention** (after TMA, or Triton 10–20% slower) → **[references/optimizing-reference.md](./references/optimizing-reference.md)** — EVEN_K fast path, transpose via pointer arithmetic, grid layout, autotune breadth, epilogue subtile; use these patterns during conversion and before perf sign-off (summarized in **optimization-strategy §2–§3**)
- **Gemma attention / GQA FMHA conversion** → **[references/optimization-strategy.md §4](./references/optimization-strategy.md#4-gemma-fmha--gemma_attention-conversion-checklist-mandatory)**
- **Blackwell optimization** (complex kernels with iterative algorithms, register pressure, loop unrolling) → **[references/optimizing-reference.md](./references/optimizing-reference.md) §9** — TMA descriptors, `loop_unroll_factor`, occupancy autotuning, TMEM-friendly block sizes, slab allocator, dual-path kernel design
- **⚠️ 10-50x REGRESSION** (catastrophic slowdown after conversion) → **[translations/workflow.md](./translations/workflow.md)** — section **CRITICAL PERFORMANCE PATTERNS (AVOID 10-50x REGRESSION)**
- **⚠️ Good perf on `transpose=True` only, collapse on `transpose=False`** (or opposite) → **[translations/advanced-patterns.md](./translations/advanced-patterns.md)** — §1 Dual layout flag; two `@triton.jit` kernels + `grid = lambda META: (... META["BLOCK_H"] ...)`
## Pre-flight Analysis (Run BEFORE converting)
```bash
# Count kernels (only main kernel gets @triton.jit, helpers stay plain def)
grep "@ct\.kernel" source.py | wc -l
# Check for patterns needing special handling
grep "ct\.transpose\|ct\.permute" source.py # → use tl.trans/tl.permute
grep "ct\.astype" source.py # → use .to(dtype)
grep "ct\.load\|ct\.store" source.py # → TMA for 2D+ (tl.make_tensor_descriptor), NOT raw tl.load(ptr+offs)
grep "ct\.launch" source.py # → bracket launch: kernel then [grid] then (args)
grep "ct\.Constant\|ct\.ConstInt" source.py # → tl.constexpr
grep "ct\.cdiv" source.py # → triton.cdiv (host) or Python (a+b-1)//b
grep "ct\.bid\|ct\.num_blocks" source.py # → tl.program_id/tl.num_programs
grep "1 << .*\.bit_length" source.py # → triton.next_power_of_2 if needed
grep "transpose\|transpose_v" source.py # → if hit, read translations/advanced-patterns.md (dual kernels + META grid)
```
## Conversion Checklist
Copy this checklist and track progress:
```
Conversion Progress:
[ ] Step 0 (attention / Gemma FMHA / GQA / soft cap / sliding window): Read [references/optimization-strategy.md](./references/optimization-strategy.md) and apply §4 checklist before inner-loop Triton
[ ] Step 1: Pre-flight — run grep commands above, note special patterns and 2D+ loads (→ TMA)
[ ] Step 2: Analyze source cuTile kernel (identify patterns, shapes, dtypes)
[ ] Step 3: Create Triton file with correct structure (see translations/file-structure.md)
[ ] Step 4: Convert kernel signature (tensor args → pointer args, Constant → constexpr)
[ ] Step 4b: TMA (MANDATORY for 2D+ loads) — use tl.make_tensor_descriptor for every 2D+ tile load/store; do NOT ship raw tl.load(ptr+offs,mask) for block-shaped access (see workflow.md § TMA OPTIMIZATION)
[ ] Step 5: Convert kernel body (apply gotchas table below + API mapping)
[ ] Step 6: Convert host wrapper (grid tuple/lambda, bracket-style launch: kernel, grid, then arguments; no ct.launch); call triton.set_allocator(alloc_fn) if using TMA
[ ] Step 7: Validate — run pytest or syntax check on Triton file
[ ] Step 8: Test — run pytest, verify X passed 0 failed
[ ] Step 9: If test fails → fix → re-validate → re-test (loop until green)
[ ] Step 10: Benchmark — run perf test, compare vs cuTile (see workflow.md § PERFORMANCE ANALYSIS)
[ ] Step 10b: If GEMM/BMM/attention and Triton >20% slower → walk [references/optimization-strategy.md](./references/optimization-strategy.md) §2–§3 then [references/optimizing-reference.md](./references/optimizing-reference.md) (EVEN_K, transpose, grid, autotune, epilogue subtile), then re-benchmark
[ ] Step 10c: If op has `transpose` / layout flag → read [translations/advanced-patterns.md](./translations/advanced-patterns.md); verify **separate kernels** per layout (not transpose-kernel + `tl.trans`); **autotuned** launches use `lambda META: (triton.cdiv(..., META["BLOCK_H"]), ...)` — no fixed `BLOCK_H`/`BLOCK_N` through `apply()` unless autotune is disabled
Post-conversion Verification (TMA is mandatory for 2D+ loads):
[ ] TMA: All 2D+ tile loads use tl.make_tensor_descriptor(...).load([...]); no raw ptr+mask for block-shaped 2D+ access (else 5x-20x regression)
[ ] Grid uses tuple or lambda (not 3-tuple required like cuTile)
[ ] Triton autotune added if cuTile op used kernel_configs/autotune (see workflow § PERFORMANCE ANALYSIS)
[ ] Host grid uses triton.cdiv where appropriate (not (a+b-1)//b only)
[ ] Pointer/offset indexing: Triton uses element offsets (ptr + offs), not block index in tl.load (or use TMA descriptor)
[ ] ct.astype(x, dtype) → x.to(dtype) in Triton
[ ] ct.mma(a, b, acc=acc) → tl.dot(a, b, acc) (no keyword in Triton)
[ ] Optional/None args: Triton allows None in kernel args if desired (cuTile required dummy+flag)
[ ] Masking applied when BLOCK_SIZE > actual dimension (same as cuTile); with TMA, masks can often be removed for full tiles
[ ] Reduction divisor uses actual_size, NOT BLOCK_SIZE
[ ] fp32/tf32: Triton defaults allow_tf32=True; match cuTile behavior if you had explicit tf32 cast
[ ] If any 2D+ load uses raw ptr+mask (exception only): document WHY TMA was not used
[ ] tl.assume() alignment hints added for strides and pointers
```
## Gotchas (Most Common Translation Errors) {#gotchas-most-common-translation-errors}
Comprehensive table of patterns that frequently break or regress when porting `@ct.kernel` to `@triton.jit` — *mma accumulator, type cast, grid, TMA usage, dtype handling, layout flags, batched matmul, etc.*
**See:** [references/gotchas.md](./references/gotchas.md) — read this BEFORE writing the Triton kernel.
## Performance Gotchas (10-50x Regression Risk) {#performance-gotchas-10-50x-regression-risk}
**⚠️ These cause CATASTROPHIC slowdowns. Check BEFORE benchmarking.**
Patterns and their impact: TMA vs raw ptr+mask (5-20×), autotune vs fixed tile sizes (2-3×), `broadcast_to + tl.dot` (10-50×), `extract_slice` chains (2-5×), and more.
**See:** [references/performance-gotchas.md](./references/performance-gotchas.md) — full regression-risk table.
**Full details:** [translations/workflow.md](./translations/workflow.md) — section **CRITICAL PERFORMANCE PATTERNS (AVOID 10-50x REGRESSION)**.
Full API mapping: [references/api-mapping.md](./references/api-mapping.md).
Triton math dtype (erf/erfc/exp/log/sqrt) and the "don't substitute erf with tanh" pattern: [references/debugging.md](./references/debugging.md) — section **Triton Math Function Dtype Requirements (CRITICAL)**.
## Optimization strategy (hub)
**File:** [references/optimization-strategy.md](./references/optimization-strategy.md)
Summarizes **[translations/advanced-patterns.md](./translations/advanced-patterns.md)** (layout flags, dual kernels, autotune+`META`, batched launch, Blackwell pointers) and **[references/optimizing-reference.md](./references/optimizing-reference.md)** (post-TMA micro-opts, §9) into **§1–§3** plus a **mandatory §4 Gemma FMHA checklist**.
**Rule:** For **attention / FMHA / Gemma-style** conversions, open **optimization-strategy** in the same session as **workflow** — do not rely on TMA alone for perf sign-off.
## Reference Documents {#reference-documents}
Read from **cuTile → Triton** perspective. Core files live in this skill under ``.
| Category | Document | Content |
|----------|----------|---------|
| **Strategy** | **[optimization-strategy.md](./references/optimization-strategy.md)** | **Ordered hub:** advanced-patterns + optimizing-reference; **§4 Gemma FMHA mandatory checklist** |
| **Workflows** | [translations/workflow.md](translations/workflow.md) | Standard c2t conversion (phases + checklist) |
| | [translations/file-structure.md](translations/file-structure.md) | Where to place Triton files when converting from cuTile |
| | **[translations/advanced-patterns.md](./translations/advanced-patterns.md)** | **Dual layout flags (transpose), autotune + `META` grid, MLA-style two kernels** |
| **API** | [api-mapping.md](./references/api-mapping.md) | cuTile → Triton mapping |
| | [optimizing-reference.md](./references/optimizing-reference.md) | **GEMM/BMM/attention optimizations** (EVEN_K, transpose, grid, autotune, epilogue subtile) |
| **Gotchas** | [gotchas.md](./references/gotchas.md) | **Common cuTile→Triton translation errors** (mma, dtype, grid, TMA, layout flags) |
| | [performance-gotchas.md](./references/performance-gotchas.md) | **10-50× regression-risk table** (TMA vs ptr+mask, broadcast_to, extract_slice chains, autotune) |
| **Testing & errors** | [references/debugging.md](./references/debugging.md) | **Triton runtime errors** (cudaErrorIllegalAddress, pointer type, stride overflow) |
## Worked Examples
Use **cutile_kernel.py as source** and **triton_kernel.py as target**:
| Example | Directory | Complexity |
|---------|-----------|------------|
| Vector Add | [examples/01_vector_add/](examples/01_vector_add/) | Basic |
| Softmax | [examples/02_softmax/](examples/02_softmax/) | Intermediate |
| LayerNorm | [examples/03_layernorm/](examples/03_layernorm/) | Intermediate |
| MatMul | [examples/04_matmul/](examples/04_matmul/) | Advanced |
| Attention | [examples/05_attention/](examples/05_attention/) | Advanced |
Read `cutile_kernel.py` first, then `triton_kernel.py`, to see the inverse mapping.
## ⚠️ MANDATORY COMPLETION CHECKLIST (DO NOT SKIP)
**A conversion is NOT COMPLETE until ALL items are checked. Copy and complete:**
```
MANDATORY COMPLETION GATES:
[ ] 1. CORRECTNESS: pytest passes with 0 failures
Command: python -m pytest {test_path} -k "test_op and triton" -vs --tb=short
Gate: "X passed, 0 failed"
[ ] 2. TMA OPTIMIZATION: All 2D+ tile loads use tl.make_tensor_descriptor
Verify: grep -n "tl.load.*mask" triton_file.py | wc -l # Should be 0 for 2D+ ops
Skip = 5-20x performance regression
[ ] 3. PERFORMANCE TEST: Triton within 20% of cuTile baseline
Command: python -m pytest {test_path} -k "test_perf" --print-record -v
OR: Run benchmark script: cd tests/benchmark && python bench_{op}.py
Gate: Triton TFLOPS >= 0.8 * CuTile TFLOPS
[ ] 4. PERFORMANCE COMPARISON RECORDED:
Document results:
| Config | Triton (TFLOPS) | CuTile (TFLOPS) | Ratio |
|--------|-----------------|-----------------|-------|
| [fill] | [fill] | [fill] | [fill]|
CONVERSION COMPLETE: All 4 gates passed? → YES / NO
```
**Why this matters:**
- Gate 1 catches functional bugs
- Gate 2 prevents catastrophic 5-20x regressions (most common mistake)
- Gate 3 validates that optimization was effective
- Gate 4 creates accountability record
**If any gate fails:** Fix and re-verify before declaring complete.
skill.oms.sig
{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json","verificationMaterial":{"x509CertificateChain":{"certificates":[{"rawBytes":"MIICgzCCAgmgAwIBAgIUKIyS7SxNteQIiWzK1dWj85E6520wCgYIKoZIzj0EAwMwVTELMAkGA1UEBhMCVVMxGzAZBgNVBAoMEk5WSURJQSBDb3Jwb3JhdGlvbjEpMCcGA1UEAwwgTlZJRElBIEFnZW50IENhcGFiaWxpdGllcyBJQ0EgMDEwHhcNMjYwNDAxMDAwMDAwWhcNMjgwNDIyMTUzMzA5WjBUMQswCQYDVQQGEwJVUzEbMBkGA1UECgwSTlZJRElBIENvcnBvcmF0aW9uMSgwJgYDVQQDDB9OVklESUEgQWdlbnQgU2tpbGxzIFNpZ25pbmcgMDAxMHYwEAYHKoZIzj0CAQYFK4EEACIDYgAEYoRM9bQl/dGlwSRNi6bTpIJUXH8Nv9GciP6LSflJYYMLCc296kpyuTSsk5ddbAWiDcFX3C/ydX3jwc+qCLYP6uHy9XphyLjOQ27Yb2J6rBLVtRBS1mgGco/Gr7fL6ODco4GaMIGXMB0GA1UdDgQWBBRQ/5ZW3nJ6lmo9SVk7I15o7UGmpTAfBgNVHSMEGDAWgBRPGpILxMBBleJSsBGjrMKsby1CgjAMBgNVHRMBAf8EAjAAMA4GA1UdDwEB/wQEAwIHgDA3BggrBgEFBQcBAQQrMCkwJwYIKwYBBQUHMAGGG2h0dHA6Ly9vY3NwLm5kaXMubnZpZGlhLmNvbTAKBggqhkjOPQQDAwNoADBlAjAUygu/GiOCIXrgGr4SmLgeEVDcEitfFUv7ALbvLVGVyMysB3mxmO/uInZfXzWcJZsCMQDxuoxj4ZmO30jhkPIcCxGFCOvnUsnfU3TfGcouYm4M6iRpbKvtVnHPiy4bi6pcKf0="},{"rawBytes":"MIICiDCCAg6gAwIBAgIUZsIuSv9NkpJCNqtYEfCouVv5BzowCgYIKoZIzj0EAwMwUTELMAkGA1UEBhMCVVMxGzAZBgNVBAoMEk5WSURJQSBDb3Jwb3JhdGlvbjElMCMGA1UEAwwcTlZJRElBIEFnZW50IENhcGFiaWxpdGllcyBDQTAgFw0yNjA0MDEwMDAwMDBaGA85OTk5MTIzMTIzNTk1OVowVTELMAkGA1UEBhMCVVMxGzAZBgNVBAoMEk5WSURJQSBDb3Jwb3JhdGlvbjEpMCcGA1UEAwwgTlZJRElBIEFnZW50IENhcGFiaWxpdGllcyBJQ0EgMDEwdjAQBgcqhkjOPQIBBgUrgQQAIgNiAASI72cR3ctKGg4VWnB3bNja6g1Z2PnOmFEopkPof+QeIcPk9rT+g9MjJnq51EQXL93a7C2GJ9J985G4o2V85VD7wJ1RaXhluHW2rf3y8bQGeAYaKMr5s/hUgn+M3/9WlWejgaAwgZ0wHQYDVR0OBBYEFE8akgvEwEGV4lKwEaOswqxvLUKCMB8GA1UdIwQYMBaAFItnoAjjfuCEUvzyvWyI2vOGvwPjMBIGA1UdEwEB/wQIMAYBAf8CAQAwDgYDVR0PAQH/BAQDAgEGMDcGCCsGAQUFBwEBBCswKTAnBggrBgEFBQcwAYYbaHR0cDovL29jc3AubmRpcy5udmlkaWEuY29tMAoGCCqGSM49BAMDA2gAMGUCMQCeIMMfAbyzPDacw2MxG+Yt1cikrJX/DVxiGfXuHmkkXn6VgSzE79+lkqDErpVO2gYCMCNEColOyvUvkzZGUEI1hQ3PfMgi3FIo9tHoBKMw4/wGBLFpu/0ubtmbBXM6/UMOEw=="},{"rawBytes":"MIICRTCCAcygAwIBAgIUeJdY3rV86EdvFmG7L8LJBsyQFYkwCgYIKoZIzj0EAwMwUTELMAkGA1UEBhMCVVMxGzAZBgNVBAoMEk5WSURJQSBDb3Jwb3JhdGlvbjElMCMGA1UEAwwcTlZJRElBIEFnZW50IENhcGFiaWxpdGllcyBDQTAgFw0yNjA0MDEwMDAwMDBaGA85OTk5MTIzMTIzNTk1OVowUTELMAkGA1UEBhMCVVMxGzAZBgNVBAoMEk5WSURJQSBDb3Jwb3JhdGlvbjElMCMGA1UEAwwcTlZJRElBIEFnZW50IENhcGFiaWxpdGllcyBDQTB2MBAGByqGSM49AgEGBSuBBAAiA2IABAYpiXCDjJ9NT2eSDhyHJVSw1Tbze18cGG2F/578oWvHxg23eQAhNRYdq88i1iOshZSO6C29doKui5Xpmo/7Ctw9Sx4PP2RzOmIuOLCuTdNtKcTRwi4GEsd5BAFvWj42M6NjMGEwHQYDVR0OBBYEFItnoAjjfuCEUvzyvWyI2vOGvwPjMB8GA1UdIwQYMBaAFItnoAjjfuCEUvzyvWyI2vOGvwPjMA8GA1UdEwEB/wQFMAMBAf8wDgYDVR0PAQH/BAQDAgEGMAoGCCqGSM49BAMDA2cAMGQCMCwtAjWLaNwgGWNCgdyNoTyvNhqWRECRJV2r3+7w8g0PL6NHLOsbkgE09BH95h8XlgIwTaQmbbUh2ChAJ5TA1wRiVDnCcvbzHlZl2jM2FcwQQZlk19LOAbyGMRixbu2Ww/rj"}]},"tlogEntries":[]},"dsseEnvelope":{"payload":"ewogICJfdHlwZSI6ICJodHRwczovL2luLXRvdG8uaW8vU3RhdGVtZW50L3YxIiwKICAic3ViamVjdCI6IFsKICAgIHsKICAgICAgIm5hbWUiOiAidGlsZWd5bS1jb252ZXJ0aW5nLWN1dGlsZS10by10cml0b24iLAogICAgICAiZGlnZXN0IjogewogICAgICAgICJzaGEyNTYiOiAiZmQxZjVkMTZjMDdmOWE5MGEwMWQxZjIzNDRjMGFkMDRlMjhjZDlkYThiYjNhNzJhYTEwYTM0OWIwNDVmMTc0MyIKICAgICAgfQogICAgfQogIF0sCiAgInByZWRpY2F0ZVR5cGUiOiAiaHR0cHM6Ly9tb2RlbF9zaWduaW5nL3NpZ25hdHVyZS92MS4wIiwKICAicHJlZGljYXRlIjogewogICAgInJlc291cmNlcyI6IFsKICAgICAgewogICAgICAgICJkaWdlc3QiOiAiYzNlNWIwY2FiZDY1OWQ4ZTFiZTJkN2ExY2Q1MWUyOGY2OGVhYjI3OGNhZDgwYTczNDMzMTU5MmYxMzAyNDJmYyIsCiAgICAgICAgImFsZ29yaXRobSI6ICJzaGEyNTYiLAogICAgICAgICJuYW1lIjogIkJFTkNITUFSSy5tZCIKICAgICAgfSwKICAgICAgewogICAgICAgICJkaWdlc3QiOiAiM2E4NTdiNzZhZmIzYjczNDM4MDUwNTg0N2Q4OWUyZDQ3Yjk5N2NiMTE4YjFjMzIyNmFiMGExODNkZjRkOTYwNiIsCiAgICAgICAgImFsZ29yaXRobSI6ICJzaGEyNTYiLAogICAgICAgICJuYW1lIjogIlNLSUxMLm1kIgogICAgICB9LAogICAgICB7CiAgICAgICAgImRpZ2VzdCI6ICIyYWM2MzM1Y2M3NzgzM2YxZWVjZTFmMDhhNWQ4Zjc2NmFkNTZiYjU3NzZhNDA3OWE0MjgxNjRkOGIwNWEwY2NmIiwKICAgICAgICAiYWxnb3JpdGhtIjogInNoYTI1NiIsCiAgICAgICAgIm5hbWUiOiAiZXZhbHMvZXZhbHMuanNvbiIKICAgICAgfSwKICAgICAgewogICAgICAgICJkaWdlc3QiOiAiNDIxMWNjYmEyNjJiNmFlNjNkYzI0NDEwMGI0MjhlOWMwYTFjYjE0YjAzZDVlYmMxOGE0OGQyYmY2YjUzMjY4ZiIsCiAgICAgICAgImFsZ29yaXRobSI6ICJzaGEyNTYiLAogICAgICAgICJuYW1lIjogImV4YW1wbGVzLzAxX3ZlY3Rvcl9hZGQvY3V0aWxlX2tlcm5lbC5weSIKICAgICAgfSwKICAgICAgewogICAgICAgICJkaWdlc3QiOiAiZWUxZjQ4NjYwNzFiMDViYmE3NTk2OWVmMjZiNjFiYmI3YjhiZjQ3NDE0YjZiMmQ0MGM1OTFjMmMxMmIwMTVlZiIsCiAgICAgICAgImFsZ29yaXRobSI6ICJzaGEyNTYiLAogICAgICAgICJuYW1lIjogImV4YW1wbGVzLzAxX3ZlY3Rvcl9hZGQvdHJpdG9uX2tlcm5lbC5weSIKICAgICAgfSwKICAgICAgewogICAgICAgICJkaWdlc3QiOiAiOTllOTJkOTE1NTZkZWI5YzhmZTIzMTRjMDdjNmE0ODNmNTEwYjAzOTVkYTg5N2Y1MGI0OTc4MWQ4YWU1MzBmYyIsCiAgICAgICAgImFsZ29yaXRobSI6ICJzaGEyNTYiLAogICAgICAgICJuYW1lIjogImV4YW1wbGVzLzAyX3NvZnRtYXgvY3V0aWxlX2tlcm5lbC5weSIKICAgICAgfSwKICAgICAgewogICAgICAgICJkaWdlc3QiOiAiM2VlYmZmMDJjMzY1NWM5NTgyZDFkODlhMmM1NDNlNDJmNzNkMjE1ZWY1MjI3NWFiOTQ1MjFiNTBkY2ExNjVjMCIsCiAgICAgICAgImFsZ29yaXRobSI6ICJzaGEyNTYiLAogICAgICAgICJuYW1lIjogImV4YW1wbGVzLzAyX3NvZnRtYXgvdHJpdG9uX2tlcm5lbC5weSIKICAgICAgfSwKICAgICAgewogICAgICAgICJkaWdlc3QiOiAiNDVhZTMwNzcyZTBiZDBiNjUwNjRjNWNjYmIwMGE0YWQ3MzA3MjJiNDZjZDJhM2FlYmVmNTFhYjA5NjVmNjE5ZSIsCiAgICAgICAgImFsZ29yaXRobSI6ICJzaGEyNTYiLAogICAgICAgICJuYW1lIjogImV4YW1wbGVzLzAzX2xheWVybm9ybS9jdXRpbGVfa2VybmVsLnB5IgogICAgICB9LAogICAgICB7CiAgICAgICAgImRpZ2VzdCI6ICJhNDQyNzRmOTg4NTBhZTFmYmU2MzZlZmU3NjVlMDlmZTYxNmE4YTA3OTQ2YzMwOGY2MjQyYTI4NmQ0YTJjZjYzIiwKICAgICAgICAiYWxnb3JpdGhtIjogInNoYTI1NiIsCiAgICAgICAgIm5hbWUiOiAiZXhhbXBsZXMvMDNfbGF5ZXJub3JtL3RyaXRvbl9rZXJuZWwucHkiCiAgICAgIH0sCiAgICAgIHsKICAgICAgICAiZGlnZXN0IjogImU4MzIwYzI1OTk4N2U3NzdjMTc3MjI4N2Q1NzhiNTg4NzJkZDQzMWEzYzA2YjI3OWI3YWQzZWE4ZTJlM2U0ZWQiLAogICAgICAgICJhbGdvcml0aG0iOiAic2hhMjU2IiwKICAgICAgICAibmFtZSI6ICJleGFtcGxlcy8wNF9tYXRtdWwvY3V0aWxlX2tlcm5lbC5weSIKICAgICAgfSwKICAgICAgewogICAgICAgICJkaWdlc3QiOiAiMGY2OGJmOWU2MjBjNjAyMjRjMDdhYmJlYWMyYzg0NmY5ZmY2MDZiM2M1NThkZTNlODRiZTE3ZDQ4MTQyNmY0YyIsCiAgICAgICAgImFsZ29yaXRobSI6ICJzaGEyNTYiLAogICAgICAgICJuYW1lIjogImV4YW1wbGVzLzA0X21hdG11bC90cml0b25fa2VybmVsLnB5IgogICAgICB9LAogICAgICB7CiAgICAgICAgImRpZ2VzdCI6ICJlMGQ2YTY4NDNkZThlMGVlMmFmYmM2MTI2MTk4MzA0MGZkZjU4MjlmODM5Nzg3MGIwZjFkNjE2ZGNmMjUxMWE3IiwKICAgICAgICAiYWxnb3JpdGhtIjogInNoYTI1NiIsCiAgICAgICAgIm5hbWUiOiAiZXhhbXBsZXMvMDVfYXR0ZW50aW9uL2N1dGlsZV9rZXJuZWwucHkiCiAgICAgIH0sCiAgICAgIHsKICAgICAgICAiZGlnZXN0IjogImQyMjgwNjI0NmY4YTIwZWYwZTU5MTRlOGViYThlM2Y0ZDNkNzY2ZDI2M2YyNTY0MWY4M2EyMzcwYWIwYTI4YWYiLAogICAgICAgICJhbGdvcml0aG0iOiAic2hhMjU2IiwKICAgICAgICAibmFtZSI6ICJleGFtcGxlcy8wNV9hdHRlbnRpb24vdHJpdG9uX2tlcm5lbC5weSIKICAgICAgfSwKICAgICAgewogICAgICAgICJkaWdlc3QiOiAiZDVhMDM5MGQ0ZmE5ZDkzMzA1YWVlNGZjMGFjNjMwYzk3ZWY3ZTJjMjg3ZjEzZGJiZGQwM2Y4MWMxMTE1ZTBiNSIsCiAgICAgICAgImFsZ29yaXRobSI6ICJzaGEyNTYiLAogICAgICAgICJuYW1lIjogInJlZmVyZW5jZXMvYXBpLW1hcHBpbmcubWQiCiAgICAgIH0sCiAgICAgIHsKICAgICAgICAiZGlnZXN0IjogIjRhNTM1MzMxYmVhNmRhZDFiMTQxYzRkZTljYjhkZDdjYWQ5ZDFiYTI0ZjA5MWVjNzhhMzNjYmQ4ZDA5NjFhNGMiLAogICAgICAgICJhbGdvcml0aG0iOiAic2hhMjU2IiwKICAgICAgICAibmFtZSI6ICJyZWZlcmVuY2VzL2RlYnVnZ2luZy5tZCIKICAgICAgfSwKICAgICAgewogICAgICAgICJkaWdlc3QiOiAiN2RiMmUwMzUyNGE2NDkwZmFiZTcyZDExZWFlZDI0MWNkYjliNGZhNjE5Y2IzNDI5MDhjMjg4ODQyMWQ0NmUzOCIsCiAgICAgICAgImFsZ29yaXRobSI6ICJzaGEyNTYiLAogICAgICAgICJuYW1lIjogInJlZmVyZW5jZXMvZ290Y2hhcy5tZCIKICAgICAgfSwKICAgICAgewogICAgICAgICJkaWdlc3QiOiAiOWZhYmE0NDY4ZTQ5ODhhYzY4MzAzNTUzNWFhYzE3MGJkYzUwM2IwNmM0NjViNTkyNDI1YmNmMzY2NWE2OTMxMyIsCiAgICAgICAgImFsZ29yaXRobSI6ICJzaGEyNTYiLAogICAgICAgICJuYW1lIjogInJlZmVyZW5jZXMvaGFybmVzcy1pbnRlZ3JhdGlvbi5tZCIKICAgICAgfSwKICAgICAgewogICAgICAgICJkaWdlc3QiOiAiZGQzYzlmMDBjNjhiODRjMDQwMDUzNDBlMzEyMjM0NjgyOTY2NzhiYTc5Njk5OWNhMGVmNTVjMDUzMzk3NDRiMiIsCiAgICAgICAgImFsZ29yaXRobSI6ICJzaGEyNTYiLAogICAgICAgICJuYW1lIjogInJlZmVyZW5jZXMvb3B0aW1pemF0aW9uLXN0cmF0ZWd5Lm1kIgogICAgICB9LAogICAgICB7CiAgICAgICAgImRpZ2VzdCI6ICJjZWUyMDY4OGIzZjc2NTFiMTg4MTNhYTg3ZGViZGY2ZmE1ZTAwMjJjZTFjZDA3ZTYzOTRhMjFhNzBlOGQyMjZhIiwKICAgICAgICAiYWxnb3JpdGhtIjogInNoYTI1NiIsCiAgICAgICAgIm5hbWUiOiAicmVmZXJlbmNlcy9vcHRpbWl6aW5nLXJlZmVyZW5jZS5tZCIKICAgICAgfSwKICAgICAgewogICAgICAgICJkaWdlc3QiOiAiNDFmNjNlYTZjM2NhNTQyY2E4YWJlN2RlYjE3ZmM0ZTM3YmE3MzgxOGY5NTc2MDEzZTQ3ZWU1MzA3ZDVlZTg0MiIsCiAgICAgICAgImFsZ29yaXRobSI6ICJzaGEyNTYiLAogICAgICAgICJuYW1lIjogInJlZmVyZW5jZXMvcGVyZm9ybWFuY2UtZ290Y2hhcy5tZCIKICAgICAgfSwKICAgICAgewogICAgICAgICJkaWdlc3QiOiAiYTMxZDIwODU4NzU5ZjAwMjI1NTAxMjdlNTI5YTJhNjlmNjMyOWEwOGU0YmRlMGY4NWE3YWI1NDNiZGFjNDZkZSIsCiAgICAgICAgImFsZ29yaXRobSI6ICJzaGEyNTYiLAogICAgICAgICJuYW1lIjogInNraWxsLWNhcmQubWQiCiAgICAgIH0sCiAgICAgIHsKICAgICAgICAiZGlnZXN0IjogImE3MjllY2QyZTJlOTE1YjZjNjZiOTdhMzMxMTBiNmY4YTkxYTgxNTQxMjZmNjUyZjZkOGU3ZjQwM2Y2NTJmN2EiLAogICAgICAgICJhbGdvcml0aG0iOiAic2hhMjU2IiwKICAgICAgICAibmFtZSI6ICJ0cmFuc2xhdGlvbnMvYWR2YW5jZWQtcGF0dGVybnMubWQiCiAgICAgIH0sCiAgICAgIHsKICAgICAgICAiZGlnZXN0IjogIjRmM2EyYjE3YWE2YTllMTM2NTczZDc3NjdiZGM5YTg1M2FhMGEzOTNmYTNhYzQyODQwMjdlNDU5ZTQ2ODFjY2QiLAogICAgICAgICJhbGdvcml0aG0iOiAic2hhMjU2IiwKICAgICAgICAibmFtZSI6ICJ0cmFuc2xhdGlvbnMvZmlsZS1zdHJ1Y3R1cmUubWQiCiAgICAgIH0sCiAgICAgIHsKICAgICAgICAiZGlnZXN0IjogImQ2MDkzMmU3NTJiOTM5YjdmYzQ2YzE1MzhiNjA4MTU4OTE3ODc3MWY1OTcyMzZhM2M1MGEwZjVhN2VhY2U4ZGYiLAogICAgICAgICJhbGdvcml0aG0iOiAic2hhMjU2IiwKICAgICAgICAibmFtZSI6ICJ0cmFuc2xhdGlvbnMvd29ya2Zsb3cubWQiCiAgICAgIH0KICAgIF0sCiAgICAic2VyaWFsaXphdGlvbiI6IHsKICAgICAgImlnbm9yZV9wYXRocyI6IFsKICAgICAgICAiLmdpdCIsCiAgICAgICAgIi5naXRodWIiLAogICAgICAgICIuZ2l0aWdub3JlIiwKICAgICAgICAiLmdpdGF0dHJpYnV0ZXMiCiAgICAgIF0sCiAgICAgICJoYXNoX3R5cGUiOiAic2hhMjU2IiwKICAgICAgIm1ldGhvZCI6ICJmaWxlcyIsCiAgICAgICJhbGxvd19zeW1saW5rcyI6IGZhbHNlCiAgICB9CiAgfQp9","payloadType":"application/vnd.in-toto+json","signatures":[{"sig":"MGUCMQC735+97Oe6H5noc68/vjRbBPwjo2AxOxfHyKJOUD6S3V++BNiFfddCPBWbDp/VgvgCMEDY6Quzy/Mn+FkzgDdHPVZTEZ/DAWSZZScoZTxTDjF9N5VZBWKzYuALKpVacPWVkA==","keyid":""}]}}translations/advanced-patterns.md
# Advanced Triton patterns (cuTile → Triton)
**Strategy hub:** For an ordered summary of this file plus [optimizing-reference.md](../references/optimizing-reference.md) (and a **mandatory Gemma FMHA checklist**), read **[references/optimization-strategy.md](../references/optimization-strategy.md)** first when converting attention or matching `gemma_attention`-class perf.
Patterns that are easy to miss in conversion but cause **large correctness or performance gaps**. Complements [translations/workflow.md](./workflow.md) (phases, TMA, autotune) and [SKILL.md](../SKILL.md) (checklist).
---
## 1. Dual layout flag (`transpose` / `transpose_v`): two kernels, not one + transposes {#dual-layout-flag}
**Applies to:** MLA-style decoding, attention variants, any op where the host exposes a boolean like `transpose=True/False` to match cuTile or framework layout.
### Failure mode (real case: MLA decoding)
| Benchmark slice | What you see |
|-----------------|--------------|
| `transpose=True` (e.g. small head counts in tests) | Converted kernel ≈ same ms as reference Triton |
| `transpose=False` | **Severe regression** (often **3–15×**, worse on fp8 / long `S_kv`) |
**Why:** The fast Triton baseline uses **different math and tensor layouts** per mode:
- **`transpose=False` (cache-friendly V along batch×seq):** Keep `qk` as **`[BLOCK_H, BLOCK_N]`**, softmax state **`l_prev` shape `[BLOCK_H]`**, `tl.max(qk, 1)`, `tl.sum(p, 1)`, and **`tl.dot(p, v, acc)`** with `p` already `[H, N]` and `v` `[N, D]` — one streamlined path.
- **`transpose=True` (V read with seq leading):** `qk` is **`[BLOCK_N, BLOCK_H]`** (e.g. `tl.dot(k, q)` with `q` as `[D, H]`), `l_prev` **`[BLOCK_N, BLOCK_H]`**, `V` needs a **separate TMA descriptor** (`shape=[S_kv, B, D]`, `strides=[stride_n, stride_b, 1]`, `block_shape=[BLOCK_N, 1, D]`) so the value load is not folded away or mis-coalesced — then **`tl.dot(v, p, acc)`** with transposed `v`/`p` (see `naive_absorb_mla_transpose` in the same file).
**Anti-pattern:** Implementing `transpose=False` by **reusing the transpose kernel’s structure** (e.g. forcing `l_prev` to `[BLOCK_N, BLOCK_H]`, transposing `qk` and `p` every KV block). That is **correctable** but **much slower** than the dedicated non-transpose kernel.
**Agent checklist:**
1. If cuTile or the PyTorch wrapper has **`transpose` (or equivalent)**, grep the Triton tree for **one** `@triton.jit` handling both — verify each branch uses the **same layout strategy** as the in-repo reference, not “transpose path + extra `tl.trans`”.
2. Add **two** `@triton.jit` kernels when the reference does (or split with `tl.constexpr` flags only if the compiler fully specializes — prefer two kernels for clarity and perf).
3. **Do not** pass fixed `BLOCK_H` / `BLOCK_N` from Python into `autograd.Function.forward` when the kernel is **`@triton.autotune`** — see §2.
---
## 2. Autotune + grid: `BLOCK_*` must come from `META`, not from the host
**Failure mode:** Host passes `BLOCK_H=16`, `BLOCK_N=128` into `apply()`, and launch uses `grid = (cdiv(heads, BLOCK_H), B)`. Autotune configs that use **64×128** or **128×128** never affect grid → **under-occupancy and wrong tuning** vs a baseline that uses:
```python
grid = lambda META: (triton.cdiv(num_head, META["BLOCK_H"]), B, 1)
# Launch kernel with that grid; pass tensor bases and META-sized BLOCK_* inside the kernel — not via apply().
```
**Rule:** For autotuned kernels, **`forward` should only pass dynamic shapes / strides / pointers**; tile sizes are **`tl.constexpr`** filled by autotune from `META`.
**Cross-check:** Run with `TRITON_PRINT_AUTOTUNING=1` and confirm the chosen config matches the problem key (`BLOCK_D`, `S_kv`, `EVEN_N`, etc.).
---
## 3. When to use host `TensorDescriptor` vs `tl.make_tensor_descriptor` in-kernel
TileGym’s `mla_decoding.py` uses **`tl.make_tensor_descriptor`** inside the JIT function with raw tensor bases + strides (plus `triton.set_allocator` for TMA metadata).
An alternate style (some LLM ports) builds **`triton.tools.tensor_descriptor.TensorDescriptor`** on the host when `sm >= 90` and passes descriptors into the kernel. That can be correct but:
- If you **freeze** tile sizes on the host descriptor to match **fixed** `BLOCK_H`/`BLOCK_N`, you **fight autotune** — either descriptors must be rebuilt per config (heavy) or you should use **in-kernel** `tl.make_tensor_descriptor` like `mla_decoding.py`.
**Preference for new conversions:** Align with **`mla_decoding.py`**: allocator + in-kernel descriptors + autotune + dual kernels for layout flags.
---
## 4. Quick diagnosis table
| Symptom | Likely cause | Action |
|---------|----------------|--------|
| Good perf when `transpose=True`, terrible when `transpose=False` | Single “transpose-style” kernel + extra transposes / wrong `l_prev` shape | Add dedicated non-transpose kernel; match `naive_absorb_mla` layout |
| Autotune “does nothing” / grid too large | `BLOCK_H`/`BLOCK_N` from Python, not `META` | `lambda META: (cdiv(..., META["BLOCK_H"]), ...)`; drop fixed blocks from `apply()` |
| `transpose=True` wrong or slow on V | V shares K descriptor | Separate **V_desc** with seq-leading shape/strides (see `mla_decoding.py` comment on optimization-out) |
---
## 5. Batched kernel launch for multi-tensor ops
**Applies to:** `cat`, `stack`, `split`, multi-tensor copy/scatter operations.
When converting a cuTile kernel that processes multiple tensors (e.g., concatenation), **do not** naively launch one Triton kernel per tensor. Instead, batch up to 4 tensors per launch using a 2D grid:
```python
# Grid: (blocks_over_elements, num_tensors_in_batch)
grid = (triton.cdiv(max_elements, BLOCK), min(4, num_remaining_tensors))
# Kernel uses program_id(1) to select which tensor to process
pid_y = tl.program_id(1)
if pid_y == 0:
in_ptr, size = in_ptr_a, size_a
elif pid_y == 1:
in_ptr, size = in_ptr_b, size_b
# ...
```
**Impact:** 2–4× speedup for small-to-medium tensors (LLM KV-cache concatenation). Kernel launch overhead (~5–10µs per launch) dominates for small tensors; batching amortizes this cost.
**Full pattern and code:** See **§8 Batched Kernel Launch** in [references/optimizing-reference.md](../references/optimizing-reference.md).
---
## 6. Blackwell Optimization Patterns
For complex kernels targeting Blackwell (sm_100+), additional optimization patterns can yield **2–10× speedups**. These are documented in **[references/optimizing-reference.md §9](../references/optimizing-reference.md)** and include:
| Pattern | Impact | When to Use |
|---------|--------|-------------|
| TMA Descriptors | 20–50% | All 2D+ loads on triton DSL |
| Loop Unroll Control | 2–5× | Iterative algorithms (matrix inversion, recurrence) |
| Occupancy Autotuning | 1.5–3× | All kernels on Blackwell |
| TMEM-Friendly Blocks | 1.5–2× | Kernels with multiple `tl.dot` operations |
| Slab Allocator | 10–30% | Kernels with multiple intermediate buffers |
| Dual-Path Design | — | Cross-platform support (triton DSL + OpenAI Triton) |
**Key insight:** On Blackwell, register pressure from loop unrolling is often the #1 performance killer. Use `tl.range(..., loop_unroll_factor=1)` for loops with >16 iterations and complex bodies.
---
## 7. Reference implementations
- **MLA decoding with transpose flag:** Implement two separate `@triton.jit` kernels (`transpose=False` path and `transpose=True` path). The `transpose=False` path: `qk` as `[BLOCK_H, BLOCK_N]`, `l_prev` as `[BLOCK_H]`, `tl.dot(p, v, acc)`. The `transpose=True` path: separate V TMA descriptor with layout matching the transposed access pattern. See §1 above and [examples/05_attention/](../examples/05_attention/) for a worked example.
- **Cat (batched):** Original TileGym `cat` implementation (adapted from [FlagGems](https://github.com/FlagOpen/FlagGems)) demonstrates batched multi-tensor kernel launch. See `cat_fwd_tensor_list_contiguous_any_dim_up_to_four_inputs_kernel` pattern or §8 of [references/optimizing-reference.md](../references/optimizing-reference.md).
translations/file-structure.md
# File Structure & Registration (cuTile → Triton)
Where to place Triton files when converting from cuTile. Inverse of the Triton→cuTile layout.
---
## Directory Structure
### Standard Mode (Directory-Based)
When converting **from** cuTile **to** Triton, create Triton files under the `triton/` mirror of the cuTile path.
There are two top-level layouts depending on whether the op is a first-party TileGym op or an external-framework suite:
```
TileGym/
├── src/tilegym/
│ ├── ops/ # First-party TileGym ops (fmha, matmul, softmax, …)
│ │ ├── triton/
│ │ │ ├── add.py # Triton (target of c2t conversion)
│ │ │ ├── softmax.py
│ │ │ └── layer_norm.py
│ │ └── cutile/
│ │ ├── add.py # Existing cuTile (source)
│ │ ├── softmax.py
│ │ └── layer_norm.py
│ └── suites/ # External-framework suites
│ └── <framework>/ # e.g. unsloth, flashinfer
│ ├── triton/
│ │ └── <op>.py # Triton conversion target
│ └── cutile/
│ └── <op>.py # Existing cuTile source
└── tests/
├── ops/
│ └── test_<op>.py # Tests for ops/ kernels
└── suites/
└── <framework>/
└── test_<op>.py # Tests for suites/ kernels
```
**Path derivation:**
```bash
# ops/ kernel: swap /cutile/ → /triton/
CUTILE_PATH="src/tilegym/ops/cutile/softmax.py"
TRITON_PATH="${CUTILE_PATH//\/cutile\//\/triton\/}"
# → src/tilegym/ops/triton/softmax.py
# suites/ kernel: same rule
CUTILE_PATH="src/tilegym/suites/<framework>/cutile/<op>.py"
TRITON_PATH="${CUTILE_PATH//\/cutile\//\/triton\/}"
# → src/tilegym/suites/<framework>/triton/<op>.py
mkdir -p $(dirname $TRITON_PATH)
```
---
## Registration Patterns
Same as the Triton→cuTile skill: register implementations by backend.
```python
from tilegym.backend import register_impl
@register_impl("op_name", backend="triton")
def op_triton(...):
...
@register_impl("op_name", backend="cutile")
def op_cutile(...):
...
```
Tests typically parametrize over `backend=["triton", "cutile"]` so both are exercised.
---
## Multi-Agent / Two-Step Workflow
| Step | Purpose |
|------|---------|
| **Step 1: Convert** | cuTile → Triton conversion |
| **Step 2: Perf** | Performance testing & comparison (Triton vs cuTile) |
Default: run both steps unless the user asks only for conversion or only for perf.
---
## Common Pitfalls
- **Wrong path:** Putting the new Triton file in `cutile/` instead of `triton/`.
- **Leftover cuTile imports:** Removing `import cuda.tile as ct` and all `ct.*` usage in the new Triton file; use `import triton.language as tl` and `triton.jit` only.
- **Launch style:** Using `ct.launch(stream, grid, kernel, args)` in the Triton host; must use `<code>kernel［grid］(launch_args)</code>` and `triton.cdiv` for grid.
translations/workflow.md
# cuTile to Triton Conversion Workflow
**Guide for converting cuTile kernels to Triton with the same rigor as the inverse (Triton→cuTile).**
---
## 🚀 TODO WORKFLOW (MANDATORY - CREATE IMMEDIATELY)
**Upon starting a cuTile→Triton conversion task, IMMEDIATELY create this todo list using `todowrite`:**
```
todowrite([
{ id: "c2t-1", content: "[Optional] Analyze test coverage - verify cuTile tests pass, identify edge cases", status: "pending", priority: "medium" },
{ id: "c2t-2", content: "Convert cuTile → Triton - apply API mapping, generate Triton file", status: "pending", priority: "high" },
{ id: "c2t-3", content: "Test correctness - run python -m pytest -k 'triton', fix errors (max 5 attempts)", status: "pending", priority: "high" },
{ id: "c2t-4", content: "TMA optimization (MANDATORY) - replace ALL 2D+ raw ptr+mask with tl.make_tensor_descriptor; add tl.assume() alignment hints; add autotuning; if transpose/layout flag, two kernels + lambda META grid (advanced-patterns.md); skip = 5-20x regression", status: "pending", priority: "high" },
{ id: "c2t-5", content: "Performance test - run pytest -k 'test_perf' --print-record, compare vs cuTile, optimize if >20% slower", status: "pending", priority: "high" }
])
```
### Workflow Execution Rules
| Rule | Description |
|------|-------------|
| **Auto-proceed** | Move to next phase automatically after success - NO user confirmation needed |
| **Single focus** | Only ONE todo `in_progress` at a time |
| **Immediate update** | Mark `completed` immediately after phase passes |
| **Skip c2t-1** | If user says "skip test checker" OR Triton already exists and tests pass |
| **Stop conditions** | Only stop on: (1) critical failure after 5 attempts, (2) all phases complete |
### Phase → Todo Mapping
| Phase | Todo ID | Success Criteria | Next Action |
|-------|---------|------------------|-------------|
| Test Coverage | c2t-1 | cuTile tests pass, edge cases identified | → c2t-2 |
| Convert | c2t-2 | Triton file created, no syntax errors | → c2t-3 |
| Test | c2t-3 | `X passed, 0 failed` | → c2t-4 |
| TMA Optimize | c2t-4 | **MANDATORY:** TMA descriptors for ALL 2D+ loads (no raw ptr+mask for block tiles), alignment hints, autotuning added | → c2t-5 |
| Performance | c2t-5 | `pytest -k test_perf` run, Triton within 20% of cuTile | → DONE |
**DO NOT ask "should I proceed?" - execute the full workflow end-to-end.**
---
## RATIONALE: Key Thresholds
**Why these values? (Aligned with Triton→cuTile skill.)**
| Threshold | Value | Rationale |
|-----------|-------|-----------|
| Max fix attempts | 5 | Most errors resolve in 1–2; after 5, likely needs human insight |
| Perf threshold | >20% slower | Below 20%, measurement noise masks real differences (5–15% variance) |
| float32 rtol/atol | 1e-3 | 7 sig digits; allows 4 digits agreement |
| float16 rtol/atol | 1e-2 | 3–4 sig digits; matches precision limit |
| bfloat16 rtol/atol | 1e-2 | Same as float16 |
**Relaxed tolerances:** Use 2× for reductions, transcendentals, chained ops.
---
## VALIDATION LOOP (MANDATORY)
**NEVER proceed until tests pass. This pattern applies to ALL test phases.**
```
┌─────────────────────────────────────────────────────────┐
│ VALIDATION LOOP │
│ │
│ ┌─────────┐ ┌─────────┐ ┌─────────┐ │
│ │ RUN │────▶│ CHECK │────▶│ PASS? │ │
│ │ TEST │ │ OUTPUT │ │ │ │
│ └─────────┘ └─────────┘ └────┬────┘ │
│ ▲ │ │
│ │ ┌────────────────┼───────────┐ │
│ │ │ │ │ │
│ │ ▼ ▼ │ │
│ ┌─────────┐ ┌─────────┐ ┌─────────┐ │ │
│ │ FIX │◀───│ NO │ │ YES │──────┘ │
│ │ ERROR │ │(attempt │ │ DONE │ │
│ └─────────┘ │ < 5) │ └─────────┘ │
│ └─────────┘ │
└─────────────────────────────────────────────────────────┘
```
**Validation Checklist** (copy for each attempt):
```
- [ ] Attempt #__: Run test command
- [ ] Check: `X passed, 0 failed`?
- [ ] No `FAILED tests/...` markers?
- [ ] No exceptions (syntax, shape, dtype)?
- [ ] If FAIL: identify error → fix → increment attempt
- [ ] If attempt >= 5: STOP, escalate to user
```
**CRITICAL:** Do NOT proceed to next phase until loop completes successfully.
---
## EXACT TEST COMMANDS (LOW FREEDOM)
**DO NOT MODIFY these commands. Flags are validated for correct output handling.**
### Correctness Test (Triton)
```bash
# DO NOT MODIFY
python -m pytest {test_path} -k "test_op and triton" -vs --tb=short 2>&1 | tail -100
```
### Performance Test
```bash
# DO NOT MODIFY
python -m pytest {test_path} -k "test_perf" --print-record -v 2>&1 | tail -50
```
### Triton Debug / Profiling
```bash
TRITON_INTERPRET=1 python script.py
TRITON_PRINT_AUTOTUNING=1 python script.py
TILEIR_DUMP_DIR=/tmp/dumping/triton python -m pytest {test_path} -k "test_op and triton" --timeout=120
```
**Flag rationale:** `-vs` (verbose + no capture), `--tb=short` (concise tracebacks), `2>&1 | tail -100` (capture stderr, limit for context).
---
## TABLE OF CONTENTS
[🚀 TODO WORKFLOW (MANDATORY - CREATE IMMEDIATELY)](#-todo-workflow-mandatory---create-immediately)
[RATIONALE: Key Thresholds](#rationale-key-thresholds)
[VALIDATION LOOP (MANDATORY)](#validation-loop-mandatory)
[EXACT TEST COMMANDS (LOW FREEDOM)](#exact-test-commands-low-freedom)
[MODE SELECTION](#mode-selection)
[COMMAND CHEAT SHEET](#command-cheat-sheet)
[TMA OPTIMIZATION (Phase c2t-4)](#tma-optimization-phase-c2t-4)
[PERFORMANCE ANALYSIS (Phase c2t-5)](#performance-analysis-phase-c2t-5)
[CRITICAL PERFORMANCE PATTERNS](#critical-performance-patterns-avoid-10-50x-regression)
[DUAL LAYOUT FLAG + AUTOTUNE GRID (MLA-style)](#dual-layout-flag--autotune-grid-mla-style)
[MEMORY LAYOUT PATTERNS](#memory-layout-patterns-avoid-50-150-regression)
[Standard Conversion Steps](#standard-conversion-steps)
[Quick Reference: cuTile → Triton](#quick-reference-cutile--triton)
---
## MODE SELECTION
### Standard Mode (cuTile → Triton)
**Use when:** Converting existing TileGym cuTile operators to Triton (e.g. for portability or comparison).
**Path convention:** `/cutile/` → `/triton/`
```bash
TRITON_PATH="${CUTILE_PATH//\/cutile\//\/triton\/}"
mkdir -p $(dirname $TRITON_PATH)
```
---
## COMMAND CHEAT SHEET
```bash
# Standard mode path derivation
TRITON_PATH="${CUTILE_PATH//\/cutile\//\/triton\/}"
mkdir -p $(dirname $TRITON_PATH)
# Correctness testing
python -m pytest {test_path} -k "test_op and triton" -vs
# Performance testing (Triton vs cuTile)
python -m pytest {test_path} -k "test_perf and (triton or cutile)" --print-record -v
# Triton profiling / autotune visibility
TRITON_PRINT_AUTOTUNING=1 python -m pytest {test_path} -k "test_op and triton"
triton.testing.do_bench(lambda: kernel［grid］(launch_args)) # In script: measure ms
```
---
## TMA OPTIMIZATION (Phase c2t-4) {#tma-optimization-phase-c2t-4}
**This phase is MANDATORY. Do not skip.** Raw pointer + mask for 2D+ tile loads are **5-20x (500%-2000%) slower** than TMA on Hopper/Blackwell. Converted kernels that use only `tl.load(ptr+offs, mask=...)` for block-shaped 2D+ access will show severe regression until TMA is added.
### When TMA is Required
- **2D+ tile loads** (GEMM, attention, convolution, any load with shape like `(BLOCK_M, BLOCK_K)`) → **ALWAYS use TMA** (`tl.make_tensor_descriptor` + `.load([...])`)
- **1D loads** (single contiguous block, elementwise, simple reductions) → raw pointer OK
### TMA Conversion Pattern
```python
# BEFORE (cuTile ct.load) - already uses TMA internally
tile = ct.load(arr, index=(bid_m, bid_k), shape=(BLOCK_M, BLOCK_K))
# WRONG (naive Triton conversion) - 10-20x SLOWER
offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
offs_k = tl.arange(0, BLOCK_K)
tile = tl.load(ptr + offs_m[:, None] * stride_m + offs_k[None, :], mask=mask)
# CORRECT (TMA tensor descriptor) - FAST
desc = tl.make_tensor_descriptor(
base=ptr,
shape=[M, K],
strides=[stride_m, 1],
block_shape=[BLOCK_M, BLOCK_K],
)
tile = desc.load([pid_m * BLOCK_M, pid_k * BLOCK_K])
```
### TMA Setup (Required Once)
```python
from typing import Optional
def alloc_fn(size: int, alignment: int, stream: Optional[int]):
return torch.empty(size, device="cuda", dtype=torch.int8)
triton.set_allocator(alloc_fn)
```
### TMA Checklist (Run BEFORE moving to c2t-5) {#tma-checklist-run-before-moving-to-c2t-5}
```
TMA Optimization Verification:
[ ] All 2D+ tile loads use tl.make_tensor_descriptor
[ ] Masks removed (TMA handles bounds automatically)
[ ] tl.assume() alignment hints added for strides: tl.assume(stride % 8 == 0)
[ ] tl.assume() alignment hints added for pointers: tl.assume(ptr.to(tl.int64) % 16 == 0)
[ ] Autotuning added with GPU-specific configs (see PERFORMANCE ANALYSIS)
```
---
## PERFORMANCE ANALYSIS (Phase c2t-5) {#performance-analysis-phase-c2t-5}
When Triton is **>20% slower** than cuTile, follow this systematic flow.
### Step 1: Benchmark
```bash
python -m pytest {test_path} -k "test_perf and (triton or cutile)" --print-record -v
```
If Triton is within 20% of cuTile → done. If not → continue.
### Step 2: Match Tile Sizes and Grid
- **Constexpr block sizes** — Use the same BLOCK_M, BLOCK_N, BLOCK_K (or BLOCK_SIZE) as the cuTile kernel as `tl.constexpr` in Triton.
- **Grid** — Use `triton.cdiv(M, BLOCK_M)` etc. so grid shape matches cuTile's `(n_m, n_n, 1)`-style launch.
### Step 3: Memory Access Pattern
- **cuTile** uses block index in `ct.load(arr, index=(...), shape=(...))`.
- **Triton** uses element offset: `offs = pid * BLOCK + tl.arange(0, BLOCK)` (and strides for 2D+). Ensure `tl.load(ptr + offs, mask=...)` or `tl.make_block_ptr` + load matches the same coalescing and alignment.
- **Cache hints** — For memory-bound kernels, try `tl.load(..., cache_modifier=".cg")` if appropriate (see performance-model.md).
### Step 4: Triton Autotuning
If the cuTile op uses autotune (tile sizes, occupancy), add `triton.autotune` to the Triton kernel so Triton can search the same space:
```python
@triton.autotune(
configs=[
triton.Config({'BLOCK_M': 128, 'BLOCK_N': 256, 'BLOCK_K': 64}, num_stages=3, num_warps=8),
triton.Config({'BLOCK_M': 64, 'BLOCK_N': 256, 'BLOCK_K': 32}, num_stages=4, num_warps=4),
# ... match or expand cuTile config space
],
key=['M', 'N', 'K'], # Retune when these change
)
@triton.jit
def kernel(a_ptr, b_ptr, c_ptr, M, N, K, ...,
BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr):
...
```
**Triton autotune knobs:** `BLOCK_*` (tile dimensions), `num_warps` (1, 2, 4, 8, 16, 32), `num_stages` (2–5 typical).
### Step 5: Profiling and Bottleneck
- Run with `TRITON_PRINT_AUTOTUNING=1` to see which config is chosen.
- Use `triton.testing.do_bench(lambda: kernel［grid］(launch_args))` for stable timings.
- For deeper analysis: Nsight Compute, or dump Triton IR (`TILEIR_DUMP_DIR=/tmp/dumping/triton`) and compare with cuTile IR (see [references/debugging.md](../references/debugging.md)) if the gap remains unexplained.
### Step 6: tf32 and Dtypes
- Triton defaults `allow_tf32=True` for `tl.dot`. If cuTile used an explicit tf32 cast before `ct.mma`, behavior should already match; use `allow_tf32=False` only if you need strict fp32.
**Bottleneck checklist:** algorithmic/grid → memory access → occupancy/warps → micro-optimizations.
---
## Standard Conversion Steps
| Step | Action |
|------|--------|
| 1 | Pre-flight: grep for ct.kernel, ct.load, ct.store, ct.launch, ct.Constant, ct.astype (see SKILL.md); if 2D+ ct.load → plan TMA |
| 2 | Create Triton file under triton/ mirror path (see file-structure.md) |
| 3 | Convert signature: tensor args → pointers + strides/shapes; ct.Constant[int] → tl.constexpr |
| 4 | Convert body: ct.load/ct.store → **for 2D+ block loads use tl.make_tensor_descriptor + .load([...]) (TMA)**; for 1D use tl.load(ptr+offs, mask=...); ct.astype → .to(); ct.mma(..., acc=acc) → tl.dot(..., acc) |
| 4b | **TMA (mandatory):** Replace any 2D+ raw ptr+mask loads with TMA; add triton.set_allocator(alloc_fn) in host; add tl.assume() alignment (see TMA OPTIMIZATION above) |
| 5 | Convert host: grid = (n,) or lambda meta: (...); <code>kernel［grid］(launch_args)</code>; use triton.cdiv for grid size |
| 6 | Test and compare numerically with compare_outputs.py |
---
## Quick Reference: cuTile → Triton
| cuTile | Triton |
|--------|--------|
| `@ct.kernel` | `@triton.jit` |
| `import cuda.tile as ct` | `import triton.language as tl` |
| `ct.bid(axis)` | `tl.program_id(axis)` |
| `ct.num_blocks(axis)` | `tl.num_programs(axis)` |
| `ct.arange(N, dtype=ct.int32)` | `tl.arange(0, N)` |
| `ct.load(arr, index=(bid,), shape=(BLOCK,))` (1D) | `tl.load(ptr + offs, mask=...)` (offs = bid * BLOCK + arange) |
| `ct.load(arr, index=(i,j), shape=(BM,BK))` (2D+) | **TMA:** `tl.make_tensor_descriptor(...).load([...])` — do NOT use raw tl.load(ptr+offs) (5-20x regression) |
| `ct.astype(x, dtype)` | `x.to(dtype)` |
| `ct.mma(a, b, acc=acc)` | `tl.dot(a, b, acc)` |
| `ct.Constant[int]` | `tl.constexpr` |
| `grid = (n, 1, 1)`; `ct.launch(stream, grid, kernel, args)` | `grid = (triton.cdiv(n, BLOCK),)`; <code>kernel［grid］(launch_args)</code> |
| `(a + b - 1) // b` (host) | `triton.cdiv(a, b)` |
Full cuTile → Triton API mapping: **[references/api-mapping.md](../references/api-mapping.md)**.
---
## CRITICAL PERFORMANCE PATTERNS (AVOID 10-50x REGRESSION) {#critical-performance-patterns-avoid-10-50x-regression}
**Case study: Group GEMM conversion showed 20-50x regression. Root causes and fixes below.**
### Performance Killer #1: Raw Pointer Arithmetic vs TMA Tensor Descriptors
**Impact: 5-20x slowdown (500%-2000% regression) — most common cause of conversion regression**
TMA (Tensor Memory Accelerator) enables async bulk memory transfers on Hopper/Blackwell. Raw pointer arithmetic with masks for 2D+ tile loads is dramatically slower. **Always use TMA for 2D+ block-shaped loads; do not ship conversion with raw ptr+mask for GEMM/attention/conv tiles.**
```python
# SLOW (10-20x regression) - Raw pointer + masks
offs_m = tile_m * TILE_M + tl.arange(0, TILE_M)
offs_k = tl.arange(0, TILE_K)
a_ptrs = A_ptr + (offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak)
a_mask = (offs_m[:, None] < M) & (offs_k[None, :] < K)
a = tl.load(a_ptrs, mask=a_mask, other=0.0) # Masked load = SLOW
# FAST - TMA tensor descriptors
a_desc = tl.make_tensor_descriptor(
base=a_base_ptr,
shape=[m, k],
strides=[lda, 1],
block_shape=[BLOCK_M, BLOCK_K],
)
a = a_desc.load([offset_am, kk * BLOCK_K]) # Bulk TMA load = FAST
```
**When to use TMA:**
- Blackwell (sm_120) or Hopper (sm_90) GPUs
- Matrix operations (GEMM, attention, convolution)
- Any kernel with 2D+ tile loads
**TMA setup requirements:**
```python
from typing import Optional
# TMA allocator (required once per kernel launch context)
def alloc_fn(size: int, alignment: int, stream: Optional[int]):
return torch.empty(size, device="cuda", dtype=torch.int8)
triton.set_allocator(alloc_fn)
```
### Performance Killer #2: Inefficient Group/Batch Iteration
**Impact: 2-5x slowdown**
For grouped operations (group GEMM, batched attention), how you find which group a tile belongs to matters.
```python
# SLOW (2-5x regression) - Linear search ALL groups per tile
@triton.jit
def _find_group_id(tile_idx, problem_offsets, num_groups: tl.constexpr):
group_id = num_groups - 1
for g in range(num_groups): # Scans ALL groups every time
offset_start = tl.load(problem_offsets + g)
offset_end = tl.load(problem_offsets + g + 1)
is_in_group = (tile_idx >= offset_start) & (tile_idx < offset_end)
group_id = tl.where(is_in_group, group_id, g) # Conditional per group
return group_id
# FAST - Natural while-loop advancement
@triton.jit
def kernel(...):
tile_idx = tl.program_id(0)
last_problem_end = 0
for g in range(group_size):
# ... load group dimensions ...
num_tiles = num_m_tiles * num_n_tiles
# Only process tiles belonging to this group
while tile_idx >= last_problem_end and tile_idx < last_problem_end + num_tiles:
# Process tile
tile_idx += num_programs # Persistent scheduling
last_problem_end += num_tiles # Advance boundary
```
### Performance Killer #3: Missing Autotuning
**Impact: 2-3x slowdown**
Fixed tile sizes vs architecture-optimized configurations.
```python
# SLOW - Fixed sizes, no tuning
TILE_M, TILE_N, TILE_K = 128, 128, 64 # May be wrong for your GPU/problem
# FAST - Autotuning with GPU-specific configs
def _get_configs():
gpu_cap = torch.cuda.get_device_capability()
if gpu_cap in [(12, 0), (12, 1)]: # Blackwell
return [
triton.Config({"BLOCK_M": BM, "BLOCK_N": BN, "BLOCK_K": BK})
for BM in [64, 128]
for BN in [64, 128, 256]
for BK in [64, 128]
]
elif gpu_cap == (9, 0): # Hopper
return [
triton.Config({"BLOCK_M": BM, "BLOCK_N": BN, "BLOCK_K": BK},
num_stages=s, num_warps=w)
for BM in [128, 256]
for BN in [128, 256]
for BK in [64, 128]
for s in [4, 5]
for w in [8]
]
# ... other architectures
@triton.autotune(configs=_get_configs(), key=["group_size", "dtype"])
@triton.jit
def kernel(...):
...
```
### Performance Killer #4: Missing Alignment Hints
**Impact: 1.5-2x slowdown**
Triton compiler can optimize better with alignment guarantees.
```python
# SLOW - No hints, compiler assumes worst case
lda = tl.load(group_strides + g * 3)
a_base_ptr = tl.load(group_a_ptrs + g).to(tl.pointer_type(dtype))
# FAST - Alignment hints enable compiler optimizations
lda = tl.load(group_strides + g * 3)
tl.assume(lda % 8 == 0) # Stride is 8-element aligned
a_base_ptr = tl.load(group_a_ptrs + g).to(tl.pointer_type(dtype))
tl.assume(a_base_ptr.to(tl.int64) % 16 == 0) # 16-byte aligned pointer
```
### Performance Killer #5: Unnecessary Masking
**Impact: 1.2-1.5x slowdown**
Masks add predication overhead. TMA handles boundaries automatically.
```python
# SLOW - Masks on every operation
a = tl.load(a_ptrs, mask=a_mask, other=0.0)
b = tl.load(b_ptrs, mask=b_mask, other=0.0)
tl.store(c_ptrs, c, mask=c_mask)
# FAST - TMA handles bounds, no masks needed
a = a_desc.load([offset_am, kk * BLOCK_K])
b = b_desc.load([kk * BLOCK_K, offset_bn])
c_desc.store([offset_cm, offset_cn], c)
```
### Performance Killer #6: K-loop Offset Recalculation
**Impact: 1.1-1.2x slowdown**
For GEMM K-loops, avoid recalculating full offsets each iteration.
```python
# SLOW - Recalculate every iteration
for k_tile in range(num_k_tiles):
k_offset = k_tile * TILE_K
a_ptrs = A_ptr + (offs_m[:, None] * stride_am + (k_offset + offs_k[None, :]) * stride_ak)
# ... full offset calculation each time
# FAST - Increment pointers (for non-TMA kernels)
a_ptrs = a_ptr + offs_am[:, None] * lda + offs_k[None, :]
for kk in range(0, tl.cdiv(k, BLOCK_K)):
tl.multiple_of(a_ptrs, [16, 16]) # Pipeline hint
a = tl.load(a_ptrs)
a_ptrs += BLOCK_K # Simple increment
```
### Performance Checklist (Run BEFORE declaring conversion complete)
```
Performance Verification:
[ ] TMA tensor descriptors used for 2D+ tile loads (Hopper/Blackwell)
[ ] Autotuning added with GPU-specific configs
[ ] Group/batch iteration uses natural loop advancement (not linear search)
[ ] tl.assume() hints added for stride and pointer alignment
[ ] Masks removed where TMA handles boundaries
[ ] K-loop uses pointer increment or TMA offset (not full recalc)
[ ] Pipeline hints (tl.multiple_of) added for non-TMA loads
[ ] num_stages and num_warps tuned in autotune configs
[ ] Benchmark shows <20% regression vs original cuTile
[ ] Memory layout matches original (transpose + contiguous, NOT 5D reshape)
[ ] Dimension params marked as tl.constexpr (bs, hd, n_heads, pad_*)
[ ] No unnecessary tensor clones (transpose + contiguous suffices)
[ ] If `transpose`/layout flag: two kernels + `grid = lambda META: (triton.cdiv(..., META["BLOCK_H"]), ...)` — see [advanced-patterns.md](advanced-patterns.md)
```
---
## DUAL LAYOUT FLAG + AUTOTUNE GRID (MLA-style)
**When:** cuTile or the public API exposes **`transpose`** (or equivalent) and perf tests show **one mode fast, the other 3–15× slow**.
**Do not** implement the non-transpose case by reusing the transpose kernel with extra **`tl.trans`** on `qk` / `p` and a **`[BLOCK_N, BLOCK_H]`** softmax state unless that matches a proven-fast reference.
**Do** implement two separate kernels following this structure:
| Mode | Kernel role (conceptually) |
|------|----------------------------|
| `transpose=False` | Head-major `qk` `[H,N]`, `l_prev` `[H]`, direct **`tl.dot(p, v, acc)`** |
| `transpose=True` | Seq-major `qk` `[N,H]`, separate **V** TMA descriptor (`shape=[S_kv, B, D]`, …), **`tl.dot(v, p, acc)`** after layout transposes |
**Autotune:** If `@triton.autotune` supplies `BLOCK_H` / `BLOCK_N`, the host must **not** pass fixed blocks into `torch.autograd.Function.apply`. Use:
```python
grid = lambda META: (triton.cdiv(num_head, META["BLOCK_H"]), B, 1)
kernel with grid then call (tensor_bases…, BLOCK_D=BLOCK_D, BLOCK_KPE=BLOCK_KPE)
```
**Full pattern, diagnosis table, and host-descriptor caveat:** [advanced-patterns.md](advanced-patterns.md).
---
## MEMORY LAYOUT PATTERNS (AVOID 50-150% REGRESSION)
**Case study: RoPE conversion showed 50-150% regression. Root cause: wrong memory layout.**
### Performance Killer #7: 5D Reshape vs Contiguous Access
**Impact: 50-150% slowdown**
When kernels access tensor halves (first `[0:dim//2]`, second `[dim//2:dim]`), naive 5D reshape breaks memory coalescing.
```python
# SLOW (50-150% regression) - 5D reshape
# Original cuTile might reshape: q.reshape(bsz, n_head, seq_len, 2, head_dim//2)
# Then access via stride_2 dimension - NON-CONTIGUOUS
q_offs_1 = batch * q_stride_b + heads[:, None] * q_stride_h + 0 * q_stride_2 + hd[None, :]
q_offs_2 = batch * q_stride_b + heads[:, None] * q_stride_h + 1 * q_stride_2 + hd[None, :]
# FAST - Transpose + contiguous + offset arithmetic
# Transpose to: (bsz, seq_len, n_head, head_dim) - head_dim CONTIGUOUS
q = q.transpose(1, 2).contiguous()
q = q + pid * q_row_stride # Row-based addressing
first_half_offs = heads[:, None] * hd + tl.arange(0, pad_hd // 2)[None, :]
second_half_offs = first_half_offs + (hd // 2) # Just add offset!
q_tile_1 = tl.load(q + first_half_offs, mask=mask)
q_tile_2 = tl.load(q + second_half_offs, mask=mask)
```
**Why this matters:**
- 5D reshape with `stride_2` creates non-sequential access: `[..., 0, :half_hd]` then `[..., 1, :half_hd]`
- Transpose + offset gives sequential access: `[:half_hd]` then `[half_hd:]` - both contiguous
- Memory coalescing difference: ~50-150% performance gap
### Performance Killer #8: Missing constexpr on Dimension Parameters
**Impact: 10-20% slowdown**
Triton compiler can optimize better when dimensions are compile-time constants.
```python
# SLOW - All parameters dynamic
def kernel(q, k, n_qh, n_kh, hd, pad_hd, BACKWARD_PASS):
# Compiler cannot unroll loops or optimize register allocation
# FAST - Dimension params as constexpr
def kernel(q, k,
n_qh: tl.constexpr,
n_kh: tl.constexpr,
hd: tl.constexpr,
pad_hd: tl.constexpr,
bs: tl.constexpr,
BACKWARD_PASS: tl.constexpr = False):
# Compiler can unroll, specialize, and optimize
```
### Performance Killer #9: Unnecessary Tensor Clones
**Impact: 10-20% slowdown**
In-place operations don't need explicit clones when transpose + contiguous already creates a copy.
```python
# SLOW - Explicit clone overhead
def forward(ctx, q, k, ...):
q = q.clone() # Unnecessary copy
k = k.clone() # Unnecessary copy
q, k = rope_forward(q, k, ...)
# FAST - Transpose + contiguous is sufficient
def rope_forward(q, k, ...):
q = q.transpose(1, 2) # Returns view
k = k.transpose(1, 2) # Returns view
q = q.contiguous() # Creates copy (if needed)
k = k.contiguous() # Creates copy (if needed)
# Now safe to modify in-place
```
### Pre-flight: Detect Split-Dimension Patterns
```bash
# Check if kernel splits a dimension
grep "reshape.*2," source.py # Reshape with 2 in shape
grep "stride_2\|stride(3)" source.py # Extra stride dimension
grep "half_\|// 2" source.py # Half-dimension patterns
```
If found, use transpose + offset pattern, NOT 5D reshape.
---
### Quick Diagnosis: Why Is My Converted Kernel Slow?
| Symptom | Likely Cause | Fix |
|---------|--------------|-----|
| 10-50x slower | No TMA, raw pointer + masks | Add `tl.make_tensor_descriptor` |
| Fast only when `transpose=True` (or only `False`) | One kernel + wrong softmax/`dot` layout for the other mode | Two kernels + autotune grid from `META` — [advanced-patterns.md](advanced-patterns.md) |
| 3-10x slower | O(N) group search per tile | Use while-loop with boundary tracking |
| 2-5x slower | Fixed tile sizes | Add `@triton.autotune` with GPU configs |
| **50-150% slower** | 5D reshape for split dims | Transpose + contiguous + offset arithmetic |
| 1.5-3x slower | No alignment hints | Add `tl.assume(stride % 8 == 0)` etc. |
| 1.2-2x slower | Masks on full tiles | Remove masks, rely on TMA bounds |
| 10-20% slower | Dynamic dimension params | Mark as `tl.constexpr` |
| 10-20% slower | Unnecessary `.clone()` | Let transpose + contiguous handle copies |