[Bugfix][ROCm] RDNA3 W4A16 MoE: handle the N-first CT weight layout

[JartX]

Purpose

#52798 ("Use canonical N-first weight format for CT WNA16 MoE", merged
September 28) changed the layout compressed-tensors WNA16 MoE weights reach
the backends in, but the RDNA3 (gfx1100) path was not updated. Since then any
compressed-tensors W4A16 MoE checkpoint fails on RDNA3 at the first forward,
with an assertion on w1_zp in the experts.

Two things broke:

• _process_weights_rdna3 still read the weights as K-first
  ([E, K // 8, N], scales [E, groups, N]), which is what
  moe_gptq_gemm_rdna3 needs, but they now arrive N-first
  ([E, N, K // 8], scales [E, N, groups]).
• CompressedTensorsWNA16MoEMethod registered the synthesized zero points
  only for the CPU backend. The RDNA3 kernel needs them too, even for
  symmetric checkpoints (the oracle synthesizes them), so w13/w2_qzeros
  never reached the layer.

What changes

• oracle/int_wna16.py: _process_weights_rdna3 transposes weights and
  scales to the kernel's K-first layout before the exllama shuffle.
• compressed_tensors_moe_wna16.py: register the synthesized zero points
  for WNA16MoEBackend.RDNA3 as well as CPU.

Nothing changes for other backends.

Not a duplicate

upstream/main (df417f7) still has the K-first read. I searched open PRs
for rdna3 moe, rdna3 w4a16, gfx1100 moe, moe_gptq_gemm_rdna3,
rdna3 compressed-tensors and 52798. The only open PR touching
int_wna16.py for RDNA3 is #59783, which adds a group-size check in backend
selection and does not touch weight processing.

Test plan

pytest tests/kernels/quantization/test_rdna3_moe_w4a16.py -v
pre-commit run --files <changed files>

• A regression test, test_process_weights_from_compressed_tensors_layout:
  it builds weights the way compressed-tensors hands them to the oracle
  (N-first, packed [E, N, K // 8], scales [E, N, groups]), runs them
  through _process_weights_rdna3 and a full MoE forward (w1, SiLU, w2,
  top-k reduce), and compares with a dequantized PyTorch reference. The
  existing tests build the kernel's K-first layout directly and call the op,
  so they never exercised weight processing, which is why this regression
  went unnoticed.
• End to end on 4x RX 7900 XTX with a symmetric compressed-tensors
  checkpoint (Qwen3.6-35B-A3B GPTQ W4A16, g32, pack-quantized): fails on
  main without this change, starts and passes GSM8K with it (below).

Test results

On gfx1100, in an image built from main (df417f7):

• Without this change, the new test fails for fp16 and bf16: the kernel
  rejects the weights (group size (K/groups = 1) must be a multiple of 32),
  because it reads N-first data as K-first.
• With this change (its int_wna16.py applied on the same main and mounted
  into that image; the C++ kernel is unchanged): the new test passes, and test_rdna3_moe_w4a16.py passes
  in full (99 passed).
• pre-commit passes, except update-dockerfile-graph, which needs
  /bin/bash and could not run on my machine.

End to end, on an image built from main (df417f7) for gfx1100, 4x RX
7900 XTX, TP4, --dtype float16, Qwen3.6-35B-A3B GPTQ W4A16 g32 (symmetric,
compressed-tensors):

• Without this change the engine fails at the first forward:
  assert self.w1_zp is not None and self.w2_zp is not None.
• With this change it selects Rdna3WNA16Experts, starts, and scores 97.5% on
  GSM8K (tests/evals/gsm8k/gsm8k_eval.py --use-chat-completions --chat-template-kwargs '{"enable_thinking": false}' --num-questions 200 --temperature 0.6 --top-p 0.95 --max-tokens 2048, 5-shot, 0% invalid).

The two files of this change were mounted over that image; the C++ kernel is
unchanged. The stock RCCL in the image could not set up P2P between these
GPUs, so it was rebuilt; this is unrelated to this PR.

AI assistance

I used AI assistance to find the cause and write this change. I have
reviewed every changed line and will run the tests above before marking it
ready.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[tjtanaa]

LGTM

[tjtanaa]

/ci run

[github-actions]

✅ Triggered Buildkite CI #93601 for commit efe51f555eef.