K3 DSpark AR fusion (#50242)
Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai>
Jee Jee Li committed
Jul 31, 2026 at 19:12 UTC
92643d68f551fe36c35a26b4a2e3bfd894d7d92a
5 files changed
+24
-8
vllm/models/common/ops/__init__.py
+2
@@ -2,8 +2,10 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
"""Ops shared across model implementations."""
4
5
+from .fused_allreduce_rms_norm import fused_allreduce_rms_norm
6
from .fused_qk_rmsnorm import fused_q_kv_rmsnorm
7
8
__all__ = [
9
+ "fused_allreduce_rms_norm",
10
"fused_q_kv_rmsnorm",
11
]
vllm/models/common/ops/fused_allreduce_rms_norm.py
renamed
+3
-3
@@ -1,9 +1,9 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""Fused ops for deepseek_v32 (eager / breakable-cudagraph path).
3
+"""Fused all-reduce + residual-add + RMSNorm for eager model paths.
4
5
-These recover fusions that vLLM's torch.compile passes would normally do but
6
-that don't fire when running eager under the breakable CUDA graph.
5
+This recovers a fusion that vLLM's torch.compile passes would normally do but
6
+that doesn't fire for models running eager (or under a breakable CUDA graph).
7
"""
8
9
import torch
vllm/models/deepseek_v32/amd/model.py
+1
-1
@@ -34,7 +34,7 @@ from vllm.model_executor.models.utils import (
34
make_empty_intermediate_tensors_factory,
35
make_layers,
36
)
37
-from vllm.models.deepseek_v32.common.fused_ops import fused_allreduce_rms_norm
37
+from vllm.models.common.ops import fused_allreduce_rms_norm
38
from vllm.sequence import IntermediateTensors
39
40
from .rocm import DeepseekV32MLAAttention
vllm/models/deepseek_v32/nvidia/model.py
+1
-1
@@ -38,8 +38,8 @@ from vllm.model_executor.models.utils import (
38
make_layers,
39
sequence_parallel_chunk,
40
)
41
+from vllm.models.common.ops import fused_allreduce_rms_norm
42
from vllm.models.deepseek_v32.attention import DeepseekV32Attention
42
-from vllm.models.deepseek_v32.common.fused_ops import fused_allreduce_rms_norm
43
from vllm.sequence import IntermediateTensors
44
45
vllm/models/kimi_k3/nvidia/dspark_mla.py
+17
-3
@@ -20,6 +20,7 @@ from vllm.model_executor.models.utils import (
20
get_draft_quant_config,
21
maybe_prefix,
22
)
23
+from vllm.models.common.ops import fused_allreduce_rms_norm
24
from vllm.models.kimi_k3.nvidia.mla import MultiHeadLatentAttention
25
from vllm.models.kimi_k3.nvidia.model import KimiMLP
26
from vllm.utils.torch_utils import is_quantized_kv_cache
@@ -54,11 +55,15 @@ class K3DSparkDecoderLayer(nn.Module):
55
use_rope=True,
56
non_causal_multi_token_decode=True,
57
)
58
+ # Both row-parallel outputs stay un-reduced; their all-reduces are fused
59
+ # into the RMSNorm that follows via fused_allreduce_rms_norm.
60
+ self.self_attn.o_proj.reduce_results = False
61
self.mlp = KimiMLP(
62
hidden_size=config.hidden_size,
63
intermediate_size=config.intermediate_size,
64
hidden_act=config.hidden_act,
65
quant_config=quant_config,
66
+ reduce_results=False,
67
prefix=maybe_prefix(prefix, f"layers.{start_layer_id + layer_idx}.mlp"),
68
)
69
self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
@@ -73,16 +78,23 @@ class K3DSparkDecoderLayer(nn.Module):
78
residual: torch.Tensor | None,
79
) -> tuple[torch.Tensor, torch.Tensor]:
80
if residual is None:
81
+ # First layer: hidden_states is the (already reduced) embedding.
82
residual = hidden_states
83
hidden_states = self.input_layernorm(hidden_states)
84
else:
79
- hidden_states, residual = self.input_layernorm(hidden_states, residual)
85
+ hidden_states, residual = fused_allreduce_rms_norm(
86
+ hidden_states, residual, self.input_layernorm
87
+ )
88
89
hidden_states = self.self_attn(
90
positions=positions,
91
hidden_states=hidden_states,
92
)
85
- hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
93
+ hidden_states, residual = fused_allreduce_rms_norm(
94
+ hidden_states, residual, self.post_attention_layernorm
95
+ )
96
+ # The MLP output is reduced by the next layer's input_layernorm (or by
97
+ # the model's final_norm).
98
hidden_states = self.mlp(hidden_states)
99
return hidden_states, residual
100
@@ -410,7 +422,9 @@ class K3DSparkModel(nn.Module):
422
hidden_states=hidden_states,
423
residual=residual,
424
)
413
- hidden_states, _ = self.final_norm(hidden_states, residual)
425
+ hidden_states, _ = fused_allreduce_rms_norm(
426
+ hidden_states, residual, self.final_norm
427
+ )
428
return hidden_states
429
430