[Bugfix] Fix stale HPC QK-norm weights after weight refit

[aoshen02]

Purpose

HpcRopeNorm kept float32 mirror Parameters (qnorm_weight / knorm_weight) of the
q_norm / k_norm weights. They were filled by process_weights_after_loading(), and that
hook was dispatched only by the generic HpcModule loop in
model_executor/model_loader/utils.py, which runs at cold start.

No reload path reaches that loop:

path	dispatcher	reaches the HPC loop?
base_loader.load_model	utils.process_weights_after_loading	yes (cold start only)
weight_cache/ipc_loader	same	yes (cold start only)
GPUModelRunner.reload_weights	initialize/finalize_layerwise_reload	no
NCCL / IPC / sharded-RDT weight transfer	same	no
wake_up after sleep	CuMemAllocator restore	no

layerwise.py only dispatches process_weights_after_loading for layers carrying a
QuantizeMethodBase and for deferred attention layers. HpcRopeNorm is neither, so after an
RL weight refit the fused kernel consumed pre-refit QK-norm weights — silently, with no
error and no warning.

Fix

Build q_norm / k_norm in float32 when the HPC path is active, and have the kernel read
them directly. No derived copy means nothing to refresh, on any load path. The checkpoint is
untouched (still bf16) — default_weight_loader casts on param.data.copy_().

Cleanup bundled in (not a fix)

HpcIHCPre / HpcIHCHead bound owner weight aliases in the same hook. Those aliases do
survive refits (verified empirically, see Test Result), so their forwards now read the owner
attributes directly purely to remove the last reason for the hook to exist. With no
HpcModule needing it, the three overrides, the base no-op and the loader loop are deleted
(4 process_weights_after_loading mounting points → 3).

The two failure modes are different and it matters:

	what was cached	after a refit
HpcIHCPre / HpcIHCHead	an alias to the owner's Parameter	pointer unchanged, alias still valid
HpcRopeNorm	an independent copy_().float() duplicate	nobody refreshes it → stale values

An alias follows its owner because layerwise reload's _copy_and_restore_kernel_tensors
copies in place and _place_kernel_tensors re-registers the original Parameter objects.
A duplicate does not.

Test Plan

# lint (all hooks, changed files)
pre-commit run --files $(git diff --name-only upstream/main...HEAD)

# unit
pytest tests/model_executor/model_loader/test_reload.py -v -m "not slow_test"

New test test_hpc_rope_norm_kernel_sees_refit_norm_weights drives the production
_forward_impl with a recording hpc stub and asserts the kernel receives the refit Q and
K weights, float32, same storage.

End-to-end on GB200 (sm100) with Tencent hpc-ops built from source, HY4
(HYV4ForCausalLM, enable_ihc: true, hc_mult: 4, hidden_size: 6144), three arms —
base, base2 (same code as base, to establish the run-to-run noise floor), pr — plus a
deliberately perturbed arm as a positive control.

Test Result

Unit — tests/model_executor/model_loader/test_reload.py -m "not slow_test", run inside
vllm/vllm-openai:nightly (aarch64, GB200):

	result
without this PR	27 passed, 2 failed
with this PR	28 passed, 1 failed

test_hpc_rope_norm_kernel_sees_refit_norm_weights FAILS without the change and PASSES with
it — it is the new test, and it does catch the bug.

The one remaining failure, test_kv_scale_reload, fails identically with and without the
change; it is pre-existing in that image and unrelated.

Applying only this PR's utils.py deletion (the HpcModule loop + its import) on top of an
otherwise untouched tree gives byte-identical results to not applying it, confirming the
removed mounting point is inert for every non-HPC path.

Lint — pre-commit run --files <changed>: all hooks pass. (mypy initially flagged
rope_norm.py:398-399 Item "None" of "Any | None" has no attribute "weight" — the same latent
crash the removed fallback_qnorm is not None guard used to prevent. Fixed by validating the
invariant in __init__ with an explicit ValueError and narrowing at the use site.)

End-to-end, HY4 on GB200 / sm100 — the HPC gate fires: HpcIHCPre: 8, HpcIHCPost: 8,
HpcIHCHead: 1, gated_mla_layers: 4, identical across arms, so the changed code really is
executed.

process_weights_after_loading mounting points, read off the live model:

base: ['HpcIHCHead', 'HpcIHCPre', 'MLAAttention', 'UnquantizedFusedMoEMethod']
pr:   [                           'MLAAttention', 'UnquantizedFusedMoEMethod']

Exactly the two HPC hooks are removed; the other mounting points are untouched.

Output equivalence — greedy, fixed prompts, token ids compared one by one:

	vs base
pr (cudagraph)	8/8 identical
pr (eager)	8/8 identical
base2 (same code, control)	8/8 identical
positive control (hc_fn.weight * 1.01 on the changed line)	6/8 diverge, first at token 56

The positive control is what makes the 8/8 meaningful: the same test does detect a
perturbation on the very line this PR touches.

Throughput, batch 256 × 512 output tokens, 7 reps, median:

mode	pr vs base	same-code noise floor (base2)
cudagraph	+0.16%	+0.15%
eager (run 1)	+1.21%	+0.05%
eager (run 2)	+2.91%	+2.47%

In eager run 2 the same code (base2) moved +2.47%, so the apparent gain tracks arm
ordering (warm-up), not the change. No measurable difference in either mode, which matches the
mechanism — a handful of extra attribute lookups per forward.

Refit-alias validity, reload_weights on the base arm (the claim that the iHC cleanup is
not a fix):

modules holding a cached alias: 9  (8x HpcIHCPre + 1x HpcIHCHead)
alias still points at the live weight after reload: 9 / 9
sample pointer before -> after: 26843826176 -> 26843826176   (unchanged)
generation identical after reload: 6 / 6

Confirms the alias half is genuinely cleanup.

GSM8K on the full 78-layer model — the 4-layer slice I first used has no
discriminating power (every arm, including a deliberately sabotaged one, scores 0.000 with
invalid_rate=0.985), so I pulled the full tencent/Hy4-preview (131 shards, 1.5 TB bf16) and
re-ran on 4 nodes / 16 GPUs, TP=4 x PP=4 with expert parallelism, via
distributed_executor_backend="external_launcher":

	base	this PR
fused_modules	HpcIHCPre: 156, HpcIHCPost: 156, HpcIHCHead: 1	identical
gated_mla_layers	78	identical
greedy output, 8 prompts	—	8/8 token-identical
GSM8K accuracy (500 q, 5-shot)	0.5660	0.5660
GSM8K invalid_rate	0.0120	0.0120

156 = 78 layers x 2 iHC boundaries per layer, so the changed path is exercised on every layer.
invalid_rate drops from 0.985 on the slice to 0.012 here, i.e. the model is actually
answering and the score is meaningful. Accuracy is identical between the arms.

(0.566 is the raw-completion score without a chat template; it is used only for the A/B
comparison, not as a claim about the model's ceiling.)

Not verified

sm90 + HY3 + HPC_ATTN refit. HpcAttentionBackend.supports_compute_capability is
capability == DeviceCapability(9, 0), so the actual fix (HpcRopeNorm, hy_v3.py) cannot be
exercised on sm100 at all — my cluster is GB200 only. What I verified end-to-end is the iHC
cleanup half plus the mounting-point removal. I'd appreciate the Tencent hpc-ops authors (or
anyone with H100/H800) running an sm90 refit-consistency check; the script must fail on
main, otherwise it isn't exercising the bug.

Reproducing the hpc-ops build

hpc-ops ships no wheel (no GitHub release, nothing on PyPI), and building it inside
vllm/vllm-openai:nightly needs three environment fixes unrelated to this PR:

1. PyPI is throttled on some clusters — prefetch build deps and install with --no-index.
2. The image has no git, and hpc-ops/setup.py:get_version() shells out to it for a short hash.
3. CUDA headers come from the nvidia-*-cu13 wheels, so cusparse.h is not under
   /usr/local/cuda/include; torch's ATen/cuda/CUDAContextLight.h fails to include it.
   export CPATH=<site-packages>/nvidia/cu13/include.

Known limitation

Sharded-state artifacts saved by an older build contain the retired
hpc_rope_norm.qnorm_weight / knorm_weight keys; ShardedStateLoader indexes state_dict
by checkpoint key and raises KeyError. Regenerate the artifact.

Worth noting this swaps the direction of an existing incompatibility rather than introducing
one: on main, loading a sharded artifact saved without HPC into an HPC-enabled build
already fails, because those two keys stay in state_dict and hit
raise ValueError(f"Missing keys {tuple(state_dict)} in loaded state!"). I chose not to touch
the shared loader from an HPC PR — happy to add a narrow tolerance here or as a prerequisite PR
if maintainers prefer.

Follow-ups deliberately left out of this PR

• HpcRopeNorm.__init__ could also reject a non-float32 fallback_qnorm.weight. The fp32
  contract now lives with the caller (hy_v3.py), and hy_v3 is the only in-tree caller, but
  nothing checks it for the next model that wires up HPC_ATTN.
• After this PR nothing does isinstance(x, HpcModule) anymore, and the base class is reduced
  to support() returning True plus a pass forward() — both overridden by every
  subclass. It is arguably dead and could be dropped.

---

Not a duplicate: no open PR addresses HPC QK-norm staleness across weight reload
(searched vllm-project/vllm for HpcRopeNorm, hpc qk norm refit,
process_weights_after_loading refit; the only nearby open PR, #47443, changes HPC attention
decode split-K selection). RFC #54477 is the umbrella for selective weight reload in RL.

AI assistance was used for this change. Every changed line was reviewed by me, and I ran the
lint, unit and end-to-end runs reported above.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[aoshen02]

@thisjiang Hi could you take a look?

[aoshen02]

/ci run

[github-actions]

✅ Triggered Buildkite CI #88695 for commit 0d4c4b9f267f.

[thisjiang]

@thisjiang Hi could you take a look?

Thanks! I'll test 'sm90 + HY3 + HPC_ATTN refit' on Hopper GPU and share the results once ready.

[aoshen02]

@thisjiang Hi could you take a look?

Thanks! I'll test 'sm90 + HY3 + HPC_ATTN refit' on Hopper GPU and share the results once ready.

Thank you!

[aoshen02]

/ci run

[github-actions]

✅ Triggered Buildkite CI #88800 for commit 6c60d6fb9e0f.

[thisjiang]

Hi~I verified 'sm90 + HY3 + HPC_ATTN refit' on Hopper GPU and confirmed the results are correct. Thanks for your fix!