diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend.py b/python/sglang/srt/layers/attention/deepseek_v4_backend.py index 618e377ec7e0..ba38d6f32d80 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend.py +++ b/python/sglang/srt/layers/attention/deepseek_v4_backend.py @@ -783,7 +783,14 @@ def refresh_for_breakable_cuda_graph_replay_(self, other: DSV4AttnMetadata) -> N if src_val is None and dst_val is None: continue assert dst_val is not None, f"{field_name=} {src_val=} {dst_val=}" - dst_val.copy_(src_val) + shape_mismatch = dst_val.shape != src_val.shape + assert not shape_mismatch or field_name in self._CP_GLOBAL_FIELDS, ( + f"Only CP-global replay metadata may use a shorter live prefix, " + f"got {field_name=} {src_val.shape=} {dst_val.shape=}" + ) + _copy_tensor_allowing_storage_alias( + dst_val, src_val, pad_value=0 if shape_mismatch else None + ) # Safe to replace: captured kernels read only the per-replay objects, or # the field is produced in-graph before the attention graph break reads it. @@ -1093,6 +1100,27 @@ def _prefill_graph_max_seq_len() -> Optional[int]: return get_exec().graph.cuda_graph_config.prefill.max_seq_len +def _copy_tensor_allowing_storage_alias( + dst: torch.Tensor, src: torch.Tensor, *, pad_value: Optional[int] = None +) -> None: + """Copy replay metadata while preserving capture-stable destination addresses.""" + if dst is src: + return + if dst.untyped_storage().data_ptr() == src.untyped_storage().data_ptr(): + src = src.clone() + if dst.shape == src.shape: + dst.copy_(src) + return + assert ( + pad_value is not None + and dst.ndim == src.ndim + and dst.shape[0] >= src.shape[0] + and dst.shape[1:] == src.shape[1:] + ), f"Cannot copy replay metadata from {src.shape=} to {dst.shape=}" + dst.fill_(pad_value) + dst[: src.shape[0]].copy_(src) + + @dataclass class DSV4Metadata: core_attn_metadata: DSV4AttnMetadata @@ -3619,8 +3647,11 @@ def _forward_low_ratio_sources_cp( q_lens_cpu = tail.local_lens_cpu req_global, pos_global = tail.req_global, tail.pos_global else: + extend_lens_cpu = _as_int_list(forward_batch.extend_seq_lens_cpu) + # BCG gathers the capture bucket; compressed state belongs to live rows. + total = sum(extend_lens_cpu) q_lens_cpu = interleave_rows_per_request( - _as_int_list(forward_batch.extend_seq_lens_cpu), + extend_lens_cpu, get_parallel().attn_cp_rank, get_parallel().attn_cp_size, ) diff --git a/python/sglang/srt/layers/cp/bcg/__init__.py b/python/sglang/srt/layers/cp/bcg/__init__.py index 6b2bb5be5982..aeb28f787eb8 100644 --- a/python/sglang/srt/layers/cp/bcg/__init__.py +++ b/python/sglang/srt/layers/cp/bcg/__init__.py @@ -23,12 +23,20 @@ ] = { (ContextParallelStrategyKind.ZIGZAG, CPAttentionBackendKind.TRTLLM_MHA), (ContextParallelStrategyKind.INTERLEAVE, CPAttentionBackendKind.DSA), + (ContextParallelStrategyKind.INTERLEAVE, CPAttentionBackendKind.DSV4), } -def get_cp_bcg_input_type(strategy: ContextParallelStrategyKind): +def get_cp_bcg_input_type( + strategy: ContextParallelStrategyKind, + attention_backend: CPAttentionBackendKind = CPAttentionBackendKind.DSA, +): # Keep runtime adapters lazy: config resolution also queries CP BCG support. if strategy == ContextParallelStrategyKind.INTERLEAVE: + if attention_backend == CPAttentionBackendKind.DSV4: + from sglang.srt.layers.cp.bcg.dsv4_bcg import DSV4CPBCGInput + + return DSV4CPBCGInput from sglang.srt.layers.cp.bcg.interleave_bcg import InterleaveCPBCGInput return InterleaveCPBCGInput @@ -61,4 +69,4 @@ def supports_prefill_cp_bcg(server_args: ServerArgs) -> bool: return False if combination not in PREFILL_CP_BCG_SUPPORTED_COMBINATIONS: return False - return get_cp_bcg_input_type(combination[0]).supports(server_args) + return get_cp_bcg_input_type(*combination).supports(server_args) diff --git a/python/sglang/srt/layers/cp/bcg/bcg.py b/python/sglang/srt/layers/cp/bcg/bcg.py index a86f65b40088..a5ee01ac3f99 100644 --- a/python/sglang/srt/layers/cp/bcg/bcg.py +++ b/python/sglang/srt/layers/cp/bcg/bcg.py @@ -21,9 +21,14 @@ import torch -from sglang.srt.arg_groups.model_override_base import resolved_view, resolving_view +from sglang.srt.arg_groups.model_override_base import ( + attention_backends_of, + resolved_view, + resolving_view, +) from sglang.srt.layers.cp.base import ( ContextParallelStrategyKind, + CPAttentionBackendKind, get_cp_strategy_kind, ) from sglang.srt.layers.cp.padding import get_cp_padding_align_size @@ -33,6 +38,10 @@ prepare_cp_forward, ) from sglang.srt.model_executor.forward_batch_info import PPProxyTensors +from sglang.srt.model_executor.model_interfaces import ( + CPBoundedTailLogits, + CPLogitsInputs, +) if TYPE_CHECKING: from sglang.srt.model_executor.forward_batch_info import ForwardBatch @@ -83,13 +92,19 @@ class PrefillCPBCGInput: positions: torch.Tensor bucket_local_tokens: Dict[int, int] = field(default_factory=dict) live_local_tokens: int = 0 + supports_max_context_size = False @classmethod def create(cls, runner: PrefillCudaGraphRunner) -> PrefillCPBCGInput: if cls is PrefillCPBCGInput: from sglang.srt.layers.cp.bcg import get_cp_bcg_input_type - return get_cp_bcg_input_type(get_cp_strategy_kind()).create(runner) + backend, _ = attention_backends_of( + resolved_view(runner.model_runner.server_args) + ) + return get_cp_bcg_input_type( + get_cp_strategy_kind(), CPAttentionBackendKind.from_string(backend) + ).create(runner) with torch.device(runner.device): return cls( input_embeds=torch.zeros( @@ -112,6 +127,9 @@ def model_positions(self, forward_batch): # No override: retain the runner's existing model/mRoPE selection. return None + def model_input_ids(self, forward_batch): + return forward_batch.input_ids + @staticmethod def supports(server_args: ServerArgs) -> bool: return True @@ -325,11 +343,22 @@ def gather_aux(aux): hidden_states, hidden_states_before_norm = hidden_states if aux_hidden_states is None: logits_kwargs["hidden_states_before_norm"] = hidden_states_before_norm - return model.logits_processor( - forward_batch.input_ids, + logits_inputs = CPLogitsInputs(forward_batch.input_ids, forward_batch) + if isinstance(model, CPBoundedTailLogits): + logits_inputs = model.prepare_cp_logits_inputs( + forward_batch.input_ids, forward_batch + ) + + output = model.logits_processor( + logits_inputs.input_ids, hidden_states, model.lm_head, - forward_batch, + logits_inputs.logits_metadata, aux_hidden_states, **logits_kwargs, ) + if logits_inputs.hidden_states_token_indices is not None: + output.hidden_states_token_indices = ( + logits_inputs.hidden_states_token_indices + ) + return output diff --git a/python/sglang/srt/layers/cp/bcg/dsv4_bcg.py b/python/sglang/srt/layers/cp/bcg/dsv4_bcg.py new file mode 100644 index 000000000000..f1cb5c3bb6f5 --- /dev/null +++ b/python/sglang/srt/layers/cp/bcg/dsv4_bcg.py @@ -0,0 +1,212 @@ +"""DeepSeek-V4.1 interleave CP inputs for breakable prefill graphs.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Optional + +import torch + +from sglang.srt.layers.cp.base import get_cp_strategy +from sglang.srt.layers.cp.bcg.bcg import PrefillCPBCGInput +from sglang.srt.layers.cp.interleave import ( + InterleaveContextParallelMetadata, + InterleaveCPStrategy, +) +from sglang.srt.layers.cp.padding import get_cp_padding_align_size +from sglang.srt.layers.cp.utils import ( + cp_interleave_input_ids, + cp_shard_hidden_states, + cp_split_before_forward, + prepare_cp_forward, +) +from sglang.srt.layers.moe.utils import get_moe_a2a_backend +from sglang.srt.runtime_context import get_parallel + +if TYPE_CHECKING: + from sglang.srt.model_executor.forward_batch_info import ForwardBatch + from sglang.srt.model_executor.runner.prefill_cuda_graph_runner import ( + PrefillCudaGraphRunner, + ) + + +@dataclass +class DSV4CPBCGInput(PrefillCPBCGInput): + input_ids: Optional[torch.Tensor] = None + global_input_ids: Optional[torch.Tensor] = None + num_token_non_padded: Optional[torch.Tensor] = None + supports_max_context_size = True + + def _max_local_tokens(self, extend_seq_lens): + size = get_parallel().attn_cp_size + return (sum(int(length) for length in extend_seq_lens) + size - 1) // size + + def model_positions(self, forward_batch): + if forward_batch.cp_positions is None: + raise RuntimeError("CP BCG positions must be prepared before capture") + return forward_batch.cp_positions + + def model_input_ids(self, forward_batch): + if forward_batch.cp_input_ids is None: + raise RuntimeError("CP BCG token IDs must be prepared before capture") + return forward_batch.cp_input_ids + + @classmethod + def create(cls, runner: PrefillCudaGraphRunner) -> PrefillCPBCGInput: + strategy = get_cp_strategy() + global_id_capacity = 0 + if ( + isinstance(strategy, InterleaveCPStrategy) + and get_moe_a2a_backend().is_none() + ): + align = get_cp_padding_align_size() + local_capacity = ( + runner.max_num_tokens + strategy.cp_size - 1 + ) // strategy.cp_size + local_capacity = (local_capacity + align - 1) // align * align + global_id_capacity = strategy.cp_size * local_capacity + with torch.device(runner.device): + return cls( + input_embeds=torch.zeros( + ( + runner.max_num_tokens, + runner.model_runner.model_config.hidden_size, + ), + dtype=runner.model_runner.dtype, + ), + positions=torch.zeros( + (runner.max_num_tokens,), + dtype=torch.int64, + ), + input_ids=torch.zeros((runner.max_num_tokens,), dtype=torch.int64), + global_input_ids=( + torch.zeros((global_id_capacity,), dtype=torch.int64) + if global_id_capacity + else None + ), + num_token_non_padded=torch.zeros((), dtype=torch.int32), + ) + + def prepare( + self, + runner: PrefillCudaGraphRunner, + forward_batch: ForwardBatch, + *, + static_num_tokens: int, + capture: bool, + ) -> None: + """Shard global prefill inputs into fixed-address CP-local buffers.""" + # Replay batches may reuse a ForwardBatch object whose metadata was + # built for a different request layout. Always rebuild before sharding. + forward_batch.attn_cp_metadata = None + forward_batch.cp_positions = None + forward_batch.cp_input_ids = None + out_cache_loc = forward_batch.out_cache_loc + prepare_cp_forward(forward_batch) + metadata = forward_batch.attn_cp_metadata + if not isinstance(metadata, InterleaveContextParallelMetadata): + raise TypeError( + "DSV4 CP BCG requires InterleaveContextParallelMetadata, " + f"got {type(metadata).__name__}" + ) + + captured_local_tokens = None + if not capture: + try: + captured_local_tokens = self.bucket_local_tokens[static_num_tokens] + except KeyError as exc: + raise RuntimeError( + "Missing CP-local capture capacity for global prefill bucket " + f"{static_num_tokens}" + ) from exc + + # Breakable graph segments retain their captured row geometry when + # a smaller live batch reuses the same global token bucket. + live_physical_tokens = max(metadata.per_rank_actual_token) + if live_physical_tokens > captured_local_tokens: + raise RuntimeError( + f"Live batch needs {live_physical_tokens} local CP rows, " + f"but global prefill bucket {static_num_tokens} has " + f"captured capacity {captured_local_tokens}" + ) + cp_size = len(metadata.per_rank_actual_token) + metadata.per_rank_actual_token = [captured_local_tokens] * cp_size + metadata.max_rank_len = [captured_local_tokens] * cp_size + + raw_tokens = int(forward_batch.extend_num_tokens) + global_input_ids = forward_batch.input_ids[:raw_tokens] + global_positions = forward_batch.positions[:raw_tokens] + local_input_ids = cp_shard_hidden_states(global_input_ids, forward_batch) + global_input_embeds = runner.model_runner.model.get_input_embeddings()( + global_input_ids + ) + local_input_embeds, local_positions = cp_split_before_forward( + global_input_embeds, + global_positions, + forward_batch, + ) + live_local_tokens = int(local_input_embeds.shape[0]) + + if capture: + captured_local_tokens = live_local_tokens + self.bucket_local_tokens[static_num_tokens] = captured_local_tokens + else: + assert captured_local_tokens is not None + if live_local_tokens > captured_local_tokens: + raise RuntimeError( + f"Live batch needs {live_local_tokens} local CP rows, but global " + f"prefill bucket {static_num_tokens} has captured capacity " + f"{captured_local_tokens}" + ) + + if captured_local_tokens > self.input_embeds.shape[0]: + raise RuntimeError( + f"CP-local capture needs {captured_local_tokens} rows, but the " + f"fixed input buffer has capacity {self.input_embeds.shape[0]}" + ) + + input_embeds = self.input_embeds[:captured_local_tokens] + positions = self.positions[:captured_local_tokens] + assert self.input_ids is not None + input_ids = self.input_ids[:captured_local_tokens] + input_embeds.zero_() + positions.zero_() + input_ids.zero_() + input_embeds[:live_local_tokens].copy_(local_input_embeds) + positions[:live_local_tokens].copy_(local_positions) + input_ids[:live_local_tokens].copy_(local_input_ids) + forward_batch.input_embeds = input_embeds + forward_batch.cp_positions = positions + # Keep the global input_ids field intact: the runner uses its length to + # select the global capture bucket. MegaMoE routes local rows, whereas + # TP-MoE all-gathers hidden rows in CP rank order before routing. + forward_batch.cp_input_ids = input_ids + if self.global_input_ids is not None: + global_ids = cp_interleave_input_ids(global_input_ids, forward_batch) + if global_ids.numel() > self.global_input_ids.numel(): + raise RuntimeError("CP global token IDs exceed the capture buffer") + forward_batch.input_ids_global = self.global_input_ids[: global_ids.numel()] + forward_batch.input_ids_global.copy_(global_ids) + else: + forward_batch.input_ids_global = input_ids + if forward_batch.num_token_non_padded is not None: + assert self.num_token_non_padded is not None + metadata = forward_batch.attn_cp_metadata + logical_tokens = ( + metadata.per_rank_logical_token or metadata.per_rank_actual_token + ) + strategy = get_cp_strategy() + assert strategy is not None + self.num_token_non_padded.fill_(logical_tokens[strategy.cp_rank]) + forward_batch.num_token_non_padded = self.num_token_non_padded + self.live_local_tokens = live_local_tokens + + if isinstance(get_cp_strategy(), InterleaveCPStrategy): + metadata = forward_batch.attn_cp_metadata + cp_size = len(metadata.per_rank_actual_token) + base, remainder = divmod(static_num_tokens, cp_size) + metadata.total_seq_lens = static_num_tokens + metadata.per_rank_logical_token = [ + base + int(rank < remainder) for rank in range(cp_size) + ] + forward_batch.out_cache_loc = out_cache_loc diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index 3a132e6ce5f9..e1eaac67e5be 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/python/sglang/srt/model_executor/forward_batch_info.py @@ -694,6 +694,9 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin): tbo_children: Optional[List[ForwardBatch]] = None attn_cp_metadata: Optional[BaseContextParallelMetadata] = None + # Capture-stable local views, populated only by PrefillCPBCGInput.prepare. + cp_positions: Optional[torch.Tensor] = None + cp_input_ids: Optional[torch.Tensor] = None # For decode context parallel. # NOTE: DecodeContextParallelMetadata is imported under TYPE_CHECKING only (see the @@ -1015,6 +1018,8 @@ def init_new( encoder_lens=batch.encoder_lens, encoder_out_cache_loc=batch.encoder_out_cache_loc, input_embeds=batch.input_embeds, + cp_positions=None, + cp_input_ids=None, replace_embeds=batch.replace_embeds, replace_positions=batch.replace_positions, # Scalar config / flags diff --git a/python/sglang/srt/model_executor/model_interfaces.py b/python/sglang/srt/model_executor/model_interfaces.py new file mode 100644 index 000000000000..31259e1152bd --- /dev/null +++ b/python/sglang/srt/model_executor/model_interfaces.py @@ -0,0 +1,60 @@ +"""Explicit model capabilities used by context-parallel runners.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from dataclasses import dataclass +from typing import TYPE_CHECKING, Optional, Tuple, Union + +if TYPE_CHECKING: + import torch + + from sglang.srt.layers.logits_processor import LogitsMetadata + from sglang.srt.model_executor.forward_batch_info import ForwardBatch + + +class LanguageModelInputPreparation: + """Opt in to preparing global model inputs before CP sharding. + + The default is the identity preparation used by text-only models. + """ + + def prepare_language_model_inputs( + self, + input_ids: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: Optional[torch.Tensor] = None, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + return input_ids, input_embeds + + +def prepare_language_model_inputs( + model, + input_ids: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: Optional[torch.Tensor] = None, +) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + """Dispatch only to models that explicitly declare input preparation.""" + if isinstance(model, LanguageModelInputPreparation): + return model.prepare_language_model_inputs( + input_ids, forward_batch, input_embeds + ) + return input_ids, input_embeds + + +@dataclass +class CPLogitsInputs: + input_ids: torch.Tensor + logits_metadata: Union[ForwardBatch, LogitsMetadata] + hidden_states_token_indices: Optional[torch.Tensor] = None + + +class CPBoundedTailLogits(ABC): + """Opt in to adapting global CP outputs to a model's bounded logits tail.""" + + @abstractmethod + def prepare_cp_logits_inputs( + self, input_ids: torch.Tensor, forward_batch: ForwardBatch + ) -> CPLogitsInputs: + """Return model-owned logits inputs and the optional output row mapping.""" + raise NotImplementedError diff --git a/python/sglang/srt/model_executor/runner/eager_runner.py b/python/sglang/srt/model_executor/runner/eager_runner.py index b3d9bda5b015..9be3aeb78946 100644 --- a/python/sglang/srt/model_executor/runner/eager_runner.py +++ b/python/sglang/srt/model_executor/runner/eager_runner.py @@ -45,6 +45,7 @@ get_token_to_kv_pool, ) from sglang.srt.model_executor.forward_mode import ForwardMode +from sglang.srt.model_executor.model_interfaces import prepare_language_model_inputs from sglang.srt.model_executor.runner.base_runner import BaseRunner from sglang.srt.model_executor.runner_utils.shared_read_event import ( maybe_publish_prefill_shared_read_done, @@ -351,14 +352,19 @@ def _execute_extend_cp( """ model = self.model_runner.model + input_ids = forward_batch.input_ids if prepare_inputs := getattr(model, "prepare_cp_inputs", None): input_embeds, positions, model_kwargs = prepare_inputs( forward_batch, **kwargs ) else: input_embeds = kwargs.get("input_embeds") + # Prepare multimodal embeddings and model IDs in global token order. + input_ids, input_embeds = prepare_language_model_inputs( + model, input_ids, forward_batch, input_embeds + ) if input_embeds is None: - input_embeds = model.get_input_embeddings()(forward_batch.input_ids) + input_embeds = model.get_input_embeddings()(input_ids) positions = forward_batch.positions model_kwargs = {} if (pp_proxy_tensors := kwargs.get("pp_proxy_tensors")) is not None: @@ -368,7 +374,7 @@ def _execute_extend_cp( input_embeds, positions, forward_batch, - forward_batch.input_ids, + input_ids, ) as (sharded_input_embeds, sharded_positions, model_input_ids): model_kwargs["input_embeds"] = sharded_input_embeds hidden_states = model.model( @@ -411,7 +417,7 @@ def _execute_extend_cp( if aux_hidden_states is None: logits_kwargs["hidden_states_before_norm"] = hidden_states_before_norm return model.logits_processor( - forward_batch.input_ids, + input_ids, hidden_states, model.lm_head, forward_batch, diff --git a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py index c7c08e4e5dd5..6186746c9498 100644 --- a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py @@ -652,14 +652,15 @@ def __init__(self, model_runner: ModelRunner): self.backend, BreakableCudaGraphBackend ) and supports_prefill_cp_bcg(server_args) if self.enable_cp_bcg_capture: - if self.max_context_size is not None: - # TODO(SYChen123): Preserve max_seq_len_override through CP's padded - # metadata preparation before enabling the fixed context limit. - self._ignore_max_context_size("CP breakable prefill CUDA graph") self.capture_num_tokens = filter_prefill_cp_bcg_capture_num_tokens( self.capture_num_tokens, server_args ) self.prefill_cp_bcg_input = PrefillCPBCGInput.create(self) + if ( + self.max_context_size is not None + and not self.prefill_cp_bcg_input.supports_max_context_size + ): + self._ignore_max_context_size("CP breakable prefill CUDA graph") logger.info( "Prefill CP breakable CUDA graph enabled: strategy=%s, cp_size=%d", get_parallel().cp_strategy, @@ -941,6 +942,8 @@ def _run_forward(self, forward_batch: ForwardBatch, num_tokens: int): # BCG / Full: capture the transformer body only. positions = self._get_layer_model_positions(forward_batch) input_ids = forward_batch.input_ids + if self.enable_cp_bcg_capture: + input_ids = self.prefill_cp_bcg_input.model_input_ids(forward_batch) kwargs = _build_layer_model_forward_kwargs( self.layer_model, forward_batch, pp_proxy_tensors ) @@ -1587,6 +1590,14 @@ def can_run_graph(self, forward_batch: ForwardBatch) -> bool: if getattr(self, "enable_cp_bcg_capture", False): if not is_cp_active(forward_batch): return False + # CP BCG embeds tokens before replay. Image spans need the vision + # encoder and ID normalization performed by the eager CP runner. + if ( + getattr(self.model_runner.model, "vision", None) is not None + and forward_batch.mm_inputs is not None + and any(item is not None for item in forward_batch.mm_inputs) + ): + return False assert self.prefill_cp_bcg_input is not None if ( self.prefill_cp_bcg_input.select_replay_bucket_for_batch( @@ -2480,9 +2491,9 @@ def execute( self._wait_for_layer_loads() if self.enable_cp_bcg_capture: - assert self.max_context_size is None, ( - "CP-v2 BCG replay does not support a fixed prefill context size" - ) + # assert self.max_context_size is None, ( + # "CP-v2 BCG replay does not support a fixed prefill context size" + # ) output = execute_prefill_cp_bcg( self, forward_batch, diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index 596d4622b3af..b88175da6ff7 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/python/sglang/srt/models/deepseek_v4.py @@ -140,6 +140,11 @@ get_token_to_kv_pool, ) from sglang.srt.model_executor.forward_mode import CaptureHiddenMode +from sglang.srt.model_executor.model_interfaces import ( + CPBoundedTailLogits, + CPLogitsInputs, + LanguageModelInputPreparation, +) from sglang.srt.model_executor.runner import ( compile_in_capture_mode, get_is_capture_mode, @@ -934,7 +939,8 @@ def _low_ratio_sources_break(attn_backend, forward_batch) -> bool: forward_batch.forward_mode.is_extend() and is_in_breakable_cuda_graph() and ( - not getattr(attn_backend, "low_ratio_prefill_graph", False) + is_cp_active(forward_batch) + or not getattr(attn_backend, "low_ratio_prefill_graph", False) or in_replay_tail_graph() ) ) @@ -4044,11 +4050,17 @@ def _forward_layers_hc_pre_from_prev( ) if self.engram_hasher is not None: if cp_extend: - # n-gram hashing needs each token's predecessors: hash the whole prompt + # N-gram hashing needs each token's predecessors, so hash the + # whole prompt before selecting this CP rank's interleaved rows. total = int(forward_batch.attn_cp_metadata.total_seq_lens) - hash_ids = self.engram_hasher( - forward_batch.input_ids[:total], forward_batch - ) + global_input_ids = forward_batch.input_ids[:total] + if is_in_breakable_cuda_graph(): + hash_ids = self.engram_hasher.new_hash_ids(global_input_ids) + bcg_deepseek_v4_engram_hash_ids( + self.engram_hasher, global_input_ids, hash_ids + ) + else: + hash_ids = self.engram_hasher(global_input_ids, forward_batch) parallel = get_parallel() hash_ids = hash_ids[parallel.attn_cp_rank :: parallel.attn_cp_size] pad_rows = hidden_states.shape[0] - hash_ids.shape[0] @@ -4681,7 +4693,14 @@ def _is_quark_checkpoint(config: DeepSeekV4Config) -> bool: return quantization_config.get("quant_method") == "quark" -class DeepseekV4ForCausalLM(nn.Module): +def _v41_vision_a2a_supported() -> bool: + backend = get_moe_a2a_backend() + return backend.is_none() or (_is_cuda and backend.is_megamoe()) + + +class DeepseekV4ForCausalLM( + nn.Module, LanguageModelInputPreparation, CPBoundedTailLogits +): supports_cuda_vmm_feature_transport = True @classmethod @@ -4733,15 +4752,11 @@ def __init__( and not getattr(config, "language_model_only", False) ): if ( - get_parallel().attn_cp_size != 1 - or get_parallel().pp_group.world_size != 1 - or not ( - get_moe_a2a_backend().is_none() - or (_is_cuda and get_moe_a2a_backend().is_megamoe()) - ) + get_parallel().pp_group.world_size != 1 + or not _v41_vision_a2a_supported() ): raise ValueError( - "V4.1 vision supports TP/EP/DP without CP or PP; " + "V4.1 vision supports TP/EP/DP and prefill CP without PP; " "MoE A2A requires CUDA MegaMoE" ) @@ -4948,6 +4963,25 @@ def num_fusable_shared_experts(cls, hf_config) -> int: """Shared experts fused as routed slots when fusion is on.""" return hf_config.n_shared_experts or 0 + def prepare_cp_logits_inputs( + self, input_ids: torch.Tensor, forward_batch: ForwardBatch + ) -> CPLogitsInputs: + if not ( + self.capture_aux_hidden_states + and self.model.late_layer_start is not None + and forward_batch.forward_mode.is_extend_without_speculative() + ): + return CPLogitsInputs(input_ids, forward_batch) + + tail = get_attn_backend().tail_forward_metadata.late_layer_tail + if tail is None: + raise RuntimeError("CP bounded-tail logits require prepared tail metadata") + logits_metadata = LogitsMetadata.from_forward_batch(forward_batch) + logits_metadata.extend_seq_lens = tail.extend_seq_lens + logits_metadata.extend_seq_lens_cpu = tail.extend_seq_lens_cpu + logits_metadata.extend_logprob_start_lens_cpu = tail.extend_seq_lens_cpu + return CPLogitsInputs(tail.rows(input_ids), logits_metadata, tail.token_indices) + def prepare_language_model_inputs( self, input_ids: torch.Tensor, diff --git a/test/registered/cp/test_cp_strategy_unit.py b/test/registered/cp/test_cp_strategy_unit.py index 3e39063f4eb9..2d3ed7ba873d 100644 --- a/test/registered/cp/test_cp_strategy_unit.py +++ b/test/registered/cp/test_cp_strategy_unit.py @@ -16,6 +16,12 @@ is_interleave, is_zigzag, ) +from sglang.srt.layers.cp.bcg import get_cp_bcg_input_type, supports_prefill_cp_bcg +from sglang.srt.layers.cp.bcg.bcg import ( + PrefillCPBCGInput, + filter_prefill_cp_bcg_capture_num_tokens, +) +from sglang.srt.layers.cp.bcg.dsv4_bcg import DSV4CPBCGInput from sglang.srt.layers.cp.bcg.interleave_bcg import InterleaveCPBCGInput from sglang.srt.layers.cp.bcg.zigzag_bcg import ZigzagCPBCGInput from sglang.srt.layers.cp.interleave import InterleaveCPStrategy @@ -32,6 +38,7 @@ from sglang.srt.layers.cp.zigzag import ZigzagCPStrategy from sglang.srt.mem_cache.memory_pool import KVWriteLoc from sglang.srt.model_executor.cuda_graph_config import Backend +from sglang.srt.model_executor.forward_batch_info import ForwardBatch from sglang.srt.model_executor.forward_context import ForwardContext, forward_context from sglang.srt.model_executor.forward_mode import CaptureHiddenMode, ForwardMode from sglang.srt.model_executor.runner.prefill_cuda_graph_runner import ( @@ -59,6 +66,103 @@ def all_gather_into_tensor(self, output, input_tensor): torch.cat(self.all_rank_tensors, dim=0, out=output) +@contextmanager +def _bcg_config( + *, strategy="interleave", backend="dsv4", cp=8, tp=8, pp=1, enabled=True +): + cfg = SimpleNamespace( + enable_prefill_cp=enabled, pp_size=pp, tp_size=tp, cp_strategy=strategy + ) + resolved = SimpleNamespace(attn_cp_size=cp) + module = "sglang.srt.layers.cp.bcg" + with ( + patch(module + ".resolving_view", return_value=cfg), + patch(module + ".resolved_view", return_value=resolved), + patch(module + ".attention_backends_of", return_value=(backend, backend)), + patch(module + ".bcg.resolving_view", return_value=cfg), + patch(module + ".bcg.resolved_view", return_value=resolved), + patch(module + ".bcg.attention_backends_of", return_value=(backend, backend)), + ): + yield cfg + + +class TestCPBCGConfiguration(CustomTestCase): + def test_interleave_backend_selects_the_matching_adapter(self): + self.assertIs( + get_cp_bcg_input_type( + ContextParallelStrategyKind.INTERLEAVE, CPAttentionBackendKind.DSV4 + ), + DSV4CPBCGInput, + ) + self.assertIs( + get_cp_bcg_input_type( + ContextParallelStrategyKind.INTERLEAVE, CPAttentionBackendKind.DSA + ), + InterleaveCPBCGInput, + ) + self.assertIs( + get_cp_bcg_input_type( + ContextParallelStrategyKind.ZIGZAG, CPAttentionBackendKind.TRTLLM_MHA + ), + ZigzagCPBCGInput, + ) + self.assertTrue(DSV4CPBCGInput.supports_max_context_size) + self.assertFalse(InterleaveCPBCGInput.supports_max_context_size) + + def test_runtime_factory_selects_dsv4_without_glm_replay_restrictions(self): + init_cp_strategy(enable_prefill_cp=True, cp_size=4, cp_strategy="interleave") + try: + with ( + _bcg_config(cp=4, tp=4) as cfg, + get_parallel().override(attn_cp_size=4, attn_cp_rank=0), + ): + runner = SimpleNamespace( + device="cpu", + max_num_tokens=16, + model_runner=SimpleNamespace( + server_args=cfg, + model_config=SimpleNamespace(hidden_size=1), + dtype=torch.float32, + ), + ) + adapter = PrefillCPBCGInput.create(runner) + self.assertIsInstance(adapter, DSV4CPBCGInput) + self.assertTrue(adapter.allows_replay(2, 12, [4, 8], True)) + glm = InterleaveCPBCGInput( + input_embeds=torch.empty(0), + positions=torch.empty(0), + bucket_local_tokens={16: 4}, + ) + self.assertFalse(glm.allows_replay(2, 12, [4, 8], True)) + finally: + init_cp_strategy(enable_prefill_cp=False, cp_size=1, cp_strategy="zigzag") + + def test_dsv41_interleave_enables_cp_bcg_and_filters_small_buckets(self): + with _bcg_config() as cfg: + self.assertTrue(supports_prefill_cp_bcg(cfg)) + self.assertTrue(supports_prefill_cp_bcg(cfg)) + self.assertEqual( + filter_prefill_cp_bcg_capture_num_tokens([1, 2, 4, 8, 16], cfg), + [8, 16], + ) + + def test_unsupported_cp_bcg_configurations_stay_disabled(self): + cases = ( + dict(strategy="zigzag"), + dict(backend="fa3"), + dict(cp=4), + dict(pp=2), + dict(enabled=False), + ) + for case in cases: + with self.subTest(case=case), _bcg_config(**case) as cfg: + self.assertFalse(supports_prefill_cp_bcg(cfg)) + + def test_trtllm_zigzag_support_is_preserved(self): + with _bcg_config(strategy="zigzag", backend="trtllm_mha") as cfg: + self.assertTrue(supports_prefill_cp_bcg(cfg)) + + class TestCPStrategyUnit(CustomTestCase): def tearDown(self): init_cp_strategy(enable_prefill_cp=False, cp_size=1, cp_strategy="zigzag") @@ -122,6 +226,7 @@ def tearDown(self): def _make_runner(self): runner = PrefillCudaGraphRunner.__new__(PrefillCudaGraphRunner) + runner.model_runner = SimpleNamespace(model=SimpleNamespace(vision=None)) runner._is_full_backend = False runner.enable_lora = False runner._capture_chunked_prefix = False @@ -187,6 +292,60 @@ def test_interleave_replay_uses_total_tokens_across_requests(self): batch.extend_prefix_lens_cpu = [0, 1, 0] self.assertTrue(runner.can_run_graph(batch)) + def test_dsv41_capture_prepares_cp_before_attention_metadata(self): + class MetadataChecked(Exception): + pass + + init_cp_strategy(enable_prefill_cp=True, cp_size=8, cp_strategy="interleave") + batch = ForwardBatch( + forward_mode=ForwardMode.EXTEND, + batch_size=2, + input_ids=torch.arange(16), + req_pool_indices=torch.tensor([0, 1]), + seq_lens=torch.tensor([8, 8]), + out_cache_loc=torch.arange(16), + seq_lens_sum=16, + positions=torch.arange(16), + seq_lens_cpu=torch.tensor([8, 8]), + extend_seq_lens_cpu=[8, 8], + extend_num_tokens=16, + max_seq_len_override=51200, + ) + runner = self._make_runner() + runner.device = "cpu" + runner.max_num_tokens = 16 + runner.model_runner = SimpleNamespace( + model=SimpleNamespace( + get_input_embeddings=lambda: lambda ids: ids[:, None].float() + ), + model_config=SimpleNamespace(hidden_size=1), + dtype=torch.float32, + ) + runner.capture_prepare = lambda size: (batch, None) + + def check_metadata(forward_batch, shape_key): + self.assertIsNotNone(forward_batch.attn_cp_metadata) + self.assertEqual(forward_batch.attn_cp_metadata.total_seq_lens, 16) + self.assertEqual( + len(forward_batch.attn_cp_metadata.per_rank_actual_token), 8 + ) + self.assertIsNotNone(forward_batch.cp_positions) + self.assertIsNotNone(forward_batch.cp_input_ids) + self.assertEqual(forward_batch.max_seq_len_override, 51200) + raise MetadataChecked + + runner._init_forward_metadata_for_capture = check_metadata + with ( + _bcg_config() as cfg, + get_parallel().override( + attn_cp_rank=0, attn_cp_size=8, enable_cp_tp_group_sharing=True + ), + ): + runner.enable_cp_bcg_capture = supports_prefill_cp_bcg(cfg) + runner.prefill_cp_bcg_input = DSV4CPBCGInput.create(runner) + with self.assertRaises(MetadataChecked): + runner.capture_one_shape(16) + def test_local_capacity_overflow_uses_next_capture_bucket(self): runner = self._make_runner() runner.capture_num_tokens.append(2560) diff --git a/test/registered/unit/layers/test_dsv41_full_topk_cp.py b/test/registered/unit/layers/test_dsv41_full_topk_cp.py new file mode 100644 index 000000000000..11612277763a --- /dev/null +++ b/test/registered/unit/layers/test_dsv41_full_topk_cp.py @@ -0,0 +1,126 @@ +"""CPU regressions for the V4.1 Torch indexer with interleaved CP rows.""" + +import unittest +from types import SimpleNamespace +from unittest.mock import patch + +import torch + +from sglang.srt.layers.attention.dsv4.v41_indexer.full_topk import FullTopKIndexer +from sglang.srt.layers.attention.dsv4.v41_indexer.scoring import prefill_k_slots +from sglang.srt.layers.attention.dsv4.v41_indexer.types import PrefillInputs +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=3, suite="base-a-test-cpu") + + +class TestDSV41FullTopKCP(CustomTestCase): + def setUp(self): + self.req_to_token = torch.arange(8 * 64).view(8, 64) * 2 + self.pool = SimpleNamespace( + get_low_ratio_index_k_dequant=lambda layer, slots: slots.float()[:, None] + ) + with patch( + "sglang.srt.layers.attention.dsv4.v41_indexer.full_topk.is_deep_select_supported", + return_value=False, + ): + self.backend = FullTopKIndexer( + token_to_kv_pool=self.pool, + req_to_token=self.req_to_token, + use_deep_gemm_prefill=False, + use_deep_gemm_decode=False, + ) + + def _inputs(self, req, pos, counts, ratio, *, explicit=False, device_counts=True): + rows = pos.numel() + indexer = SimpleNamespace( + index_topk=3, + queries=lambda q, freqs: q[:, None, :], + head_weights=lambda x: x.new_ones((x.shape[0], 1)), + scores=lambda q, k, weights: q[:, 0, 0, None] * k[None, :, 0], + ) + req_pool_indices = torch.tensor([5, 2, 7]) + seq_lens_cpu = [20, 0, 20] + return PrefillInputs( + indexer=indexer, + layer_id=0, + compress_ratio=ratio, + freqs_cis=torch.zeros(64, 1), + x=torch.ones(rows, 1), + q_lora=torch.ones(rows, 1), + positions=pos, + req_rows=req if explicit else None, + req_pool_indices=req_pool_indices, + k_slots=prefill_k_slots( + req_to_token=self.req_to_token, + req_pool_indices=req_pool_indices, + seq_lens_cpu=seq_lens_cpu, + ratio=ratio, + ), + index_plans=None, + kv_page_table=torch.zeros(rows, 1, dtype=torch.int32), + seq_lens_cpu=seq_lens_cpu, + rows_per_request=counts, + rows_per_request_device=( + torch.tensor(counts, dtype=torch.int32) if device_counts else None + ), + out_raw_indices=torch.full((rows + 2, 3), 99, dtype=torch.int32), + out_page_indices=torch.full((rows + 2, 3), 99, dtype=torch.int32), + ) + + def _check(self, inputs, req): + self.backend.topk_prefill(inputs) + expected_raw = torch.full_like(inputs.out_raw_indices, -1) + expected_page = torch.full_like(inputs.out_page_indices, -1) + for row, (request, position) in enumerate( + zip(req.tolist(), inputs.positions.tolist()) + ): + visible = (position + 1) // inputs.compress_ratio + picks = torch.arange(max(0, visible - 3), visible) + expected_raw[row, : len(picks)] = picks.to(torch.int32) + slots = ( + self.req_to_token[request, picks * inputs.compress_ratio] + // inputs.compress_ratio + ) + expected_page[row, : len(picks)] = slots.to(torch.int32) + torch.testing.assert_close(inputs.out_raw_indices, expected_raw) + torch.testing.assert_close(inputs.out_page_indices, expected_page) + + def test_interleaved_cp_topk_with_prefix_and_empty_requests(self): + lengths = [3, 0, 9] + req = torch.repeat_interleave(torch.tensor([5, 2, 7]), torch.tensor(lengths)) + pos = torch.cat([torch.arange(7, 10), torch.arange(11, 20)]) + for ratio in (1, 2): + for cp in (2, 8): + for rank in range(cp): + with self.subTest(ratio=ratio, cp=cp, rank=rank): + counts = [int((req[rank::cp] == r).sum()) for r in (5, 2, 7)] + inputs = self._inputs( + req[rank::cp], pos[rank::cp], counts, ratio + ) + self._check(inputs, req[rank::cp]) + + def test_small_prompt_has_empty_ranks_and_no_visible_compressed_keys(self): + req = torch.tensor([5, 7, 7]) + pos = torch.tensor([0, 0, 1]) + for rank in range(8): + with self.subTest(rank=rank): + counts = [int((req[rank::8] == r).sum()) for r in (5, 2, 7)] + self._check( + self._inputs(req[rank::8], pos[rank::8], counts, 2), req[rank::8] + ) + + def test_cp_supports_host_row_counts(self): + req = torch.tensor([5, 7, 7]) + pos = torch.tensor([8, 13, 15]) + self._check(self._inputs(req, pos, [1, 0, 2], 2, device_counts=False), req) + + def test_explicit_request_rows_keep_existing_torch_behavior(self): + req = torch.tensor([7, 5, 7]) + pos = torch.tensor([12, 8, 14]) + self._check(self._inputs(req, pos, [1, 0, 2], 2, explicit=True), req) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/model_executor/test_cp_model_interfaces.py b/test/registered/unit/model_executor/test_cp_model_interfaces.py new file mode 100644 index 000000000000..c9c2d676a8d2 --- /dev/null +++ b/test/registered/unit/model_executor/test_cp_model_interfaces.py @@ -0,0 +1,242 @@ +"""CPU regressions for explicit CP batch fields and model capabilities.""" + +import unittest +from contextlib import nullcontext +from types import SimpleNamespace as NS +from unittest.mock import Mock, patch + +import torch + +import sglang.srt.model_executor.runner.prefill_cuda_graph_runner as prefill +from sglang.srt.layers.cp.bcg.bcg import execute_prefill_cp_bcg +from sglang.srt.layers.cp.bcg.dsv4_bcg import DSV4CPBCGInput +from sglang.srt.model_executor.forward_batch_info import ForwardBatch +from sglang.srt.model_executor.forward_mode import ForwardMode +from sglang.srt.model_executor.model_interfaces import ( + CPBoundedTailLogits, + LanguageModelInputPreparation, + prepare_language_model_inputs, +) +from sglang.srt.models.deepseek_v4 import DeepseekV4ForCausalLM +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + +MODEL = "sglang.srt.models.deepseek_v4" +BCG = "sglang.srt.layers.cp.bcg.bcg" + + +class _PreparedModel(NS, LanguageModelInputPreparation): + pass + + +class _TailModel(NS, CPBoundedTailLogits): + prepare_cp_logits_inputs = DeepseekV4ForCausalLM.prepare_cp_logits_inputs + + +def make_batch(): + return ForwardBatch( + forward_mode=ForwardMode.EXTEND, + batch_size=1, + input_ids=torch.arange(6), + req_pool_indices=torch.tensor([0]), + seq_lens=torch.tensor([6]), + out_cache_loc=torch.arange(6), + seq_lens_sum=6, + positions=torch.arange(6), + extend_seq_lens=torch.tensor([6]), + extend_seq_lens_cpu=[6], + extend_prefix_lens_cpu=[0], + dp_padding_mode=NS(is_max_len=lambda: True), + ) + + +class TestCPModelInterfaces(CustomTestCase): + def test_text_model_preparation_is_identity(self): + batch = make_batch() + embeds = torch.ones(6, 2) + for model in (object(), LanguageModelInputPreparation()): + ids, prepared = prepare_language_model_inputs( + model, batch.input_ids, batch, embeds + ) + self.assertIs(ids, batch.input_ids) + self.assertIs(prepared, embeds) + + def test_method_presence_does_not_declare_preparation(self): + batch = make_batch() + model = NS(prepare_language_model_inputs=Mock(side_effect=AssertionError)) + ids, embeds = prepare_language_model_inputs(model, batch.input_ids, batch) + self.assertIs(ids, batch.input_ids) + self.assertIsNone(embeds) + model.prepare_language_model_inputs.assert_not_called() + + def test_declared_preparation_runs_and_propagates_failures(self): + batch = make_batch() + normalized = batch.input_ids + 100 + embeds = torch.ones(6, 2) + model = _PreparedModel( + prepare_language_model_inputs=Mock(return_value=(normalized, embeds)) + ) + ids, prepared = prepare_language_model_inputs(model, batch.input_ids, batch) + self.assertIs(ids, normalized) + self.assertIs(prepared, embeds) + model.prepare_language_model_inputs.assert_called_once_with( + batch.input_ids, batch, None + ) + model.prepare_language_model_inputs.side_effect = ValueError("vision failed") + with self.assertRaisesRegex(ValueError, "vision failed"): + prepare_language_model_inputs(model, batch.input_ids, batch) + + def _make_capture_runner(self, cp): + runner = prefill.PrefillCudaGraphRunner.__new__(prefill.PrefillCudaGraphRunner) + runner.enable_cp_bcg_capture = cp + runner.prefill_cp_bcg_input = DSV4CPBCGInput( + input_embeds=torch.empty(0), positions=torch.empty(0) + ) + runner._prefill_forward_context = lambda *a, **kw: nullcontext() + runner._capture_pp_proxy_tensors = lambda n: None + runner._uses_eager_prefill_tail = lambda: True + runner._pack_qwen_bcg_hc_output = lambda output: output + runner.layer_model = NS(forward=Mock(return_value=torch.ones(3, 2))) + return runner + + def _capture(self, runner, batch): + with ( + patch.object(prefill, "set_dp_buffer_len"), + patch.object(prefill, "set_is_extend_in_batch"), + patch.object(prefill, "_build_layer_model_forward_kwargs", return_value={}), + ): + return runner._run_forward(batch, 6) + + def test_non_cp_capture_keeps_global_inputs(self): + batch = make_batch() + self.assertIsNone(batch.cp_positions) + self.assertIsNone(batch.cp_input_ids) + runner = self._make_capture_runner(cp=False) + self._capture(runner, batch) + args = runner.layer_model.forward.call_args.args + self.assertIs(args[0], batch.input_ids) + self.assertIs(args[1], batch.positions) + + def test_cp_capture_requires_prepared_positions_and_ids(self): + batch = make_batch() + runner = self._make_capture_runner(cp=True) + with self.assertRaisesRegex(RuntimeError, "positions must be prepared"): + self._capture(runner, batch) + batch.cp_positions = batch.positions[::2] + with self.assertRaisesRegex(RuntimeError, "token IDs must be prepared"): + self._capture(runner, batch) + runner.layer_model.forward.assert_not_called() + + def test_cp_capture_uses_local_fields_without_mutating_global_inputs(self): + batch = make_batch() + global_ids = batch.input_ids.clone() + global_positions = batch.positions.clone() + batch.cp_positions = batch.positions[::2].clone() + batch.cp_input_ids = batch.input_ids[::2].clone() + runner = self._make_capture_runner(cp=True) + self._capture(runner, batch) + args = runner.layer_model.forward.call_args.args + self.assertIs(args[0], batch.cp_input_ids) + self.assertIs(args[1], batch.cp_positions) + torch.testing.assert_close(batch.input_ids, global_ids) + torch.testing.assert_close(batch.positions, global_positions) + + def _tail_model(self, **kwargs): + return _TailModel( + capture_aux_hidden_states=True, + model=NS(late_layer_start=1), + pp_group=NS(is_last_rank=True), + lm_head=object(), + logits_processor=Mock(return_value=NS()), + **kwargs, + ) + + def _replay(self, model, batch, local_output): + runner = NS( + prefill_cp_bcg_input=NS(live_local_tokens=6), + model_runner=NS(model=model), + backend=NS(replay=Mock(return_value=local_output)), + _prefill_forward_context=lambda *a, **kw: nullcontext(), + ) + with ( + patch(BCG + ".cp_gather_after_forward", side_effect=lambda x, *a: x), + patch(BCG + ".torch.cuda.current_stream", return_value=None), + ): + return execute_prefill_cp_bcg(runner, batch, batch, 6, 6) + + def test_bounded_tail_metadata_and_output_mapping_survive_cp_replay(self): + batch = make_batch() + indices = torch.tensor([1, 4]) + tail = NS( + rows=lambda x: x[indices], + extend_seq_lens=torch.tensor([2]), + extend_seq_lens_cpu=[2], + token_indices=indices, + ) + backend = NS(tail_forward_metadata=NS(late_layer_tail=tail)) + model = self._tail_model() + hidden = torch.ones(2, 3) + aux = [torch.full((2, 3), 2.0)] + with ( + patch(MODEL + ".get_attn_backend", return_value=backend), + patch(MODEL + ".LogitsMetadata.from_forward_batch", return_value=NS()), + ): + output = self._replay(model, batch, (hidden, aux)) + args = model.logits_processor.call_args.args + torch.testing.assert_close(args[0], batch.input_ids[indices]) + torch.testing.assert_close(args[1], hidden) + self.assertEqual(args[1].data_ptr(), hidden.data_ptr()) + self.assertIs(args[3].extend_seq_lens, tail.extend_seq_lens) + self.assertEqual(args[3].extend_seq_lens_cpu, [2]) + self.assertEqual(args[3].extend_logprob_start_lens_cpu, [2]) + torch.testing.assert_close(args[4][0], aux[0]) + self.assertIs(output.hidden_states_token_indices, indices) + self.assertEqual(batch.extend_seq_lens_cpu, [6]) + + def test_bounded_tail_requires_metadata_when_enabled(self): + batch = make_batch() + model = self._tail_model() + backend = NS(tail_forward_metadata=NS(late_layer_tail=None)) + with patch(MODEL + ".get_attn_backend", return_value=backend): + with self.assertRaisesRegex(RuntimeError, "prepared tail metadata"): + model.prepare_cp_logits_inputs(batch.input_ids, batch) + + def test_bounded_tail_is_identity_outside_aux_extend(self): + batch = make_batch() + model = self._tail_model() + cases = ( + (False, 1, ForwardMode.EXTEND), + (True, None, ForwardMode.EXTEND), + (True, 1, ForwardMode.DECODE), + (True, 1, ForwardMode.TARGET_VERIFY), + ) + for capture, start, mode in cases: + with self.subTest(capture=capture, start=start, mode=mode): + model.capture_aux_hidden_states = capture + model.model.late_layer_start = start + batch.forward_mode = mode + with patch(MODEL + ".get_attn_backend") as get_backend: + inputs = model.prepare_cp_logits_inputs(batch.input_ids, batch) + self.assertIs(inputs.input_ids, batch.input_ids) + self.assertIs(inputs.logits_metadata, batch) + self.assertIsNone(inputs.hidden_states_token_indices) + get_backend.assert_not_called() + + def test_cp_logits_without_tail_capability_keep_global_metadata(self): + batch = make_batch() + model = NS( + capture_aux_hidden_states=False, + pp_group=NS(is_last_rank=True), + lm_head=object(), + logits_processor=Mock(return_value=NS()), + ) + self._replay(model, batch, torch.ones(6, 3)) + args = model.logits_processor.call_args.args + self.assertIs(args[0], batch.input_ids) + self.assertIs(args[3], batch) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/models/test_deepseek_v41_bcg_multimodal.py b/test/registered/unit/models/test_deepseek_v41_bcg_multimodal.py index d4cb6fa063df..98acbdb36698 100644 --- a/test/registered/unit/models/test_deepseek_v41_bcg_multimodal.py +++ b/test/registered/unit/models/test_deepseek_v41_bcg_multimodal.py @@ -3,7 +3,7 @@ import unittest from contextlib import nullcontext from types import SimpleNamespace -from unittest.mock import patch +from unittest.mock import Mock, patch import torch @@ -20,6 +20,105 @@ class TestDeepseekV41BCGMultimodal(CustomTestCase): + def test_cp_engram_graph_writes_global_hash_buffer_before_sharding(self): + class ReachedEngram(Exception): + pass + + for cp, rank, local_rows in ((2, 0, 4), (2, 1, 4), (8, 0, 1), (8, 7, 1)): + with self.subTest(cp=cp, rank=rank): + ids = torch.arange(10, 16) + output = torch.empty((6, 1), dtype=torch.long) + hasher = Mock(side_effect=lambda ids, batch: ids[:, None]) + hasher.new_hash_ids.return_value = output + expected = torch.zeros((local_rows, 1), dtype=torch.long) + live = ids[rank::cp] + expected[: live.numel(), 0] = live + + def read_engram(residual, hashes, batch, **kwargs): + torch.testing.assert_close(hashes, expected[:, 0]) + self.assertTrue(kwargs["cp_all_tokens"]) + raise ReachedEngram + + engram = Mock(side_effect=read_engram) + engram.layer_hash_index = 0 + model = SimpleNamespace( + pp_group=SimpleNamespace(world_size=1), + engram_hasher=hasher, + late_layer_start=None, + decoder_replay_tail=None, + start_layer=0, + end_layer=1, + layers=[SimpleNamespace(engram=engram)], + ) + batch = SimpleNamespace( + forward_mode=ForwardMode.EXTEND, + input_ids=ids, + attn_cp_metadata=SimpleNamespace(total_seq_lens=6), + ) + with ( + patch.object(dsv4, "is_cp_active", return_value=True), + patch.object(dsv4, "is_in_breakable_cuda_graph", return_value=True), + patch.object(dsv4, "_is_cuda", False), + patch.object(dsv4, "_is_hip", False), + patch.object( + dsv4, + "get_parallel", + return_value=SimpleNamespace( + attn_cp_size=cp, attn_cp_rank=rank, attn_dp_size=1 + ), + ), + patch.object( + dsv4, + "get_forward", + return_value=SimpleNamespace(sp_active=False), + ), + patch.object( + dsv4, + "get_tc_piecewise_forward_context", + return_value=SimpleNamespace(forward_batch=batch), + ), + patch.object( + dsv4, + "bcg_deepseek_v4_engram_hash_ids", + side_effect=dsv4.deepseek_v4_engram_hash_ids, + ) as graph_hash, + self.assertRaises(ReachedEngram), + ): + dsv4.DeepseekV4Model._forward_layers_hc_pre_from_prev( + model, + torch.arange(local_rows), + torch.zeros(local_rows, 1, 2), + batch, + live, + ids, + False, + None, + ) + self.assertEqual(graph_hash.call_count, 1) + self.assertIs(graph_hash.call_args.args[2], output) + torch.testing.assert_close(output[:, 0], ids) + hasher.new_hash_ids.assert_called_once() + hasher.assert_called_once() + + def test_cp_low_ratio_sources_keep_an_eager_break_with_graph_backend(self): + batch = SimpleNamespace(forward_mode=ForwardMode.EXTEND) + backend = SimpleNamespace(low_ratio_prefill_graph=True) + for graph, cp, tail, expected in ( + (True, True, False, True), + (True, False, False, False), + (True, False, True, True), + (False, True, False, False), + ): + with ( + self.subTest(graph=graph, cp=cp, tail=tail), + patch.object(dsv4, "is_in_breakable_cuda_graph", return_value=graph), + patch.object(dsv4, "is_cp_active", return_value=cp), + patch.object(dsv4, "in_replay_tail_graph", return_value=tail), + ): + self.assertEqual( + dsv4._low_ratio_sources_break(backend, batch), expected + ) + def test_replay_uses_prepared_image_ids_without_mutating_scheduler_ids(self): for keyword in (False, True): with self.subTest(keyword=keyword): @@ -85,9 +184,14 @@ def __call__(self, residual, *args, input_ids, **kwargs): residual, kv, weight, weight, 1e-6, 1e-6, input_ids, image_id ) - layer = SimpleNamespace( + class Layer(SimpleNamespace): + def __call__(self, *, hidden_states, **kwargs): + return hidden_states + + layer = Layer( engram=Engram(), hc_cfg=None, + boundary_mhc_fusions=False, forward_hc_pre_from_prev=lambda **kw: kw["state"], ) model = SimpleNamespace( diff --git a/test/registered/unit/models/test_deepseek_v41_multimodal_cp.py b/test/registered/unit/models/test_deepseek_v41_multimodal_cp.py new file mode 100644 index 000000000000..eb34d015ef75 --- /dev/null +++ b/test/registered/unit/models/test_deepseek_v41_multimodal_cp.py @@ -0,0 +1,559 @@ +"""V4.1 image/text CP input contracts; vision and model compute are mocked.""" + +import unittest +from contextlib import ExitStack, contextmanager, nullcontext +from types import SimpleNamespace as NS +from unittest.mock import Mock, patch + +import torch + +from sglang.srt.layers.cp.bcg.dsv4_bcg import DSV4CPBCGInput +from sglang.srt.layers.cp.interleave import InterleaveCPStrategy +from sglang.srt.layers.cp.padding import pad_logical_token_to_physical +from sglang.srt.layers.cp.zigzag import ZigzagContextParallelMetadata +from sglang.srt.managers.req_types import MM_PAD_SHIFT_VALUE, MultimodalInputs +from sglang.srt.model_executor.forward_mode import ForwardMode +from sglang.srt.model_executor.model_interfaces import LanguageModelInputPreparation +from sglang.srt.model_executor.runner.eager_runner import EagerRunner +from sglang.srt.model_executor.runner.prefill_cuda_graph_runner import ( + PrefillCudaGraphRunner, +) +from sglang.srt.models.deepseek_v4 import ( + DeepseekV4ForCausalLM, + _v41_vision_a2a_supported, +) +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + + +MODEL = "sglang.srt.models.deepseek_v4" +RUNNER = "sglang.srt.model_executor.runner.eager_runner" +IMAGE_ID = 129264 +BCG = "sglang.srt.layers.cp.bcg.dsv4_bcg" +PREFILL_GRAPH = "sglang.srt.model_executor.runner.prefill_cuda_graph_runner" + + +CP = "sglang.srt.layers.cp" + + +@contextmanager +def cp_context(size, rank, lengths=(3, 6), prefix_lengths=(7, 13)): + """Keep real interleave indexing/padding; replace only runtime context.""" + strategy = InterleaveCPStrategy(size) + parallel = NS(attn_cp_size=size, attn_cp_rank=rank, attn_cp_group=None) + batch = NS( + forward_mode=ForwardMode.EXTEND, + input_ids=torch.arange(1, sum(lengths) + 1), + positions=torch.cat( + [ + torch.arange(prefix, prefix + length) + for prefix, length in zip(prefix_lengths, lengths) + ] + ), + extend_seq_lens_cpu=list(lengths), + extend_prefix_lens_cpu=list(prefix_lengths), + mm_inputs=None, + spec_info=None, + ) + batch.attn_cp_metadata = strategy.build_metadata( + sum(lengths), [p + n for p, n in zip(prefix_lengths, lengths)], list(lengths) + ) + with ExitStack() as stack: + for module in ("base", "utils", "padding", "interleave"): + stack.enter_context( + patch(CP + "." + module + ".get_parallel", return_value=parallel) + ) + stack.enter_context(patch(CP + ".utils.get_cp_strategy", return_value=strategy)) + stack.enter_context( + patch(CP + ".padding.get_cp_padding_align_size", return_value=size) + ) + stack.enter_context( + patch( + CP + ".utils.get_moe_a2a_backend", return_value=NS(is_none=lambda: True) + ) + ) + pad_logical_token_to_physical(batch.attn_cp_metadata) + yield strategy, batch + + +@contextmanager +def simulated_collective(strategy, batch, global_tensor): + """Inject peer buffers into all-gather; retain production unpadding/reordering.""" + physical = max(batch.attn_cp_metadata.per_rank_actual_token) + buffers = [] + for rank in range(strategy.cp_size): + buf = global_tensor.new_zeros((physical, *global_tensor.shape[1:])) + local = global_tensor[rank :: strategy.cp_size] + buf[: len(local)] = local + buffers.append(buf) + + def gather(output, local): + torch.testing.assert_close(local, buffers[strategy.cp_rank], rtol=0, atol=0) + output.copy_(torch.cat(buffers)) + + with ( + patch( + CP + ".interleave.use_symmetric_memory", + side_effect=lambda *a, **k: nullcontext(), + ), + patch(CP + ".interleave.is_allocation_symmetric", return_value=False), + patch(CP + ".interleave.attn_cp_all_gather_into_tensor", side_effect=gather), + ): + yield + + +class _PreparedModel(NS, LanguageModelInputPreparation): + pass + + +class TestDSV41MultimodalCP(CustomTestCase): + def test_cp_bcg_defers_image_batches_to_eager_vision_path(self): + runner = NS( + max_context_size=None, + max_seq_len=None, + enable_lora=False, + enable_cp_bcg_capture=True, + model_runner=NS(model=NS(vision=object())), + can_replay_locally=Mock(return_value=True), + _has_inactive_dp_rank=Mock(return_value=False), + prefill_cp_bcg_input=NS(select_replay_bucket_for_batch=lambda **_: 8), + capture_num_tokens=[8], + _backend_can_run_prefill_cuda_graph=None, + ) + batch = NS( + global_num_tokens_cpu=None, + batch_size=1, + input_ids=torch.arange(8), + input_embeds=None, + replace_embeds=None, + extend_prefix_lens_cpu=[0], + extend_seq_lens_cpu=[8], + forward_mode=ForwardMode.EXTEND, + global_forward_mode=ForwardMode.EXTEND, + capture_hidden_mode=None, + return_logprob=False, + mm_inputs=[MultimodalInputs(mm_items=[])], + ) + batch.contains_mm_inputs = lambda: any(x is not None for x in batch.mm_inputs) + with patch(PREFILL_GRAPH + ".is_cp_active", return_value=True): + self.assertFalse(PrefillCudaGraphRunner.can_run_graph(runner, batch)) + batch.mm_inputs = [None] + self.assertTrue(PrefillCudaGraphRunner.can_run_graph(runner, batch)) + + def test_vision_megamoe_prefill_is_admitted(self): + for backend_name, expected in ( + ("none", True), + ("megamoe", True), + ("deepep", False), + ): + with self.subTest(backend=backend_name): + backend = NS( + is_none=lambda: backend_name == "none", + is_megamoe=lambda: backend_name == "megamoe", + ) + with ( + patch(MODEL + ".get_moe_a2a_backend", return_value=backend), + patch(MODEL + "._is_cuda", True), + ): + self.assertEqual(_v41_vision_a2a_supported(), expected) + + def test_bcg_prepare_rejects_non_interleave_metadata(self): + cp_input = DSV4CPBCGInput( + input_embeds=torch.empty(4, 1), + positions=torch.empty(4, dtype=torch.int64), + bucket_local_tokens={8: 4}, + ) + for capture in (True, False): + for metadata in ( + None, + NS(per_rank_actual_token=[4, 4]), + ZigzagContextParallelMetadata(per_rank_actual_token=[4, 4]), + ): + with self.subTest(capture=capture, metadata=type(metadata).__name__): + batch = NS(out_cache_loc=torch.arange(8)) + + def prepare_metadata(forward_batch): + forward_batch.attn_cp_metadata = metadata + + with ( + patch( + BCG + ".prepare_cp_forward", side_effect=prepare_metadata + ), + patch(BCG + ".cp_shard_hidden_states") as shard, + self.assertRaisesRegex( + TypeError, + "DSV4 CP BCG requires InterleaveContextParallelMetadata, " + f"got {type(metadata).__name__}", + ), + ): + cp_input.prepare( + NS(), batch, static_num_tokens=8, capture=capture + ) + shard.assert_not_called() + + def test_bcg_router_ids_follow_moe_row_layout_across_replays(self): + for backend_name in ("none", "megamoe"): + with ( + self.subTest(backend=backend_name), + cp_context(2, 1) as ( + strategy, + batch, + ), + ): + backend = NS(is_none=lambda: backend_name == "none") + model = NS( + get_input_embeddings=lambda: lambda ids: ids[:, None].float() + ) + runner = NS( + max_num_tokens=16, + device="cpu", + model_runner=NS( + model=model, + model_config=NS(hidden_size=1), + dtype=torch.float32, + ), + ) + + def prepare_metadata(fb): + fb.attn_cp_metadata = strategy.build_metadata( + len(fb.input_ids), None, fb.extend_seq_lens_cpu + ) + pad_logical_token_to_physical(fb.attn_cp_metadata) + + batch.out_cache_loc = torch.arange(len(batch.input_ids)) + batch.extend_num_tokens = len(batch.input_ids) + batch.num_token_non_padded = None + with ( + patch(BCG + ".get_cp_strategy", return_value=strategy), + patch(BCG + ".get_moe_a2a_backend", return_value=backend), + patch(BCG + ".prepare_cp_forward", side_effect=prepare_metadata), + ): + cp_input = DSV4CPBCGInput.create(runner) + cp_input.prepare(runner, batch, static_num_tokens=9, capture=True) + captured_ids = batch.input_ids_global + captured_positions = batch.cp_positions + captured_local_ids = batch.cp_input_ids + self.assertEqual( + captured_positions.data_ptr(), cp_input.positions.data_ptr() + ) + self.assertEqual( + captured_local_ids.data_ptr(), cp_input.input_ids.data_ptr() + ) + self.assertEqual( + captured_ids.shape[0], 12 if backend_name == "none" else 6 + ) + + batch.input_ids = torch.tensor([7, IMAGE_ID, 9, 10, 11, 12, 13]) + batch.positions = torch.arange(7) + batch.extend_seq_lens_cpu = [3, 4] + batch.extend_num_tokens = 7 + cp_input.prepare(runner, batch, static_num_tokens=9, capture=False) + self.assertEqual( + batch.input_ids_global.data_ptr(), captured_ids.data_ptr() + ) + self.assertEqual( + batch.cp_positions.data_ptr(), captured_positions.data_ptr() + ) + self.assertEqual( + batch.cp_input_ids.data_ptr(), captured_local_ids.data_ptr() + ) + if backend_name == "none": + expected = torch.tensor( + [7, 9, 11, 13, 0, 0, IMAGE_ID, 10, 12, 0, 0, 0] + ) + else: + expected = torch.tensor([IMAGE_ID, 10, 12, 0, 0, 0]) + torch.testing.assert_close(batch.input_ids_global, expected) + + def test_image_spans_cross_ranks_before_shard_and_gather(self): + # Two image spans with distinct cache hashes; first request is text-only. + original = torch.tensor( + [ + 7, + 8, + 9, + MM_PAD_SHIFT_VALUE + 11, + MM_PAD_SHIFT_VALUE + 11, + 10, + MM_PAD_SHIFT_VALUE + 23, + MM_PAD_SHIFT_VALUE + 23, + 12, + ] + ) + normalized = torch.tensor( + [7, 8, 9, IMAGE_ID, IMAGE_ID, 10, IMAGE_ID, IMAGE_ID, 12] + ) + full = torch.arange(36, dtype=torch.float32).reshape(9, 4) + # Distinct image features expose using text embeddings or wrong row order. + full[3:5] += 1000 + full[6:8] += 2000 + for size in (2, 4): + for rank in range(size): + with ( + self.subTest(size=size, rank=rank), + cp_context(size, rank) as (strategy, batch), + ): + batch.input_ids = original.clone() + batch.mm_inputs = [None, MultimodalInputs(mm_items=[])] + model = _PreparedModel( + vision=object(), + config=NS(image_token_id=IMAGE_ID), + get_input_embeddings=Mock( + side_effect=AssertionError( + "Raw image hashes entered text embeddings" + ) + ), + _prepare_mm_embeddings=Mock(return_value=full), + capture_aux_hidden_states=False, + pp_group=NS(is_last_rank=True), + lm_head=object(), + logits_processor=Mock(return_value="ok"), + ) + model.prepare_language_model_inputs = lambda ids, fb, emb: ( + DeepseekV4ForCausalLM.prepare_language_model_inputs( + model, ids, fb, emb + ) + ) + + def body(ids, positions, fb, input_embeds): + model._prepare_mm_embeddings.assert_called_once_with( + batch.input_ids, batch + ) + n = len(normalized[rank::size]) + torch.testing.assert_close(ids[:n], normalized[rank::size]) + torch.testing.assert_close(input_embeds[:n], full[rank::size]) + torch.testing.assert_close( + positions[:n], batch.positions[rank::size] + ) + self.assertFalse( + (fb.input_ids_global >= MM_PAD_SHIFT_VALUE).any().item() + ) + return input_embeds + + model.model = body + with ( + simulated_collective(strategy, batch, full), + patch(RUNNER + ".torch.cuda.current_stream", return_value=None), + ): + result = EagerRunner._execute_extend_cp( + NS(model_runner=NS(model=model)), batch, {} + ) + self.assertEqual(result, "ok") + torch.testing.assert_close( + model.logits_processor.call_args.args[0], normalized + ) + torch.testing.assert_close( + model.logits_processor.call_args.args[1], full + ) + torch.testing.assert_close(batch.input_ids, original) + self.assertNotIn("input_ids_global", vars(batch)) + + def test_chunk_prefix_metadata_and_scheduler_hashes_survive_embedder(self): + for prefixes, lengths in (([0, 0], [3, 6]), ([16384, 127], [3, 6])): + with self.subTest(prefixes=prefixes): + ids = torch.tensor([7, 8, 9] + [MM_PAD_SHIFT_VALUE + 17] * 6) + original = ids.clone() + image = MultimodalInputs(mm_items=[]) + batch = NS( + mm_inputs=[None, image], + extend_prefix_lens_cpu=prefixes, + extend_seq_lens_cpu=lengths, + ) + full = torch.arange(27, dtype=torch.float32).reshape(9, 3) + embedding = object() + model = NS(get_input_embeddings=lambda: embedding) + + def embed(**kwargs): + self.assertEqual(kwargs["extend_prefix_lens"], prefixes) + self.assertEqual(kwargs["extend_seq_lens"], lengths) + self.assertIs(kwargs["mm_inputs_list"][1], image) + self.assertEqual(kwargs["mm_inputs_list"][0].mm_items, []) + self.assertIs(kwargs["input_embedding"], embedding) + self.assertNotEqual(kwargs["input_ids"].data_ptr(), ids.data_ptr()) + kwargs["input_ids"].zero_() + return full, {} + + with patch(MODEL + ".embed_mm_inputs", side_effect=embed) as mocked: + result = DeepseekV4ForCausalLM._prepare_mm_embeddings( + model, ids, batch + ) + mocked.assert_called_once() + self.assertIs(result, full) + self.assertIs(batch.mm_input_embeds, full) + torch.testing.assert_close(ids, original) + + def test_vision_enabled_text_batch_skips_image_encoder(self): + for mm_inputs in (None, [None, None], []): + with self.subTest(mm_inputs=mm_inputs): + ids = torch.tensor([4, 5, 6]) + model = NS( + vision=object(), + config=NS(image_token_id=IMAGE_ID), + _prepare_mm_embeddings=Mock(), + ) + batch = NS(forward_mode=ForwardMode.EXTEND, mm_inputs=mm_inputs) + result, embeds = DeepseekV4ForCausalLM.prepare_language_model_inputs( + model, ids, batch + ) + torch.testing.assert_close(result, ids) + self.assertIsNone(embeds) + model._prepare_mm_embeddings.assert_not_called() + + def test_decode_idle_and_verify_preserve_vocab_ids(self): + for mode in (ForwardMode.DECODE, ForwardMode.IDLE, ForwardMode.TARGET_VERIFY): + with self.subTest(mode=mode): + ids = torch.tensor([4, IMAGE_ID, 6]) + model = NS( + vision=object(), + config=NS(image_token_id=IMAGE_ID), + _prepare_mm_embeddings=Mock(), + ) + batch = NS(forward_mode=mode, mm_inputs=None) + result, embeds = DeepseekV4ForCausalLM.prepare_language_model_inputs( + model, ids, batch + ) + self.assertIs(result, ids) + self.assertIsNone(embeds) + model._prepare_mm_embeddings.assert_not_called() + + def test_image_embedding_failure_does_not_mutate_scheduler_ids(self): + ids = torch.tensor([7, MM_PAD_SHIFT_VALUE + 12, 8]) + original = ids.clone() + batch = NS( + mm_inputs=[MultimodalInputs(mm_items=[])], + extend_prefix_lens_cpu=[0], + extend_seq_lens_cpu=[3], + ) + model = NS(get_input_embeddings=lambda: object()) + + def fail(**kwargs): + kwargs["input_ids"].zero_() + raise RuntimeError("vision failure") + + with patch(MODEL + ".embed_mm_inputs", side_effect=fail): + with self.assertRaisesRegex(RuntimeError, "vision failure"): + DeepseekV4ForCausalLM._prepare_mm_embeddings(model, ids, batch) + torch.testing.assert_close(ids, original) + self.assertNotIn("mm_input_embeds", vars(batch)) + + +class TestDSV41MultimodalInputs(CustomTestCase): + def setUp(self): + self.ids = torch.tensor( + [7, MM_PAD_SHIFT_VALUE + 12, MM_PAD_SHIFT_VALUE + 12, 9, 10] + ) + self.original = self.ids.clone() + self.embeds = torch.arange(15, dtype=torch.float32).reshape(5, 3) + self.model = _PreparedModel( + vision=object(), + config=NS(image_token_id=129264), + _prepare_mm_embeddings=Mock(return_value=self.embeds), + ) + self.batch = NS( + input_ids=self.ids, + forward_mode=ForwardMode.EXTEND, + mm_inputs=[MultimodalInputs(mm_items=[])], + ) + + def test_prepare_global_embeddings_and_normalized_ids(self): + ids, embeds = DeepseekV4ForCausalLM.prepare_language_model_inputs( + self.model, self.ids, self.batch + ) + self.assertEqual(ids.tolist(), [7, 129264, 129264, 9, 10]) + self.assertIs(embeds, self.embeds) + self.model._prepare_mm_embeddings.assert_called_once_with(self.ids, self.batch) + self.assertTrue(torch.equal(self.ids, self.original)) + + def test_reject_preembedded_images(self): + with self.assertRaisesRegex(ValueError, "Cannot combine"): + DeepseekV4ForCausalLM.prepare_language_model_inputs( + self.model, self.ids, self.batch, self.embeds + ) + + def test_text_only_model_keeps_existing_embeddings(self): + self.model.vision = None + ids, embeds = DeepseekV4ForCausalLM.prepare_language_model_inputs( + self.model, self.ids, self.batch, self.embeds + ) + self.assertIs(ids, self.ids) + self.assertIs(embeds, self.embeds) + self.model._prepare_mm_embeddings.assert_not_called() + + def test_text_subclass_without_vision_module(self): + del self.model.vision + ids, embeds = DeepseekV4ForCausalLM.prepare_language_model_inputs( + self.model, self.ids, self.batch, self.embeds + ) + self.assertIs(ids, self.ids) + self.assertIs(embeds, self.embeds) + + def test_decode_keeps_vocabulary_ids(self): + self.batch.forward_mode = ForwardMode.DECODE + ids, embeds = DeepseekV4ForCausalLM.prepare_language_model_inputs( + self.model, self.ids, self.batch + ) + self.assertIs(ids, self.ids) + self.assertIsNone(embeds) + self.model._prepare_mm_embeddings.assert_not_called() + + def test_embedding_does_not_mutate_scheduler_hashes(self): + self.batch.extend_prefix_lens_cpu = [0] + self.batch.extend_seq_lens_cpu = [5] + self.model.get_input_embeddings = lambda: None + + def embed(**kwargs): + kwargs["input_ids"].zero_() + return (self.embeds, {}) + + with patch("sglang.srt.models.deepseek_v4.embed_mm_inputs", side_effect=embed): + result = DeepseekV4ForCausalLM._prepare_mm_embeddings( + self.model, self.ids, self.batch + ) + self.assertTrue(torch.equal(self.ids, self.original)) + self.assertIs(result, self.batch.mm_input_embeds) + + def test_cp_runner_prepares_before_sharding_and_uses_model_ids_for_logits(self): + normalized = torch.tensor([7, 129264, 129264, 9, 10]) + self.batch.positions = torch.arange(5) + self.model.prepare_language_model_inputs = lambda ids, batch, emb: ( + DeepseekV4ForCausalLM.prepare_language_model_inputs( + self.model, ids, batch, emb + ) + ) + self.model.get_input_embeddings = Mock( + side_effect=AssertionError("Raw hashes must not enter text embedding") + ) + self.model.model = Mock(return_value=self.embeds[1::4]) + self.model.capture_aux_hidden_states = False + self.model.pp_group = NS(is_last_rank=True) + self.model.lm_head = object() + self.model.logits_processor = Mock(return_value="ok") + + @contextmanager + def shard(embeds, positions, batch, ids): + self.assertIs(embeds, self.embeds) + self.assertTrue(torch.equal(ids, normalized)) + yield (embeds[1::4], positions[1::4], ids[1::4]) + + runner = NS(model_runner=NS(model=self.model)) + module = "sglang.srt.model_executor.runner.eager_runner" + with ( + patch(module + ".cp_shard_model_inputs", side_effect=shard), + patch(module + ".cp_gather_after_forward", return_value=self.embeds), + patch(module + ".torch.cuda.current_stream", return_value=None), + ): + result = EagerRunner._execute_extend_cp(runner, self.batch, {}) + self.assertEqual(result, "ok") + args, kwargs = self.model.model.call_args + self.assertEqual(args[0].tolist(), [129264]) + self.assertTrue(torch.equal(kwargs["input_embeds"], self.embeds[1::4])) + self.assertTrue( + torch.equal(self.model.logits_processor.call_args.args[0], normalized) + ) + self.assertTrue(torch.equal(self.batch.input_ids, self.original)) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/models/test_deepseek_v4_shared_expert_fusion.py b/test/registered/unit/models/test_deepseek_v4_shared_expert_fusion.py index 06bf5b06923f..45385acd974d 100644 --- a/test/registered/unit/models/test_deepseek_v4_shared_expert_fusion.py +++ b/test/registered/unit/models/test_deepseek_v4_shared_expert_fusion.py @@ -85,8 +85,8 @@ def test_disables_shared_fusion_without_enforce(self): def test_v41_vision_admits_cuda_megamoe_with_dp_attention(self): """A vision-capable checkpoint must initialize for text or image serving - with DP attention and CUDA MegaMoE. CP, PP and other A2A paths still lack - vision routing support and must fail before constructing the tower. + with DP attention and prefill CP, using no A2A or CUDA MegaMoE. + PP and unsupported A2A paths must fail before constructing the tower. """ self._publish(enforce=False) config = DeepSeekV4Config( @@ -112,7 +112,7 @@ def test_v41_vision_admits_cuda_megamoe_with_dp_attention(self): (MoeA2ABackend.NONE, True, 1, 1, True), (MoeA2ABackend.MEGAMOE, False, 1, 1, False), (MoeA2ABackend.DEEPEP, True, 1, 1, False), - (MoeA2ABackend.MEGAMOE, True, 2, 1, False), + (MoeA2ABackend.DEEPEP, True, 2, 1, False), (MoeA2ABackend.MEGAMOE, True, 1, 2, False), ) for backend, cuda, cp, pp, admitted in cases: