-
Notifications
You must be signed in to change notification settings - Fork 47
TileLang Case Studies
Detailed records of real debugging and optimization sessions on TileOps kernels. Each case follows the pattern: context → symptoms → investigation → root cause → fix → results → lessons.
These cases are extracted from hands-on work and represent the kinds of problems you are likely to encounter when writing or tuning TileLang kernels.
Skill files reference these cases with a short summary. Full technical detail lives here.
| Case | Topic | Date | Related Skill |
|---|---|---|---|
| A |
num_stages=0 sentinel — correctness failure across all configs |
2026-02-27 | kernel-debug |
| B | GEMV uncoalesced B access — achieved only ~3% of peak BW | 2026-02-27 | tune-multiplication |
| B.1 | GEMV O3 — pipelined B loads via cp.async: does it help? | 2026-02-27 | tune-multiplication |
| B.2 | GEMV O2 — explicit shmem for a: negligible on H200 |
2026-02-27 | tune-multiplication |
| B.3 | GEMV large-shape validation with LLM production shapes | 2026-02-27 | tune-multiplication |
| B.4 | GEMV forward() Python closure: 11 ms wall-clock vs 70 µs GPU |
2026-02-27 | tune-multiplication |
| C | nsys profiling: autotune contamination makes traces useless | 2026-02-27 | tune |
| D | Cross-pipeline shared memory race in GLA decode — intermittent state corruption | 2026-03-17 | kernel-debug |
| E | Builder function lru_cache — eliminating 28–200ms per-call overhead |
2026-03-19 | perf |
| F | DeltaNet backward correctness & performance — from wrong to faster than FLA | 2026-03-23 | kernel-debug, perf |
| G | Lane folding in small logical tiles — correct lowering with duplicated lane work | 2026-06-15 | kernel-debug, perf |
Date: 2026-02-27 Skill: kernel-debug
When refactoring the GEMV kernel from a register-only design (O1) to a pipelined shared-memory design (O3), a num_stages parameter was added to the autotune search space. To preserve backward compatibility with the old register-only path, num_stages=0 was introduced as a sentinel value meaning "skip pipeline, use registers". This seemed reasonable at the Python level — but it was fundamentally incompatible with how TileLang compiles kernels.
-
Shape under test:
(n=18432, k=7168, fp16, tune=True) -
max_err = 0.125–0.25across every autotune config, without exception - The error value was suspiciously round (a power of 2), suggesting a systematic offset rather than random numerical imprecision
- Clearing the JIT cache did not help
- Running a different shape worked fine
block_n= 1 rt=32 ns=0 max_err=0.1250 [FAIL]
block_n= 4 rt=32 ns=0 max_err=0.1250 [FAIL]
block_n= 4 rt=32 ns=2 max_err=0.1250 [FAIL]
block_n= 8 rt=32 ns=3 max_err=0.2500 [FAIL]
→ identical or near-identical errors across all configs
The first instinct was to look for a config-specific bug — maybe the autotune winner was wrong. However, the scan above showed that even configs known to work on other shapes were failing here. This pointed away from config-level bugs and toward something that was wrong regardless of hyperparameters.
The key diagnostic question is: does the error vary with config, or is it the same for all?
Here, it was the same → design-level flaw.
The kernel contained this branching pattern inside @T.prim_func:
# WRONG — num_stages=0 used as sentinel
if num_stages > 0:
for bk in T.Pipelined(T.ceildiv(k, block_k), num_stages=num_stages):
T.copy(b[...], b_shared, disable_tma=True)
...
else:
for bk in T.serial(T.ceildiv(k, block_k)):
# register-only path
...T.Pipelined requires num_stages >= 1. Using 0 as a sentinel violated this contract. More critically, the if num_stages > 0 conditional inside @T.prim_func was being processed by TileLang's TVM-based tracing, which may not treat it as a purely compile-time branch. This led to both branches potentially being included in the traced computation graph, causing interference. Additionally, when the kernel's parameter signature changed (from no num_stages in O1 to having num_stages in O3), stale cached binaries for ns=0 were reused incorrectly.
Remove the else branch (register-only path) entirely. The register-only fallback was not actually needed for production use — num_stages=1 through shared memory is at worst neutral and always valid.
# CORRECT — always use T.Pipelined; num_stages=1 means sequential, no overlap
for bk in T.Pipelined(T.ceildiv(k, block_k), num_stages=num_stages):
T.copy(b[bn * block_n, bk * block_k], b_shared, disable_tma=True)
...
# Autotune search space — no more 0:
for ns in [1, 2, 3]: ... # 1=sequential, 2=double-buffer, 3=triple-buffer-
Identical
max_erracross all configs → the design concept is wrong, not an implementation detail. A config-level bug (e.g., wrong index formula for a specificblock_n) would show selective failures. A design-level flaw breaks everything uniformly. Fix the concept before debugging details. -
Never use
num_stages=0as a sentinel forT.Pipelined. The API requiresnum_stages >= 1. Use separate Python-level flags to switch between code paths if needed, but keep TileLang primitives semantically valid. -
Powers-of-2 error values (0.125, 0.25, 0.5) signal systematic offset bugs, not rounding. This is a useful early signal to redirect investigation toward design-level issues rather than numerical precision.
See also: Anti-Patterns #25, Validation Rules — Error Pattern Interpretation.
Date: 2026-02-27 Skill: tune-multiplication Issue: tile-ai/TileOPs#232
The GEMV kernel computes c = B @ a where B is (n, k) and a is (k,). The initial implementation was functionally correct but extremely slow on large shapes. Performance profiling revealed the culprit: the thread layout caused the GPU to issue 32-way strided memory accesses to matrix B.
-
Affected shapes:
(n=7168, k=16384),(n=18432, k=7168)in fp16 and bf16 - Measured effective bandwidth: ~110 GB/s, i.e., ~3% of H200 peak (4.8 TB/s = 4800 GB/s)
- The kernel was functionally correct; only performance was affected
The kernel defined threads=(block_n, reduce_threads). In TileLang, threads=(dim_x, dim_y) maps threadIdx.x to the first dimension (fast-varying within a warp) and threadIdx.y to the second. With this layout:
-
tn = threadIdx.x— the row index of B -
tk = threadIdx.y— the column index of B
Threads within a warp (consecutive threadIdx.x values) therefore accessed different rows of B at the same column:
warp: B[row+0, col], B[row+1, col], B[row+2, col], ..., B[row+31, col]
stride between consecutive accesses = K × sizeof(fp16) ≈ 32 KB (for k=16384)
→ 32 completely non-contiguous memory locations per warp → 32 separate cache lines
This is the worst possible memory access pattern: 32-way fully strided access, achieving only ~3% of peak bandwidth.
Swap the thread dimensions: threads=(reduce_threads, block_n). Now:
-
tk = threadIdx.x— the column index (fast-varying) -
tn = threadIdx.y— the row index
Threads within a warp now access consecutive columns of the same row:
warp: B[row, col+0], B[row, col+1], ..., B[row, col+31]
→ 64 contiguous bytes per thread pair → one 512-byte fully coalesced transaction per warp
A coalesced 512-byte transaction saturates a single L2 cache line and approaches the hardware's maximum memory throughput.
| Shape (n, k) | tileops BW | torch baseline | speedup |
|---|---|---|---|
| (7168, 16384) | 3.65 GB/s | 3.34 GB/s | 1.09× |
| (18432, 7168) | 3.85 GB/s | 3.34 GB/s | 1.15× |
| (1024, 1024) | 0.76 GB/s | 0.26 GB/s | 2.9× |
H200 peak: 4.8 TB/s → achieved ~80% utilization on large shapes after the fix.
Autotune best configs:
-
(n=7168, k=16384):block_n=8, reduce_threads=32 -
(n=18432, k=7168):block_n=1, reduce_threads=32
-
For row-major matrices, the fast-varying thread dimension must correspond to the column index. This is the single most impactful optimization for memory-bound kernels — coalescing trumps all other micro-optimizations.
-
Always use a full warp (32 threads) for the reduction dimension.
reduce_threads < 32causes cross-row accesses within a warp, breaking coalescing. TileLang'stvm_thread_allreducemaps to efficient warp shuffles whenreduce_threads=32. -
The 2.9× gain at small shapes reflects
torchdispatch overhead, not kernel efficiency. For large shapes both implementations approach the roofline together (~3.3–3.9 TB/s).
See also: Hardware Constraints — Roofline Model, Anti-Patterns #1.
Date: 2026-02-27 Continuation of: Case B
After the coalescing fix (O1), the kernel ran at ~80% of H200 peak. The remaining ~20% gap was hypothesized to be due to HBM3e latency. The O3 optimization attempted to hide this latency by adding a software pipeline: while one tile of B is being consumed, the next tile is being prefetched into shared memory using cp.async.
The T.serial + direct register-load loop was replaced with T.Pipelined + T.copy(disable_tma=True) + b_shared:
b_shared = T.alloc_shared((block_n, block_k), dtype)
for bk in T.Pipelined(T.ceildiv(k, block_k), num_stages=num_stages):
T.copy(b[bn * block_n, bk * block_k], b_shared, disable_tma=True)
# ... compute using b_shared ...disable_tma=True is required on SM90: T.copy defaults to TMA, which requires mbarrier layout inference that TileLang cannot infer for manually-indexed shared memory in non-WGMMA kernels. Using disable_tma=True switches to cp.async, which still hides HBM latency when inside T.Pipelined.
Pipeline depths in the autotune space: num_stages ∈ {1, 2, 3} (1 = sequential, 2 = double-buffer, 3 = triple-buffer).
| Shape (n, k) | dtype | tileops BW | torch baseline | speedup |
|---|---|---|---|---|
| (7168, 16384) | fp16 | 3.47 TB/s | 3.34 TB/s | 1.04× |
| (18432, 7168) | fp16 | 3.83 TB/s | 3.34 TB/s | 1.15× |
| (7168, 16384) | bf16 | 3.62 TB/s | 3.35 TB/s | 1.08× |
| (18432, 7168) | bf16 | 3.78 TB/s | 3.35 TB/s | 1.13× |
| (1024, 1024) | fp16 | 0.57 TB/s | 0.26 TB/s | 2.2× |
Autotune winners: All large shapes converged to block_n=1, num_stages=3.
O3 is slightly slower than O1 for (n=7168, k=16384) (3.47 vs 3.65 TB/s). This was counterintuitive: adding a pipeline should help. The explanation:
- For large K (16384), the GPU's own out-of-order execution already hides most of the HBM latency for simple register loads. The pipeline provides diminishing returns.
- O3 adds overhead: shared memory allocation (TLB pressure),
cp.asyncsetup cost per tile, and potential bank conflicts inb_shared. For very smallblock_n=1, these costs are not amortized. - For smaller K (7168), O3 does help (+15%) because shorter K means fewer in-flight transactions and less natural latency hiding — the pipeline adds real benefit.
-
Pipelining does not always help, even for memory-bound kernels. When the GPU's warp scheduler already hides latency through warp switching, explicit software pipelining may add overhead without benefit. Profile before assuming.
-
disable_tma=Trueis required forT.copyinto manually-indexed shared memory on SM90. TMA requires layout inference that is only available for WGMMA kernels. For all other kernels, usecp.asyncviadisable_tma=True. See Anti-Patterns #12.
Date: 2026-02-27 Continuation of: Case B
In GEMV, each of the block_n rows in a block reads the same vector a from global memory. If a is loaded once into shared memory by the first warp and then read by all warps, the global traffic for a is reduced by a factor of block_n. For block_n=8, this is a 7/8 reduction in a reads. The hypothesis was that this would provide a measurable bandwidth improvement.
| Shape (n, k) | O1 BW | O1+O2 BW | Delta | Best config |
|---|---|---|---|---|
| (7168, 16384) | 3.65 GB/s | 3.67 GB/s | +0.5% | block_n=8, rt=32 |
| (18432, 7168) | 3.85 GB/s | 3.85 GB/s | 0% | block_n=1, rt=32 |
| (1024, 1024) | 0.76 GB/s | 0.71 GB/s | −7% | (no tune) |
The O2 optimization delivered negligible benefit on large shapes and caused a regression on small shapes.
Vector a has size k × sizeof(dtype) = 14–32 KB (for k=7168–16384 in fp16). The L1 data cache per SM on H200 is 256 KB. When block_n=8 threads all access the same a vector, the hardware L1 cache automatically deduplicates those accesses — the second through eighth warps find a already in L1. The explicit shared memory copy adds overhead (allocation, T.sync_threads() calls) without providing any additional locality beyond what L1 already gives.
For small shapes (n=1024), the two T.sync_threads() calls per outer loop iteration added measurable overhead (-7%), because the kernel is otherwise very fast and each sync represents a larger fraction of total execution time.
-
Explicit shared memory only helps when the working set exceeds L1 capacity. For inputs that fit comfortably in L1 (< 128 KB), the hardware handles deduplication automatically.
-
Always measure before assuming an optimization helps. The theoretical savings (reducing
atraffic by 7/8) looked compelling on paper but yielded nothing in practice on H200. -
T.sync_threads()has non-zero cost. On fast kernels or with smallblock_n, the synchronization overhead can dominate any data-traffic savings.
Date: 2026-02-27 Continuation of: Case B
After O1/O3 tuning, the benchmark covered shapes like (7168, 16384) and (18432, 7168). These are representative but not the largest shapes that appear in production LLM inference. Adding shapes from real models verifies that the optimizations scale and that no regressions exist at higher problem sizes.
| Shape (n, k) | Source model | Layer type |
|---|---|---|
| (28672, 8192) | Llama-3 70B | MLP gate/up projection |
| (57344, 7168) | DeepSeek-V3 | MoE aggregated output (8 experts × 7168) |
| Shape (n, k) | dtype | tileops BW | baseline BW | speedup | H200 util. |
|---|---|---|---|---|---|
| (7168, 16384) | fp16 | 3.47 TB/s | 3.34 TB/s | 1.04× | 72% |
| (18432, 7168) | fp16 | 3.80 TB/s | 3.34 TB/s | 1.14× | 79% |
| (28672, 8192) | fp16 | 4.02 TB/s | 3.80 TB/s | 1.06× | 84% |
| (57344, 7168) | fp16 | 4.26 TB/s | 3.92 TB/s | 1.09× | 89% |
| (7168, 16384) | bf16 | 3.61 TB/s | 3.35 TB/s | 1.08× | 75% |
| (18432, 7168) | bf16 | 3.75 TB/s | 3.34 TB/s | 1.12× | 78% |
| (28672, 8192) | bf16 | 4.08 TB/s | 3.81 TB/s | 1.07× | 85% |
| (57344, 7168) | bf16 | 4.30 TB/s | 3.93 TB/s | 1.09× | 90% |
Utilization increases steadily with problem size: 72% at n=7168 → 89% at n=57344. The reason: n=57344 places approximately 13.5 thread blocks per SM (vs ~3.4 for n=7168) on the 132 SMs of H200. More warps per SM means more in-flight HBM3e transactions, which gives the memory controller more opportunity to reorder requests and achieve higher effective throughput.
-
For GEMV on Hopper, practical bandwidth utilization rises with problem size because increased warp-level parallelism improves HBM3e latency hiding. Small shapes (n=1024) are far below peak due to insufficient occupancy.
-
Always validate on production-scale shapes before finalizing a kernel. An optimization that looks good at
n=7168may behave differently atn=57344due to occupancy changes.
Date: 2026-02-27 Continuation of: Case B
After all kernel-level optimizations, profiling GemvOp.forward() in a tight Python loop showed ~11 ms per call in wall-clock time. The GPU kernel itself runs in ~70 µs — a factor of ~150 discrepancy.
forward() was calling _gemv_wrapped_kernel, a torch.library.custom_op wrapper. Internally this called:
_gemv_kernel(n, k, dtype)(block_n, reduce_threads, num_stages)Each call chain:
- Creates a new Python function object (closure) capturing
n,k,dtype - Looks up the JIT cache by key (involves hashing)
- Returns a callable that then invokes the compiled kernel
Steps 1 and 2 happen on every single forward() call, adding ~10 ms of pure Python overhead per invocation regardless of how fast the GPU kernel is.
Populate self.kernel in __init__ (after init_config / autotune completes) and call it directly in forward():
class GemvOp:
def __init__(self, n, k, dtype, ...):
...
self.kernel = _gemv_kernel(n, k, dtype) # compiled once at init time
def forward(self, a, b):
a = a.flatten().contiguous()
return self.kernel(
self.config["block_n"],
self.config["reduce_threads"],
self.config["num_stages"],
)(a, b)self.kernel is the closure already compiled and cached; calling it hits the in-memory JIT cache with near-zero overhead. _gemv_wrapped_kernel is kept alive for torch.compile compatibility but is no longer called in the hot path.
-
GPU time via CUDA events (
tilelang.profiler.do_bench) is the only valid metric for kernel evaluation. Wall-clock time includes Python dispatch, closure creation, and JIT lookup — all of which are irrelevant to kernel performance. Never benchmark withtime.time()ordatetime.now(). -
Compile once at init time, cache the result. Any operation that creates objects or does lookups proportional to call frequency (not problem size) will dominate at high call rates. Move it to
__init__.
Date: 2026-02-27
Skill: tune
Shape: (n=7168, k=16384, fp16, H200)
After completing kernel-level optimizations, nsys profile was used to get a precise breakdown of GPU time and to compare against cuBLAS. The first profiling run was done naively — with autotune enabled — and the results were uninterpretable. This case documents the problem and the fix.
Time (%) ... Avg (ns) Med (ns) StdDev (ns) Name
28.9 ... 91 701 85 472 24 658 _gemv_main_kernel
nsys stats aggregates all launches of a given kernel name within the profiling window. With autotune enabled, this meant 15 configs × 20 trial runs = 300 slow trial launches were mixed with the 200 steady-state benchmark calls. The result:
- Average was inflated by slow trial runs (different
block_n, suboptimal configs) - StdDev of 24 µs was enormous relative to the ~64 µs median — a sign of highly heterogeneous launches
- The 200 steady-state calls (which represent actual performance) were invisible in the noise
Pass the best config directly at construction time to bypass autotune:
op = GemvOp(
n, k, dtype=dtype,
config={"block_n": 1, "reduce_threads": 32, "num_stages": 3}
)Then profile:
nsys profile --trace=cuda --output=/tmp/gemv_fixed \
python -m pytest benchmarks/ops/bench_gemv.py::test_case -vvs
nsys stats /tmp/gemv_fixed.nsys-rep --report cuda_gpu_kern_sum Time (%) ... Avg (ns) Med (ns) StdDev (ns) Name
47.3 ... 64 261 66 208 3 733 _gemv_main_kernel ← tileops
45.8 ... 62 107 63 500 1 200 cutlass_splitK_... ← cuBLAS (pass 1)
3.1 ... 2 354 2 300 180 cutlass_reduce_... ← cuBLAS (pass 2)
- tileops: median 66 µs → 3.55 TB/s (74% of H200 peak 4.8 TB/s)
- cuBLAS split-K: 63.5 + 2.3 = 65.8 µs total → 3.64 TB/s
- Both within 5%; CUPTI benchmark (
do_bench) shows tileops slightly ahead due to measurement variance
Note: cuBLAS uses split-K, which issues two kernels — the main GEMM pass and a reduction pass. Always sum their latencies for a fair comparison.
Even with a good kernel, H200 rarely achieves 100% of rated bandwidth for GEMV. The main reasons:
| Factor | Explanation |
|---|---|
| Warp occupancy: 54/64 = 84% | Determined by problem size (n=7168 / 132 SMs). Cannot be increased without changing the shape. |
| HBM3e row buffer effects | DRAM rows are opened/closed with each access; latency is not zero even with perfect coalescing |
| Realistic BW utilisation | H200 achieves 60–80% of its rated 4.8 TB/s in practice for GEMV |
| shmem bank conflicts | 4-way conflicts in b_shared; but shmem reads are NOT the bottleneck — HBM loading dominates |
74% efficiency is close to the practical ceiling for this problem size. Further improvements would require increasing n to raise occupancy (see Case B.3).
-
Always disable autotune and fix the best config before running nsys. The trace must contain only steady-state launches to be interpretable.
-
Use median, not average, when comparing kernels. Average is vulnerable to outlier launches. StdDev > 10% of median is a sign of heterogeneous launches (likely autotune contamination or cache variance).
-
cuBLAS split-K issues two kernels. When you see a
splitK+reducepair, sum their latencies for the true cuBLAS cost. -
Use
do_bench(CUDA events / CUPTI) for the authoritative benchmark number. nsys is for understanding why;do_benchis for comparing how fast.
See also: Hardware Constraints — H200 Bandwidth.
Date: 2026-03-17
Kernel: tileops/kernels/linear_attn/gla/gla_decode.py — GLADecodeKernel (TC path)
Issue: tile-ai/TileOPs#550
The GLA (Gated Linear Attention) decode kernel computes a single-step recurrence:
S_new = diag(exp(gk)) @ S + outer(k, v)
o = scale * q^T @ S_new
The TC (tensor-core) kernel uses warpgroup specialization with two passes:
-
Pass 1:
T.Pipelined+T.gemm— computes outputovia fused matvec -
Pass 2:
T.Pipelined— computesnew_statevia element-wise state update
Both passes used the same shared memory buffer h_tile for TMA-pipelined state tile prefetch.
- bf16/fp16 TC kernel: intermittent failures, ~5–10% of test runs
-
fp32 kernel: always correct (uses scalar accumulation, no
T.gemm, no warpgroup specialization) - Error magnitude: ~0.2–0.4 (far exceeding bf16 tolerance of 0.02), indicating completely wrong data
-
Error location: always
new_state(Pass 2 output);o(Pass 1 output) was always correct - Flaky in CI: running the same test 10 times → only 1/10 fully green
The investigation followed a systematic elimination process:
Step 1 — Isolate the trigger.
Running the kernel function directly in a tight loop: 0/500 failures. Running through GLADecodeOp (which creates the kernel via tilelang JIT cache each time): ~7% failure rate. This initially pointed at the Op/cache path, but further tests showed that even inserting torch.cuda.synchronize() between calls triggered ~5% failures. Any CPU-side delay changed the failure rate, confirming a timing-dependent race condition.
Step 2 — Identify the affected region.
Collecting 50+ failures revealed: corruption was always in tile 1 (dk=16:32) — the second k-dimension tile. This corresponds to buffer 1 in the double-buffered T.Pipelined pipeline (num_stages=2). The number of corrupted elements varied (23 to 400 out of 1024 per tile), and the affected rows were uniformly distributed, consistent with partially-completed TMA loads.
Step 3 — Understand the generated code.
Inspecting the CUDA source (kernel.get_kernel_source()) revealed the warpgroup-specialized structure:
threads 128-255 (producer): threads 0-127 (consumer):
Pass 1: TMA load loop Pass 1: gemm loop
__sync_partial (producer) __sync_partial (consumer)
Pass 2: TMA load loop output computation
Pass 2: state update loop
Key observation: __sync_thread_partial only synchronizes within each warpgroup — there is no inter-warpgroup barrier between Pass 1 and Pass 2.
When TileLang compiles two consecutive T.Pipelined loops sharing the same shared memory buffer (h_tile), the producer warpgroup for Pass 2 starts TMA loads immediately — its first consumer-done barrier wait (mbarrier[6].wait(1)) returns instantly because the barrier is in its initial state (parity 0 ≠ 1).
Meanwhile, the consumer warpgroup may still be in Pass 1's final gemm iterations, reading from h_tile. The producer's Pass 2 TMA loads overwrite h_tile with new data while the consumer is still reading it for Pass 1's gemm.
The double-buffering makes tile 1 (buffer 1) the most vulnerable: it is the last buffer written by Pass 1 (at kt=3) and the first buffer overwritten by Pass 2's producer (at kt_1=1).
Timeline (race scenario):
Producer Pass 1 done ──┐
├─ __sync_partial (producer only)
Producer Pass 2 start ─┘
kt_1=0: TMA → buffer 0 ─┐
kt_1=1: TMA → buffer 1 ────┤── overwrites h_tile!
│
Consumer Pass 1 still running:│
kt_2=3: gemm reads buffer 1 ←── RACE: reads partially overwritten data
...
Consumer Pass 1 done ──┐
├─ __sync_partial (consumer only)
Consumer Pass 2 start ─┘
kt_3=1: reads buffer 1 ←── may contain stale/mixed data
Allocate a separate shared memory buffer for Pass 2's pipeline, eliminating the cross-pipeline data hazard:
# Before (buggy): both passes share h_tile
h_tile = T.alloc_shared([k_tile, dim_v], dtype)
# Pass 1
for kt in T.Pipelined(...):
T.copy(state[...], h_tile)
T.gemm(q_tile, h_tile, acc, ...)
# Pass 2
for kt in T.Pipelined(...):
T.copy(state[...], h_tile) # ← overwrites Pass 1's buffer
... = h_tile[kk, j]
# After (fixed): separate buffers
h_tile = T.alloc_shared([k_tile, dim_v], dtype) # Pass 1 only
h_tile2 = T.alloc_shared([k_tile, dim_v], dtype) # Pass 2 only
# Pass 1
for kt in T.Pipelined(...):
T.copy(state[...], h_tile)
T.gemm(q_tile, h_tile, acc, ...)
# Pass 2
for kt in T.Pipelined(...):
T.copy(state[...], h_tile2) # ← independent buffer
... = h_tile2[kk, j]Trade-off: +2 KB shared memory per thread block (one extra [16, 64] bf16 tile). Negligible impact on occupancy.
| Metric | Before | After |
|---|---|---|
| 10× full test suite (21 tests each) | 1/10 all-pass, avg 2 failures/run | 10/10 all-pass |
| 500-trial stress test (with trigger) | ~35 failures | 0 failures |
| Performance | — | No regression (same kernel structure) |
-
Never share a
T.alloc_sharedbuffer across twoT.Pipelinedloops in a warpgroup-specialized kernel. TileLang'sT.Pipelinedgenerates a producer-consumer warpgroup split with TMA async copies. Two consecutive pipelines reusing the same buffer lack inter-warpgroup synchronization at the boundary — the second pipeline's producer can start overwriting the buffer before the first pipeline's consumer finishes reading it. Always allocate a separate buffer for each pipeline. -
Intermittent correctness failures that depend on timing (sleep, sync, CPU overhead) almost certainly indicate a shared-memory or async-copy race condition. The diagnostic fingerprint: (a) non-deterministic, (b) affected by unrelated CPU-side activity, (c) large errors (~0.3) rather than precision drift (~0.001). When you see this pattern, focus on shared memory lifecycle across pipeline boundaries.
-
Systematic tile-level error analysis pinpoints the race. When 50+ failures all corrupt the same tile index (here: always tile 1 = buffer 1), the double-buffer index directly implicates the pipeline synchronization mechanism. Cross-reference with the generated CUDA source to trace which warpgroup writes/reads that buffer and when.
-
compute-sanitizerdoes not catch TMA/mbarrier races. The CUDA memory checker and race checker reported zero errors for this bug. For TMA pipeline races, manual code inspection of the generated CUDA source (viakernel.get_kernel_source()) and systematic empirical bisection are the only reliable debugging approaches. -
Bisect by replacing
T.PipelinedwithT.SerialandT.gemmwith scalarT.Parallel. When debugging a suspected pipeline or gemm issue, systematically downgrade one component at a time to isolate which primitive is involved. In this case, the bug disappeared when either pass used a non-pipelined loop, confirming the cross-pipeline interaction as the root cause.
See also: Anti-Patterns — Sharing Buffers Across Pipelines, Patterns — Warpgroup-Specialized Pipelines.
Date: 2026-03-19 PRs: #597, #600 Issues: #571, #599
TileOPs kernels follow a two-level architecture: builder functions (_*_kernel in tileops/kernels/) construct the TileLang program and return a @tilelang.jit callable, while dispatch wrappers (@torch.library.custom_op) call the builder on every forward() invocation with the same shape parameters.
Even though tilelang has an internal KernelCache._memory_cache, each builder call still performs TIR generation → serialization → SHA-256 lookup before hitting the cache. This overhead was measured at 28–200 ms per call, completely dominating the actual GPU kernel execution time (~0.07 ms for a typical GEMV).
- Wall-clock profiling of
forward()showed 30–200 ms per call, while CUDA event profiling (do_bench) showed ~70 µs GPU time - The overhead was proportional to kernel complexity (more TIR nodes → slower serialization): simple kernels ~28 ms, complex chunkwise kernels ~200 ms
- The overhead was identical for repeated calls with the same parameters — no amortization across calls
The call chain on every forward():
forward() → @torch.library.custom_op wrapper → _*_kernel(shape_params)
↓
TIR program construction
↓
@tilelang.jit compilation
↓
KernelCache: serialize → SHA-256 → lookup
↓
cache hit → return JITImpl
Steps 2–5 execute every call, even though the result is always the same JITImpl for identical shape_params. The builder function is a pure function of its parameters — it always returns the same program — but was not cached at the Python level.
Add @functools.lru_cache(maxsize=32) to every builder function:
import functools
@functools.lru_cache(maxsize=32)
def _gemv_kernel(n: int, k: int, dtype: str = "float16") -> Callable:
@tilelang.jit(out_idx=[2])
def kernel(block_n, reduce_threads, num_stages):
@T.prim_func
def main(...): ...
return main
return kernelWith the decorator, the second call with the same (n, k, dtype) returns the cached JITImpl directly — skipping TIR generation, serialization, and SHA-256 entirely.
Before applying the fix across 51 files, we verified the following invariants:
| Concern | Status | Reason |
|---|---|---|
| Mutable return values | Safe | Builders return @tilelang.jit callables — stateless, never mutated by callers |
| Shared state across Op instances | Safe |
config (threads, num_stages) is per-Op-instance via self.config, not baked into the cached object |
tune parameter conflict |
Safe |
tune is handled at the Op layer (init_config), not passed to the builder |
| Unhashable parameters | Safe | All builder params are int, str, bool, float, torch.dtype — all hashable |
| Memory retention | Acceptable |
maxsize=32 bounds memory; typical models use <32 distinct shapes |
| Thread safety | Safe | CPython GIL protects lru_cache dict operations |
| Builder | Without cache | With cache | Speedup |
|---|---|---|---|
_h_recurrence_tl |
36.08 ms | 0.0003 ms | ~120,000× |
_output_o_tl |
17.83 ms | 0.0002 ms | ~89,000× |
fused_prepare_compute_w_u_tl |
38.37 ms | 0.0002 ms | ~192,000× |
Total builder overhead per forward call: ~92 ms → ~0.0007 ms
The fix was applied in two rounds:
| PR | Scope | Files | Builder functions |
|---|---|---|---|
| #597 | gated_deltanet, flash_attn, flash_decode | 11 | 21 |
| #600 | Remaining 13 kernel families | 40 | 80+ |
| Total | 51 | 100+ |
The pattern had accumulated across the entire codebase because there was no coding standard requiring it.
This case is the generalization of Case B.4, which identified the same overhead pattern in the GEMV kernel. Case B.4 fixed it by caching in self.kernel at __init__ time. This case applies lru_cache at the builder function level — a more fundamental fix that caches the builder output globally, benefiting all callers (including @torch.library.custom_op wrappers used by torch.compile).
-
Cache builder functions, not just kernel objects. The builder function is a pure function of shape parameters.
@functools.lru_cache(maxsize=32)eliminates the entire TIR → serialize → hash → lookup chain. This is now a mandatory coding standard inDEVELOPMENT.md. -
Distinguish builder functions from dispatch wrappers.
_*_kernel(shape_params)is the builder (cacheable, hashable params).@torch.library.custom_opis the dispatch wrapper (takes tensors, not cacheable). Only the former getslru_cache. -
Two-level caching is complementary.
lru_cache(Python level) avoids re-entering the builder entirely.KernelCache._memory_cache(tilelang level) avoids recompilation if the builder is entered with new params. Both are needed; neither replaces the other. -
Coding standards prevent accumulation. This issue grew from 0 to 51 files over months because nothing enforced the pattern. Adding it to
DEVELOPMENT.mdand.claude/rules/code-style.mdensures new kernels include the decorator from day one. -
maxsize=32is a safe default. Most LLM training/inference workloads use fewer than 32 distinct shape combinations per kernel family. The bound prevents unbounded memory growth while covering realistic usage.
See also: Case B.4 (GEMV-specific instance of this pattern).
Date: 2026-03-23 Skills: kernel-debug, perf
Ungated DeltaNet chunkwise operators were ported from the existing gated DeltaNet implementation by removing the gate parameter g and all exp(g) scaling. The forward pass and recurrence (decode) passed tests immediately, but the backward pass had numerical errors in dk and dbeta that exceeded the tolerance.
FAILED test_deltanet_bwd[2-64-2-64-64-32-dtype0-False]
AssertionError: dk: Tensor-likes are not close!
Mismatched elements: 8 / 16384 (0.0%)
Greatest absolute difference: 0.01564 at index (1, 1, 7, 8) (up to 0.01 allowed)
- Only
dkanddbetafailed;dqanddvpassed easily (errors ~1e-4) - Error was small (0.0156 vs tolerance 0.01) and affected <0.1% of elements
- Larger test (S=128, 4 chunks) had larger error (0.0226)
The backward computes dk = dk_partial + dk_corr + dk_wu from three separate stages. A debug script computed each component and compared against autograd reference:
dk_partial + dk_corr: correct (error ~1e-4)
dk_wu: WRONG (error 0.0157)
dAw max abs: 0.0086 (discarded!)
dAu max abs: 0.0877 (discarded!)
compute_w_u_bwd output dAw and dAu (gradients w.r.t. the WY matrix A_inv) but they were discarded with _. This was inherited from the gated DeltaNet code.
Deriving the full backward from the forward equations h_{c+1} = h_c + k^T @ v_new, v_new = u - w @ h, w = A_inv @ (k * beta) revealed three missing contributions:
| Missing path | Formula | Impact |
|---|---|---|
| A_inv backward |
dA = -A_inv^T @ dA_inv @ A_inv^T → dP → dk_A, dbeta_A
|
dk_wu wrong by 0.0157 |
| dh recurrence coefficient |
dh_c = dh_local + (I - W^T @ K) @ dh_{c+1}, not dh_c = dh_local + dh_{c+1}
|
dh propagation wrong for >2 chunks |
| dw cross-chunk | dw_corr = -(du_corr @ S^T) |
dw missing cross-chunk contribution |
The first bug (A_inv backward) caused the S=64 failure. The second (dh coefficient) was masked for 2 chunks but caused the S=128 failure (0.0226 error). All three were needed for correctness.
The gated variant has the same missing paths but its tests passed because:
- The gate decay
exp(g) < 1naturally dampens error propagation across chunks - The A_inv gradient contribution is smaller when weighted by the gate
Three changes to deltanet_bwd.py:
1. Corrected dh recurrence — The state update h_{c+1} = h_c + k^T @ (u - w @ h_c) means ∂h_{c+1}/∂h_c = I - k^T @ w, not I. Added W^T @ K @ dh correction term:
# Before (wrong):
dh_fp32 = dh_fp32 + dh_loc
# After (correct):
wk_dh = w_c^T @ (k_c @ dh_buf) # [DK,DV]
dh_fp32 = dh_fp32 + dh_loc - wk_dh2. dw cross-chunk correction — du_corr flows through v_new = u - w @ h to produce dw_corr = -(du_corr @ S^T).
3. A_inv backward — Backpropagate dAw + dAu through the Neumann-series inverse:
dA = -(A_inv^T @ (dAw + dAu) @ A_inv^T)
dP = strictLower(dA)
dk_A = beta * (dP @ k) + dP^T @ (k * beta)
dbeta_A = (dP * KK^T).sum(-1)After all three fixes, errors dropped from 0.0156 → 0.00015 (100x improvement) and tolerances were tightened 10x.
The correctness fix introduced significant backward overhead. Initial BWD was 2.47x slower than FLA at S=16384. A profiling breakdown revealed:
| Stage | Time | % | Type |
|---|---|---|---|
| dh_recurrence | 0.881ms | 48% | TileLang sequential (W^T@K@dh added 1 gemm/step) |
| A_inv_bwd | 0.568ms | 31% | Python batched matmul |
| wu_bwd | 0.114ms | 6% | TileLang parallel |
| fused_prepare (recompute w,u) | 0.099ms | 5% | TileLang parallel |
| bwd_parallel | 0.093ms | 5% | TileLang parallel |
| tensor_adds | 0.079ms | 4% | Python element-wise |
Four optimizations were applied sequentially:
Opt 1: Move dw_corr out of sequential recurrence.
The dw_corr = -(k @ dh_buf) @ h^T gemm ran inside the serial loop (grid = B×H = 8 threads). Moved to a batched matmul after the recurrence. Savings: ~10% at S=16384.
Opt 2: Fuse A_inv backward into compute_w_u_bwd TileLang kernel.
The Python-level A_inv backward did 6+ CUDA kernel launches per call (reshape, two batched matmuls, tril, element-wise ops). Fusing all operations into the existing compute_w_u_bwd TileLang kernel (which already loads Aw, k, beta) eliminates launch overhead. One TileLang subtlety: reusing a fragment across two T.gemm calls with different transpose patterns causes a layout conflict — the fix is to allocate a fresh dA_frag2 for the second gemm. Savings: ~30%.
Opt 3: Forward saves w, u; backward skips recomputation.
Forward already computes w, u but discarded them. Now returns (o, S, Aw, Au, w, u). Backward receives these from autograd.Function.ctx instead of rerunning fused_prepare_compute_w_u. Savings: ~8%.
Opt 4: Fuse dw_corr, du merge, and dk merge into compute_w_u_bwd.
The wrapper did 3 Python-level tensor additions (du = du_partial + du_corr, dw += dw_corr, dk = dk_partial + dk_corr + dk_wu) plus the dw_corr batched matmul. All fused into the TileLang kernel by passing raw pieces as inputs. The kernel now accepts 11 inputs and produces final (dk, dv, dbeta) with all corrections applied. Savings: ~17%.
Correctness:
| dtype | fwd atol | bwd atol | actual max error |
|---|---|---|---|
| fp32 | 1e-3 | 1e-3 | ~2e-4 |
| fp16 | 2e-2 | 5e-3 | ~1e-4 |
| bf16 | 5e-2 | 2e-2 | ~1e-3 |
Performance (BWD only, B=2 H=4 fp16):
| S | Initial | After 4 opts | FLA | Speedup | vs FLA |
|---|---|---|---|---|---|
| 2048 | 0.416ms | 0.120ms | 0.402ms | 3.5x | 0.30x |
| 4096 | 0.689ms | 0.234ms | 0.396ms | 2.9x | 0.59x |
| 8192 | 0.955ms | 0.441ms | 0.407ms | 2.2x | 1.08x |
| 16384 | 1.813ms | 0.863ms | 0.731ms | 2.1x | 1.18x |
Performance (FWD+BWD combined):
| S | TileOPs | FLA | vs FLA |
|---|---|---|---|
| 2048 | 0.279ms | 0.772ms | 0.36x |
| 4096 | 0.336ms | 0.779ms | 0.43x |
| 8192 | 0.661ms | 0.781ms | 0.85x |
| 16384 | 1.293ms | 1.109ms | 1.17x |
Autotune was tested but produced no further improvement — the default config (num_stages=2, threads=256) was already optimal for the small search space ([128, 256] threads × [1, 2] stages). The search space is too narrow for this kernel because the bottleneck is the sequential recurrence loop, not thread-block-level parallelism.
-
When porting gated → ungated, re-derive the adjoint equations from scratch. Removing the gate changes
∂h_{c+1}/∂h_cfromexp(g) · I - k_scaled^T @ PtoI - k^T @ w. The coefficientI(simple accumulation) is wrong for the ungated case — thev_new = u - w @ hdependency creates a coupling term that doesn't exist in the gated version because the gate decay dampens it. -
Discarded gradients (
_dAw, _dAu = ...) are a code smell. If a backward kernel computes a gradient and you throw it away, either the forward doesn't depend on that variable (verify!) or you're missing a gradient path. In this case,A_invdepends onkandbetathrough the Gram matrix — the discardeddAw/dAushould have flowed back. -
Profile before optimizing. The initial instinct was to optimize the recurrence kernel (48% of time). But the A_inv backward (31%) was pure Python overhead easily eliminated by TileLang fusion. The profiling breakdown
[recurrence=48%, A_inv=31%, wu_bwd=6%, fused=5%, parallel=5%, adds=4%]directly prioritized the work. -
TileLang fragment layout conflicts. Reusing a fragment in two
T.gemmcalls with different transpose patterns (T.gemm(A, B, frag)thenT.gemm(C, D, frag, transpose_A=True)) causesInternalError: Get different layout for frag. Solution: allocate a fresh fragment for the second gemm. This is a TileLang-specific pitfall not documented elsewhere. -
Fuse Python-level tensor ops into TileLang kernels. Each Python tensor operation (
+,@,.reshape()) launches a separate CUDA kernel. For the backward pipeline, 3 tensor additions + 1 batched matmul cost 0.25ms at S=16384 — 20% of total time. Fusing them into an existing parallel TileLang kernel (which already has the data in shared memory) is nearly free. -
Forward should save intermediate results for backward. The
fused_prepare_compute_w_ukernel computesw, uin the forward pass, but the original design discarded them (only savingAw, Au, S). The backward then recomputed them from scratch. Passingw, uthroughautograd.Function.ctxsaved ~8% of backward time at zero forward cost. -
Autotune has diminishing returns when the bottleneck is algorithmic. The sequential dh recurrence runs on a grid of
B×H(= 8 threads), iteratingNCtimes. No thread/stage tuning can parallelize this loop — the bottleneck is theW^T @ K @ dhgemm per step, which is mathematically required for correctness. Autotune is most effective for parallel kernels where thread count and pipeline depth affect occupancy. -
Neumann series vs forward substitution trade-off. TileOPs uses Neumann series (
A_inv ≈ I + P + P² + ...) for the WY inverse because it maps naturally toT.gemminlog₂(chunk_size)rounds. FLA uses forward substitution (sequential row solve). The Neumann approach enables GPU parallelism but requires an explicitA_invbackward; forward substitution is inherently differentiable through autograd. This design choice is the root cause of the performance gap at long sequences.
Date: 2026-06-15 Skill: kernel-debug, perf
While optimizing the Gated DeltaNet decode kernel on Hopper, we tried a TileLang implementation that looked logically clean: one block handled one (batch, head, V tile) and the kernel used T.Parallel(v_tile) or T.Parallel(dim_k, v_tile) to work on the V tile.
The tested configuration used:
threads=32dim_k=128dim_v=128v_tile=8bf16
At the TileLang source level, there was no intentional duplicate work:
with T.Kernel(batch, head, dim_v // v_tile, threads=threads):
v_base = vid * v_tile
h_frag = T.alloc_fragment([dim_k, v_tile], accum_dtype)
for kk, j in T.Parallel(dim_k, v_tile):
h_frag[kk, j] = state[bid, hid, kk, v_base + j]
for kk in T.Serial(dim_k):
for j in T.Parallel(v_tile):
old_val[j] += h_frag[kk, j] * k[bid, hid, kk]The kernel was correct, but slower than expected. For the representative B=1, H=32, DK=DV=128, bf16 decode workload:
| Variant | Latency | NCU Duration | Registers / Thread |
|---|---|---|---|
TileLang FLA-style prototype (serial001) |
0.006636 ms | 7.62 us | 255 |
Explicit lane-group kernel (serial067) |
0.003818 ms | 4.74 us | 104 |
The prototype also showed high register pressure. The generated CUDA revealed the important clue:
h_frag[i] = state[..., (((int)threadIdx.x) & 7)];
v_new[0] = v[..., (((int)threadIdx.x) & 7)] - old_val[0];
new_state[..., (((int)threadIdx.x) & 7)] = ...;
if ((((int)threadIdx.x) >> 3) == 0) {
o[...] = ...;
}threadIdx.x & 7 means the 32 physical warp lanes were folded onto only 8 logical V lanes. Lanes 0, 8, 16, 24 mapped to the same V element, lanes 1, 9, 17, 25 mapped to the next V element, and so on.
This was not a correctness issue. It was a performance trap caused by a mismatch between the physical thread count and the logical parallel space exposed to TileLang.
The source said "parallel over v_tile=8", while the kernel launched 32 threads. TileLang lowered the extra lanes by folding them back onto the same small logical tile. That preserves the program's semantics, but it does not automatically turn the spare lanes into a cooperative split over the K reduction dimension.
The important diagnostic pattern was:
-
threadsis much larger than the most important logical lane count. - Generated CUDA contains modulo or bit-mask indexing such as
threadIdx.x & (tile - 1). - Only a subset of lanes commits output, or multiple lanes address the same logical value.
- NCU shows high register pressure and disappointing latency even though occupancy is not the obvious limiter.
T.Parallel(v_tile) expresses the logical iteration space, not a guarantee that extra hardware lanes will be used to split a different dimension.
In this case, the useful computation for each output value required a reduction over dim_k=128. The compiler did not infer that the spare lanes should cooperate on that reduction. Instead, it mapped 32 lanes onto 8 V lanes. The resulting code was semantically valid, but it caused repeated lane work and large live state per lane.
The optimized kernel made the lane decomposition explicit. Instead of hoping the lowering would assign spare lanes to useful reduction work, each lane got a defined role:
const int lane = threadIdx.x;
const int v_lane = lane / kGroupSize;
const int k_rank = lane - v_lane * kGroupSize;
const int v_idx = vid * kVTile + v_lane;
const int k_begin = k_rank * kKChunk;
float old_partial = 0.0f;
for (int ii = 0; ii < kKChunk; ++ii) {
const int kk = k_begin + ii;
old_partial += h_val * k_value;
}
float old_val = old_partial;
for (int offset = kGroupSize / 2; offset > 0; offset >>= 1) {
old_val += __shfl_down_sync(0xffffffff, old_val, offset, kGroupSize);
}The production fast path used:
v_tile=16raw_group_size=2threads=32- one warp per block
- two lanes cooperatively reduce one V element
The explicit lane-group mapping improved latency even though achieved occupancy did not increase:
| Metric | TileLang FLA-style prototype | Explicit lane-group kernel |
|---|---|---|
| Grid blocks | 512 | 256 |
| Registers / thread | 255 | 104 |
| Theoretical active warps / SM | 8 | 16 |
| Theoretical occupancy | 12.5% | 25.0% |
| Achieved active warps / SM | 3.80 | 1.91 |
| Achieved occupancy | 5.94% | 2.99% |
| NCU duration | 7.62 us | 4.74 us |
| End-to-end latency | 0.006636 ms | 0.003818 ms |
The speedup came from better warp-internal work decomposition, less repeated work, and lower register pressure. It did not come from simply launching more active warps.
-
Inspect generated CUDA when
threadsexceeds the logicalT.Parallelspace. A correct TileLang program can still lower to a lane mapping that wastes hardware lanes. -
Modulo or bit-mask lane indexing is a smell for small logical tiles. Patterns like
threadIdx.x & 7orthreadIdx.x % 8should trigger a check: are the extra lanes doing useful cooperative work, or are they repeating the same logical lane? -
T.Parallel(v_tile)does not imply automatic reduction splitting. If spare lanes should splitdim_k, express that explicitly with lane roles and warp-level reductions. -
Occupancy is not the whole story. The optimized kernel had lower achieved occupancy but faster latency because each lane carried less live state and performed more useful work.
-
For small-tile decode kernels, validate the lowering as part of tuning. Correctness and benchmark numbers are necessary, but generated code inspection or NCU evidence is often needed to explain why a seemingly natural TileLang expression underperforms.
Weekly Reports
Maintain Dashboard
AI Knowledge Base