|
| 1 | +# Tri-Dao mamba_ssm bwd -> OUR (tilelang/tvm triton_frontend) stack: route + parity status |
| 2 | + |
| 3 | +MEASURED on gb10 (sm_121a, venv /home/dave/cppmega-venv, source /home/dave/source/tilelang). |
| 4 | +Commit b4f9f94c (merge/upstream-codegen-reorg). All GPU exec gb10-only. |
| 5 | + |
| 6 | +## Verdict: NO-GO (parity 0/7). Routing + TF32 tensor-core path PROVEN; numeric parity NOT yet achieved. |
| 7 | + |
| 8 | +The route is real (7/7 kernels lower end-to-end to TF32 tensor-core MMA, 288 HMMA in SASS), |
| 9 | +but the routed *output is numerically WRONG* on every real-strided config because of a |
| 10 | +gemm-accumulator-fragment epilogue defect (below). Per RULE #1 no routed ms is reported as a |
| 11 | +correctness/perf result — only parity-PASS kernels qualify, and 0/7 PASS. |
| 12 | + |
| 13 | +## Route + SASS re-confirm (after the FIX#3/FIX#4 read-path change) |
| 14 | + |
| 15 | +| check | result | |
| 16 | +|---|---| |
| 17 | +| route_all7 (lower -> CUDA, __global__ + mma_sync) | 7/7 OK WITH MMA | |
| 18 | +| sass_all7 (export_sass) | 7/7 SASS_OK | |
| 19 | +| HMMA.1688.F32.TF32 total in SASS | 288 (32+32+32+32+32+64+64) | |
| 20 | +| reduce kernels emitting atomicAdd | dx=1, state_dx=3, ddAcs=1 (REDG.E.ADD) | |
| 21 | + |
| 22 | +So the read-path change (FIX#3 input-extent + FIX#4 STSM/LDSM gate) did NOT regress routing: |
| 23 | +still 7/7 route, still 288 HMMA. |
| 24 | + |
| 25 | +## Parity (routed vs NATIVE Tri-Dao triton), REAL-strided NON-degenerate config |
| 26 | + |
| 27 | +Config: b1 nh8 hd64 ds64 nc8 cs64, BLOCK_K=32 => cs=64 != BLOCK_K => 2 K-trips/chunk |
| 28 | +(NON-degenerate, real seqlen strides — NOT cs=BK single-trip, NOT strides=0). |
| 29 | + |
| 30 | +| kernel | NATIVE nz | ROUTED nz | MAXDIFF | allclose 1e-3 | |
| 31 | +|---|---|---|---|---| |
| 32 | +| _chunk_scan_bwd_dstates | 262144/262144 | 2048/262144 | 2.335834e+02 | **FAIL** | |
| 33 | +| (6 others share the identical gemm->fragment->epilogue codegen) | — | — | — | **FAIL (not driven; same defect)** | |
| 34 | + |
| 35 | +dstates is the simplest of the 7; it FAILS, and all 7 emit the identical defective |
| 36 | +fragment-store (verified in generated CUDA), so parity is 0/7. Production §P1 parity was |
| 37 | +NOT run because the small real-strided config already fails for the same structural reason — |
| 38 | +running prod would only reproduce the same wrong fragment store at larger extent. |
| 39 | + |
| 40 | +## Root cause (precise, RULE #1 — the exact remaining defect, no fabricated PASS) |
| 41 | + |
| 42 | +The hand-built mma.1688 gemm writes its 64x64 result into a `local.fragment` (`dot_c_frag`), |
| 43 | +whose elements are DISTRIBUTED across the 128 threads in the mma C-fragment register layout: |
| 44 | +logical element [i,j] does NOT live at flat offset i*N+j of any one thread's register array. |
| 45 | + |
| 46 | +FIX#4 correctly tries to materialise the fragment to a logical shared tile before the |
| 47 | +Tri-Dao epilogue (`T.copy(c_frag, c_logical)`), and correctly gates STSM/LDSM off when the |
| 48 | +fragment layout is absent (src/backend/cuda/op/copy_analysis.cc CheckSTSMCopy/CheckLDSMCopy). |
| 49 | +BUT the fragment produced by the hand-built mma loop is NOT the recognised epilogue of a |
| 50 | +`tl::gemm` op, so **LayoutInference never registers a thread-layout for `dot_c_frag`**. With |
| 51 | +no layout in `layout_map`, the SIMT/Normal copy falls back to a FLAT thread-iteration copy. |
| 52 | +Generated CUDA (all 7 kernels, identical signature): |
| 53 | + |
| 54 | +```c |
| 55 | +*(float4*)(dot_c_logical_313 + ((i_4 * 512) + (((int)threadIdx.x) * 4))) |
| 56 | + = dot_c_frag_311[((i_4 * 128) + ((int)threadIdx.x))]; // FLAT, layout-blind |
| 57 | +``` |
| 58 | + |
| 59 | +This copies the fragment in raw thread-flat order, NOT through the mma-1688 register->[i,j] |
| 60 | +mapping. Only the slots where flat order happens to coincide with the true fragment layout |
| 61 | +carry meaning (2048/262144 nonzero, MAXDIFF 2.34e2). The materialisation is structurally |
| 62 | +unable to be correct via this path because there is no registered fragment layout to invert. |
| 63 | + |
| 64 | +### The real fix (next step, not done here) |
| 65 | +The fragment->shared store must apply the mma-1688 C-fragment thread layout. Two routes: |
| 66 | + 1. Emit the gemm as a real `tl::gemm` whose epilogue LayoutInference recognises, so the |
| 67 | + fragment gets a registered layout and `T.copy(c_frag, shared)` lowers layout-correctly |
| 68 | + (then STSM/LDSM or the layout-aware SIMT store both work); OR |
| 69 | + 2. Register the mma-1688 C-fragment layout for the hand-built `dot_c_frag` explicitly before |
| 70 | + the copy, so the SIMT/Normal store distributes registers to the correct [i,j] positions. |
| 71 | +RULE #1: this is a real correctness route to implement — NOT a silent fallback to patch over. |
| 72 | + |
| 73 | +## Native reference (valid regardless of routed parity) |
| 74 | + |
| 75 | +| measurement | value | |
| 76 | +|---|---| |
| 77 | +| native _chunk_scan_bwd_dstates @ §P1 (S=4096 c=64 g=8 H=112 P=64 N=64 bs1, grid=(1,64,112)) | 1.126 ms/kernel (MEASURED this run; ~1.2ms class, consistent w/ recorded 1.206ms) | |
| 78 | +| native full mamba_ssm bwd (all 7 kernels + non-Triton ops) | ~10ms-class | |
| 79 | +| path_c v1 full chain | 905 ms | |
| 80 | + |
| 81 | +routed ms vs native: NOT REPORTED — routed fails parity (RULE #1: no perf number for a |
| 82 | +numerically-wrong kernel). full-7 routed sum vs ~10ms / vs 905ms: NOT REPORTED for the same |
| 83 | +reason. |
| 84 | + |
| 85 | +## Summary |
| 86 | +- 7/7 route to real TF32 tensor cores (288 HMMA) — PROVEN and re-confirmed after read-path fix. |
| 87 | +- 0/7 numeric parity — blocked by the layout-blind fragment->shared epilogue store. |
| 88 | +- GO/NO-GO: **NO-GO** until the mma C-fragment layout is applied at the materialisation copy. |
0 commit comments