cuda: compact Blackwell NVFP4 MoE work scheduling (+10% to +15% prefill, NVFP4-only)

[coder543]

⚠️ This PR does not meet the contributor requirements, this is draft only unless someone else takes ownership. ⚠️

If someone wants to take ownership of this PR and drive it to completion, or just borrow ideas from it, they are welcome to. This PR is one of two that came out of a multi-hour GPT-5.6-Sol optimization session, as I wanted to experiment with it and see if it could actually find any low hanging fruit in the CUDA kernels. I was obviously involved in guiding the agent, and ensured that it did proper testing against a variety of models, but I have not had time to review every line of change.

I figured it was better to make this available in case anyone else cares enough to drive it forward, rather than throwing away the experiments entirely, but I am not requesting any reviews on this.

Overview

It makes NVFP4 10% to 15% faster on my DGX Spark for the two NVFP4 models that were tested.

Requirements

• I have read and agree with the contributing guidelines: Yes, which is why this is not being submitted for review.
• AI usage disclosure: Yes, all of it.

Everything below this point is written by Codex/GPT-5.6-Sol:

---

The native NVFP4 MUL_MAT_ID path launches a rectangular expert grid. With sparse routing, many of those blocks immediately discover that their expert has no rows to process. Build a compact (expert, tile) worklist and stripe only useful jobs across the GPU instead.

Also fold the per-expert NVFP4 output scale into MMQ writeback. This avoids the GET_ROWS and post-matmul MUL sequence without changing multiplication order. Both changes are gated on native NVFP4 MMQ with Blackwell MMA.

DGX Spark / GB10 pp8192 benchmark:

llama-bench -p 8192 -n 0 -b 16384 -ub 2048 -r 7 -fa 1 -lm none

Values are median prompt tokens/s. Each row compares this commit with its upstream parent:

north-mini-code NVFP4 3185.2 -> 3677.4 (+15.5%)
qwen3.6-35b-a3b NVFP4 2863.9 -> 3151.1 (+10.0%)
laguna-xs Q4_K 2847.4 -> 2856.7 (+0.3%, noise)
qwen3.6-35b-a3b BF16 1630.6 -> 1631.1 (+0.0%)
qwen3.6-27b BF16 dense 942.7 -> 951.3 (+0.9%, noise)
gemma-4-26b-a4b Q4_K 3114.3 -> 3113.8 (-0.0%)
gpt-oss-20b MXFP4 4691.4 -> 4694.8 (+0.1%)
gpt-oss-120b MXFP4 2399.6 -> 2393.9 (-0.2%, noise)

The two NVFP4 MoE models show the intended benefit. The other formats do not enter this path and show no measurable regression. GPT-OSS-20B was repeated in reverse order to exclude an initial frequency-ordering effect.

Deterministic output validation used a 31.9 KiB README prompt, 32 generated tokens, temperature 0, seed 12345, batch 16384, ubatch 2048, Flash Attention, and load mode none. Output was byte-identical to the parent for all eight models above. Laguna and Gemma used their Jinja chat templates.

The standalone commit also passes all 73 NVFP4 MUL_MAT_ID backend tests.

[am17an]

Also fold the per-expert NVFP4 output scale into MMQ writeback. This avoids the GET_ROWS and post-matmul MUL sequence without changing multiplication order

My guess is most of the benefit comes from here. The other optimization is already implemented for all data types. cc @ORippler incase you were doing something similar

[ORippler]

@ORippler incase you were doing something similar

yeah we are working on this as well, and are/were going to file a PR in that direction "soonish"