Speculative decoding support on XPU (#23180)

This commit is contained in:
ANSHUMAN TRIPATHY
2026-07-02 13:23:24 +08:00
committed by GitHub
parent 790156c98f
commit 4fffc6448b
9 changed files with 562 additions and 16 deletions
@@ -7,15 +7,26 @@ spec server (sequential -- one model resident at a time; see SpecParityKit).
import unittest
from sglang.srt.environ import envs
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.srt.utils import is_xpu
from sglang.test.ci.ci_register import register_cuda_ci, register_xpu_ci
from sglang.test.kits.spec_server_kits import SpecParityKit
from sglang.test.server_fixtures.spec_eagle_fixture import Eagle3Base
register_cuda_ci(est_time=360, stage="base-b", runner_config="1-gpu-large")
register_xpu_ci(est_time=360, stage="stage-b", runner_config="1-gpu-xpu")
_is_xpu = is_xpu()
class TestEagle3Parity(SpecParityKit, Eagle3Base):
"""EAGLE3 spec v2 (flashinfer) greedy output == non-spec reference.
class _Eagle3ParityBase(Eagle3Base):
"""Shared knobs for EAGLE3 parity variants; no test methods."""
env_overrides = ((envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY, 1),)
@unittest.skipIf(_is_xpu, "CUDA runner only")
class TestEagle3ParityCUDA(SpecParityKit, _Eagle3ParityBase):
"""EAGLE3 spec v2 (flashinfer, overlap) greedy output == non-spec reference.
SpecParityKit is first so its setUpClass runs the reference server (and tears
it down) before the fixture launches the spec server -- sequential, one model
@@ -23,7 +34,14 @@ class TestEagle3Parity(SpecParityKit, Eagle3Base):
"""
disable_overlap = False
env_overrides = ((envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY, 1),)
@unittest.skipUnless(_is_xpu, "XPU runner only")
class TestEagle3ParityXPU(SpecParityKit, _Eagle3ParityBase):
"""EAGLE3 parity on XPU (triton, no overlap, deterministic)."""
disable_overlap = False
attention_backend = "triton"
if __name__ == "__main__":
@@ -215,19 +215,25 @@ class TestBuildEagleTree(unittest.TestCase):
]
parents_list = [
torch.tensor(
[[-1, 0, 1, 2, 3], [-1, 0, 1, 2, 3]], dtype=torch.int64, device="cuda"
[[-1, 0, 1, 2, 3], [-1, 0, 1, 2, 3]],
dtype=torch.int64,
device=get_device(),
),
torch.tensor(
[[4, 8, 9, 10], [4, 5, 6, 7]], dtype=torch.int64, device="cuda"
[[4, 8, 9, 10], [4, 5, 6, 7]], dtype=torch.int64, device=get_device()
),
torch.tensor(
[[20, 24, 21, 28], [24, 28, 20, 21]], dtype=torch.int64, device="cuda"
[[20, 24, 21, 28], [24, 28, 20, 21]],
dtype=torch.int64,
device=get_device(),
),
torch.tensor(
[[36, 40, 41, 44], [36, 40, 44, 45]], dtype=torch.int64, device="cuda"
[[36, 40, 41, 44], [36, 40, 44, 45]],
dtype=torch.int64,
device=get_device(),
),
]
seq_lens = torch.tensor([5, 10], dtype=torch.int64, device="cuda")
seq_lens = torch.tensor([5, 10], dtype=torch.int64, device=get_device())
topk = 4
depth = 4
num_draft_token = 8