diff --git a/vllm/models/deepseek_v41/attention.py b/vllm/models/deepseek_v41/attention.py index 78114f0161..a6ae7cb7a5 100644 --- a/vllm/models/deepseek_v41/attention.py +++ b/vllm/models/deepseek_v41/attention.py @@ -27,7 +27,11 @@ from vllm.model_executor.layers.sparse_attn_indexer import SparseAttnIndexer from vllm.model_executor.layers.sparse_mqa_indexer import SparseMQAIndexer from vllm.model_executor.model_loader.weight_utils import make_attention_sink from vllm.models.common.ops import fused_q_kv_rmsnorm -from vllm.models.common.ops.sequence_parallel import sp_reduce_scatter +from vllm.models.common.ops.sequence_parallel import ( + sp_all_gather, + sp_reduce_scatter, + sp_shard, +) from vllm.models.deepseek_v41.common.ops import ( MXFP4_BLOCK_SIZE, fused_indexer_q_rope_quant, @@ -82,6 +86,10 @@ from vllm.v1.kv_cache_interface import ( logger = init_logger(__name__) +# From this many tokens, a layer with shard_fused_wqa_wkv set projects one token +# shard per TP rank and all-gathers the result. +_SHARD_FUSED_WQA_WKV_MIN_TOKENS = 2048 + def _replace_layer_index(prefix: str, layer_id: int) -> str: """Swap the layer index inside a ``...layers....`` prefix.""" @@ -402,6 +410,8 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): # Set by ``bind_gemm_rs`` when the decoder layer runs sequence # parallel and the fused GEMM + reduce-scatter kernel accepts wo_b. self.gemm_rs: GemmRsAr | None = None + # Set by the NVIDIA decoder layer from _shard_fused_wqa_wkv. + self.shard_fused_wqa_wkv = False # Initialize rotary embedding before the indexer/compressor consume it. self.rotary_emb = build_deepseek_v4_rope( @@ -921,7 +931,11 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): def _fused_wqa_wkv_gemm(self, hidden_states: torch.Tensor) -> torch.Tensor: # Override point: the ROCm layer preshuffles this weight in place, so # it cannot go through fused_wqa_wkv directly. + num_tokens = hidden_states.shape[0] # MergedColumnParallelLinear returns (output, bias); bias is None. + if self.shard_fused_wqa_wkv and num_tokens >= _SHARD_FUSED_WQA_WKV_MIN_TOKENS: + qr_kv, _ = self.fused_wqa_wkv(sp_shard(hidden_states)) + return sp_all_gather(qr_kv)[:num_tokens] qr_kv, _ = self.fused_wqa_wkv(hidden_states) return qr_kv diff --git a/vllm/models/deepseek_v41/nvidia/model.py b/vllm/models/deepseek_v41/nvidia/model.py index 3757809dbc..3946058de0 100644 --- a/vllm/models/deepseek_v41/nvidia/model.py +++ b/vllm/models/deepseek_v41/nvidia/model.py @@ -16,6 +16,7 @@ from vllm.distributed import ( get_engram_dp_size, get_pp_group, get_tensor_model_parallel_world_size, + get_tp_group, tensor_model_parallel_all_reduce, ) from vllm.forward_context import ( @@ -255,6 +256,21 @@ def _use_sequence_parallel(vllm_config: VllmConfig) -> bool: ) +def _shard_fused_wqa_wkv(vllm_config: VllmConfig) -> bool: + """Whether large batches shard fused_wqa_wkv by token and all-gather the output.""" + ca_comm = getattr(get_tp_group().device_communicator, "ca_comm", None) + return ( + current_platform.is_device_capability(90) + and not _use_sequence_parallel(vllm_config) + and vllm_config.lora_config is None + and not vllm_config.parallel_config.use_ubatching + and not envs.VLLM_BATCH_INVARIANT + and ca_comm is not None + and not ca_comm.disabled + and ca_comm.fully_connected + ) + + def maybe_init_gemm_rs(vllm_config: VllmConfig, use_sequence_parallel: bool) -> bool: """Set up the fused ``wo_b`` GEMM + reduce-scatter when opted in. @@ -329,6 +345,7 @@ class DeepseekV4DecoderLayer(nn.Module): # Binds only when wo_b's kernel and shape qualify; otherwise # forward keeps the separate reduce-scatter below. self.attn.bind_gemm_rs() + self.attn.shard_fused_wqa_wkv = _shard_fused_wqa_wkv(vllm_config) self.ffn = DeepseekV4MoE( vllm_config, prefix=f"{prefix}.ffn",