[xiangyuT]

Scope and current state

This PR is anchored to the exact Muse-Glimmer v0.29 source baseline:

• Base branch: baseline/muse-glimmer-v029-20260921
• Base commit: 98dff2a81d747d1dba01a47f939f48c3526d4206
• Current head: 48f43667b0bd5e7bec3ed87c185987e44c915d9b
• Companion kernels baseline/dev branches: 6d92b1bfbf32767ecda8e819613eb151e70030ad
  (0.1.14.1 tag); no kernels source change is included.

The earlier two model-local MuseGlimmer RMSNorm commits were intentionally
reverted. The current implementation is a standard-layer integration and
generic XPU dispatch change; no worker monkeypatch is included.

Upstream source conclusion

Current upstream main still defines MuseGlimmerRMSNorm as a plain
torch.nn.Module in vllm/model_executor/models/muse_glimmer.py. Its forward
performs FP32 cast, reduction, rsqrt and weight multiplication directly. The
Muse decoder calls these modules directly, including QK norm.

The generic vLLM RMSNorm and GemmaRMSNorm classes do have CustomOp/XPU
dispatch paths, but MuseGlimmer does not instantiate those classes. Therefore:

• changing only vllm-xpu-kernels cannot intercept the current Muse Python norm
  chain;
• changing only generic layernorm.py dispatch does not affect MuseGlimmer;
• embedding a model-specific XPU fast path in models/muse_glimmer.py is the
  wrong boundary; a minimal conversion to standard norm-layer declarations is
  a separate modeling cleanup required to expose the generic dispatch.

References:

• Muse model-local norm definition: https://github.com/vllm-project/vllm/blob/e9f169d16b9408bb9ae44f75072b91a5521d733c/vllm/model_executor/models/muse_glimmer.py#L503-L533
• Muse QK and sandwich-norm construction: https://github.com/vllm-project/vllm/blob/e9f169d16b9408bb9ae44f75072b91a5521d733c/vllm/model_executor/models/muse_glimmer.py#L1073-L1079
• Muse attention/decoder norm calls: https://github.com/vllm-project/vllm/blob/e9f169d16b9408bb9ae44f75072b91a5521d733c/vllm/model_executor/models/muse_glimmer.py#L1131-L1212
• Generic RMSNorm, including has_weight=False: https://github.com/vllm-project/vllm/blob/e9f169d16b9408bb9ae44f75072b91a5521d733c/vllm/model_executor/layers/layernorm.py#L30-L116
• Upstream XPU GemmaRMSNorm implementation: https://github.com/vllm-project/vllm/blob/a0d3e5c16e96b62ec5c7f741952cfddc99534faf/vllm/model_executor/layers/layernorm.py#L130-L194
• Upstream XPU Gemma op wrappers and registration: https://github.com/vllm-project/vllm/blob/a0d3e5c16e96b62ec5c7f741952cfddc99534faf/vllm/_xpu_ops.py#L116-L151
• Upstream XPU Gemma op registration: https://github.com/vllm-project/vllm/blob/a0d3e5c16e96b62ec5c7f741952cfddc99534faf/vllm/_xpu_ops.py#L1255-L1273
• Upstream Gemma modeling import and construction: https://github.com/vllm-project/vllm/blob/a0d3e5c16e96b62ec5c7f741952cfddc99534faf/vllm/model_executor/models/gemma.py#L31-L39
• Upstream Gemma norm construction: https://github.com/vllm-project/vllm/blob/a0d3e5c16e96b62ec5c7f741952cfddc99534faf/vllm/model_executor/models/gemma.py#L235-L238
• Upstream Qwen3.5 modeling import and construction: https://github.com/vllm-project/vllm/blob/e9f169d16b9408bb9ae44f75072b91a5521d733c/vllm/model_executor/models/qwen3_5.py#L37-L38
• Upstream Qwen3.5 norm construction: https://github.com/vllm-project/vllm/blob/e9f169d16b9408bb9ae44f75072b91a5521d733c/vllm/model_executor/models/qwen3_5.py#L170-L175
• Exact XPU kernel op schemas: https://github.com/xiangyuT/vllm-xpu-kernels/blob/6d92b1bfbf32767ecda8e819613eb151e70030ad/csrc/torch_bindings.cpp#L22-L49
• Exact XPU Gemma kernel implementation: https://github.com/xiangyuT/vllm-xpu-kernels/blob/6d92b1bfbf32767ecda8e819613eb151e70030ad/csrc/layernorm.cpp#L845-L884

Implemented design

The worker-level monkeypatch draft is not the recommended implementation. A
more idiomatic path is to make MuseGlimmer declare the standard vLLM norm
layers, exactly as Gemma and Qwen3.5 do:

• use GemmaRMSNorm for the four sandwich norms (weight + 1 semantics);
• use RMSNorm(..., has_weight=False) for QK and embedding norms;
• use ordinary RMSNorm for the final scale norm.

Those are model semantics, not model-specific fast-path code. They preserve the
checkpoint parameter names and allow all backends to use the same fallback
behavior. The implementation backports upstream's XPU
GemmaRMSNorm.forward_xpu and its _xpu_ops.py wrappers into the v0.29 branch.
The exact kernels commit 6d92b1bfbf32767ecda8e819613eb151e70030ad already
provides gemma_rms_norm and fused_add_gemma_rms_norm; no kernels source
change is needed for the ordinary norm optimization. The existing generic XPU
RMSNorm path already accepts an optional weight for weightless norms.

This design removes the model-local Python reduction from the optimized path,
keeps backend dispatch in layernorm.py, and remains compatible with native
and non-XPU execution. A strict prohibition on any muse_glimmer.py wiring
change would force an implicit worker transform and is not aligned with the
other vLLM model implementations.

Current validation

Commit 48f43667b0bd5e7bec3ed87c185987e44c915d9b contains the implementation
and XPU dispatch tests. Static compilation and diff checks passed.

On vllm/vllm-openai-xpu:v0.29.0, using devices 2/3 and the exact Muse model:

• direct gemma_rms_norm probe: max error 0;
• fused Gemma residual probe: max error 0.0009765625, aliases preserved;
• weightless RMSNorm probe: max error 0;
• V1 runner, TP2 FP8, FP16, 32768/512, two warmups, seeds 42/43/44:
  36.609, 36.791, 36.701 ms/token, mean 36.700;
• all three generated-text hashes match the retained upstream baseline;
• serial 5-shot GSM8K: 10/10, invalid 0.

The default V2 runner measured 39.029 ms/token; it is retained separately
because the historical 34.156 ms/token candidate used V1. A unitrace target
request completed, but the profiler output lacked decode-step/raw kernel events,
so no structural trace claim is made.

Residual-fusion boundary

The current MuseGlimmer sandwich order is norm(attention_output) + residual,
while vLLM's existing fused_add_rms_norm is
norm(attention_output + residual). Enabling the existing op would change the
model math. The downstream norm_add_norm path is a separate primitive that
fuses norm(x) * w1 + residual followed by a second norm. If pursued, it needs
a generic XPU op/layer with a native fallback and guarded shape/dtype dispatch;
it must not be represented as fused_add_rms_norm. It should be reviewed and
validated separately from the standard norm conversion.

Historical evidence (not a current PR claim)

The previously tested model-local candidate, which is no longer in the effective
PR diff, achieved:

• upstream TPOT: 60.758 ms/token
• standalone RMSNorm candidate: 45.258 ms/token
• standalone + weightless QK norm candidate: 34.156 ms/token
• trace calls: 3843 -> 1347
• norm-related calls: 2825 -> 329

Both historical candidate stages passed the available E2E output checks. The
final historical candidate also passed the serial 5-shot GSM8K smoke:
10/10, invalid 0, with all ten output text hashes matching upstream.
These numbers are retained as diagnostic evidence only; they must be reproduced
after implementation is moved to the accepted XPU-only boundary.

Remaining validation and limitations

1. Run the permanent XPU dispatch tests in the supported CI/device test
   environment, including non-XPU, BF16, prefill, non-contiguous and fallback
   paths.
2. Rework or replace the current unitrace setup so a valid dual-rank decode
   kernel trace is captured; the current target request is retained but lacks
   decode-step/raw kernel events.
3. Residual norm-add-norm fusion remains intentionally out of scope for this
   implementation and requires a separate generic primitive.

No formal acceptance, long-duration stability, or full accuracy claim is made by
this PR. The detailed local evidence is retained under
codex_output/muse-glimmer-decode-trace-20260921/.