Skip to content

cuda: warp-per-row PTQ1_0 mat-vec for MUL_MAT_ID when K <= 4096 - #415

Merged
bri-prism merged 1 commit into
prismfrom
cuda/ptq1-moe-warprow
Oct 29, 2026
Merged

bri-prism merged 1 commit into
prismfrom
cuda/ptq1-moe-warprow

Conversation

@bri-prism

Copy link
Copy Markdown
Collaborator

What

  • PTQ1_0 MUL_MAT_ID mat-vec with one column and K <= 4096 (at most 32 blocks per row): each of the 4 warps in the
    CTA now owns one row, and its lanes stride that row's 128-weight blocks on the planar layout.

Why

  • The generic loop spreads the blocks of a row over all 128 threads of the CTA. With 32 blocks or fewer per row, only
    the first warp loads anything, so MoE expert mat-vecs at these shapes ran at about 0.5 TB/s on H200.

How

  • Same per-row accumulation order, same warp_reduce_sum and the same epilogue (bias, gate, GLU), so outputs are
    bit-identical to the current path.
  • Other types, other K, the non-ids path, HIP and MUSA are unchanged.

Testing

  • test-backend-ops MUL_MAT_ID, MUL_MAT_VEC_FUSION and MUL_MAT pass (CUDA vs CPU).
  • Greedy output is byte-identical before and after on several PTQ1_0 models, MoE and dense.
  • H200: a fused gate/up expert mat-vec at K = 4096 goes from 63 to 32 us. End-to-end MoE decode improves by 2-15%,
    depending on how much of the token time the expert mat-vecs take. No change on dense PTQ1_0 models.
  • Tested on H200 only.

The generic loop strides the 128-weight blocks over all threads of the CTA, so
at K <= 4096 (<= 32 blocks per row) only the first warp loads. Give each warp
its own row and let the lanes stride the blocks, with the same per-row
accumulation order and epilogue, so outputs are unchanged.
@bri-prism
bri-prism merged commit 9d9efea into prism Oct 29, 2026
8 checks passed
bri-prism added a commit that referenced this pull request Oct 29, 2026
Resolve the prism.hadamard loader conflict by keeping load_hparams_hadamard()
and adding Gemma 4 to its arch whitelist (#413). Add the SWIGLU_CLAMP case to
the PTQ1_0 warp-per-row epilogue from #415, since this branch has that op.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant