@tduncan / Nvidia-Optimization-Stack-VLLM / commits / eb61707851

[Bugfix] Fix stale HPC QK-norm weights after weight refit (#56706)

Signed-off-by: aoshen02 <aoshen02@users.noreply.github.com>

aoshen02 committed Sep 15, 2026 at 12:34 UTC eb61707851f35281aa82c3ba7714f1ce9c1860e0

6 files changed +143 -98

tests/model_executor/model_loader/test_reload.py

+107

@@ -1,7 +1,10 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 import gc

4 +import importlib.machinery

5 import inspect

6 +import sys

7 +import types

8 from unittest.mock import Mock

9 from weakref import WeakKeyDictionary, ref

@@ -817,6 +820,110 @@ def test_layerwise_reload_updates_loaded_non_persistent_buffers(monkeypatch):

820 assert "0.scale" not in model.state_dict()

821

822

823 +@pytest.fixture

824 +def hpc_rope_norm(monkeypatch, default_vllm_config):

825 + """Import HpcRopeNorm with the external ``hpc`` package stubbed out."""

826 + if "hpc" not in sys.modules:

827 + stub = types.ModuleType("hpc")

828 + stub.__spec__ = importlib.machinery.ModuleSpec("hpc", loader=None)

829 + stub.QuantType = types.SimpleNamespace( # type: ignore[attr-defined]

830 + QPERTOKEN_PERHEAD_KPERTENSOR_VPERTENSOR=types.SimpleNamespace(value=0)

831 + )

832 + monkeypatch.setitem(sys.modules, "hpc", stub)

833 + from vllm.model_executor.layers.hpc import rope_norm

834 +

835 + monkeypatch.setattr(rope_norm, "_hpc_rope_norm_instances", {})

836 + return rope_norm

837 +

838 +

839 +def test_hpc_rope_norm_kernel_sees_refit_norm_weights(monkeypatch, hpc_rope_norm):

840 + """The fused HPC kernel is handed the live QK-norm weights after a refit.

841 +

842 + Drives the production ``_forward_impl`` with a recording ``hpc`` stub. The

843 + Q/K norm weights it receives must be the model's own float32 parameters,

844 + so a layerwise reload that rewrites them in place (same storage) is what

845 + the kernel sees. Previously the kernel read separate mirrors that no

846 + reload path refreshed.

847 + """

848 + from vllm.model_executor.layers.layernorm import RMSNorm

849 +

850 + head_dim, num_heads, num_kv_heads, block_size = 128, 8, 1, 4

851 + layer = torch.nn.Module()

852 + layer.q_norm = RMSNorm(head_dim, 1e-6, dtype=torch.float32)

853 + layer.k_norm = RMSNorm(head_dim, 1e-6, dtype=torch.float32)

854 + layer.hpc_rope_norm = hpc_rope_norm.HpcRopeNorm(

855 + num_heads=num_heads,

856 + num_kv_heads=num_kv_heads,

857 + head_dim=head_dim,

858 + cos_sin_cache=torch.ones(16, head_dim),

859 + use_qk_norm=True,

860 + fallback_qnorm=layer.q_norm,

861 + fallback_knorm=layer.k_norm,

862 + kv_cache_dtype="auto",

863 + layer_name="hpc_test_layer",

864 + )

865 + model = torch.nn.Sequential(layer)

866 + rnorm = layer.hpc_rope_norm

867 +

868 + calls: list[dict] = []

869 +

870 + def record(*args, **kwargs):

871 + calls.append(kwargs)

872 +

873 + monkeypatch.setattr(sys.modules["hpc"], "rope_norm_store_kv", record, raising=False)

874 +

875 + def kernel_norm_weights():

876 + q_size, kv_size = num_heads * head_dim, num_kv_heads * head_dim

877 + qkv = torch.zeros(1, q_size + 2 * kv_size, dtype=torch.bfloat16)

878 + kv_cache = torch.zeros(

879 + 2, num_kv_heads, block_size, 2 * head_dim, dtype=torch.bfloat16

880 + )

881 + attn_layer = types.SimpleNamespace(

882 + _k_scale=torch.ones(1), _v_scale=torch.ones(1)

883 + )

884 + attn_metadata = types.SimpleNamespace(

885 + num_actual_tokens=1,

886 + num_decodes=1,

887 + num_decode_tokens=1,

888 + num_prefills=0,

889 + num_prefill_tokens=0,

890 + max_query_len=1,

891 + decode_query_len=1,

892 + qo_indptr=None,

893 + qo_indptr_decode=None,

894 + slot_mapping=torch.tensor([4]),

895 + seq_lens=torch.tensor([1]),

896 + block_table_tensor=torch.tensor([[1]]),

897 + hpc_kv_written=False,

898 + )

899 + output = torch.zeros(1, q_size, dtype=torch.bfloat16)

900 + rnorm._forward_impl(qkv, kv_cache, attn_metadata, attn_layer, output)

901 + return calls[-1]["q_norm_weight"], calls[-1]["k_norm_weight"]

902 +

903 + def loaded(value):

904 + return torch.full((head_dim,), value, dtype=torch.bfloat16)

905 +

906 + default_weight_loader(layer.q_norm.weight, loaded(0.5))

907 + default_weight_loader(layer.k_norm.weight, loaded(0.25))

908 + q, k = kernel_norm_weights()

909 + assert torch.equal(q, loaded(0.5).float())

910 + assert torch.equal(k, loaded(0.25).float())

911 + q_ptr, k_ptr = q.data_ptr(), k.data_ptr()

912 +

913 + record_metadata_for_reloading(model)

914 + initialize_layerwise_reload(model)

915 + layer.q_norm.weight.weight_loader(layer.q_norm.weight, loaded(2.0))

916 + layer.k_norm.weight.weight_loader(layer.k_norm.weight, loaded(3.0))

917 + finalize_layerwise_reload(model, model_config=None)

918 +

919 + q, k = kernel_norm_weights()

920 + assert q.dtype == k.dtype == torch.float32

921 + assert torch.equal(q, loaded(2.0).float())

922 + assert torch.equal(k, loaded(3.0).float())

923 + assert (q.data_ptr(), k.data_ptr()) == (q_ptr, k_ptr)

924 + assert not hasattr(rnorm, "qnorm_weight")

925 +

926 +

927 @pytest.mark.parametrize(

928 "tp_size", [pytest.param(1), pytest.param(2, marks=[pytest.mark.slow_test])]

929 )

vllm/model_executor/layers/hpc/hpc_ihc.py

+13 -38

@@ -20,9 +20,8 @@ Constraints:

20 - x / residual must be bfloat16; hc_fn weights stay float32 (the checkpoint

21 keeps them out of fp8 quantization via modules_to_not_convert)

23 -The fused modules own no parameters: they hold references to the weights of the

24 -eager layer they replace, bound in ``process_weights_after_loading`` (invoked

25 -generically for every ``HpcModule`` by the model loader).

23 +The fused modules own no parameters: their forwards read the weights of the

24 +eager layer they replace directly via the captured owner reference.

26 NOTE: The reference implementation also offers a cross-layer post+pre fusion

27 (``HpcIHCPostPre``) that folds one segment's post into the next segment's pre.

@@ -147,41 +146,26 @@ class HpcIHCPre(HpcModule):

146 object.__setattr__(self, "_fallback_op", fallback_op)

147 # Same anti-cycle idiom for the fused RMSNorm owner (a sibling module).

148 object.__setattr__(self, "_norm_owner", norm_owner)

150 - self.weight = None

151 - self.hc_scale = None

152 - self.hc_base = None

153 - self.rms_weight = None

154 - self.rms_eps = 0.0

149

150 @classmethod

151 def support(cls, hc_mult: int, hidden_size: int) -> bool:

152 return _ihc_supported(hc_mult, hidden_size)

153

160 - def process_weights_after_loading(self, model=None):

161 - # hc_fn is a ReplicatedLinear whose weight is [2 * hc_mult, hc_mult * d].

162 - # It stays float32 -- the HPC kernel consumes fp32 weights directly, so

163 - # there is no conversion (and no precision loss) here.

164 - owner = self._fallback_op

165 - self.weight = owner.hc_fn.weight

166 - self.hc_scale = owner.hc_scale

167 - self.hc_base = owner.hc_base

168 - if self._norm_owner is not None:

169 - self.rms_weight = self._norm_owner.weight

170 - self.rms_eps = self._norm_owner.variance_epsilon

171 -

154 def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:

155 import hpc

156

157 + owner = self._fallback_op

158 + norm_owner = self._norm_owner

159 return hpc.fuse_ihc_pre(

160 x,

177 - self.weight,

178 - self.hc_scale,

179 - self.hc_base,

161 + owner.hc_fn.weight,

162 + owner.hc_scale,

163 + owner.hc_base,

164 self.norm_eps,

165 self.hc_eps,

166 self.magnitude,

183 - self.rms_weight,

184 - self.rms_eps,

167 + norm_owner.weight if norm_owner is not None else None,

168 + norm_owner.variance_epsilon if norm_owner is not None else 0.0,

169 )

170

171

@@ -233,29 +217,20 @@ class HpcIHCHead(HpcModule):

217 self.norm_eps = norm_eps

218 # See HpcIHCPre: keep the owner out of the module tree.

219 object.__setattr__(self, "_fallback_op", fallback_op)

236 - self.weight = None

237 - self.hc_scale = None

238 - self.hc_base = None

220

221 @classmethod

222 def support(cls, hc_mult: int, hidden_size: int) -> bool:

223 return _ihc_supported(hc_mult, hidden_size)

224

244 - def process_weights_after_loading(self, model=None):

245 - # hc_head_fn weight is [hc_mult, hc_mult * d], float32.

246 - owner = self._fallback_op

247 - self.weight = owner.hc_head_fn.weight

248 - self.hc_scale = owner.hc_head_scale

249 - self.hc_base = owner.hc_head_base

250 -

225 def forward(self, x: torch.Tensor) -> torch.Tensor:

226 import hpc

227

228 + owner = self._fallback_op

229 return hpc.fuse_ihc_head(

230 x,

256 - self.weight,

257 - self.hc_scale,

258 - self.hc_base,

231 + owner.hc_head_fn.weight,

232 + owner.hc_head_scale,

233 + owner.hc_head_base,

234 self.norm_eps,

235 self.hc_eps,

236 )

vllm/model_executor/layers/hpc/hpc_module.py

-3

@@ -11,8 +11,5 @@ class HpcModule(nn.Module):

11 def support(cls, *args, **kwargs):

12 return True

14 - def process_weights_after_loading(self, model):

15 - pass

16 -

14 def forward(self, *args, **kwargs):

15 pass

vllm/model_executor/layers/hpc/rope_norm.py

+18 -45

@@ -90,13 +90,13 @@ class HpcRopeNorm(CustomOp, HpcModule):

90 """HPC fused RoPE + QK-Norm + KV-Cache-Write (+ optional FP8 Q quant).

92 Registered as a sub-module in model layers (e.g. HunYuanAttention).

93 - Norm weights are extracted from fallback norm modules via

94 - process_weights_after_loading() after all weights are loaded.

93 + The QK-Norm weights are read directly from the fallback norm modules

94 + (built in float32 when HPC is active), eliminating the derived copy.

96 forward() is dispatched by CustomOp framework:

97 - In compiled mode: forward_cuda() calls torch.ops.vllm.hpc_rope_norm_forward

98 as a splitting point — internal Python control flow is opaque

99 - to torch.compile and not captured by CUDA Graph.

99 + to torch.compile (its kernel launches can still be CUDA-graph captured).

100 - In eager/native mode: forward_native() falls back to forward_cuda().

101 """

102

@@ -124,8 +124,6 @@ class HpcRopeNorm(CustomOp, HpcModule):

124 self.num_kv_heads = num_kv_heads

125 self.head_dim = head_dim

126

127 - self.use_qk_norm = use_qk_norm

128 -

127 self.q_size = num_heads * head_dim

128 self.kv_size = num_kv_heads * head_dim

129

@@ -139,31 +137,18 @@ class HpcRopeNorm(CustomOp, HpcModule):

137

138 self.head_per_group = num_heads // num_kv_heads

139

142 - # Pre-allocate norm weight tensors as Parameters so they are tracked by

143 - # CuMemAllocator (for sleep/wake_up) and have stable addresses for CUDA

144 - # Graph replay. process_weights_after_loading() updates them inplace via

145 - # copy_() so refit does not invalidate captured graph tensor pointers.

146 - # Shape is [head_dim] to match the HPC kernel's q/k_norm_weight layout.

147 - if use_qk_norm and fallback_qnorm is not None:

148 - self.qnorm_weight: torch.nn.Parameter | None = torch.nn.Parameter(

149 - torch.empty(head_dim, dtype=torch.float32),

150 - requires_grad=False,

151 - )

152 - else:

153 - self.qnorm_weight = None

154 - if use_qk_norm and fallback_knorm is not None:

155 - self.knorm_weight: torch.nn.Parameter | None = torch.nn.Parameter(

156 - torch.empty(head_dim, dtype=torch.float32),

157 - requires_grad=False,

158 - )

159 - else:

160 - self.knorm_weight = None

161 -

140 self.use_fp8 = "fp8" in kv_cache_dtype

141 # The RMSNorm/RoPE ordering is model dependent (e.g. HunYuan V3 applies

142 # QK-Norm before RoPE -> NORM_THEN_ROPE), so it is supplied by the

143 # caller. When QK-Norm is disabled the policy is forced to NONE.

144 self.qk_norm_policy = qk_norm_policy if use_qk_norm else QkNormPolicy.NONE

145 + if self.qk_norm_policy != QkNormPolicy.NONE and (

146 + fallback_qnorm is None or fallback_knorm is None

147 + ):

148 + raise ValueError(

149 + "QK-Norm is enabled but fallback_qnorm/fallback_knorm were not "

150 + "supplied; the fused kernel reads their weights directly."

151 + )

152

153 # Register layer_name + add self to the global instance registry so the

154 # module-level custom op (hpc_rope_norm_forward) can route back here.

@@ -215,20 +200,6 @@ class HpcRopeNorm(CustomOp, HpcModule):

200 logger.info_once("enable hpc rope_norm")

201 return True

202

218 - def process_weights_after_loading(self, model: torch.nn.Module = None) -> None:

219 - """Copy norm weights (float32) from fallback norm modules inplace.

220 -

221 - Uses copy_() to preserve tensor addresses for CUDA Graph / refit

222 - compatibility. Called by the model's load_weights() after all weights

223 - are loaded (and generically from the model loader for DummyModelLoader

224 - / sleep-wake_up reload paths).

225 - """

226 - if self.use_qk_norm:

227 - if self.fallback_qnorm is not None and self.qnorm_weight is not None:

228 - self.qnorm_weight.data.copy_(self.fallback_qnorm.weight.data.float())

229 - if self.fallback_knorm is not None and self.knorm_weight is not None:

230 - self.knorm_weight.data.copy_(self.fallback_knorm.weight.data.float())

231 -

203 def register_layer_name(self, layer_name: str) -> None:

204 """Register layer_name and add self to the global registry.

205

@@ -430,12 +401,14 @@ class HpcRopeNorm(CustomOp, HpcModule):

401 k_scale = attn_layer._k_scale.reshape(1)

402 v_scale = attn_layer._v_scale.reshape(1)

403

433 - q_norm_weight = (

434 - self.qnorm_weight if self.qk_norm_policy != QkNormPolicy.NONE else None

435 - )

436 - k_norm_weight = (

437 - self.knorm_weight if self.qk_norm_policy != QkNormPolicy.NONE else None

438 - )

404 + if self.qk_norm_policy != QkNormPolicy.NONE:

405 + assert self.fallback_qnorm is not None

406 + assert self.fallback_knorm is not None

407 + q_norm_weight = self.fallback_qnorm.weight

408 + k_norm_weight = self.fallback_knorm.weight

409 + else:

410 + q_norm_weight = None

411 + k_norm_weight = None

412

413 # --- Prefill ---

414 if num_prefill_reqs > 0:

vllm/model_executor/model_loader/utils.py

-10

@@ -15,7 +15,6 @@ import vllm.envs as envs

15 from vllm.config import ModelConfig, VllmConfig, set_current_vllm_config

16 from vllm.logger import init_logger

17 from vllm.model_executor.layers.attention import is_deferred_attention_layer

18 -from vllm.model_executor.layers.hpc import HpcModule

18 from vllm.model_executor.layers.quantization.base_config import (

19 QuantizationConfig,

20 QuantizeMethodBase,

@@ -154,15 +153,6 @@ def process_weights_after_loading(

153 with device_loading_context(module, target_device):

154 module.process_weights_after_loading(model_config.dtype)

155

157 - # Process HPC modules (HpcRopeNorm, etc.) that rely on

158 - # process_weights_after_loading being called from the model's

159 - # load_weights(). When using DummyModelLoader (e.g. profiling or

160 - # sleep/wake_up reload), the model's load_weights() is not called, so we

161 - # must handle HPC modules here generically.

162 - for _, module in model.named_modules():

163 - if isinstance(module, HpcModule):

164 - module.process_weights_after_loading(model)

165 -

156 # Model-level post-load hook, after the per-layer quant finalize.

157 if hasattr(model, "process_weights_after_loading"):

158 model.process_weights_after_loading()

vllm/model_executor/models/hy_v3.py

+5 -2

@@ -300,8 +300,11 @@ class HYV3Attention(nn.Module):

300 prefix=f"{prefix}.attn",

301 )

302 if self.use_qk_norm:

303 - self.q_norm = RMSNorm(self.head_dim, rms_norm_eps)

304 - self.k_norm = RMSNorm(self.head_dim, rms_norm_eps)

303 + # The HPC fused kernel reads norm weights directly in float32;

304 + # the fallback path keeps the model default dtype.

305 + norm_dtype = torch.float32 if rope_support else None

306 + self.q_norm = RMSNorm(self.head_dim, rms_norm_eps, dtype=norm_dtype)

307 + self.k_norm = RMSNorm(self.head_dim, rms_norm_eps, dtype=norm_dtype)

308

309 # HPC fused RoPE + QK-Norm + KV-Cache-Write (+ optional FP8 Q quant).

310 # HunYuan V3 applies QK-Norm *before* RoPE, so NORM_THEN_ROPE.