[Perf][DSv4.1] Split fused_wqa_wkv by token across TP ranks on SM90 (#58872)

Signed-off-by: Shuolei Wang <shuoleiwang123@gmail.com>

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.<idx>...`` 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
 

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",