fix(peagle): validate the packed document lengths instead of reusing them per row

tests/integration/test_peagle_backend_contract.py

@@ -294,6 +294,57 @@ def _model(**kwargs):
     assert float(out["accuracy"]) == pytest.approx(0.75)
 
 
+def _peagle_training_model_and_batch(batch_size: int, seq_len: int, seq_lengths):
+    import torch
+
+    from verl_speco.backends.peagle_trainer_backend import PEagleTrainingModel
+    from verl_speco.models.peagle import LlamaForCausalLMPeagle
+
+    config = _tiny_peagle_config()
+    model = PEagleTrainingModel(
+        LlamaForCausalLMPeagle(config), num_depths=2, down_sample_ratio=0.5
+    )
+    kwargs = dict(
+        input_ids=torch.zeros(batch_size, seq_len, dtype=torch.long),
+        aux_hidden=torch.zeros(batch_size, seq_len, config.target_hidden_size * 3),
+        loss_mask=torch.ones(batch_size, seq_len),
+        attention_mask=torch.ones(batch_size, seq_len, dtype=torch.long),
+        target_logits=torch.zeros(batch_size, seq_len, config.vocab_size),
+        seq_lengths=seq_lengths,
+    )
+    return model, kwargs
+
+
+def test_peagle_rejects_packed_lengths_for_a_multi_row_batch() -> None:
+    """seq_lengths describes one packed sequence, so it cannot span rows.
+
+    The forward loops over the batch but applies the whole seq_lengths tensor to
+    every row, which would silently give row 1 row 0's document layout.
+    """
+    torch = pytest.importorskip("torch")
+    pytest.importorskip("transformers")
+
+    model, kwargs = _peagle_training_model_and_batch(
+        batch_size=2, seq_len=8, seq_lengths=torch.tensor([4, 4])
+    )
+
+    with pytest.raises(ValueError, match="single packed sequence"):
+        model(**kwargs)
+
+
+def test_peagle_rejects_packed_lengths_that_do_not_cover_the_sequence() -> None:
+    """A short seq_lengths silently turns the tail into unattendable padding."""
+    torch = pytest.importorskip("torch")
+    pytest.importorskip("transformers")
+
+    model, kwargs = _peagle_training_model_and_batch(
+        batch_size=1, seq_len=8, seq_lengths=torch.tensor([3, 2])
+    )
+
+    with pytest.raises(ValueError, match="sum to 5"):
+        model(**kwargs)
+
+
 def test_peagle_checkpoint_export_unwraps_the_training_model() -> None:
     pytest.importorskip("torch")
     from types import SimpleNamespace

verl_speco/backends/peagle_trainer_backend.py

@@ -122,6 +122,25 @@ def forward(
 
         batch_size, seq_len = input_ids.shape
         device = input_ids.device
+        if seq_lengths is not None:
+            # seq_lengths is a flat list of document lengths for ONE packed
+            # sequence, which is what base_trainer builds for P-EAGLE. There is no
+            # per-row structure to index, so a multi-row batch would silently
+            # apply one row's document layout to all of them.
+            if batch_size > 1:
+                raise ValueError(
+                    "P-EAGLE seq_lengths describe a single packed sequence, but the batch "
+                    f"has {batch_size} rows; pack the documents into one row or drop seq_lengths"
+                )
+            # Any tail past sum(seq_lengths) gets document id -1 in the COD mask,
+            # which makes those queries attend to nothing at all rather than
+            # failing, so check the invariant instead of drafting on garbage.
+            total_length = int(seq_lengths.sum())
+            if total_length != seq_len:
+                raise ValueError(
+                    f"P-EAGLE seq_lengths sum to {total_length} but the packed sequence is "
+                    f"{seq_len} tokens long"
+                )
         loss_num = torch.zeros((), device=device, dtype=torch.float32)
         loss_den = torch.zeros((), device=device, dtype=torch.float32)
         correct = torch.zeros((), device=device, dtype=torch.float32)
@@ -137,9 +156,11 @@ def forward(
             )
             orig_positions = anchor_pos + depth
             if seq_lengths is not None:
-                row_length = seq_lengths.to(device)
+                document_lengths = seq_lengths.to(device)
             else:
-                row_length = attention_mask[b].sum().clamp_min(1).reshape(1).to(device)
+                document_lengths = (
+                    attention_mask[b].sum().clamp_min(1).reshape(1).to(device)
+                )
             loss_positions = row_loss_mask[0, orig_positions].bool()
 
             is_depth0 = depth == 0
@@ -161,7 +182,7 @@ def forward(
             block_mask = draft.build_peagle_block_mask(
                 anchor_pos=anchor_pos,
                 depth=depth,
-                lengths=row_length,
+                lengths=document_lengths,
                 total_seq_len=seq_len,
             )
             hidden = draft.forward_peagle(