[WebGPU] Optimize MatMulNBits wide-tile shader

[daijh]

Description

• Store the A workgroup tile as [kTileM][kBlockVecs/2][2] pairs of vecs so the two vecs consumed per dequantized weight column are read in one access.
• Accumulate directly in output_element_t (e.g., f16 for f16 models) instead of hardcoding to f32.

Intel Panther Lake

	Prefill Length	Prefill TPS	Optimized Prefill TPS	Improvement
gpt-oss-20b-ONNX	128	305.70	347.71	114%
gpt-oss-20b-ONNX	1024	396.50	451.51	114%
Phi-4-mini-instruct-ONNX	128	515.90	597.81	116%
Phi-4-mini-instruct-ONNX	1024	615.39	747.07	121%

[1] https://huggingface.co/onnx-community/gpt-oss-20b-ONNX
[2] https://huggingface.co/onnx-community/Phi-4-mini-instruct-ONNX

Motivation and Context

See above.

[daijh]

@qjia7 @hariharans29 PTAL

[qjia7]

Accumulate directly in output_element_t (e.g., f16 for f16 models) instead of hardcoding to f32.

Please check #29599, I think using f32 accumulators is the right direction.

Could you split the refactoring changes into a separate PR? That would make it much easier to review and focus on the optimization work.

[daijh]

Let's leave this refactoring for a follow-up PR.

[hariharans29]

Accumulate directly in output_element_t (e.g., f16 for f16 models) instead of hardcoding to f32.

Please check #29599, I think using f32 accumulators is the right direction.

Could you split the refactoring changes into a separate PR? That would make it much easier to review and focus on the optimization work.

I agree - isn't it safer to accumulate in fp32 ? Accumulating in fp16 universally across all workloads (if the op is fp16) seems risky ?

[RobertoReale]

Since #29599 was referenced here — a data point on the accumulator part of this change (the tile-layout optimization itself looks like a clear win and is orthogonal):

The f32 accumulator in the wide-tile kernel isn't just precision hygiene — it's load-bearing for correctness. With f16 accumulation, partial dot-product sums overflow 65504 and saturate to +Inf once K approaches ~2048 (deterministically, not as a gradual accuracy loss), which then propagates NaN through LayerNorm/Softmax. That's the failure class of #26732. Notably, both benchmark models here are in that range (gpt-oss-20b K≈2880, Phi-4-mini K=3072) — the table reports prefill TPS, but was output correctness checked against the f32-accumulator baseline on those runs?

Also, since the PR changes tile layout and accumulator precision together, it would be useful to benchmark them separately — the memory-access refactoring may account for most of the gain, in which case you could keep the speedup without reintroducing the overflow.

If f32 accumulation does turn out to be a measurable cost on Intel in isolation, gating the promotion on K above a threshold (overflow risk scales with reduction length) was floated in #29599 as a compromise that keeps small-K workloads untouched — happy to coordinate there.

[daijh]

Accumulate directly in output_element_t (e.g., f16 for f16 models) instead of hardcoding to f32.

Please check #29599, I think using f32 accumulators is the right direction.
Could you split the refactoring changes into a separate PR? That would make it much easier to review and focus on the optimization work.

I agree - isn't it safer to accumulate in fp32 ? Accumulating in fp16 universally across all workloads (if the op is fp16) seems risky ?

I think it would be helpful to reproduce the overflow issue in the LLM first, before we decide on our next steps.

Alternatively, we could use the accuracy_level attribute as a hint if accuracy or overflow is a concern. E.g. set accuracy_level to fp32 to enforce fp32 accumulation.
[1] https://github.com/microsoft/onnxruntime/blob/main/docs/ContribOperators.md#commicrosoftmatmulnbits

[RobertoReale]

I think it would be helpful to reproduce the overflow issue in the LLM first, before we decide on our next steps.

Done — details in #29599 (comment) (the non-repro had a concrete cause: the Hub's fp16 Gemma weights were re-exported on 2025-12-02 with Clip nodes inserted as a model-side workaround for this exact overflow). Summary: pinning onnx-community/gemma-3-270m-it-ONNX to the pre-mitigation revision 2950c41f, stock onnxruntime-web on WebGPU produces garbage output while the same revision is correct on WASM — and in my runs the WebGPU adapter was an Intel Iris Xe (gen-12lp), so the overflow is live on Intel hardware too.

On accuracy_level: using it as an explicit opt-down sounds workable — a model that knowingly tolerates f16 accumulation can set accuracy_level = 2 and keep the fastest path, similar to how the dp4a kernels already gate on accuracy_level = 4. But two constraints for it to actually fix #26732:

1. The unset default (0) must mean f32 accumulation. None of the affected models in the wild set the attribute (transformers.js exports don't), so an f16-when-unset default leaves all of them broken.
2. It only covers MatMulNBits. The plain fp16 MatMul path — which is what the Gemma fp16 case goes through — has no accuracy_level attribute, so it needs the f32 accumulator unconditionally regardless.

[RobertoReale]

Data point relevant to the f16-accumulator question here: following up on the repro request, I ran raw-WGSL f16 accumulation probes across adapters and backends — full write-up in #29599 (comment).

Short version: on Intel (D3D12) the shader compiler evaluates straight-line/unrolled f16 acc += chains at higher precision and only rounds across loop iterations, so ORT's unrolled native kernels get f32 accumulation silently from the driver — a synthetic MatMul whose f16 partial sums provably exceed 65504 still returns the exact result on the native EP on my Iris Xe. On NVIDIA (Vulkan) the same two-term f16 expression rounds strictly and yields Inf. Both behaviors are WGSL-conformant.

The implication for this PR: correctness results obtained with f16 accumulators on Intel don't transfer to other vendors (or even to Intel on a different backend, or to looped codegen). Demoting the wide-tile accumulator to output_element_t may well benchmark clean and produce correct outputs on Panther Lake, while overflowing on hardware that rounds f16 strictly. An explicit f32 accumulator is the only portable guarantee — and on drivers that already promote intermediates it should be close to free, which matches the neutral perf discussion in #29599.

[daijh]

As #29599 (comment)

Since this is an NVIDIA-specific issue, could you gate the fix behind a vendor string check?
Please let me know if any issues arise on Intel hardware, and I'd be happy to follow up.

Since the NVIDIA-specific logic requires fp32 accumulators and will be gated by a vendor string check.
I am fine moving forward with the current PR.
We can leave the fp32 accumulators implementation to PR #29599.

[RobertoReale]

Flagging one interaction for whoever reviews this PR, since the accumulator part is being deferred to #29599 (full reasoning in #29599 (comment)):

The proposal on #29599 is to gate the f32 accumulator behind an NVIDIA vendor-string check. If that lands and this PR demotes the wide-tile accumulator to output_element_t, then q4f16 models routed through the wide-tile kernel overflow on any non-NVIDIA config that rounds f16 strictly — with no f32 fallback. That's the SmolLM2 / Gemma-q4f16 failure class of #26732, and it's not NVIDIA-specific: my original repro of that overflow was on an Intel Iris Xe (via JSEP's looped codegen), and the same Iris Xe rounds strictly on the Vulkan backend. f16 accumulation happening to be safe on Intel/D3D12 here is a property of the driver promoting unrolled chains, not of the vendor.

The tile-layout half of this PR looks like a clean, orthogonal win. The only ask is to keep the accumulator in f32 (or make the demotion opt-in via accuracy_level, defaulting to f32) so the two PRs don't combine into a wide-tile overflow regression. Happy to help with a separate layout-only-vs-accumulator benchmark if that's useful.