[rmsnorm] Support FP32 weights with FP16/BF16 activations

[jhinpan]

Summary

Add the mixed-weight training contract used by Quack: FP16/BF16 activations with FP32 RMSNorm weights, while preserving all existing same-dtype paths.

• Use independent activation and weight element types in plain and fused-add forward kernels, atomic backward, and staged backward.
• Keep out/dx in activation dtype and dweight in the original weight dtype.
• Centralize the allowed dtype pairs: matching f32/f16/bf16, or f16/bf16 activation with f32 weight.
• Keep FP32 weights vectorized as two 128-bit, four-element loads per vec8 activation tile.
• Include weight dtype in forward/backward wrapper caches, direct JIT specialization, and the RMSNorm autotune key. The offline artifact identity from [3/5] autotune: add offline config artifacts (#770) #786 is derived from this same key.
• Preserve the original positional rmsnorm_direct(..., dtype_str, BLOCK_THREADS, stream) API.
• Add mixed forward/plain-backward/fused-backward benchmark rows to the existing MI325 (gfx942) and MI355 (gfx950) benchmark workflow.

Verification

• gfx950 single-GPU RMSNorm suite: 28 passed.
• gfx950 RMSNorm autotune suite: 5 passed.
• gfx950 multi-GPU RMSNorm suite: 2 passed.
• Focused mixed-dtype matrix (small/generic/vec8, atomic/staged, plain/fused, unaligned FP32 weight): 13 passed.
• gfx942 compile-only: BF16/FP32 and FP16/FP32 plain/fused forward plus atomic/staged backward passed.
• gfx950 benchmark workflow: all 7 RMSNorm rows passed and survived CSV conversion with stable identities.
  • BF16/FP32: 5.393 TB/s forward, 1.387 TB/s backward, 1.393 TB/s fused backward at 4096x4096.
  • FP16/FP32: 1.661 TB/s forward, 0.167 TB/s backward, 0.170 TB/s fused backward at 512x4096.
• Python style, Ruff, shell syntax, py_compile, and git diff --check: passed.

Runtime correctness and benchmark coverage on gfx942 will run in the existing MI325 PR jobs.

Refs #749 and Dao-AILab/quack#178.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.