feat(peagle): add random anchor subsampling (max_anchors)

scripts/train.py

@@ -948,8 +948,9 @@ def parse_args():
     parser.add_argument(
         "--max-anchors",
         type=int,
-        default=256,
-        help="Maximum anchor positions for DFlash training (default: 256)",
+        default=3072,
+        help="Maximum anchor positions for DFlash, DSpark, "
+        "and P-EAGLE training (default: 3072).",
     )
     parser.add_argument(
         "--dflash-decay-gamma",

src/speculators/models/dflash/core.py

@@ -176,7 +176,9 @@ def _build_base_config_kwargs(
             "transformer_layer_config": verifier_config,
             "draft_vocab_size": kwargs["draft_vocab_size"],
             "block_size": block_size,
-            "max_anchors": kwargs.get("max_anchors", 3072),
+            "max_anchors": (
+                3072 if kwargs.get("max_anchors") is None else kwargs["max_anchors"]
+            ),
             "aux_hidden_state_layer_ids": target_layer_ids,
             "mask_token_id": kwargs.get("mask_token_id"),
             "sliding_window_non_causal": kwargs.get("sliding_window_non_causal", False),

src/speculators/models/peagle/core.py

@@ -63,6 +63,7 @@ def forward(
         loss_mask: torch.Tensor | None = None,
         verifier_last_hidden_states: torch.Tensor | None = None,
         loss_config: LossConfig | None = None,
+        max_anchors: int | None = None,
         **kwargs,
     ):
         """
@@ -97,6 +98,7 @@ def forward(
             num_depths=self.num_depths,
             down_sample_ratio=self.down_sample_ratio,
             down_sample_ratio_min=self.down_sample_ratio_min,
+            max_anchors=max_anchors,
         )
         total_sampled = anchor_pos.shape[0]
 
@@ -267,4 +269,6 @@ def get_trainer_kwargs(**kwargs) -> tuple[dict, dict]:
             Tuple of (train_call_kwargs, val_call_kwargs)
         """
         loss_config = resolve_loss_config(kwargs["loss_fn"])
-        return {"loss_config": loss_config}, {"loss_config": loss_config}
+        max_anchors = kwargs.get("max_anchors")
+        shared = {"loss_config": loss_config, "max_anchors": max_anchors}
+        return dict(shared), dict(shared)

src/speculators/models/peagle/data.py

@@ -10,6 +10,7 @@ def generate_cod_sample_indices(
     down_sample_ratio: float = 0.7,
     down_sample_ratio_min: float = 0.2,
     filter_position_zero: bool = True,
+    max_anchors: int | None = None,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """
     Generate sampling indices for parallel sequences using COD sampling.
@@ -36,12 +37,18 @@ def generate_cod_sample_indices(
     device = loss_mask.device
     all_valid_indices = torch.where(loss_mask == 1)[0]
 
+    if max_anchors is not None and all_valid_indices.shape[0] > max_anchors:
+        perm = torch.randperm(all_valid_indices.shape[0], device=device)[:max_anchors]
+        seed_indices = all_valid_indices[perm].sort()[0]
+    else:
+        seed_indices = all_valid_indices
+
     sample_indices = [torch.arange(seq_length, device=device)]
     n_per_depth = [seq_length]
-    prev_indices = all_valid_indices
+    prev_indices = seed_indices
 
     for d in range(1, num_depths):
-        valid_length = max(0, all_valid_indices.shape[0] - d)
+        valid_length = max(0, seed_indices.shape[0] - d)
         ratio = max(down_sample_ratio**d, down_sample_ratio_min)
         sample_size = int(valid_length * ratio)
 

tests/unit/models/test_peagle_data.py

@@ -0,0 +1,104 @@
+"""Unit tests for P-EAGLE COD sampling with max_anchors."""
+
+import torch
+
+from speculators.models.peagle.data import generate_cod_sample_indices
+
+
+def _loss_mask(seq_length: int) -> torch.Tensor:
+    return torch.ones(1, seq_length, dtype=torch.float32)
+
+
+class TestMaxAnchors:
+    def test_depth0_is_full_sequence(self):
+        """Depth 0 should always be the full sequence regardless of max_anchors."""
+        seq_len = 32
+        anchor_pos, depth = generate_cod_sample_indices(
+            seq_length=seq_len,
+            loss_mask=_loss_mask(seq_len),
+            num_depths=4,
+            max_anchors=4,
+        )
+        depth0_positions = anchor_pos[depth == 0]
+        assert depth0_positions.shape[0] == seq_len
+        assert torch.equal(depth0_positions, torch.arange(seq_len))
+
+    def test_max_anchors_caps_chains(self):
+        """With max_anchors, depth-1+ chains should not exceed max_anchors."""
+        seq_len = 64
+        max_anchors = 8
+        anchor_pos, depth = generate_cod_sample_indices(
+            seq_length=seq_len,
+            loss_mask=_loss_mask(seq_len),
+            num_depths=4,
+            max_anchors=max_anchors,
+        )
+        for d in range(1, 4):
+            assert (depth == d).sum().item() <= max_anchors
+
+    def test_max_anchors_preserves_full_depth0(self):
+        """Depth 0 count should equal seq_length even with small max_anchors."""
+        seq_len = 128
+        anchor_pos, depth = generate_cod_sample_indices(
+            seq_length=seq_len,
+            loss_mask=_loss_mask(seq_len),
+            num_depths=8,
+            max_anchors=4,
+        )
+        assert (depth == 0).sum().item() == seq_len
+
+    def test_max_anchors_none_uses_all(self):
+        """max_anchors=None should use all valid positions (default behavior)."""
+        seq_len = 32
+        torch.manual_seed(42)
+        anchor_pos_none, depth_none = generate_cod_sample_indices(
+            seq_length=seq_len,
+            loss_mask=_loss_mask(seq_len),
+            num_depths=4,
+            max_anchors=None,
+        )
+        torch.manual_seed(42)
+        anchor_pos_default, depth_default = generate_cod_sample_indices(
+            seq_length=seq_len,
+            loss_mask=_loss_mask(seq_len),
+            num_depths=4,
+        )
+        assert torch.equal(anchor_pos_none, anchor_pos_default)
+        assert torch.equal(depth_none, depth_default)
+
+    def test_max_anchors_fewer_valid_than_cap(self):
+        """When valid positions < max_anchors, all valid positions are used."""
+        seq_len = 16
+        loss_mask = torch.zeros(1, seq_len)
+        loss_mask[0, :5] = 1
+        torch.manual_seed(42)
+        anchor_pos_capped, depth_capped = generate_cod_sample_indices(
+            seq_length=seq_len,
+            loss_mask=loss_mask,
+            num_depths=4,
+            max_anchors=100,
+        )
+        torch.manual_seed(42)
+        anchor_pos_none, depth_none = generate_cod_sample_indices(
+            seq_length=seq_len,
+            loss_mask=loss_mask,
+            num_depths=4,
+            max_anchors=None,
+        )
+        assert (depth_capped == 0).sum().item() == seq_len
+        assert torch.equal(anchor_pos_capped, anchor_pos_none)
+        assert torch.equal(depth_capped, depth_none)
+
+    def test_max_anchors_sorted_order(self):
+        """Subsampled anchors should be in sorted order for causal masking."""
+        seq_len = 64
+        anchor_pos, depth = generate_cod_sample_indices(
+            seq_length=seq_len,
+            loss_mask=_loss_mask(seq_len),
+            num_depths=4,
+            max_anchors=8,
+        )
+        for d in range(1, 4):
+            d_anchors = anchor_pos[depth == d]
+            if d_anchors.shape[0] > 1:
+                assert torch.all(d_anchors[1:] >= d_anchors[:-1])