Skip to content

HSLM dW kernel: split-K variant (kernel 3) to close the 11x gap to CUTLASS s8 at K=1782 #10895

Open

Description

@gHashTag

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

  1. added 20 commits that reference this issue on Oct 19, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions