Skip to content

fix(qwen36): restore MoE device-tilemap (MOE_GPU) correctness — masked down-dequant no-op#602

Closed
James-CUDA wants to merge 1 commit into
gittensor-ai-lab:mainfrom
James-CUDA:fix/qwen36-moe-gpu-mask-dequant
Closed

fix(qwen36): restore MoE device-tilemap (MOE_GPU) correctness — masked down-dequant no-op#602
James-CUDA wants to merge 1 commit into
gittensor-ai-lab:mainfrom
James-CUDA:fix/qwen36-moe-gpu-mask-dequant

Conversation

@James-CUDA

Copy link
Copy Markdown
Contributor

Summary

Correctness fix (not a speed change). SPARKINFER_PREFILL_MOE_GPU=1 — the Qwen3.6 device-tilemap prefill path that skips the per-layer D2H counts sync — is silently broken on main: it produces garbage output.

Root cause: the MoE down projection's masked Q→int8 dequant runs with cols = moe_ffn = 512, which the vectorized fast-mask kernel declines (it only accepts cols that are a multiple of 1024). launch_gguf_dequant_rows_i8_mask ignored that false return, so the int8 down weights were never written and the down GEMM ran on stale/garbage data. This is the #586 corruption that forced the path default-off.

Fix: launch_gguf_dequant_rows_i8_mask now falls back to a correct masked slow kernel (single-pass, register-resident, identical dequant/amax/round to the host deq_rows_i8_kernel; dead experts exit via the counts mask) when the fast path declines. The default host-tilemap path is untouched, and SPARKINFER_PREFILL_MOE_GPU stays default OFF (the device path is not currently a prefill speedup — it must dequant all experts per group rather than the host path's live-only coalesced runs — so this fixes the opt-in flag rather than changing perf).

Proof (correctness, not speed)

Not a speedup PR — no bench.sh before/after and the RTX-5090 speedup box is intentionally left unticked. Evidence is the prefill fidelity gate, on RTX 5090 / Qwen3.6-35B-A3B-UD-Q4_K_M, qwen3_gguf_prefill_check <model> 512 16 with SPARKINFER_PREFILL_MOE_GPU=1:

TOP1 KL
before (main) 0.375 (6/16) 0.349
after (this PR) 0.9375 (15/16) 0.010

The host-tilemap default path (MOE_GPU=0) is unchanged (TOP1 1.0 @512, byte-identical to before).

  • Tested on RTX 5090 (sm_120) — correctness only; no speed claim
# main (MOE_GPU=1):   TOP1 6/16 0.3750   KL 0.34935
# this PR (MOE_GPU=1): TOP1 15/16 0.9375  KL 0.00974

…d down dequant no-op

SPARKINFER_PREFILL_MOE_GPU=1 (device tilemap, skips the per-layer D2H counts
sync) is silently broken on main: the down projection's masked Q->int8 dequant
has cols=mffn=512, which the vectorized fast-mask kernel declines (cols not a
multiple of 1024), and launch_gguf_dequant_rows_i8_mask ignored that false
return — so the int8 down weights were never written and the down GEMM ran on
stale/garbage data. Batched-vs-token-loop prefill_check collapses to TOP1 ~0.375
/ KL ~0.35 @512 (this is the gittensor-ai-lab#586 corruption that forced the path default-off).

Fix: launch_gguf_dequant_rows_i8_mask now falls back to a correct masked slow
kernel (single-pass, register-resident, same dequant/amax/round as the host
deq_rows_i8_kernel; dead experts exit via the counts mask) when the fast path
declines. Restores prefill_check to TOP1 0.94 / KL ~0.01 @512 with MOE_GPU=1.

Keeps SPARKINFER_PREFILL_MOE_GPU default OFF: the device path is not currently a
prefill speedup (it must dequant all experts per group rather than the host
path's live-only coalesced runs), so this is a correctness fix for the opt-in
flag, not a perf change. The default host-tilemap path is unaffected.
@github-actions

Copy link
Copy Markdown

Closed — RTX 5090 checkbox not ticked

This PR was auto-closed because the template includes Tested on RTX 5090 as - [ ] (unchecked). Evaluation is opt-in — tick the box only after a real 5090 run.

To submit for review:

  1. Edit this PR description (you can edit while closed): change to - [x] Tested on RTX 5090
  2. Fill the decode and/or prefill before → after tables with real bench/scripts/bench.sh numbers showing improvement
  3. Reopen this PR

If this PR does not need GPU eval (e.g. docs-only), remove the proof-of-speedup section from the description instead of leaving an unchecked box.

CONTRIBUTING.md

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

@github-actions github-actions Bot closed this Jul 24, 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