Skip to content

TileLang Case Studies

SuperAngGao edited this page Jun 15, 2026 · 5 revisions

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.


Index

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

Case A: num_stages=0 Sentinel

Date: 2026-02-27 Skill: kernel-debug

Background

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.

Symptoms

  • Shape under test: (n=18432, k=7168, fp16, tune=True)
  • max_err = 0.1250.25 across 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

Investigation

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.

Root Cause

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.

Fix

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

Lessons

  1. Identical max_err across all configs → the design concept is wrong, not an implementation detail. A config-level bug (e.g., wrong index formula for a specific block_n) would show selective failures. A design-level flaw breaks everything uniformly. Fix the concept before debugging details.

  2. Never use num_stages=0 as a sentinel for T.Pipelined. The API requires num_stages >= 1. Use separate Python-level flags to switch between code paths if needed, but keep TileLang primitives semantically valid.

  3. 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.


Case B: GEMV Uncoalesced B Access

Date: 2026-02-27 Skill: tune-multiplication Issue: tile-ai/TileOPs#232

Background

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.

Symptoms

  • 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

Root Cause

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.

Fix (O1)

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.

Benchmark Results (H200, fp16, tuned)

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

Lessons

  1. 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.

  2. Always use a full warp (32 threads) for the reduction dimension. reduce_threads < 32 causes cross-row accesses within a warp, breaking coalescing. TileLang's tvm_thread_allreduce maps to efficient warp shuffles when reduce_threads=32.

  3. The 2.9× gain at small shapes reflects torch dispatch 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.


Case B.1: O3 Pipelined B Loads

Date: 2026-02-27 Continuation of: Case B

Motivation

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.

Implementation

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).

Benchmark Results (H200, tuned)

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.

Surprising Finding

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.async setup cost per tile, and potential bank conflicts in b_shared. For very small block_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.

Lessons

  1. 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.

  2. disable_tma=True is required for T.copy into manually-indexed shared memory on SM90. TMA requires layout inference that is only available for WGMMA kernels. For all other kernels, use cp.async via disable_tma=True. See Anti-Patterns #12.


Case B.2: O2 Shared Memory for a

Date: 2026-02-27 Continuation of: Case B

Motivation

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.

Results (O1+O2 vs O1 only, H200, fp16)

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.

Why It Did Not Help

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.

Lessons

  1. 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.

  2. Always measure before assuming an optimization helps. The theoretical savings (reducing a traffic by 7/8) looked compelling on paper but yielded nothing in practice on H200.

  3. T.sync_threads() has non-zero cost. On fast kernels or with small block_n, the synchronization overhead can dominate any data-traffic savings.


Case B.3: Large Shape Validation

Date: 2026-02-27 Continuation of: Case B

Motivation

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.

Shapes Added

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)

Benchmark Results (benchmarks/ops/bench_gemv.py, H200, tuned)

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%

Key Trend

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.

Lessons

  1. 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.

  2. Always validate on production-scale shapes before finalizing a kernel. An optimization that looks good at n=7168 may behave differently at n=57344 due to occupancy changes.


Case B.4: forward() Python Closure Overhead

Date: 2026-02-27 Continuation of: Case B

Symptom

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.

Root Cause

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:

  1. Creates a new Python function object (closure) capturing n, k, dtype
  2. Looks up the JIT cache by key (involves hashing)
  3. 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.

Fix

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.

Lessons

  1. 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 with time.time() or datetime.now().

  2. 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__.


Case C: nsys Profiling with Autotune Contamination

Date: 2026-02-27 Skill: tune Shape: (n=7168, k=16384, fp16, H200)

Background

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.

Problem: Contaminated Trace (Autotune Enabled)

 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

Fix: Disable Autotune, Fix Best Config

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

Result: Clean Trace

 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 total3.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.

Why 74% and Not Higher

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).

Lessons

  1. Always disable autotune and fix the best config before running nsys. The trace must contain only steady-state launches to be interpretable.

  2. 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).

  3. cuBLAS split-K issues two kernels. When you see a splitK + reduce pair, sum their latencies for the true cuBLAS cost.

  4. Use do_bench (CUDA events / CUPTI) for the authoritative benchmark number. nsys is for understanding why; do_bench is for comparing how fast.

See also: Hardware Constraints — H200 Bandwidth.


Case D: Cross-Pipeline Shared Memory Race in GLA Decode

Date: 2026-03-17 Kernel: tileops/kernels/linear_attn/gla/gla_decode.pyGLADecodeKernel (TC path) Issue: tile-ai/TileOPs#550

Background

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 output o via fused matvec
  • Pass 2: T.Pipelined — computes new_state via element-wise state update

Both passes used the same shared memory buffer h_tile for TMA-pipelined state tile prefetch.

Symptoms

  • 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

Investigation

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.

Root Cause

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

Fix

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.

Results

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)

Lessons

  1. Never share a T.alloc_shared buffer across two T.Pipelined loops in a warpgroup-specialized kernel. TileLang's T.Pipelined generates 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.

  2. 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.

  3. 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.

  4. compute-sanitizer does 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 (via kernel.get_kernel_source()) and systematic empirical bisection are the only reliable debugging approaches.

  5. Bisect by replacing T.Pipelined with T.Serial and T.gemm with scalar T.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.


Case E: Builder Function lru_cache — Eliminating 28–200ms Per-Call Overhead

Date: 2026-03-19 PRs: #597, #600 Issues: #571, #599

Background

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).

Symptoms

  • 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

Root Cause

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.

Fix

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 kernel

With the decorator, the second call with the same (n, k, dtype) returns the cached JITImpl directly — skipping TIR generation, serialization, and SHA-256 entirely.

Safety Analysis

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

Benchmark (gated_deltanet forward, 3 builders)

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

Scale of the Problem

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.

Relationship to Case B.4

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).

Lessons

  1. 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 in DEVELOPMENT.md.

  2. Distinguish builder functions from dispatch wrappers. _*_kernel(shape_params) is the builder (cacheable, hashable params). @torch.library.custom_op is the dispatch wrapper (takes tensors, not cacheable). Only the former gets lru_cache.

  3. 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.

  4. Coding standards prevent accumulation. This issue grew from 0 to 51 files over months because nothing enforced the pattern. Adding it to DEVELOPMENT.md and .claude/rules/code-style.md ensures new kernels include the decorator from day one.

  5. maxsize=32 is 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).


Case F: DeltaNet Backward Correctness and Performance — From Wrong to Faster than FLA

Date: 2026-03-23 Skills: kernel-debug, perf

Background

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.

Symptoms

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 dk and dbeta failed; dq and dv passed 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)

Investigation

Step 1: Isolate which dk component is wrong

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.

Step 2: Three missing gradient paths identified

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^TdPdk_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.

Step 3: Why gated DeltaNet wasn't affected

The gated variant has the same missing paths but its tests passed because:

  1. The gate decay exp(g) < 1 naturally dampens error propagation across chunks
  2. The A_inv gradient contribution is smaller when weighted by the gate

Fix (Correctness)

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_dh

2. dw cross-chunk correctiondu_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.

Fix (Performance)

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%.

Results

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.

Lessons

  1. When porting gated → ungated, re-derive the adjoint equations from scratch. Removing the gate changes ∂h_{c+1}/∂h_c from exp(g) · I - k_scaled^T @ P to I - k^T @ w. The coefficient I (simple accumulation) is wrong for the ungated case — the v_new = u - w @ h dependency creates a coupling term that doesn't exist in the gated version because the gate decay dampens it.

  2. 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_inv depends on k and beta through the Gram matrix — the discarded dAw/dAu should have flowed back.

  3. 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.

  4. TileLang fragment layout conflicts. Reusing a fragment in two T.gemm calls with different transpose patterns (T.gemm(A, B, frag) then T.gemm(C, D, frag, transpose_A=True)) causes InternalError: Get different layout for frag. Solution: allocate a fresh fragment for the second gemm. This is a TileLang-specific pitfall not documented elsewhere.

  5. 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.

  6. Forward should save intermediate results for backward. The fused_prepare_compute_w_u kernel computes w, u in the forward pass, but the original design discarded them (only saving Aw, Au, S). The backward then recomputed them from scratch. Passing w, u through autograd.Function.ctx saved ~8% of backward time at zero forward cost.

  7. Autotune has diminishing returns when the bottleneck is algorithmic. The sequential dh recurrence runs on a grid of B×H (= 8 threads), iterating NC times. No thread/stage tuning can parallelize this loop — the bottleneck is the W^T @ K @ dh gemm per step, which is mathematically required for correctness. Autotune is most effective for parallel kernels where thread count and pipeline depth affect occupancy.

  8. Neumann series vs forward substitution trade-off. TileOPs uses Neumann series (A_inv ≈ I + P + P² + ...) for the WY inverse because it maps naturally to T.gemm in log₂(chunk_size) rounds. FLA uses forward substitution (sequential row solve). The Neumann approach enables GPU parallelism but requires an explicit A_inv backward; forward substitution is inherently differentiable through autograd. This design choice is the root cause of the performance gap at long sequences.


Case G: Lane Folding in Small Logical Tiles

Date: 2026-06-15 Skill: kernel-debug, perf

Background

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=32
  • dim_k=128
  • dim_v=128
  • v_tile=8
  • bf16

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]

Symptoms

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.

Investigation

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:

  • threads is 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.

Root Cause

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.

Fix

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=16
  • raw_group_size=2
  • threads=32
  • one warp per block
  • two lanes cooperatively reduce one V element

Results

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.

Lessons

  1. Inspect generated CUDA when threads exceeds the logical T.Parallel space. A correct TileLang program can still lower to a lane mapping that wastes hardware lanes.

  2. Modulo or bit-mask lane indexing is a smell for small logical tiles. Patterns like threadIdx.x & 7 or threadIdx.x % 8 should trigger a check: are the extra lanes doing useful cooperative work, or are they repeating the same logical lane?

  3. T.Parallel(v_tile) does not imply automatic reduction splitting. If spare lanes should split dim_k, express that explicitly with lane roles and warp-level reductions.

  4. 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.

  5. 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.

Clone this wiki locally