feat(eagle3): Qwen3 EAGLE-3 drafter forward pass (single-step + chain rollout)#707
Open
scatyf3 wants to merge 22 commits into
Open
feat(eagle3): Qwen3 EAGLE-3 drafter forward pass (single-step + chain rollout)#707scatyf3 wants to merge 22 commits into
scatyf3 wants to merge 22 commits into
Conversation
Add the single-sequence custom-mask attention kernels the EAGLE-3 drafter forward needs (prefill_attention.cu / paged_attention.cu) with their FFI declarations and safe openinfer-kernels / openinfer-core op wrappers. No caller yet; the drafter loading and forward land in follow-up PRs. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…isibility Address review on the new attention entries (openinfer-project#647): - Wrap single_prefill_nhd_causal_cuda / single_decode_nhd_cuda in OPENINFER_FFI_GUARD_BEGIN/END(-1) like their siblings — FlashInfer's DISPATCH_GQA_GROUP_SIZE throws for group sizes outside {1,2,3,4,8}, and without the guard the C++ exception crossed the extern "C" boundary and aborted the process instead of returning -1. The two Rust wrappers now append ffi_exception_message(result) to their bail messages. - Reject num_qo_heads % num_kv_heads != 0 in both entries: FlashInfer truncates group_size = qo/kv, so a non-divisible pair silently served fewer heads and returned cudaSuccess, leaving stale/uninitialized output. - Assert sin_cache.data.len() == cos_cache.data.len() in eagle3_rope_into; cos_max_pos was derived from cos_cache alone, so a shorter sin_cache let the kernel read out of bounds. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…AGLE-3 wrappers Second review round on openinfer-project#647 — soundness of the safe pub fn entry points: - Add checked_row_offset (checked_add for offset+span, checked_mul for the byte offset) to the two row-offset wrappers eagle3_rope_into and single_prefill_nhd_causal_into. Release builds leave overflow-checks off, so a usize::MAX offset previously wrapped past the range assert and the pointer arithmetic, faulting only at the next sync as CUDA_ERROR_ILLEGAL_ADDRESS; it now returns Err before entering FFI. - Add ensure_hidden_capacity (require data.len() >= hidden_dim * seq_len — >=, not ==, to keep the capacity-backed HiddenStates convention where buffers are allocated at a max and .seq_len is rewritten per step) across the Q/K/V/output buffers of all three new wrappers. A safe caller inflating .seq_len past a small allocation now returns Err instead of silently reading out of bounds. Pre-existing siblings (dflash_qk_norm_rope_into, single_prefill_nhd_noncausal_into) share both gaps; left as follow-up. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Add the EAGLE-3 drafter model skeleton: Eagle3Config (parse + validate against the target's geometry/vocab/rope), the Eagle3DraftModel / Eagle3Layer weight structs, safetensors loading (fc fusion, midlayer, draft lm_head, d2t/t2d vocab maps, rope caches), and the drafter KV memory reservation. The forward pass is not present yet (module is dead-code-allowed); it lands in the next PR. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Two lints broke the `openinfer-server` CUDA clippy job (`-D warnings`): - `chunks_exact_to_as_chunks` in `load_tensor_i64_host`: use `as_chunks::<8>()` over the byte slice, which also drops the `try_into().unwrap()` since each chunk is already `[u8; 8]`. - `unused_imports` on the `Eagle3MemoryReservation` re-export: the module's `#![allow(dead_code)]` doesn't cover unused imports, and the reservation is only wired into the KV budget by the forward-pass PR. Mark the re-export `#[allow(unused_imports)]` with a note. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
… tautological test aux_hidden_state_layers returned the reference's pre-layer gate (2, N/2, N-3) but fed it into openinfer's post-layer / 0-based capture path (prefill.rs, verify_graph.rs) — one block late for all three fc inputs. SafeAILab traineagle3 appends the hidden state *before* the block (vLLM matches it via a post-layer idx+1 capture), so the equivalent post-layer indices are (1, N/2-1, N-4) = [1, 17, 32] for the 36-layer Qwen3-4B target. Also drop aux_layers_qwen3_4b: it was a constant-true echo of the formula. Unlike DFlash (whose target_layer_ids are parsed from the checkpoint and cross-checked), EAGLE-3's AngelSlim config carries no capture-layer field, so there is no external value to validate against. End-to-end correctness is covered by the drafter golden gate. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
`-p openinfer-qwen3 --all-targets` compiles the lib-test target, which the upstream `-p openinfer-server` clippy job never reaches — so two lints slipped through the earlier drafter-loader clippy pass: - items_after_test_module (eagle3.rs): the `#[cfg(test)] mod tests` sat above Eagle3Layer/Eagle3DraftModel. Move it to the bottom of the file. - unreadable_literal (reservation.rs): separators on the 5+ digit fields of the Qwen3-4B eagle3 test config (151_936 / 32_000 / 40_960). Test code only; no behavior change. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
`load_tensor_i64_host` / `load_tensor_bool_host` only checked byte-length divisibility, so a mistyped tensor (e.g. F64, also 8-byte) was silently reinterpreted as I64 and a U8 tensor as BOOL; a >1D shape passed too. Enforce `Dtype::I64` / `Dtype::BOOL` and rank-1 before the raw-byte read. Addresses PR openinfer-project#662 review (FeathBow). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
… exactly Two review gaps in the drafter loader: - Only `fc` was shape-checked; q/k/v, o, gate/up/down, all norms and lm_head were loaded blind, and `q_dim`/`kv_dim` were derived from the loaded Q/K tensors rather than the config. A malformed checkpoint (e.g. a `v_proj` with the wrong row count) could build an internally inconsistent draft, since `vstack` only requires equal columns. Derive the dims from config and check every tensor's shape against it before stacking or storing. - The two BF16 RoPE caches `[max_position_embeddings, head_dim]` were only covered by the generic 10% weight slack. `validate_for_target` only lower-bounds the draft limit to the target's, so a 131072-position draft overruns the slack by ~23 MB. Bill the exact term in `fixed_bytes` and pin the delta in the test. Addresses PR openinfer-project#662 review (FeathBow). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Implement the EAGLE-3 drafter forward on the loaded model: per-request state/scratch, the single `midlayer` step (fc-fuse captured target hidden states, dual-input attention via the custom-mask kernels, MLP, draft lm_head, d2t remap), batched teacher-forcing prefill, and the top-1 chain rollout (draft_chain / chain_round / reseed_after_verify). Not yet wired into the executor; correctness gate lands in the next PR. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
- Rename the single-char locals in `prefill_batched` (n/q/k/v/o ->
num_tokens/query/key/value/attn_proj) to clear `many_single_char_names`.
- `#[allow(unused_imports)]` on the `forward::{Eagle3RequestState,
Eagle3Scratch}` re-export (consumed by the later scheduler PR), mirroring
the reservation export.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Align the boundary target-feature field with vLLM's EAGLE-3 naming: it holds the low/mid/high aux hidden state (pre-`fc`) at the last committed position, which vLLM calls `aux_hidden_states`. No behaviour change. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
The forward pass writes K/V and runs the midlayer through *unchecked* gemms (`gemm_into`/`gemm_rows_into` don't validate dims), trusting that the passed `Eagle3RequestState`/`Eagle3Scratch` were allocated for this drafter. A state or scratch from a differently-shaped drafter would silently corrupt the KV cache / produce garbage rather than error — the same "internally inconsistent" class flagged on the loader (openinfer-project#662). - `ensure_state_geometry` / `ensure_scratch_geometry`: one-shot dim asserts at the leaf consumers (`draft_step`, `prefill_batched`, `seed_hidden_from_context`). - `draft_chain`: fail up front when `start + k` overruns `max_cache_len`, before writing any speculative KV, instead of mid-chain in `draft_step`. - `debug_assert` the `aux_hidden_states` boundary dim where it feeds `fc`. No behaviour change on valid inputs; the host-side integer checks are dwarfed by the per-step device sync already in `draft_chain`. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Add `// 1.`..`// 14.` step comments mirroring `draft_step`, so the two forward paths read side by side. prefill_batched has two extra numbered stages the single-token path lacks: the fc teacher-forcing residual (2) and the boundary last-hidden seed (14). Comment-only. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Restructure the dense doc block into rule -> concrete mapping (bulleted) -> params -> v1 constraint. Same content, readable. Comment-only. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
The field holds the *last* committed position's aux hidden state (`[3*hidden,1]`, pre-fc) — the name now says both "last/boundary" and "aux (3-layer)". Drop the "seed" wording for it: it was conflated with the post-fc *chain seed* (`fc(..)` output that seeds the draft residual), which keeps its name. The local that captures it is now `boundary`/`new_boundary`. Comment-and-name only. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
The chain seed passed by `chain_round` is `fc(last_aux_hidden_states)` — the fused target boundary hidden `[hidden,1]`, not "the prefill's last decoder output" the doc claimed. Rename the param and correct the description. Also note on `prefill_batched` that its `(logits, last_hidden)` return is not consumed yet (reserved for the scheduler PR). No behaviour change. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…hidden naming The EAGLE papers and the reference/vLLM/SGLang implementations describe this step as the draft's "feature" / "input hidden", not a "seed". Align the eagle3 draft vocabulary with the field: - fn seed_hidden_from_context -> fuse_input_hidden_from_context - fn reseed_after_verify -> rebuild_after_verify - var seed -> fused_target_hidden (matches draft_chain's param) - comments/errors: chain seed / re-seed -> input hidden / rebuild Docs-and-naming only; no behavior change. Both renamed fns have no callers (WIP scaffolding), so the rename is contained to this file. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
clippy::doc_comment_double_space_linebreaks fired on the draft_chain doc's trailing two-space line break. Remove the trailing whitespace. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Resolve add/add conflict in openinfer-qwen3/src/eagle3.rs: keep the forward-pass module (mod forward + Eagle3RequestState/Eagle3Scratch exports) this branch adds on top of main's openinfer-project#662 loading/reservation skeleton. Comment wording ("scheduler PR" for the reservation export) kept from this branch since the forward pass now lands here. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Collaborator
|
Could you account for the extra committed token in the capacity check? |
Contributor
Author
|
Thank you for catching the bug, I already fix that |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This is an follow up PR for #662
What
Adds the EAGLE-3 drafter forward path for the Qwen3 line, on top of the drafter config/loading skeleton. The core change is in
openinfer-qwen3/src/eagle3/forward.rsCode
draft_step: one drafter decoding stepprefill_prompt/prefill_batched: teacher-forced prefill that builds thedrafter KV and the boundary target feature from captured hidden states.
draft_chain: autoregressive chain of γ draft tokens from thefused target boundary hidden; v1 syncs logits to host per step for the argmax.
chain_round: one speculative round: init the residual stream fromfc(boundary feature), draft γ tokens, then rewind the speculative KV so theround is side-effect-free except for the returned span.
rebuild_after_verify: after a verify step, teacher-force the acceptedprefix back into the drafter KV and set the new boundary target feature for the next round.
Testing
cargo test --release -p openinfer-qwen3 --lib— reservation geometry unittest passes;
cargo fmt/clippy clean.scheduler PR that consumes these entry points and perform accuracy check vs hf golden gate.