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

Co-authored-by: Qiaolin-Yu <liin1211@outlook.com>

python/sglang/srt/arg_groups/fields/spec.py

@@ -73,6 +73,10 @@ class Spec:
         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.",

python/sglang/srt/models/dflash.py

@@ -635,6 +635,35 @@ class DFlashDraftModel(nn.Module):
             )
         self.hidden_norm = RMSNorm(hidden_size, eps=rms_norm_eps)
 
+        # The model loader calls load_weights() before set_block_size(). Build
+        # Domino projector modules here so their parameters are present while
+        # checkpoint weights are loaded.
+        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 set_block_size(self, block_size: int) -> None:
         """Adopt the block size the worker resolved.
 
@@ -728,6 +757,7 @@ class DFlashDraftModel(nn.Module):
         ]
 
         params_dict = dict(self.named_parameters())
+        loaded_params = set()
 
         # Alias the native export's "encoder." names.
         _VENDOR_ENCODER_ALIASES = {
@@ -752,6 +782,14 @@ class DFlashDraftModel(nn.Module):
             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
@@ -762,6 +800,7 @@ class DFlashDraftModel(nn.Module):
                 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)
@@ -796,8 +835,31 @@ class DFlashDraftModel(nn.Module):
                             f"(num_context_features={self.num_context_features}, hidden_size={int(self.config.hidden_size)}), "
                             f"but got {loaded_shape} for weight '{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}."
+                    )
                 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):

python/sglang/srt/speculative/dflash_utils.py

@@ -538,6 +538,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:
@@ -698,6 +707,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,
@@ -711,6 +786,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,
     )
 
 

python/sglang/srt/speculative/dflash_worker_v2.py

@@ -60,6 +60,10 @@ from sglang.srt.speculative.dflash_utils import (
     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,
@@ -280,6 +284,55 @@ class _SelectorDraftSampler:
         self.q_out[:bs].copy_(q_rows)
 
 
+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,
+        candidate_pool_size,
+    ):
+        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)
+        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)
+        bonus_tokens = input_ids.view(bs, self.block_size)[:, 0]
+        proposals = domino_greedy_rollout(
+            draft_hidden=draft_hidden,
+            bonus_tokens=bonus_tokens,
+            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,
+        )
+        self.out[: bs * (self.block_size - 1)].copy_(proposals.reshape(-1))
+
+
 class DFlashWorkerV2(BaseSpecWorker):
     """DFLASH speculative decoding worker (spec-v2).
 
@@ -333,6 +386,36 @@ class DFlashWorkerV2(BaseSpecWorker):
         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(
+            get_spec().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),
+                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 get_spec().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))
@@ -351,6 +434,11 @@ class DFlashWorkerV2(BaseSpecWorker):
                 )
         self.draft_model.set_block_size(self.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
@@ -373,6 +461,11 @@ class DFlashWorkerV2(BaseSpecWorker):
                 self.draft_window_size,
                 self.use_compact_draft_cache,
             )
+            if self._is_domino:
+                logger.info(
+                    "DFLASH Domino rollout enabled (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, noise_embed_scale=%s",
                 self._mask_token,
@@ -661,6 +754,28 @@ class DFlashWorkerV2(BaseSpecWorker):
             # 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(get_exec().graph.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:
                 # No shard metadata to recover per-rank vocab offsets from.
@@ -2154,8 +2269,35 @@ class DFlashWorkerV2(BaseSpecWorker):
             draft_out = self.draft_model_runner.forward(forward_batch)
         draft_logits_output = draft_out.logits_output
 
-        folded = self._draft_sampler is not None and draft_out.can_run_graph
-        if folded:
+        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.")
+            draft_next = domino_greedy_rollout(
+                draft_hidden=draft_hidden,
+                bonus_tokens=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,
+            )
+        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)

python/sglang/srt/speculative/domino_utils.py

@@ -0,0 +1,212 @@
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+
+def _domino_gru_cell(
+    prefix_gru: nn.GRU, input: torch.Tensor, hidden: torch.Tensor
+) -> torch.Tensor:
+    """Run one feedback token without cuDNN's per-call 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,
+    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,
+    bonus_tokens: 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,
+) -> 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 bonus_tokens.shape != (batch_size,):
+        raise ValueError(
+            f"bonus_tokens must have shape ({batch_size},), got {tuple(bonus_tokens.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."
+        )
+
+    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)
+
+    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):
+        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((bonus_tokens, 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)