Add max_anchors random chain sampling to P-EAGLE

scripts/train.py

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

src/speculators/models/dflash/core.py

@@ -159,7 +159,7 @@ def from_training_args(
             transformer_layer_config=verifier_config,
             draft_vocab_size=kwargs["draft_vocab_size"],
             block_size=kwargs.get("block_size", 8),
-            max_anchors=kwargs.get("max_anchors", 3072),
+            max_anchors=kwargs.get("max_anchors") or 3072,
             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/config.py

@@ -56,6 +56,17 @@ class PEagleSpeculatorConfig(Eagle3SpeculatorConfig):
         description="Token ID used for padding unused positions in parallel groups",
     )
 
+    max_anchors: int | None = Field(
+        default=None,
+        description=(
+            "Maximum number of COD chain starting points. When set, "
+            "randomly subsamples valid positions for depth-1+ chains. "
+            "Depth 0 always retains the full sequence. "
+            "None means use all valid positions."
+        ),
+        ge=1,
+    )
+
     sink_size: int | None = Field(
         default=None,
         description=(

src/speculators/models/peagle/core.py

@@ -90,6 +90,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=self.config.max_anchors,
         )
         total_sampled = anchor_pos.shape[0]
 
@@ -224,6 +225,7 @@ def from_training_args(
             down_sample_ratio=kwargs.get("down_sample_ratio", 0.7),
             down_sample_ratio_min=kwargs.get("down_sample_ratio_min", 0.2),
             mask_token_id=kwargs.get("mask_token_id"),
+            max_anchors=kwargs.get("max_anchors"),
             sink_size=kwargs.get("sink_size"),
             max_context_window=kwargs.get("max_context_window"),
             speculators_config=SpeculatorsConfig(

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,6 +37,10 @@ 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]
+        all_valid_indices = all_valid_indices[perm].sort()[0]
+
     sample_indices = [torch.arange(seq_length, device=device)]
     n_per_depth = [seq_length]
     prev_indices = all_valid_indices

tests/unit/models/test_peagle_data.py

@@ -0,0 +1,94 @@
+"""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,
+        )
+        depth1_count = (depth == 1).sum().item()
+        assert depth1_count <= 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
+        anchor_pos, depth = generate_cod_sample_indices(
+            seq_length=seq_len,
+            loss_mask=loss_mask,
+            num_depths=4,
+            max_anchors=100,
+        )
+        assert anchor_pos.shape[0] > 0
+
+    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])