Skip compile SP pass under enforce_eager (V2 runner)

sp_threshold force-enabled compilation_config.pass_config.enable_sp, which the
V2 model runner rejects ("does not yet support: sequence parallelism"). The
eager in-forward SP path reads sp_threshold directly and needs no compile pass;
under enforce_eager there is no compilation to attach it to. Guard the mapping.

Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>

vllm/config/parallel.py

@@ -207,6 +207,24 @@ class ParallelConfig:
 
     enable_dbo: bool = False
     """Enable dual batch overlap for the model executor."""
+
+    sp_threshold: int | None = None
+    """Sequence-parallelism engage threshold, in tokens. A single umbrella knob
+    covering both mechanisms:
+
+    - the compile-based `SequenceParallelismPass` (drives
+      `compilation_config.pass_config.enable_sp` / `sp_min_token_num`), used by
+      torch.compile models, and
+    - the eager, in-forward SP applied by models that run without
+      torch.compile (e.g. DeepSeek V4's per-token mHC residual stream).
+
+    Values:
+    - `None` (default): sequence parallelism is not used. For the compile pass
+      this defers to the optimization-level default.
+    - `N > 0`: use SP, engaging only on forwards with at least `N` tokens (so
+      smaller forwards such as decode run plain TP).
+
+    Only takes effect when `tensor_parallel_size > 1`."""
     ubatch_size: int = Field(default=0, ge=0)
     """Number of ubatch size."""
 
@@ -630,6 +648,14 @@ class ParallelConfig:
         assert last_exc is not None
         raise last_exc
 
+    @property
+    def enable_sp(self) -> bool:
+        return self.tensor_parallel_size > 1 and self.sp_threshold is not None
+
+    @property
+    def sequence_parallel_size(self) -> int:
+        return self.tensor_parallel_size if self.enable_sp else 1
+
     # The all_reduce at the end of attention (during o_proj) means that
     # inputs are replicated across each rank of the tensor parallel group.
     # If using expert-parallelism with DeepEP All2All ops, replicated

vllm/config/vllm.py

@@ -1201,6 +1201,19 @@ class VllmConfig:
 
         # async tp is built on top of sequence parallelism and requires it.
         pass_config = self.compilation_config.pass_config
+        # Map the umbrella `parallel_config.sp_threshold` onto the compile-based
+        # SP pass (the eager in-forward SP path reads sp_threshold directly):
+        #   None  -> defer to the optimization-level default
+        #   N > 0 -> force the pass on with N as its min-token threshold
+        sp_threshold = self.parallel_config.sp_threshold
+        # Only enable the compile-based SP pass when compilation actually runs.
+        # The eager in-forward SP path reads sp_threshold directly and runs the
+        # V2 model runner, which rejects the compile SP pass; under enforce_eager
+        # there is no compilation to attach it to anyway.
+        if sp_threshold is not None and not self.model_config.enforce_eager:
+            pass_config.enable_sp = True
+            if pass_config.sp_min_token_num is None:
+                pass_config.sp_min_token_num = sp_threshold
         if pass_config.fuse_gemm_comms:
             pass_config.enable_sp = True
         if pass_config.enable_sp:

vllm/engine/arg_utils.py

@@ -490,6 +490,7 @@ class EngineArgs:
     all2all_backend: All2AllBackend = ParallelConfig.all2all_backend
     enable_elastic_ep: bool = ParallelConfig.enable_elastic_ep
     enable_dbo: bool = ParallelConfig.enable_dbo
+    sp_threshold: int | None = ParallelConfig.sp_threshold
     ubatch_size: int = ParallelConfig.ubatch_size
     dbo_decode_token_threshold: int = ParallelConfig.dbo_decode_token_threshold
     dbo_prefill_token_threshold: int = ParallelConfig.dbo_prefill_token_threshold
@@ -1097,6 +1098,10 @@ class EngineArgs:
             "--all2all-backend", **parallel_kwargs["all2all_backend"]
         )
         parallel_group.add_argument("--enable-dbo", **parallel_kwargs["enable_dbo"])
+        parallel_group.add_argument(
+            "--sp-threshold",
+            **parallel_kwargs["sp_threshold"],
+        )
         parallel_group.add_argument(
             "--ubatch-size",
             **parallel_kwargs["ubatch_size"],
@@ -2092,6 +2097,7 @@ class EngineArgs:
             all2all_backend=self.all2all_backend,
             enable_elastic_ep=self.enable_elastic_ep,
             enable_dbo=self.enable_dbo,
+            sp_threshold=self.sp_threshold,
             ubatch_size=self.ubatch_size,
             dbo_decode_token_threshold=self.dbo_decode_token_threshold,
             dbo_prefill_token_threshold=self.dbo_prefill_token_threshold,

vllm/models/deepseek_v4/nvidia/model.py

@@ -15,6 +15,9 @@ from vllm.distributed import (
     get_pp_group,
     get_tensor_model_parallel_rank,
     get_tensor_model_parallel_world_size,
+    tensor_model_parallel_all_gather,
+    tensor_model_parallel_all_reduce,
+    tensor_model_parallel_reduce_scatter,
 )
 from vllm.distributed.eplb.eplb_state import EplbLayerState
 from vllm.forward_context import get_forward_context, is_forward_context_available
@@ -57,6 +60,7 @@ from vllm.model_executor.models.utils import (
     is_pp_missing_parameter,
     make_layers,
     maybe_prefix,
+    sequence_parallel_chunk,
 )
 from vllm.model_executor.utils import set_weight_attrs
 from vllm.models.deepseek_v4.attention import DeepseekV4Attention
@@ -524,6 +528,12 @@ class DeepseekV4MoE(nn.Module):
                 "Enable it with --enable-expert-parallel, or pick a different "
                 "moe backend."
             )
+        self.enable_sp = vllm_config.parallel_config.enable_sp
+        if self.enable_sp and self.use_mega_moe:
+            raise NotImplementedError(
+                "mHC sequence parallelism requires the FusedMoE path; "
+                "MegaMoE + sequence parallelism is not supported."
+            )
 
         self.routed_scaling_factor = getattr(config, "routed_scaling_factor", 1.0)
         self.hidden_size = config.hidden_size
@@ -682,6 +692,7 @@ class DeepseekV4MoE(nn.Module):
             router_logits_dtype=torch.float32,
             enable_eplb=parallel_config.enable_eplb,
             num_redundant_experts=eplb_config.num_redundant_experts,
+            is_sequence_parallel=self.enable_sp,
         )
 
     def forward(
@@ -797,6 +808,7 @@ class DeepseekV4DecoderLayer(nn.Module):
 
         config = vllm_config.model_config.hf_config
         self.hidden_size = config.hidden_size
+        self.enable_sp = vllm_config.parallel_config.enable_sp
 
         self.rms_norm_eps = config.rms_norm_eps
         self.attn = _select_dsv4_attn_cls(vllm_config)(
@@ -805,6 +817,8 @@ class DeepseekV4DecoderLayer(nn.Module):
             topk_indices_buffer=topk_indices_buffer,
             aux_stream_list=aux_stream_list,
         )
+        if self.enable_sp:
+            self.attn.wo_b.reduce_results = False
         self.ffn = DeepseekV4MoE(vllm_config, prefix=f"{prefix}.ffn")
 
         self.attn_norm = RMSNorm(self.hidden_size, self.rms_norm_eps)
@@ -866,6 +880,7 @@ class DeepseekV4DecoderLayer(nn.Module):
         post_mix: torch.Tensor | None = None,
         res_mix: torch.Tensor | None = None,
         residual: torch.Tensor | None = None,
+        is_sp_sharded: bool = False,
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
         attn_norm_weight = self.attn_norm.weight.data
         attn_norm_eps = self.attn_norm.variance_epsilon
@@ -904,9 +919,18 @@ class DeepseekV4DecoderLayer(nn.Module):
                 norm_weight=attn_norm_weight,
                 norm_eps=attn_norm_eps,
             )
+        if is_sp_sharded:
+            x = tensor_model_parallel_all_gather(x, dim=0)
 
-        # attn_norm is fused into mhc_pre_tilelang / mhc_fused_post_pre above.
         x = self.attn(positions, x, None)
+        if self.enable_sp:
+            # When SP is enabled, attn o-proj returns an un-reduced per-rank partial.
+            if is_sp_sharded:
+                # `num_tokens` meets the sp_threshold, so we do reduce-scatter.
+                x = tensor_model_parallel_reduce_scatter(x, dim=0)
+            else:
+                # `num_tokens` does not meet the sp_threshold, so we do all-reduce.
+                x = tensor_model_parallel_all_reduce(x)
 
         ffn_norm_weight = self.ffn_norm.weight.data
         ffn_norm_eps = self.ffn_norm.variance_epsilon
@@ -928,8 +952,18 @@ class DeepseekV4DecoderLayer(nn.Module):
             norm_weight=ffn_norm_weight,
             norm_eps=ffn_norm_eps,
         )
+        if is_sp_sharded:
+            x = tensor_model_parallel_all_gather(x, dim=0)
 
         x = self.ffn(x, input_ids)
+        if self.enable_sp:
+            # When SP is enabled, MLP/MoE returns an un-reduced per-rank partial.
+            if is_sp_sharded:
+                # `num_tokens` meets the sp_threshold, so we do reduce-scatter.
+                x = tensor_model_parallel_reduce_scatter(x, dim=0)
+            else:
+                # `num_tokens` does not meet the sp_threshold, so we do all-reduce.
+                x = tensor_model_parallel_all_reduce(x)
         return x, residual, post_mix, res_mix
 
 
@@ -1033,6 +1067,19 @@ class DeepseekV4Model(nn.Module):
     def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
         return self.embed_tokens(input_ids)
 
+    def _is_sp_sharded(self, num_tokens: int) -> bool:
+        # Engage SP once the (tp-padded) token count clears the threshold. The
+        # runner pads eligible batches to a multiple of tp_size, so the chunk is
+        # exact. enable_sp guarantees sp_threshold is not None. Never shard
+        # fewer tokens than tp ranks (can't give each rank a token) -> plain TP.
+        sp_threshold = self.parallel_config.sp_threshold
+        return (
+            self.parallel_config.enable_sp
+            and sp_threshold is not None
+            and num_tokens >= sp_threshold
+            and num_tokens >= self.parallel_config.tensor_parallel_size
+        )
+
     def make_empty_intermediate_tensors(
         self,
         batch_size: int,
@@ -1043,10 +1090,15 @@ class DeepseekV4Model(nn.Module):
         # of shape (num_tokens, hc_mult, hidden_size) — V4 expands the
         # token embedding to hc_mult streams before the first decoder
         # layer and keeps that shape until hc_head() collapses it.
+        # Under mHC sequence parallelism the stream is token-sharded across the
+        # TP group, so the per-rank PP buffer holds num_tokens / tp_size rows.
+        num_tokens = batch_size
+        if self._is_sp_sharded(batch_size):
+            num_tokens = cdiv(batch_size, get_tensor_model_parallel_world_size())
         return IntermediateTensors(
             {
                 "hidden_states": torch.zeros(
-                    (batch_size, self.hc_mult, self.config.hidden_size),
+                    (num_tokens, self.hc_mult, self.config.hidden_size),
                     dtype=dtype,
                     device=device,
                 ),
@@ -1060,15 +1112,32 @@ class DeepseekV4Model(nn.Module):
         intermediate_tensors: IntermediateTensors | None,
         inputs_embeds: torch.Tensor | None = None,
     ) -> torch.Tensor | IntermediateTensors:
+        # Whether to token-shard the mHC stream this forward (eager SP). A pure
+        # function of the full (padded) token count + threshold, so it is stable
+        # between cudagraph capture and replay for a given bucket. (The runner
+        # pads eligible batches to a multiple of tp_size, so an engaged forward's
+        # count is tp-aligned and the chunk is exact.)
         if get_pp_group().is_first_rank:
             if inputs_embeds is not None:
                 hidden_states = inputs_embeds
             else:
                 hidden_states = self.embed_input_ids(input_ids)
             hidden_states = hidden_states.unsqueeze(-2).repeat(1, self.hc_mult, 1)
+            # Full token count is hidden_states.shape[0] here (before the chunk).
+            is_sp_sharded = self._is_sp_sharded(hidden_states.shape[0])
+            # Shard the mHC residual stream along the token dim; it stays
+            # sharded through every layer. Non-first PP ranks already receive a
+            # sharded stream from the previous rank.
+            if is_sp_sharded:
+                hidden_states = sequence_parallel_chunk(hidden_states)
         else:
             assert intermediate_tensors is not None
             hidden_states = intermediate_tensors["hidden_states"]
+            # Match the producing rank from the full (padded) token count on the
+            # forward context (rank-independent).
+            batch_descriptor = get_forward_context().batch_descriptor
+            assert batch_descriptor is not None
+            is_sp_sharded = self._is_sp_sharded(batch_descriptor.num_tokens)
 
         if self.use_mega_moe:
             input_ids = input_ids.to(torch.int64)
@@ -1082,6 +1151,7 @@ class DeepseekV4Model(nn.Module):
                 post_mix,
                 res_mix,
                 residual,
+                is_sp_sharded,
             )
         if layer is not None:
             hidden_states = mhc_post_tilelang(
@@ -1091,6 +1161,12 @@ class DeepseekV4Model(nn.Module):
         if not get_pp_group().is_last_rank:
             return IntermediateTensors({"hidden_states": hidden_states})
 
+        # Reassemble the full token stream before the MTP buffer copy and the
+        # hc_head collapse (both need all tokens on this rank). Only when this
+        # forward actually sharded; otherwise the stream is already full.
+        if is_sp_sharded:
+            hidden_states = tensor_model_parallel_all_gather(hidden_states, dim=0)
+
         # Stash pre-hc_head residual for the MTP draft (captured copy_).
         num_tokens = hidden_states.shape[0]
         self._mtp_hidden_buffer[:num_tokens].copy_(hidden_states.flatten(1))

vllm/models/deepseek_v4/nvidia/mtp.py

@@ -22,6 +22,7 @@ from vllm.config import VllmConfig
 from vllm.distributed import (
     get_tensor_model_parallel_rank,
     get_tensor_model_parallel_world_size,
+    tensor_model_parallel_all_gather,
 )
 from vllm.logger import init_logger
 from vllm.model_executor.kernels.mhc.tilelang import (
@@ -40,7 +41,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 from vllm.model_executor.models.deepseek_mtp import SharedHead
 from vllm.model_executor.models.deepseek_v2 import get_spec_layer_idx_from_weight_name
-from vllm.model_executor.models.utils import maybe_prefix
+from vllm.model_executor.models.utils import maybe_prefix, sequence_parallel_chunk
 from vllm.models.deepseek_v4.common.ops import (
     fused_mtp_input_rmsnorm,
     mtp_shared_head_rmsnorm,
@@ -80,6 +81,7 @@ class DeepSeekV4MultiTokenPredictorLayer(nn.Module):
         self.config = config
         quant_config = vllm_config.quant_config
         self.rms_norm_eps = config.rms_norm_eps
+        self.parallel_config = vllm_config.parallel_config
 
         self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
         self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
@@ -122,6 +124,8 @@ class DeepSeekV4MultiTokenPredictorLayer(nn.Module):
         self.shared_head = SharedHead(
             config=config, prefix=prefix, quant_config=quant_config
         )
+        # The draft decoder block participates in mHC SP like the target's
+        # layers (sharded-in / sharded-out when the forward engages SP).
         self.mtp_block = DeepseekV4DecoderLayer(
             vllm_config,
             prefix,
@@ -157,10 +161,26 @@ class DeepSeekV4MultiTokenPredictorLayer(nn.Module):
         hidden_states = self.h_proj(previous_hidden_states) + self.e_proj(
             inputs_embeds
         ).unsqueeze(-2)
+
+        sp_threshold = self.parallel_config.sp_threshold
+        tp_size = self.parallel_config.tensor_parallel_size
+        # Never shard fewer tokens than tp ranks -> plain TP.
+        is_sp_sharded = (
+            self.parallel_config.enable_sp
+            and hidden_states.shape[0] >= sp_threshold
+            and hidden_states.shape[0] >= tp_size
+        )
+        if is_sp_sharded:
+            hidden_states = sequence_parallel_chunk(hidden_states)
         hidden_states, residual, post_mix, res_mix = self.mtp_block(
-            positions=positions, x=hidden_states, input_ids=None
+            positions=positions,
+            x=hidden_states,
+            input_ids=None,
+            is_sp_sharded=is_sp_sharded,
         )
         hidden_states = mhc_post_tilelang(hidden_states, residual, post_mix, res_mix)
+        if is_sp_sharded:
+            hidden_states = tensor_model_parallel_all_gather(hidden_states, dim=0)
         # Return the flat pre-hc_head residual so it can be re-fed as the
         # next spec step's `previous_hidden_states` when
         # num_speculative_tokens > 1. hc_head is deferred to compute_logits.

vllm/v1/worker/gpu/model_runner.py

@@ -45,7 +45,7 @@ from vllm.model_executor.model_loader import get_model_loader
 from vllm.multimodal import MULTIMODAL_REGISTRY
 from vllm.sequence import IntermediateTensors
 from vllm.tasks import SupportedTask
-from vllm.utils.math_utils import cdiv
+from vllm.utils.math_utils import cdiv, round_up
 from vllm.utils.mem_utils import DeviceMemoryProfiler, format_gib
 from vllm.utils.torch_utils import PIN_MEMORY, STR_DTYPE_TO_TORCH_DTYPE
 from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
@@ -464,6 +464,21 @@ class GPUModelRunner(LoRAModelRunnerMixin):
             self.kv_cache_config,
             self.max_num_reqs,
         )
+
+        sp_threshold = self.parallel_config.sp_threshold
+        max_cg_size = self.compilation_config.max_cudagraph_capture_size
+        if (
+            sp_threshold is not None
+            and self.parallel_config.tensor_parallel_size > 1
+            and max_cg_size is not None
+            and sp_threshold <= max_cg_size
+        ):
+            raise ValueError(
+                f"sp_threshold ({sp_threshold}) must be larger than the max "
+                f"cudagraph capture size ({max_cg_size}) so that sequence "
+                "parallelism never engages on a captured cudagraph size."
+            )
+
         self.cudagraph_manager = ModelCudaGraphManager(
             self.vllm_config,
             self.device,
@@ -1143,6 +1158,14 @@ class GPUModelRunner(LoRAModelRunnerMixin):
         max_query_len = max(scheduler_output.num_scheduled_tokens.values())
         uniform_tok_count = get_uniform_token_count(num_reqs, num_toks, max_query_len)
 
+        # mHC sequence parallelism shards the token dim across the TP group.
+        # so a forward that engages SP (>= sp_threshold tokens) must run
+        # on a multiple of tp_size tokens.
+        tp_size = self.parallel_config.tensor_parallel_size
+        sp_threshold = self.parallel_config.sp_threshold
+        if sp_threshold is not None and num_toks >= sp_threshold:
+            num_toks = round_up(num_toks, tp_size)
+
         num_active_loras = 0
         if self.lora_config:
             req_ids = list(scheduler_output.num_scheduled_tokens.keys())