ROCm RmsNormPlusAdd (and DualRmsNormPlusRes) still weld every operand dtype to one T
Nobody has claimed this yet.
- Dominant language
- C++
- Stars
- 423
- Forks
- 53
- Avg merge
- 20h 26m
- Merged PRs (30d)
- 310
Description
Row: MODEL-MM-QWEN4-EXP
Spec: rmsnorm-gamma-dtype-twins.md, under ## Owed.
Found while closing #2492, which names two twins of the weld and not these. Not fixed in that
flow: every additional blind edit to a file nothing can compile is unverifiable surface, which
is the same reason #2477 left the ROCm half alone.
The weld, third and fourth instances
#2492 fixed RmsNormRowKernel in src/vt/rocm/rocm_rmsnorm.hip and
RmsNormQuantFp8RowKernel in src/vt/cuda/cuda_ops.cu. Two more kernels in that same
ROCm file still take one T for every operand:
RmsNormPlusAddKernel<T>---__global__ void RmsNormPlusAddKernel(T* out, const T* x, const T* w, const T* add, ...),
dispatched byRmsNormPlusAddRocm, which holds
VT_CHECK(w.dtype == x.dtype && out.dtype == x.dtype && addend.dtype == x.dtype, "rocm rmsnorm+add: dtype match").DualRmsNormPlusResKernel<T>--- three gammas and two activations on oneT,
dispatched byDualRmsNormPlusResRocm, which holds
VT_CHECK(out.dtype == x1.dtype && x1.dtype == x2.dtype && x1.dtype == residual.dtype, ...).
Why the first one is a real divergence and not merely a strict signature
vt::RmsNormPlusAdd (src/vt/fused_ops.cpp:19-31) dispatches to ROCm when the queue is
ROCm and otherwise runs a composed reference:
RmsNorm(q, out, x, w, args);
Add(q, out, out, addend);
vt::RmsNorm admits IsFloat(weight.dtype) and, since #2477 and #2492, every device arm
serves an independent gamma dtype. So the composed non-ROCm path ACCEPTS a mixed gamma and
the ROCm path REFUSES it: a device arm refusing what its sibling accepts, which
cuda_qwen4_exp.cu:60-62 says must not stand unrecorded.
DualRmsNormPlusResRocm is different in kind --- its non-ROCm sibling throws
"vt::DualRmsNormPlusRes: CPU compose not yet wired" rather than accepting anything --- so
there is no arm to diverge from yet. It is named here so the file's inventory is complete.
Is it live?
Not today. vt::RmsNormPlusAdd's only consumers are src/vllm/model_executor/models/gemma4.cpp:643
and :668, whose gammas come from the dense safetensors loader and are bf16 alongside a bf16
activation, so the equality holds. It becomes live the first time a mixed gamma reaches that
op on an AMD device.
What would close this
Thread a separate Tw through RmsNormPlusAddKernel and narrow RmsNormPlusAddRocm's
VT_CHECK to a supported-gamma-dtype refusal that names what it got, exactly as #2492 did
for RmsNormRowKernel in the same file. It needs an AMD device, or an explicit statement
that it lands ungated --- and, better than either, a CI lane that merely COMPILES
src/vt/rocm/, which today nothing does.
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start with src/vt/rocm/rocm_rmsnorm.hip and compare RmsNormPlusAddKernel/RmsNormPlusAddRocm with the changes from #2492. Read src/vt/fused_ops.cpp:19-31 and the Gemma4 call sites at src/vllm/model_executor/models/gemma4.cpp:643 and :668 to confirm the dispatch contract. Done means mixed gamma dtypes are handled consistently, unsupported dtypes are reported, and the ROCm path is compiled or explicitly gated.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp
- Domain
- backend
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 45/100