Improve CUDA graph and speculative execution output handling (#37329)

Co-authored-by: jiayisuse <jiayisuse@fb.com>
Co-authored-by: Yinghai Lu <yinghai@meta.com>
Co-authored-by: Hao Zhang <zhisbug@users.noreply.github.com>
Co-authored-by: Yichao Fu <yichaofu@meta.com>

test/registered/unit/model_executor/test_prefill_cuda_graph_runner.py

@@ -121,6 +121,13 @@ class TestPrefillCudaGraphRunnerChunkedPrefix(CustomTestCase):
             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),
+            get_cuda_graph_layers=lambda _layer_model: (
+                [object()],
+                [],
+                [],
+                [],
+                [None],
+            ),
         )
         language_model = SimpleNamespace(layers=[object()])
 
@@ -129,11 +136,6 @@ class TestPrefillCudaGraphRunnerChunkedPrefix(CustomTestCase):
             patch.object(
                 graph_setup, "resolve_language_model", return_value=language_model
             ),
-            patch.object(
-                graph_setup,
-                "compute_attention_and_moe_layers",
-                return_value=([object()], [], [], [], [None]),
-            ),
             patch.object(
                 graph_setup,
                 "get_available_gpu_memory",

test/registered/unit/spec/test_dspark_target_hidden_projection.py

@@ -0,0 +1,81 @@
+import unittest
+from types import MethodType, SimpleNamespace
+from unittest import mock
+
+import torch
+
+from sglang.srt.layers.aux_hidden_states import pack_aux_hidden_states
+from sglang.srt.models.dspark import DSparkDraftMixin
+from sglang.srt.speculative.dspark_components.dspark_kv_inject import (
+    TargetHiddenKvInjector,
+)
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=1, suite="base-a-test-cpu")
+
+
+class _Attention:
+    num_kv_heads = 1
+    head_dim = 3
+
+    def __init__(self) -> None:
+        self.attn = SimpleNamespace(k_scale=0.5, v_scale=0.25)
+        self.input = None
+
+    def kv_proj_only(self, hidden_states: torch.Tensor):
+        self.input = hidden_states
+        return hidden_states + 1, hidden_states + 2
+
+    def apply_k_norm(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return hidden_states + 3
+
+    def apply_k_rope(
+        self, _positions: torch.Tensor, hidden_states: torch.Tensor
+    ) -> torch.Tensor:
+        return hidden_states + 4
+
+
+class DSparkTargetHiddenProjectionTest(CustomTestCase):
+    def test_single_aux_hidden_state_is_returned_without_copy(self) -> None:
+        hidden_states = torch.empty(2, 3)
+
+        self.assertIs(hidden_states, pack_aux_hidden_states([hidden_states]))
+
+    def test_preprojected_hidden_is_not_projected_again(self) -> None:
+        attention = _Attention()
+        draft_model = SimpleNamespace(
+            layers=[SimpleNamespace(self_attn=attention)],
+            project_target_hidden=mock.Mock(
+                side_effect=AssertionError("projection must not run twice")
+            ),
+            _fused_kv_write_bundle=lambda _pool: None,
+            _stacked_ctx_kv_params=lambda: None,
+        )
+        draft_model.write_target_hidden_kv = MethodType(
+            DSparkDraftMixin.write_target_hidden_kv, draft_model
+        )
+        pool = SimpleNamespace(set_kv_buffer=mock.Mock())
+        injector = TargetHiddenKvInjector(
+            draft_model=draft_model,
+            draft_model_runner=SimpleNamespace(token_to_kv_pool=pool),
+            model_runner=SimpleNamespace(device=torch.device("cpu")),
+            device=torch.device("cpu"),
+            verify_num_draft_tokens=2,
+            block_pos_offsets=torch.arange(2),
+        )
+        projected_hidden = torch.arange(6, dtype=torch.float32).reshape(2, 3)
+
+        injector.inject_target_hidden(
+            target_hidden=projected_hidden,
+            cache_loc=torch.arange(2),
+            positions=torch.arange(2),
+            target_hidden_is_projected=True,
+        )
+
+        draft_model.project_target_hidden.assert_not_called()
+        self.assertIs(attention.input, projected_hidden)
+
+
+if __name__ == "__main__":
+    unittest.main()