From 92643d68f551fe36c35a26b4a2e3bfd894d7d92a Mon Sep 17 00:00:00 2001 From: Jee Jee Li Date: Fri, 31 Jul 2026 19:12:03 +0800 Subject: [PATCH] K3 DSpark AR fusion (#50242) Signed-off-by: Jee Jee Li --- vllm/models/common/ops/__init__.py | 2 ++ .../ops/fused_allreduce_rms_norm.py} | 6 +++--- vllm/models/deepseek_v32/amd/model.py | 2 +- vllm/models/deepseek_v32/nvidia/model.py | 2 +- vllm/models/kimi_k3/nvidia/dspark_mla.py | 20 ++++++++++++++++--- 5 files changed, 24 insertions(+), 8 deletions(-) rename vllm/models/{deepseek_v32/common/fused_ops.py => common/ops/fused_allreduce_rms_norm.py} (91%) diff --git a/vllm/models/common/ops/__init__.py b/vllm/models/common/ops/__init__.py index 0d0fe8de26..f577a3c4c2 100644 --- a/vllm/models/common/ops/__init__.py +++ b/vllm/models/common/ops/__init__.py @@ -2,8 +2,10 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Ops shared across model implementations.""" +from .fused_allreduce_rms_norm import fused_allreduce_rms_norm from .fused_qk_rmsnorm import fused_q_kv_rmsnorm __all__ = [ + "fused_allreduce_rms_norm", "fused_q_kv_rmsnorm", ] diff --git a/vllm/models/deepseek_v32/common/fused_ops.py b/vllm/models/common/ops/fused_allreduce_rms_norm.py similarity index 91% rename from vllm/models/deepseek_v32/common/fused_ops.py rename to vllm/models/common/ops/fused_allreduce_rms_norm.py index 6a795e2a15..443c8f3eec 100644 --- a/vllm/models/deepseek_v32/common/fused_ops.py +++ b/vllm/models/common/ops/fused_allreduce_rms_norm.py @@ -1,9 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Fused ops for deepseek_v32 (eager / breakable-cudagraph path). +"""Fused all-reduce + residual-add + RMSNorm for eager model paths. -These recover fusions that vLLM's torch.compile passes would normally do but -that don't fire when running eager under the breakable CUDA graph. +This recovers a fusion that vLLM's torch.compile passes would normally do but +that doesn't fire for models running eager (or under a breakable CUDA graph). """ import torch diff --git a/vllm/models/deepseek_v32/amd/model.py b/vllm/models/deepseek_v32/amd/model.py index 2b710f8203..7efa52fd95 100644 --- a/vllm/models/deepseek_v32/amd/model.py +++ b/vllm/models/deepseek_v32/amd/model.py @@ -34,7 +34,7 @@ from vllm.model_executor.models.utils import ( make_empty_intermediate_tensors_factory, make_layers, ) -from vllm.models.deepseek_v32.common.fused_ops import fused_allreduce_rms_norm +from vllm.models.common.ops import fused_allreduce_rms_norm from vllm.sequence import IntermediateTensors from .rocm import DeepseekV32MLAAttention diff --git a/vllm/models/deepseek_v32/nvidia/model.py b/vllm/models/deepseek_v32/nvidia/model.py index d16f4451ad..e7a7db2b59 100644 --- a/vllm/models/deepseek_v32/nvidia/model.py +++ b/vllm/models/deepseek_v32/nvidia/model.py @@ -38,8 +38,8 @@ from vllm.model_executor.models.utils import ( make_layers, sequence_parallel_chunk, ) +from vllm.models.common.ops import fused_allreduce_rms_norm from vllm.models.deepseek_v32.attention import DeepseekV32Attention -from vllm.models.deepseek_v32.common.fused_ops import fused_allreduce_rms_norm from vllm.sequence import IntermediateTensors diff --git a/vllm/models/kimi_k3/nvidia/dspark_mla.py b/vllm/models/kimi_k3/nvidia/dspark_mla.py index 33a4012443..8643789034 100644 --- a/vllm/models/kimi_k3/nvidia/dspark_mla.py +++ b/vllm/models/kimi_k3/nvidia/dspark_mla.py @@ -20,6 +20,7 @@ from vllm.model_executor.models.utils import ( get_draft_quant_config, maybe_prefix, ) +from vllm.models.common.ops import fused_allreduce_rms_norm from vllm.models.kimi_k3.nvidia.mla import MultiHeadLatentAttention from vllm.models.kimi_k3.nvidia.model import KimiMLP from vllm.utils.torch_utils import is_quantized_kv_cache @@ -54,11 +55,15 @@ class K3DSparkDecoderLayer(nn.Module): use_rope=True, non_causal_multi_token_decode=True, ) + # Both row-parallel outputs stay un-reduced; their all-reduces are fused + # into the RMSNorm that follows via fused_allreduce_rms_norm. + self.self_attn.o_proj.reduce_results = False self.mlp = KimiMLP( hidden_size=config.hidden_size, intermediate_size=config.intermediate_size, hidden_act=config.hidden_act, quant_config=quant_config, + reduce_results=False, prefix=maybe_prefix(prefix, f"layers.{start_layer_id + layer_idx}.mlp"), ) self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) @@ -73,16 +78,23 @@ class K3DSparkDecoderLayer(nn.Module): residual: torch.Tensor | None, ) -> tuple[torch.Tensor, torch.Tensor]: if residual is None: + # First layer: hidden_states is the (already reduced) embedding. residual = hidden_states hidden_states = self.input_layernorm(hidden_states) else: - hidden_states, residual = self.input_layernorm(hidden_states, residual) + hidden_states, residual = fused_allreduce_rms_norm( + hidden_states, residual, self.input_layernorm + ) hidden_states = self.self_attn( positions=positions, hidden_states=hidden_states, ) - hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) + hidden_states, residual = fused_allreduce_rms_norm( + hidden_states, residual, self.post_attention_layernorm + ) + # The MLP output is reduced by the next layer's input_layernorm (or by + # the model's final_norm). hidden_states = self.mlp(hidden_states) return hidden_states, residual @@ -410,7 +422,9 @@ class K3DSparkModel(nn.Module): hidden_states=hidden_states, residual=residual, ) - hidden_states, _ = self.final_norm(hidden_states, residual) + hidden_states, _ = fused_allreduce_rms_norm( + hidden_states, residual, self.final_norm + ) return hidden_states