[Feat] [NVIDIA] deepseek-v4.1: Support prefill CP with BCG in PD serving (2/N)

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

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)

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

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