You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
{{ message }}
Repository navigation
HSLM dW kernel: split-K variant (kernel 3) to close the 11x gap to CUTLASS s8 at K=1782 #10895
The dW kernel of PR #10855 (hslm_dw_s8, kernel 2 of specs/igla/coder/hslm_bwd_ptx.t27) is 11x behind CUTLASS s8 at K = 1782 rows, the batched HSLM step, and fine at K = 27. On an RTX 4090 it takes 0.075 ms against 0.0066-0.0068 ms for the best CUTLASS s8 kernel. The cause: a 256 x 256 or 256 x 736 output is only 4-12 blocks of 128 x 128 on a 128-SM part, and each block walks all 56 k-steps alone. This issue adds a split-K variant as a new kernel id.
Boundary
specs/igla/coder/hslm_bwd_ptx.t27
bootstrap/gen/rust/igla/coder/hslm_bwd_ptx.rs
.trinity/seals/coder_igla-coder-hslm-bwd-ptx.json
bootstrap/src/host/driver.rs
User Scenarios
Given the spec, when the t27c test runner executes it on the Railway lab, then every test passes: the text pins of the new kernel, the slice bounds (every k belongs to exactly one slice), and the integer reference summed over slices equals the unsplit reference.
Given an RTX 4090, when t27c ptx-run --kernel 3 runs dW at m x n = 256 x 256, 256 x 736 and 736 x 256 over k = 1782, then it prints mismatches=0 and the time per call.
Given the same shapes, when the time is set against the best CUTLASS s8 kernel measured on the same pod, then kernel 3 is within about 1.5x of it.
Requirements
FR-001: A new kernel id (3, hslm_dw_splitk_s8) MUST compute the same C as kernel 2 (C[i][j] = sum_r X[r][i] * dY[r][j], i32, k <= 131071), with grid z = the number of k slices and each block summing its slice of k.
FR-002: The partial sums of the slices MUST be combined exactly in i32 (red/atom global add, or a second pass); the result MUST equal the integer CPU reference bit for bit.
FR-003: The slice count MUST come from a spec function of (m, n, k) under which no slice is empty, and 1 slice MUST reproduce kernel 2's result.
FR-004: Kernels 0, 1 and 2, their CPU references and their text pins MUST stay unchanged.
FR-005: The hand-written foreign glue MUST NOT grow: the host change is made in place on existing lines of bootstrap/src/host/driver.rs (the launch grid z and a zero fill of C before each launch).
Success Criteria
The spec's tests pass on the lab, and gen-rust and the seal are regenerated there.
t27c ptx-run --kernel 3 prints mismatches=0 for 256x256, 256x736 and 736x256 at k = 1782 (and at k = 27 as a regression check) on an RTX 4090.
The PR body has a table of ms per call for kernel 2, kernel 3 and CUTLASS s8, with kernel 3 within about 1.5x of CUTLASS s8 at k = 1782, or an honest statement of the gap if it is not.
Activity
spec(igla): hslm_bwd_ptx -- pin the kernel 3 text length (Refs #10895)
gen+seal: hslm_bwd_ptx with kernel 3, regenerated on the lab (Refs #1…
spec(igla): hslm_bwd_ptx -- kernel 3 maps a warp onto one chunk and 3…
gen+seal: hslm_bwd_ptx with the kernel 3 thread remap, regenerated on…
spec(igla): hslm_bwd_ptx -- kernel 3 aims at 64 blocks, the floor of …
gen+seal: hslm_bwd_ptx with target 64, regenerated on the lab (Refs #…
spec(igla): hslm_bwd_ptx -- kernel 3, dW split along k with red.globa…
spec(igla): hslm_bwd_ptx -- pin the kernel 3 text length (Refs #10895)
spec(igla): hslm_bwd_ptx -- kernel 3 maps a warp onto one chunk and 3…
spec(igla): hslm_bwd_ptx -- kernel 3 aims at 64 blocks, the floor of …
gen+seal: hslm_bwd_ptx with kernel 3, regenerated on the lab (Refs #1…
spec(igla): hslm_bwd_ptx -- kernel 3, dW split along k with red.globa…
spec(igla): hslm_bwd_ptx -- pin the kernel 3 text length (Refs #10895)
spec(igla): hslm_bwd_ptx -- kernel 3 maps a warp onto one chunk and 3…
spec(igla): hslm_bwd_ptx -- kernel 3 aims at 64 blocks, the floor of …
gen+seal: hslm_bwd_ptx with kernel 3, regenerated on the lab (Refs #1…
spec(igla): hslm_bwd_ptx -- kernel 3, dW split along k with red.globa…
spec(igla): hslm_bwd_ptx -- pin the kernel 3 text length (Refs #10895)
spec(igla): hslm_bwd_ptx -- kernel 3 maps a warp onto one chunk and 3…
spec(igla): hslm_bwd_ptx -- kernel 3 aims at 64 blocks, the floor of …
gen+seal: hslm_bwd_ptx with kernel 3, regenerated on the lab (Refs #1…
spec(igla): hslm_bwd_ptx -- kernel 3, dW split along k with red.globa…
spec(igla): hslm_bwd_ptx -- pin the kernel 3 text length (Refs #10895)
spec(igla): hslm_bwd_ptx -- kernel 3 maps a warp onto one chunk and 3…
spec(igla): hslm_bwd_ptx -- kernel 3 aims at 64 blocks, the floor of …
gen+seal: hslm_bwd_ptx with kernel 3 on top of ptx_emit, regenerated …