Speculative decoding support on XPU (#23180)
This commit is contained in:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user