From cfca748c021c727a67e14ab80b5612b6bafd3d17 Mon Sep 17 00:00:00 2001 From: Dmitrii Vasilev Date: Mon, 19 Oct 2026 12:57:51 +0700 Subject: [PATCH 1/5] spec(igla): hslm_bwd_ptx -- kernel 3, dW split along k with red.global.add.s32 partials, grid z slices (Refs #10895) Co-Authored-By: Claude Sonnet 5.5 --- bootstrap/src/host/driver.rs | 6 +- specs/igla/coder/hslm_bwd_ptx.t27 | 296 +++++++++++++++++++++++++++++- 2 files changed, 291 insertions(+), 11 deletions(-) diff --git a/bootstrap/src/host/driver.rs b/bootstrap/src/host/driver.rs index 34dededc3c..cb81f2bcc5 100644 --- a/bootstrap/src/host/driver.rs +++ b/bootstrap/src/host/driver.rs @@ -299,14 +299,14 @@ mod tests { } } -// `t27c ptx-run` (#10404, #10847): load the PTX that ternary_gemm_ptx.t27 (kernel 0) or hslm_bwd_ptx.t27 (1, 2) wrote, run it, count mismatches with its reference. +// `t27c ptx-run` (#10404, #10847): load the PTX that ternary_gemm_ptx.t27 (kernel 0) or hslm_bwd_ptx.t27 (1, 2, 3) wrote, run it, count mismatches with its reference. #[cfg(feature = "cuda")] #[link(name = "cuda")] extern "C" { fn cuInit(flags: u32) -> i32; fn cuDeviceGet(dev: *mut i32, ordinal: i32) -> i32; fn cuCtxCreate_v2(ctx: *mut usize, flags: u32, dev: i32) -> i32; fn cuCtxSynchronize() -> i32; fn cuModuleLoadData(module: *mut usize, image: *const u8) -> i32; fn cuModuleGetFunction(func: *mut usize, module: usize, name: *const u8) -> i32; fn cuMemAlloc_v2(ptr: *mut u64, bytes: usize) -> i32; fn cuMemcpyHtoD_v2(dst: u64, src: *const u8, bytes: usize) -> i32; - fn cuMemcpyDtoH_v2(dst: *mut u8, src: u64, bytes: usize) -> i32; + fn cuMemcpyDtoH_v2(dst: *mut u8, src: u64, bytes: usize) -> i32; fn cuMemsetD32_v2(dst: u64, v: u32, n: usize) -> i32; fn cuLaunchKernel(f: usize, gx: u32, gy: u32, gz: u32, bx: u32, by: u32, bz: u32, smem: u32, stream: usize, params: *mut *mut u8, extra: usize) -> i32; } #[cfg(feature = "cuda")] @@ -328,7 +328,7 @@ pub fn ptx_run(kernel: u32, m: usize, n: usize, k: usize, iters: u32, stride: us ck(cuMemcpyHtoD_v2(da, a.as_ptr() as *const u8, a.len()))?; ck(cuMemcpyHtoD_v2(dw, wp.as_ptr(), wp.len()))?; let mut args: [*mut u8; 6] = [&mut da as *mut u64 as *mut u8, &mut dw as *mut u64 as *mut u8, &mut dc as *mut u64 as *mut u8, &mut pm as *mut u32 as *mut u8, &mut pn as *mut u32 as *mut u8, &mut pk as *mut u32 as *mut u8]; - let mut launch = || ck(cuLaunchKernel(f, t::grid_dim(n), t::grid_dim(m), 1, t::THREADS, 1, 1, 0, 0, args.as_mut_ptr(), 0)); + let mut launch = || { if u::zero_first(kernel) { ck(cuMemsetD32_v2(dc, 0, m * n))?; } ck(cuLaunchKernel(f, t::grid_dim(n), t::grid_dim(m), u::grid_z(kernel, m, n, k), t::THREADS, 1, 1, 0, 0, args.as_mut_ptr(), 0)) }; launch()?; ck(cuCtxSynchronize())?; ck(cuMemcpyDtoH_v2(c.as_mut_ptr() as *mut u8, dc, c.len() * 4))?; let t0 = std::time::Instant::now(); for _ in 0..iters { launch()?; } ck(cuCtxSynchronize())?; diff --git a/specs/igla/coder/hslm_bwd_ptx.t27 b/specs/igla/coder/hslm_bwd_ptx.t27 index d922270533..82dc4bffed 100644 --- a/specs/igla/coder/hslm_bwd_ptx.t27 +++ b/specs/igla/coder/hslm_bwd_ptx.t27 @@ -25,6 +25,12 @@ ; i8), B(r, j) = dY[r][j] (rows x out_dim s8). Both reductions run down a column, so each thread ; loads 16 bytes of one row (row t/8 of the k-step, elements (t%8)*16..+15) and stores them ; transposed, one byte per sA/sB row. +; hslm_dw_splitk_s8 (kernel 3, issue #10895). Kernel 2 split along k. At k = 1782 a 256 x 736 output is +; 12 blocks of 128 x 128 on a 128-SM part and each walks 56 k-steps alone (0.075 ms on an RTX 4090, +; 11x CUTLASS s8). Kernel 3 launches grid z = grid_z() slices; slice z sums k in [z * per * 32, +; (z + 1) * per * 32), per = ceil(ceil(k / 32) / nz) read from %nctaid.z, and ends with +; red.global.add.s32 of its i32 partials into C, which the host MUST zero before each launch +; (zero_first()). Integer addition is associative, so C equals kernel 2's bit for bit. ; Scales (the host's side, stated here so the integer kernel has one meaning): ; dX: dYq[r][i] = round(dy[r][i] / s_r), s_r = max_i |dy[r][i]| / 127 (per row); dx[r][j] = s_r * C. ; The trainer's sum has no alpha in it; a per-matrix alpha is one scalar and commutes with the sum. @@ -47,6 +53,10 @@ use igla::coder::ptx_emit; pub const KERNEL_FWD : u32 = 0; pub const KERNEL_DX : u32 = 1; pub const KERNEL_DW : u32 = 2; +/// Kernel 3 is kernel 2 split along k (issue #10895): grid z = slices, i32 partials added into a zeroed C. +pub const KERNEL_DW_SPLITK : u32 = 3; +/// Blocks the split aims at: two 256-thread blocks per SM of a 128-SM part. +pub const SPLIT_TARGET_BLOCKS : usize = 256; /// Block tile, k-step and thread count, as in the forward kernel. pub const BM : u32 = 128; pub const BN : u32 = 128; @@ -75,6 +85,7 @@ fn put_name_at(out : []u8, cap : usize, w : usize, kernel : u32) -> usize { if (kernel == KERNEL_FWD) { return put_s(out, cap, w, "ternary_gemm_s8"); } if (kernel == KERNEL_DX) { return put_s(out, cap, w, "hslm_dx_s8"); } if (kernel == KERNEL_DW) { return put_s(out, cap, w, "hslm_dw_s8"); } + if (kernel == KERNEL_DW_SPLITK) { return put_s(out, cap, w, "hslm_dw_splitk_s8"); } return w; } @@ -315,10 +326,80 @@ fn emit_loads_dw(out : []u8, cap : usize, w : usize) -> usize { return emit_transposed(out, cap, x, 50, "%r16", "%r37"); } +/// Slice bounds of a split-K block (kernel 3). z = %ctaid.z of nz = %nctaid.z slices: every slice runs +/// per = ceil(ceil(k / 32) / nz) k-steps, from r49 = z * per * 32 up to r3 = min(k, (z + 1) * per * 32), +/// which overwrites the k bound the loop and the loads compare against. A block with nothing to sum +/// returns before the first barrier. +fn emit_slice(out : []u8, cap : usize, w : usize) -> usize { + var x : usize = ins(out, cap, w, "mov.u32 %r46, %ctaid.z;"); + x = ins(out, cap, x, "mov.u32 %r47, %nctaid.z;"); + x = ins(out, cap, x, "add.u32 %r48, %r3, 31;"); + x = ins(out, cap, x, "shr.u32 %r48, %r48, 5;"); + x = ins(out, cap, x, "add.u32 %r48, %r48, %r47;"); + x = ins(out, cap, x, "sub.u32 %r48, %r48, 1;"); + x = ins(out, cap, x, "div.u32 %r48, %r48, %r47;"); + x = ins(out, cap, x, "shl.b32 %r48, %r48, 5;"); + x = ins(out, cap, x, "mul.lo.u32 %r49, %r46, %r48;"); + x = ins(out, cap, x, "add.u32 %r46, %r46, 1;"); + x = ins(out, cap, x, "mul.lo.u32 %r46, %r46, %r48;"); + x = ins(out, cap, x, "min.u32 %r3, %r3, %r46;"); + x = ins(out, cap, x, "setp.ge.u32 %p9, %r49, %r3;"); + return ins(out, cap, x, "@%p9 ret;"); +} + +/// One pair of accumulators b, b + 1 of row dr, column dc of the warp tile as two guarded i32 adds +/// into C (the split-K epilogue; ptx_emit's emit_store writes the pair instead). +fn emit_red_store(out : []u8, cap : usize, w : usize, dr : u32, dc : u32, b : u32) -> usize { + var x : usize = put_s(out, cap, w, " add.u32 %r41, %r39, "); + x = put_u(out, cap, x, dr); + x = put_s(out, cap, x, ";\n"); + x = put_s(out, cap, x, " add.u32 %r42, %r40, "); + x = put_u(out, cap, x, dc); + x = put_s(out, cap, x, ";\n"); + x = ins(out, cap, x, "setp.lt.u32 %p6, %r41, %r1;"); + x = ins(out, cap, x, "setp.lt.u32 %p7, %r42, %r2;"); + x = ins(out, cap, x, "and.pred %p6, %p6, %p7;"); + x = ins(out, cap, x, "cvt.u64.u32 %rd10, %r42;"); + x = ins(out, cap, x, "mad.wide.u32 %rd10, %r41, %r2, %rd10;"); + x = ins(out, cap, x, "shl.b64 %rd10, %rd10, 2;"); + x = ins(out, cap, x, "add.u64 %rd10, %rd10, %rd3;"); + x = put_s(out, cap, x, " @%p6 red.global.add.s32 [%rd10], %acc"); + x = put_u(out, cap, x, b); + x = put_s(out, cap, x, ";\n @%p6 red.global.add.s32 [%rd10+4], %acc"); + x = put_u(out, cap, x, b + 1); + return put_s(out, cap, x, ";\n"); +} + +/// The split-K epilogue: ptx_emit's row and column base setup, then 32 guarded pairs of adds. +fn emit_red_epilogue(out : []u8, cap : usize, w : usize) -> usize { + var x : usize = ins(out, cap, w, "shl.b32 %r39, %r6, 7;"); + x = ins(out, cap, x, "shl.b32 %r41, %r23, 6;"); + x = ins(out, cap, x, "add.u32 %r39, %r39, %r41;"); + x = ins(out, cap, x, "add.u32 %r39, %r39, %r21;"); + x = ins(out, cap, x, "shl.b32 %r40, %r5, 7;"); + x = ins(out, cap, x, "shl.b32 %r41, %r24, 5;"); + x = ins(out, cap, x, "add.u32 %r40, %r40, %r41;"); + x = ins(out, cap, x, "shl.b32 %r41, %r22, 1;"); + x = ins(out, cap, x, "add.u32 %r40, %r40, %r41;"); + var mi : u32 = 0; + while (mi < 4) { + var ni : u32 = 0; + while (ni < 4) { + const b : u32 = (mi * 4 + ni) * 4; + x = emit_red_store(out, cap, x, mi * 16, ni * 8, b); + x = emit_red_store(out, cap, x, mi * 16 + 8, ni * 8, b + 2); + ni += 1; + } + mi += 1; + } + return x; +} + /// The whole text of kernel 1 or 2 into out[0..cap). Returns the bytes it needs (more than cap means /// a clipped text); an unknown id writes nothing and returns 0. pub fn emit_ptx(kernel : u32, out : []u8, cap : usize) -> usize { - if (kernel != KERNEL_DX and kernel != KERNEL_DW) { return 0; } + if (kernel != KERNEL_DX and kernel != KERNEL_DW and kernel != KERNEL_DW_SPLITK) { return 0; } + const split : bool = kernel == KERNEL_DW_SPLITK; var x : usize = emit_header(out, cap, 0, kernel); x = emit_params(out, cap, x); if (kernel == KERNEL_DX) { @@ -326,7 +407,9 @@ pub fn emit_ptx(kernel : u32, out : []u8, cap : usize) -> usize { } else { x = emit_setup_dw(out, cap, x); } + if (split) { x = emit_slice(out, cap, x); } x = emit_warp_setup(out, cap, x); + if (split) { x = ins(out, cap, x, "mov.u32 %r28, %r49;"); } x = put_s(out, cap, x, "LOOP_K:\n"); x = ins(out, cap, x, "setp.ge.u32 %p3, %r28, %r3;"); x = ins(out, cap, x, "@%p3 bra DONE_K;"); @@ -341,7 +424,11 @@ pub fn emit_ptx(kernel : u32, out : []u8, cap : usize) -> usize { x = ins(out, cap, x, "add.u32 %r28, %r28, 32;"); x = ins(out, cap, x, "bra LOOP_K;"); x = put_s(out, cap, x, "DONE_K:\n"); - x = emit_epilogue(out, cap, x); + if (split) { + x = emit_red_epilogue(out, cap, x); + } else { + x = emit_epilogue(out, cap, x); + } x = ins(out, cap, x, "ret;"); return put_s(out, cap, x, "}\n"); } @@ -355,19 +442,52 @@ pub fn grid_dim(x : usize) -> u32 { return ((x + 127) / 128) as u32; } -/// Whether kernel 1 or 2 can run an m x n output over k. +/// Slices (grid z) of a kernel over an m x n output and k: 1 for every kernel but 3. Kernel 3 aims at +/// SPLIT_TARGET_BLOCKS blocks, never more slices than k-steps, and then drops the slices the rounding of +/// the per-slice step count would leave empty, so no block is idle. +pub fn grid_z(kernel : u32, m : usize, n : usize, k : usize) -> u32 { + if (kernel != KERNEL_DW_SPLITK) { return 1; } + const steps : usize = (k + 31) / 32; + const blocks : usize = ((m + 127) / 128) * ((n + 127) / 128); + var s : usize = (SPLIT_TARGET_BLOCKS + blocks - 1) / blocks; + if (s > steps) { s = steps; } + if (s < 1) { s = 1; } + const per : usize = (steps + s - 1) / s; + return ((steps + per - 1) / per) as u32; +} + +/// Whether C must be zero before a launch: kernel 3 adds its partials into it. +pub fn zero_first(kernel : u32) -> bool { + return kernel == KERNEL_DW_SPLITK; +} + +/// The k range [lo, hi) slice z of nz sums, as the kernel computes it: per = ceil(ceil(k / 32) / nz) +/// steps of 32. Empty (lo >= hi) when z * per * 32 >= k. +pub fn slice_lo(z : usize, nz : usize, k : usize) -> usize { + const per : usize = ((k + 31) / 32 + nz - 1) / nz; + return z * per * 32; +} + +pub fn slice_hi(z : usize, nz : usize, k : usize) -> usize { + const per : usize = ((k + 31) / 32 + nz - 1) / nz; + const hi : usize = (z + 1) * per * 32; + if (hi < k) { return hi; } + return k; +} + +/// Whether kernel 1, 2 or 3 can run an m x n output over k. pub fn shape_ok(kernel : u32, m : usize, n : usize, k : usize) -> bool { if (m < 1 or n < 16 or n % 16 != 0) { return false; } if ((m + 127) / 128 > GRID_Y_MAX) { return false; } if (kernel == KERNEL_DX) { return k >= 16 and k % 16 == 0 and k <= K_MAX_DX; } - if (kernel == KERNEL_DW) { return m % 16 == 0 and k >= 1 and k <= K_MAX_DW; } + if (kernel == KERNEL_DW or kernel == KERNEL_DW_SPLITK) { return m % 16 == 0 and k >= 1 and k <= K_MAX_DW; } return false; } -/// Bytes of the B operand: k x n 2-bit codes for kernels 0 and 1, k x n s8 for kernel 2. A is m * k bytes +/// Bytes of the B operand: k x n 2-bit codes for kernels 0 and 1, k x n s8 for kernels 2 and 3. A is m * k bytes /// for every kernel (dX: dYq rows x k; dW: X as k rows x m). pub fn b_len(kernel : u32, n : usize, k : usize) -> usize { - if (kernel == KERNEL_DW) { return k * n; } + if (kernel == KERNEL_DW or kernel == KERNEL_DW_SPLITK) { return k * n; } return k * n / 4; } @@ -522,9 +642,9 @@ test both_kernels_open_with_their_header_and_entry { test an_unknown_kernel_writes_nothing { var out : [64]u8 = [7; 64]; assert emit_ptx(KERNEL_FWD, &out, 64) == 0; - assert emit_ptx(3, &out, 64) == 0; + assert emit_ptx(4, &out, 64) == 0; assert out[0] == 7; - assert put_name(3, &out, 64) == 0; + assert put_name(4, &out, 64) == 0; } test the_names_are_the_entry_names { @@ -921,6 +1041,166 @@ test the_limits_of_each_kernel { assert grid_dim(HSLM_FF_PAD) == 6; } +test the_split_k_kernel_has_its_own_header_and_name { + var out : [32768]u8 = undefined; + const m : usize = emit_ptx(KERNEL_DW_SPLITK, &out, PTX_CAP); + assert m > 0; + assert at_text(&out, m, 0, "// hslm_dw_splitk_s8 -- emitted by t27c from specs/igla/coder/hslm_bwd_ptx.t27 (#10847). Do not edit.\n.version 7.0\n.target sm_80\n.address_size 64\n\n"); + assert count_text(&out, m, ".visible .entry hslm_dw_splitk_s8(\n") == 1; + var nm : [32]u8 = undefined; + const c : usize = put_name(KERNEL_DW_SPLITK, &nm, 32); + assert c == 17; + assert at_text(&nm, c, 0, "hslm_dw_splitk_s8"); + assert at_text(&out, m, m - 11, " ret;\n}\n"); +} + +test the_split_k_kernel_adds_its_partials_into_c_and_slices_k_by_the_grid { + var out : [32768]u8 = undefined; + const m : usize = emit_ptx(KERNEL_DW_SPLITK, &out, PTX_CAP); + // The same 16 int8 mma per k-step and the same transposing dW loads as kernel 2. + assert count_text(&out, m, "mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 ") == 16; + assert count_text(&out, m, "ld.global.v4.u32 ") == 2; + assert count_text(&out, m, "st.shared.u8 ") == 32; + assert count_text(&out, m, "bar.sync 0;") == 2; + // The epilogue is 64 guarded i32 adds and no plain store: C must be zero before the launch. + assert count_text(&out, m, "red.global.add.s32 ") == 64; + assert count_text(&out, m, " @%p6 red.global.add.s32 [%rd10], %acc0;\n @%p6 red.global.add.s32 [%rd10+4], %acc1;\n") == 1; + assert count_text(&out, m, " @%p6 red.global.add.s32 [%rd10], %acc62;\n @%p6 red.global.add.s32 [%rd10+4], %acc63;\n") == 1; + assert count_text(&out, m, "st.global") == 0; + assert count_text(&out, m, "atom.") == 0; + // The slice code: z and nz from the grid, per = ceil(ceil(k / 32) / nz) steps of 32, an empty slice returns. + assert count_text(&out, m, " mov.u32 %r46, %ctaid.z;\n mov.u32 %r47, %nctaid.z;\n") == 1; + assert count_text(&out, m, " div.u32 %r48, %r48, %r47;\n shl.b32 %r48, %r48, 5;\n") == 1; + assert count_text(&out, m, " min.u32 %r3, %r3, %r46;\n") == 1; + assert count_text(&out, m, " setp.ge.u32 %p9, %r49, %r3;\n @%p9 ret;\n") == 1; + // The k loop starts at the slice's own lower bound, once, after the warp setup. + assert count_text(&out, m, " mov.u32 %r28, %r49;\nLOOP_K:\n") == 1; + assert count_text(&out, m, "ctaid.z") == 2; + // Kernel 2 has none of it. + const d : usize = emit_ptx(KERNEL_DW, &out, PTX_CAP); + assert count_text(&out, d, "ctaid.z") == 0; + assert count_text(&out, d, "red.global") == 0; + assert count_text(&out, d, "%p9") == 0; +} + +test the_split_k_text_is_ascii_fits_and_is_stable { + var a : [32768]u8 = undefined; + var b : [32768]u8 = undefined; + const m : usize = emit_ptx(KERNEL_DW_SPLITK, &a, PTX_CAP); + // The length ptx-run wrote on the RTX 4090 pod (wc -c ptx-run-3.ptx). + assert m == 0; + assert emit_ptx(KERNEL_DW_SPLITK, &b, PTX_CAP) == m; + var i : usize = 0; + while (i < m) { + assert a[i] == b[i]; + assert a[i] == '\n' or (a[i] >= 32 and a[i] < 127); + i += 1; + } +} + +test every_k_belongs_to_exactly_one_slice { + var k : usize = 1; + while (k <= 1800) { + var nz : usize = 1; + while (nz <= 70) { + // Slices tile [0, k): each starts where the one before ended, the last ends at k. + var z : usize = 0; + var at : usize = 0; + while (z < nz) { + const lo : usize = slice_lo(z, nz, k); + const hi : usize = slice_hi(z, nz, k); + if (lo < k) { + assert lo == at; + assert hi > lo; + at = hi; + } + z += 1; + } + assert at == k; + nz += 3; + } + k += 37; + } + // k = 1782 over 19 slices: 56 steps, 3 per slice, the last slice has 2. + assert slice_lo(0, 19, 1782) == 0; + assert slice_hi(0, 19, 1782) == 96; + assert slice_lo(18, 19, 1782) == 1728; + assert slice_hi(18, 19, 1782) == 1782; +} + +test the_slice_count_leaves_no_block_idle { + assert grid_z(KERNEL_FWD, 256, 736, 1782) == 1; + assert grid_z(KERNEL_DX, 256, 736, 1792) == 1; + assert grid_z(KERNEL_DW, 256, 736, 1782) == 1; + assert grid_z(KERNEL_DW_SPLITK, 256, 736, 1782) == 19; + assert grid_z(KERNEL_DW_SPLITK, 256, 256, 1782) == 56; + assert grid_z(KERNEL_DW_SPLITK, 736, 256, 1782) == 19; + // One step of k cannot be split; a big output needs no split. + assert grid_z(KERNEL_DW_SPLITK, 256, 256, 27) == 1; + assert grid_z(KERNEL_DW_SPLITK, 4096, 4096, 1782) == 1; + var k : usize = 1; + while (k <= 3000) { + var mm : usize = 16; + while (mm <= 1024) { + const nz : usize = grid_z(KERNEL_DW_SPLITK, mm, 256, k) as usize; + assert nz >= 1; + // The last slice is not empty, so none is. + assert slice_lo(nz - 1, nz, k) < k; + mm += 240; + } + k += 61; + } + assert zero_first(KERNEL_DW_SPLITK) == true; + assert zero_first(KERNEL_DW) == false; + assert zero_first(KERNEL_DX) == false; +} + +test the_slices_summed_are_the_unsplit_reference { + // m = 16, n = 16, k = 100: C[i][j] = sum over rows r of X[r][i] * dY[r][j], by slices of 4 and by one. + var a : [1600]i8 = [0; 1600]; + var b : [1600]u8 = [0; 1600]; + var s : u64 = 5; + var t : usize = 0; + while (t < 1600) { + s = lcg(s); + a[t] = ((((s >> 8) % 256) as i32) - 128) as i8; + s = lcg(s); + b[t] = ((s >> 8) & 255) as u8; + t += 1; + } + var whole : [16]i32 = [0; 16]; + ref_row(KERNEL_DW_SPLITK, &a, &b, 16, 16, 100, 3, &whole); + var sum : [16]i32 = [0; 16]; + var z : usize = 0; + while (z < 4) { + const lo : usize = slice_lo(z, 4, 100); + const hi : usize = slice_hi(z, 4, 100); + var j : usize = 0; + while (j < 16) { + var r : usize = lo; + while (r < hi) { + sum[j] += (a[r * 16 + 3] as i32) * sbyte(b[r * 16 + j]); + r += 1; + } + j += 1; + } + z += 1; + } + var j2 : usize = 0; + while (j2 < 16) { + assert sum[j2] == whole[j2]; + j2 += 1; + } +} + +test the_split_k_kernel_takes_the_shapes_of_kernel_2 { + assert shape_ok(KERNEL_DW_SPLITK, HSLM_FF_PAD, HSLM_EMBED_PAD, HSLM_ROWS) == true; + assert shape_ok(KERNEL_DW_SPLITK, 16, 16, 131071) == true; + assert shape_ok(KERNEL_DW_SPLITK, 16, 16, 131072) == false; + assert shape_ok(KERNEL_DW_SPLITK, 8, 16, 4) == false; + assert b_len(KERNEL_DW_SPLITK, 736, 1782) == 1311552; +} + invariant dx_accumulator_cannot_overflow { 128 * K_MAX_DX <= 2147483647 } invariant dw_accumulator_cannot_overflow { 16384 * K_MAX_DW <= 2147483647 and 16384 * (K_MAX_DW + 1) > 2147483647 } From fddbc331edc5ca89795e09de6958a39ace9a142e Mon Sep 17 00:00:00 2001 From: Dmitrii Vasilev Date: Mon, 19 Oct 2026 13:02:12 +0700 Subject: [PATCH 2/5] spec(igla): hslm_bwd_ptx -- pin the kernel 3 text length (Refs #10895) Co-Authored-By: Claude Sonnet 5.5 --- specs/igla/coder/hslm_bwd_ptx.t27 | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/specs/igla/coder/hslm_bwd_ptx.t27 b/specs/igla/coder/hslm_bwd_ptx.t27 index 82dc4bffed..df45492124 100644 --- a/specs/igla/coder/hslm_bwd_ptx.t27 +++ b/specs/igla/coder/hslm_bwd_ptx.t27 @@ -1088,7 +1088,7 @@ test the_split_k_text_is_ascii_fits_and_is_stable { var b : [32768]u8 = undefined; const m : usize = emit_ptx(KERNEL_DW_SPLITK, &a, PTX_CAP); // The length ptx-run wrote on the RTX 4090 pod (wc -c ptx-run-3.ptx). - assert m == 0; + assert m == 21819; assert emit_ptx(KERNEL_DW_SPLITK, &b, PTX_CAP) == m; var i : usize = 0; while (i < m) { From 1ad9ec5ae61d6cbea9cbe35d941ef88e672e44b7 Mon Sep 17 00:00:00 2001 From: Dmitrii Vasilev Date: Mon, 19 Oct 2026 13:15:20 +0700 Subject: [PATCH 3/5] spec(igla): hslm_bwd_ptx -- kernel 3 maps a warp onto one chunk and 32 k rows, so the transposed byte stores stop hitting one bank (Refs #10895) Co-Authored-By: Claude Sonnet 5.5 --- specs/igla/coder/hslm_bwd_ptx.t27 | 18 +++++++++++++----- 1 file changed, 13 insertions(+), 5 deletions(-) diff --git a/specs/igla/coder/hslm_bwd_ptx.t27 b/specs/igla/coder/hslm_bwd_ptx.t27 index df45492124..e034e0d497 100644 --- a/specs/igla/coder/hslm_bwd_ptx.t27 +++ b/specs/igla/coder/hslm_bwd_ptx.t27 @@ -169,9 +169,17 @@ fn emit_setup_dx(out : []u8, cap : usize, w : usize) -> usize { } /// dW: k-row r12 = tid/8 of both tiles, element chunk r13 = (tid%8)*16; i = r9 along m, j = r14 along n. -fn emit_setup_dw(out : []u8, cap : usize, w : usize) -> usize { - var x : usize = ins(out, cap, w, "shr.u32 %r12, %r4, 3;"); - x = ins(out, cap, x, "and.b32 %r13, %r4, 7;"); +/// Kernel 3 maps a warp onto one chunk and 32 k rows (r12 = tid%32, r13 = (tid/32)*16): the 32 lanes of a +/// byte store then land in 8 banks, 4 lanes to a word, where the 8 lanes per row of kernel 2 hit one bank. +fn emit_setup_dw(out : []u8, cap : usize, w : usize, split : bool) -> usize { + var x : usize = w; + if (split) { + x = ins(out, cap, x, "and.b32 %r12, %r4, 31;"); + x = ins(out, cap, x, "shr.u32 %r13, %r4, 5;"); + } else { + x = ins(out, cap, x, "shr.u32 %r12, %r4, 3;"); + x = ins(out, cap, x, "and.b32 %r13, %r4, 7;"); + } x = ins(out, cap, x, "shl.b32 %r13, %r13, 4;"); x = ins(out, cap, x, "shl.b32 %r9, %r6, 7;"); x = ins(out, cap, x, "add.u32 %r9, %r9, %r13;"); @@ -405,7 +413,7 @@ pub fn emit_ptx(kernel : u32, out : []u8, cap : usize) -> usize { if (kernel == KERNEL_DX) { x = emit_setup_dx(out, cap, x); } else { - x = emit_setup_dw(out, cap, x); + x = emit_setup_dw(out, cap, x, split); } if (split) { x = emit_slice(out, cap, x); } x = emit_warp_setup(out, cap, x); @@ -1088,7 +1096,7 @@ test the_split_k_text_is_ascii_fits_and_is_stable { var b : [32768]u8 = undefined; const m : usize = emit_ptx(KERNEL_DW_SPLITK, &a, PTX_CAP); // The length ptx-run wrote on the RTX 4090 pod (wc -c ptx-run-3.ptx). - assert m == 21819; + assert m == 21820; assert emit_ptx(KERNEL_DW_SPLITK, &b, PTX_CAP) == m; var i : usize = 0; while (i < m) { From ce14e2ba363dd3d8cca62e45494d372ea6c62f05 Mon Sep 17 00:00:00 2001 From: Dmitrii Vasilev Date: Mon, 19 Oct 2026 13:26:00 +0700 Subject: [PATCH 4/5] spec(igla): hslm_bwd_ptx -- kernel 3 aims at 64 blocks, the floor of the RTX 4090 sweep (Refs #10895) Co-Authored-By: Claude Sonnet 5.5 --- specs/igla/coder/hslm_bwd_ptx.t27 | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/specs/igla/coder/hslm_bwd_ptx.t27 b/specs/igla/coder/hslm_bwd_ptx.t27 index e034e0d497..238daf9775 100644 --- a/specs/igla/coder/hslm_bwd_ptx.t27 +++ b/specs/igla/coder/hslm_bwd_ptx.t27 @@ -31,6 +31,8 @@ ; (z + 1) * per * 32), per = ceil(ceil(k / 32) / nz) read from %nctaid.z, and ends with ; red.global.add.s32 of its i32 partials into C, which the host MUST zero before each launch ; (zero_first()). Integer addition is associative, so C equals kernel 2's bit for bit. +; It also maps a warp onto one 16-byte chunk and 32 k rows (kernel 2: 8 lanes per row), which makes the +; transposed byte stores of a warp land in 8 banks instead of 1. ; Scales (the host's side, stated here so the integer kernel has one meaning): ; dX: dYq[r][i] = round(dy[r][i] / s_r), s_r = max_i |dy[r][i]| / 127 (per row); dx[r][j] = s_r * C. ; The trainer's sum has no alpha in it; a per-matrix alpha is one scalar and commutes with the sum. @@ -55,8 +57,9 @@ pub const KERNEL_DX : u32 = 1; pub const KERNEL_DW : u32 = 2; /// Kernel 3 is kernel 2 split along k (issue #10895): grid z = slices, i32 partials added into a zeroed C. pub const KERNEL_DW_SPLITK : u32 = 3; -/// Blocks the split aims at: two 256-thread blocks per SM of a 128-SM part. -pub const SPLIT_TARGET_BLOCKS : usize = 256; +/// Blocks the split aims at. Every slice ends in one red.global.add.s32 per C element, so more slices buy +/// parallelism and pay in red traffic; on an RTX 4090 the sweep of 8..256 was U-shaped with its floor at 48..64. +pub const SPLIT_TARGET_BLOCKS : usize = 64; /// Block tile, k-step and thread count, as in the forward kernel. pub const BM : u32 = 128; pub const BN : u32 = 128; @@ -1140,9 +1143,9 @@ test the_slice_count_leaves_no_block_idle { assert grid_z(KERNEL_FWD, 256, 736, 1782) == 1; assert grid_z(KERNEL_DX, 256, 736, 1792) == 1; assert grid_z(KERNEL_DW, 256, 736, 1782) == 1; - assert grid_z(KERNEL_DW_SPLITK, 256, 736, 1782) == 19; - assert grid_z(KERNEL_DW_SPLITK, 256, 256, 1782) == 56; - assert grid_z(KERNEL_DW_SPLITK, 736, 256, 1782) == 19; + assert grid_z(KERNEL_DW_SPLITK, 256, 736, 1782) == 6; + assert grid_z(KERNEL_DW_SPLITK, 256, 256, 1782) == 14; + assert grid_z(KERNEL_DW_SPLITK, 736, 256, 1782) == 6; // One step of k cannot be split; a big output needs no split. assert grid_z(KERNEL_DW_SPLITK, 256, 256, 27) == 1; assert grid_z(KERNEL_DW_SPLITK, 4096, 4096, 1782) == 1; From b0010203f78cf4e350d7755a7efc7f7ac26db7ea Mon Sep 17 00:00:00 2001 From: Dmitrii Vasilev Date: Mon, 19 Oct 2026 15:36:29 +0700 Subject: [PATCH 5/5] gen+seal: hslm_bwd_ptx with kernel 3 on top of ptx_emit, regenerated on the lab (Refs #10895) test-report 26/26, 0 vacuous; gen-rust and seal --save on the Railway lab. Co-Authored-By: Claude Sonnet 5.5 --- .../seals/coder_igla-coder-hslm-bwd-ptx.json | 18 +-- bootstrap/gen/rust/igla/coder/hslm_bwd_ptx.rs | 139 ++++++++++++++++-- 2 files changed, 139 insertions(+), 18 deletions(-) diff --git a/.trinity/seals/coder_igla-coder-hslm-bwd-ptx.json b/.trinity/seals/coder_igla-coder-hslm-bwd-ptx.json index 38325ab3ac..591d0b787d 100644 --- a/.trinity/seals/coder_igla-coder-hslm-bwd-ptx.json +++ b/.trinity/seals/coder_igla-coder-hslm-bwd-ptx.json @@ -1,5 +1,5 @@ { - "built_by": "t27c-bootstrap@0.5.2+4e6abf0ec", + "built_by": "t27c-bootstrap@0.5.2+ce14e2b", "closure": [ { "spec": "specs/igla/coder/ptx_emit.t27", @@ -7,23 +7,23 @@ } ], "config": "gen=c,rust,verilog,zig;tests=zig", - "gen_hash_c": "sha256:5680fdfb166ad3d0f26bb2e3caae8e1e1896df8271ee20fe162c8c7c5ef52aa7", - "gen_hash_rust": "sha256:c6e4550b86cf46546df827e1e1a6427242b9bdd1755e64c73e0600f04d3923b4", - "gen_hash_verilog": "sha256:e5234465d9e51e5a2e244a964ef6aabd16ab80f3c09d193da971a30cc1469fc1", - "gen_hash_zig": "sha256:0217dc3f12608f12adaa5b79d9d655e7d750ff2d6f1652de4b39f95bac52f30e", + "gen_hash_c": "sha256:22aa0e3bed6a5dc83976ed5094e7f8e434e85fb0ae7ea2ee3ab28ec96cb9ebba", + "gen_hash_rust": "sha256:bd4040475901d80ad7316c3c29303e10de70924ec6046fb9dc2d2958681fa109", + "gen_hash_verilog": "sha256:634e8992f0c93c2c59316a11d1b66beccb433b150692b5289cccbacc01c3e917", + "gen_hash_zig": "sha256:cb73ef03b2a6b6fce7b2c9d4889dc1b53921a326627cd0251e54b14a3b556664", "module": "igla-coder-hslm-bwd-ptx", "ring": 12, "seal_schema": 2, - "sealed_at": "2026-10-19T08:01:23Z", + "sealed_at": "2026-10-19T08:35:21Z", "sealed_by": "t27c-bootstrap@0.5.2", - "spec_hash": "sha256:adc270a216694ead6050bcd7d1df1f64f0f99e5fc3447556399d50b0ce3e7f19", + "spec_hash": "sha256:c6e574a517e847a24751ecc3b85956b6fc6d0ba09516bafbb3fa61fe450eac44", "spec_path": "specs/igla/coder/hslm_bwd_ptx.t27", "tests": { "failed": 0, "failing": [], "forced": false, - "passed": 19, - "total": 19, + "passed": 26, + "total": 26, "vacuous": 0 }, "toolchain": { diff --git a/bootstrap/gen/rust/igla/coder/hslm_bwd_ptx.rs b/bootstrap/gen/rust/igla/coder/hslm_bwd_ptx.rs index 01d024dee5..3559e77429 100644 --- a/bootstrap/gen/rust/igla/coder/hslm_bwd_ptx.rs +++ b/bootstrap/gen/rust/igla/coder/hslm_bwd_ptx.rs @@ -1,6 +1,6 @@ // Generated from .t27 spec // DO NOT EDIT — generated by t27c -// NOT LOWERED BY THIS BACKEND: 19 test(s), 7 invariant(s), 0 function(s) with no body. +// NOT LOWERED BY THIS BACKEND: 26 test(s), 7 invariant(s), 0 function(s) with no body. // This backend emits declarations only. The spec's checks live in // the Zig and Verilog outputs; do not read this file as verified. @@ -10,6 +10,10 @@ pub const KERNEL_DX: u32 = 1; pub const KERNEL_DW: u32 = 2; +pub const KERNEL_DW_SPLITK: u32 = 3; + +pub const SPLIT_TARGET_BLOCKS: usize = 64; + pub const BM: u32 = 128; pub const BN: u32 = 128; @@ -56,6 +60,9 @@ pub fn put_name_at(out: &mut [u8], cap: usize, w: usize, kernel: u32) -> usize { if (kernel == KERNEL_DW) { return put_s(out, cap, w, "hslm_dw_s8"); } + if (kernel == KERNEL_DW_SPLITK) { + return put_s(out, cap, w, "hslm_dw_splitk_s8"); + } return w; } @@ -129,9 +136,15 @@ pub fn emit_setup_dx(out: &mut [u8], cap: usize, w: usize) -> usize { return ins(out, cap, x, "mov.b32 %r45, 0;"); } -pub fn emit_setup_dw(out: &mut [u8], cap: usize, w: usize) -> usize { - let mut x: usize = ins(out, cap, w, "shr.u32 %r12, %r4, 3;"); - x = ins(out, cap, x, "and.b32 %r13, %r4, 7;"); +pub fn emit_setup_dw(out: &mut [u8], cap: usize, w: usize, split: bool) -> usize { + let mut x: usize = w; + if split { + x = ins(out, cap, x, "and.b32 %r12, %r4, 31;"); + x = ins(out, cap, x, "shr.u32 %r13, %r4, 5;"); + } else { + x = ins(out, cap, x, "shr.u32 %r12, %r4, 3;"); + x = ins(out, cap, x, "and.b32 %r13, %r4, 7;"); + } x = ins(out, cap, x, "shl.b32 %r13, %r13, 4;"); x = ins(out, cap, x, "shl.b32 %r9, %r6, 7;"); x = ins(out, cap, x, "add.u32 %r9, %r9, %r13;"); @@ -280,18 +293,87 @@ pub fn emit_loads_dw(out: &mut [u8], cap: usize, w: usize) -> usize { return emit_transposed(out, cap, x, 50, "%r16", "%r37"); } +pub fn emit_slice(out: &mut [u8], cap: usize, w: usize) -> usize { + let mut x: usize = ins(out, cap, w, "mov.u32 %r46, %ctaid.z;"); + x = ins(out, cap, x, "mov.u32 %r47, %nctaid.z;"); + x = ins(out, cap, x, "add.u32 %r48, %r3, 31;"); + x = ins(out, cap, x, "shr.u32 %r48, %r48, 5;"); + x = ins(out, cap, x, "add.u32 %r48, %r48, %r47;"); + x = ins(out, cap, x, "sub.u32 %r48, %r48, 1;"); + x = ins(out, cap, x, "div.u32 %r48, %r48, %r47;"); + x = ins(out, cap, x, "shl.b32 %r48, %r48, 5;"); + x = ins(out, cap, x, "mul.lo.u32 %r49, %r46, %r48;"); + x = ins(out, cap, x, "add.u32 %r46, %r46, 1;"); + x = ins(out, cap, x, "mul.lo.u32 %r46, %r46, %r48;"); + x = ins(out, cap, x, "min.u32 %r3, %r3, %r46;"); + x = ins(out, cap, x, "setp.ge.u32 %p9, %r49, %r3;"); + return ins(out, cap, x, "@%p9 ret;"); +} + +pub fn emit_red_store(out: &mut [u8], cap: usize, w: usize, dr: u32, dc: u32, b: u32) -> usize { + let mut x: usize = put_s(out, cap, w, " add.u32 %r41, %r39, "); + x = put_u(out, cap, x, dr); + x = put_s(out, cap, x, ";\n"); + x = put_s(out, cap, x, " add.u32 %r42, %r40, "); + x = put_u(out, cap, x, dc); + x = put_s(out, cap, x, ";\n"); + x = ins(out, cap, x, "setp.lt.u32 %p6, %r41, %r1;"); + x = ins(out, cap, x, "setp.lt.u32 %p7, %r42, %r2;"); + x = ins(out, cap, x, "and.pred %p6, %p6, %p7;"); + x = ins(out, cap, x, "cvt.u64.u32 %rd10, %r42;"); + x = ins(out, cap, x, "mad.wide.u32 %rd10, %r41, %r2, %rd10;"); + x = ins(out, cap, x, "shl.b64 %rd10, %rd10, 2;"); + x = ins(out, cap, x, "add.u64 %rd10, %rd10, %rd3;"); + x = put_s(out, cap, x, " @%p6 red.global.add.s32 [%rd10], %acc"); + x = put_u(out, cap, x, b); + x = put_s(out, cap, x, ";\n @%p6 red.global.add.s32 [%rd10+4], %acc"); + x = put_u(out, cap, x, (b + 1)); + return put_s(out, cap, x, ";\n"); +} + +pub fn emit_red_epilogue(out: &mut [u8], cap: usize, w: usize) -> usize { + let mut x: usize = ins(out, cap, w, "shl.b32 %r39, %r6, 7;"); + x = ins(out, cap, x, "shl.b32 %r41, %r23, 6;"); + x = ins(out, cap, x, "add.u32 %r39, %r39, %r41;"); + x = ins(out, cap, x, "add.u32 %r39, %r39, %r21;"); + x = ins(out, cap, x, "shl.b32 %r40, %r5, 7;"); + x = ins(out, cap, x, "shl.b32 %r41, %r24, 5;"); + x = ins(out, cap, x, "add.u32 %r40, %r40, %r41;"); + x = ins(out, cap, x, "shl.b32 %r41, %r22, 1;"); + x = ins(out, cap, x, "add.u32 %r40, %r40, %r41;"); + let mut mi: u32 = 0; + while (mi < 4) { + let mut ni: u32 = 0; + while (ni < 4) { + let b: u32 = (((mi * 4) + ni) * 4); + x = emit_red_store(out, cap, x, (mi * 16), (ni * 8), b); + x = emit_red_store(out, cap, x, ((mi * 16) + 8), (ni * 8), (b + 2)); + ni += 1; + } + mi += 1; + } + return x; +} + pub fn emit_ptx(kernel: u32, out: &mut [u8], cap: usize) -> usize { - if ((kernel != KERNEL_DX) && (kernel != KERNEL_DW)) { + if (((kernel != KERNEL_DX) && (kernel != KERNEL_DW)) && (kernel != KERNEL_DW_SPLITK)) { return 0; } + let split: bool = (kernel == KERNEL_DW_SPLITK); let mut x: usize = emit_header(out, cap, 0, kernel); x = emit_params(out, cap, x); if (kernel == KERNEL_DX) { x = emit_setup_dx(out, cap, x); } else { - x = emit_setup_dw(out, cap, x); + x = emit_setup_dw(out, cap, x, split); + } + if split { + x = emit_slice(out, cap, x); } x = emit_warp_setup(out, cap, x); + if split { + x = ins(out, cap, x, "mov.u32 %r28, %r49;"); + } x = put_s(out, cap, x, "LOOP_K:\n"); x = ins(out, cap, x, "setp.ge.u32 %p3, %r28, %r3;"); x = ins(out, cap, x, "@%p3 bra DONE_K;"); @@ -306,7 +388,11 @@ pub fn emit_ptx(kernel: u32, out: &mut [u8], cap: usize) -> usize { x = ins(out, cap, x, "add.u32 %r28, %r28, 32;"); x = ins(out, cap, x, "bra LOOP_K;"); x = put_s(out, cap, x, "DONE_K:\n"); - x = emit_epilogue(out, cap, x); + if split { + x = emit_red_epilogue(out, cap, x); + } else { + x = emit_epilogue(out, cap, x); + } x = ins(out, cap, x, "ret;"); return put_s(out, cap, x, "}\n"); } @@ -315,6 +401,41 @@ pub fn grid_dim(x: usize) -> u32 { return (((x + 127) / 128) as u32); } +pub fn grid_z(kernel: u32, m: usize, n: usize, k: usize) -> u32 { + if (kernel != KERNEL_DW_SPLITK) { + return 1; + } + let steps: usize = ((k + 31) / 32); + let blocks: usize = (((m + 127) / 128) * ((n + 127) / 128)); + let mut s: usize = (((SPLIT_TARGET_BLOCKS + blocks) - 1) / blocks); + if (s > steps) { + s = steps; + } + if (s < 1) { + s = 1; + } + let per: usize = (((steps + s) - 1) / s); + return ((((steps + per) - 1) / per) as u32); +} + +pub fn zero_first(kernel: u32) -> bool { + return (kernel == KERNEL_DW_SPLITK); +} + +pub fn slice_lo(z: usize, nz: usize, k: usize) -> usize { + let per: usize = (((((k + 31) / 32) + nz) - 1) / nz); + return ((z * per) * 32); +} + +pub fn slice_hi(z: usize, nz: usize, k: usize) -> usize { + let per: usize = (((((k + 31) / 32) + nz) - 1) / nz); + let hi: usize = (((z + 1) * per) * 32); + if (hi < k) { + return hi; + } + return k; +} + pub fn shape_ok(kernel: u32, m: usize, n: usize, k: usize) -> bool { if (((m < 1) || (n < 16)) || ((n % 16) != 0)) { return false; @@ -325,14 +446,14 @@ pub fn shape_ok(kernel: u32, m: usize, n: usize, k: usize) -> bool { if (kernel == KERNEL_DX) { return (((k >= 16) && ((k % 16) == 0)) && (k <= K_MAX_DX)); } - if (kernel == KERNEL_DW) { + if ((kernel == KERNEL_DW) || (kernel == KERNEL_DW_SPLITK)) { return ((((m % 16) == 0) && (k >= 1)) && (k <= K_MAX_DW)); } return false; } pub fn b_len(kernel: u32, n: usize, k: usize) -> usize { - if (kernel == KERNEL_DW) { + if ((kernel == KERNEL_DW) || (kernel == KERNEL_DW_SPLITK)) { return (k * n); } return ((k * n) / 4);