minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

feat: add optimized Domino rollout to DFlash V2 (#36899)

Browse Source 
Co-authored-by: Qiaolin-Yu <liin1211@outlook.com>
This commit is contained in:
Francis
2026-09-09 14:12:15 -07:00
committed by GitHub
co-authored by Qiaolin-Yu
parent 2948a62a6f
commit a84ffd1326
8 changed files with 861 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/e2e/speculative/test_dflash_domino.py
+104
View File
@@ -0,0 +1,104 @@
import unittest
from pathlib import Path
from tempfile import NamedTemporaryFile
import requests
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.kits.eval_accuracy_kit import GSM8KMixin
from sglang.test.test_utils import (
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
DEFAULT_URL_FOR_TEST,
CustomTestCase,
kill_process_tree,
popen_launch_server,
)
register_cuda_ci(est_time=600, stage="base-b", runner_config="1-gpu-small")
class TestDFlashDominoFullVocab(CustomTestCase):
model = "Qwen/Qwen3-8B"
draft_model = "Huang2020/Qwen3-8B-Domino-b16"
candidate_pool_size = 0
@classmethod
def setUpClass(cls):
cls.base_url = DEFAULT_URL_FOR_TEST
cls.server_log = NamedTemporaryFile(mode="w+", suffix="-domino.log")
cls.addClassCleanup(cls.server_log.close)
print(f"Domino server log: {cls.server_log.name}", flush=True)
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
return_stdout_stderr=(cls.server_log, cls.server_log),
other_args=[
"--trust-remote-code",
"--dtype",
"bfloat16",
"--tp-size",
"1",
"--attention-backend",
"triton",
"--speculative-algorithm",
"DFLASH",
"--speculative-draft-model-path",
cls.draft_model,
"--speculative-domino-candidate-pool-size",
str(cls.candidate_pool_size),
"--speculative-draft-attention-backend",
"triton",
"--cuda-graph-backend-decode",
"full",
"--cuda-graph-max-bs-decode",
"64",
"--max-running-requests",
"64",
"--mem-fraction-static",
"0.7",
],
)
def test_domino_runtime(self):
response = requests.get(self.base_url + "/server_info", timeout=10)
response.raise_for_status()
state = response.json()["internal_states"][0]
self.assertEqual(state["speculative_num_draft_tokens"], 16)
self.assertFalse(state["disable_overlap_schedule"])
self.assertEqual(
state["speculative_domino_candidate_pool_size"], self.candidate_pool_size
)
log = Path(self.server_log.name).read_text()
self.assertIn(
"DFLASH Domino rollout enabled (BF16, TP=1, "
f"block-shared candidate pool size={self.candidate_pool_size}).",
log,
)
self.assertIn("Domino rollout folded into the draft cuda graph", log)
self.assertIn(
"Capture draft verify CUDA graph begin. backend=full, num_tokens_per_req=16,",
log,
)
self.assertIn(
"Capture target verify CUDA graph begin. backend=full, num_tokens_per_req=16,",
log,
)
@classmethod
def tearDownClass(cls):
if hasattr(cls, "process") and cls.process:
kill_process_tree(cls.process.pid)
class TestDFlashDomino(TestDFlashDominoFullVocab, GSM8KMixin):
gsm8k_score_threshold = 0.90
gsm8k_num_examples = 200
gsm8k_accept_length_thres = 4.0
gsm8k_num_threads = 128
gsm8k_num_shots = 5
candidate_pool_size = 2048
if __name__ == "__main__":
unittest.main()
test/registered/kernel/speculative/test_dflash_domino.py
+125
View File
@@ -0,0 +1,125 @@
import unittest
import torch
from torch import nn
from sglang.srt.speculative.dflash_worker_v2 import _DominoDraftSampler
from sglang.srt.speculative.domino_utils import _domino_gru_cell, domino_greedy_rollout
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.test_utils import CustomTestCase
register_cuda_ci(est_time=10, stage="base-b-kernel-unit", runner_config="1-gpu-large")
@unittest.skipUnless(torch.cuda.is_available(), "CUDA is required")
class TestDFlashDominoRollout(CustomTestCase):
def setUp(self):
torch.manual_seed(0)
self.embedding = nn.Embedding(31, 8, device="cuda", dtype=torch.bfloat16)
self.prefix_gru = nn.GRU(8, 4, batch_first=True, bias=False).cuda().bfloat16()
self.embed_proj = (
nn.Sequential(
nn.Linear(12, 5, bias=False), nn.SiLU(), nn.Linear(5, 31, bias=False)
)
.cuda()
.bfloat16()
)
self.lm_head_weight = torch.randn(31, 8, device="cuda", dtype=torch.bfloat16)
self.hidden = torch.randn(3, 16, 8, device="cuda", dtype=torch.bfloat16)
self.bonus_tokens = torch.tensor([1, 4, 9], device="cuda")
def rollout(self, hidden, bonus_tokens, pool_size=5, shift_label=True):
return domino_greedy_rollout(
draft_hidden=hidden,
bonus_tokens=bonus_tokens,
target_embedding=self.embedding,
lm_head_weight=self.lm_head_weight,
prefix_gru=self.prefix_gru,
embed_proj=self.embed_proj,
vocab_size=31,
shift_label=shift_label,
candidate_pool_size=pool_size,
)
def test_gru_feedback_matches_sequence(self):
embeddings = self.embedding(torch.tensor([[1, 2, 3], [4, 5, 6]], device="cuda"))
_, expected = self.prefix_gru(embeddings)
state = torch.zeros(2, 4, device="cuda", dtype=torch.bfloat16)
for step in embeddings.unbind(dim=1):
state = _domino_gru_cell(self.prefix_gru, step, state)
torch.testing.assert_close(state, expected[0], rtol=0.02, atol=0.002)
def test_candidate_pool_boundaries(self):
for shift_label in (True, False):
with self.subTest(shift_label=shift_label):
full = self.rollout(self.hidden, self.bonus_tokens, 0, shift_label)
for pool_size in (31, 32):
actual = self.rollout(
self.hidden, self.bonus_tokens, pool_size, shift_label
)
torch.testing.assert_close(actual, full, rtol=0, atol=0)
first_hidden = self.hidden[:, 0 if shift_label else 1]
expected_first = (first_hidden @ self.lm_head_weight.T).argmax(dim=-1)
for block_size in (2, 16):
limited = self.rollout(
self.hidden[:, :block_size], self.bonus_tokens, 1, shift_label
)
self.assertEqual(limited.shape, (3, block_size - 1))
torch.testing.assert_close(limited[:, 0], expected_first)
if block_size > 2:
torch.testing.assert_close(
limited[:, 1:], limited[:, 1:2].expand_as(limited[:, 1:])
)
def test_batch_matches_individual_requests(self):
for pool_size in (0, 5):
with self.subTest(pool_size=pool_size):
batched = self.rollout(self.hidden, self.bonus_tokens, pool_size)
individual = torch.cat(
[
self.rollout(hidden[None], bonus[None], pool_size)
for hidden, bonus in zip(self.hidden, self.bonus_tokens)
]
)
torch.testing.assert_close(batched, individual, rtol=0, atol=0)
def test_sampler_replays_with_new_inputs(self):
sampler = _DominoDraftSampler(
target_embedding=self.embedding,
lm_head_weight=self.lm_head_weight,
prefix_gru=self.prefix_gru,
embed_proj=self.embed_proj,
vocab_size=31,
block_size=16,
shift_label=True,
max_bs=3,
candidate_pool_size=5,
)
block_ids = torch.zeros(3, 16, device="cuda", dtype=torch.long)
block_ids[:, 0].copy_(self.bonus_tokens)
def sample():
sampler(self.hidden.flatten(0, 1), block_ids.flatten())
warmup = torch.cuda.Stream()
warmup.wait_stream(torch.cuda.current_stream())
with torch.cuda.stream(warmup):
sample()
torch.cuda.current_stream().wait_stream(warmup)
graph = torch.cuda.CUDAGraph()
with torch.cuda.graph(graph):
sample()
for _ in range(2):
self.hidden.copy_(torch.randn_like(self.hidden))
block_ids[:, 0].copy_(torch.randint(31, (3,), device="cuda"))
expected = self.rollout(self.hidden, block_ids[:, 0])
graph.replay()
torch.cuda.synchronize()
torch.testing.assert_close(
sampler.out.view(3, 15), expected, rtol=0, atol=0
)
if __name__ == "__main__":
unittest.main()
test/registered/unit/spec/test_dflash_domino.py
+130
View File
@@ -0,0 +1,130 @@
import unittest
from types import SimpleNamespace
import torch
from torch import nn
from sglang.srt.models.dflash import DFlashDraftModel
from sglang.srt.speculative.dflash_utils import parse_dflash_draft_config
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=10, suite="base-a-test-cpu")
def _domino_config(**overrides):
dflash_config = {
"projector_type": "domino",
"mask_token_id": 29,
"shift_label": True,
"target_layer_ids": [1, 3],
"pure_draft_prefix_len": 1,
"gru_hidden_dim": 4,
"emb_dim": 5,
}
dflash_config.update(overrides.pop("dflash_config", {}))
fields = {
"num_hidden_layers": 2,
"num_target_layers": 4,
"block_size": 16,
"hidden_size": 8,
"vocab_size": 31,
"emb_dim": 5,
"dflash_config": dflash_config,
}
fields.update(overrides)
return SimpleNamespace(**fields)
def _projector_model(projector_type="domino"):
model = DFlashDraftModel.__new__(DFlashDraftModel)
nn.Module.__init__(model)
model.projector_type = projector_type
model.config = SimpleNamespace(hidden_size=8)
if projector_type == "domino":
model.prefix_gru = nn.GRU(8, 4, batch_first=True, bias=False)
model.embed_proj = nn.Sequential(
nn.Linear(12, 5, bias=False),
nn.SiLU(),
nn.Linear(5, 31, bias=False),
)
else:
model.prefix_gru = None
model.embed_proj = None
return model
def _projector_weights(model):
return {
"prefix_gru.weight_ih_l0": torch.randn_like(model.prefix_gru.weight_ih_l0),
"prefix_gru.weight_hh_l0": torch.randn_like(model.prefix_gru.weight_hh_l0),
"embed_proj.0.weight": torch.randn_like(model.embed_proj[0].weight),
"embed_proj.2.weight": torch.randn_like(model.embed_proj[2].weight),
}
class TestDFlashDominoConfig(CustomTestCase):
def test_top_level_emb_dim_fallback(self):
config = _domino_config()
del config.dflash_config["emb_dim"]
self.assertEqual(parse_dflash_draft_config(draft_hf_config=config).emb_dim, 5)
def test_invalid_domino_config_fails_fast(self):
cases = {
"shift_label": {"shift_label": 1},
"pure_draft_prefix_len": {"pure_draft_prefix_len": 2},
"gru_hidden_dim": {"gru_hidden_dim": None},
}
for expected, updates in cases.items():
with self.subTest(expected=expected):
with self.assertRaisesRegex(ValueError, expected):
parse_dflash_draft_config(
draft_hf_config=_domino_config(dflash_config=updates)
)
config = _domino_config(dflash_config={"emb_dim": None}, emb_dim=None)
with self.assertRaisesRegex(ValueError, "emb_dim"):
parse_dflash_draft_config(draft_hf_config=config)
with self.assertRaisesRegex(ValueError, "block_size > 1"):
parse_dflash_draft_config(draft_hf_config=_domino_config(block_size=1))
def test_conflicting_emb_dim_fails(self):
with self.assertRaisesRegex(ValueError, "emb_dim differs"):
parse_dflash_draft_config(draft_hf_config=_domino_config(emb_dim=6))
class TestDFlashDominoWeights(CustomTestCase):
def test_projector_weights_load_exactly(self):
model = _projector_model()
weights = _projector_weights(model)
model.load_weights(weights.items())
for name, expected in weights.items():
torch.testing.assert_close(
dict(model.named_parameters())[name], expected, rtol=0, atol=0
)
def test_each_required_projector_weight_is_checked(self):
for missing_name in _projector_weights(_projector_model()):
with self.subTest(missing_name=missing_name):
model = _projector_model()
weights = _projector_weights(model)
del weights[missing_name]
with self.assertRaisesRegex(ValueError, missing_name):
model.load_weights(weights.items())
def test_projector_shape_mismatch_fails(self):
model = _projector_model()
weights = _projector_weights(model)
weights["embed_proj.2.weight"] = torch.empty(30, 5)
with self.assertRaisesRegex(ValueError, "shape mismatch"):
model.load_weights(weights.items())
def test_projector_weights_require_domino_config(self):
model = _projector_model(projector_type="domnio")
with self.assertRaisesRegex(ValueError, "projector_type"):
model.load_weights([("prefix_gru.weight_ih_l0", torch.empty(12, 8))])
if __name__ == "__main__":
unittest.main()