fix(qmoe): compute down_proj (fc2) in fp32 to avoid fp16 overflow

[AMDmoore]

Summary

A single QMoE expert's raw down_proj (fc2) channel can exceed the fp16 max (65504) and overflow to inf before the routing weight scales it back into range. Narrowing that raw GEMM result to fp16 at store time therefore corrupts the MoE output (observed as NaN logits / perplexity blow-up on gpt-oss-120b, e.g. layer 34).

This PR keeps the fc2 GEMM result in fp32 until the weighted cross-expert accumulation, then narrows to fp16 exactly once — so w_e * (W_dn·h_e + b_e) never rounds/overflows through fp16.

What changed

• matmul_nbits_kernel.hip — template GemmFp16U4Impl and the two WMMA entry wrappers (MatMulNBitsWMMA_ZP / _NoZP) on an output type OutT (default _Float16, so every existing fp16 instantiation is byte-identical — only the final store narrowing changes; the fp32 accumulator was always fp32). hip_matmul_nbits_fp32out now dispatches on M exactly like the fp16 path, with fixed configs (no autotune, no autotune cache touched):
  • M >= 16 && K % 32 == 0 → fused-dequant WMMA fp32 (128×128, bounds-checked; needs fp16 zeros)
  • otherwise → row-major int4 GEMV fp32 (valid for any M; needs uint8 zeros)
• hip_custom_kernels.h — add zero_points_fp16 parameter to hip_matmul_nbits_fp32out.
• qmoe.cpp (prefill / multi-token path) — run fc2 into an fp32 scratch with no routing weight and no bias; prepare fp16 zeros only when the WMMA path is reachable (count >= 16 && inter % 32 == 0, identical to the launcher's WMMA gate so asym never silently falls back to zp=8). The scatter applies the routing weight to the full (GEMM + bias) result and accumulates across experts in fp32.
• qmoe_kernel.hip — scatter_add_fp32 reads fp32 expert output and applies the routing weight to the whole (GEMM + bias); the old pre-scaling helper was removed.

Scope / notes

• Decode (num_tokens == 1) uses the separate fused-decode kernel and is unchanged by this PR.
• fp16 paths (Llama / gpt-oss decode, all non-fp32out callers) are byte-identical — OutT defaults preserve every existing instantiation.

Test plan

• gpt-oss-120b: confirm NaN logits / PPL blow-up is resolved (prefill path), baseline 4b4dfa98 vs this branch.
• Confirm prefill uses WMMA fp32 for large per-expert count (long prompt) and GEMV fp32 for small count.
• Regression: Llama-8B / gpt-oss-20b logits unchanged vs baseline (fp16 paths byte-identical).

[github-actions]

Thanks for opening a PR!

This project follows LLVM's incremental-development and AI-tool-use
guidance. See CONTRIBUTING.md
for the project workflow.

Before requesting review, please check that:

1. The change is focused. Substantial work links the relevant issue
   or design discussion.
2. The PR documents relevant test results and updates affected
   documentation.
3. If AI tools provided substantial assistance, the description
   explains what was assisted and how it was validated, and commit
   trailers identify the tool. The contributor has reviewed and
   understands the result.

Reviewers are assigned through
CODEOWNERS where ownership
is configured.

[github-actions]

L2 Accuracy Results (EP vs CPU)

Model	Combined L2	Total Elems	Skipped NaN/Inf
conv_test_hybrid	4.8668E-07	64	0
GroupQueryAttention_seq256	25.2366	2621440	0
MatMulNBits_o_seq128	259.906	368640	0
QMoE_seq128	17.9848	368640	0

Threshold: 0.01 | Run: 3701 - Commit: 181ee1a

[github-actions]

MorphiZen EP Performance Results

Model	QPS	Session (s)	1st Infer (ms)	CPU%	Mem (MB)
full_model_seq128	7.50	5.98	362	3	1245
GroupQueryAttention_seq128	4383.03	1.68355	11	6	311
matmul_down_seq128	520.66	2.31	73	3	352

EPContext Export Performance

Model	QPS	Session (s)	1st Infer (ms)	CPU%	Mem (MB)
full_model_seq128	7.55	44.42	360	3	15590

EPContext Import Performance

Model	QPS	Session (s)	1st Infer (ms)	CPU%	Mem (MB)
full_model_seq128	7.53	9.75	362	3	15761

OGA Benchmark Results

Model	Warmup	Reps	Prompt Len	Gen Tokens	TTFT (ms)	TPS	Peak Mem (GB)	GPU Mem (GB)
gpt-oss-20b-webgpu-int4-rtn-block-32	1	5	128	128	192.2	80.1	1.33	13.53
Llama-3.1-8B-awq-g128-int4-asym-fp16-onnx-dml	1	5	128	128	316.0	40.8	1.22	6.43

OGA Wheel Smoke (Python benchmark_e2e.py)

Model	TTFT (ms)	TPS
Llama-3.1-8B-awq-g128-int4-asym-fp16-onnx-dml	193	39.7

Run: 3701 - Commit: 181ee1a