mudler / mudler/vllm.cpp

ROCm RmsNormPlusAdd (and DualRmsNormPlusRes) still weld every operand dtype to one T

Open
#2,543 0 comments 0 reactions 0 assignees View on GitHub

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 by RmsNormPlusAddRocm, 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 one T,
    dispatched by DualRmsNormPlusResRocm, 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

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.