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/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_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..4d908fb6064b 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, @@ -147,6 +151,68 @@ def __call__(self, hidden_states, input_ids=None): self.out[:n].copy_(selected.view(-1)) +class _DominoDraftSampler: + """Capture-safe 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, + 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 + 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) + 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 + ) + + 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, + 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, + prefer_tp_candidate_pool=bs > 1, + ) + self.out[: bs * (self.block_size - 1)].copy_(proposals.reshape(-1)) + + class DFlashWorkerV2(BaseSpecWorker): """DFLASH speculative decoding worker (spec-v2). @@ -195,6 +261,37 @@ 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 + 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) + 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), + 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), + 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 +309,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 +330,12 @@ def __init__( self.draft_window_size, self.use_compact_draft_cache, ) + if self._is_domino: + logger.info( + "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( "DFLASH draft runner ready. mask_token=%s, mask_token_id=%s, mask_token_id_override=%s", self._mask_token, @@ -369,6 +477,38 @@ 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: + 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=%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, + 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), + 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: # No shard metadata to recover per-rank vocab offsets from. @@ -1613,7 +1753,47 @@ 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 + 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.") + 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.") + 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], + 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), + 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[ : 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..de99f5077c5e --- /dev/null +++ b/python/sglang/srt/speculative/domino_utils.py @@ -0,0 +1,481 @@ +from __future__ import annotations + +import torch +import torch.nn.functional as F +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( + 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 _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, + 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}.") + 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}." + ) + + 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}." + ) + + 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, + 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, + 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: + raise ValueError( + f"draft_hidden must have shape [batch, block, hidden], got {tuple(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)}." + ) + + 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: + raise ValueError( + "Domino draft hidden states do not contain enough proposal positions." + ) + + 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) + ) + local_logits = F.linear(logits_input, weight).view( + num_proposals, batch_size, local_vocab_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) + ) + 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)) + ) + + 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_weight = None + 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), + 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_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( + prefix_gru, target_embedding(next_ids), gru_hidden[0] + )[None] + + 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..41384c02abe6 --- /dev/null +++ b/test/registered/unit/spec/test_dflash_domino.py @@ -0,0 +1,887 @@ +import unittest +from types import SimpleNamespace +from unittest import mock + +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, 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] + 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] + + 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): + 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) + 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_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 + self.full_logit_gathers = 0 + self.candidate_all_reduces = 0 + + 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: + self.full_logit_gathers += 1 + 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): + 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( + 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 + 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:] + 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, + 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, + 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=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, + ) + 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, + verified_ids, + shift_label=shift_label, + 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 + draft_hidden = torch.randn( + 3, + block_size, + self.hidden_size, + device=self.device, + dtype=self.dtype, + ) + verified_ids = torch.tensor([1, 4, 9], device=self.device) + 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 + ) + 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, + ) + 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 + ) + 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 + + 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, + candidate_pool_size=5, + ) + 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, + 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 + + 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, + candidate_pool_size=5, + ) + 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, + candidate_pool_size=5, + ) + 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): + 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 _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() + ) + args = { + "device": torch.device("cuda"), + "tp_size": 1, + "tp_rank": 0, + "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_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) + + 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()