diff --git a/scripts/train.py b/scripts/train.py index 0bb4cf403..e3b192dcb 100644 --- a/scripts/train.py +++ b/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" diff --git a/src/speculators/models/dflash/core.py b/src/speculators/models/dflash/core.py index ed79e471d..ec1fc14b6 100644 --- a/src/speculators/models/dflash/core.py +++ b/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( diff --git a/src/speculators/models/peagle/attention.py b/src/speculators/models/peagle/attention.py index ee4e9f9ae..9e92003b4 100644 --- a/src/speculators/models/peagle/attention.py +++ b/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 diff --git a/src/speculators/models/peagle/config.py b/src/speculators/models/peagle/config.py index 6c860720a..20274ee51 100644 --- a/src/speculators/models/peagle/config.py +++ b/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 diff --git a/src/speculators/models/peagle/core.py b/src/speculators/models/peagle/core.py index 037d7b652..28569503b 100644 --- a/src/speculators/models/peagle/core.py +++ b/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=[ diff --git a/src/speculators/models/peagle/data.py b/src/speculators/models/peagle/data.py index f06b72e0a..49bf638fd 100644 --- a/src/speculators/models/peagle/data.py +++ b/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 diff --git a/tests/integration/models/test_peagle_streaming.py b/tests/integration/models/test_peagle_streaming.py new file mode 100644 index 000000000..35424e159 --- /dev/null +++ b/tests/integration/models/test_peagle_streaming.py @@ -0,0 +1,142 @@ +"""Integration test for P-EAGLE with StreamingLLM (attention sinks).""" + +import copy + +import pytest +import torch + +from speculators import SpeculatorsConfig, VerifierConfig +from speculators.models.eagle3.data import shift_batch +from speculators.models.peagle.config import PEagleSpeculatorConfig +from speculators.models.peagle.core import PEagleDraftModel +from speculators.proposals.greedy import GreedyTokenProposalConfig +from speculators.train.data import create_collate_fn +from tests.conftest import requires_cuda + +_TINY_LLAMA_CONFIG = pytest.importorskip( + "transformers.models.llama.configuration_llama" +).LlamaConfig( + vocab_size=128, + hidden_size=64, + intermediate_size=256, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=4, + head_dim=16, + max_position_embeddings=256, + rms_norm_eps=1e-6, + tie_word_embeddings=False, + _attn_implementation="simple_flex_attention", +) + + +def _make_peagle_model( + sink_size=None, + max_context_window=None, + device="cuda:0", + dtype=torch.bfloat16, +): + transformer_config = copy.deepcopy(_TINY_LLAMA_CONFIG) + config = PEagleSpeculatorConfig( + transformer_layer_config=transformer_config, + draft_vocab_size=64, + norm_before_residual=False, + embed_requires_grad=True, + num_depths=4, + down_sample_ratio=0.7, + down_sample_ratio_min=0.2, + mask_token_id=0, + sink_size=sink_size, + max_context_window=max_context_window, + speculators_config=SpeculatorsConfig( + algorithm="peagle", + proposal_methods=[GreedyTokenProposalConfig(speculative_tokens=4)], + default_proposal_method="greedy", + verifier=VerifierConfig( + name_or_path=None, + architectures=["LlamaForCausalLM"], + ), + ), + ) + model = PEagleDraftModel(config) + with torch.no_grad(): + for param in model.parameters(): + if param.isnan().any(): + torch.nn.init.normal_(param, mean=0.0, std=0.02) + for buf in model.buffers(): + if buf.is_floating_point() and buf.isnan().any(): + buf.zero_() + return model.to(device=device, dtype=dtype) + + +def _make_batch(seq_lengths, hidden_size=64, max_len=128, device="cuda:0"): + samples = [] + for sl in seq_lengths: + samples.append( + { + "hidden_states": torch.randn(sl, 3 * hidden_size, dtype=torch.bfloat16), + "input_ids": torch.randint(0, 128, (sl,)), + "verifier_last_hidden_states": torch.randn( + sl, hidden_size, dtype=torch.bfloat16 + ), + "loss_mask": torch.ones(sl, dtype=torch.bfloat16), + "lengths": torch.tensor([sl], dtype=torch.long), + "position_ids": torch.arange(sl, dtype=torch.long), + } + ) + collate_fn = create_collate_fn(max_len, hidden_size, preprocess=shift_batch) + batch = collate_fn(samples) + return {k: v.to(device) for k, v in batch.items()} + + +@requires_cuda +class TestPEagleStreaming: + def test_forward_backward_streaming(self): + model = _make_peagle_model(sink_size=4, max_context_window=32) + batch = _make_batch([128]) + _, loss, metrics = model(**batch) + + assert loss.isfinite(), f"Loss is not finite: {loss.item()}" + assert "loss_sum" in metrics + loss.backward() + + def test_forward_backward_no_streaming(self): + """Baseline: verify forward+backward works without streaming too.""" + model = _make_peagle_model() + batch = _make_batch([128]) + _, loss, metrics = model(**batch) + + assert loss.isfinite(), f"Loss is not finite: {loss.item()}" + loss.backward() + + def test_forward_backward_streaming_multi_doc(self): + model = _make_peagle_model(sink_size=4, max_context_window=16) + batch = _make_batch([64, 64]) + _, loss, metrics = model(**batch) + + assert loss.isfinite(), f"Loss is not finite: {loss.item()}" + loss.backward() + + def test_streaming_config_roundtrip(self, tmp_path): + """Config with streaming params survives save/load.""" + model = _make_peagle_model(sink_size=8, max_context_window=64) + model.save_pretrained(str(tmp_path)) + + loaded_config = PEagleSpeculatorConfig.from_pretrained(str(tmp_path)) + assert loaded_config.sink_size == 8 + assert loaded_config.max_context_window == 64 + + def test_config_validation_both_or_neither(self): + """Setting only one of sink_size/max_context_window raises.""" + with pytest.raises(ValueError, match="both be set or both be None"): + PEagleSpeculatorConfig( + transformer_layer_config=copy.deepcopy(_TINY_LLAMA_CONFIG), + sink_size=4, + max_context_window=None, + ) + with pytest.raises(ValueError, match="both be set or both be None"): + PEagleSpeculatorConfig( + transformer_layer_config=copy.deepcopy(_TINY_LLAMA_CONFIG), + sink_size=None, + max_context_window=32, + ) diff --git a/tests/unit/models/test_peagle_attention.py b/tests/unit/models/test_peagle_attention.py new file mode 100644 index 000000000..552e9eba9 --- /dev/null +++ b/tests/unit/models/test_peagle_attention.py @@ -0,0 +1,334 @@ +"""Unit tests for P-EAGLE attention mask with StreamingLLM (attention sinks).""" + +import torch + +from speculators.models.peagle.attention import ( + _build_document_ids, + _compute_doc_start_positions, + create_peagle_mask_mod, +) + + +def _lengths_to_document_ids(lengths, total_seq_len): + return _build_document_ids(lengths, total_seq_len) + + +def _evaluate_mask(mask_mod, total_sampled): + """Evaluate mask_mod element-wise over the q x kv grid.""" + zero = torch.zeros((), dtype=torch.long) + mask = torch.zeros(total_sampled, total_sampled, dtype=torch.bool) + for q in range(total_sampled): + for kv in range(total_sampled): + mask[q, kv] = bool(mask_mod(zero, zero, torch.tensor(q), torch.tensor(kv))) + return mask + + +def _simple_cod_indices(seq_length, depths_to_sample=None): + """Create simple COD-like indices for testing. + + Depth 0: all positions [0, seq_length) + Depth 1: even positions + """ + anchor_pos = list(range(seq_length)) + depth_vals = [0] * seq_length + + if depths_to_sample is not None: + for d, positions in depths_to_sample.items(): + for p in positions: + anchor_pos.append(p) + depth_vals.append(d) + + return ( + torch.tensor(anchor_pos, dtype=torch.long), + torch.tensor(depth_vals, dtype=torch.long), + ) + + +class TestComputeDocStartPositions: + def test_single_document(self): + doc_ids = torch.tensor([0, 0, 0, 0, 0]) + result = _compute_doc_start_positions(doc_ids) + assert result.tolist() == [0, 0, 0, 0, 0] + + def test_two_documents(self): + doc_ids = torch.tensor([0, 0, 0, 1, 1, 1]) + result = _compute_doc_start_positions(doc_ids) + assert result.tolist() == [0, 0, 0, 3, 3, 3] + + def test_with_padding(self): + doc_ids = torch.tensor([0, 0, 0, 1, 1, -1, -1]) + result = _compute_doc_start_positions(doc_ids) + assert result.tolist() == [0, 0, 0, 3, 3, -1, -1] + + def test_all_padding(self): + doc_ids = torch.tensor([-1, -1, -1]) + result = _compute_doc_start_positions(doc_ids) + assert result.tolist() == [-1, -1, -1] + + def test_three_documents(self): + doc_ids = torch.tensor([0, 0, 1, 1, 1, 2, 2]) + result = _compute_doc_start_positions(doc_ids) + assert result.tolist() == [0, 0, 2, 2, 2, 5, 5] + + +class TestNoStreamingUnchanged: + """When sink_size=None (no StreamingLLM), mask is identical to original.""" + + def test_single_doc_no_streaming(self): + seq_len = 8 + lengths = torch.tensor([seq_len]) + anchor_pos, depth = _simple_cod_indices(seq_len) + + mask_with = create_peagle_mask_mod( + anchor_pos, depth, lengths, seq_len, sink_size=None, max_context_window=None + ) + mask_without = create_peagle_mask_mod(anchor_pos, depth, lengths, seq_len) + + total = anchor_pos.shape[0] + dense_with = _evaluate_mask(mask_with, total) + dense_without = _evaluate_mask(mask_without, total) + assert torch.equal(dense_with, dense_without) + + +class TestStreamingSingleDoc: + """StreamingLLM with a single document.""" + + def test_sink_and_window(self): + seq_len = 20 + sink_size = 4 + window = 6 + lengths = torch.tensor([seq_len]) + anchor_pos, depth = _simple_cod_indices(seq_len) + + mask_mod = create_peagle_mask_mod( + anchor_pos, + depth, + lengths, + seq_len, + sink_size=sink_size, + max_context_window=window, + ) + mask = _evaluate_mask(mask_mod, anchor_pos.shape[0]) + + # Query at position 15 (depth 0, index 15) attending to depth-0 KVs: + q_idx = 15 + for kv_idx in range(seq_len): + kv_pos = kv_idx + is_sink = kv_pos < sink_size # 0,1,2,3 + in_window = kv_pos >= (q_idx - window) # 9..15 + is_causal = kv_pos <= q_idx + + expected = is_causal and (is_sink or in_window) + assert mask[q_idx, kv_idx] == expected, ( + f"q={q_idx}, kv={kv_idx}: expected {expected}, " + f"got {mask[q_idx, kv_idx]}" + ) + + def test_gap_positions_blocked(self): + """Positions in the middle (between sink and window) should be blocked.""" + seq_len = 20 + sink_size = 3 + window = 4 + lengths = torch.tensor([seq_len]) + anchor_pos, depth = _simple_cod_indices(seq_len) + + mask_mod = create_peagle_mask_mod( + anchor_pos, + depth, + lengths, + seq_len, + sink_size=sink_size, + max_context_window=window, + ) + mask = _evaluate_mask(mask_mod, anchor_pos.shape[0]) + + q_idx = 15 + # Gap positions: 3..10 (not sinks, not in window [11..15]) + for kv_idx in range(sink_size, q_idx - window): + assert not mask[q_idx, kv_idx], ( + f"Gap position kv={kv_idx} should be blocked for q={q_idx}" + ) + + def test_early_query_no_gap(self): + """When query is near start, window and sink overlap — all positions visible.""" + seq_len = 20 + sink_size = 5 + window = 8 + lengths = torch.tensor([seq_len]) + anchor_pos, depth = _simple_cod_indices(seq_len) + + mask_mod = create_peagle_mask_mod( + anchor_pos, + depth, + lengths, + seq_len, + sink_size=sink_size, + max_context_window=window, + ) + mask = _evaluate_mask(mask_mod, anchor_pos.shape[0]) + + # Query at position 6: sinks cover 0-4, window covers max(0,6-8)=0..6 + # So all positions 0..6 should be visible + q_idx = 6 + for kv_idx in range(q_idx + 1): + assert mask[q_idx, kv_idx], ( + f"Position kv={kv_idx} should be visible for q={q_idx}" + ) + + +class TestStreamingMultiDoc: + """StreamingLLM respects document boundaries in packed sequences.""" + + def test_cross_doc_sinks_blocked(self): + """Sinks from doc0 should not be visible to queries in doc1.""" + doc0_len = 10 + doc1_len = 10 + total = doc0_len + doc1_len + lengths = torch.tensor([doc0_len, doc1_len]) + anchor_pos, depth = _simple_cod_indices(total) + + mask_mod = create_peagle_mask_mod( + anchor_pos, + depth, + lengths, + total, + sink_size=4, + max_context_window=6, + ) + mask = _evaluate_mask(mask_mod, anchor_pos.shape[0]) + + # Query in doc1 at position 18 + q_idx = 18 + # Doc0 positions (0-9) should all be blocked + for kv_idx in range(doc0_len): + assert not mask[q_idx, kv_idx], ( + f"Doc0 position kv={kv_idx} should be blocked for doc1 q={q_idx}" + ) + + def test_doc1_has_own_sinks(self): + """Doc1's sink tokens are relative to doc1's start position.""" + doc0_len = 10 + doc1_len = 15 + total = doc0_len + doc1_len + sink_size = 3 + window = 4 + lengths = torch.tensor([doc0_len, doc1_len]) + anchor_pos, depth = _simple_cod_indices(total) + + mask_mod = create_peagle_mask_mod( + anchor_pos, + depth, + lengths, + total, + sink_size=sink_size, + max_context_window=window, + ) + mask = _evaluate_mask(mask_mod, anchor_pos.shape[0]) + + # Query at position 22 (doc1, offset 12 within doc1) + q_idx = 22 + # Doc1 sinks: positions 10, 11, 12 (first 3 of doc1) + for kv_idx in [10, 11, 12]: + assert mask[q_idx, kv_idx], ( + f"Doc1 sink kv={kv_idx} should be visible for q={q_idx}" + ) + # Doc1 window: positions 18..22 + for kv_idx in range(18, 23): + assert mask[q_idx, kv_idx], ( + f"Doc1 window kv={kv_idx} should be visible for q={q_idx}" + ) + # Doc1 gap: positions 13..17 + for kv_idx in range(13, 18): + assert not mask[q_idx, kv_idx], ( + f"Doc1 gap kv={kv_idx} should be blocked for q={q_idx}" + ) + + +class TestWindowCoversAll: + """When window >= seq_len, equivalent to full causal attention.""" + + def test_large_window_matches_no_streaming(self): + seq_len = 12 + lengths = torch.tensor([seq_len]) + anchor_pos, depth = _simple_cod_indices(seq_len) + + mask_full = create_peagle_mask_mod(anchor_pos, depth, lengths, seq_len) + mask_streaming = create_peagle_mask_mod( + anchor_pos, + depth, + lengths, + seq_len, + sink_size=1, + max_context_window=seq_len, + ) + + total = anchor_pos.shape[0] + dense_full = _evaluate_mask(mask_full, total) + dense_streaming = _evaluate_mask(mask_streaming, total) + assert torch.equal(dense_full, dense_streaming) + + +class TestRolloutUnaffected: + """Within-rollout attention (depth > 0) is not affected by StreamingLLM.""" + + def test_rollout_attention_preserved(self): + seq_len = 10 + depths = {1: [2, 5, 8], 2: [2]} + lengths = torch.tensor([seq_len]) + anchor_pos, depth = _simple_cod_indices(seq_len, depths_to_sample=depths) + + mask_full = create_peagle_mask_mod(anchor_pos, depth, lengths, seq_len) + mask_streaming = create_peagle_mask_mod( + anchor_pos, + depth, + lengths, + seq_len, + sink_size=2, + max_context_window=3, + ) + + total = anchor_pos.shape[0] + dense_full = _evaluate_mask(mask_full, total) + dense_streaming = _evaluate_mask(mask_streaming, total) + + # For rollout queries (depth > 0), check that their attention to + # same-rollout KVs is identical + for q in range(total): + if depth[q] == 0: + continue + for kv in range(total): + if anchor_pos[q] == anchor_pos[kv] and depth[kv] > 0: + assert dense_full[q, kv] == dense_streaming[q, kv], ( + f"Rollout attention changed at q={q}, kv={kv}" + ) + + +class TestPaddingExcluded: + """Padding positions (document_id = -1) should never be attended to.""" + + def test_padding_blocked(self): + doc_len = 6 + total_seq_len = 10 # 4 padding positions + lengths = torch.tensor([doc_len]) + anchor_pos, depth = _simple_cod_indices(doc_len) + + mask_mod = create_peagle_mask_mod( + anchor_pos, + depth, + lengths, + total_seq_len, + sink_size=2, + max_context_window=3, + ) + _evaluate_mask(mask_mod, anchor_pos.shape[0]) + + # All entries should be within doc boundaries (no padding access) + # Since anchor_pos only covers positions 0..5, and all are in doc0, + # verify that no query can see positions >= doc_len (padding) + # (This is implicitly handled since anchor_pos doesn't include padding, + # but we verify the mask_mod logic itself) + zero = torch.zeros((), dtype=torch.long) + for q in range(doc_len): + # Query a position that would be in padding range if it existed + result = mask_mod(zero, zero, torch.tensor(q), torch.tensor(q)) + assert result # self-attention should work for valid positions diff --git a/tests/unit/models/test_peagle_data.py b/tests/unit/models/test_peagle_data.py new file mode 100644 index 000000000..8f206c932 --- /dev/null +++ b/tests/unit/models/test_peagle_data.py @@ -0,0 +1,95 @@ +"""Unit tests for P-EAGLE COD sampling logic.""" + +import torch + +from speculators.models.peagle.data import generate_cod_sample_indices + + +class TestGenerateCodSampleIndices: + def test_depth0_is_full_sequence(self): + seq_len = 64 + loss_mask = torch.ones(1, seq_len) + anchor_pos, depth = generate_cod_sample_indices( + seq_length=seq_len, loss_mask=loss_mask, num_depths=4 + ) + depth0_anchors = anchor_pos[depth == 0] + assert depth0_anchors.shape[0] == seq_len + assert torch.equal(depth0_anchors, torch.arange(seq_len)) + + def test_max_anchors_caps_chains(self): + seq_len = 128 + loss_mask = torch.ones(1, seq_len) + max_anchors = 16 + anchor_pos, depth = generate_cod_sample_indices( + seq_length=seq_len, + loss_mask=loss_mask, + num_depths=4, + max_anchors=max_anchors, + ) + # Depth 0 must still be the full sequence + assert anchor_pos[depth == 0].shape[0] == seq_len + + # Depth 1 chains should be capped at max_anchors + depth1_count = (depth == 1).sum().item() + assert depth1_count <= max_anchors + + def test_max_anchors_preserves_full_depth0(self): + seq_len = 256 + loss_mask = torch.ones(1, seq_len) + anchor_pos, depth = generate_cod_sample_indices( + seq_length=seq_len, + loss_mask=loss_mask, + num_depths=4, + max_anchors=8, + ) + depth0_anchors = anchor_pos[depth == 0] + assert depth0_anchors.shape[0] == seq_len + assert torch.equal(depth0_anchors, torch.arange(seq_len)) + + def test_max_anchors_none_uses_all(self): + seq_len = 64 + loss_mask = torch.ones(1, seq_len) + _, depth_limited = generate_cod_sample_indices( + seq_length=seq_len, + loss_mask=loss_mask, + num_depths=4, + max_anchors=None, + ) + _, depth_default = generate_cod_sample_indices( + seq_length=seq_len, + loss_mask=loss_mask, + num_depths=4, + ) + # Both should produce depth-0 with full sequence + assert (depth_limited == 0).sum() == (depth_default == 0).sum() == seq_len + + def test_max_anchors_fewer_valid_than_cap(self): + seq_len = 64 + loss_mask = torch.zeros(1, seq_len) + loss_mask[0, 10:20] = 1 # only 10 valid positions + anchor_pos, depth = generate_cod_sample_indices( + seq_length=seq_len, + loss_mask=loss_mask, + num_depths=4, + max_anchors=32, + ) + # Depth 0 still full sequence + assert anchor_pos[depth == 0].shape[0] == seq_len + # Chains should use all 10 valid positions (< max_anchors) + depth1_count = (depth == 1).sum().item() + assert depth1_count <= 10 + + def test_max_anchors_sorted_order(self): + seq_len = 128 + loss_mask = torch.ones(1, seq_len) + anchor_pos, depth = generate_cod_sample_indices( + seq_length=seq_len, + loss_mask=loss_mask, + num_depths=4, + max_anchors=16, + ) + for d in range(1, 4): + d_anchors = anchor_pos[depth == d] + if d_anchors.shape[0] > 1: + diffs = d_anchors[1:] - d_anchors[:-1] + assert (diffs >= 0).all(), f"Depth {d} anchors not sorted"