|
| 1 | +# Triton-frontend iteration 3 — coalesced async K-loop loads (LDGSTS) |
| 2 | + |
| 3 | +Status: **HONEST PARTIAL / NO-GO on LDGSTS this iteration.** Pipeline-annotation |
| 4 | +infrastructure landed and committed on the routed-triton path; the generated `.cu` |
| 5 | +is byte-identical OFF==ON, so cp.async/LDGSTS stays 0 and there is **no ms drop** |
| 6 | +from the 1102 ms iteration-2 baseline. The blocker is precisely localized in |
| 7 | +TVM's `pipeline_planning.cc` and verified by reading the source. |
| 8 | + |
| 9 | +All numbers below are MEASURED-FROM-COMMITTED-HEAD on gb10 (sm_121a), regenerating |
| 10 | +the `_chunk_scan_bwd_dstates` `.cu` from HEAD with `async_loads` OFF vs ON, |
| 11 | +interleaved, CUDA-event timed. |
| 12 | + |
| 13 | +- tilelang committed HEAD: `53be06b321b78844ab12873980cfbf5a0bbbf742` |
| 14 | + (gb10 working tree byte-identical to HEAD for the 4 iteration-3 files; md5-verified). |
| 15 | + |
| 16 | +## What landed (committed, iteration 3) |
| 17 | + |
| 18 | +- `WalkerCtx.routed_triton_async_loads` gate (ON by default on the routed |
| 19 | + prologue-opt path) + a shared `_gmem_shared_copies` counter list. |
| 20 | +- `_emit_load_copy` / `_emit_ptrstate_tile_load_tir` append to the counter when |
| 21 | + they emit a global->shared shared-scope copy. |
| 22 | +- `control.map_scf_for._maybe_pipeline()` stamps the serial K-loop `For` with the |
| 23 | + `num_stages` software-pipeline annotation (mirroring `T.Pipelined`) when the |
| 24 | + body emitted >=1 shared copy and the body is a `SeqStmt`. |
| 25 | + |
| 26 | +## Measured (gb10, `_chunk_scan_bwd_dstates`, §P1) |
| 27 | + |
| 28 | +§P1 config: b=1 nh=112 hd=64 ds=64 nc=64 cs=64, BM/BN/BK = 64/64/32. |
| 29 | + |
| 30 | +| metric | async OFF | async ON | Triton ref | |
| 31 | +|---|---|---|---| |
| 32 | +| generated `.cu` md5 | `99c232418b10c2793928ed63776b76c2` | `99c232418b10c2793928ed63776b76c2` | — | |
| 33 | +| cp.async in `.cu` | 0 | 0 | — | |
| 34 | +| SASS LDGSTS | 0 | 0 | 75 | |
| 35 | +| SASS spill (STL+LDL) | 306 (168+138) | 306 (168+138) | 0 | |
| 36 | +| SASS HMMA | 32 | 32 | 256 | |
| 37 | +| plain LDG | 38 | 38 | — | |
| 38 | +| §P1 parity MAXDIFF | 4.882812e-04 | 4.882812e-04 | — | |
| 39 | +| ALLCLOSE 1e-3 | PASS | PASS | — | |
| 40 | +| EXEC ms/kernel (CUDA events, N=50) | ~1102.3 | ~1102.3 (byte-identical) | — | |
| 41 | +| native ms/kernel | — | — | ~1.12 | |
| 42 | + |
| 43 | +The `.cu` is byte-identical OFF==ON, so timing is identical: rep medians 1102.14 / |
| 44 | +1102.31 / 1102.43 / 1102.64 ms over interleaved reps. **No real drop from 1102.** |
| 45 | +Remaining gap to native ~1.12 ms is ~984x and is dominated by LDGSTS=0 |
| 46 | +(memory-bound plain LDG) plus HMMA 32 vs 256 tensor-core under-util. |
| 47 | + |
| 48 | +Other paths: route_all7 = 5/7 build with mma intact |
| 49 | +(dstates/dc/db/state_dx/ddAcs, mma=2 each); dcb/dx fail IDENTICALLY to baseline |
| 50 | +(pre-existing PtrState/undefined-arg failures, not a regression). GEMM HMMA=32 |
| 51 | +intact in the routed kernels. |
| 52 | + |
| 53 | +## Why LDGSTS did not rise — precisely localized blocker (source-verified) |
| 54 | + |
| 55 | +The K-loop producer for the dout/C/dA tiles is a **predicated manual `tir.For` |
| 56 | +grid**: |
| 57 | + |
| 58 | +``` |
| 59 | +shared[i,j] = T.if_then_else(mask, global[flat_idx], other) |
| 60 | +``` |
| 61 | + |
| 62 | +TVM's async-pipeline planner only creates a cp.async stage for a *proven copy* |
| 63 | +stage. In `src/transform/pipeline_planning.cc`: |
| 64 | + |
| 65 | +- line ~973: `pure_copy_stage = collector.GetGlobalCopyPattern() && IsPureCopyStmt(block->body);` |
| 66 | +- `IsPureCopyStmt` (line ~652) requires the store VALUE to be a *pure raw copy |
| 67 | + value* — `is_pure_raw_copy_value` accepts only a `BufferLoad` of a global-like |
| 68 | + buffer, optionally wrapped in a `Cast`. A `tir.if_then_else` is a `CallNode`, |
| 69 | + so it returns `false`, which sets `saw_non_copy_buffer_store = true` and makes |
| 70 | + `IsPureCopyStmt` return `false`. |
| 71 | + |
| 72 | +Result: `copy_stage = false` -> no async stage created -> no cp.async emitted |
| 73 | +even though the `num_stages` annotation is present. The byte-identical `.cu` |
| 74 | +confirms the annotation alone is inert. |
| 75 | + |
| 76 | +Converting the predicated `For` to a clean `T.copy` `CopyNode` is further blocked |
| 77 | +because the structured 2D tile was lowered to a 1D `BufferLoad(src, [flat_idx])` |
| 78 | +with SYMBOLIC per-axis strides: the dout/C last-axis stride is |
| 79 | +`stride_dout_hdim` / `stride_c_dstate` (runtime args), NOT a literal 1, so the |
| 80 | +block-level axis-contiguity Triton's `Coalesce.cpp` keys on was discarded in the |
| 81 | +flat-index lowering. |
| 82 | + |
| 83 | +## RULE #1 compliance |
| 84 | + |
| 85 | +No silent uncoalesced fallback was introduced. The kernel still emits correct |
| 86 | +(plain LDG) code and parity holds at §P1 (MAXDIFF 4.88e-04). The honest reason |
| 87 | +LDGSTS stayed 0 is stated above and verified against the TVM source. The real fix |
| 88 | +(next iteration): preserve a 2D strided source region, split the predicate off the |
| 89 | +copy, emit a clean `T.copy` `CopyNode` global->shared so `IsPureCopyStmt` passes, |
| 90 | +then re-verify §P1 parity and re-measure ms/LDGSTS. |
| 91 | + |
| 92 | +## GO / NO-GO |
| 93 | + |
| 94 | +NO-GO on LDGSTS for iteration 3. Infrastructure (num_stages annotation, gate, |
| 95 | +counter) is landed and committed with full no-regression. The next iteration must |
| 96 | +attack the `IsPureCopyStmt` gate by producing a CopyNode, not a predicated |
| 97 | +`if_then_else` store. Remaining gap to native ~1.12 ms is ~984x, dominated by the |
| 98 | +uncoalesced loads and tensor-core under-utilization. |
0 commit comments