Ornith-1.5-35B-A3B-MTP

Ornith-1.5-35B-A3B-MTP

STANDALONE MERGED BF16 CHECKPOINT. This repository combines the complete BF16 ornith-ai/Ornith-1.5-35B-A3B target with our replacement BF16 MTP draft head. Point vLLM at this repository directly; no sidecar merge is required.

Related uploads:

• FP8 Ornith target + the same BF16 MTP head
• BF16 MTP head only

The target transformer, embeddings, LM head, tokenizer, chat template, and multimodal components come from the official BF16 Ornith-1.5 checkpoint. The replacement head was initialized from the compatible MoE MTP head in Qwen/Qwen3.6-35B-A3B, then re-aligned to Ornith-1.5's hidden states with full-vocabulary KL distillation. The final 12K run is a two-epoch, code-heavy warm-start from our earlier 5K KL-distilled head.

Why replace the shipped Ornith-1.5 MTP head?

The official Ornith-1.5 checkpoint includes a native 785-tensor MTP head, but it performed poorly in our matched vLLM tests. On ShareGPT-style MTP3 rows it accepted only 21.7% of proposed draft tokens and reached 0.71x/0.68x the AR throughput at concurrency 1/4.

This is consistent with the community report “mtp.* tensors look like random init, not trained weights”. Our own sample of the released BF16 tensors also found initializer-like statistics: q_proj std 0.01993 with kurtosis 2.997; expert-0 projections std 0.01996–0.01998 with kurtosis approximately 3; and mtp.norm.weight centered at 0.02281 rather than near 1. These observations do not prove provenance, but the weight statistics and serving behavior point in the same direction.

The low acceptance is not explained by target quantization:

The acceptance profile is effectively identical, so the native-head weakness is intrinsic to the head rather than an FP8 artifact.

Acceptance accuracy: native vs graft vs KL distillation

"Accuracy" below means speculative draft acceptance, not downstream task accuracy. vLLM's target verifier is still responsible for model correctness. The canonical gate used MTP3, temperature 0, concurrency 1, 10 prompts × 3 repetitions, and up to 384 generated tokens. The suite contains code, English-general, Japanese-general, and mixed Japanese/English prompts; the code subset has 12 requests.

The final head is strongest on code. Relative to the zero-training graft, code acceptance improves by 19.08 percentage points and overall suite acceptance by 12.17 points. The 12K run adds another 2.28 points of code acceptance over the 5K distill, mostly at the first two speculative positions.

Performance context on the same code gate: AR produced 220.47 client output tok/s at 1737.80 ms p50 E2E; the direct graft produced 247.10 tok/s at 1531.69 ms; this head produced 273.09 tok/s at 1299.62 ms. These are local single-GPU measurements, not universal throughput claims.

The code-heavy tradeoff is visible on matched ShareGPT rows: the 12K head's acceptance was 47.19% at c=1 and 48.08% at c=4, essentially flat versus the 5K head (48.71%/47.79%), while output throughput was 2.0%/1.1% lower. This head is therefore a code-priority variant, not an across-the-board replacement.

Training scale

The final checkpoint accumulates two stages: an initial 5K code distill from the Qwen3.6 graft, followed by the broader 12K mixed-data warm-start.

For the final stage, 11,964 teacher trajectories were generated from 12,000 prompts (36 duplicate prompt keys); 11,963 parsed into a length-stratified split of 10,763 train and 1,200 validation rows. Of the training split, 10,463 rows had completion room under the 2,048-token training cap, and 10,460 were used so the batch packed evenly.

Data mix

The designed 12K prompt mix is 50% code, plus agentic, multilingual general, and translation coverage:

The 2,400 synthetic coding prompts include 500 Python, 400 HIP, 300 Bash, 300 Go, 300 C++, 250 TypeScript, 250 Rust, and 100 JavaScript tasks. Translation includes 150 en→ja, 150 ja→en, 50 ko→en, 40 zh→en, 40 ja→ko, and smaller en→ko/zh/European-language buckets.

The private prompt and trajectory corpus is not included in this repository. It contains local code/repository-derived material. The table documents the mix, but it is not a claim that the training data is public, clean-room, or fully redistributable.

Corpus generation and KL process

1. Generate target trajectories. We sampled the official Ornith-1.5 FP8 target offline with reasoning enabled, temperature 0.7, top-p 0.95, and up to 2,048 output tokens. The 11,964 trajectories contain about 22.45M generated tokens (about 1,914 per row on average).
2. Length-stratified split. We split 10,763/1,200 train/validation across ten total-token buckets, then selected completion positions that fit the 2,048 training window.
3. Freeze the target. The BF16 ornith-ai/Ornith-1.5-35B-A3B target and shared LM head are frozen. Only the 844.6M MTP-head parameters are optimized.
4. Match the full target distribution. At completion positions, the MTP head consumes Ornith's post-norm hidden state plus the next-token embedding. Its hidden output and the target hidden state are projected through the same frozen LM head, and we minimize full-vocabulary KL divergence at T=1.0. This follows the key result from protoLabsAI/Ornith-1.0-9B-MTP: hard CE on sampled tokens can sharpen the wrong distribution, while speculative acceptance rewards matching the target distribution.
5. Warm-start and export. We warm-start from the 5K KL head, train two epochs at LR 5e-5 with a 5% warmup and cosine decay, then export 19 BF16 fused mtp.* tensors for vLLM. Target weights stay frozen throughout training; this merged release packages the official BF16 target.

Training efficiency: 4.25× faster

We adapted the original protoLabs technique first for shisa-ai/Ornith-1.0-35B-FP8-BLOCK-MTP, then optimized the 1.5/35B MoE training path:

• pack four 2,048-token rows per forward while preserving the same 8,192-token optimizer batch;
• project only completion positions into the 248K vocabulary;
• use an FLA/Triton fused full-vocabulary KL that computes gradients only for student hidden states, skipping the frozen 1 GB LM-head gradient;
• use BF16 for LM-head projection and FP32 for softmax/KL reduction;
• enable high matmul precision/TF32 and fused AdamW.

On one RTX PRO 6000 Blackwell 96 GB GPU, our own measured update time fell from 12.0 to 2.82 seconds/update (4.25×) at 84.24 GiB peak allocation. The full 5,230-update run finished in about four hours instead of the naive ~17.4-hour projection. This is a before/after comparison of our implementation, not a cross-hardware benchmark against protoLabs.

Why this merged checkpoint keeps the BF16 head

We also tested MTP-head-only FP8_BLOCK post-training quantization with the same Triton draft-MoE backend:

FP8 saved 0.78 GiB of loaded model memory and added 3.2% logical KV capacity, but regressed code throughput by 5.6% and ShareGPT c4 by 9.5% against the backend-matched BF16 control. Acceptance did not collapse, so the current cost appears to be the draft FP8 kernel path rather than head quality. This merged checkpoint therefore keeps the MTP head in BF16.

Serving

This is a complete Transformers-format checkpoint. The replacement-aware graft/distillation workflow is based on the recipe published with protoLabsAI/Ornith-1.0-9B-MTP, but users of this repository do not need to perform that merge themselves.

The locally tested vLLM 0.27.1 command shape is:

MODEL=./Ornith-1.5-35B-A3B-MTP

vllm serve "$MODEL" \
  --served-model-name ornith-1.5-35b-mtp \
  --trust-remote-code \
  --language-model-only \
  --max-model-len 262144 \
  --gpu-memory-utilization 0.93 \
  --max-num-seqs 32 \
  --max-num-batched-tokens 32768 \
  --max-cudagraph-capture-size 32 \
  --attention-backend flashinfer \
  --kv-cache-dtype fp8 \
  --calculate-kv-scales \
  --generation-config vllm \
  --no-enable-prefix-caching \
  --speculative-config '{"method":"mtp","num_speculative_tokens":3}'

Use gpu-memory-utilization=0.93 as the starting point on a 96 GB GPU. At 0.95, all tested MTP depths hit reproducible temporary GDN/FLA workspace OOMs on the first 32K-prefill row. The 0.93 matched recovery passed 80/80 rows. Re-benchmark memory utilization, depth, and concurrency on your own hardware.

Limitations and correctness status

• Experimental; not promoted as a production default. On vLLM 0.27.1, selected long-output and c16 greedy rows differed reproducibly between AR and speculative modes. One unchanged diagnostic prefix showed a material target decision shift, not just an exact tie. We also observed deterministic output sensitivity to max_tokens and to OpenAI logprob instrumentation.
• Exact AR/spec divergence was established only for vLLM. Our SGLang DFlash rows did not include a matched same-engine SGLang AR token-capture baseline, so cross-engine hash differences must not be presented as proof of the same SGLang correctness failure.
• Speculative decoding is intended to preserve the target distribution, so these differences are a runtime correctness investigation—not an accuracy/speed tradeoff that this model card silently accepts.
• Automatic prefix caching was disabled because the tested hybrid GDN + MTP path has known local correctness issues.
• The head was trained with a 2,048-token window. Long-context acceptance was measured, but the training distribution does not directly cover 256K input.
• The data mix is code-heavy and private/synthetic. ShareGPT acceptance did not improve over the 5K head.
• Evaluation was text-only (--language-model-only). Multimodal MTP behavior was not tested.
• Acceptance benchmarks are small local diagnostics, not downstream Ornith task benchmark scores. For base-model capabilities and risks, see the official Ornith-1.5 model card.

DFlash context-length cross-check

The community discussion also suggests DFlash as an alternative and notes that its acceptance falls at longer context. We reproduced the general long-context decline in both engines, with a clear DFlash decline by 64K:

Our retained output-length sweeps suggest DFlash declines more than MTP, with half-context being the apparent cross-over point where MTP outperforms DFlash:

This speed degradation is separate from output correctness: low draft acceptance should normally make speculation slower, not change the target result.

Provenance, prior work, and license

• Target model: ornith-ai/Ornith-1.5-35B-A3B
• MTP initialization donor: Qwen/Qwen3.6-35B-A3B
• Original graft/KL-distillation technique and recipe: protoLabsAI/Ornith-1.0-9B-MTP
• Our prior 35B adaptation: shisa-ai/Ornith-1.0-35B-FP8-BLOCK-MTP
• Native-head investigation: Ornith discussion #10

The replacement MTP head is a modified derivative of the Apache-2.0-licensed Qwen/Qwen3.6-35B-A3B MTP head: the Qwen donor weights were used as initialization and then changed by KL distillation against the MIT-licensed Ornith-1.5 target. Accordingly, this merged repository is distributed under the Apache License 2.0; see the local LICENSE and Qwen's upstream license.

The included Ornith target weights remain under Ornith's MIT license. This merged distribution contains both components and preserves the applicable Apache-2.0 attribution for the Qwen-derived head and the MIT notice for the Ornith target; see NOTICE.