[Prototype] Add vLLM Apriel2 model with plugin-based registration

[tscholak]

Summary

Adds vLLM-optimized Apriel2 model implementation to fast_llm_external_models

• Uses vLLM's ModelRegistry.register_model() for runtime registration (no vLLM patching required)
• Supports hybrid architectures: attention, mamba, GDN, and KDA mixers

Attribution

Model implementation based on work by @nandahkrishna from the apriel2-vllm branch. This PR adapts that implementation for plugin-based registration as an alternative to patching vLLM directly.

Goal

Evaluate whether vLLM's plugin/registration mechanism can work for us as a short-term solution, avoiding the need to maintain a patched vLLM fork.

Usage

from fast_llm_external_models.apriel2.vllm import register
from vllm import LLM

register()
llm = LLM(model="path/to/apriel2/checkpoint")

---

vLLM vs Transformers Alignment Verification

Statistical comparison using test_apriel2.py stats command with:

• 64 prompts from C4 dataset (deterministic sampling with seed=42)
• 128 tokens prompt length
• 16 decode steps
• Identical token IDs sent to both backends (controlled tokenization)
• Per-position logprob comparison with percentile statistics

Models Tested

Model	Description
pure-gdn	100% GDN layers
attn-swa	100% attention (sliding window)
every5th-kda	80% attention + 20% KDA

---

Results Summary

Model	Mode	Match%	Mean Diff	p95 Diff	Max Diff	Outliers
GDN	no-compile	84.4%	1.05	7.83	22.80	142 (14.0%)
GDN	compiled	83.4%	1.07	7.87	13.99	155 (15.4%)
SWA	no-compile	87.5%	0.83	7.55	18.59	111 (10.8%)
SWA	compiled	84.7%	1.14	8.85	23.38	140 (13.7%)
KDA	no-compile	87.1%	0.83	7.30	15.70	120 (11.7%)
KDA	compiled	84.8%	1.03	8.39	15.15	141 (13.8%)

Per-Position Token Match Rate (no-compile mode)

Position	GDN	SWA	KDA
prefill	95.3%	96.9%	98.4%
decode1	96.9%	93.8%	95.3%
decode2	92.2%	93.8%	89.1%
decode3	92.2%	92.2%	90.6%
decode4	90.6%	92.2%	90.6%
decode5	90.6%	90.6%	87.5%
decode6	85.7%	93.8%	89.1%
decode7	85.7%	89.1%	89.1%
decode8	84.1%	87.5%	89.1%
decode9	79.4%	82.8%	85.9%
decode10	81.0%	85.9%	84.4%
decode11	79.4%	84.4%	82.8%
decode12	73.0%	81.2%	82.8%
decode13	74.6%	79.7%	82.8%
decode14	74.2%	79.7%	78.1%
decode15	73.8%	76.6%	78.1%

---

Key Findings

1. Divergence is NOT mixer-specific

All models (GDN, SWA, KDA) show similar divergence patterns between vLLM and Transformers. This indicates the issue is in shared model code (RMSNorm, MLP, embeddings) rather than mixer implementations.

2. torch.compile has minimal impact

Compile vs no-compile produces nearly identical results:

• GDN: 84.4% vs 83.4% match
• SWA: 87.5% vs 84.7% match
• KDA: 87.1% vs 84.8% match

Previous reports of GDN torch.compile issues appear to have been measurement artifacts.

3. Divergence accumulates over decode steps

• Prefill: 95-98% token match rate
• Decode15: 73-78% token match rate

Small numerical differences compound during autoregressive generation, causing progressive divergence.

4. Prefill is well-aligned

All models show excellent prefill alignment (95-98% match, avg diff ~0.04), making them reliable for likelihood-based evaluation (MMLU, etc.).

---

Implications

For likelihood-based evaluation (MMLU)

✅ All models reliable - prefill-only evaluation shows 95-98% alignment

For generative evaluation (GSM8K)

⚠️ All models show accumulating divergence - vLLM and Transformers will produce different outputs over long generations, regardless of mixer type or compilation mode

Root Cause Investigation Needed

The divergence affects all model types equally, suggesting the issue is in:

• RMSNorm implementation differences
• MLP/SwiGLU numerical precision
• Embedding layer handling
• KV cache management differences

---

Test Configuration

# Run statistical comparison
python test_apriel2.py stats /path/to/model \
    --num-prompts 64 \
    --prompt-length 128 \
    --decode-length 16 \
    --batch-size 1 \
    --dtype bfloat16 \
    --tf-kernels vllm \
    [--no-compile]  # Add for no-compile mode

Test plan

• Test registration mechanism with vLLM
• Verify model loads correctly
• Statistical comparison of vLLM vs Transformers (GDN, SWA, KDA)
• Tested compile vs no-compile modes
• Per-position analysis of divergence patterns
• Investigate shared code divergence (RMSNorm, MLP, embeddings)

🤖 Generated with Claude Code