feat: add Domino drafter training backend

[khazic]

What

Adds a Domino drafter training backend to the SpeCo overlay (part of #16). Select it for drafter training with:

actor_rollout_ref.rollout.drafter.speculative_algorithm=DOMINO

Design

Domino is a DFlash variant, so DominoTrainerBackend subclasses DFlashTrainerBackend exactly like DSpark does; only build_model and the training forward differ. The DFlash parallel block backbone drafts a whole block in one non-causal forward, so each predicted position is blind to the block's earlier (drafted) tokens. Domino adds a causal correction head that fixes this:

• prefix_gru: a single-layer GRU over the block's previous-token embeddings, producing a causal state per block position.
• embed_proj: a low-rank MLP over [backbone hidden | GRU state] emitting a full-vocabulary logit delta added to the parallel base logits.

Training jointly supervises the refined (final) and backbone-only (base) logits with a base-anchor curriculum loss = (1 - lambda) * final + lambda * base, where lambda decays from its start value to 0 so early steps keep the parallel backbone strong and later steps let the correction head take over.

The shifted-label alignment (target x[a+1:a+1+block], prev [x[a], labels[:-1]], every position supervised) reuses the DSpark alignment, which equals the AutoModel reference's shift_label Domino path. Domino reuses the DFlash block-drafter plumbing (anchor sampling, noise block, block attention mask, dflash_aux hidden layout) and is wired through the worker dispatch, base_trainer block-drafter gates, auto config routing, old-logprob aux layers, and config keys (domino_*, falling back to dflash_*).

Serving

Domino is not an engine-level speculative algorithm. Engines expose it as a projector_type sub-mode of DFlash: the method stays dflash, and the causal correction head (prefix_gru + embed_proj) is enabled from the checkpoint's dflash_config.projector_type="domino". DOMINO is therefore never a valid engine algorithm string, so both the vLLM and the SGLang config builders raise with a clear message directing the rollout/serve path to DFLASH, which keeps the Domino correction head active on engines that carry the projector.

Engine-side support is in flight and not merged yet (vllm-project/vllm#48241 for vLLM, sgl-project/sglang#31328 for SGLang), so a Domino checkpoint cannot be served on a released engine today. This PR only adds the training backend plus the two guardrails; no serving claim is made here.

Not a duplicate

Checked the open issues/PRs (gh pr list): no open PR touches Domino. #13 (Separate Draft Model Training) and #10 (draft weight loading) are unrelated.

Testing

A human reviewed every changed line and ran the tests below.

CPU contract tests (tests/integration/test_domino_backend_contract.py), 7 passed: projector-head shapes, the lambda-base curriculum schedule, block-drafter classification, auto-config routing for DominoDraftModel, old-logprob DFlash-aux routing, and the DOMINO serve guardrail on both the vLLM and the SGLang config builders. The full suite is 149 passed, 2 skipped, 2 failed; both failures pre-exist on main (a tests/heavy runner-env check and a stale metrics assertion) and are unrelated to this change.

pytest tests/integration/test_domino_backend_contract.py -q
# 7 passed

GPU hardware smoke (ci/domino_gpu_smoke.py) on a real target (Qwen3-4B, 1x A100): it collects the real DFlash-style multi-layer context hidden states from the frozen target, builds the Domino draft via DominoTrainerBackend.build_model, and runs 120 optimizer steps through compute_loss (the block-drafter forward with the GRU correction head and dual-logit curriculum). The draft is cold-started, so the signal is convergence: both the final (Domino-refined) and base (backbone-only) CE fall and accuracy rises. Crucially, the final loss stays below the base loss throughout, i.e. the causal correction head genuinely improves on the parallel backbone, and lambda_base decays from ~1 to 0 as designed.

python ci/domino_gpu_smoke.py --target /path/to/Qwen3-4B --steps 120 --lambda-decay-steps 60

Domino GPU training log (Qwen3-4B target, cold start, 120 steps)

[smoke] loading target /path/to/Qwen3-4B
[smoke] context layers=[1, 9, 17, 25, 33] (of 36)
[smoke] batch seq_len=91 hidden=12800
[smoke] block_size=8 gru=1024 emb_dim=256 trainable_params=174,696,608
[smoke] step   0  final_loss=11.9464  base_loss=11.9907  final_acc=0.0000  base_acc=0.0000  lambda_base=0.983
[smoke] step  10  final_loss=4.9976  base_loss=5.1292  final_acc=0.0679  base_acc=0.0679  lambda_base=0.817
[smoke] step  20  final_loss=2.2520  base_loss=2.6272  final_acc=0.4090  base_acc=0.4220  lambda_base=0.650
[smoke] step  30  final_loss=0.5375  base_loss=0.9535  final_acc=0.7991  base_acc=0.7991  lambda_base=0.483
[smoke] step  40  final_loss=0.0794  base_loss=0.3261  final_acc=0.9899  base_acc=0.9393  lambda_base=0.317
[smoke] step  50  final_loss=0.0340  base_loss=0.1515  final_acc=0.9971  base_acc=0.9928  lambda_base=0.150
[smoke] step  60  final_loss=0.0240  base_loss=0.1278  final_acc=1.0000  base_acc=1.0000  lambda_base=0.000
[smoke] step  70  final_loss=0.0232  base_loss=0.1251  final_acc=1.0000  base_acc=1.0000  lambda_base=0.000
[smoke] step  80  final_loss=0.0232  base_loss=0.1249  final_acc=1.0000  base_acc=1.0000  lambda_base=0.000
[smoke] step  90  final_loss=0.0227  base_loss=0.1243  final_acc=1.0000  base_acc=1.0000  lambda_base=0.000
[smoke] step 100  final_loss=0.0227  base_loss=0.1242  final_acc=1.0000  base_acc=1.0000  lambda_base=0.000
[smoke] step 110  final_loss=0.0227  base_loss=0.1239  final_acc=1.0000  base_acc=1.0000  lambda_base=0.000
[smoke] step 119  final_loss=0.0226  base_loss=0.1239  final_acc=1.0000  base_acc=1.0000  lambda_base=0.000
[smoke] DONE  final_loss 11.9464->0.0226  base_loss 11.9907->0.1239  final_acc 0.0000->1.0000

Known follow-ups

• Domino trains a full_vocab CE head; restricted/sampled-CE (as DFlash/DSpark support) would need to index the embed_proj output rows and is left as a follow-up.
• Serving is gated on the engine PRs above landing. The checkpoint this backend writes matches their config contract on paper (projector_type, gru_hidden_dim, pure_draft_prefix_len, shift_label, emb_dim, and the prefix_gru.* / embed_proj.* weight names), but an end-to-end load of a SpeCo-trained Domino checkpoint on those branches has not been validated yet.

[gemini-code-assist]

Code Review

This pull request introduces the Domino drafter training backend, a training-only speculative decoding algorithm that extends the DFlash backbone with a causal GRU correction head and a dual-logit base-anchor curriculum. The changes include the core backend implementation, configuration files, integration into the training pipeline, and comprehensive tests. The review feedback highlights three key improvement opportunities in domino_trainer_backend.py: adding a guard to prevent a TypeError crash when spec_model_path is None, registering the _forward_count step counter as a PyTorch buffer to ensure it is correctly saved and restored during checkpointing, and retrieving its scalar value safely using .item().

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

If spec_model_path is None (which is common when starting training from scratch without pre-existing drafter weights), os.path.join will raise a TypeError and crash the initialization. Adding a guard ensures robust fallback behavior.

Suggested change

	config_path = os.path.join(spec_model_path, "config.json")
	config_path = os.path.join(spec_model_path, "config.json") if spec_model_path else None

[gemini-code-assist]

The step counter _forward_count is currently defined as a plain Python integer. This means it will not be saved in the model's state dict and will be reset to 0 when resuming training from a checkpoint, which resets the curriculum schedule (lambda_base) and can cause training instability. Registering it as a buffer ensures it is correctly saved and restored.

Suggested change

	self._forward_count = 0
	self.register_buffer("_forward_count", torch.tensor(0, dtype=torch.long))

[gemini-code-assist]

Since _forward_count is registered as a buffer tensor, we should retrieve its scalar value using .item() and cast it to an integer before passing it to get_lambda_base to avoid device or type conversion issues.

Suggested change

	return get_lambda_base(self._forward_count, self.lambda_base_decay_steps, self.lambda_base_start)
	return get_lambda_base(int(self._forward_count.item()), self.lambda_base_decay_steps, self.lambda_base_start)

[khazic]

Thanks for the review. Addressed the findings:

spec_model_path=None guard (high): Fixed in b23c03c. config_path now falls back to None, and the existing if config_path and os.path.exists(...) / short-circuited checks route a from-scratch run to the fallback config instead of crashing.

_forward_count as a registered buffer (medium x2): Not applied, because in this repo the wrapper's counter is never checkpointed or restored, so a buffer would not survive resume:

• Hot publish (_get_trainable_state_dict) skips non-floating tensors, so a long buffer is excluded anyway.
• Offline export (_get_pretrained_export_state_dict) strips to draft_model.-prefixed names only; _forward_count lives on the DominoTrainingModel wrapper, not under draft_model., so it is dropped.
• Resume goes through build_model -> _load_draft_checkpoint, which loads only into draft_model; the wrapper is reconstructed fresh either way.

Drafter resume reloads the HF drafter weights with a fresh optimizer and schedule, so the per-run counter restarting at 0 is consistent with the rest of the training state, not a desync. Keeping it a plain int avoids a per-step device sync with no behavioral gain.

[tpx818]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces the Domino drafter training backend, a DFlash variant that incorporates a causal GRU correction head and a dual-logit base-anchor curriculum loss. The implementation includes new model architectures, configuration schemas, integration with runtime engines, and smoke/contract tests. The review feedback highlights three key issues: a potential shape mismatch crash when loading pretrained models with conflicting context layer configurations, an uncalculated top-5 accuracy metric that remains at zero, and potential TypeErrors when optional configuration parameters are explicitly set to null in the YAML config.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

If a pretrained Domino model is loaded from spec_model_path, its num_context_layers is already defined in the checkpoint's config.json. If the user configures a different domino_num_target_layers in rollout.drafter.training, _normalize_dflash_config will keep the pretrained model's num_context_layers (due to the is None guard), but resolve_oldlogprob_aux_layer_ids will prioritize domino_num_target_layers from the training config. This leads to a mismatch between the number of collected context layers and the number of layers the model's projection weights expect, causing a shape mismatch crash during the forward pass. We should raise a ValueError if they mismatch.

Suggested change

	def _normalize_dflash_config(self, drafter_config, target_hf_config, normalized_state, spec_model_path):
	training_cfg = self.config.rollout.drafter.training
	if training_cfg.get("domino_num_target_layers", None) is not None:
	if getattr(drafter_config, "num_context_layers", None) is None:
	drafter_config.num_context_layers = int(training_cfg["domino_num_target_layers"])
	return super()._normalize_dflash_config(drafter_config, target_hf_config, normalized_state, spec_model_path)
	def _normalize_dflash_config(self, drafter_config, target_hf_config, normalized_state, spec_model_path):
	training_cfg = self.config.rollout.drafter.training
	if training_cfg.get("domino_num_target_layers", None) is not None:
	num_target_layers = int(training_cfg["domino_num_target_layers"])
	if getattr(drafter_config, "num_context_layers", None) is None:
	drafter_config.num_context_layers = num_target_layers
	elif drafter_config.num_context_layers != num_target_layers:
	raise ValueError(
	f"Configured domino_num_target_layers ({num_target_layers}) does not match "
	f"the loaded model's num_context_layers ({drafter_config.num_context_layers})."
	)
	return super()._normalize_dflash_config(drafter_config, target_hf_config, normalized_state, spec_model_path)

[khazic]

The divergence is real, but it is not introduced here and not Domino-specific: dspark_trainer_backend.py:393 is the same code with the same is None guard, and _dflash_num_context_layers gives dspark_num_target_layers the same priority over the checkpoint value. Domino mirrors DSpark verbatim by design.

Adding the raise only to Domino would make the two backends behave differently on the same misconfiguration. Worth fixing, but as one change covering DSpark and Domino together rather than inside this PR.

[gemini-code-assist]

The top5_correct tensor is initialized to 0.0 but is never updated with the actual top-5 predictions. This causes the top5_acc metric reported in get_training_metrics to always be 0.0. We should compute the top-5 accuracy by checking if the target token is within the top-5 predicted logits for both suffix and non-suffix positions.

Suggested change

	top1_correct = torch.zeros((), dtype=torch.float32, device=device)
	top5_correct = torch.zeros((), dtype=torch.float32, device=device)
	quality_token_count = torch.zeros((), dtype=torch.float32, device=device)
	if active_final_pred is not None and active_targets.numel() > 0:
	active_correct = active_final_pred.eq(active_targets)
	correct[active_mask] = active_correct
	base_correct[active_mask] = active_base_pred.eq(active_targets)
	top1_correct = active_correct.float().sum()
	quality_token_count = active_targets.new_tensor(float(active_targets.numel()), dtype=torch.float32)
	top1_correct = torch.zeros((), dtype=torch.float32, device=device)
	top5_correct = torch.zeros((), dtype=torch.float32, device=device)
	quality_token_count = torch.zeros((), dtype=torch.float32, device=device)
	if active_final_pred is not None and active_targets.numel() > 0:
	active_correct = active_final_pred.eq(active_targets)
	correct[active_mask] = active_correct
	base_correct[active_mask] = active_base_pred.eq(active_targets)
	top1_correct = active_correct.float().sum()
	quality_token_count = active_targets.new_tensor(float(active_targets.numel()), dtype=torch.float32)
	top5_correct_mask = torch.zeros_like(active_correct, dtype=torch.bool)
	if active_final_logits is not None and suffix_mask.any():
	top5_final = active_final_logits.topk(5, dim=-1).indices
	top5_correct_mask[suffix_mask] = top5_final.eq(active_targets[suffix_mask].unsqueeze(-1)).any(dim=-1)
	non_suffix_mask = ~suffix_mask
	if non_suffix_mask.any():
	top5_base = base_logits[non_suffix_mask].topk(5, dim=-1).indices
	top5_correct_mask[non_suffix_mask] = top5_base.eq(active_targets[non_suffix_mask].unsqueeze(-1)).any(dim=-1)
	top5_correct = top5_correct_mask.float().sum()

[khazic]

Good catch, fixed in e2ea1a4. DFlash and DSpark both reduce this counter and Domino did not, so top5_acc (derived in base_trainer) was pinned at 0.

I took a different shape than the suggestion: mirrored the argmax idiom already in this block (topk over base_logits, suffix rows overwritten from active_final_logits) so no second [num_active, vocab] tensor is materialized, which this forward deliberately avoids. Also guarded with topk = min(5, vocab) like DSpark, since the suggestion hardcodes topk(5) and would crash on the small-vocab configs the contract tests use. Added a CPU regression test that fails without the fix (top5=0 vs top1=2).

[gemini-code-assist]

If any of the domino_* parameters are explicitly set to null (None) in the YAML config, training_cfg.get(...) will return None instead of the default value. Passing None to int() or float() will raise a TypeError. We should use a helper function to safely fall back to the default value when the configured value is None.

Suggested change

	return DominoTrainingModel(
	draft_model=draft_model,
	block_size=int(training_cfg.get("domino_block_size", getattr(drafter_config, "block_size", 16))),
	num_anchors=int(training_cfg.get("domino_num_anchors", getattr(drafter_config, "num_anchors", 512))),
	loss_decay_gamma=float(training_cfg.get("domino_loss_decay_gamma", getattr(drafter_config, "loss_decay_gamma", 7.0))),
	pure_draft_prefix_len=int(training_cfg.get("domino_pure_draft_prefix_len", getattr(drafter_config, "pure_draft_prefix_len", 1))),
	lambda_base_start=float(training_cfg.get("domino_lambda_base_start", getattr(drafter_config, "lambda_base_start", 1.0))),
	lambda_base_decay_steps=int(training_cfg.get("domino_lambda_base_decay_steps", getattr(drafter_config, "lambda_base_decay_steps", 2000))),
	def get_val(key, default):
	val = training_cfg.get(key)
	return val if val is not None else default
	return DominoTrainingModel(
	draft_model=draft_model,
	block_size=int(get_val("domino_block_size", getattr(drafter_config, "block_size", 16))),
	num_anchors=int(get_val("domino_num_anchors", getattr(drafter_config, "num_anchors", 512))),
	loss_decay_gamma=float(get_val("domino_loss_decay_gamma", getattr(drafter_config, "loss_decay_gamma", 7.0))),
	pure_draft_prefix_len=int(get_val("domino_pure_draft_prefix_len", getattr(drafter_config, "pure_draft_prefix_len", 1))),
	lambda_base_start=float(get_val("domino_lambda_base_start", getattr(drafter_config, "lambda_base_start", 1.0))),
	lambda_base_decay_steps=int(get_val("domino_lambda_base_decay_steps", getattr(drafter_config, "lambda_base_decay_steps", 2000))),
	), drafter_config

[khazic]

Same reasoning: int(training_cfg.get(key, default)) is the pattern DSpark (435-442, 477-478) and DFlash (692-698) already use, and the shipped speco_base.yaml sets all six of these keys to non-null values, so this only fires if someone explicitly writes domino_block_size: null. dspark_block_size: null has the identical hazard today.

Hardening only Domino would make it inconsistent with its siblings, so this belongs in a repo-wide change if we want it.