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>

test/registered/unit/model_executor/model_runner_components/test_cuda_graph_setup.py

@@ -7,12 +7,47 @@ from sglang.srt.model_executor.model_runner_components import cuda_graph_setup
 from sglang.srt.model_executor.model_runner_components.cuda_graph_setup import (
     _align_pipeline_layers,
     capture_decode_graph,
+    has_standard_gqa_for_all_local_layers,
+    index_attention_layers_by_global_id,
 )
 from sglang.test.ci.ci_register import register_cpu_ci
 
 register_cpu_ci(est_time=5, suite="base-a-test-cpu")
 
 
+def test_standard_gqa_gate_uses_pipeline_local_layer_range():
+    # PP rank owns layers [23, 46), while the full model has 92 layers.
+    assert has_standard_gqa_for_all_local_layers(
+        attention_layer_count=23, start_layer=23, end_layer=46
+    )
+    assert not has_standard_gqa_for_all_local_layers(
+        attention_layer_count=22, start_layer=23, end_layer=46
+    )
+
+
+def test_standard_gqa_gate_is_unchanged_without_pipeline_parallelism():
+    assert has_standard_gqa_for_all_local_layers(
+        attention_layer_count=92, start_layer=0, end_layer=92
+    )
+
+
+def test_pipeline_attention_metadata_is_indexed_by_global_layer_id():
+    layer23 = SimpleNamespace(layer_id=23)
+    layer24 = SimpleNamespace(layer_id=24)
+    companion24 = object()
+
+    attention, companions = index_attention_layers_by_global_id(
+        [layer23, layer24], [None, companion24]
+    )
+
+    assert len(attention) == 25
+    assert all(layer is None for layer in attention[:23])
+    assert attention[23] is layer23
+    assert attention[24] is layer24
+    assert companions[23] is None
+    assert companions[24] is companion24
+
+
 def test_model_runner_can_override_decode_graph_runner(monkeypatch):
     from sglang.srt.runtime_context import get_context
 

test/registered/unit/model_executor/test_cuda_graph_buffer_registry.py

@@ -1195,6 +1195,76 @@ class TestBuildPrefillRegistry(unittest.TestCase):
         reg.fill_from(fb, raw_bs=2, padded_bs=2, raw_num_tokens=3, padded_num_tokens=8)
         self.assertTrue(torch.equal(idx, torch.tensor([3, 4], dtype=torch.int64)))
 
+    def test_pp_proxy_token_slots_copy_head_and_zero_bucket_tail(self):
+        from sglang.srt.model_executor.cuda_graph_buffer_registry import (
+            build_prefill_registry,
+        )
+        from sglang.srt.model_executor.forward_batch_info import PPProxyTensors
+
+        hidden = torch.full((16, 4), 7.0)
+        residual = torch.full((16, 4), 7.0)
+        src = self._src(
+            pp_proxy_tensors={
+                "hidden_states": hidden,
+                "residual": residual,
+            }
+        )
+        reg = build_prefill_registry(
+            device=torch.device("cpu"),
+            max_bs=1,
+            max_num_token=16,
+            cache_loc_dtype=torch.int64,
+            source=src,
+        )
+        self.assertTrue(reg.has_slot("pp_proxy_tensors.hidden_states"))
+        fb = _MiniForwardBatch(
+            input_ids=torch.zeros(3, dtype=torch.int64),
+            positions=torch.zeros(3, dtype=torch.int64),
+            out_cache_loc=torch.zeros(3, dtype=torch.int64),
+        )
+        pp_proxy = PPProxyTensors(
+            {
+                "hidden_states": torch.ones((3, 4)),
+                "residual": torch.full((3, 4), 2.0),
+            }
+        )
+        reg.fill_from(
+            fb,
+            raw_bs=1,
+            padded_bs=1,
+            raw_num_tokens=3,
+            padded_num_tokens=8,
+            pp_proxy_tensors=pp_proxy,
+        )
+        self.assertTrue(torch.all(hidden[:3] == 1.0))
+        self.assertTrue(torch.all(residual[:3] == 2.0))
+        self.assertTrue(torch.all(hidden[3:8] == 0.0))
+        self.assertTrue(torch.all(residual[3:8] == 0.0))
+        self.assertTrue(torch.all(hidden[8:] == 7.0))
+
+    def test_prefill_input_buffers_allocate_pp_proxy_by_token(self):
+        from sglang.srt.model_executor.runner_utils.buffers import (
+            PrefillInputBuffers,
+        )
+
+        buffers = PrefillInputBuffers.create(
+            device=torch.device("cpu"),
+            max_bs=4,
+            max_num_tokens=16,
+            cache_loc_dtype=torch.int64,
+            is_multimodal=False,
+            hidden_size=8,
+            dtype=torch.bfloat16,
+            enable_mamba_track=False,
+            pp_size=2,
+            pp_proxy_topk_size=3,
+        )
+        self.assertEqual(
+            tuple(buffers.pp_proxy_tensors["hidden_states"].shape), (16, 8)
+        )
+        self.assertEqual(tuple(buffers.pp_proxy_tensors["residual"].shape), (16, 8))
+        self.assertEqual(tuple(buffers.pp_proxy_tensors["topk_indices"].shape), (16, 3))
+
     def test_source_none_owns_allocated_buffers(self):
         # source=None -> the registry allocates (owns) every slot.
         from sglang.srt.model_executor.cuda_graph_buffer_registry import (

test/registered/unit/model_executor/test_prefill_cuda_graph_runner.py

@@ -9,7 +9,12 @@ import torch
 import sglang.srt.model_executor.model_runner_components.cuda_graph_setup as graph_setup
 import sglang.srt.model_executor.runner.prefill_cuda_graph_runner as runner_module
 from sglang.srt.model_executor.cuda_graph_config import Backend
-from sglang.srt.model_executor.forward_batch_info import CaptureHiddenMode
+from sglang.srt.model_executor.forward_batch_info import (
+    CaptureHiddenMode,
+    ForwardBatch,
+    ForwardMode,
+    PPProxyTensors,
+)
 from sglang.srt.model_executor.model_runner_components.cuda_graph_setup import (
     capture_prefill_graph,
 )
@@ -61,6 +66,32 @@ class _FakeKVIndexKernel:
         return run
 
 
+class _FakeGraphSlot:
+    def __init__(self, buffer):
+        self.buffer = buffer
+
+    def slice_for(self, _batch_size, num_tokens):
+        return self.buffer[:num_tokens]
+
+
+class _FakeBatchRegistry:
+    def __init__(self):
+        self.slots = {
+            "input_ids": _FakeGraphSlot(torch.arange(4, dtype=torch.int64)),
+            "positions": _FakeGraphSlot(torch.arange(4, dtype=torch.int64)),
+            "out_cache_loc": _FakeGraphSlot(torch.arange(4, dtype=torch.int64)),
+        }
+
+    def fill_from(self, *_args, **_kwargs):
+        return None
+
+    def has_slot(self, name):
+        return name in self.slots
+
+    def get_slot(self, name):
+        return self.slots[name]
+
+
 class TestPrefillCudaGraphRunnerChunkedPrefix(CustomTestCase):
     def test_low_free_memory_still_captures_prefill_graph(self):
         eager_runner = object()
@@ -86,6 +117,7 @@ class TestPrefillCudaGraphRunnerChunkedPrefix(CustomTestCase):
             server_args=SimpleNamespace(),
             model=SimpleNamespace(),
             model_config=SimpleNamespace(context_len=8192, num_hidden_layers=1),
+            layer_info=SimpleNamespace(start_layer=0, end_layer=1),
             req_to_token_pool=SimpleNamespace(size=1),
         )
         language_model = SimpleNamespace(layers=[object()])
@@ -98,7 +130,7 @@ class TestPrefillCudaGraphRunnerChunkedPrefix(CustomTestCase):
             patch.object(
                 graph_setup,
                 "compute_attention_and_moe_layers",
-                return_value=([object()], [], [], [], []),
+                return_value=([object()], [], [], [], [None]),
             ),
             patch.object(
                 graph_setup,
@@ -145,6 +177,63 @@ class TestPrefillCudaGraphRunnerChunkedPrefix(CustomTestCase):
 
         self.assertIs(capture.runner, eager_runner)
 
+    def test_pp_proxy_output_is_trimmed_to_raw_prefill_tokens(self):
+        runner = PrefillCudaGraphRunner.__new__(PrefillCudaGraphRunner)
+        runner.raw_num_tokens = 3
+        output = PPProxyTensors(
+            {
+                "hidden_states": torch.arange(32).view(8, 4),
+                "residual": torch.arange(32, 64).view(8, 4),
+            }
+        )
+
+        trimmed = runner._finalize_execute_output(output)
+
+        self.assertIsInstance(trimmed, PPProxyTensors)
+        self.assertEqual(tuple(trimmed["hidden_states"].shape), (3, 4))
+        self.assertEqual(tuple(trimmed["residual"].shape), (3, 4))
+
+    def test_static_batch_preserves_consumed_multimodal_embeddings(self):
+        runner = PrefillCudaGraphRunner.__new__(PrefillCudaGraphRunner)
+        runner.capture_num_tokens = [4]
+        runner.buffer_registry = _FakeBatchRegistry()
+        runner.enable_cp_v2_bcg_capture = False
+        runner._is_full_backend = False
+        runner.backend = SimpleNamespace()
+        runner.has_mha_companion_layers = False
+        runner._prefill_static_buffers = None
+        runner.static_draft_hidden_states = None
+        runner.capture_return_pooled_hidden_states = False
+        runner._next_token_logits_buffer = lambda _rows: None
+        runner._prefill_logits_buffer_rows = lambda _batch: 1
+        runner._prepare_forward_metadata_for_replay = lambda *_args: None
+
+        mm_input_embeds = torch.randn(3, 8)
+        forward_batch = ForwardBatch(
+            forward_mode=ForwardMode.EXTEND,
+            batch_size=1,
+            input_ids=torch.arange(3, dtype=torch.int64),
+            req_pool_indices=torch.zeros(1, dtype=torch.int64),
+            seq_lens=torch.tensor([3], dtype=torch.int32),
+            out_cache_loc=torch.arange(3, dtype=torch.int64),
+            seq_lens_sum=3,
+            positions=torch.arange(3, dtype=torch.int64),
+            seq_lens_cpu=torch.tensor([3], dtype=torch.int32),
+            extend_seq_lens=torch.tensor([3], dtype=torch.int32),
+            extend_prefix_lens=torch.zeros(1, dtype=torch.int32),
+            extend_start_loc=torch.zeros(1, dtype=torch.int32),
+            extend_seq_lens_cpu=[3],
+            extend_prefix_lens_cpu=[0],
+            mm_inputs=None,
+            mm_input_embeds=mm_input_embeds,
+            capture_hidden_mode=CaptureHiddenMode.NULL,
+            global_forward_mode=ForwardMode.EXTEND,
+        )
+
+        static_batch = runner.load_batch(forward_batch)
+
+        self.assertIs(static_batch.mm_input_embeds, mm_input_embeds)
+
     def test_prefix_chunk_capacity_is_aggregate_and_can_be_overridden(self):
         graph_config = SimpleNamespace(
             prefill=SimpleNamespace(full_prefill_prefix_chunk_tokens=None, max_bs=8)