Cycle-level profiling: IntraWG/InterWG overlap breakdown (TileOPs vs FA3)

[superAngGao]

Summary

Cycle-level profiling of the WS GQA forward kernel's IntraWGOverlap + InterWG scheduling pipeline, measured on H200 (locked freq, GPU 7). Shape: B=4 S=4096 H=64 Hkv=8 D=128 fp16, block_m=128 block_n=128, non-causal.

Goal: quantify where time goes inside the steady-state K-loop iteration, and compare with FA3.

Measurement methods

1. clock64 instrumentation (CUDA Core side)

Insert clock64() probes at key points in WG1's steady-state loop via TileLang T.call_extern("int64", "clk::read_clock"), accumulate via atomicAdd across all CTAs.

Probed regions in the original pipeline:

t0 → issue QK[N]  (8× wgmma_ss)
     rescale O
     barrier_wait(v_full)
     issue PV[N-1] (8× wgmma_rs)
t1 → wait_wgmma<1>  (QK done)
t2 → softmax (reduce_max, exp2f, reduce_sum, rescale ls)
t3 → wait_wgmma<0>  (PV done)
t4

2. NCU metrics (Tensor Core utilization)

ncu --kernel-name regex:"main_kernel" \
  --metrics sm__pipe_tensor_op_hmma_cycles_active.avg.pct_of_peak_sustained_elapsed,\
            sm__pipe_tensor_op_hmma_cycles_active.avg,\
            sm__cycles_elapsed.avg \
  --launch-skip 5 --launch-count 1 python3 bench.py

3. In-pipeline PV measurement

Move wait_wgmma<0> before softmax (kills IntraWGOverlap for measurement only) to directly time PV execution on TC:

issue QK → issue PV → wait<1> → wait<0> → softmax
                                 ↑ TC PV = this wait

TileOPs results (threadbind kernel, _test_ws_fa3_v2_threadbind.py)

Full pipeline (clock64)

Region	Cycles	%
CC issue (QK+rescale+waitV+PV)	657	31.7%
wait<1> (QK done on TC)	59	2.9%
softmax	1347	65.1%
wait<0> (PV residual)	7	0.3%
TOTAL	2070	100%

wait<0> ≈ 0 → PV finishes before softmax ends → softmax is the pacing bottleneck.

TC GEMM times

Metric	Value	Source
TC utilization	65%	NCU hmma_cycles_active.pct_of_peak_sustained_elapsed
TC PV (wgmma_rs)	450 cyc	In-pipeline wait<0> measurement
TC QK (wgmma_ss)	≈222 cyc	Derived: (2070 × 0.65)/2 − 450
TC per WG (QK+PV)	≈672 cyc
TC idle per iter	≈726 cyc	During softmax, no WGMMA issued

NCU stall breakdown

Stall	Ratio
long_scoreboard	3.09
wait	1.57
short_scoreboard	0.32
math_pipe_throttle	0.02

Timeline diagram

Three rows: WG1 CC, WG2 CC, shared Tensor Core.

• Red (softmax) dominates CC time at 65%
• TC has 726-cycle idle gap each iteration while both WGs do softmax
• Scheduler barrier (yellow arrows) serializes WG1/WG2 WGMMA issue

Key findings

1. Softmax is the bottleneck (65% of loop, 1347 cycles). TC is 65% utilized with 35% idle time.
2. IntraWGOverlap works correctly: PV GEMM executes on TC in parallel with softmax on CC. wait<0> residual ≈ 0.
3. InterWG overlap works: scheduler barrier alternates WGMMA issue between WG1/WG2. One WG's softmax overlaps with the other's WGMMA.
4. Softmax bottleneck comes from AllReduce (cross-warp reduce_max + reduce_sum with named barrier sync).
5. long_scoreboard = 3.09 (vs FA3's 0.19) — this is the K/V TMA bandwidth gap, separate from the softmax bottleneck. FA3 uses cluster multicast TMA to halve this.

TODO: FA3 comparison

Run the same NCU + clock analysis on FA3's kernel for the same shape to get:

• FA3 TC utilization
• FA3 softmax time
• FA3 stall breakdown
• Side-by-side comparison to identify where the remaining 6-16% gap comes from

Files

• _bench_clock_intra_wg.py — full pipeline clock64 instrumentation
• _bench_clock_tc_gemm.py — no-softmax TC GEMM measurement (dual WG)
• _bench_clock_tc_single_wg.py — single WG TC GEMM (no contention)
• _bench_clock_tc_in_pipeline.py — in-pipeline PV measurement (wait<0> before softmax)
• _plot_intra_wg_timeline.py — timeline visualization
• _intra_wg_timeline.png — output diagram

All on branch fix/ws-fa3-v2-epilogue-fence.

[superAngGao]

FA3 NCU comparison (same shape, non-causal)

FA3 kernel: FlashAttnFwdSm90, block_m=128, block_n=176, cluster=2×1×1, persistent (132 SMs), MmaPV_is_RS=true, IntraWGOverlap=true

Side-by-side NCU metrics

Metric	TileOPs (bn=128)	FA3 (bn=176)	Notes
Grid	(32,64,4)×384	(132,1,1)×384	FA3 persistent
block_n	128	176
Cluster	1	2	FA3 multicast
SM elapsed	6,262,542	4,985,248	FA3 20% faster
TC util	65.0%	84.1%	+19pp
FMA pipe util	13.9%	15.6%
FMA cycles	6,981,322	6,223,305	ours 12% more
ALU pipe util	18.8%	19.1%
long_scoreboard	3.09	0.19	16× — multicast TMA
short_scoreboard	0.32	0.24
math_pipe_throttle	0.02	0.03
barrier	—	0.80	cluster barrier
wait	1.57	1.67	wgmma wait
mio_throttle	—	0.15

Instruction mix (per smsp avg)

Pipe	TileOPs	FA3	Ratio
Total inst	2,231,622	1,839,980	1.21×
FMA inst	872,541	777,509	1.12×
Tensor inst	63,550	56,599	1.12×
LSU inst	33,963	14,804	2.29×

Analysis

1. TC util gap (+19pp): FA3's block_n=176 increases GEMM compute per iter by (176/128)=1.375× while softmax scales sub-linearly. This shifts the GEMM/softmax ratio in favor of TC → less idle time.

2. long_scoreboard gap (16×): FA3 uses cluster=2 multicast TMA, halving K/V DRAM bandwidth. However, TMA latency is NOT on the critical path — it runs on the producer WG in parallel with consumer's softmax (1347 cyc). The high long_scoreboard only impacts producer stalls, not the consumer's pacing bottleneck.

3. FMA cycles (softmax proxy): Our 6.98M vs FA3's 6.22M. Per K-iter, FA3 executes more FMA (larger block_n), but fewer total iters (24 vs 32). Net effect: FA3 spends ~12% fewer FMA cycles total.

4. LSU inst 2.29×: We issue 2× more TMA load commands per K/V block (D=128 split into 2×64 due to TileLang swizzle limit). FA3's CUTLASS handles this with a single multi-dimensional TMA descriptor.

Conclusion

FA3's 20% speed advantage comes primarily from:

• block_n=176 → better TC util (84% vs 65%), fewer iterations
• Multicast TMA → lower long_scoreboard, but NOT currently on critical path

Our bottleneck remains softmax (65% of loop) → optimize softmax first, then multicast TMA unlocks further gains.

[superAngGao]

Multicast TMA: unblocked, root cause, and L2 analysis

Root cause of the hang

The cluster kernel (_test_ws_fa3_v2_tb_cluster.py) hung because T.dec_max_nreg(24) was commented out in the producer WG (line 363):

# T.dec_max_nreg(24)  # DEBUG: skip register dealloc  ← THIS

setmaxnreg.dec not executing → producer doesn't release registers → consumer's setmaxnreg.inc(240) blocks forever waiting for registers → deadlock.

Fix: uncomment the line. Kernel now passes correctness (diff=0.0001).

Named barrier → mbarrier replacement

Also replaced ALL named barriers (bar.sync/bar.arrive) with mbarrier (T.alloc_barrier + T.barrier_wait/T.barrier_arrive) to eliminate potential named-barrier ↔ cluster-barrier interaction. File: _test_ws_fa3_v2_tb_cluster_mbar.py.

The original cluster kernel also had a broken scheduler pattern — both WG1 and WG2 used barrier_id=1 for both sync and arrive (should be cross: WG1 sync=1/arrive=2, WG2 sync=2/arrive=1). The mbarrier version fixes this with explicit sched_12 (WG1→WG2) and sched_21 (WG2→WG1) barriers.

NCU results: multicast doesn't help for this shape

Metric	No multicast (bn=128)	Multicast (bn=128)	FA3 (bn=176)
SM cycles	6,261,355	6,519,406	4,985,248
TC util	65.0%	62.4%	84.1%
long_scoreboard	3.09	3.39	0.19
DRAM sectors	10,488,792	—	10,487,592

Multicast didn't reduce long_scoreboard (3.39 vs 3.09). Cluster overhead slightly increased it.

DRAM analysis: bandwidth is not the bottleneck

Kernel	DRAM sectors read	long_scoreboard
Threadbind (non-persistent)	10,488,792	3.09
Our persistent (bn=128)	16,780,680 (+60%!)	4.76
FA3 persistent (bn=176)	10,487,592	0.19

Key findings:

1. FA3 and our threadbind read identical DRAM (~10.49M sectors) — same total HBM traffic
2. Our persistent reads 60% more — likely a T.Persistent tile ordering or codegen issue
3. FA3's 16× lower long_scoreboard is NOT from less DRAM bandwidth — it's from better latency hiding

Why multicast doesn't help

GQA with H=64, Hkv=8: 8 Q heads share each KV head. K+V per head = 2 MB. Total K+V = 64 MB ≈ H200 L2 (60 MB). Most K/V accesses are L2 hits regardless of multicast. The ~10.49M DRAM sectors represent cold misses that both FA3 and we pay equally.

FA3's tile scheduling

FA3 non-causal uses StaticPersistentTileScheduler — simple linear stride (same as ours):

tile_idx = blockIdx.x + iter * gridDim.x
(m_block, head, batch) = unravel(tile_idx)  // M fastest

FA3 does NOT use L2 swizzle for non-causal. The DynamicPersistentTileScheduler (with L2-aware sectioning) is only used for causal.

Remaining gap: latency hiding

Same DRAM, same tile ordering, but FA3 long_scoreboard = 0.19 vs ours 3.09. Hypotheses:

1. EVICT_LAST cache hint on TMA loads — FA3 uses TMA::CacheHintSm90::EVICT_LAST, we don't
2. bn=176 — 25% fewer iterations → fewer TMA load operations → less TMA unit contention
3. Better TMA/compute overlap — FA3's CUTLASS code may schedule TMA loads more efficiently

Next step

Test EVICT_LAST cache hint on our kernel to see if it reduces long_scoreboard.

Files

• _test_ws_fa3_v2_tb_cluster_mbar.py — cluster + multicast + mbarrier scheduler (PASS, 556.8 TFLOPS)
• Original hang root cause: dec_max_nreg commented out in _test_ws_fa3_v2_tb_cluster.py:363

[superAngGao]

Corrected benchmark: all-in-one session (2026-04-13, GPU 7, locked 1830 MHz, 32°C)

Previous measurements compared numbers from different sessions. This run measures ALL configs in the same Python process with identical warmup.

Shape: B=4 S=4096 H=64 Hkv=8 D=128 fp16 non-causal, 50 warmup + 200 timed iterations.

Config	TFLOPS	% FA3
Threadbind bn=128 (no cluster)	539.6	83.3%
Threadbind bn=176 (no cluster)	560.2	86.5%
Cluster+multicast bn=128	529.2	81.7%
Cluster+multicast bn=176	551.6	85.2%
FA3 (bn=176, cluster=2, persistent)	647.5	100%

Key findings

1. bn=176 vs 128: +4% (560 vs 540). Fewer iterations (24 vs 32), better TC utilization.
2. Multicast cluster: −1.5% vs threadbind (551 vs 560). Cluster overhead (cluster_sync, mbarrier complexity) outweighs L2 savings.
3. Gap to FA3: ~15% (560 vs 648). Previous report of 94.6% was from comparing different sessions with different conditions.

Memory hierarchy analysis (NCU, same session)

Metric	Ours bn=128	Ours bn=176	Ours multicast bn=128	FA3 bn=176
DRAM read	335.6 MB	335.7 MB	335.8 MB	335.6 MB
L2 total	12.16 GB	11.55 GB	9.52 GB	10.15 GB
L2 read sectors	366.9M	348.0M	283.1M	291.5M
long_scoreboard	3.09	3.14	3.39	0.19
TC util	65.0%	70.1%	62.4%	84.1%

• DRAM reads identical across all configs (~336 MB). HBM bandwidth is NOT the bottleneck.
• Multicast reduces L2 traffic by 23% (367M→283M sectors), even below FA3 (291M). But no perf gain — L2 is not the bottleneck either.
• EVICT_LAST cache hint: tested, zero effect on long_scoreboard (3.09→3.09).
• long_scoreboard 3.09 vs FA3 0.19: mostly a statistical artifact — our 4 producer warps all actively poll barriers and show up as stalled. FA3 sleeps 3 of its 4 producer warps (sleeping=0.34), removing them from stall statistics. Actual compute impact is minimal since stalled warps don't steal scheduler slots.

Cycle breakdown (clock64 instrumentation)

Region	Cycles	% of iter
CC issue (QK+rescale+waitV+PV)	657	31.7%
wait<1> (QK done)	59	2.9%
softmax	1347	65.1%
wait<0> (PV residual)	7	0.3%
Total	2070	100%

TC GEMM (measured): QK ≈ 222 cyc, PV = 450 cyc. TC util = 65% (NCU).

Softmax dominates loop time. AllReduce implementation is identical to FA3 (quad shfl, no barrier).

What we ruled out

Hypothesis	Result
Softmax reduce implementation	Identical (quad shfl)
Multicast TMA saves DRAM bandwidth	DRAM identical, no effect
Multicast TMA saves L2 traffic	L2 reduced 23%, no perf effect
EVICT_LAST cache hint	Zero effect
Tile scheduling (L2 swizzle)	FA3 non-causal doesn't use it
long_scoreboard = real bottleneck	Mostly statistical (producer warp polling)

Remaining 15% gap to FA3

The gap is likely from:

1. Code quality: TileLang TIR→CUDA→nvcc pipeline vs CUTLASS hand-tuned C++ templates. FA3's CUTLASS code has tighter instruction scheduling, fewer register spills, better pipeline interleaving.
2. Producer thread efficiency: FA3 uses 1 warp (32 threads) for TMA, other 3 producer warps sleep. We use all 128 threads actively polling.
3. Persistent scheduling overhead: FA3 uses StaticPersistentTileScheduler with minimal overhead. Our T.Persistent may have higher per-tile overhead.
4. 2× TMA loads per K/V block: TileLang splits D=128 into 2×64 (swizzle limit). FA3's CUTLASS TMA descriptor may handle this in a single load.

Files

• _bench_clock_intra_wg.py — full pipeline clock64
• _bench_clock_tc_gemm.py — TC GEMM timing (no softmax)
• _bench_clock_tc_single_wg.py — single WG TC (no contention)
• _bench_clock_tc_in_pipeline.py — in-pipeline PV measurement
• _plot_intra_wg_timeline.py — timeline visualization
• _test_ws_fa3_v2_tb_cluster_mbar.py — cluster+multicast+mbarrier (PASS)

[superAngGao]

Optimization attempt: move clear(acc_s) off critical path

Hypothesis

Between barrier_wait(k_full) and the first wgmma_ss, our kernel has extra work that FA3 doesn't:

Step	Ours	FA3
Zero-init acc_s	T.clear(acc_s) — 64 reg writes	ScaleOut::Zero on first wgmma (HW zero-init, free)
Q descriptor	Re-computed every iter (initialize_wgmma_descriptor)	Pre-computed once outside K-loop
K descriptor	Re-computed every iter (2 branches × 2 descs)	Pre-computed tensor partition, indexed by stage

What we changed

Moved T.clear(acc_s) from after barrier_wait (critical path) to end of previous iteration (overlaps with next TMA load). File: _test_ws_fa3_v2_threadbind_opt.py.

Result

Original threadbind bn=176:                562.2 TFLOPS = 87.4% FA3
Optimized (clear off critical path) bn=176: 566.3 TFLOPS = 88.0% FA3  (+0.7%)
FA3:                                        643.3 TFLOPS = 100%

Correctness: diff=0.0051 PASS.

Limitations

1. Descriptor re-init not fixed: TileLang's T.wgmma_gemm codegen always emits initialize_wgmma_descriptor before each WGMMA batch. Q descriptor is redundantly recomputed every iteration (Q smem address never changes). 8 descriptor inits per steady-state iteration.

2. No HW zero-init: TileLang doesn't expose wgmma's ScaleOut::Zero (scale_d=0) parameter. The clear is still there, just moved off the critical path.

Remaining gap analysis

+0.7% from clear relocation. 12% gap remains. The descriptor overhead is ~5 register ops each × 4 per iter = ~20 cycles — too small to explain 12%.

Next step: SASS-level comparison between our kernel and FA3 to identify the actual instruction scheduling differences.

[superAngGao]

SASS analysis & loop unrolling investigation

SASS code size comparison

	TileOPs (bn=176)	FA3 (bn=176)	Ratio
SASS lines	35,053	5,493	6.4×
WGMMA	21	8	2.6×
MUFU (exp2f)	490	182	2.7×
FFMA	1,072	180	6.0×
Instructions executed (per SM)	8,141,801	7,359,864	1.11×
I-cache requests	5,197,764	626,174	8.3×
I-cache hit rate	99.91%	99.97%	OK

Root cause: nvcc loop unrolling

Our CUDA has K-loop: for (int n_idx = 0; n_idx < 24; ++n_idx). nvcc at -O3 fully unrolls this → 24 copies of the loop body → 35K SASS.

FA3 uses #pragma unroll 1 on all K-loops → 5.5K SASS.

Does unrolling hurt?

No. Testing with #pragma unroll 1 injected into our CUDA:

• With unroll: 563.4 TFLOPS (35K SASS)
• Without unroll: 549.9 TFLOPS (7.5K SASS compiled manually)

Unrolling is FASTER because it eliminates branch overhead and enables better instruction scheduling. I-cache hit rate is 99.91% either way.

FA3 softmax/rescale alignment attempt

Aligned three differences with FA3:

1. Simplified softmax (removed fill -inf, monotonic max, NaN guard)
2. Moved rescale_o to after wait<0> + release V (FA3 position)
3. Moved clear(acc_s) off critical path

Result: ~0% improvement (562.4 vs 563.4).

Conclusion

The 12% gap (560 vs 644 TFLOPS) comes from 11% more executed instructions (8.1M vs 7.4M per SM). This is structural to TileLang's codegen:

• if/else branches for double-buffer stage selection (both paths compiled)
• Redundant descriptor re-initialization every iteration
• No wgmma ScaleOut::Zero (explicit register clear instead)

These cannot be fixed without modifying TileLang's wgmma codegen.

[superAngGao]

Persistent kernel regression: 10-20% slower than threadbind

Data (all bn=128, same session)

S	B	Threadbind TFLOPS	Persistent TFLOPS	FA3 TFLOPS	TB%	PS%
4096	4	568.5	506.8	663.8	85.6%	76.4%
8192	1	595.4	514.3	675.5	88.1%	76.1%
16384	1	585.7	461.7	682.5	85.8%	67.6%
32768	1	524.5	457.9	669.6	78.3%	68.4%

Persistent is 10-20% slower than threadbind and getting worse with longer seqlen.

NCU diagnosis (B=4 S=4096)

Metric	Threadbind	Persistent	FA3
DRAM sectors read	10,488,792	16,780,680 (+60%)	10,487,592
long_scoreboard	3.09	4.76	0.19

Persistent reads 60% more DRAM. Excess ≈ 201 MB ≈ 75% of Q data (268 MB).

Root cause analysis

1. Q load per tile: Persistent loads Q via TMA inside the persistent loop (per-tile). Threadbind loads Q via T.copy once at kernel start.

2. Tile ordering: SM 0 visits tiles (b=0,h=0,m=0), (b=0,h=4,m=4), (b=0,h=8,m=8)... — Q address jumps 8.4 MB per wave across 268 MB address space. 132 SMs simultaneously scan different parts → L2 thrashing.

3. EVICT_FIRST on Q: Tested, no effect on DRAM (16,780,080 vs 16,780,680). Not a cache hint issue.

4. T.Persistent overhead: Per-tile Q load + acc_o clear + ls clear + sm fill + output store + lse write. These are fixed costs per tile that threadbind doesn't have (or has once per kernel).

TODO

Investigate T.Persistent's codegen for hidden synchronization or redundant global memory accesses at tile boundaries.

[superAngGao]

T.Persistent DRAM regression: root cause isolated

Key finding

Even with 1 SM (NUM_SMS=1, zero L2 contention), persistent reads 545 MB vs threadbind's 336 MB. The 60% excess DRAM is inherent to T.Persistent's codegen, NOT from L2 contention.

NUM_SMS	Persistent DRAM read
1	544.8 MB
8	538.5 MB
32	537.2 MB
132	537.0 MB
Threadbind (8192 CTAs)	336.0 MB

What we ruled out

Hypothesis	Test	Result
Q TMA load polluting L2	Removed Q loads	DRAM unchanged (537 MB)
Output TMA store polluting L2	Removed output stores	DRAM unchanged
lse writes	Removed lse writes	DRAM unchanged
EVICT_FIRST cache hint	Patched all TMA → EVICT_FIRST (verified in PTX)	DRAM unchanged
L2 contention between SMs	NUM_SMS=1	DRAM still 545 MB
Tile ordering	Same ordering as threadbind	N/A
TMA descriptor	Identical params verified	N/A
PTX instructions	Same cp.async.bulk, same cache_hint, same expect_tx	N/A

What's different

The ONLY difference between threadbind and persistent K/V TMA loads:

1. Smem base address: threadbind at offset 16384, persistent at offset 0
2. Coordinate computation: threadbind uses blockIdx.y/z (hardware registers), persistent computes from w * 132 + blockIdx.x (runtime arithmetic)
3. Loop structure: threadbind has a flat for-loop, persistent has a nested persistent-loop with break conditions

Conclusion

T.Persistent introduces a 60% DRAM read regression that is independent of L2 contention. This is likely a TileLang codegen issue where the persistent loop structure prevents nvcc from optimizing the TMA load pattern for L2 reuse. Filed as a TileLang framework issue.

Impact

Persistent is 10-20% slower than threadbind due to this regression. For now, threadbind is the better choice for the GQA kernel. Threadbind achieves 85-100% FA3 depending on sequence length.

[superAngGao]

Persistent DRAM regression: root cause confirmed — tile ordering breaks K/V L2 reuse

Bisection result

Config	Tiles	TB DRAM	PS DRAM	Ratio
B=1 S=4096 H=8 Hkv=1	256	10.57 MB	10.56 MB	1.00×
B=1 S=4096 H=64 Hkv=8	2048	83.98 MB	128.62 MB	1.53×
B=4 S=4096 H=8 Hkv=1	1024	42.03 MB	42.47 MB	1.01×
B=4 S=4096 H=64 Hkv=8	8192	335.78 MB	536.95 MB	1.60×

DRAM regression only appears with Hkv > 1 (multiple KV heads). When all Q heads share a single KV head (Hkv=1), persistent and threadbind have identical DRAM.

Root cause

T.Persistent([B, H, M_blocks]) with tile ordering: M varies fastest → H → B.

With stride=132 across 132 SMs:

• SM 0 processes tiles 0, 132, 264, ...
• tile 0: (b=0, h=0, m=0) → head_kv=0
• tile 132: (b=0, h=4, m=4) → head_kv=0
• tile 264: (b=0, h=8, m=8) → head_kv=1 ← switches KV head after ~2 tiles

Each SM switches KV head every ~2 tiles. K/V for old KV head gets evicted from L2 before all 8 Q heads finish using it. With 8 KV heads × 4 batches = 32 KV head groups, total K+V unique = 67 MB ≈ L2 size (60 MB). L2 thrashing causes K/V to be re-read from DRAM.

Threadbind grid (M_blocks, H, B) schedules CTAs in order: all CTAs for h=0..7 (head_kv=0) run before h=8..15 (head_kv=1). CUDA runtime naturally groups same-KV-head CTAs together → K/V stays in L2.

Previous dead ends ruled out

• EVICT_FIRST cache hint: no effect (verified in PTX)
• Removing Q loads / output stores / lse writes: no effect
• 1 SM test: same regression (not L2 contention)
• PTX/SASS comparison: identical TMA instructions, same cache hint, same descriptor
• Register spills: zero (STL/LDL = 0 in both)
• H=8 Hkv=1 (single KV head): no regression

Fix

Change tile ordering to keep same KV head tiles together:
T.Persistent([B, Hkv, groups, M_blocks]) or equivalent reordering where M_blocks and groups (Q heads per KV head) vary fastest, keeping Hkv stable per SM.

[superAngGao]

Tile ordering fix: persistent DRAM regression resolved

Fix

Changed T.Persistent([B, H, M_blocks]) to T.Persistent([B, Hkv, M_blocks, groups]).

This keeps same-KV-head tiles together: M_blocks and groups (Q heads per KV head) vary fastest, Hkv varies slowest. K/V data stays in L2 across all Q heads sharing the same KV head.

Results (bn=176)

S	B	Threadbind	PS original	PS reordered	FA3	PS reorder % FA3
4096	4	585.4	483.7	586.7	661.9	88.6%
8192	1	620.9	—	603.6	668.6	90.3%
16384	1	614.6	—	608.0	681.6	89.2%
32768	1	567.9	—	554.1	681.6	81.3%
65536	1	563.3	—	555.0	658.7	84.3%

NCU verification (B=4 S=4096 bn=128):

	PS original	PS reordered	Threadbind	FA3
DRAM read	537 MB	335.6 MB ✓	335.8 MB	335.6 MB

Key outcomes

1. DRAM regression eliminated: 537 → 336 MB (matches theoretical minimum)
2. S=4096: persistent now matches threadbind (586.7 vs 585.4 TFLOPS)
3. Long seqlen: persistent 2-3% behind threadbind (per-tile overhead: Q load, clear, output store)
4. Both ~12% behind FA3 across all seqlen — likely same root cause (per-tile overhead vs FA3's lighter overhead)

File

_test_ws_fa3_v2_persistent_reorder.py — reordered tile ordering fix

Next step

Compare per-tile overhead between our kernel and FA3 to close the remaining 12% gap. This overhead affects both threadbind and persistent equally.

[superAngGao]

Overhead decomposition: per-iter vs per-tile

Method

Measured cycles/tile across multiple seqlen (B=1 bn=176), fit linear model:
cycles_per_tile = overhead + n_iters × per_iter_cost

Raw data

S	K-iters	TB cycles/tile	FA3 cycles/tile	Delta	Delta/iter
2048	12	58,084	49,412	8,672	723
4096	24	110,155	98,691	11,464	478
8192	47	215,069	195,116	19,952	425
16384	94	435,794	398,378	37,416	398

Linear regression

	Per-iter cost	Per-tile overhead
Threadbind	4,620 cycles	~352 cycles
FA3	4,263 cycles	~-3,222 cycles (tile overlap)

Analysis

Per-iter delta: 357 cycles (8.4%) — this is the dominant gap.

Source: TileLang codegen per-iteration overhead:

• Descriptor re-initialization (Q + K + V = 3× per iter)
• clear(acc_s) (88 registers for bn=176)
• if/else double branches for stage selection (only one taken)
• No wgmma ScaleOut::Zero (explicit clear instead)
• 11% more executed instructions (8.14M vs 7.36M per SM)

FA3's negative overhead suggests tile-to-tile pipelining (output store overlaps with next Q load + K prefetch), which we don't have.

Summary of the 12% gap

Component	Cycles/tile (24 iters)	% of total
Per-iter codegen overhead	357 × 24 = 8,568	~8.9%
Per-tile overhead	~3,500	~3.6%
Total	~12,000	~12.4%

Path forward

1. Per-iter: Requires TileLang codegen changes (descriptor caching, ScaleOut::Zero, single-branch stage select)
2. Per-tile: Overlap output store with next tile Q load (like FA3's barrier_O pattern)

[superAngGao]

Per-iter clock breakdown: bn=176

Measured body (clock64, WG1 steady-state)

Region	bn=128	bn=176	Ratio
wgmma_issue (QK+rescale+waitV+PV)	657	1034	1.57×
wait<1> (QK done)	59	166	2.81×
softmax	1347	2213	1.64×
wait<0> (PV residual)	7	6	—
Measured total	2070	3419	1.65×

Full per-iter breakdown (B=4 S=4096 bn=176)

Component	Ours	FA3
Full per-iter (NCU derived)	4017 cyc	3347 cyc
Clock measured body	3419 cyc	(not measurable)
Unmeasured (barrier_wait_k + scheduler + overhead)	598 cyc	≤ 3347-body
Delta	+670 cyc/iter (+20%)

Analysis

The 670 cycles/iter delta spans BOTH the measured body and the unmeasured overhead. Since measured body = 3419 and FA3 total = 3347, FA3's body must be significantly shorter than ours.

The per-iter body overhead comes from:

1. Descriptor re-init: 3× initialize_wgmma_descriptor per iter (Q + K + V), each with smem pointer computation. FA3 pre-computes descriptors.
2. if/else stage selection: Both branches compiled, only one taken. Instruction footprint 2× larger, potential i-cache inefficiency.
3. No ScaleOut::Zero: Explicit clear(acc_s) (88 registers for bn=176) vs FA3's hardware zero-init.
4. 11% more total instructions: 8.14M vs 7.36M per SM (NCU confirmed).

These are structural to TileLang's wgmma codegen and cannot be fixed without modifying the framework.

Tile ordering fix summary

The tile ordering fix ([B, Hkv, M_blocks, groups]) resolved the 60% DRAM regression in persistent mode. Combined with the clock analysis, the remaining gap to FA3 is:

• ~9% from per-iter codegen overhead (357 cycles × 24 iters)
• ~3% from per-tile overhead (Q load, acc_o clear, output store)

[superAngGao]

Root cause confirmed: producer K delivery latency

Full per-iter clock breakdown (bn=176, B=4 S=4096)

Region	Cycles	%
barrier_wait(k_full) + scheduler_sync + clear	698	17.2%
wgmma_issue (QK + rescale + waitV + PV)	969	23.9%
wait<1> (QK done)	157	3.9%
softmax	2218	54.8%
wait<0> (PV residual)	5	0.1%
FULL ITER	4048	100%

Comparison with FA3

Component	Ours	FA3	Delta
Pre-body overhead	698 cyc	~0 cyc	~698 cyc
Body (wgmma + softmax)	3349 cyc	~3347 cyc	~0 cyc
Full iter	4048 cyc	3347 cyc	701 cyc

The consumer body (WGMMA + softmax) matches FA3 almost exactly (3349 vs 3347 cycles). The entire 12% gap comes from the 698-cycle overhead before each iteration's compute begins.

Breakdown of the 698 cycles

• barrier_wait(k_full) ≈ 535 cycles: Consumer waiting for producer to deliver K data
• scheduler_sync ≈ 75 cycles: Named barrier sync between WG1 and WG2
• clear(acc_s) ≈ 88 cycles: Zero-init 88 registers (bn=176)

Why producer K delivery is slow

Our producer (128 threads) per K-loop iteration:

1. barrier_wait(k_empty) — 128 threads poll mbarrier (wait for consumer to free buffer)
2. if (gi_kp % 2 == 0) — runtime branch for stage selection
3. T.tma_copy(k[...], k_smem_0/1) — 2× TMA load (D split)
4. barrier_arrive(k_full) — 128 threads arrive

FA3 producer (32 threads, 1 warp):

1. pipeline_k.producer_acquire() — CUTLASS pipeline, 1 elected thread
2. copy(tma_load_K.with(...)) — single TMA copy call (CUTLASS handles D internally)
3. Pipeline commit — integrated with TMA

Key differences:

• 128 vs 32 threads polling barriers: 4× more barrier traffic
• if/else stage selection: runtime branch vs pre-computed pipeline state index
• 2 TMA loads per K block (D=128 split): vs potentially 1 in FA3 (CUTLASS multi-dim TMA)
• Separate barrier_arrive from TMA: vs CUTLASS integrated TMA+barrier

Previous wrong hypothesis

Earlier analysis attributed the 12% gap to "11% more instructions" and "codegen overhead." While instruction count IS higher (8.14M vs 7.36M), the actual compute body takes the same time (3349 ≈ 3347). The extra instructions execute during barrier wait time and don't add to the critical path. The real bottleneck is producer TMA pipeline latency.

Next step

Optimize producer to reduce K delivery latency:

1. Reduce producer thread count (only 1 warp participates, other 3 can sleep)
2. Remove if/else stage branching (use descriptor arithmetic)
3. Investigate single TMA load for full D=128

[superAngGao]

ScaleOut::Zero optimization + pipeline analysis deep-dive (2026-04-13)

1. ScaleOut::Zero (clear_accum=True)

TileLang's T.wgmma_gemm supports clear_accum=True, which uses WGMMA's ScaleOut::Zero (scaleD=0) to zero the accumulator in hardware, eliminating the separate T.clear(acc_s).

Change: Replace T.clear(acc_s) + T.wgmma_gemm(Q, K, acc_s) with T.wgmma_gemm(Q, K, acc_s, clear_accum=True) for QK GEMM in both WG1 and WG2.

Result (non-causal, bn=176):

Shape	Original TF	ScaleOut TF	Speedup
B=4 S=4096	584.0	598.6	1.025x
B=1 S=8192	597.0	610.2	1.022x
B=1 S=16384	602.2	618.9	1.028x

Consistent 2-3% improvement. Less than the theoretical 10% because:

• clear was partially hidden by warp scheduler (not fully on critical path)
• The real bottleneck is TC utilization, not instruction count

2. Per-instruction stall analysis (NCU source-level)

Used ncu --page source to get per-SASS-instruction stall breakdown.

Key finding: Dominant stall is WARPGROUP.DEPBAR.LE (wait_wgmma), NOT mbarrier polling.

Metric	Original	ScaleOut	FA3
total samples	375,494	370,211	316,995
barrier stall %	15.9%	17.7%	14.5%
wait<0> barrier samples	54,480	58,830	3,859
wait<1> barrier samples	1,371	4,403	9,754

Critical observation: FA3 has 14x fewer wait<0> stalls but 7x more wait<1> stalls compared to us. This indicates a fundamentally different WGMMA pipeline ordering.

3. Root cause: prologue (n_idx==0) full TC stall

Our n_idx==0 (first K-loop iteration per tile):

QK issue → wait_wgmma(0) → softmax  // full ~800 cycle TC stall!

Our n_idx>0 (steady state):

QK issue → rescale → wait_V → PV issue → wait_wgmma(1) → softmax → wait_wgmma(0)
// QK+PV overlap, wait<0> ≈ 5 cycles

The 54,480 wait<0> barrier samples are dominated by the prologue path where wait<0> stalls for the full QK TC time (~800 cycles) with zero overlap.

FA3 avoids this by using ScaleOut::Zero for the first PV WGMMA, enabling QK+PV pipeline from the very first iteration.

4. SCHED_OFF analysis

SCHED_OFF = time between WG1 and WG2 starting WGMMA = time WG2 waits before issuing QK.

	Original	ScaleOut	FA3 (est.)
pre-body overhead	698 cyc (BWAIT+sched+clear)	318 cyc (BWAIT+sched)	~50 cyc
SCHED_OFF	1667 cyc	~1287 cyc	~1019 cyc
TC idle/iter	1210 cyc (30%)	~900 cyc	~562 cyc (16%)
TC util	70.1%	~73%	84.1%

5. Next steps: investigate FA3's WGMMA pipeline ordering

FA3's stall pattern (low wait<0>, high wait<1>) suggests a different WGMMA commit/wait ordering:

Our pipeline (QK-first commit):

QK commit (group N) → PV commit (group N+1)
wait<1>: waits for QK (oldest) → PV still flies during softmax
wait<0>: waits for PV → usually done (5 cyc), but prologue stalls ~800 cyc

FA3's pipeline (possibly PV-first commit or different grouping):

PV commit (group N) → QK commit (group N+1)  
wait<1>: waits for PV (oldest) → stalls more (9,754 samples)
wait<0>: waits for QK → almost always done (3,859 samples)

Action items:

1. ✅ ScaleOut::Zero implemented and validated (+2.5% non-causal)
2. 🔲 Examine FA3 SASS to determine exact WGMMA commit ordering
3. 🔲 Determine if FA3 uses a unified prologue (no separate n_idx==0 path)
4. 🔲 Test reversing our commit order (PV-first) to match FA3's stall pattern
5. 🔲 Handle causal masking with clear_accum=True (mask fill + accumulate mode)

File: _test_ws_fa3_v2_persistent_scaleout.py

[superAngGao]

Update: source-level FA3-order alignment experiment on GPU7 (locked freq), non-causal B=4 S=4096 H=64 Hkv=8 D=128 fp16, block_m=128 block_n=176

What changed

I created a new experimental kernel file in the TileOPs workspace:

• _test_ws_fa3_v2_threadbind_fa3_order.py

This keeps the original threadbind WS kernel untouched, and introduces a new builder:

• build_fa3_v2_fa3_order(...)

The goal was to align the TileLang/CUTLASS-level steady-state statement order with FA3's Hopper forward mainloop before touching lower-level codegen.

Baseline TileOPs steady-state order

In the original threadbind kernel, the steady-state order is effectively:

wait K[next]
scheduler sync
issue QK[next]
rescale O
wait V[cur]
issue PV[cur]
scheduler arrive
wait_wgmma(1)
release K
softmax(next)
wait_wgmma(0)
release V
copy P

FA3-aligned order used in the new experiment

The experimental kernel changes the steady-state order to:

wait K[next]
scheduler sync
issue QK[next]
wait V[cur]
issue PV[cur]
scheduler arrive
wait_wgmma(1)
release K
softmax(next)
wait_wgmma(0)
release V
copy P
rescale O

The important change is moving:

rescale_1(acc_o_1, ss_1)
rescale_2(acc_o_2, ss_2)

from before wait V / issue PV to after:

• wait_wgmma(0)
• barrier_arrive(v_empty)
• T.copy(acc_s_*, acc_s_cast_*)

This is meant to match the FA3 source-level common path where RescaleOBeforeGemm == false for the standard hdim=128 route.

Why this experiment matters

At the TileLang/CUTLASS semantic level, we could not find a decisive object-model difference between TileOPs and FA3. Both expose the same three logical paths:

• score accumulator
• packed P path
• output/PV accumulator

So the natural next question was: if we align the source-level consumer steady-state order exactly, how much performance do we recover before going deeper into lowering / SASS / compiler behavior?

Measurement method

To avoid cross-talk and to keep the comparison tight:

• GPU: H200 GPU7, locked frequency
• Shape: B=4 S=4096 H=64 Hkv=8 D=128 fp16
• Mode: non-causal
• Tile: block_m=128 block_n=176
• Measurement tool: Nsight Compute (ncu)
• Measurement style:
  • compile + warmup first
  • then use cudaProfilerStart() / cudaProfilerStop()
  • capture exactly one steady-state kernel launch
  • run kernels serially, not in parallel

TileOPs measurement command pattern

CUDA_VISIBLE_DEVICES=7 TMPDIR=/home/ga/.tmpncu \
  ncu --target-processes all --devices 0 \
      --kernel-name-base demangled \
      --launch-count 1 \
      --section LaunchStats --page raw --csv \
      --profile-from-start off --clock-control none \
      python ...

Inside Python:

• build kernel
• run once for compile/warmup
• cudaProfilerStart()
• run once
• cudaProfilerStop()

The same capture style was used for official FA3.

Results

Single-kernel steady-state time (gpu__time_duration.sum) from NCU LaunchStats:

Kernel	Time
TileOPs baseline threadbind WS (bn=176)	3.48256 ms
TileOPs fa3_order experiment (bn=176)	3.45856 ms
Official FA3 (bn=176, persistent)	3.087648 ms

Relative numbers

• fa3_order / baseline = 0.9931x latency
  • about 0.69% faster than baseline
• fa3_order / FA3 = 1.1201x latency
  • about 12.0% slower than FA3
• equivalently, TileOPs fa3_order reaches about 89.3% of FA3 on this setup

Interpretation

This experiment is useful because it narrows the problem substantially.

What it says

1. Source-level steady-state order does matter, because even after much of the structure already matched, aligning the order still gave a measurable improvement.
2. But the gain is small at block_n=176: only ~0.69%.
3. Therefore, once bn=176 and source-level order are aligned, the remaining gap to FA3 is not primarily explained by the high-level TileLang/CUTLASS statement order.

What it does not say

It does not mean the problem is solved. There is still ~12% latency gap to FA3 under this measurement setup.

The remaining gap is more likely in lower layers, e.g.:

• lowering of wait_wgmma(1) / wait_wgmma(0)
• placement of BAR.ARV / WARPGROUP.DEPBAR
• scheduler barrier realization
• persistent scheduler / launch shape
• final SASS instruction scheduling

Cross-check with earlier bn=128

Earlier A/B measurements on the same experimental kernel family showed a much larger improvement (~7.6%) when moving from baseline order to fa3_order at block_n=128.

That contrast is informative:

• when bn=128, source-level order still materially affects overlap quality
• when bn=176, after matching the FA3 tile shape, this layer is much closer to exhausted

So for bn=176, we should treat this step as mostly closing the semantic-order gap, not the final performance gap.

Takeaway

This update strengthens the case that:

• TileLang/CUTLASS semantic alignment is necessary but not sufficient
• after aligning statement order, the remaining work should focus on the lowered kernel / SASS-level behavior, not more high-level reordering alone

Next step: compare baseline bn176 vs fa3_order bn176 vs official FA3 at the SASS / NCU stall level, especially around:

• wait<1>
• BAR.ARV
• WARPGROUP.DEPBAR
• the softmax window between QK and PV completion