[EXTERNAL] 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

[Copilot]

Pull request overview

This PR adds vLLM-optimized Apriel2 model implementation using vLLM's plugin-based registration mechanism, avoiding the need to maintain a patched vLLM fork. The implementation supports hybrid architectures with attention, mamba, GDN, and KDA mixers, with runtime mixer switching capability for stochastic layers.

Changes:

• Adds vLLM model implementation with plugin registration via entry points
• Implements comprehensive test suite comparing vLLM vs Transformers outputs
• Adds debug flags for kernel comparison between vLLM and FLA implementations

Reviewed changes

Copilot reviewed 9 out of 10 changed files in this pull request and generated 34 comments.

Show a summary per file

File	Description
setup.cfg	Adds vLLM plugin entry point for automatic model registration
tests/utils/model_configs.py	Updates test config with AutoModelForImageTextToText class
fast_llm_external_models/apriel2/vllm/init.py	Module initialization exporting Apriel2ForCausalLM
fast_llm_external_models/apriel2/vllm/README.md	Documentation for vLLM usage
fast_llm_external_models/apriel2/vllm/config_convertor.py	Config converter for nested Apriel2 decoder format
fast_llm_external_models/apriel2/vllm/modeling_apriel2.py	Main vLLM model implementation (~2863 lines)
fast_llm_external_models/apriel2/vllm/test_apriel2.py	Manual test script for vLLM vs Transformers comparison
fast_llm_external_models/apriel2/modeling_apriel2.py	Transformers model with debug flags for kernel comparison
fast_llm_external_models/apriel2/examples/pure_gdn_step*.yaml	Example configs for GDN model conversion

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Lines 580-626 implement vLLM convolution path that is never executed because USE_VLLM_CONV is hardcoded to False. This is 47 lines of dead code that adds unnecessary complexity and maintenance burden. Either remove this code or make the flag configurable if this functionality is actually needed.

[Copilot]

Instance-level debug flags (_debug_enabled, _debug_layer, _debug_state, _debug_output) are defined within the class but hardcoded to False. These add dead debug code throughout the forward method. The debug methods (_debug_tensor, _debug_print, _debug_state_stats) check these flags at runtime, creating unnecessary overhead. Consider removing the debug infrastructure or making it opt-in through a proper logging framework.

Suggested change

	if not self._debug_enabled:
	return
	if t is None:
	print(f"[TF-GDN layer={self.layer_idx}] {name}: None")
	return
	try:
	flat = t.flatten()[:8]
	vals = ", ".join(f"{v:.6f}" for v in flat.float().tolist())
	print(f"[TF-GDN layer={self.layer_idx}] {name}: shape={t.shape}, dtype={t.dtype}, "
	f"mean={t.float().mean().item():.6f}, std={t.float().std().item():.6f}, "
	f"first8=[{vals}]")
	except Exception as e:
	print(f"[TF-GDN layer={self.layer_idx}] {name}: ERROR accessing tensor: {e}")
	def _debug_print(self, msg: str):
	if not self._debug_enabled:
	return
	print(f"[TF-GDN layer={self.layer_idx}] {msg}")
	def _debug_state_stats(self, name: str, state: torch.Tensor, seq_len: int):
	"""Debug recurrent state with statistics."""
	if not self._debug_state or state is None:
	return
	try:
	flat = state.flatten()
	first8 = ", ".join(f"{v:.6f}" for v in flat[:8].float().tolist())
	print(f"[TF-GDN L{self.layer_idx}] {name} (seq_len={seq_len}): shape={state.shape}, "
	f"mean={state.float().mean().item():.6f}, std={state.float().std().item():.6f}, "
	f"min={state.float().min().item():.6f}, max={state.float().max().item():.6f}, "
	f"first8=[{first8}]")
	except Exception as e:
	print(f"[TF-GDN L{self.layer_idx}] {name}: ERROR accessing state: {e}")
	"""
	No-op debug hook for tensor inspection.
	This method is intentionally empty to avoid runtime overhead from
	disabled debug printing while preserving the public interface.
	"""
	return
	def _debug_print(self, msg: str):
	"""
	No-op debug hook for printing diagnostic messages.
	This method is intentionally empty to avoid runtime overhead from
	disabled debug printing while preserving the public interface.
	"""
	return
	def _debug_state_stats(self, name: str, state: torch.Tensor, seq_len: int):
	"""
	No-op debug hook for recurrent state statistics.
	This method is intentionally empty to avoid runtime overhead from
	disabled debug printing while preserving the public interface.
	"""
	return

[Copilot]

The Apriel2StochasticMixer loads ALL sub-mixer weights into memory even though only one mixer is active at a time (line 2352). For models with multiple mixer types per layer, this could significantly increase memory usage. While the PR description mentions this is for "runtime switching", consider documenting the memory implications or providing a mode to only load the active mixer weights if memory efficiency is a concern.

[Copilot]

Monkey-patching the vLLM Worker class (lines 2843-2863) at module import time is fragile and could cause issues with vLLM version changes or multi-process scenarios. The patch is applied globally when the module is imported (line 2863), which happens automatically when vLLM loads the model. Consider:

1. Adding a version check to ensure compatibility
2. Providing a way to disable this patching if it causes issues
3. Documenting this behavior clearly as it modifies vLLM's internal classes

This pattern may break with vLLM updates and makes debugging more difficult.

[Copilot]

The auto_model_class is set to AutoModelForImageTextToText, but based on the PR description and code, this appears to be a text-only model (Apriel2TextConfig). The configuration change suggests this might support multimodal variants, but there's no indication in the PR that image-text functionality is being added. This could cause incorrect model loading behavior.

Suggested change

	auto_model_class=transformers.AutoModelForImageTextToText,
	auto_model_class=transformers.AutoModelForCausalLM,

[Copilot]

This statement is unreachable.

Suggested change

	if USE_VLLM_MAMBA_OPS:
	raise ImportError("vLLM mamba ops not yet wrapped")
	else:
	from mamba_ssm.ops.selective_scan_interface import selective_scan_fn
	from mamba_ssm.ops.triton.selective_state_update import selective_state_update
	# vLLM mamba ops are not yet wrapped; use mamba_ssm implementation when available.
	from mamba_ssm.ops.selective_scan_interface import selective_scan_fn
	from mamba_ssm.ops.triton.selective_state_update import selective_state_update

[Copilot]

This statement is unreachable.

[Copilot]

This statement is unreachable.

Suggested change

	# Debug final norm
	batch_size, seq_len = hidden_states.shape[:2]
	_debug_final = False # seq_len <= 10
	if _debug_final:
	# Show LAST token (to match vLLM)
	last_token = hidden_states[0, -1, :8]
	vals = ", ".join(f"{v:.6f}" for v in last_token.float().tolist())
	print(f"[TF Final] hidden_states (before norm): shape={hidden_states.shape}, last_token_first8=[{vals}]")
	print(f"[TF Final] norm.weight: first8=[{', '.join(f'{v:.6f}' for v in self.norm.weight.flatten()[:8].float().tolist())}]")
	print(f"[TF Final] norm.variance_epsilon={self.norm.variance_epsilon}")
	hidden_states = self.norm(hidden_states)
	if _debug_final:
	last_token = hidden_states[0, -1, :8]
	vals = ", ".join(f"{v:.6f}" for v in last_token.float().tolist())
	print(f"[TF Final] hidden_states (after norm): shape={hidden_states.shape}, last_token_first8=[{vals}]")
	batch_size, seq_len = hidden_states.shape[:2]
	hidden_states = self.norm(hidden_states)

[Copilot]

This statement is unreachable.

[Copilot]

This statement is unreachable.