feat(spec): SpecInput data contract + host/device boundary annotations

[zorrofox]

Part of #1053 (P1-5a). Draft for early interface review — P1-2 (BaseSpecWorker/BaseDraftWorker class refactor) builds on this contract, so opening early per discussion.

Note

Depends on #1063 (P1-0). This branch is stacked on it; the actual P1-5a diff is the last commit only (4 files, +201/−26). Will rebase onto main once #1063 merges.

What

Defines the spec-decode data contract from the #1053 RFC tables so the P1-2 abstraction layer has a fixed interface to target:

• spec_info.SpecInput — @runtime_checkable Protocol exposing:
  • is_draft_input() / is_verify_input()
  • get_logical_token_num() / get_allocated_token_num() / get_verify_token_num() — the three token counts the RFC requires to be separated (scheduler advance / KV trimming / verify-metadata+DP-accounting respectively)
  • get_spec_adjust_token_coefficient() / filter_batch() / merge_batch()
  • Docstring fixes the host/device boundary rules and DP-padded layout (Route 1: target+draft both DP) semantics
• EagleDraftInput — implements SpecInput; per-field docstrings annotate device vs host, shape, JIT-cache-key participation, and the allocate_lens ≠ accept_length ≠ new_seq_lens distinction. Multi-layer MTP cross-round hidden_states semantics noted.
• EagleVerifyInput — implements SpecInput; per-field docstrings annotate device vs host and static-metadata fields; filter_batch/merge_batch raise (verify input is single-round).
• test_spec_info.py — Protocol conformance + three-token-count distinction.

Non-changes

• No behavior change to the existing dp=1 path (pure additions: docstrings, Protocol, new methods).
• GenerationBatchResult already has all RFC-required fields except optional new_seq_lens (derivable from seq_lens + accept_lens); not adding it here.
• Type annotations on device fields changed np.ndarray → jax.Array | None to reflect actual runtime types after fix(eagle): restore dp=1 dense EAGLE3 path on main + add E2E to CI #1063's reshard changes — annotations only, no code path change.

Open questions for P1-2 author

1. Should SpecInput be a Protocol (current, duck-typed) or an ABC base class? Protocol avoids inheritance coupling but ABC lets BaseSpecWorker type-hint more strictly.
2. get_logical_token_num currently returns a scalar int (sum across batch). Should it return per-request np.ndarray instead for DP per-rank accounting?
3. EagleDraftInput.merge_batch currently uses np.concatenate (host-side). RFC says these are device arrays — should merge happen on device (jnp.concatenate)?

Test plan

• test_spec_info.py 3/3 passed (CPU)
• pre-commit clean
• dp=1 EAGLE E2E unchanged (will verify after rebasing onto merged fix(eagle): restore dp=1 dense EAGLE3 path on main + add E2E to CI #1063)

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[pengchengneo]

LGTM