PR #41431: fix(llama): use weightless RMSNorm for FlashNorm-folded checkpoints (has_weight=False)

• Repository: vllm-project/vllm
• Author: fm1320
• State: open | merged: False
• Link: https://github.com/vllm-project/vllm/pull/41431

Description (problem / solution / changelog)

Summary

Adds the same one-line has_weight=False plumbing pattern as PR #40117 (fix(gemma4): use weightless k_norm for KV-shared layers) to the Llama loader so that FlashNorm-folded HuggingFace checkpoints load cleanly. Three call sites in vllm/model_executor/models/llama.py get one extra keyword argument each, gated by a flashnorm_folded flag on the model config.

This is purely a loader-side fix. It unblocks loading the weightless-rmsnorm collection on HuggingFace (Llama, Qwen, SmolLM, etc.; this PR covers the Llama path) but does not deliver any speedup by itself, because forward_cuda does not yet branch on has_weight=False. That kernel-side gap is tracked in a companion issue.

Problem

Loading open-machine/Llama-3.2-1B-FlashNorm (or any other FlashNorm-folded Llama checkpoint with "flashnorm_folded": true in config.json) in stock vLLM fails at engine init:

from vllm import LLM
LLM(model='open-machine/Llama-3.2-1B-FlashNorm', dtype='bfloat16', enforce_eager=True)
# RuntimeError: Engine core initialization failed.

Root cause: vllm/model_executor/models/llama.py always calls RMSNorm(...) with the default has_weight=True, so params_dict expects every per-layer *_layernorm.weight and model.norm.weight tensor. FlashNorm-folded checkpoints either (a) keep those weights as all-ones for HF compatibility or (b) omit them; neither layout loads cleanly today.

Background: FlashNorm and the precedent

FlashNorm (Graef et al., 2024) folds the per-channel RMSNorm weight into the next linear, leaving an RMSNorm whose weight is mathematically all ones. A runtime that recognizes the all-ones case can skip the per-channel multiply and recover an end-to-end speedup.

vLLM's RMSNorm class at vllm/model_executor/layers/layernorm.py already supports has_weight=False, added in #40117 for Gemma-4 KV-shared k_norm. This PR extends the same pattern to the Llama path.

Proposed change

Three RMSNorm call sites in vllm/model_executor/models/llama.py get one keyword argument each, gated by a flashnorm_folded flag read from the model's config:

flashnorm_folded = getattr(config, "flashnorm_folded", False)

# In LlamaDecoderLayer.__init__:
self.input_layernorm = RMSNorm(
    config.hidden_size, eps=config.rms_norm_eps,
    has_weight=not flashnorm_folded,
)
self.post_attention_layernorm = RMSNorm(
    config.hidden_size, eps=config.rms_norm_eps,
    has_weight=not flashnorm_folded,
)

# In LlamaModel.__init__:
self.norm = RMSNorm(
    config.hidden_size, eps=config.rms_norm_eps,
    has_weight=not flashnorm_folded,
)

The flashnorm_folded flag is read from the HF config via getattr(...) with a False default, so existing non-flashified Llama checkpoints are unaffected. No schema changes anywhere.

What this PR does and does not do

• Does: makes the FlashNorm-folded Llama checkpoints on the HF Hub loadable in vLLM. Engine init succeeds, the model runs, and outputs are numerically equivalent to the unfolded baseline (multiply by ones is a no-op).
• Does not: deliver any speedup. The CUDA dispatch through vllm/kernels/vllm_c.py:rms_norm substitutes a ones tensor when weight is None (the C++ kernel _custom_ops.rms_norm requires a weight), and forward_cuda's residual path calls _custom_ops.fused_add_rms_norm which similarly requires weight. An adjacent TODO(luka) at vllm/model_executor/layers/layernorm.py:240 acknowledges the broader weight=None passing inconsistency.

The two pieces are sequenced: this PR is the loader-side compatibility fix; the kernel-side performance fix is described in the companion issue (linked below) and would unlock the actual speedup.

Empirical verification

We tested the runtime equivalent of this source-level patch (a monkey-patch on RMSNorm.__init__ plus a small loader-side filter for the all-ones tensors that the HF-compat checkpoint still carries) on Colab A100, vllm 0.9.2, bf16, 256-token greedy decode, 20 trials, enforce_eager=True, tensor_parallel_size=1:

The model loads cleanly, runs, and produces correct outputs. The -0.04% delta confirms that the kernel does the same work in both cases (within measurement noise) and isolates the remaining speedup to the kernel-side fix.

For comparison, the equivalent path in HuggingFace Transformers, where torch's F.rms_norm(weight=None) already implements the weight-null branch, gives +12.77% end-to-end on the same model and hardware. That is the gain the kernel-side companion would close.

Companion issue

The kernel-side fix that actually delivers the speedup is described in: https://github.com/vllm-project/vllm/issues/41430

Together, this PR plus the kernel issue make FlashNorm a first-class option in vLLM deployments.

Other architectures

The same one-line pattern applies to mistral.py, qwen2.py, qwen3.py, phi3.py, and any other RMSNorm-based architecture whose checkpoint can be flashified. Happy to follow up with separate small PRs for those if there is interest; this PR is scoped to Llama for review focus.

References

• FlashNorm paper: https://arxiv.org/abs/2407.09577
• vLLM PR #40117 (precedent): https://github.com/vllm-project/vllm/pull/40117
• vLLM Issue #39370 (related RFC, orthogonal): https://github.com/vllm-project/vllm/issues/39370
• llama.cpp Issue #22486 (the equivalent ask in llama.cpp): https://github.com/ggml-org/llama.cpp/issues/22486
• HuggingFace Transformers reproducibility for the +12.77% number: https://github.com/OpenMachine-ai/transformer-tricks/blob/main/notebooks/flashNorm_hf_a100.ipynb

Changed files

• vllm/model_executor/layers/layernorm.py (modified, +7/-2)
• vllm/model_executor/models/llama.py (modified, +27/-3)

PR #14: FlashNorm: realized end-to-end speedup in HuggingFace Transformers (+12.77% on Llama-3.2-1B / A100 / bf16)

• Repository: OpenMachine-ai/transformer-tricks
• Author: fm1320
• State: closed | merged: True
• Link: https://github.com/OpenMachine-ai/transformer-tricks/pull/14

Description (problem / solution / changelog)

Summary

Adds a new measurement to the FlashNorm paper: the realized end-to-end speedup of FlashNorm Proposition 1 in stock HuggingFace Transformers, via PyTorch's existing F.rms_norm(weight=None) weightless C++ code path applied as a runtime monkey-patch on LlamaRMSNorm.forward.

Headline: +12.77% on open-machine/Llama-3.2-1B-FlashNorm at bf16, NVIDIA A100 (40 GB), 256-token greedy decode, median of 20 trials with 3 warmup. Same model both variants; the gain isolates kernel choice from folding.

This composes with the decode-profile appendix from #13 (now merged): that one establishes the upper bound at 24.6%; this PR reports how much is realized today via kernel choice alone (roughly half).

Companion upstream contributions (filed; details in the "Other-runtime integrations" section below):

• vLLM loader PR: vllm-project/vllm#41431 (fix(llama): use weightless RMSNorm for FlashNorm-folded checkpoints (has_weight=False)).
• vLLM kernel issue: vllm-project/vllm#41430 (fused_add_rms_norm does not branch on has_weight=False; needed for the GPU-path speedup).
• llama.cpp tracking issue: ggml-org/llama.cpp#22486 (weightless RMSNorm for FlashNorm; runtime kernel exists, loader needs TENSOR_NOT_REQUIRED).

What's in this PR

No new bib entries required.

Bundled corrections to app:rmsnorm-fraction (merged in #13)

Three small follow-ups derived from existing reported values (no new measurements):

1. tab:rmsnorm-fraction 4-bit row: norm fraction 15.01% -> 9.31% (the arithmetically consistent value: 1650 ms / 17720 ms = 0.0931). The norm time and total decode are kept; they are the values the invariance claim of 1640 / 1675 / 1650 ms relies on.
2. Findings paragraph: 24.62% -> 8.19% -> 15.01% -> 24.62% -> 8.19% -> 9.31%. The "non-monotonic across precisions" narrative still holds; the rebound at 4-bit is smaller than the original table indicated.
3. Table caption: A100, fp16 inputs -> A100, bf16 inputs. The notebook loads the model at torch.bfloat16 and the row is labelled bf16.

Reproducibility

notebooks/flashNorm_hf_a100.ipynb runs end to end on a Colab A100 in ~3 minutes:

1. Pip-installs the latest transformers and accelerate (same install pattern as flashNorm_decode_profile.ipynb).
2. Loads open-machine/Llama-3.2-1B-FlashNorm at bf16 with eager attention.
3. Benchmarks variant A (HF default LlamaRMSNorm.forward, multiplies by the all-ones weight) and variant B (runtime monkey-patch to F.rms_norm(x, shape, None, eps), skips the multiply).
4. Saves a JSON results file and downloads it in Colab.

No vLLM, no llama.cpp, no version-pinning conflicts with the rest of the Colab kernel.

Other-runtime integrations (work in progress, not in this PR)

We separately tested the same FlashNorm-folded checkpoint in vLLM and llama.cpp; both are integration-blocked today, with concrete upstream paths to land them. Out of scope for this PR; mentioned here so reviewers see the broader picture.

• vLLM: applied the runtime equivalent of the proposed source-level loader fix (has_weight=False plumbing modeled on PR vllm-project/vllm#40117, which added the same pattern for Gemma-4 KV-shared k_norm). With those patches, the FlashNorm-folded checkpoint loads cleanly and runs in vLLM with correct outputs. The speedup itself is not yet realized: the CUDA RMSNorm kernels require a weight tensor (the vllm_c.py impl substitutes ones for None, and _custom_ops.fused_add_rms_norm similarly takes a required weight), so the multiply still happens on the GPU path; an adjacent TODO(luka) at vllm/model_executor/layers/layernorm.py:240 acknowledges this. Measured delta: -0.04% (within noise). Two upstream contributions filed: loader PR at vllm-project/vllm#41431, kernel issue at vllm-project/vllm#41430.
• llama.cpp: the converter convert_hf_to_gguf.py correctly recognizes flashnorm_folded: true in the source HF config.json and drops the per-layer norm tensors when writing the GGUF. The runtime then fails to load the resulting GGUF with error loading model: missing tensor 'blk.0.attn_norm.weight'. The runtime model-load path requires the per-layer norm tensors to be present and there is no "treat absent norm as identity" branch. Tracked in upstream issue ggml-org/llama.cpp#22486; we will file a comment there with the empirical reproduction recipe and propose a runtime-side fix.

These integrations will land as separate upstream contributions in those projects' repos. This PR is deliberately scoped to the HF result that is realizable today.

Status

Draft. Open for review.

Changed files

• notebooks/README.md (modified, +1/-0)
• notebooks/flashNorm_hf_a100.ipynb (added, +222/-0)
• tex/flashNorm.tex (modified, +3/-3)