Support speculative decoding on CPU (#27862)

Co-authored-by: Valentine233 <xuan.liao@intel.com>

test/registered/cpu/test_spec_eagle_cpu.py

@@ -0,0 +1,65 @@
+"""EAGLE spec-decoding core on CPU: the standard config (topk=1, page_size=1)
+on the synchronous (non-overlap) path. topk > 1 tree drafting is covered in
+test_spec_eagle_topk_cpu.py (split to stay under the per-file CI timeout).
+"""
+
+import unittest
+
+from sglang.srt.environ import envs
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.kits.matched_stop_kit import MatchedStopMixin
+from sglang.test.kits.spec_server_kits import (
+    SpecAccuracyKit,
+    SpecCorrectnessKit,
+    SpecFeatureKit,
+    SpecLogprobKit,
+    SpecPenaltyKit,
+)
+from sglang.test.server_fixtures.spec_eagle_fixture import EagleLlama2Base
+
+# Measured 780s all-green on a 40-core GNR socket (1 launch + 18 methods).
+register_cpu_ci(est_time=800, suite="base-b-test-cpu")
+
+_KITS = (
+    SpecCorrectnessKit,
+    SpecAccuracyKit,
+    SpecLogprobKit,
+    SpecPenaltyKit,
+    SpecFeatureKit,
+    MatchedStopMixin,
+)
+
+
+class _Core(EagleLlama2Base):
+    """EAGLE (Llama-2) preset on CPU."""
+
+    attention_backend = "intel_amx"
+    disable_overlap = True
+    mem_fraction_static = 0.3
+    gsm8k_num_examples = 64
+    env_overrides = ((envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY, 1),)
+
+
+class TestEagleLlama2NoOverlap(_Core, *_KITS):
+    """Spec v1 (overlap scheduler off) -- the only mode reachable on CPU."""
+
+    # Standard chain config (topk=1, page_size=1), same shape as the CUDA core.
+    spec_steps = 5
+    spec_topk = 1
+    spec_tokens = 6
+    # EAGLE/Llama-2 topk=1 accepts modestly; tune against CI if needed.
+    acc_length_thres = 1.6
+    batch_accept_len_thres = 1.3
+    gsm8k_accept_len_thres = 1.3
+
+    @unittest.skip(
+        "constrained decoding on CPU needs a vocab-mask CPU branch in the "
+        "xgrammar backend (upstream gap, not spec-specific); the other grammar "
+        "backends lack the rollback spec verification requires"
+    )
+    def test_constrained_decoding(self):
+        pass
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/cpu/test_spec_eagle_parity_cpu.py

@@ -0,0 +1,23 @@
+import unittest
+
+from sglang.srt.environ import envs
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.kits.spec_server_kits import SpecParityKit
+from sglang.test.server_fixtures.spec_eagle_fixture import Eagle3Base
+
+# Estimated: 2 sequential 8B server launches + one 4-prompt greedy method
+# (CUDA sibling: 360); tune from CI TIMINGS once it has run there.
+register_cpu_ci(est_time=480, suite="base-b-test-cpu")
+
+
+class TestEagle3ParityCPU(SpecParityKit, Eagle3Base):
+    """EAGLE3 spec (intel_amx) greedy output == non-spec reference."""
+
+    attention_backend = "intel_amx"
+    disable_overlap = True
+    mem_fraction_static = 0.3
+    env_overrides = ((envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY, 1),)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/cpu/test_spec_eagle_topk_cpu.py

@@ -0,0 +1,62 @@
+"""EAGLE topk > 1 tree drafting on CPU (Llama-2 topk=4, synchronous path).
+
+Split from test_spec_eagle_cpu.py, mirroring the CUDA test_spec_eagle.py /
+test_spec_eagle_topk.py layout, so each file stays under the per-file CI
+timeout.
+"""
+
+import unittest
+
+from sglang.srt.environ import envs
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.kits.spec_server_kits import (
+    SpecAccuracyKit,
+    SpecCorrectnessKit,
+    SpecFeatureKit,
+    SpecLogprobKit,
+    SpecPenaltyKit,
+)
+from sglang.test.server_fixtures.spec_eagle_fixture import EagleLlama2Base
+
+# Measured 830s all-green on a 40-core GNR socket (1 launch + 14 methods).
+register_cpu_ci(est_time=850, suite="base-b-test-cpu")
+
+
+class _Core(EagleLlama2Base):
+    """EAGLE (Llama-2) preset on CPU."""
+
+    attention_backend = "intel_amx"
+    disable_overlap = True
+    mem_fraction_static = 0.3
+    gsm8k_num_examples = 64
+    env_overrides = ((envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY, 1),)
+
+
+class TestEagleLlama2Topk4(
+    _Core,
+    SpecCorrectnessKit,
+    SpecAccuracyKit,
+    SpecLogprobKit,
+    SpecPenaltyKit,
+    SpecFeatureKit,
+):
+    """EAGLE/Llama-2 topk=4 tree coverage (kits listed in bases)."""
+
+    spec_steps = 3
+    spec_topk = 4
+    spec_tokens = 8
+    acc_length_thres = 2.4
+    batch_accept_len_thres = 1.6
+    gsm8k_accept_len_thres = 2.0
+
+    @unittest.skip(
+        "constrained decoding on CPU needs a vocab-mask CPU branch in the "
+        "xgrammar backend (upstream gap, not spec-specific); the other grammar "
+        "backends lack the rollback spec verification requires"
+    )
+    def test_constrained_decoding(self):
+        pass
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/spec/test_eagle_worker_v2_topk1_fastpath.py

@@ -15,13 +15,13 @@ import torch
 from sglang.srt.speculative.adaptive_runtime_state import SpecRuntimeState
 from sglang.srt.speculative.eagle_utils import organize_draft_results
 from sglang.srt.speculative.eagle_worker_v2 import EagleDraftWorker, EAGLEWorkerV2
-from sglang.srt.utils import get_device
-from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.ci.ci_register import register_cpu_ci, register_cuda_ci
 from sglang.test.test_utils import CustomTestCase
 
 register_cuda_ci(est_time=20, stage="base-b", runner_config="1-gpu-small")
+register_cpu_ci(est_time=20, suite="base-a-test-cpu")
 
-DEVICE = get_device()
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
 
 def _fake_server_args(**fields):

test/registered/unit/spec/test_spec_cpu_overlap_constraint.py

@@ -0,0 +1,75 @@
+import unittest
+from types import SimpleNamespace
+
+from sglang.srt.arg_groups.speculative_hook import handle_speculative_decoding
+from sglang.srt.server_args import ServerArgs
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=20, suite="base-a-test-cpu")
+
+
+def _make_spec_args(device: str, algorithm: str = "EAGLE", **overrides) -> ServerArgs:
+    # model_path="dummy" short-circuits ServerArgs.__post_init__; invoke the
+    # speculative hook directly (same pattern as the unit/server_args tests).
+    args = ServerArgs(model_path="dummy")
+    args.speculative_algorithm = algorithm
+    args.device = device
+    # Fully specify the chain config so the hook doesn't auto-choose params.
+    args.speculative_num_steps = 3
+    args.speculative_eagle_topk = 1
+    args.speculative_num_draft_tokens = 4
+    args.get_model_config = lambda: SimpleNamespace(
+        hf_config=SimpleNamespace(
+            architectures=["LlamaForCausalLM"],
+            get_text_config=lambda: SimpleNamespace(),
+        )
+    )
+    for key, value in overrides.items():
+        setattr(args, key, value)
+    return args
+
+
+class TestSpecCPUOverlapConstraint(CustomTestCase):
+    def test_cpu_eagle_forces_disable_overlap_schedule(self):
+        args = _make_spec_args(device="cpu")
+        self.assertFalse(args.disable_overlap_schedule)
+
+        handle_speculative_decoding(args)
+
+        self.assertTrue(args.disable_overlap_schedule)
+
+    def test_cpu_eagle3_forces_disable_overlap_schedule(self):
+        args = _make_spec_args(device="cpu", algorithm="EAGLE3")
+
+        handle_speculative_decoding(args)
+
+        self.assertTrue(args.disable_overlap_schedule)
+
+    def test_cpu_explicit_disable_overlap_is_preserved(self):
+        args = _make_spec_args(device="cpu", disable_overlap_schedule=True)
+
+        # Already disabled: the hook must not flip the flag, and (unlike the
+        # forced-disable cases) must not warn about overriding it.
+        with self.assertLogs(
+            "sglang.srt.arg_groups.speculative_hook", "WARNING"
+        ) as logs:
+            handle_speculative_decoding(args)
+
+        self.assertTrue(args.disable_overlap_schedule)
+        self.assertFalse(
+            any("Overlap schedule" in message for message in logs.output),
+            f"hook warned about overriding an already-disabled overlap: {logs.output}",
+        )
+
+    def test_cuda_eagle_keeps_overlap_schedule(self):
+        # Guard the constraint's scope: the hook must not touch non-CPU devices.
+        args = _make_spec_args(device="cuda")
+
+        handle_speculative_decoding(args)
+
+        self.assertFalse(args.disable_overlap_schedule)
+
+
+if __name__ == "__main__":
+    unittest.main()