Add max_anchors random chain sampling to P-EAGLE

scripts/train.py

@@ -618,8 +618,10 @@ 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 DFlash: caps training anchors "
+        "(default: 3072). For P-EAGLE: caps COD chain starting points; "
+        "depth 0 always retains the full sequence (default: unlimited).",
     )
     # P-EAGLE specific parameters
     parser.add_argument(
@@ -640,6 +642,22 @@ def parse_args():
         default=0.2,
         help="Minimum retention ratio for COD sampling in P-EAGLE (default: 0.2)",
     )
+    parser.add_argument(
+        "--sink-size",
+        type=int,
+        default=None,
+        help="Number of initial tokens per document to retain as attention sinks "
+        "(StreamingLLM). Must be set together with --max-context-window. "
+        "When unset, P-EAGLE uses full causal attention for depth-0 tokens.",
+    )
+    parser.add_argument(
+        "--max-context-window",
+        type=int,
+        default=None,
+        help="Size of the local sliding window for depth-0 KV attention "
+        "(StreamingLLM). Must be set together with --sink-size. "
+        "When unset, P-EAGLE uses full causal attention for depth-0 tokens.",
+    )
     # Dataloader parameters
     parser.add_argument(
         "--num-workers", type=int, default=12, help="Number of dataloader workers"

src/speculators/models/dflash/core.py

@@ -138,7 +138,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"),
             speculators_config=SpeculatorsConfig(

src/speculators/models/peagle/attention.py

@@ -3,11 +3,71 @@
 import torch
 
 
+def _build_document_ids(
+    lengths: torch.Tensor,
+    total_seq_len: int,
+) -> torch.Tensor:
+    """Build a ``document_ids`` tensor from packed document lengths.
+
+    Returns:
+        [total_seq_len] tensor mapping each position to its document index,
+        with -1 for padding positions.
+    """
+    return torch.cat(
+        [
+            torch.repeat_interleave(
+                torch.arange(
+                    lengths.shape[0], device=lengths.device, dtype=torch.long
+                ),
+                lengths,
+            ),
+            -1
+            * torch.ones(
+                total_seq_len - lengths.sum(),
+                device=lengths.device,
+                dtype=torch.long,
+            ),
+        ]
+    ).contiguous()
+
+
+@torch.compiler.disable
+def _compute_doc_start_positions(document_ids: torch.Tensor) -> torch.Tensor:
+    """Compute the start position of each document in a packed sequence.
+
+    Args:
+        document_ids: [total_seq_len] maps each position to its doc index,
+            -1 for padding.
+
+    Returns:
+        [total_seq_len] tensor where entry *i* is the first position of the
+        document that position *i* belongs to.  Padding positions get -1.
+    """
+    valid = document_ids != -1
+    result = torch.full_like(document_ids, -1)
+    if not valid.any():
+        return result
+
+    positions = torch.arange(len(document_ids), device=document_ids.device)
+    unique_docs, inverse = torch.unique(document_ids[valid], return_inverse=True)
+    first_pos = torch.full(
+        (unique_docs.shape[0],),
+        len(document_ids),
+        device=document_ids.device,
+        dtype=torch.long,
+    )
+    first_pos.scatter_reduce_(0, inverse, positions[valid], reduce="amin")
+    result[valid] = first_pos[inverse]
+    return result
+
+
 def create_peagle_mask_mod(
     anchor_pos: torch.Tensor,  # shape: [total_sampled]
     depth: torch.Tensor,  # shape: [total_sampled]
     lengths: torch.Tensor,  # shape: [batch_size]
     total_seq_len: int,
+    sink_size: int | None = None,
+    max_context_window: int | None = None,
 ):
     """
     Create a flex attention mask modifier for P-EAGLE parallel groups.
@@ -18,6 +78,10 @@ def create_peagle_mask_mod(
     This function creates a mask where each element can attend to only to previous
     elements in the same sampling chain/rollout and previous context in the base sample.
 
+    When ``sink_size`` and ``max_context_window`` are both set (StreamingLLM mode),
+    depth-0 causal attention is restricted to:
+    - the first ``sink_size`` positions of each document (attention sinks), and
+    - the most recent ``max_context_window`` positions before the query.
 
     Args:
         anchor_pos: The starting position in the original sequence the current
@@ -26,6 +90,10 @@ def create_peagle_mask_mod(
         lengths: The length of each document. Used to produce a document mask to prevent
             cross contamination
         total_seq_len: int, combined padded length of the original sequences
+        sink_size: Number of initial tokens per document retained as attention
+            sinks. Must be set together with max_context_window.
+        max_context_window: Size of the local sliding window for depth-0 KV
+            attention. Must be set together with sink_size.
 
     Args example:
 
@@ -45,22 +113,11 @@ def create_peagle_mask_mod(
         A mask_mod function compatible with flex_attention create_block_mask
     """
 
-    # Generate sample_ids to prevent cross-sample attention
-    document_ids = torch.repeat_interleave(
-        torch.arange(lengths.shape[0], device=lengths.device, dtype=torch.long), lengths
-    )
-    # Pad ids with -1 to indicate padding
-    document_ids = torch.cat(
-        [
-            document_ids,
-            -1
-            * torch.ones(
-                total_seq_len - document_ids.shape[0],
-                device=lengths.device,
-                dtype=torch.long,
-            ),
-        ]
-    ).contiguous()
+    document_ids = _build_document_ids(lengths, total_seq_len)
+
+    use_streaming = sink_size is not None and max_context_window is not None
+    if use_streaming:
+        doc_start_positions = _compute_doc_start_positions(document_ids)
 
     def peagle_mask_mod(_b, _h, q_idx, kv_idx):
         q_anchor_pos = anchor_pos[q_idx]
@@ -75,10 +132,18 @@ def peagle_mask_mod(_b, _h, q_idx, kv_idx):
         in_depth_order = q_depth >= kv_depth
         is_anchor_causal = q_anchor_pos >= kv_anchor_pos
 
+        depth0_attention = kv_depth0 & is_anchor_causal
+
+        if use_streaming:
+            kv_doc_start = doc_start_positions[kv_anchor_pos]
+            is_sink = (kv_anchor_pos - kv_doc_start) < sink_size
+            in_window = kv_anchor_pos >= (q_anchor_pos - max_context_window)
+            depth0_attention = depth0_attention & (is_sink | in_window)
+
         return (
             is_not_padding
             & same_document
-            & ((kv_depth0 & is_anchor_causal) | (same_rollout & in_depth_order))
+            & (depth0_attention | (same_rollout & in_depth_order))
         )
 
     return peagle_mask_mod

src/speculators/models/peagle/config.py

@@ -1,6 +1,6 @@
 from typing import Literal
 
-from pydantic import Field
+from pydantic import Field, model_validator
 
 from speculators import SpeculatorModelConfig
 from speculators.models.eagle3.config import Eagle3SpeculatorConfig
@@ -56,6 +56,37 @@ 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=(
+            "Number of initial tokens per document retained as attention "
+            "sinks (StreamingLLM). Must be set together with "
+            "max_context_window. When None, full causal attention is used."
+        ),
+        ge=1,
+    )
+
+    max_context_window: int | None = Field(
+        default=None,
+        description=(
+            "Size of the local sliding window for depth-0 KV attention "
+            "(StreamingLLM). Must be set together with sink_size. "
+            "When None, full causal attention is used."
+        ),
+        ge=1,
+    )
+
     # Override Eagle3 default: P-EAGLE requires trainable embeddings
     # (matches p-eagle-train)
     embed_requires_grad: bool = Field(
@@ -65,3 +96,11 @@ class PEagleSpeculatorConfig(Eagle3SpeculatorConfig):
             "training (True for P-EAGLE)"
         ),
     )
+
+    @model_validator(mode="after")
+    def _validate_streaming_llm(self) -> "PEagleSpeculatorConfig":
+        if (self.sink_size is None) != (self.max_context_window is None):
+            raise ValueError(
+                "sink_size and max_context_window must both be set or both be None"
+            )
+        return self

src/speculators/models/peagle/core.py

@@ -93,6 +93,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]
 
@@ -133,6 +134,8 @@ def forward(
             depth=depth,
             lengths=lengths,
             total_seq_len=seq_length,
+            sink_size=self.config.sink_size,
+            max_context_window=self.config.max_context_window,
         )
 
         attention_mask = create_block_mask(  # type: ignore[assignment]
@@ -216,7 +219,10 @@ 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"),
             mask_token_id=kwargs.get("mask_token_id"),
+            sink_size=kwargs.get("sink_size"),
+            max_context_window=kwargs.get("max_context_window"),
             speculators_config=SpeculatorsConfig(
                 algorithm="peagle",
                 proposal_methods=[

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.
@@ -25,6 +26,9 @@ def generate_cod_sample_indices(
         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. When set and
+            there are more valid positions, a random subset is sampled. Depth 0
+            always retains the full sequence. None means use all positions.
 
     Returns:
         Tuple of:
@@ -36,6 +40,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