fix(domino): address PR #773 review — fix stale test params and lint

- tests/e2e: --domino-lambda-decay-steps → --domino-lambda-decay-ratio
- tests/integration: lambda_base_decay_steps → lambda_base_decay_ratio
- tests/unit: remove phantom shift_targets/shift_label/max_anchors params,
  restructure tests to verify domino vs dflash behavior directly,
  add loss_config to lambda decay test forward calls
- core.py: accept LossConfig | None in _compute_domino_metrics,
  add j_star/lambda_base observability metrics, ruff format fixes
- domino.py: fix E501 docstring line lengths
- dspark/metrics.py: ruff format fix

Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
Signed-off-by: Orestis Zambounis <orestis.zambounis@gmail.com>

Author: orestis-z

src/speculators/models/dflash/core.py

@@ -251,7 +251,8 @@ def _auf_mask(
         base_mask: torch.Tensor,
         num_anchors: int,
         block_size: int,
-    ) -> torch.Tensor:
+        return_j_star: bool = False,
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
         """Compute the Accept-Until-Fail (AUF) loss mask for the base branch.
 
         Truncates the cross-entropy support at the first greedy prediction error
@@ -272,21 +273,43 @@ def _auf_mask(
             base_mask: Boolean validity mask [1, T] (e.g. domino_loss_mask).
             num_anchors: Number of anchor blocks.
             block_size: Tokens per block.
+            return_j_star: If True, also return j* (first error position) per block.
 
         Returns:
-            Boolean mask [1, T] with the same shape as base_mask, zeroed out
-            for all positions strictly after the first error in each block.
+            If return_j_star is False:
+                Boolean mask [1, T] with the same shape as base_mask, zeroed out
+                for all positions strictly after the first error in each block.
+            If return_j_star is True:
+                Tuple of (mask, j_star) where j_star is [num_anchors] tensor
+                containing the first error position (0-indexed) per block.
+                j* = block_size means no error (all accepted).
         """
-        base_preds_4d = logits.detach().argmax(dim=-1).reshape(1, num_anchors, block_size)
-        target_ids_4d = targets.detach().argmax(dim=-1).reshape(1, num_anchors, block_size)
+        base_preds_4d = (
+            logits.detach().argmax(dim=-1).reshape(1, num_anchors, block_size)
+        )
+        target_ids_4d = (
+            targets.detach().argmax(dim=-1).reshape(1, num_anchors, block_size)
+        )
         mask_4d = base_mask.reshape(1, num_anchors, block_size)
         errors = (base_preds_4d != target_ids_4d) & mask_4d
         error_floats = errors.float()
         running_errors = error_floats.cumsum(dim=-1)
         # errors_before == 0: keep accepted prefix + breaker token j*
         # errors_before >= 1: zero the unreachable suffix
         errors_before = running_errors - error_floats
-        return (mask_4d & (errors_before == 0)).reshape_as(base_mask)
+        auf_mask = (mask_4d & (errors_before == 0)).reshape_as(base_mask)
+
+        if not return_j_star:
+            return auf_mask
+
+        # j* = first error position per block (0-indexed within block)
+        # If no error, j* = block_size (all positions accepted)
+        first_error = errors.float().argmax(dim=-1)  # [1, num_anchors]
+        has_error = errors.any(dim=-1)  # [1, num_anchors]
+        j_star = torch.where(
+            has_error, first_error, torch.full_like(first_error, block_size)
+        )
+        return auf_mask, j_star.squeeze(0)  # [num_anchors]
 
     @property
     def mask_token_id(self) -> int:
@@ -421,7 +444,9 @@ def _backbone_forward(
             )
             # Shift right by 1 so verifier_logits[i] predicts token at position i
             verifier_logits = torch.roll(verifier_logits, 1, dims=1)
-            target_indices = anchored_block_indices + (1 if self.config.projector_type == "domino" else 0)
+            target_indices = anchored_block_indices + (
+                1 if self.config.projector_type == "domino" else 0
+            )
             target_indices = target_indices.clamp(max=verifier_logits.shape[1] - 1)
             targets = verifier_logits[:, target_indices]
             # shape: [1, num_anchors*block_size, draft_vocab_size]
@@ -471,7 +496,7 @@ def _compute_domino_metrics(
         loss_mask: torch.Tensor,
         aligned_loss_mask: torch.Tensor,
         anchored_block_indices: torch.Tensor,
-        loss_config: "LossConfig",
+        loss_config: "LossConfig | None",
         gamma: float,
         normalize_by_decay: bool,
         global_step: int,
@@ -506,38 +531,67 @@ def _compute_domino_metrics(
         anchor_pos_in_mask = anchored_block_indices[:: self.block_size]
         domino_loss_mask[:, :: self.block_size] = loss_mask[:, anchor_pos_in_mask]
 
-        # B-AUF+D (arXiv 2607.01893): L_final keeps full mask; L_base is optionally truncated.
+        # B-AUF+D (arXiv 2607.01893): L_final keeps full mask;
+        # L_base is optionally truncated.
+        j_star = None
         if self.config.domino_auf:
-            base_mask = self._auf_mask(
-                logits, targets, domino_loss_mask, num_anchors, self.block_size
+            base_mask, j_star = self._auf_mask(
+                logits,
+                targets,
+                domino_loss_mask,
+                num_anchors,
+                self.block_size,
+                return_j_star=True,
             )
         else:
             base_mask = domino_loss_mask
 
         base_loss, base_metrics = compute_metrics(
-            logits, targets, base_mask, self.block_size,
-            gamma=gamma, loss_config=loss_config,
-            normalize_by_decay=normalize_by_decay, decay_mode="domino",
+            logits,
+            targets,
+            base_mask,
+            self.block_size,
+            gamma=gamma,
+            loss_config=loss_config,
+            normalize_by_decay=normalize_by_decay,
+            decay_mode="domino",
         )
         final_loss, final_metrics = compute_metrics(
-            refined_logits, targets, domino_loss_mask, self.block_size,
-            gamma=gamma, loss_config=loss_config,
-            normalize_by_decay=normalize_by_decay, decay_mode="domino",
+            refined_logits,
+            targets,
+            domino_loss_mask,
+            self.block_size,
+            gamma=gamma,
+            loss_config=loss_config,
+            normalize_by_decay=normalize_by_decay,
+            decay_mode="domino",
         )
         loss = (1.0 - lambda_base) * final_loss + lambda_base * base_loss
+
+        ones = torch.tensor(1.0, device=logits.device)
         metrics = {
             "loss_sum": loss.detach().clone(),
-            "loss_total": torch.tensor(1.0, device=logits.device),
+            "loss_total": ones,
             "base_loss_sum": base_loss.detach().clone(),
-            "base_loss_total": torch.tensor(1.0, device=logits.device),
+            "base_loss_total": ones,
             "full_acc_sum": base_metrics["full_acc_sum"],
             "full_acc_total": base_metrics["full_acc_total"],
             "final_loss_sum": final_loss.detach().clone(),
-            "final_loss_total": torch.tensor(1.0, device=logits.device),
+            "final_loss_total": ones,
             "final_full_acc_sum": final_metrics["full_acc_sum"],
             "final_full_acc_total": final_metrics["full_acc_total"],
+            "lambda_base_sum": torch.tensor(lambda_base, device=logits.device),
+            "lambda_base_total": ones,
             **{k: v for k, v in final_metrics.items() if k.startswith("position_")},
         }
+
+        # AUF observability: mean j* (first error position) - proxy for acceptance
+        if j_star is not None:
+            metrics["j_star_sum"] = j_star.float().sum()
+            metrics["j_star_total"] = torch.tensor(
+                float(j_star.numel()), device=logits.device
+            )
+
         return loss, metrics
 
     @conditional_torch_compile

src/speculators/models/dflash/domino.py

@@ -45,9 +45,9 @@ def forward(
         """Apply GRU-based prefix refinement to logits.
 
         Parameters:
-            hidden_states_4d: Tensor of shape ``(batch, num_anchors, block_size, hidden_size)``.
-            base_logits_4d: Tensor of shape ``(batch, num_anchors, block_size, vocab_size)``.
-            prev_token_ids: Tensor of shape ``(batch, num_anchors * block_size)`` with previous token IDs.
+            hidden_states_4d: ``(batch, num_anchors, block_size, hidden_size)``.
+            base_logits_4d: ``(batch, num_anchors, block_size, vocab_size)``.
+            prev_token_ids: ``(batch, num_anchors * block_size)`` previous token IDs.
             suffix_start: Position index within a block where the suffix begins.
             embed_tokens: Embedding layer for token lookup.
 

src/speculators/models/dspark/metrics.py

@@ -40,9 +40,7 @@ def _masked_decayed_mean(
     weighted = elementwise * loss_mask
     denominator = loss_mask.sum(dim=1) + _EPS
     if decay_fn is not None:
-        decay_mult = decay_fn(
-            pos_idx.to(weighted.dtype), elementwise_loss=elementwise
-        )
+        decay_mult = decay_fn(pos_idx.to(weighted.dtype), elementwise_loss=elementwise)
         weighted = weighted * decay_mult
         if normalize_by_decay:
             denominator = (loss_mask * decay_mult).sum(dim=1) + _EPS

tests/e2e/smoke/test_offline_training.py

@@ -54,8 +54,8 @@
                 "3",
                 "--projector-type",
                 "domino",
-                "--domino-lambda-decay-steps",
-                "100",
+                "--domino-lambda-decay-ratio",
+                "1.0",
             ],
             [1, 13, 25],
         ),  # Domino (DFlash + Domino correction head)

tests/integration/models/test_model_forward.py

@@ -328,18 +328,18 @@ def test_domino_head_receives_gradient(self):
         assert grad.abs().sum().item() > 0
 
     @pytest.mark.parametrize(
-        ("lambda_base_start", "lambda_base_decay_steps"),
+        ("lambda_base_start", "lambda_base_decay_ratio"),
         [
-            (1.0, 100),
-            (0.5, 0),
-            (0.0, 0),
+            (1.0, 1.0),
+            (0.5, 0.0),
+            (0.0, 0.0),
         ],
     )
-    def test_lambda_schedule_variants(self, lambda_base_start, lambda_base_decay_steps):
+    def test_lambda_schedule_variants(self, lambda_base_start, lambda_base_decay_ratio):
         model = make_dflash_model(
             projector_type="domino",
             lambda_base_start=lambda_base_start,
-            lambda_base_decay_steps=lambda_base_decay_steps,
+            lambda_base_decay_ratio=lambda_base_decay_ratio,
         )
         samples = _make_samples([128])
         batch = make_batch(max_len=MAX_LEN, samples=samples, hidden_size=HIDDEN_SIZE)