[Bugfix] Fix DeepSeek V4 mHC broadcast buffer for weight sync (#52626)
Signed-off-by: Hollow Man <hollowman@opensuse.org>
ℍ𝕠𝕝𝕝𝕠𝕨 𝕄𝕒𝕟 committed
Aug 18, 2026 at 05:55 UTC
d5f5de7a7db4e13ea9f94ce8662b55b6b5e97d02
2 files changed
+66
-1
tests/kernels/test_mhc_kernels.py
+61
@@ -1,7 +1,10 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from types import SimpleNamespace
4
+
5
import pytest
6
import torch
7
+import torch.nn as nn
8
9
import vllm.model_executor.kernels.mhc # noqa: F401
10
from vllm.model_executor.kernels.mhc.tilelang import (
@@ -9,6 +12,10 @@ from vllm.model_executor.kernels.mhc.tilelang import (
12
_torch_hc_prenorm_gemm,
13
)
14
from vllm.model_executor.layers.mhc import HAS_TILELANG_MHC
15
+from vllm.models.deepseek_v4.nvidia.model import (
16
+ DeepseekV4DecoderLayer,
17
+ DeepseekV4Model,
18
+)
19
from vllm.platforms import current_platform
20
from vllm.utils.torch_utils import set_random_seed
21
@@ -355,3 +362,57 @@ def test_hc_head_tilelang(num_tokens, hidden_size, hc_mult):
362
363
out_ref = hc_head_ref(residual, fn, hc_scale, hc_base, rms_eps, hc_eps)
364
torch.testing.assert_close(out, out_ref, atol=5e-2, rtol=1e-2)
365
+
366
+
367
+def _make_mhc_decoder_layer(hc_mult: int, hidden_size: int) -> DeepseekV4DecoderLayer:
368
+ layer = DeepseekV4DecoderLayer.__new__(DeepseekV4DecoderLayer)
369
+ nn.Module.__init__(layer)
370
+ layer.hc_mult = hc_mult
371
+ layer.hidden_size = hidden_size
372
+ mix_hc = (2 + hc_mult) * hc_mult
373
+ layer.hc_attn_fn = nn.Parameter(
374
+ torch.randn(mix_hc, hc_mult * hidden_size, dtype=torch.float32),
375
+ requires_grad=False,
376
+ )
377
+ layer.hc_attn_fn_broadcast = None
378
+ return layer
379
+
380
+
381
+def _patch_first_rank_pp_group(monkeypatch):
382
+ monkeypatch.setattr(
383
+ "vllm.models.deepseek_v4.nvidia.model.get_pp_group",
384
+ lambda: SimpleNamespace(is_first_rank=True),
385
+ )
386
+
387
+
388
+def test_deepseek_v4_mhc_broadcast_finalize_sums_hc_streams(monkeypatch):
389
+ """First finalize (at the end of load_weights) allocates
390
+ hc_attn_fn_broadcast as hc_attn_fn summed over hc streams."""
391
+ _patch_first_rank_pp_group(monkeypatch)
392
+ layer = _make_mhc_decoder_layer(hc_mult=2, hidden_size=8)
393
+ model = SimpleNamespace(start_layer=0, end_layer=1, layers=[layer])
394
+
395
+ DeepseekV4Model.finalize_mhc_broadcast_weights(model)
396
+
397
+ assert layer.hc_attn_fn_broadcast is not None
398
+ expected = layer.hc_attn_fn.detach().view(-1, 2, 8).sum(dim=1)
399
+ assert torch.equal(layer.hc_attn_fn_broadcast, expected)
400
+
401
+
402
+def test_deepseek_v4_mhc_broadcast_refit_refreshes_in_place(monkeypatch):
403
+ """Re-finalizing after a weight refit must copy into the existing
404
+ broadcast tensor so its address stays stable for captured CUDA graphs,
405
+ while picking up the new hc_attn_fn values."""
406
+ _patch_first_rank_pp_group(monkeypatch)
407
+ layer = _make_mhc_decoder_layer(hc_mult=2, hidden_size=8)
408
+ model = SimpleNamespace(start_layer=0, end_layer=1, layers=[layer])
409
+
410
+ DeepseekV4Model.finalize_mhc_broadcast_weights(model)
411
+ buffer = layer.hc_attn_fn_broadcast
412
+
413
+ layer.hc_attn_fn.add_(1.0)
414
+ DeepseekV4Model.finalize_mhc_broadcast_weights(model)
415
+
416
+ assert layer.hc_attn_fn_broadcast is buffer
417
+ expected = layer.hc_attn_fn.detach().view(-1, 2, 8).sum(dim=1)
418
+ assert torch.equal(layer.hc_attn_fn_broadcast, expected)
vllm/models/deepseek_v4/nvidia/model.py
+5
-1
@@ -1372,11 +1372,15 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
1372
return
1373
layer = self.layers[self.start_layer]
1374
if isinstance(layer, DeepseekV4DecoderLayer):
1375
- layer.hc_attn_fn_broadcast = (
1375
+ broadcast = (
1376
layer.hc_attn_fn.detach()
1377
.view(-1, layer.hc_mult, layer.hidden_size)
1378
.sum(dim=1)
1379
)
1380
+ if layer.hc_attn_fn_broadcast is None:
1381
+ layer.hc_attn_fn_broadcast = broadcast
1382
+ else:
1383
+ layer.hc_attn_fn_broadcast.copy_(broadcast)
1384
1385
1386
def _make_deepseek_v4_weights_mapper(expert_dtype: str) -> WeightsMapper: