Skip to content

fix(gemv): instantiate mmvq_q4k_rows for K=6656 (Muse Glimmer) - #879

Closed
Tet-9 wants to merge 1 commit into
gittensor-ai-lab:mainfrom
Tet-9:main
Closed

fix(gemv): instantiate mmvq_q4k_rows for K=6656 (Muse Glimmer)#879
Tet-9 wants to merge 1 commit into
gittensor-ai-lab:mainfrom
Tet-9:main

Conversation

@Tet-9

@Tet-9 Tet-9 commented Aug 18, 2026

Copy link
Copy Markdown

What

The batched-row Q4_K launchers (launch_mmvq_q4k_rows / launch_mmvq_rows_f32 in
kernels/csrc/cuda/gemm/gemv.cu) already supported K=2048/4096/5120/6144 but not
K=6656 — Muse Glimmer's hidden size — even though the single-row launch_mmvq_q4k
already handles K=6656 via si_mmvq_q4k_kfixed_kernel<__nv_bfloat16, 26>.

Any batched verify-path call (e.g. speculative decode, M>1) against Muse Glimmer
hit the unsupported-K guard and either silently declined or fell back to N
serial single-row calls instead of the fused batched path.

Change

  • Adds NSUPER=26 explicit instantiations of si_mmvq_q4k_rows_exact_kernel
    for {__nv_bfloat16, float} × {MMAX=6, MMAX=8}, matching the existing
    pattern for K=2048/4096.
  • Extends the K guard and dispatch branch in launch_mmvq_q4k_rows to cover
    K=6656.
  • launch_mmvq_rows_f32 already had a K=6656 guard entry but its fallback
    dispatch macro silently mapped any unlisted K (including 6656) to KB=24 —
    fixed to dispatch KB=26 explicitly for K=6656.
  • Adds a standalone correctness+timing check,
    runtime/examples/mmvq_q4k_rows_6656_check.cpp, isolating this kernel at
    K=6656 from the model loader and KV cache (same pattern as
    nvfp4_gemm_check.cpp).

Verification

Tested on RTX 5090 (sm_120):

  • Batched output is byte-identical to N independent serial launch_mmvq_q4k
    calls at K=6656.
  • 1.79x faster than the serial fallback (0.0312ms batched vs 0.0558ms serial
    for N=4 rows, MN=8192).
    K=6656 rows: batch 0.0312 ms, 4 serial calls 0.0558 ms, 1.79x
    [PASS] K=6656 batched-row Q4_K matches serial reference

Also re-ran the existing dflash_gdn_checkpoint_gpu_test (K=2048/4096 paths)
to confirm no regression — still passes.

The batched-row Q4_K launcher (launch_mmvq_q4k_rows / launch_mmvq_rows_f32)
already supported K=2048/4096/5120/6144 but not K=6656 (Muse Glimmer's
hidden size), even though the single-row launch_mmvq_q4k already handles
K=6656 via si_mmvq_q4k_kfixed_kernel<__nv_bfloat16, 26>. Any batched
verify-path call (e.g. speculative decode) against Muse Glimmer at M>1
silently fell back or declined.

Adds NSUPER=26 instantiations of si_mmvq_q4k_rows_exact_kernel and
extends both dispatch functions' K guards and branches to cover 6656.

Verified on RTX 5090: batched output byte-identical to N serial
single-row calls, 1.79x faster than the serial fallback at K=6656
(runtime/examples/mmvq_q4k_rows_6656_check.cpp).
@github-actions

Copy link
Copy Markdown

Closed — RTX 5090 checkbox not ticked

This PR touches runtime/ and was auto-closed because Tested on RTX 5090 is not ticked (- [x]). Runtime changes must show a real 5090 before/after benchmark to enter the eval queue.

To submit for review:

  1. Edit this PR description (you can edit while closed): add or restore the proof-of-speedup section and tick - [x] Tested on RTX 5090
  2. Fill the decode and/or prefill before → after tables with real bench/scripts/bench.sh numbers
  3. Reopen this PR

Non-speed runtime fixes that should not be evaluated need a maintainer hold label.

CONTRIBUTING.md

Automated by rtx5090-required CI. Maintainers and hold PRs are exempt.

@github-actions github-actions Bot closed this Aug 18, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant