From b42b05a5ea4304181e9d9736977ab4bea40a728c Mon Sep 17 00:00:00 2001 From: jianuo-huang <1170180988@qq.com> Date: Wed, 15 Jul 2026 19:19:50 +0800 Subject: [PATCH 1/7] feat: add correctness-first Domino support --- python/sglang/srt/models/dflash.py | 58 +++ python/sglang/srt/speculative/dflash_utils.py | 80 ++++ .../srt/speculative/dflash_worker_v2.py | 58 ++- python/sglang/srt/speculative/domino_utils.py | 154 ++++++++ .../spec/dflash/test_dflash_domino.py | 65 ++++ .../unit/spec/test_dflash_domino.py | 357 ++++++++++++++++++ 6 files changed, 771 insertions(+), 1 deletion(-) create mode 100644 python/sglang/srt/speculative/domino_utils.py create mode 100644 test/registered/spec/dflash/test_dflash_domino.py create mode 100644 test/registered/unit/spec/test_dflash_domino.py diff --git a/python/sglang/srt/models/dflash.py b/python/sglang/srt/models/dflash.py index ffef278fb833..3955910f7097 100644 --- a/python/sglang/srt/models/dflash.py +++ b/python/sglang/srt/models/dflash.py @@ -370,6 +370,31 @@ def __init__(self, config, quant_config=None, prefix: str = "") -> None: self.hidden_norm = RMSNorm(hidden_size, eps=rms_norm_eps) self.block_size = draft_config.resolve_block_size(default=16) + self.projector_type = draft_config.projector_type + self.shift_label = draft_config.shift_label + self.prefix_gru: Optional[nn.GRU] = None + self.embed_proj: Optional[nn.Sequential] = None + if draft_config.is_domino: + assert draft_config.gru_hidden_dim is not None + assert draft_config.emb_dim is not None + self.prefix_gru = nn.GRU( + input_size=hidden_size, + hidden_size=int(draft_config.gru_hidden_dim), + num_layers=1, + batch_first=True, + bias=False, + ) + self.embed_proj = nn.Sequential( + nn.Linear( + hidden_size + int(draft_config.gru_hidden_dim), + int(draft_config.emb_dim), + bias=False, + ), + nn.SiLU(), + nn.Linear( + int(draft_config.emb_dim), int(config.vocab_size), bias=False + ), + ) def get_attention_sliding_window_size(self) -> Optional[int]: return get_dflash_attention_sliding_window_size(self.config) @@ -441,6 +466,7 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): ] params_dict = dict(self.named_parameters()) + loaded_params = set() def resolve_param_name(name: str) -> Optional[str]: if name in params_dict: @@ -456,6 +482,14 @@ def resolve_param_name(name: str) -> Optional[str]: return None for name, loaded_weight in weights: + unprefixed_name = name.removeprefix("model.") + if self.projector_type != "domino" and unprefixed_name.startswith( + ("prefix_gru.", "embed_proj.") + ): + raise ValueError( + "DFLASH checkpoint contains Domino projector weights but " + f"projector_type={self.projector_type!r}." + ) for param_name, weight_name, shard_id in stacked_params_mapping: if f".{weight_name}." not in name: continue @@ -466,6 +500,7 @@ def resolve_param_name(name: str) -> Optional[str]: param = params_dict[resolved_name] weight_loader = getattr(param, "weight_loader", default_weight_loader) weight_loader(param, loaded_weight, shard_id) + loaded_params.add(resolved_name) break else: resolved_name = resolve_param_name(name) @@ -473,6 +508,14 @@ def resolve_param_name(name: str) -> Optional[str]: # Ignore unexpected weights (e.g., HF rotary caches). continue param = params_dict[resolved_name] + if resolved_name.startswith(("prefix_gru.", "embed_proj.")) and tuple( + loaded_weight.shape + ) != tuple(param.shape): + raise ValueError( + "DFLASH Domino projector weight shape mismatch: " + f"expected {resolved_name}{tuple(param.shape)}, got " + f"{tuple(loaded_weight.shape)} from {name!r}." + ) if resolved_name.endswith("fc.weight") and tuple( loaded_weight.shape ) != tuple(param.shape): @@ -485,6 +528,21 @@ def resolve_param_name(name: str) -> Optional[str]: ) weight_loader = getattr(param, "weight_loader", default_weight_loader) weight_loader(param, loaded_weight) + loaded_params.add(resolved_name) + + if self.projector_type == "domino": + required = { + "prefix_gru.weight_ih_l0", + "prefix_gru.weight_hh_l0", + "embed_proj.0.weight", + "embed_proj.2.weight", + } + missing = required - loaded_params + if missing: + raise ValueError( + "DFLASH Domino checkpoint is missing required projector weights: " + f"{sorted(missing)}." + ) class DFlashLagunaAttention(DFlashAttention): diff --git a/python/sglang/srt/speculative/dflash_utils.py b/python/sglang/srt/speculative/dflash_utils.py index 11d3938ae23f..b76539a787b7 100644 --- a/python/sglang/srt/speculative/dflash_utils.py +++ b/python/sglang/srt/speculative/dflash_utils.py @@ -392,6 +392,15 @@ class DFlashDraftConfig: target_layer_ids: Optional[List[int]] mask_token: str mask_token_id: Optional[int] + projector_type: Optional[str] + shift_label: Optional[bool] + pure_draft_prefix_len: Optional[int] + gru_hidden_dim: Optional[int] + emb_dim: Optional[int] + + @property + def is_domino(self) -> bool: + return self.projector_type == "domino" def require_num_layers(self) -> int: if self.num_hidden_layers is None: @@ -510,6 +519,72 @@ def parse_dflash_draft_config(*, draft_hf_config: Any) -> DFlashDraftConfig: f"got {mask_token_id}." ) + projector_type = dflash_cfg.get( + "projector_type", _cfg_get(draft_hf_config, "projector_type", None) + ) + shift_label = None + pure_draft_prefix_len = None + gru_hidden_dim = None + emb_dim = None + if projector_type == "domino": + shift_label = dflash_cfg.get( + "shift_label", _cfg_get(draft_hf_config, "shift_label", None) + ) + pure_draft_prefix_len = _parse_optional_int( + dflash_cfg.get( + "pure_draft_prefix_len", + _cfg_get(draft_hf_config, "pure_draft_prefix_len", None), + ), + field_name="DFLASH Domino pure_draft_prefix_len", + min_value=0, + ) + gru_hidden_dim = _parse_optional_int( + dflash_cfg.get( + "gru_hidden_dim", _cfg_get(draft_hf_config, "gru_hidden_dim", None) + ), + field_name="DFLASH Domino gru_hidden_dim", + min_value=1, + ) + nested_emb_dim = _parse_optional_int( + dflash_cfg.get("emb_dim", None), + field_name="DFLASH Domino dflash_config.emb_dim", + min_value=1, + ) + top_level_emb_dim = _parse_optional_int( + _cfg_get(draft_hf_config, "emb_dim", None), + field_name="DFLASH Domino top-level emb_dim", + min_value=1, + ) + if ( + nested_emb_dim is not None + and top_level_emb_dim is not None + and nested_emb_dim != top_level_emb_dim + ): + raise ValueError( + "DFLASH Domino emb_dim differs between dflash_config and the " + f"top-level config: {nested_emb_dim} != {top_level_emb_dim}." + ) + emb_dim = nested_emb_dim if nested_emb_dim is not None else top_level_emb_dim + + if not isinstance(shift_label, bool): + raise ValueError( + "DFLASH Domino requires dflash_config.shift_label to be a bool, " + f"got {shift_label!r}." + ) + if pure_draft_prefix_len != 1: + raise ValueError( + "DFLASH Domino currently requires pure_draft_prefix_len=1, " + f"got {pure_draft_prefix_len!r}." + ) + if gru_hidden_dim is None: + raise ValueError("DFLASH Domino requires dflash_config.gru_hidden_dim.") + if emb_dim is None: + raise ValueError("DFLASH Domino requires dflash_config.emb_dim.") + if block_size is not None and block_size <= 1: + raise ValueError( + f"DFLASH Domino requires block_size > 1, got {block_size}." + ) + return DFlashDraftConfig( num_hidden_layers=num_hidden_layers, num_target_layers=num_target_layers, @@ -517,6 +592,11 @@ def parse_dflash_draft_config(*, draft_hf_config: Any) -> DFlashDraftConfig: target_layer_ids=parsed_target_layer_ids, mask_token=mask_token, mask_token_id=mask_token_id, + projector_type=projector_type, + shift_label=shift_label, + pure_draft_prefix_len=pure_draft_prefix_len, + gru_hidden_dim=gru_hidden_dim, + emb_dim=emb_dim, ) diff --git a/python/sglang/srt/speculative/dflash_worker_v2.py b/python/sglang/srt/speculative/dflash_worker_v2.py index a1095828f93a..b9c48c775b41 100644 --- a/python/sglang/srt/speculative/dflash_worker_v2.py +++ b/python/sglang/srt/speculative/dflash_worker_v2.py @@ -39,6 +39,10 @@ is_dflash_sampling_verify_available, parse_dflash_draft_config, ) +from sglang.srt.speculative.domino_utils import ( + domino_greedy_rollout, + validate_domino_runtime, +) from sglang.srt.speculative.draft_worker_common import ( build_block_pos_offsets, build_draft_tp_worker, @@ -195,6 +199,28 @@ def __init__( draft_config = parse_dflash_draft_config( draft_hf_config=self.draft_model_runner.model_config.hf_config ) + self._is_domino = draft_config.is_domino + if self._is_domino: + target_model = self.target_worker.model_runner.model + target_embedding = target_model.get_input_embeddings() + lm_head = getattr(target_model, "lm_head", None) + prefix_gru = getattr(self.draft_model, "prefix_gru", None) + embed_proj = getattr(self.draft_model, "embed_proj", None) + if lm_head is None or prefix_gru is None or embed_proj is None: + raise ValueError( + "DFLASH Domino requires target lm_head and loaded Domino projector modules." + ) + validate_domino_runtime( + device=torch.device(self.device), + tp_size=int(get_tp_group().world_size), + target_vocab_size=int(self.model_runner.model_config.vocab_size), + draft_vocab_size=int(self.draft_model_runner.model_config.vocab_size), + hidden_size=int(self.draft_model.config.hidden_size), + target_embedding=target_embedding, + lm_head=lm_head, + prefix_gru=prefix_gru, + embed_proj=embed_proj, + ) if server_args.speculative_num_draft_tokens is None: # Should not happen (ServerArgs should have inferred it), but keep a fallback. self.block_size = int(draft_config.resolve_block_size(default=16)) @@ -212,6 +238,11 @@ def __init__( model_block_size, ) self.speculative_num_draft_tokens = int(self.block_size) + if self._is_domino and self.block_size <= 1: + raise ValueError( + "DFLASH Domino requires speculative_num_draft_tokens > 1, " + f"got {self.block_size}." + ) self._mask_token = draft_config.mask_token self._mask_token_id_override = draft_config.mask_token_id @@ -228,6 +259,10 @@ def __init__( self.draft_window_size, self.use_compact_draft_cache, ) + if self._is_domino: + logger.info( + "DFLASH Domino rollout enabled (eager BF16, TP=1, full-vocabulary greedy)." + ) logger.info( "DFLASH draft runner ready. mask_token=%s, mask_token_id=%s, mask_token_id_override=%s", self._mask_token, @@ -359,6 +394,8 @@ def _eager(reason): if envs.SGLANG_DFLASH_EAGER_DRAFT_SAMPLER.get(): return _eager("SGLANG_DFLASH_EAGER_DRAFT_SAMPLER=1") + if self._is_domino: + return _eager("Domino uses sequential eager rollout") if self.block_size <= 1: return _eager("block_size<=1") target_model = self._target_worker.model_runner.model @@ -1613,7 +1650,26 @@ def forward_batch_generation( draft_out = self.draft_model_runner.forward(forward_batch) draft_logits_output = draft_out.logits_output - if self._draft_sampler is not None and draft_out.can_run_graph: + if self._is_domino: + draft_hidden = draft_logits_output.hidden_states + if draft_hidden is None: + raise RuntimeError("DFLASH draft model returned no hidden states.") + draft_hidden = draft_hidden.view(bs, int(self.block_size), -1) + prefix_gru = self.draft_model.prefix_gru + embed_proj = self.draft_model.embed_proj + if prefix_gru is None or embed_proj is None: + raise RuntimeError("DFLASH Domino projector modules are unavailable.") + draft_next = domino_greedy_rollout( + draft_hidden=draft_hidden, + verified_ids=block_ids[:, 0], + target_embedding=embed_module, + lm_head_weight=lm_head.weight, + prefix_gru=prefix_gru, + embed_proj=embed_proj, + vocab_size=int(self.model_runner.model_config.vocab_size), + shift_label=bool(self.draft_model.shift_label), + ) + elif self._draft_sampler is not None and draft_out.can_run_graph: draft_next = self._draft_sampler.out[ : bs * (int(self.block_size) - 1) ].view(bs, int(self.block_size) - 1) diff --git a/python/sglang/srt/speculative/domino_utils.py b/python/sglang/srt/speculative/domino_utils.py new file mode 100644 index 000000000000..6b60d666b106 --- /dev/null +++ b/python/sglang/srt/speculative/domino_utils.py @@ -0,0 +1,154 @@ +from __future__ import annotations + +import torch +import torch.nn.functional as F +from torch import nn + + +def validate_domino_runtime( + *, + device: torch.device, + tp_size: int, + target_vocab_size: int, + draft_vocab_size: int, + hidden_size: int, + target_embedding: nn.Module, + lm_head: nn.Module, + prefix_gru: nn.GRU, + embed_proj: nn.Sequential, +) -> None: + """Validate the deliberately narrow correctness-first Domino runtime.""" + if device.type != "cuda": + raise ValueError(f"DFLASH Domino currently requires CUDA, got {device}.") + if int(tp_size) != 1: + raise ValueError(f"DFLASH Domino currently requires TP=1, got TP={tp_size}.") + if int(target_vocab_size) != int(draft_vocab_size): + raise ValueError( + "DFLASH Domino requires identical target and draft vocab sizes, " + f"got target={target_vocab_size}, draft={draft_vocab_size}." + ) + + embedding_weight = getattr(target_embedding, "weight", None) + lm_head_weight = getattr(lm_head, "weight", None) + if embedding_weight is None or lm_head_weight is None: + raise ValueError( + "DFLASH Domino requires target embedding and lm_head weight tensors." + ) + + shard = getattr(lm_head, "shard_indices", None) + if shard is not None: + if int(shard.num_added_elements) != 0: + raise ValueError( + "DFLASH Domino does not support added-vocab lm_head shards." + ) + if int(shard.org_vocab_start_index) != 0 or int(shard.num_org_elements) != int( + target_vocab_size + ): + raise ValueError( + "DFLASH Domino requires the complete target vocabulary on TP=1." + ) + elif int(lm_head_weight.shape[0]) != int(target_vocab_size): + raise ValueError( + "DFLASH Domino lm_head row count must equal the target vocab size, " + f"got rows={int(lm_head_weight.shape[0])}, vocab={target_vocab_size}." + ) + + if int(embedding_weight.shape[0]) < int(target_vocab_size): + raise ValueError( + "DFLASH Domino target embedding has fewer rows than the target vocab " + f"size: rows={int(embedding_weight.shape[0])}, vocab={target_vocab_size}." + ) + + if int(embedding_weight.shape[-1]) != int(hidden_size) or int( + lm_head_weight.shape[-1] + ) != int(hidden_size): + raise ValueError( + "DFLASH Domino target embedding/lm_head hidden size does not match " + f"the draft hidden size {hidden_size}." + ) + if int(prefix_gru.input_size) != int(hidden_size): + raise ValueError( + "DFLASH Domino GRU input size does not match the draft hidden size." + ) + if int(embed_proj[0].in_features) != int(hidden_size + prefix_gru.hidden_size): + raise ValueError("DFLASH Domino projector input shape is inconsistent.") + if int(embed_proj[2].out_features) != int(target_vocab_size): + raise ValueError("DFLASH Domino projector output vocab size is inconsistent.") + + weights = ( + embedding_weight, + lm_head_weight, + prefix_gru.weight_ih_l0, + prefix_gru.weight_hh_l0, + embed_proj[0].weight, + embed_proj[2].weight, + ) + non_bf16 = [ + str(weight.dtype) for weight in weights if weight.dtype != torch.bfloat16 + ] + if non_bf16: + raise ValueError( + "DFLASH Domino currently requires BF16 target and projector weights; " + f"found {non_bf16}." + ) + + +@torch.no_grad() +def domino_greedy_rollout( + *, + draft_hidden: torch.Tensor, + verified_ids: torch.Tensor, + target_embedding: nn.Module, + lm_head_weight: torch.Tensor, + prefix_gru: nn.GRU, + embed_proj: nn.Sequential, + vocab_size: int, + shift_label: bool, +) -> torch.Tensor: + """Generate a Domino chain with native GRU and full-vocabulary greedy logits.""" + if draft_hidden.ndim != 3: + raise ValueError( + f"draft_hidden must have shape [batch, block, hidden], got {tuple(draft_hidden.shape)}." + ) + batch_size, block_size, _ = draft_hidden.shape + if verified_ids.shape != (batch_size,): + raise ValueError( + f"verified_ids must have shape ({batch_size},), got {tuple(verified_ids.shape)}." + ) + + num_proposals = int(block_size) - 1 + if num_proposals < 1: + raise ValueError(f"Domino requires block_size > 1, got {block_size}.") + start = 0 if shift_label else 1 + z = draft_hidden[:, start : start + num_proposals, :] + if int(z.shape[1]) != num_proposals: + raise ValueError( + "Domino draft hidden states do not contain enough proposal positions." + ) + + weight = lm_head_weight[: int(vocab_size)] + + def base_logits(hidden: torch.Tensor) -> torch.Tensor: + if hidden.dtype != weight.dtype: + hidden = hidden.to(weight.dtype) + return F.linear(hidden, weight) + + first_ids = torch.argmax(base_logits(z[:, 0, :]), dim=-1).to(torch.long) + proposals = [first_ids] + if num_proposals == 1: + return first_ids[:, None] + + prefix_ids = torch.stack((verified_ids, first_ids), dim=1) + _, gru_hidden = prefix_gru(target_embedding(prefix_ids)) + + for index in range(1, num_proposals): + step_hidden = z[:, index, :] + correction = embed_proj(torch.cat((step_hidden, gru_hidden[0]), dim=-1)) + next_ids = torch.argmax(base_logits(step_hidden) + correction, dim=-1).to( + torch.long + ) + proposals.append(next_ids) + if index + 1 < num_proposals: + _, gru_hidden = prefix_gru(target_embedding(next_ids[:, None]), gru_hidden) + + return torch.stack(proposals, dim=1) diff --git a/test/registered/spec/dflash/test_dflash_domino.py b/test/registered/spec/dflash/test_dflash_domino.py new file mode 100644 index 000000000000..b55d84faf2ae --- /dev/null +++ b/test/registered/spec/dflash/test_dflash_domino.py @@ -0,0 +1,65 @@ +import unittest + +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=360, stage="base-b", runner_config="1-gpu-small") + + +class TestDFlashDomino(CustomTestCase, GSM8KMixin): + model = "Qwen/Qwen3-8B" + draft_model = "Huang2020/Qwen3-8B-Domino-b16" + gsm8k_score_threshold = 0.90 + gsm8k_num_examples = 200 + gsm8k_accept_length_thres = 4.0 + gsm8k_num_threads = 128 + gsm8k_num_shots = 5 + + @classmethod + def setUpClass(cls): + cls.base_url = DEFAULT_URL_FOR_TEST + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=[ + "--trust-remote-code", + "--dtype", + "bfloat16", + "--tp-size", + "1", + "--attention-backend", + "triton", + "--speculative-algorithm", + "DFLASH", + "--speculative-draft-model-path", + cls.draft_model, + "--speculative-draft-attention-backend", + "triton", + "--cuda-graph-backend-decode", + "disabled", + "--cuda-graph-backend-prefill", + "disabled", + "--disable-overlap-schedule", + "--max-running-requests", + "64", + "--mem-fraction-static", + "0.7", + ], + ) + + @classmethod + def tearDownClass(cls): + if hasattr(cls, "process") and cls.process: + kill_process_tree(cls.process.pid) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/spec/test_dflash_domino.py b/test/registered/unit/spec/test_dflash_domino.py new file mode 100644 index 000000000000..74d0a5cdba56 --- /dev/null +++ b/test/registered/unit/spec/test_dflash_domino.py @@ -0,0 +1,357 @@ +import unittest +from types import SimpleNamespace + +import torch +import torch.nn.functional as F +from torch import nn + +from sglang.srt.models.dflash import DFlashDraftModel +from sglang.srt.speculative.dflash_utils import parse_dflash_draft_config +from sglang.srt.speculative.domino_utils import ( + domino_greedy_rollout, + validate_domino_runtime, +) +from sglang.test.ci.ci_register import register_cpu_ci, register_cuda_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=10, suite="base-a-test-cpu") +register_cuda_ci(est_time=10, stage="base-b", runner_config="1-gpu-small") + + +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_public_config_fields(self): + parsed = parse_dflash_draft_config(draft_hf_config=_domino_config()) + self.assertTrue(parsed.is_domino) + self.assertTrue(parsed.shift_label) + self.assertEqual(parsed.pure_draft_prefix_len, 1) + self.assertEqual(parsed.gru_hidden_dim, 4) + self.assertEqual(parsed.emb_dim, 5) + self.assertEqual(parsed.block_size, 16) + + 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_plain_dflash_and_other_projectors_are_unchanged(self): + for projector_type in (None, "linear", "dspark"): + with self.subTest(projector_type=projector_type): + config = _domino_config( + dflash_config={ + "projector_type": projector_type, + "shift_label": None, + "pure_draft_prefix_len": None, + "gru_hidden_dim": None, + "emb_dim": None, + }, + emb_dim=None, + ) + parsed = parse_dflash_draft_config(draft_hf_config=config) + self.assertFalse(parsed.is_domino) + + 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_plain_dflash_does_not_require_projector_weights(self): + _projector_model(projector_type=None).load_weights([]) + + 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))]) + + +class TestDFlashDominoRollout(CustomTestCase): + def setUp(self): + torch.manual_seed(0) + self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + self.dtype = torch.bfloat16 if self.device.type == "cuda" else torch.float32 + self.hidden_size = 8 + self.gru_hidden_size = 4 + self.vocab_size = 31 + self.embedding = nn.Embedding( + self.vocab_size, self.hidden_size, device=self.device, dtype=self.dtype + ) + self.prefix_gru = nn.GRU( + self.hidden_size, + self.gru_hidden_size, + batch_first=True, + bias=False, + device=self.device, + dtype=self.dtype, + ) + self.embed_proj = nn.Sequential( + nn.Linear( + self.hidden_size + self.gru_hidden_size, + 5, + bias=False, + device=self.device, + dtype=self.dtype, + ), + nn.SiLU(), + nn.Linear( + 5, + self.vocab_size, + bias=False, + device=self.device, + dtype=self.dtype, + ), + ) + self.lm_head_weight = torch.randn( + self.vocab_size, + self.hidden_size, + device=self.device, + dtype=self.dtype, + ) + + def _oracle(self, draft_hidden, verified_ids, shift_label): + num_proposals = draft_hidden.shape[1] - 1 + start = 0 if shift_label else 1 + z = draft_hidden[:, start : start + num_proposals] + base_logits = F.linear(z, self.lm_head_weight) + first_ids = torch.argmax(base_logits[:, 0], dim=-1) + output = [first_ids] + if num_proposals == 1: + return first_ids[:, None] + + prefix_ids = torch.cat((verified_ids[:, None], first_ids[:, None]), dim=1) + _, state = self.prefix_gru(self.embedding(prefix_ids)) + for index in range(1, num_proposals): + bias = self.embed_proj(torch.cat((z[:, index], state[0]), dim=-1)) + next_ids = torch.argmax(base_logits[:, index] + bias, dim=-1) + output.append(next_ids) + if index + 1 < num_proposals: + _, state = self.prefix_gru(self.embedding(next_ids[:, None]), state) + return torch.stack(output, dim=1) + + def test_full_chain_matches_native_oracle(self): + for block_size in (2, 16): + for shift_label in (True, False): + with self.subTest(block_size=block_size, shift_label=shift_label): + draft_hidden = torch.randn( + 2, + block_size, + self.hidden_size, + device=self.device, + dtype=self.dtype, + ) + verified_ids = torch.tensor([2, 7], device=self.device) + actual = domino_greedy_rollout( + draft_hidden=draft_hidden, + verified_ids=verified_ids, + target_embedding=self.embedding, + lm_head_weight=self.lm_head_weight, + prefix_gru=self.prefix_gru, + embed_proj=self.embed_proj, + vocab_size=self.vocab_size, + shift_label=shift_label, + ) + expected = self._oracle( + draft_hidden, verified_ids, shift_label=shift_label + ) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + self.assertEqual(actual.shape, (2, block_size - 1)) + + first_hidden = draft_hidden[:, 0 if shift_label else 1] + first_expected = torch.argmax( + F.linear(first_hidden, self.lm_head_weight), dim=-1 + ) + torch.testing.assert_close( + actual[:, 0], first_expected, rtol=0, atol=0 + ) + + def test_batch_rollout_matches_each_individual_row(self): + draft_hidden = torch.randn( + 3, 7, self.hidden_size, device=self.device, dtype=self.dtype + ) + verified_ids = torch.tensor([1, 4, 9], device=self.device) + batched = domino_greedy_rollout( + draft_hidden=draft_hidden, + verified_ids=verified_ids, + target_embedding=self.embedding, + lm_head_weight=self.lm_head_weight, + prefix_gru=self.prefix_gru, + embed_proj=self.embed_proj, + vocab_size=self.vocab_size, + shift_label=True, + ) + individual = torch.cat( + [ + domino_greedy_rollout( + draft_hidden=draft_hidden[index : index + 1], + verified_ids=verified_ids[index : index + 1], + target_embedding=self.embedding, + lm_head_weight=self.lm_head_weight, + prefix_gru=self.prefix_gru, + embed_proj=self.embed_proj, + vocab_size=self.vocab_size, + shift_label=True, + ) + for index in range(3) + ], + dim=0, + ) + torch.testing.assert_close(batched, individual, rtol=0, atol=0) + + +class TestDFlashDominoRuntimeValidation(CustomTestCase): + def _modules(self, dtype=torch.bfloat16): + embedding = nn.Embedding(31, 8, dtype=dtype) + lm_head = nn.Linear(8, 31, bias=False, dtype=dtype) + prefix_gru = nn.GRU(8, 4, batch_first=True, bias=False, dtype=dtype) + embed_proj = nn.Sequential( + nn.Linear(12, 5, bias=False, dtype=dtype), + nn.SiLU(), + nn.Linear(5, 31, bias=False, dtype=dtype), + ) + return embedding, lm_head, prefix_gru, embed_proj + + def _validate(self, **overrides): + embedding, lm_head, prefix_gru, embed_proj = overrides.pop( + "modules", self._modules() + ) + args = { + "device": torch.device("cuda"), + "tp_size": 1, + "target_vocab_size": 31, + "draft_vocab_size": 31, + "hidden_size": 8, + "target_embedding": embedding, + "lm_head": lm_head, + "prefix_gru": prefix_gru, + "embed_proj": embed_proj, + } + args.update(overrides) + validate_domino_runtime(**args) + + def test_supported_runtime(self): + self._validate() + + def test_tp_and_vocab_mismatches_fail(self): + with self.assertRaisesRegex(ValueError, "TP=1"): + self._validate(tp_size=2) + with self.assertRaisesRegex(ValueError, "identical target and draft"): + self._validate(draft_vocab_size=30) + + def test_added_vocab_fails(self): + modules = self._modules() + modules[1].shard_indices = SimpleNamespace( + num_added_elements=1, + org_vocab_start_index=0, + num_org_elements=31, + ) + with self.assertRaisesRegex(ValueError, "added-vocab"): + self._validate(modules=modules) + + def test_embedding_vocab_mismatch_fails(self): + _, lm_head, prefix_gru, embed_proj = self._modules() + embedding = nn.Embedding(30, 8, dtype=torch.bfloat16) + with self.assertRaisesRegex(ValueError, "fewer rows"): + self._validate(modules=(embedding, lm_head, prefix_gru, embed_proj)) + + def test_non_bf16_fails(self): + with self.assertRaisesRegex(ValueError, "BF16"): + self._validate(modules=self._modules(dtype=torch.float32)) + + +if __name__ == "__main__": + unittest.main() From be8f681894b4d1c5284e1a7123de2af455852101 Mon Sep 17 00:00:00 2001 From: jianuo-huang <1170180988@qq.com> Date: Thu, 16 Jul 2026 00:34:20 +0800 Subject: [PATCH 2/7] perf: precompute Domino base logits --- python/sglang/srt/speculative/domino_utils.py | 20 +++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/python/sglang/srt/speculative/domino_utils.py b/python/sglang/srt/speculative/domino_utils.py index 6b60d666b106..816a892e3864 100644 --- a/python/sglang/srt/speculative/domino_utils.py +++ b/python/sglang/srt/speculative/domino_utils.py @@ -110,7 +110,7 @@ def domino_greedy_rollout( raise ValueError( f"draft_hidden must have shape [batch, block, hidden], got {tuple(draft_hidden.shape)}." ) - batch_size, block_size, _ = draft_hidden.shape + batch_size, block_size, hidden_size = draft_hidden.shape if verified_ids.shape != (batch_size,): raise ValueError( f"verified_ids must have shape ({batch_size},), got {tuple(verified_ids.shape)}." @@ -127,13 +127,15 @@ def domino_greedy_rollout( ) weight = lm_head_weight[: int(vocab_size)] + z_for_logits = z.to(weight.dtype) if z.dtype != weight.dtype else z + logits_input = ( + z_for_logits.transpose(0, 1) + .contiguous() + .view(num_proposals * batch_size, hidden_size) + ) + base_logits = F.linear(logits_input, weight).view(num_proposals, batch_size, -1) - def base_logits(hidden: torch.Tensor) -> torch.Tensor: - if hidden.dtype != weight.dtype: - hidden = hidden.to(weight.dtype) - return F.linear(hidden, weight) - - first_ids = torch.argmax(base_logits(z[:, 0, :]), dim=-1).to(torch.long) + first_ids = torch.argmax(base_logits[0], dim=-1).to(torch.long) proposals = [first_ids] if num_proposals == 1: return first_ids[:, None] @@ -144,9 +146,7 @@ def base_logits(hidden: torch.Tensor) -> torch.Tensor: for index in range(1, num_proposals): step_hidden = z[:, index, :] correction = embed_proj(torch.cat((step_hidden, gru_hidden[0]), dim=-1)) - next_ids = torch.argmax(base_logits(step_hidden) + correction, dim=-1).to( - torch.long - ) + next_ids = torch.argmax(base_logits[index] + correction, dim=-1).to(torch.long) proposals.append(next_ids) if index + 1 < num_proposals: _, gru_hidden = prefix_gru(target_embedding(next_ids[:, None]), gru_hidden) From b0faa37ac2d2a12d6e6600e9a3347dc60e1d3797 Mon Sep 17 00:00:00 2001 From: jianuo-huang <1170180988@qq.com> Date: Tue, 21 Jul 2026 13:50:06 +0800 Subject: [PATCH 3/7] feat: capture Domino rollout in draft CUDA graph --- .../srt/speculative/dflash_worker_v2.py | 79 +++++++++++- .../unit/spec/test_dflash_domino.py | 112 ++++++++++++++++++ 2 files changed, 188 insertions(+), 3 deletions(-) diff --git a/python/sglang/srt/speculative/dflash_worker_v2.py b/python/sglang/srt/speculative/dflash_worker_v2.py index b9c48c775b41..02a808b85e2b 100644 --- a/python/sglang/srt/speculative/dflash_worker_v2.py +++ b/python/sglang/srt/speculative/dflash_worker_v2.py @@ -151,6 +151,52 @@ def __call__(self, hidden_states, input_ids=None): self.out[:n].copy_(selected.view(-1)) +class _DominoDraftSampler: + """Capture-safe TP=1 Domino rollout over a fixed-size draft block.""" + + def __init__( + self, + *, + target_embedding, + lm_head_weight, + prefix_gru, + embed_proj, + vocab_size, + block_size, + shift_label, + max_bs, + ): + self.target_embedding = target_embedding + self.lm_head_weight = lm_head_weight + self.prefix_gru = prefix_gru + self.embed_proj = embed_proj + self.vocab_size = int(vocab_size) + self.block_size = int(block_size) + self.shift_label = bool(shift_label) + max_tokens = int(max_bs) * (self.block_size - 1) + self.out = torch.empty( + (max_tokens,), dtype=torch.int64, device=lm_head_weight.device + ) + + def __call__(self, hidden_states, input_ids=None): + if input_ids is None: + raise RuntimeError("Domino draft sampler requires block input_ids.") + bs = hidden_states.shape[0] // self.block_size + draft_hidden = hidden_states.view(bs, self.block_size, -1) + verified_ids = input_ids.view(bs, self.block_size)[:, 0] + proposals = domino_greedy_rollout( + draft_hidden=draft_hidden, + verified_ids=verified_ids, + target_embedding=self.target_embedding, + lm_head_weight=self.lm_head_weight, + prefix_gru=self.prefix_gru, + embed_proj=self.embed_proj, + vocab_size=self.vocab_size, + shift_label=self.shift_label, + ) + self.out[: bs * (self.block_size - 1)].copy_(proposals.reshape(-1)) + + class DFlashWorkerV2(BaseSpecWorker): """DFLASH speculative decoding worker (spec-v2). @@ -394,8 +440,6 @@ def _eager(reason): if envs.SGLANG_DFLASH_EAGER_DRAFT_SAMPLER.get(): return _eager("SGLANG_DFLASH_EAGER_DRAFT_SAMPLER=1") - if self._is_domino: - return _eager("Domino uses sequential eager rollout") if self.block_size <= 1: return _eager("block_size<=1") target_model = self._target_worker.model_runner.model @@ -406,6 +450,27 @@ def _eager(reason): # Quantized lm_head (FP8/INT) would break the static matmul. return _eager("quantized lm_head") tp_group = get_tp_group() + if self._is_domino: + if tp_group.world_size != 1: + return _eager("Domino cuda graph currently requires tp=1") + prefix_gru = self.draft_model.prefix_gru + embed_proj = self.draft_model.embed_proj + if prefix_gru is None or embed_proj is None: + return _eager("Domino projector modules are unavailable") + if self.ps.tp_rank == 0: + logger.info( + "DFLASH Domino rollout folded into the draft cuda graph (tp=1)." + ) + return _DominoDraftSampler( + target_embedding=target_model.get_input_embeddings(), + lm_head_weight=lm_head.weight, + prefix_gru=prefix_gru, + embed_proj=embed_proj, + vocab_size=int(self.model_runner.model_config.vocab_size), + block_size=self.block_size, + shift_label=self.draft_model.shift_label, + max_bs=max(self.server_args.cuda_graph_config.decode.bs), + ) if not hasattr(lm_head, "shard_indices"): if tp_group.world_size != 1: # No shard metadata to recover per-rank vocab offsets from. @@ -1650,7 +1715,15 @@ def forward_batch_generation( draft_out = self.draft_model_runner.forward(forward_batch) draft_logits_output = draft_out.logits_output - if self._is_domino: + if ( + self._is_domino + and self._draft_sampler is not None + and draft_out.can_run_graph + ): + draft_next = self._draft_sampler.out[ + : bs * (int(self.block_size) - 1) + ].view(bs, int(self.block_size) - 1) + elif self._is_domino: draft_hidden = draft_logits_output.hidden_states if draft_hidden is None: raise RuntimeError("DFLASH draft model returned no hidden states.") diff --git a/test/registered/unit/spec/test_dflash_domino.py b/test/registered/unit/spec/test_dflash_domino.py index 74d0a5cdba56..beb51582d4d5 100644 --- a/test/registered/unit/spec/test_dflash_domino.py +++ b/test/registered/unit/spec/test_dflash_domino.py @@ -292,6 +292,118 @@ def test_batch_rollout_matches_each_individual_row(self): ) torch.testing.assert_close(batched, individual, rtol=0, atol=0) + def test_capture_sampler_matches_eager_rollout(self): + from sglang.srt.speculative.dflash_worker_v2 import _DominoDraftSampler + + block_size = 7 + batch_size = 3 + for shift_label in (True, False): + with self.subTest(shift_label=shift_label): + draft_hidden = torch.randn( + batch_size, + block_size, + self.hidden_size, + device=self.device, + dtype=self.dtype, + ) + verified_ids = torch.tensor([1, 4, 9], device=self.device) + block_ids = torch.zeros( + batch_size, + block_size, + device=self.device, + dtype=torch.long, + ) + block_ids[:, 0].copy_(verified_ids) + 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=self.vocab_size, + block_size=block_size, + shift_label=shift_label, + max_bs=batch_size, + ) + sampler(draft_hidden.reshape(-1, self.hidden_size), block_ids.flatten()) + expected = domino_greedy_rollout( + draft_hidden=draft_hidden, + verified_ids=verified_ids, + target_embedding=self.embedding, + lm_head_weight=self.lm_head_weight, + prefix_gru=self.prefix_gru, + embed_proj=self.embed_proj, + vocab_size=self.vocab_size, + shift_label=shift_label, + ) + actual = sampler.out[: batch_size * (block_size - 1)].view( + batch_size, block_size - 1 + ) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + @unittest.skipUnless(torch.cuda.is_available(), "CUDA is required") + def test_capture_sampler_replays_with_new_inputs(self): + from sglang.srt.speculative.dflash_worker_v2 import _DominoDraftSampler + + block_size = 4 + batch_size = 2 + 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=self.vocab_size, + block_size=block_size, + shift_label=True, + max_bs=batch_size, + ) + static_hidden = torch.randn( + batch_size * block_size, + self.hidden_size, + device=self.device, + dtype=self.dtype, + ) + static_ids = torch.zeros( + batch_size * block_size, device=self.device, dtype=torch.long + ) + static_ids.view(batch_size, block_size)[:, 0].copy_( + torch.tensor([2, 7], device=self.device) + ) + + warmup_stream = torch.cuda.Stream() + warmup_stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(warmup_stream): + sampler(static_hidden, static_ids) + torch.cuda.current_stream().wait_stream(warmup_stream) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + sampler(static_hidden, static_ids) + + new_hidden = torch.randn_like(static_hidden) + new_ids = torch.zeros_like(static_ids) + new_ids.view(batch_size, block_size)[:, 0].copy_( + torch.tensor([3, 11], device=self.device) + ) + static_hidden.copy_(new_hidden) + static_ids.copy_(new_ids) + expected = domino_greedy_rollout( + draft_hidden=new_hidden.view(batch_size, block_size, self.hidden_size), + verified_ids=new_ids.view(batch_size, block_size)[:, 0], + target_embedding=self.embedding, + lm_head_weight=self.lm_head_weight, + prefix_gru=self.prefix_gru, + embed_proj=self.embed_proj, + vocab_size=self.vocab_size, + shift_label=True, + ) + graph.replay() + torch.cuda.synchronize() + + actual = sampler.out[: batch_size * (block_size - 1)].view( + batch_size, block_size - 1 + ) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + class TestDFlashDominoRuntimeValidation(CustomTestCase): def _modules(self, dtype=torch.bfloat16): From ea4b2f81b150d20c6bcffc17a9e32863cd4bd7d2 Mon Sep 17 00:00:00 2001 From: jianuo-huang <1170180988@qq.com> Date: Tue, 21 Jul 2026 16:31:28 +0800 Subject: [PATCH 4/7] perf: use GRU cell for Domino feedback --- python/sglang/srt/speculative/domino_utils.py | 18 +++++++++++++++++- 1 file changed, 17 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/speculative/domino_utils.py b/python/sglang/srt/speculative/domino_utils.py index 816a892e3864..597c7e3f03b5 100644 --- a/python/sglang/srt/speculative/domino_utils.py +++ b/python/sglang/srt/speculative/domino_utils.py @@ -5,6 +5,20 @@ from torch import nn +def _domino_gru_cell( + prefix_gru: nn.GRU, input: torch.Tensor, hidden: torch.Tensor +) -> torch.Tensor: + """Run one feedback step without cuDNN's per-call RNN weight packing.""" + return torch.ops.aten.gru_cell.default( + input, + hidden, + prefix_gru.weight_ih_l0, + prefix_gru.weight_hh_l0, + prefix_gru.bias_ih_l0 if prefix_gru.bias else None, + prefix_gru.bias_hh_l0 if prefix_gru.bias else None, + ) + + def validate_domino_runtime( *, device: torch.device, @@ -149,6 +163,8 @@ def domino_greedy_rollout( next_ids = torch.argmax(base_logits[index] + correction, dim=-1).to(torch.long) proposals.append(next_ids) if index + 1 < num_proposals: - _, gru_hidden = prefix_gru(target_embedding(next_ids[:, None]), gru_hidden) + gru_hidden = _domino_gru_cell( + prefix_gru, target_embedding(next_ids), gru_hidden[0] + )[None] return torch.stack(proposals, dim=1) From d7fb3f755c41e8bcea25894f4695ba19bf1cd011 Mon Sep 17 00:00:00 2001 From: jianuo-huang <1170180988@qq.com> Date: Tue, 21 Jul 2026 22:17:02 +0800 Subject: [PATCH 5/7] perf: add Domino candidate vocabulary pool --- python/sglang/srt/server_args.py | 4 + .../srt/speculative/dflash_worker_v2.py | 16 +- python/sglang/srt/speculative/domino_utils.py | 50 ++++- .../unit/spec/test_dflash_domino.py | 203 ++++++++++++++++-- 4 files changed, 246 insertions(+), 27 deletions(-) diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 9c8801baef1f..5488776db5b0 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -1742,6 +1742,10 @@ class ServerArgs: Optional[int], "DFLASH only. Block size (verify window length). Alias of --speculative-num-draft-tokens for DFLASH.", ] = None + speculative_domino_candidate_pool_size: A[ + int, + "Domino only. Size of the approximate block-shared base-logit candidate pool. Set to 0 to score the full vocabulary.", + ] = 2048 speculative_dspark_block_size: A[ Optional[int], "DSPARK only. Draft block size gamma (number of proposed draft tokens). The verify window is gamma + 1, so this sets --speculative-num-draft-tokens = gamma + 1. Omit to auto-infer gamma from the draft checkpoint block_size.", diff --git a/python/sglang/srt/speculative/dflash_worker_v2.py b/python/sglang/srt/speculative/dflash_worker_v2.py index 02a808b85e2b..a7e3ff25cf6d 100644 --- a/python/sglang/srt/speculative/dflash_worker_v2.py +++ b/python/sglang/srt/speculative/dflash_worker_v2.py @@ -165,6 +165,7 @@ def __init__( block_size, shift_label, max_bs, + candidate_pool_size=2048, ): self.target_embedding = target_embedding self.lm_head_weight = lm_head_weight @@ -173,6 +174,7 @@ def __init__( self.vocab_size = int(vocab_size) self.block_size = int(block_size) self.shift_label = bool(shift_label) + self.candidate_pool_size = int(candidate_pool_size) max_tokens = int(max_bs) * (self.block_size - 1) self.out = torch.empty( (max_tokens,), dtype=torch.int64, device=lm_head_weight.device @@ -193,6 +195,7 @@ def __call__(self, hidden_states, input_ids=None): embed_proj=self.embed_proj, vocab_size=self.vocab_size, shift_label=self.shift_label, + candidate_pool_size=self.candidate_pool_size, ) self.out[: bs * (self.block_size - 1)].copy_(proposals.reshape(-1)) @@ -246,7 +249,15 @@ def __init__( draft_hf_config=self.draft_model_runner.model_config.hf_config ) self._is_domino = draft_config.is_domino + self.domino_candidate_pool_size = int( + server_args.speculative_domino_candidate_pool_size + ) if self._is_domino: + if self.domino_candidate_pool_size < 0: + raise ValueError( + "--speculative-domino-candidate-pool-size must be non-negative, " + f"got {self.domino_candidate_pool_size}." + ) target_model = self.target_worker.model_runner.model target_embedding = target_model.get_input_embeddings() lm_head = getattr(target_model, "lm_head", None) @@ -307,7 +318,8 @@ def __init__( ) if self._is_domino: logger.info( - "DFLASH Domino rollout enabled (eager BF16, TP=1, full-vocabulary greedy)." + "DFLASH Domino rollout enabled (eager BF16, TP=1, block-shared candidate pool size=%s).", + self.domino_candidate_pool_size, ) logger.info( "DFLASH draft runner ready. mask_token=%s, mask_token_id=%s, mask_token_id_override=%s", @@ -470,6 +482,7 @@ def _eager(reason): block_size=self.block_size, shift_label=self.draft_model.shift_label, max_bs=max(self.server_args.cuda_graph_config.decode.bs), + candidate_pool_size=self.domino_candidate_pool_size, ) if not hasattr(lm_head, "shard_indices"): if tp_group.world_size != 1: @@ -1741,6 +1754,7 @@ def forward_batch_generation( embed_proj=embed_proj, vocab_size=int(self.model_runner.model_config.vocab_size), shift_label=bool(self.draft_model.shift_label), + candidate_pool_size=self.domino_candidate_pool_size, ) elif self._draft_sampler is not None and draft_out.can_run_graph: draft_next = self._draft_sampler.out[ diff --git a/python/sglang/srt/speculative/domino_utils.py b/python/sglang/srt/speculative/domino_utils.py index 597c7e3f03b5..90a0b83f3719 100644 --- a/python/sglang/srt/speculative/domino_utils.py +++ b/python/sglang/srt/speculative/domino_utils.py @@ -4,6 +4,8 @@ import torch.nn.functional as F from torch import nn +_DOMINO_CANDIDATE_POOL_SIZE = 2048 + def _domino_gru_cell( prefix_gru: nn.GRU, input: torch.Tensor, hidden: torch.Tensor @@ -118,8 +120,9 @@ def domino_greedy_rollout( embed_proj: nn.Sequential, vocab_size: int, shift_label: bool, + candidate_pool_size: int = _DOMINO_CANDIDATE_POOL_SIZE, ) -> torch.Tensor: - """Generate a Domino chain with native GRU and full-vocabulary greedy logits.""" + """Generate a Domino chain using one block-shared base-logit candidate pool.""" if draft_hidden.ndim != 3: raise ValueError( f"draft_hidden must have shape [batch, block, hidden], got {tuple(draft_hidden.shape)}." @@ -133,6 +136,13 @@ def domino_greedy_rollout( num_proposals = int(block_size) - 1 if num_proposals < 1: raise ValueError(f"Domino requires block_size > 1, got {block_size}.") + candidate_pool_size = int(candidate_pool_size) + if candidate_pool_size < 0: + raise ValueError( + "Domino candidate_pool_size must be non-negative, " + f"got {candidate_pool_size}." + ) + candidate_pool_size = min(candidate_pool_size, int(vocab_size)) start = 0 if shift_label else 1 z = draft_hidden[:, start : start + num_proposals, :] if int(z.shape[1]) != num_proposals: @@ -154,13 +164,47 @@ def domino_greedy_rollout( if num_proposals == 1: return first_ids[:, None] + candidate_ids = None + candidate_base = None + candidate_weight = None + if 0 < candidate_pool_size < int(vocab_size): + feedback_logits = base_logits[1:] + candidate_ids = torch.topk( + feedback_logits.amax(dim=0), + k=candidate_pool_size, + dim=-1, + sorted=False, + ).indices.contiguous() + candidate_base = torch.gather( + feedback_logits.transpose(0, 1), + 2, + candidate_ids[:, None, :].expand(-1, num_proposals - 1, -1), + ).transpose(0, 1) + candidate_weight = F.embedding(candidate_ids, embed_proj[2].weight) + prefix_ids = torch.stack((verified_ids, first_ids), dim=1) _, gru_hidden = prefix_gru(target_embedding(prefix_ids)) for index in range(1, num_proposals): step_hidden = z[:, index, :] - correction = embed_proj(torch.cat((step_hidden, gru_hidden[0]), dim=-1)) - next_ids = torch.argmax(base_logits[index] + correction, dim=-1).to(torch.long) + correction_hidden = embed_proj[1]( + embed_proj[0](torch.cat((step_hidden, gru_hidden[0]), dim=-1)) + ) + if candidate_ids is None: + correction = embed_proj[2](correction_hidden) + next_ids = torch.argmax(base_logits[index] + correction, dim=-1).to( + torch.long + ) + else: + correction = torch.bmm( + candidate_weight, correction_hidden.unsqueeze(-1) + ).squeeze(-1) + candidate_position = torch.argmax( + candidate_base[index - 1] + correction, dim=-1 + ) + next_ids = torch.gather( + candidate_ids, 1, candidate_position[:, None] + ).squeeze(1) proposals.append(next_ids) if index + 1 < num_proposals: gru_hidden = _domino_gru_cell( diff --git a/test/registered/unit/spec/test_dflash_domino.py b/test/registered/unit/spec/test_dflash_domino.py index beb51582d4d5..9a484e4d7987 100644 --- a/test/registered/unit/spec/test_dflash_domino.py +++ b/test/registered/unit/spec/test_dflash_domino.py @@ -203,7 +203,9 @@ def setUp(self): dtype=self.dtype, ) - def _oracle(self, draft_hidden, verified_ids, shift_label): + def _oracle( + self, draft_hidden, verified_ids, shift_label, candidate_pool_size=None + ): num_proposals = draft_hidden.shape[1] - 1 start = 0 if shift_label else 1 z = draft_hidden[:, start : start + num_proposals] @@ -213,11 +215,45 @@ def _oracle(self, draft_hidden, verified_ids, shift_label): if num_proposals == 1: return first_ids[:, None] + candidate_ids = None + candidate_base = None + candidate_weight = None + if candidate_pool_size is not None: + candidate_pool_size = min(candidate_pool_size, self.vocab_size) + if 0 < candidate_pool_size < self.vocab_size: + feedback_logits = base_logits[:, 1:] + candidate_ids = torch.topk( + feedback_logits.amax(dim=1), + k=candidate_pool_size, + dim=-1, + sorted=False, + ).indices + candidate_base = torch.gather( + feedback_logits, + 2, + candidate_ids[:, None, :].expand(-1, num_proposals - 1, -1), + ) + candidate_weight = F.embedding(candidate_ids, self.embed_proj[2].weight) + prefix_ids = torch.cat((verified_ids[:, None], first_ids[:, None]), dim=1) _, state = self.prefix_gru(self.embedding(prefix_ids)) for index in range(1, num_proposals): - bias = self.embed_proj(torch.cat((z[:, index], state[0]), dim=-1)) - next_ids = torch.argmax(base_logits[:, index] + bias, dim=-1) + correction_hidden = self.embed_proj[1]( + self.embed_proj[0](torch.cat((z[:, index], state[0]), dim=-1)) + ) + if candidate_ids is None: + bias = self.embed_proj[2](correction_hidden) + next_ids = torch.argmax(base_logits[:, index] + bias, dim=-1) + else: + bias = torch.bmm( + candidate_weight, correction_hidden.unsqueeze(-1) + ).squeeze(-1) + candidate_position = torch.argmax( + candidate_base[:, index - 1] + bias, dim=-1 + ) + next_ids = torch.gather( + candidate_ids, 1, candidate_position[:, None] + ).squeeze(1) output.append(next_ids) if index + 1 < num_proposals: _, state = self.prefix_gru(self.embedding(next_ids[:, None]), state) @@ -259,38 +295,112 @@ def test_full_chain_matches_native_oracle(self): actual[:, 0], first_expected, rtol=0, atol=0 ) - def test_batch_rollout_matches_each_individual_row(self): + def test_block_candidate_pool_matches_oracle(self): + block_size = 7 draft_hidden = torch.randn( - 3, 7, self.hidden_size, device=self.device, dtype=self.dtype + 3, + block_size, + self.hidden_size, + device=self.device, + dtype=self.dtype, ) verified_ids = torch.tensor([1, 4, 9], device=self.device) - batched = domino_greedy_rollout( - draft_hidden=draft_hidden, - verified_ids=verified_ids, - target_embedding=self.embedding, - lm_head_weight=self.lm_head_weight, - prefix_gru=self.prefix_gru, - embed_proj=self.embed_proj, - vocab_size=self.vocab_size, - shift_label=True, + for shift_label in (True, False): + with self.subTest(shift_label=shift_label): + actual = domino_greedy_rollout( + draft_hidden=draft_hidden, + verified_ids=verified_ids, + target_embedding=self.embedding, + lm_head_weight=self.lm_head_weight, + prefix_gru=self.prefix_gru, + embed_proj=self.embed_proj, + vocab_size=self.vocab_size, + shift_label=shift_label, + candidate_pool_size=5, + ) + expected = self._oracle( + draft_hidden, + verified_ids, + shift_label=shift_label, + candidate_pool_size=5, + ) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + first_hidden = draft_hidden[:, 0 if shift_label else 1] + first_expected = torch.argmax( + F.linear(first_hidden, self.lm_head_weight), dim=-1 + ) + torch.testing.assert_close(actual[:, 0], first_expected, rtol=0, atol=0) + + def test_candidate_pool_boundaries(self): + draft_hidden = torch.randn( + 2, 7, self.hidden_size, device=self.device, dtype=self.dtype ) - individual = torch.cat( - [ - domino_greedy_rollout( - draft_hidden=draft_hidden[index : index + 1], - verified_ids=verified_ids[index : index + 1], + verified_ids = torch.tensor([2, 7], device=self.device) + expected = self._oracle(draft_hidden, verified_ids, shift_label=True) + for candidate_pool_size in (0, self.vocab_size, self.vocab_size + 1): + with self.subTest(candidate_pool_size=candidate_pool_size): + actual = domino_greedy_rollout( + draft_hidden=draft_hidden, + verified_ids=verified_ids, target_embedding=self.embedding, lm_head_weight=self.lm_head_weight, prefix_gru=self.prefix_gru, embed_proj=self.embed_proj, vocab_size=self.vocab_size, shift_label=True, + candidate_pool_size=candidate_pool_size, ) - for index in range(3) - ], - dim=0, + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + with self.assertRaisesRegex(ValueError, "non-negative"): + domino_greedy_rollout( + draft_hidden=draft_hidden, + verified_ids=verified_ids, + target_embedding=self.embedding, + lm_head_weight=self.lm_head_weight, + prefix_gru=self.prefix_gru, + embed_proj=self.embed_proj, + vocab_size=self.vocab_size, + shift_label=True, + candidate_pool_size=-1, + ) + + def test_batch_rollout_matches_each_individual_row(self): + draft_hidden = torch.randn( + 3, 7, self.hidden_size, device=self.device, dtype=self.dtype ) - torch.testing.assert_close(batched, individual, rtol=0, atol=0) + verified_ids = torch.tensor([1, 4, 9], device=self.device) + for candidate_pool_size in (0, 5): + with self.subTest(candidate_pool_size=candidate_pool_size): + batched = domino_greedy_rollout( + draft_hidden=draft_hidden, + verified_ids=verified_ids, + target_embedding=self.embedding, + lm_head_weight=self.lm_head_weight, + prefix_gru=self.prefix_gru, + embed_proj=self.embed_proj, + vocab_size=self.vocab_size, + shift_label=True, + candidate_pool_size=candidate_pool_size, + ) + individual = torch.cat( + [ + domino_greedy_rollout( + draft_hidden=draft_hidden[index : index + 1], + verified_ids=verified_ids[index : index + 1], + target_embedding=self.embedding, + lm_head_weight=self.lm_head_weight, + prefix_gru=self.prefix_gru, + embed_proj=self.embed_proj, + vocab_size=self.vocab_size, + shift_label=True, + candidate_pool_size=candidate_pool_size, + ) + for index in range(3) + ], + dim=0, + ) + torch.testing.assert_close(batched, individual, rtol=0, atol=0) def test_capture_sampler_matches_eager_rollout(self): from sglang.srt.speculative.dflash_worker_v2 import _DominoDraftSampler @@ -323,6 +433,7 @@ def test_capture_sampler_matches_eager_rollout(self): block_size=block_size, shift_label=shift_label, max_bs=batch_size, + candidate_pool_size=5, ) sampler(draft_hidden.reshape(-1, self.hidden_size), block_ids.flatten()) expected = domino_greedy_rollout( @@ -334,12 +445,56 @@ def test_capture_sampler_matches_eager_rollout(self): embed_proj=self.embed_proj, vocab_size=self.vocab_size, shift_label=shift_label, + candidate_pool_size=5, ) actual = sampler.out[: batch_size * (block_size - 1)].view( batch_size, block_size - 1 ) torch.testing.assert_close(actual, expected, rtol=0, atol=0) + @unittest.skipUnless(torch.cuda.is_available(), "CUDA is required") + def test_candidate_pool_cuda_graph_replay(self): + block_size = 7 + batch_size = 3 + static_hidden = torch.randn( + batch_size, + block_size, + self.hidden_size, + device=self.device, + dtype=self.dtype, + ) + static_ids = torch.tensor([1, 4, 9], device=self.device) + + def rollout(): + return domino_greedy_rollout( + draft_hidden=static_hidden, + verified_ids=static_ids, + target_embedding=self.embedding, + lm_head_weight=self.lm_head_weight, + prefix_gru=self.prefix_gru, + embed_proj=self.embed_proj, + vocab_size=self.vocab_size, + shift_label=True, + candidate_pool_size=5, + ) + + rollout() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual = rollout() + + static_hidden.copy_(torch.randn_like(static_hidden)) + static_ids.copy_(torch.tensor([2, 7, 11], device=self.device)) + expected = self._oracle( + static_hidden, + static_ids, + shift_label=True, + candidate_pool_size=5, + ) + graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + @unittest.skipUnless(torch.cuda.is_available(), "CUDA is required") def test_capture_sampler_replays_with_new_inputs(self): from sglang.srt.speculative.dflash_worker_v2 import _DominoDraftSampler @@ -355,6 +510,7 @@ def test_capture_sampler_replays_with_new_inputs(self): block_size=block_size, shift_label=True, max_bs=batch_size, + candidate_pool_size=5, ) static_hidden = torch.randn( batch_size * block_size, @@ -395,6 +551,7 @@ def test_capture_sampler_replays_with_new_inputs(self): embed_proj=self.embed_proj, vocab_size=self.vocab_size, shift_label=True, + candidate_pool_size=5, ) graph.replay() torch.cuda.synchronize() From 886231493108ed818e50eac7b58f62d1c935f71a Mon Sep 17 00:00:00 2001 From: jianuo-huang <1170180988@qq.com> Date: Wed, 22 Jul 2026 03:51:28 +0800 Subject: [PATCH 6/7] feat: support tensor-parallel Domino rollout --- .../srt/speculative/dflash_worker_v2.py | 46 ++- python/sglang/srt/speculative/domino_utils.py | 298 ++++++++++++++++-- .../unit/spec/test_dflash_domino.py | 200 +++++++++++- 3 files changed, 519 insertions(+), 25 deletions(-) diff --git a/python/sglang/srt/speculative/dflash_worker_v2.py b/python/sglang/srt/speculative/dflash_worker_v2.py index a7e3ff25cf6d..fa1cd1fab908 100644 --- a/python/sglang/srt/speculative/dflash_worker_v2.py +++ b/python/sglang/srt/speculative/dflash_worker_v2.py @@ -152,7 +152,7 @@ def __call__(self, hidden_states, input_ids=None): class _DominoDraftSampler: - """Capture-safe TP=1 Domino rollout over a fixed-size draft block.""" + """Capture-safe Domino rollout over a fixed-size draft block.""" def __init__( self, @@ -166,6 +166,10 @@ def __init__( shift_label, max_bs, candidate_pool_size=2048, + tp_group=None, + lm_head_org_vocab_start=0, + lm_head_num_org=None, + lm_head_num_org_padded=None, ): self.target_embedding = target_embedding self.lm_head_weight = lm_head_weight @@ -175,6 +179,10 @@ def __init__( self.block_size = int(block_size) self.shift_label = bool(shift_label) self.candidate_pool_size = int(candidate_pool_size) + self.tp_group = tp_group + self.lm_head_org_vocab_start = int(lm_head_org_vocab_start) + self.lm_head_num_org = lm_head_num_org + self.lm_head_num_org_padded = lm_head_num_org_padded max_tokens = int(max_bs) * (self.block_size - 1) self.out = torch.empty( (max_tokens,), dtype=torch.int64, device=lm_head_weight.device @@ -196,6 +204,10 @@ def __call__(self, hidden_states, input_ids=None): vocab_size=self.vocab_size, shift_label=self.shift_label, candidate_pool_size=self.candidate_pool_size, + tp_group=self.tp_group, + lm_head_org_vocab_start=self.lm_head_org_vocab_start, + lm_head_num_org=self.lm_head_num_org, + lm_head_num_org_padded=self.lm_head_num_org_padded, ) self.out[: bs * (self.block_size - 1)].copy_(proposals.reshape(-1)) @@ -270,6 +282,7 @@ def __init__( validate_domino_runtime( device=torch.device(self.device), tp_size=int(get_tp_group().world_size), + tp_rank=int(self.ps.tp_rank), target_vocab_size=int(self.model_runner.model_config.vocab_size), draft_vocab_size=int(self.draft_model_runner.model_config.vocab_size), hidden_size=int(self.draft_model.config.hidden_size), @@ -318,7 +331,8 @@ def __init__( ) if self._is_domino: logger.info( - "DFLASH Domino rollout enabled (eager BF16, TP=1, block-shared candidate pool size=%s).", + "DFLASH Domino rollout enabled (eager BF16, TP=%s, block-shared candidate pool size=%s).", + int(get_tp_group().world_size), self.domino_candidate_pool_size, ) logger.info( @@ -463,16 +477,16 @@ def _eager(reason): return _eager("quantized lm_head") tp_group = get_tp_group() if self._is_domino: - if tp_group.world_size != 1: - return _eager("Domino cuda graph currently requires tp=1") prefix_gru = self.draft_model.prefix_gru embed_proj = self.draft_model.embed_proj if prefix_gru is None or embed_proj is None: return _eager("Domino projector modules are unavailable") if self.ps.tp_rank == 0: logger.info( - "DFLASH Domino rollout folded into the draft cuda graph (tp=1)." + "DFLASH Domino rollout folded into the draft cuda graph (tp=%s).", + int(tp_group.world_size), ) + shard = getattr(lm_head, "shard_indices", None) return _DominoDraftSampler( target_embedding=target_model.get_input_embeddings(), lm_head_weight=lm_head.weight, @@ -483,6 +497,16 @@ def _eager(reason): shift_label=self.draft_model.shift_label, max_bs=max(self.server_args.cuda_graph_config.decode.bs), candidate_pool_size=self.domino_candidate_pool_size, + tp_group=tp_group, + lm_head_org_vocab_start=( + int(shard.org_vocab_start_index) if shard is not None else 0 + ), + lm_head_num_org=( + int(shard.num_org_elements) if shard is not None else None + ), + lm_head_num_org_padded=( + int(shard.num_org_elements_padded) if shard is not None else None + ), ) if not hasattr(lm_head, "shard_indices"): if tp_group.world_size != 1: @@ -1745,6 +1769,8 @@ def forward_batch_generation( embed_proj = self.draft_model.embed_proj if prefix_gru is None or embed_proj is None: raise RuntimeError("DFLASH Domino projector modules are unavailable.") + tp_group = get_tp_group() + shard = getattr(lm_head, "shard_indices", None) draft_next = domino_greedy_rollout( draft_hidden=draft_hidden, verified_ids=block_ids[:, 0], @@ -1755,6 +1781,16 @@ def forward_batch_generation( vocab_size=int(self.model_runner.model_config.vocab_size), shift_label=bool(self.draft_model.shift_label), candidate_pool_size=self.domino_candidate_pool_size, + tp_group=tp_group, + lm_head_org_vocab_start=( + int(shard.org_vocab_start_index) if shard is not None else 0 + ), + lm_head_num_org=( + int(shard.num_org_elements) if shard is not None else None + ), + lm_head_num_org_padded=( + int(shard.num_org_elements_padded) if shard is not None else None + ), ) elif self._draft_sampler is not None and draft_out.can_run_graph: draft_next = self._draft_sampler.out[ diff --git a/python/sglang/srt/speculative/domino_utils.py b/python/sglang/srt/speculative/domino_utils.py index 90a0b83f3719..ced42f4df39a 100644 --- a/python/sglang/srt/speculative/domino_utils.py +++ b/python/sglang/srt/speculative/domino_utils.py @@ -5,6 +5,8 @@ from torch import nn _DOMINO_CANDIDATE_POOL_SIZE = 2048 +# This is a throughput policy for the logical gathered tensor, not a peak-memory cap. +_DOMINO_TP_FULL_BASE_LOGITS_MAX_BYTES = 96 * 1024 * 1024 def _domino_gru_cell( @@ -21,10 +23,108 @@ def _domino_gru_cell( ) +def _domino_tp_first_ids( + local_logits: torch.Tensor, + *, + org_vocab_start: int, + num_org: int, + tp_group, +) -> torch.Tensor: + """Select the full-vocab first argmax from contiguous vocab shards.""" + local_max, local_arg = torch.max(local_logits[:, :num_org], dim=-1) + local_ids = local_arg.to(torch.int64) + int(org_vocab_start) + tp_size = int(tp_group.world_size) + batch_size = int(local_logits.shape[0]) + gathered_max = torch.empty( + (tp_size * batch_size,), dtype=local_max.dtype, device=local_max.device + ) + gathered_ids = torch.empty( + (tp_size * batch_size,), dtype=torch.int64, device=local_max.device + ) + tp_group.all_gather_into_tensor(gathered_max, local_max.contiguous()) + tp_group.all_gather_into_tensor(gathered_ids, local_ids.contiguous()) + gathered_max = gathered_max.view(tp_size, batch_size) + gathered_ids = gathered_ids.view(tp_size, batch_size) + best_rank = torch.argmax(gathered_max, dim=0, keepdim=True) + return torch.gather(gathered_ids, 0, best_rank).squeeze(0) + + +def _domino_tp_candidate_state( + local_feedback_logits: torch.Tensor, + *, + candidate_pool_size: int, + org_vocab_start: int, + num_org: int, + tp_group, +) -> tuple[torch.Tensor, torch.Tensor]: + """Build one strict global candidate pool from vocab-sharded base logits.""" + num_steps, batch_size, local_vocab_size = local_feedback_logits.shape + tp_size = int(tp_group.world_size) + local_k = min(int(candidate_pool_size), int(local_vocab_size)) + if int(candidate_pool_size) > tp_size * local_k: + raise ValueError("Domino TP shards do not cover the requested candidate pool.") + + pool_scores = local_feedback_logits.amax(dim=0) + if num_org < local_vocab_size: + pool_scores[:, num_org:].fill_(float("-inf")) + local_scores, local_positions = torch.topk( + pool_scores, k=local_k, dim=-1, sorted=False + ) + local_ids = local_positions.to(torch.int64) + int(org_vocab_start) + + gathered_scores = torch.empty( + (tp_size * batch_size, local_k), + dtype=local_scores.dtype, + device=local_scores.device, + ) + gathered_ids = torch.empty( + (tp_size * batch_size, local_k), + dtype=torch.int64, + device=local_scores.device, + ) + tp_group.all_gather_into_tensor(gathered_scores, local_scores.contiguous()) + tp_group.all_gather_into_tensor(gathered_ids, local_ids.contiguous()) + gathered_scores = ( + gathered_scores.view(tp_size, batch_size, local_k) + .permute(1, 0, 2) + .reshape(batch_size, tp_size * local_k) + ) + gathered_ids = ( + gathered_ids.view(tp_size, batch_size, local_k) + .permute(1, 0, 2) + .reshape(batch_size, tp_size * local_k) + ) + global_positions = torch.topk( + gathered_scores, + k=int(candidate_pool_size), + dim=-1, + sorted=False, + ).indices + candidate_ids = torch.gather(gathered_ids, 1, global_positions).contiguous() + + owned = (candidate_ids >= int(org_vocab_start)) & ( + candidate_ids < int(org_vocab_start + num_org) + ) + local_positions = (candidate_ids - int(org_vocab_start)).clamp( + 0, max(num_org - 1, 0) + ) + candidate_base = torch.gather( + local_feedback_logits.transpose(0, 1), + 2, + local_positions[:, None, :].expand(-1, num_steps, -1), + ) + # Each global candidate is owned by exactly one vocab shard, so SUM + # reconstructs its base logit without gathering the full vocabulary. + candidate_base.masked_fill_(~owned[:, None, :], 0) + candidate_base = tp_group.all_reduce(candidate_base.contiguous()) + return candidate_ids, candidate_base.transpose(0, 1) + + def validate_domino_runtime( *, device: torch.device, tp_size: int, + tp_rank: int, target_vocab_size: int, draft_vocab_size: int, hidden_size: int, @@ -36,8 +136,14 @@ def validate_domino_runtime( """Validate the deliberately narrow correctness-first Domino runtime.""" if device.type != "cuda": raise ValueError(f"DFLASH Domino currently requires CUDA, got {device}.") - if int(tp_size) != 1: - raise ValueError(f"DFLASH Domino currently requires TP=1, got TP={tp_size}.") + tp_size = int(tp_size) + if tp_size < 1: + raise ValueError(f"DFLASH Domino requires TP>=1, got TP={tp_size}.") + tp_rank = int(tp_rank) + if not 0 <= tp_rank < tp_size: + raise ValueError( + f"DFLASH Domino requires 0<=TP rank int(target_vocab_size) + or num_org_padded < num_org + ): + raise ValueError("DFLASH Domino lm_head original-vocab shard is invalid.") + if int(lm_head_weight.shape[0]) < num_org_padded: + raise ValueError( + "DFLASH Domino lm_head weight is smaller than its padded vocab shard." + ) + expected_start = tp_rank * num_org_padded + expected_end = min(expected_start + num_org_padded, int(target_vocab_size)) + if ( + org_vocab_start != expected_start + or org_vocab_end != expected_end + or num_org != expected_end - expected_start + ): + raise ValueError( + "DFLASH Domino lm_head vocab shard does not match its TP rank." + ) + else: + if tp_size != 1: + raise ValueError("DFLASH Domino requires lm_head shard metadata for TP>1.") + if int(lm_head_weight.shape[0]) != int(target_vocab_size): + raise ValueError( + "DFLASH Domino lm_head row count must equal the target vocab size, " + f"got rows={int(lm_head_weight.shape[0])}, vocab={target_vocab_size}." + ) - if int(embedding_weight.shape[0]) < int(target_vocab_size): + embedding_shard = getattr(target_embedding, "shard_indices", None) + if embedding_shard is not None: + if ( + int(getattr(target_embedding, "num_added_embeddings", 0)) != 0 + or int(embedding_shard.num_added_elements) != 0 + ): + raise ValueError( + "DFLASH Domino does not support added-vocab embedding shards." + ) + if int(getattr(target_embedding, "org_vocab_size", target_vocab_size)) != int( + target_vocab_size + ): + raise ValueError( + "DFLASH Domino embedding original vocab size does not match the target." + ) + embedding_tp_size = int(getattr(target_embedding, "tp_size", tp_size)) + if embedding_tp_size not in (1, tp_size): + raise ValueError( + "DFLASH Domino embedding TP size does not match the runtime TP size." + ) + required_embedding_rows = ( + int(target_vocab_size) + if embedding_tp_size == 1 + else int(embedding_shard.num_org_elements_padded) + ) + if int(embedding_weight.shape[0]) < required_embedding_rows: + raise ValueError( + "DFLASH Domino embedding weight is smaller than its padded vocab shard." + ) + elif int(embedding_weight.shape[0]) < int(target_vocab_size): + if tp_size != 1: + raise ValueError( + "DFLASH Domino requires embedding shard metadata for TP>1." + ) raise ValueError( "DFLASH Domino target embedding has fewer rows than the target vocab " f"size: rows={int(embedding_weight.shape[0])}, vocab={target_vocab_size}." @@ -121,6 +310,10 @@ def domino_greedy_rollout( vocab_size: int, shift_label: bool, candidate_pool_size: int = _DOMINO_CANDIDATE_POOL_SIZE, + tp_group=None, + lm_head_org_vocab_start: int = 0, + lm_head_num_org: int | None = None, + lm_head_num_org_padded: int | None = None, ) -> torch.Tensor: """Generate a Domino chain using one block-shared base-logit candidate pool.""" if draft_hidden.ndim != 3: @@ -150,24 +343,93 @@ def domino_greedy_rollout( "Domino draft hidden states do not contain enough proposal positions." ) - weight = lm_head_weight[: int(vocab_size)] + tp_size = int(tp_group.world_size) if tp_group is not None else 1 + if tp_size > 1 and (lm_head_num_org is None or lm_head_num_org_padded is None): + raise ValueError( + "Domino TP rollout requires local lm_head vocab shard metadata." + ) + local_vocab_size = int(lm_head_num_org_padded) if tp_size > 1 else int(vocab_size) + if tp_size > 1: + num_org = int(lm_head_num_org) + org_vocab_start = int(lm_head_org_vocab_start) + if ( + num_org <= 0 + or num_org > local_vocab_size + or org_vocab_start < 0 + or org_vocab_start + num_org > int(vocab_size) + ): + raise ValueError( + "Domino TP rollout received an invalid lm_head vocab shard." + ) + if int(lm_head_weight.shape[0]) < local_vocab_size: + raise ValueError( + "Domino lm_head weight is smaller than its padded vocab shard." + ) + weight = lm_head_weight[:local_vocab_size] z_for_logits = z.to(weight.dtype) if z.dtype != weight.dtype else z logits_input = ( z_for_logits.transpose(0, 1) .contiguous() .view(num_proposals * batch_size, hidden_size) ) - base_logits = F.linear(logits_input, weight).view(num_proposals, batch_size, -1) + local_logits = F.linear(logits_input, weight).view( + num_proposals, batch_size, local_vocab_size + ) + full_base_logits_bytes = ( + num_proposals * batch_size * int(vocab_size) * local_logits.element_size() + ) + use_tp_candidate_pool = ( + tp_size > 1 + and 0 < candidate_pool_size < int(vocab_size) + and full_base_logits_bytes > _DOMINO_TP_FULL_BASE_LOGITS_MAX_BYTES + ) + first_ids = None + candidate_ids = None + candidate_base = None + if tp_size == 1: + base_logits = local_logits[:, :, : int(vocab_size)] + elif use_tp_candidate_pool: + first_ids = _domino_tp_first_ids( + local_logits[0], + org_vocab_start=int(lm_head_org_vocab_start), + num_org=int(lm_head_num_org), + tp_group=tp_group, + ) + if num_proposals > 1: + candidate_ids, candidate_base = _domino_tp_candidate_state( + local_logits[1:], + candidate_pool_size=candidate_pool_size, + org_vocab_start=int(lm_head_org_vocab_start), + num_org=int(lm_head_num_org), + tp_group=tp_group, + ) + base_logits = None + else: + local_logits_t = local_logits.view( + num_proposals * batch_size, local_vocab_size + ).T.contiguous() + gathered_logits = torch.empty( + (tp_size * local_vocab_size, num_proposals * batch_size), + dtype=local_logits.dtype, + device=local_logits.device, + ) + tp_group.all_gather_into_tensor(gathered_logits, local_logits_t) + base_logits = ( + gathered_logits.T[:, : int(vocab_size)] + .contiguous() + .view(num_proposals, batch_size, int(vocab_size)) + ) - first_ids = torch.argmax(base_logits[0], dim=-1).to(torch.long) + if first_ids is None: + first_ids = torch.argmax(base_logits[0], dim=-1).to(torch.long) proposals = [first_ids] if num_proposals == 1: return first_ids[:, None] - candidate_ids = None - candidate_base = None candidate_weight = None - if 0 < candidate_pool_size < int(vocab_size): + if candidate_ids is not None: + candidate_weight = F.embedding(candidate_ids, embed_proj[2].weight) + elif 0 < candidate_pool_size < int(vocab_size): feedback_logits = base_logits[1:] candidate_ids = torch.topk( feedback_logits.amax(dim=0), diff --git a/test/registered/unit/spec/test_dflash_domino.py b/test/registered/unit/spec/test_dflash_domino.py index 9a484e4d7987..e1d0341ff8ff 100644 --- a/test/registered/unit/spec/test_dflash_domino.py +++ b/test/registered/unit/spec/test_dflash_domino.py @@ -1,5 +1,6 @@ import unittest from types import SimpleNamespace +from unittest import mock import torch import torch.nn.functional as F @@ -295,6 +296,146 @@ def test_full_chain_matches_native_oracle(self): actual[:, 0], first_expected, rtol=0, atol=0 ) + def test_tp2_gathered_base_matches_full_vocab_rollout(self): + class FakeTpGroup: + world_size = 2 + + def __init__(self, remote_logits): + self.remote_logits = remote_logits + self.candidate_ids = None + self.local_top_scores = None + self.remote_top_scores = None + + def all_gather_into_tensor(self, output, local): + output[: local.shape[0]].copy_(local) + if local.ndim == 1: + remote_max, remote_pos = self.remote_logits[0, :, :15].max(-1) + remote = ( + remote_pos + 16 if local.dtype == torch.long else remote_max + ) + elif local.shape[0] == local_vocab_size: + remote = self.remote_logits.reshape(-1, local_vocab_size).T + else: + k = local.shape[1] + remote_scores, remote_pos = torch.topk( + self.remote_logits[1:, :, :15].amax(0), + k=k, + dim=-1, + sorted=False, + ) + if local.dtype == torch.long: + remote = remote_pos + 16 + all_scores = ( + torch.stack( + (self.local_top_scores, self.remote_top_scores), dim=0 + ) + .permute(1, 0, 2) + .reshape(local.shape[0], -1) + ) + all_ids = ( + torch.stack((local, remote), dim=0) + .permute(1, 0, 2) + .reshape(local.shape[0], -1) + ) + global_pos = torch.topk( + all_scores, k=k, dim=-1, sorted=False + ).indices + self.candidate_ids = torch.gather(all_ids, 1, global_pos) + else: + remote = remote_scores + self.local_top_scores = local.clone() + self.remote_top_scores = remote_scores + output[local.shape[0] :].copy_(remote) + + def all_reduce(self, local): + remote_owned = (self.candidate_ids >= 16) & (self.candidate_ids < 31) + remote_pos = (self.candidate_ids - 16).clamp(0, 14) + remote = torch.gather( + self.remote_logits[1:].transpose(0, 1), + 2, + remote_pos[:, None, :].expand( + -1, self.remote_logits.shape[0] - 1, -1 + ), + ) + remote.masked_fill_(~remote_owned[:, None, :], 0) + local.add_(remote) + return local + + block_size = 7 + batch_size = 3 + local_vocab_size = 16 + padded_weight = torch.cat( + ( + self.lm_head_weight, + torch.zeros( + 1, + self.hidden_size, + device=self.device, + dtype=self.dtype, + ).fill_(1000), + ) + ) + local_weight = padded_weight[:local_vocab_size] + remote_weight = padded_weight[local_vocab_size:] + verified_ids = torch.tensor([1, 4, 9], device=self.device) + + for shift_label in (True, False): + draft_hidden = torch.randn( + batch_size, + block_size, + self.hidden_size, + device=self.device, + dtype=self.dtype, + ) + start = 0 if shift_label else 1 + z = draft_hidden[:, start : start + block_size - 1] + logits_input = ( + z.transpose(0, 1) + .contiguous() + .view((block_size - 1) * batch_size, self.hidden_size) + ) + remote_logits = F.linear(logits_input, remote_weight).view( + block_size - 1, batch_size, local_vocab_size + ) + + cases = ((0, False), (5, False), (5, True)) + for candidate_pool_size, force_compact in cases: + with self.subTest( + shift_label=shift_label, + candidate_pool_size=candidate_pool_size, + force_compact=force_compact, + ): + rollout_kwargs = dict( + draft_hidden=draft_hidden, + verified_ids=verified_ids, + target_embedding=self.embedding, + lm_head_weight=local_weight, + prefix_gru=self.prefix_gru, + embed_proj=self.embed_proj, + vocab_size=self.vocab_size, + shift_label=shift_label, + candidate_pool_size=candidate_pool_size, + tp_group=FakeTpGroup(remote_logits), + lm_head_num_org=local_vocab_size, + lm_head_num_org_padded=local_vocab_size, + ) + if force_compact: + with mock.patch( + "sglang.srt.speculative.domino_utils." + "_DOMINO_TP_FULL_BASE_LOGITS_MAX_BYTES", + 0, + ): + actual = domino_greedy_rollout(**rollout_kwargs) + else: + actual = domino_greedy_rollout(**rollout_kwargs) + expected = self._oracle( + draft_hidden, + verified_ids, + shift_label=shift_label, + candidate_pool_size=candidate_pool_size, + ) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + def test_block_candidate_pool_matches_oracle(self): block_size = 7 draft_hidden = torch.randn( @@ -574,6 +715,24 @@ def _modules(self, dtype=torch.bfloat16): ) return embedding, lm_head, prefix_gru, embed_proj + def _tp2_modules(self): + embedding, lm_head, prefix_gru, embed_proj = self._modules() + embedding = nn.Embedding(16, 8, dtype=torch.bfloat16) + lm_head = nn.Linear(8, 16, bias=False, dtype=torch.bfloat16) + shard = SimpleNamespace( + num_added_elements=0, + org_vocab_start_index=0, + org_vocab_end_index=16, + num_org_elements=16, + num_org_elements_padded=16, + ) + for module in (embedding, lm_head): + module.shard_indices = shard + module.org_vocab_size = 31 + module.tp_size = 2 + module.num_added_embeddings = 0 + return embedding, lm_head, prefix_gru, embed_proj + def _validate(self, **overrides): embedding, lm_head, prefix_gru, embed_proj = overrides.pop( "modules", self._modules() @@ -581,6 +740,7 @@ def _validate(self, **overrides): args = { "device": torch.device("cuda"), "tp_size": 1, + "tp_rank": 0, "target_vocab_size": 31, "draft_vocab_size": 31, "hidden_size": 8, @@ -595,9 +755,45 @@ def _validate(self, **overrides): def test_supported_runtime(self): self._validate() - def test_tp_and_vocab_mismatches_fail(self): - with self.assertRaisesRegex(ValueError, "TP=1"): + def test_tp_requires_vocab_shard_metadata(self): + with self.assertRaisesRegex(ValueError, "lm_head shard metadata"): self._validate(tp_size=2) + + def test_tp2_vocab_shards_supported(self): + self._validate(tp_size=2, modules=self._tp2_modules()) + + def test_tp2_incomplete_lm_head_shard_fails(self): + modules = self._tp2_modules() + modules[1].shard_indices = SimpleNamespace( + num_added_elements=0, + num_org_elements_padded=16, + ) + with self.assertRaisesRegex(ValueError, "shard metadata is missing"): + self._validate(tp_size=2, modules=modules) + + def test_tp_vocab_shard_must_match_rank(self): + modules = self._tp2_modules() + modules[1].shard_indices.org_vocab_start_index = 1 + modules[1].shard_indices.org_vocab_end_index = 17 + with self.assertRaisesRegex(ValueError, "does not match its TP rank"): + self._validate(tp_size=2, modules=modules) + + def test_tp1_requires_complete_vocab_shard(self): + modules = self._modules() + modules[1].shard_indices = SimpleNamespace( + num_added_elements=0, + org_vocab_start_index=0, + org_vocab_end_index=30, + num_org_elements=30, + num_org_elements_padded=31, + ) + modules[1].org_vocab_size = 31 + modules[1].tp_size = 1 + modules[1].num_added_embeddings = 0 + with self.assertRaisesRegex(ValueError, "does not match its TP rank"): + self._validate(modules=modules) + + def test_vocab_mismatch_fails(self): with self.assertRaisesRegex(ValueError, "identical target and draft"): self._validate(draft_vocab_size=30) From be8c97b28790a2555bcff3b30ceac9c743708964 Mon Sep 17 00:00:00 2001 From: jianuo-huang <1170180988@qq.com> Date: Wed, 22 Jul 2026 15:29:47 +0800 Subject: [PATCH 7/7] perf: select compact Domino TP path for CUDA Graph --- .../srt/speculative/dflash_worker_v2.py | 1 + python/sglang/srt/speculative/domino_utils.py | 13 +- .../unit/spec/test_dflash_domino.py | 133 +++++++++++++----- 3 files changed, 109 insertions(+), 38 deletions(-) diff --git a/python/sglang/srt/speculative/dflash_worker_v2.py b/python/sglang/srt/speculative/dflash_worker_v2.py index fa1cd1fab908..4d908fb6064b 100644 --- a/python/sglang/srt/speculative/dflash_worker_v2.py +++ b/python/sglang/srt/speculative/dflash_worker_v2.py @@ -208,6 +208,7 @@ def __call__(self, hidden_states, input_ids=None): lm_head_org_vocab_start=self.lm_head_org_vocab_start, lm_head_num_org=self.lm_head_num_org, lm_head_num_org_padded=self.lm_head_num_org_padded, + prefer_tp_candidate_pool=bs > 1, ) self.out[: bs * (self.block_size - 1)].copy_(proposals.reshape(-1)) diff --git a/python/sglang/srt/speculative/domino_utils.py b/python/sglang/srt/speculative/domino_utils.py index ced42f4df39a..de99f5077c5e 100644 --- a/python/sglang/srt/speculative/domino_utils.py +++ b/python/sglang/srt/speculative/domino_utils.py @@ -314,6 +314,7 @@ def domino_greedy_rollout( lm_head_org_vocab_start: int = 0, lm_head_num_org: int | None = None, lm_head_num_org_padded: int | None = None, + prefer_tp_candidate_pool: bool | None = None, ) -> torch.Tensor: """Generate a Domino chain using one block-shared base-logit candidate pool.""" if draft_hidden.ndim != 3: @@ -375,13 +376,17 @@ def domino_greedy_rollout( local_logits = F.linear(logits_input, weight).view( num_proposals, batch_size, local_vocab_size ) - full_base_logits_bytes = ( - num_proposals * batch_size * int(vocab_size) * local_logits.element_size() - ) + if prefer_tp_candidate_pool is None: + full_base_logits_bytes = ( + num_proposals * batch_size * int(vocab_size) * local_logits.element_size() + ) + prefer_tp_candidate_pool = ( + full_base_logits_bytes > _DOMINO_TP_FULL_BASE_LOGITS_MAX_BYTES + ) use_tp_candidate_pool = ( tp_size > 1 + and prefer_tp_candidate_pool and 0 < candidate_pool_size < int(vocab_size) - and full_base_logits_bytes > _DOMINO_TP_FULL_BASE_LOGITS_MAX_BYTES ) first_ids = None candidate_ids = None diff --git a/test/registered/unit/spec/test_dflash_domino.py b/test/registered/unit/spec/test_dflash_domino.py index e1d0341ff8ff..41384c02abe6 100644 --- a/test/registered/unit/spec/test_dflash_domino.py +++ b/test/registered/unit/spec/test_dflash_domino.py @@ -305,6 +305,8 @@ def __init__(self, remote_logits): self.candidate_ids = None self.local_top_scores = None self.remote_top_scores = None + self.full_logit_gathers = 0 + self.candidate_all_reduces = 0 def all_gather_into_tensor(self, output, local): output[: local.shape[0]].copy_(local) @@ -314,6 +316,7 @@ def all_gather_into_tensor(self, output, local): remote_pos + 16 if local.dtype == torch.long else remote_max ) elif local.shape[0] == local_vocab_size: + self.full_logit_gathers += 1 remote = self.remote_logits.reshape(-1, local_vocab_size).T else: k = local.shape[1] @@ -348,6 +351,7 @@ def all_gather_into_tensor(self, output, local): output[local.shape[0] :].copy_(remote) def all_reduce(self, local): + self.candidate_all_reduces += 1 remote_owned = (self.candidate_ids >= 16) & (self.candidate_ids < 31) remote_pos = (self.candidate_ids - 16).clamp(0, 14) remote = torch.gather( @@ -362,7 +366,6 @@ def all_reduce(self, local): return local block_size = 7 - batch_size = 3 local_vocab_size = 16 padded_weight = torch.cat( ( @@ -377,34 +380,48 @@ def all_reduce(self, local): ) local_weight = padded_weight[:local_vocab_size] remote_weight = padded_weight[local_vocab_size:] - verified_ids = torch.tensor([1, 4, 9], device=self.device) - - for shift_label in (True, False): - draft_hidden = torch.randn( - batch_size, - block_size, - self.hidden_size, - device=self.device, - dtype=self.dtype, - ) - start = 0 if shift_label else 1 - z = draft_hidden[:, start : start + block_size - 1] - logits_input = ( - z.transpose(0, 1) - .contiguous() - .view((block_size - 1) * batch_size, self.hidden_size) - ) - remote_logits = F.linear(logits_input, remote_weight).view( - block_size - 1, batch_size, local_vocab_size - ) - - cases = ((0, False), (5, False), (5, True)) - for candidate_pool_size, force_compact in cases: + cases = ( + (8, 5, None, 1 << 60, False), + (8, 5, None, 0, True), + (1, 5, False, 1 << 60, False), + (8, 5, True, 1 << 60, True), + (8, 0, True, 1 << 60, False), + ) + for ( + batch_size, + candidate_pool_size, + prefer_tp_candidate_pool, + full_base_logits_max_bytes, + expect_compact, + ) in cases: + verified_ids = torch.arange(batch_size, device=self.device) + for shift_label in (True, False): with self.subTest( + batch_size=batch_size, shift_label=shift_label, candidate_pool_size=candidate_pool_size, - force_compact=force_compact, + prefer_tp_candidate_pool=prefer_tp_candidate_pool, + full_base_logits_max_bytes=full_base_logits_max_bytes, + expect_compact=expect_compact, ): + draft_hidden = torch.randn( + batch_size, + block_size, + self.hidden_size, + device=self.device, + dtype=self.dtype, + ) + start = 0 if shift_label else 1 + z = draft_hidden[:, start : start + block_size - 1] + logits_input = ( + z.transpose(0, 1) + .contiguous() + .view((block_size - 1) * batch_size, self.hidden_size) + ) + remote_logits = F.linear(logits_input, remote_weight).view( + block_size - 1, batch_size, local_vocab_size + ) + tp_group = FakeTpGroup(remote_logits) rollout_kwargs = dict( draft_hidden=draft_hidden, verified_ids=verified_ids, @@ -415,18 +432,16 @@ def all_reduce(self, local): vocab_size=self.vocab_size, shift_label=shift_label, candidate_pool_size=candidate_pool_size, - tp_group=FakeTpGroup(remote_logits), + tp_group=tp_group, lm_head_num_org=local_vocab_size, lm_head_num_org_padded=local_vocab_size, + prefer_tp_candidate_pool=prefer_tp_candidate_pool, ) - if force_compact: - with mock.patch( - "sglang.srt.speculative.domino_utils." - "_DOMINO_TP_FULL_BASE_LOGITS_MAX_BYTES", - 0, - ): - actual = domino_greedy_rollout(**rollout_kwargs) - else: + with mock.patch( + "sglang.srt.speculative.domino_utils." + "_DOMINO_TP_FULL_BASE_LOGITS_MAX_BYTES", + full_base_logits_max_bytes, + ): actual = domino_greedy_rollout(**rollout_kwargs) expected = self._oracle( draft_hidden, @@ -435,6 +450,56 @@ def all_reduce(self, local): candidate_pool_size=candidate_pool_size, ) torch.testing.assert_close(actual, expected, rtol=0, atol=0) + self.assertEqual( + tp_group.candidate_all_reduces, int(expect_compact) + ) + self.assertEqual( + tp_group.full_logit_gathers, int(not expect_compact) + ) + + def test_capture_sampler_uses_capture_bucket_policy(self): + from sglang.srt.speculative.dflash_worker_v2 import _DominoDraftSampler + + block_size = 7 + 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=self.vocab_size, + block_size=block_size, + shift_label=True, + max_bs=8, + candidate_pool_size=5, + ) + for batch_size, expect_compact in ((1, False), (8, True)): + with self.subTest(capture_bucket=batch_size, expect_compact=expect_compact): + hidden_states = torch.randn( + batch_size * block_size, + self.hidden_size, + device=self.device, + dtype=self.dtype, + ) + input_ids = torch.zeros( + batch_size * block_size, + device=self.device, + dtype=torch.long, + ) + proposals = torch.zeros( + batch_size, + block_size - 1, + device=self.device, + dtype=torch.long, + ) + with mock.patch( + "sglang.srt.speculative.dflash_worker_v2.domino_greedy_rollout", + return_value=proposals, + ) as rollout: + sampler(hidden_states, input_ids) + self.assertEqual( + rollout.call_args.kwargs["prefer_tp_candidate_pool"], + expect_compact, + ) def test_block_candidate_pool_matches_oracle(self): block_size = 7