EAGLE-3.1 spec-decode head for Chandra-2 (1.36x single-stream, lossless)

README.md

@@ -0,0 +1,58 @@
+---
+license: apache-2.0
+base_model: datalab-to/chandra-ocr-2
+tags:
+- speculative-decoding
+- eagle
+- eagle3
+- ocr
+- vllm
+library_name: vllm
+---
+
+# joecr-chandra-2-eagle3.1
+
+An **EAGLE‑3.1 speculative‑decoding draft head** for [`datalab-to/chandra-ocr-2`](https://huggingface.co/datalab-to/chandra-ocr-2) (a 5B Qwen3.5‑based vision‑language OCR model). Drop it into vLLM as the speculator to accelerate single‑stream (latency‑bound) OCR decoding **losslessly** — every drafted token is verified by the target, so output quality is preserved.
+
+## Results
+
+Measured in vLLM on olmOCR‑bench pages (single‑stream, greedy, CUDA graphs, RTX 3090 Ti @ 450 W):
+
+| | tok/s | speedup |
+|---|---|---|
+| baseline (no spec) | 84.1 | 1.0× |
+| **+ this head (num_spec=5)** | **114.3** | **1.36×** |
+
+Mean accepted length ≈ 2.0–2.5; per‑position acceptance ≈ [0.58, 0.37, 0.25, 0.20, 0.12]. Quality: median char‑similarity vs baseline 0.9991 (differences confined to bbox‑pixel jitter / equivalent figure captions — i.e. bf16 greedy non‑determinism, not degradation).
+
+## Architecture
+
+EAGLE‑3.1 (single decoder layer) over Chandra‑2's Qwen3.5 text backbone:
+- `hidden_size` 2560, `head_dim` 256, GQA 16/4 — matched to the target.
+- **EAGLE 3.1**: `fc_norm` (per‑aux‑layer RMSNorm) + `norm_output` (post‑norm recurrence) — fixes attention drift, holds acceptance at depth.
+- **Vocab pruning**: `draft_vocab_size` 32768 (from 248320; 99.99% token coverage) → ~7.6× smaller lm_head, cheaper draft forward. `d2t` buffer maps pruned ids back to the full vocab.
+- Aux hidden states captured from the target's **full‑attention** layers `[3, 15, 27]` (Chandra‑2 is a hybrid 3:1 linear/full‑attention model).
+
+## Usage (vLLM)
+
+```python
+from vllm import LLM
+llm = LLM(
+    model="datalab-to/chandra-ocr-2",
+    speculative_config={
+        "model": "jbarrow/joecr-chandra-2-eagle3.1",
+        "method": "eagle3",
+        "num_speculative_tokens": 5,
+    },
+    limit_mm_per_prompt={"image": 1},
+)
+```
+
+Use the Chandra `ocr_layout` prompt in the user turn (image + instruction). Spec‑decode helps most at **low concurrency**; at high batch the GPU saturates and a smaller `num_speculative_tokens` (1–2) is usually better.
+
+## Training
+
+- **On‑policy data**: Chandra‑2's own HTML‑with‑layout OCR over ~11k document pages.
+- **Features**: hidden states extracted offline from the frozen 5B target (decoupled from draft training so it fits 24 GB and trains with true DDP gradient sharing).
+- **Objective**: forward‑KL distillation with TTT (length 7), trained in vLLM's EAGLE input‑shift convention so it serves correctly.
+- 4× RTX 3090 Ti, DDP, ~8k steps.

config.json

@@ -0,0 +1,35 @@
+{
+  "architectures": [
+    "Eagle3LlamaForCausalLM"
+  ],
+  "model_type": "llama",
+  "num_hidden_layers": 1,
+  "hidden_size": 2560,
+  "intermediate_size": 9216,
+  "num_attention_heads": 16,
+  "num_key_value_heads": 4,
+  "head_dim": 256,
+  "hidden_act": "silu",
+  "rms_norm_eps": 1e-06,
+  "rope_theta": 10000000,
+  "rope_scaling": null,
+  "max_position_embeddings": 32768,
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "initializer_range": 0.02,
+  "bos_token_id": 248043,
+  "eos_token_id": 248044,
+  "tie_word_embeddings": false,
+  "torch_dtype": "bfloat16",
+  "use_cache": true,
+  "vocab_size": 248320,
+  "draft_vocab_size": 32768,
+  "fc_norm": true,
+  "norm_output": true,
+  "eagle_aux_hidden_state_layer_ids": [
+    3,
+    15,
+    27
+  ],
+  "target_hidden_size": 2560
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6beeccb1b3359cc516f491a1e556f528201ba5bad92ce220e3cb5193ec425883
+size 1704235808