fix(eagle3,peagle): default d2t to the offset identity instead of arange

[khazic]

Problem

d2t holds offsets, not absolute ids: the target id of draft id i is i + d2t[i].

That is what this repo's own mapping generator emits:

# verl_speco/data/preprocessing.py
d2t = [used_tokens[i] - i for i in range(len(used_tokens))]

and it is what the serving engine applies:

# vLLM, model_executor/models/llama_eagle3.py
base = torch.arange(self.config.draft_vocab_size, device=logits.device)
targets = base + self.draft_id_to_target_id
logits_new[:, targets] = logits

vLLM even initializes its own copy of the buffer to torch.zeros, which is the identity under this convention.

Both draft models here default it to torch.arange instead:

d2t = torch.arange(self.draft_vocab_size, dtype=torch.int64)   # llama_eagle.py, modeling_peagle.py

Under the offset convention that resolves draft id i to target id 2 * i. For a full-vocabulary draft the top of the range lands outside the target vocabulary entirely, so the engine's logits_new[:, targets] = logits indexes out of bounds.

Scope

Nothing inside this repo reads d2t; the backend only validates its shape, and the loss uses t2d. The exposure is entirely on the export path:

• Hot publish is unaffected. _get_trainable_state_dict drops non-floating tensors, so d2t never goes over the wire and the engine keeps its own zero-initialized copy.
• _get_full_export_state_dict deliberately keeps "persistent buffers such as vocab mappings", so every checkpoint written to disk carries the value.

Neither load_vocab_mapping nor generate_vocab_mapping_file has a caller in-tree, so a freshly trained drafter always exports the constructor default rather than a frequency-derived mapping.

Fix

Default the buffer to torch.zeros in both draft models and state the convention in a comment next to it.

Checkpoints already exported with the arange buffer keep it on reload, since from_pretrained overwrites the default. They need their d2t rewritten (all zeros for a full-vocabulary draft).

Validation

Ran a before/after repro on both main and this branch. It shows what the repo's own generator emits, then applies vLLM's exact formula to the default buffer of both drafts.

Before/after repro output

=================== before: main (333b754) ===================
[what the repo's own mapping generator emits]
          used_tokens          = [3, 9, 17, 21]
          generated d2t        = [3, 8, 15, 18]
          arange(4) + d2t      = [3, 9, 17, 21]
          so d2t holds OFFSETS, and the identity mapping is all zeros

[full-vocabulary draft, no mapping supplied]
          EAGLE-3
            d2t[:8]              = [0, 1, 2, 3, 4, 5, 6, 7]
            engine target ids[:8]= [0, 2, 4, 6, 8, 10, 12, 14]
            max target id        = 62 (vocab_size=32)
            identity mapping     = False
            all ids in range     = False
          P-EAGLE
            d2t[:8]              = [0, 1, 2, 3, 4, 5, 6, 7]
            engine target ids[:8]= [0, 2, 4, 6, 8, 10, 12, 14]
            max target id        = 62 (vocab_size=32)
            identity mapping     = False
            all ids in range     = False

[verdict]
          RESULT: the default d2t resolves to the wrong target ids (bug present)

=================== after: this branch ===================
[what the repo's own mapping generator emits]
          used_tokens          = [3, 9, 17, 21]
          generated d2t        = [3, 8, 15, 18]
          arange(4) + d2t      = [3, 9, 17, 21]
          so d2t holds OFFSETS, and the identity mapping is all zeros

[full-vocabulary draft, no mapping supplied]
          EAGLE-3
            d2t[:8]              = [0, 0, 0, 0, 0, 0, 0, 0]
            engine target ids[:8]= [0, 1, 2, 3, 4, 5, 6, 7]
            max target id        = 31 (vocab_size=32)
            identity mapping     = True
            all ids in range     = True
          P-EAGLE
            d2t[:8]              = [0, 0, 0, 0, 0, 0, 0, 0]
            engine target ids[:8]= [0, 1, 2, 3, 4, 5, 6, 7]
            max target id        = 31 (vocab_size=32)
            identity mapping     = True
            all ids in range     = True

[verdict]
          RESULT: the default d2t is the offset identity (fixed)

max target id = 62 against vocab_size = 32 is the out-of-bounds write described above.

Tests

New file tests/integration/test_draft_vocab_mapping_contract.py, five cases:

• test_default_d2t_is_the_offset_identity for both drafts, asserting the engine formula recovers arange.
• test_default_d2t_stays_inside_the_target_vocabulary for both drafts with a reduced draft vocabulary.
• test_generated_mapping_round_trips_through_the_offset_convention pins the convention against process_token_dict_to_mappings, so the two producers cannot drift apart again.

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, 250 passed, 2 warnings in 21.26s

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 +5 on this branch are the new tests above.

[khazic]

The NPU vLLM eagle3 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: 35191, 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).