qwen 3.8 rebase (#35758)

Co-authored-by: cherichy <cherichy@outlook.com>
Co-authored-by: guangyunh-nv <guangyunh@nvidia.com>
Co-authored-by: jiahanc <jiahanc@nvidia.com>
Co-authored-by: jinyangyuan-nvidia <joyuan@nvidia.com>
Co-authored-by: Cheng Hang <chang@nvidia.com>
Co-authored-by: Yicheng Qiang <yqiang@nvidia.com>
Co-authored-by: Sam Li <lsam@nvidia.com>
Co-authored-by: Tom-Zheng <tizheng@nvidia.com>
Co-authored-by: Yangmin Li <yangminl@nvidia.com>
Co-authored-by: xiaoweiw-nv <xiaoweiw@nvidia.com>
Co-authored-by: Zheng Li <lizheng.cs@zju.edu.cn>
Co-authored-by: yizhang2077 <1109276519@qq.com>
Co-authored-by: Ke Bao <ispobaoke@gmail.com>
Co-authored-by: Xinyuan Tong <115166877+JustinTong0323@users.noreply.github.com>
Co-authored-by: Yuhao Yang <47235274+yhyang201@users.noreply.github.com>
Co-authored-by: Zijie Xia <zijie.xia@radixark.ai>
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

python/sglang/srt/model_executor/cuda_graph_buffer_registry.py

@@ -939,6 +939,8 @@ def build_prefill_registry(
                 )
         reg.register_slot(slot, bind=bind)
 
+    # PP stage inputs live outside ForwardBatch; adopt runner-owned buffers for
+    # stable addresses and clear padding because prefill executes every bucket row.
     if source is not None:
         pp = getattr(source, "pp_proxy_tensors", None)
         if pp is not None:
@@ -954,7 +956,10 @@ def build_prefill_registry(
                 reg.register_slot(
                     GraphSlot(
                         name=f"pp_proxy_tensors.{_key}",
-                        shape_fn=lambda _bs, _mt, _s=tuple(_backing.shape): _s,
+                        shape_fn=lambda _bs, mt, _tail=tuple(_backing.shape[1:]): (
+                            mt,
+                            *_tail,
+                        ),
                         dtype=_backing.dtype,
                         axis="tokens",
                         padding_policy=PaddingPolicy.ZERO,

python/sglang/srt/model_executor/model_runner_components/cuda_graph_setup.py

@@ -3,7 +3,7 @@ from __future__ import annotations
 import logging
 import time
 from collections import defaultdict
-from typing import TYPE_CHECKING, Optional
+from typing import TYPE_CHECKING, Any, Optional
 
 import msgspec
 
@@ -80,6 +80,45 @@ def _align_pipeline_layers(layers: list, layer_model) -> list:
     )
 
 
+def has_standard_gqa_for_all_local_layers(
+    *, attention_layer_count: int, start_layer: int, end_layer: int
+) -> bool:
+    """Check the layers materialized on this pipeline rank, not the full model."""
+    return attention_layer_count >= end_layer - start_layer
+
+
+def index_attention_layers_by_global_id(
+    attention_layers: list[Any],
+    mha_companion_layers: list[Any],
+    layer_model=None,
+) -> tuple[list[Any], list[Any]]:
+    """Pad PP-local attention metadata so global layer_id remains a valid index."""
+    if len(attention_layers) != len(mha_companion_layers):
+        raise ValueError("attention and MHA companion metadata must be parallel")
+    populated = [layer for layer in attention_layers if layer is not None]
+    if not populated or any(not hasattr(layer, "layer_id") for layer in populated):
+        if layer_model is not None:
+            return (
+                _align_pipeline_layers(attention_layers, layer_model),
+                _align_pipeline_layers(mha_companion_layers, layer_model),
+            )
+        return attention_layers, mha_companion_layers
+    max_layer_id = max(int(layer.layer_id) for layer in populated)
+    indexed_attention = [None] * (max_layer_id + 1)
+    indexed_companions = [None] * (max_layer_id + 1)
+    for attention, companion in zip(attention_layers, mha_companion_layers):
+        if attention is None:
+            if companion is not None:
+                raise ValueError("MHA companion has no primary attention layer")
+            continue
+        layer_id = int(attention.layer_id)
+        if layer_id < 0 or indexed_attention[layer_id] is not None:
+            raise ValueError(f"invalid or duplicate attention layer_id: {layer_id}")
+        indexed_attention[layer_id] = attention
+        indexed_companions[layer_id] = companion
+    return indexed_attention, indexed_companions
+
+
 class GraphCapture(msgspec.Struct, frozen=True, kw_only=True):
     runner: Optional[BaseRunner]
     memory_phase: str
@@ -392,11 +431,22 @@ def capture_prefill_graph(
         model_runner.dsa_indexers,
         model_runner.mha_companion_layers,
     ) = compute_attention_and_moe_layers(layer_model)
-
-    model_runner.attention_layers = _align_pipeline_layers(
-        model_runner.attention_layers, layer_model
+    (
+        model_runner.attention_layers,
+        model_runner.mha_companion_layers,
+    ) = index_attention_layers_by_global_id(
+        model_runner.attention_layers,
+        model_runner.mha_companion_layers,
+        layer_model,
     )
-    if len(model_runner.attention_layers) < model_runner.model_config.num_hidden_layers:
+
+    if not has_standard_gqa_for_all_local_layers(
+        attention_layer_count=sum(
+            layer is not None for layer in model_runner.attention_layers
+        ),
+        start_layer=model_runner.layer_info.start_layer,
+        end_layer=model_runner.layer_info.end_layer,
+    ):
         # TODO(yuwei): support Non-Standard GQA
         log_info_on_rank0(
             logger,
@@ -404,10 +454,6 @@ def capture_prefill_graph(
         )
         return result(None)
 
-    model_runner.mha_companion_layers = _align_pipeline_layers(
-        model_runner.mha_companion_layers, layer_model
-    )
-
     tic = time.perf_counter()
     before_mem = get_available_gpu_memory(model_runner.device, model_runner.gpu_id)
     role = "draft" if model_runner.is_draft_worker else "target"

python/sglang/srt/model_executor/runner/base_runner.py

@@ -134,7 +134,7 @@ def _allocate_decode_buffers(
                 residual_shape = (
                     (max_num_token, pp_proxy_residual_num_blocks, hidden_size)
                     if pp_proxy_residual_num_blocks is not None
-                    else (max_bs, hidden_size)
+                    else (max_num_token, hidden_size)
                 )
                 pp_proxy_tensors["residual"] = torch.zeros(residual_shape, dtype=dtype)
             if pp_proxy_topk_size is not None:
@@ -244,6 +244,14 @@ class BaseRunner(ABC):
         self._pre_initialize_flashinfer_allreduce_workspace()
         self._pre_initialize_fi_a2a_workspace()
 
+        # Model-owned communication resources may depend on the resolved
+        # request pool and must be compiled/allocated before graph capture.
+        prepare_model_resources = getattr(
+            mr.model, "prepare_before_cuda_graph_capture", None
+        )
+        if prepare_model_resources is not None:
+            prepare_model_resources(mr)
+
         if should_run_flashinfer_autotune(self.model_runner):
             buffers, batch_size = self._autotune_buffers()
             assert (

python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py

@@ -38,6 +38,7 @@ Backend selection comes from cuda_graph_config.prefill:
 from __future__ import annotations
 
 import copy
+import dataclasses
 import inspect
 import logging
 from collections.abc import Sequence
@@ -328,10 +329,12 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
             dtype=self.model_runner.dtype,
             enable_mamba_track=self.mamba_track_enabled,
             pp_size=self.pp_size,
-            hc_hidden_size=model_runner.model_config.hc_hidden_size,
-            pp_proxy_topk_size=model_runner.get_pp_proxy_topk_size(),
+            hc_hidden_size=getattr(
+                self.model_runner.model_config, "hc_hidden_size", None
+            ),
+            pp_proxy_topk_size=self.model_runner.get_pp_proxy_topk_size(),
             pp_proxy_residual_num_blocks=(
-                model_runner.get_pp_proxy_residual_num_blocks()
+                self.model_runner.get_pp_proxy_residual_num_blocks()
             ),
         )
         self.buffers.share_buffers()
@@ -385,11 +388,6 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
         self._capture_lora = False
         self.enable_cp_v2_bcg_capture = False
         self.prefill_cp_bcg_input: Optional[PrefillCPBCGInput] = None
-        self._static_pp_proxy_tensors = (
-            PPProxyTensors(self.buffers.pp_proxy_tensors)
-            if self.buffers.pp_proxy_tensors is not None
-            else None
-        )
         # TcPiecewise does its compile pass during backend construction.
         # Wrap only that path with the prefill CUDA graph failure hint.
         try:
@@ -646,6 +644,14 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
 
         return forward_batch.positions
 
+    def _static_pp_proxy_tensors(self, num_tokens: int) -> Optional[PPProxyTensors]:
+        buffers = self.buffers.pp_proxy_tensors
+        if buffers is None:
+            return None
+        return PPProxyTensors(
+            {key: value[:num_tokens] for key, value in buffers.items()}
+        )
+
     @contextmanager
     def _prefill_forward_context(
         self,
@@ -702,31 +708,25 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
         set_is_extend_in_batch(False)
 
         with self._prefill_forward_context(forward_batch):
+            pp_kwargs = self.model_runner._pp_kwargs(
+                self._static_pp_proxy_tensors(num_tokens)
+            )
             if self._uses_eager_prefill_tail():
                 # BCG / Full: capture the transformer body only.
                 positions = self._get_layer_model_positions(forward_batch)
-                input_ids = forward_batch.input_ids
-                input_embeds = forward_batch.input_embeds
-                layer_kwargs = {}
-                if self._static_pp_proxy_tensors is not None:
-                    layer_kwargs["pp_proxy_tensors"] = self._static_pp_proxy_tensors[
-                        :num_tokens
-                    ]
-                    if not self.model_runner.pp_group.is_first_rank:
-                        input_ids = None
-                        input_embeds = None
                 return self.layer_model.forward(
-                    input_ids,
+                    forward_batch.input_ids,
                     positions,
                     forward_batch,
-                    input_embeds,
-                    **layer_kwargs,
+                    forward_batch.input_embeds,
+                    **pp_kwargs,
                 )
             # tc_piecewise: compile/capture the outer model.forward path.
             return self.model_runner.model.forward(
                 forward_batch.input_ids,
                 forward_batch.positions,
                 forward_batch,
+                **pp_kwargs,
             )
 
     def _run_dummy_forward(self, num_tokens: int) -> None:
@@ -1568,6 +1568,19 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
             else forward_batch.global_forward_mode
         )
 
+        # The draft tail concatenates hidden states with padded input embeddings;
+        # expose the bucket-sized static view and refresh its live prefix below.
+        padded_spec_info = forward_batch.spec_info
+        if (
+            self.static_draft_hidden_states is not None
+            and padded_spec_info is not None
+            and getattr(padded_spec_info, "hidden_states", None) is not None
+        ):
+            padded_spec_info = dataclasses.replace(
+                padded_spec_info,
+                hidden_states=self.static_draft_hidden_states[:static_num_tokens],
+            )
+
         static_forward_batch = ForwardBatch(
             forward_mode=pcg_forward_mode,
             batch_size=bs,
@@ -1611,7 +1624,7 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
             global_dp_buffer_len=forward_batch.global_dp_buffer_len,
             mrope_positions=mrope_positions,
             spec_algorithm=forward_batch.spec_algorithm,
-            spec_info=forward_batch.spec_info,
+            spec_info=padded_spec_info,
             capture_hidden_mode=forward_batch.capture_hidden_mode,
             num_token_non_padded=num_token_non_padded,
             num_token_non_padded_cpu=forward_batch.num_token_non_padded_cpu,
@@ -1619,6 +1632,9 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
             lora_ids=forward_batch.lora_ids,
             sampling_info=forward_batch.sampling_info,
             mm_inputs=forward_batch.mm_inputs,
+            # Multimodal preparation consumes mm_inputs but retains embeddings for
+            # later MTP draft extend, so copy that side channel into the replay view.
+            mm_input_embeds=forward_batch.mm_input_embeds,
             temperature=forward_batch.temperature,
             top_p=forward_batch.top_p,
             dimensions=forward_batch.dimensions,

python/sglang/srt/model_executor/runner_utils/buffers.py

@@ -133,7 +133,7 @@ class DecodeInputBuffers(ForwardInputBuffers):
                 is_mhc = hc_hidden_size is not None
                 hs = hc_hidden_size if is_mhc else hidden_size
                 pp_proxy_tensors = {
-                    "hidden_states": torch.zeros((max_bs, hs), dtype=dtype),
+                    "hidden_states": torch.zeros((max_num_token, hs), dtype=dtype),
                 }
                 if not is_mhc:
                     # Only Kimi K3 supplies num_blocks: its PP bank is token-major
@@ -141,7 +141,7 @@ class DecodeInputBuffers(ForwardInputBuffers):
                     residual_shape = (
                         (max_num_token, pp_proxy_residual_num_blocks, hidden_size)
                         if pp_proxy_residual_num_blocks is not None
-                        else (max_bs, hidden_size)
+                        else (max_num_token, hidden_size)
                     )
                     pp_proxy_tensors["residual"] = torch.zeros(
                         residual_shape, dtype=dtype
@@ -356,7 +356,7 @@ class PrefillInputBuffers(ForwardInputBuffers):
         hidden_size: int,
         dtype: torch.dtype,
         enable_mamba_track: bool,
-        pp_size: int,
+        pp_size: int = 1,
         hc_hidden_size: Optional[int] = None,
         pp_proxy_topk_size: Optional[int] = None,
         pp_proxy_residual_num_blocks: Optional[int] = None,