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

[khazic]

Problem

seq_lengths is a flat list of document lengths for one packed sequence. base_trainer builds it that way:

batch["seq_lengths"] = torch.tensor(
    [chunk.size(0) for chunk in input_id_chunks], dtype=torch.long, device=dev
)

The P-EAGLE forward loops over the batch but hands the whole tensor to every row:

for b in range(batch_size):
    ...
    if seq_lengths is not None:
        row_length = seq_lengths.to(device)          # not indexed by b
    else:
        row_length = attention_mask[b].sum()...      # this branch does index by b

Two gaps:

1. A multi-row batch would silently give every row row zero's document layout. base_trainer always produces a single flat row today, so this is latent, but the loop is written for batch_size > 1 and the fallback branch does index per row.
2. Nothing checks that the lengths cover the sequence. The COD mask builds document_ids from seq_lengths and pads the remainder with -1:

document_ids = torch.cat([document_ids, torch.full((total_seq_len - document_ids.shape[0],), -1, ...)])
...
is_not_padding = document_ids[q_anchor_pos] != -1

so any tail past sum(seq_lengths) becomes a query that attends to nothing at all. That is not an error, it is a silently degenerate attention row that still flows into the loss.

Fix

Reject both, and rename the local to document_lengths, which is what the mask consumes.

The length check costs one comparison per forward on a path that already synchronizes several times (generate_cod_sample_indices and create_block_mask both do).

Validation

Ran a before/after repro on both main and this branch, feeding the training model each malformed shape.

Before/after repro output

=================== before: main (333b754) ===================
[batch of 2 rows, seq_lengths=[4, 4] for an 8-token sequence]
          ACCEPTED (drafted on a mismatched document layout)

[1 row of 8 tokens, seq_lengths=[3, 2] covers only 5]
          ACCEPTED (drafted on a mismatched document layout)

[verdict]
          RESULT: 2 of 2 accepted silently (bug present)

=================== after: this branch ===================
[batch of 2 rows, seq_lengths=[4, 4] for an 8-token sequence]
          REJECTED: P-EAGLE seq_lengths describe a single packed sequence, but the batch has 2 rows; pack the documents into one row or drop seq_lengths

[1 row of 8 tokens, seq_lengths=[3, 2] covers only 5]
          REJECTED: P-EAGLE seq_lengths sum to 5 but the packed sequence is 8 tokens long

[verdict]
          RESULT: malformed packed lengths are rejected (fixed)

Tests

Two new tests in tests/integration/test_peagle_backend_contract.py:

• test_peagle_rejects_packed_lengths_for_a_multi_row_batch
• test_peagle_rejects_packed_lengths_that_do_not_cover_the_sequence

Full CPU suite (tests/integration tests/compat tests/config tests/examples) run on both sides:

Test suite before/after

### BASELINE (origin/main, 333b754) ###
FAILED tests/integration/test_drafter_runtime_control_contract.py::test_target_head_sync_defers_for_all_lm_head_drafters[DSPARK-veomni-npu-veomni_lm_head_full]
FAILED tests/integration/test_drafter_runtime_control_contract.py::test_target_head_sync_defers_for_all_lm_head_drafters[DFLASH-veomni-npu-veomni_lm_head_sparse]
FAILED tests/integration/test_drafter_runtime_control_contract.py::test_target_head_sync_defers_for_all_lm_head_drafters[EAGLE3-veomni-cuda-veomni_lm_head_full]
FAILED tests/integration/test_drafter_runtime_control_contract.py::test_target_head_sync_defers_for_all_lm_head_drafters[EAGLE1-fsdp-npu-engine_full_param]
FAILED tests/integration/test_drafter_runtime_control_contract.py::test_target_head_sync_defers_for_all_lm_head_drafters[DOMINO-fsdp2-cuda-engine_full_param]
FAILED tests/integration/test_drafter_runtime_control_contract.py::test_target_head_transfer_waits_after_actor_update
FAILED tests/integration/test_drafter_runtime_control_contract.py::test_async_publish_sets_pending_ref_and_waits_before_next_publish
FAILED tests/integration/test_dspark_trainer_backend.py::test_dspark_checkpoint_preserves_source_config_and_vllm_weight_names
FAILED tests/integration/test_verl_npu_vllm_compat.py::test_factory_fused_moe_survives_verl_npu_patch_import
9 failed, 245 passed, 2 warnings in 33.60s

### THIS BRANCH ###
(same 9 failures)
9 failed, 247 passed, 2 warnings in 23.06s

The same 9 tests fail on main and on this branch. They need optional dependencies (VeOmni, the NPU vLLM stack) that are absent in this environment, so they are pre-existing and unrelated. The +2 on this branch are the new tests above.

[khazic]

The NPU vLLM dflash example failure here is runner port contention, not this change. The job died during engine init with:

torch.distributed.DistNetworkError: The server socket has failed to listen on any local network address. port: 39403, useIpv6: false, code: -98

code: -98 is EADDRINUSE. The same failure hit four PRs in this batch, on four different ports, and the job that failed does not line up with what each PR touches:

PR	files changed	failing job	port
#66	peagle_trainer_backend.py, one error-message string in eagle3_trainer_backend.py	eagle3	36693
#67	llama_eagle.py, modeling_peagle.py	eagle3	35191
#68	eagle3_trainer_backend.py	dspark	38527
#70	peagle_trainer_backend.py	dflash	39403

#68 and #70 are the clearest: neither touches the dspark or the dflash code path, yet those are the jobs that failed.

Within a single run the three example jobs are serialized (on #63 they ran 11:25:05 to 11:42:21, 11:42:46 to 12:01:01, 12:01:27 to 12:19:55), so the contention comes from runs of different PRs overlapping. The failures cluster in the window where three PRs had example jobs in flight at once, and every job that started after that queue drained passed. The failed jobs also died in about 12 minutes against roughly 17 for a successful one, consistent with dying at engine init rather than during real work.

CPU unit tests and pre-commit pass on this PR.

Could a maintainer re-run the failed job? I do not have the permission to (gh run rerun returns Must have admin rights to Repository).