Add max_anchors and max_context_window support to P-EAGLE

scripts/train.py

@@ -748,8 +748,11 @@ 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 for training. "
+            "Defaults to 3072 for DFlash, unlimited for P-EAGLE."
+        ),
     )
     parser.add_argument(
         "--draft-attn-impl",
@@ -778,6 +781,15 @@ def parse_args():
         default=0.2,
         help="Minimum retention ratio for COD sampling in P-EAGLE (default: 0.2)",
     )
+    parser.add_argument(
+        "--max-context-window",
+        type=int,
+        default=4096,
+        help=(
+            "Hard cap on contiguous window size when max-anchors is set "
+            "for P-EAGLE (default: 4096)"
+        ),
+    )
     parser.add_argument(
         "--sliding-window",
         type=int,

src/speculators/models/dflash/core.py

@@ -156,7 +156,9 @@ 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=(
+                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/config.py

@@ -51,6 +51,25 @@ class PEagleSpeculatorConfig(Eagle3SpeculatorConfig):
         le=1.0,
     )
 
+    max_anchors: int | None = Field(
+        default=None,
+        description=(
+            "Maximum number of COD chain starting points (loss_mask=1 tokens) "
+            "to use during training. When set, selects a contiguous window of "
+            "the sequence to cap memory usage. None means use all positions."
+        ),
+        ge=1,
+    )
+
+    max_context_window: int = Field(
+        default=4096,
+        description=(
+            "Hard cap on the contiguous window size when max_anchors is set. "
+            "Prevents sparse loss masks from reinflating the window."
+        ),
+        ge=1,
+    )
+
     mask_token_id: int | None = Field(
         default=None,
         description="Token ID used for padding unused positions in parallel groups",

src/speculators/models/peagle/core.py

@@ -42,6 +42,8 @@ def __init__(
         self.num_depths = config.num_depths
         self.down_sample_ratio = config.down_sample_ratio
         self.down_sample_ratio_min = config.down_sample_ratio_min
+        self.max_anchors = config.max_anchors
+        self.max_context_window = config.max_context_window
         self.mask_token_id = config.mask_token_id
 
         # Learnable mask_hidden parameter for padding unsampled positions
@@ -91,6 +93,8 @@ 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.max_anchors,
+            max_context_window=self.max_context_window,
         )
         total_sampled = anchor_pos.shape[0]
 
@@ -214,6 +218,8 @@ def from_training_args(
             num_depths=kwargs.get("num_depths", 8),
             down_sample_ratio=kwargs.get("down_sample_ratio", 0.7),
             down_sample_ratio_min=kwargs.get("down_sample_ratio_min", 0.2),
+            max_anchors=kwargs.get("max_anchors"),
+            max_context_window=kwargs.get("max_context_window", 4096),
             mask_token_id=kwargs.get("mask_token_id"),
             speculators_config=SpeculatorsConfig(
                 algorithm="peagle",

src/speculators/models/peagle/data.py

@@ -10,6 +10,8 @@ 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,
+    max_context_window: int = 4096,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """
     Generate sampling indices for parallel sequences using COD sampling.
@@ -18,13 +20,22 @@ def generate_cod_sample_indices(
     decay: depth 0 retains all n positions, depth 1 retains n*r positions,
     depth 2 retains n*r^2 positions, etc.
 
+    When max_anchors is set, selects a contiguous window of the original sequence
+    containing up to max_anchors valid (loss_mask=1) positions as COD starting
+    points. The window preserves all intervening tokens (including prompts)
+    to maintain context.
+
     Args:
         seq_length: Length of the sequence
         loss_mask: Binary mask indicating valid training positions [batch, seq_len]
         num_depths: Number of parallel prediction groups (K)
         down_sample_ratio: Geometric decay ratio r in (0,1)
         down_sample_ratio_min: Minimum retention ratio floor
         filter_position_zero: Whether to filter out position 0 from candidates
+        max_anchors: Maximum number of COD chain starting points. None means
+            use all positions.
+        max_context_window: Hard cap on contiguous window size to prevent
+            sparse loss masks from reinflating the window.
 
     Returns:
         Tuple of:
@@ -36,8 +47,30 @@ def generate_cod_sample_indices(
     device = loss_mask.device
     all_valid_indices = torch.where(loss_mask == 1)[0]
 
-    sample_indices = [torch.arange(seq_length, device=device)]
-    n_per_depth = [seq_length]
+    if max_anchors is not None and all_valid_indices.shape[0] > 0:
+        if all_valid_indices.shape[0] > max_anchors:
+            n_valid = all_valid_indices.shape[0]
+            start_idx = int(torch.randint(0, n_valid - max_anchors + 1, (1,)).item())
+            selected_valid = all_valid_indices[start_idx : start_idx + max_anchors]
+        else:
+            selected_valid = all_valid_indices
+
+        window_start = int(selected_valid[0].item())
+        window_end = int(selected_valid[-1].item()) + 1
+
+        if (window_end - window_start) > max_context_window:
+            window_end = min(window_start + max_context_window, seq_length)
+            selected_valid = selected_valid[
+                (selected_valid >= window_start) & (selected_valid < window_end)
+            ]
+        all_valid_indices = selected_valid
+
+        sample_indices = [torch.arange(window_start, window_end, device=device)]
+        n_per_depth = [window_end - window_start]
+    else:
+        sample_indices = [torch.arange(seq_length, device=device)]
+        n_per_depth = [seq_length]
+
     prev_indices = all_valid_indices
 
     for d in range(1, num_depths):

tests/integration/conftest.py

@@ -163,6 +163,8 @@ def make_peagle_model(
     draft_vocab_size: int = 64,
     num_depths: int = 4,
     down_sample_ratio: float = 0.7,
+    max_anchors: int | None = None,
+    max_context_window: int = 4096,
     device: str = "cuda:0",
     dtype: torch.dtype = torch.bfloat16,
 ) -> PEagleDraftModel:
@@ -175,6 +177,8 @@ def make_peagle_model(
         num_depths=num_depths,
         down_sample_ratio=down_sample_ratio,
         down_sample_ratio_min=0.2,
+        max_anchors=max_anchors,
+        max_context_window=max_context_window,
         mask_token_id=0,
         speculators_config=SpeculatorsConfig(
             algorithm="peagle",

tests/integration/models/test_model_forward.py

@@ -338,6 +338,25 @@ def test_varying_down_sample_ratio(self, down_sample_ratio):
         assert loss.isfinite()
         loss.backward()
 
+    @pytest.mark.parametrize("max_anchors", [4, 16, None])
+    def test_varying_max_anchors(self, max_anchors):
+        model = make_peagle_model(max_anchors=max_anchors)
+        samples = _make_samples([128])
+        batch = make_batch(max_len=MAX_LEN, samples=samples, hidden_size=HIDDEN_SIZE)
+        draft_tokens, loss, metrics = model(**batch)
+
+        assert loss.isfinite()
+        loss.backward()
+
+    def test_max_context_window(self):
+        model = make_peagle_model(max_anchors=16, max_context_window=32)
+        samples = _make_samples([128])
+        batch = make_batch(max_len=MAX_LEN, samples=samples, hidden_size=HIDDEN_SIZE)
+        draft_tokens, loss, metrics = model(**batch)
+
+        assert loss.isfinite()
+        loss.backward()
+
 
 @requires_cuda
 @requires_transformers_version("5.2.0")

tests/unit/models/test_peagle_data.py

@@ -0,0 +1,59 @@
+"""Unit tests for P-EAGLE COD sampling logic."""
+
+import torch
+
+from speculators.models.peagle.data import generate_cod_sample_indices
+
+
+class TestMaxContextWindow:
+    def test_applied_when_anchors_within_limit(self):
+        """max_context_window caps the window even when anchors <= max_anchors."""
+        seq_length = 256
+        loss_mask = torch.zeros(1, seq_length)
+        # Place 4 valid positions spread far apart (indices 10, 80, 160, 240)
+        valid_positions = [10, 80, 160, 240]
+        for pos in valid_positions:
+            loss_mask[0, pos] = 1
+
+        anchor_pos, depth = generate_cod_sample_indices(
+            seq_length=seq_length,
+            loss_mask=loss_mask,
+            max_anchors=8,
+            max_context_window=64,
+        )
+
+        depth_0_mask = depth == 0
+        depth_0_positions = anchor_pos[depth_0_mask]
+        window_size = depth_0_positions.shape[0]
+        assert window_size <= 64
+
+    def test_applied_when_anchors_exceed_limit(self):
+        """max_context_window caps the window when anchors > max_anchors."""
+        seq_length = 512
+        loss_mask = torch.ones(1, seq_length)
+
+        anchor_pos, depth = generate_cod_sample_indices(
+            seq_length=seq_length,
+            loss_mask=loss_mask,
+            max_anchors=16,
+            max_context_window=32,
+        )
+
+        depth_0_mask = depth == 0
+        depth_0_positions = anchor_pos[depth_0_mask]
+        window_size = depth_0_positions.shape[0]
+        assert window_size <= 32
+
+    def test_no_windowing_without_max_anchors(self):
+        """Without max_anchors, full sequence is used."""
+        seq_length = 128
+        loss_mask = torch.ones(1, seq_length)
+
+        anchor_pos, depth = generate_cod_sample_indices(
+            seq_length=seq_length,
+            loss_mask=loss_mask,
+            max_anchors=None,
+        )
+
+        depth_0_mask = depth == 0
+        assert anchor_pos[depth_0_mask].shape[0] == seq_length