perf(llama3): project only the sampled row through the LM head at prefill #1968

Problem / Background

Follow-up to #1955, which did this for Cohere2 (+14% prefill). Llama3Model in src/models/llama3.rs (which also serves Mistral and Qwen2/Qwen2.5 through the registry) uses the LanguageModel::forward_last_logits default in src/lib/mlxcel-core/src/generate.rs, so every single-sequence prefill (CLI, bench, server) computes [1, L, vocab] logits through the LM head and keeps one row. At 512 tokens that is about 0.5 TFLOP of wasted head work for a 128k (Llama 3) or 152k (Qwen2.5) vocabulary.

Separately, VisionLanguageModel in src/vision/mod.rs does not forward the three forward_last_logits* entry points to its text model, so any text-model override is skipped behind the vision wrapper (Gemma 3 VLM, Llama 4 VLM, LLaVA, FastVLM).

Proposed Solution

• Split the two duplicated forward bodies of Llama3Model into hidden_states (embeddings, layers, final norm) plus the head.
• Override forward_last_logits, forward_last_logits_with_sequence_id, and forward_last_logits_with_embeddings_and_sequence_id to slice the hidden state at last_pos before the head; make generate::logits_at_position public for that.
• Forward the three entry points through VisionLanguageModel to the text model.

Measurements

M1 Ultra, mlxcel-bench-decode with bench_decode.sh arguments, ABBA: Llama 3.1 8B 4-bit prefill 771 to 826 tok/s at pp512 and 764 to 819 at pp2048 (+7.1%); Qwen2.5 7B 4-bit 817 to 881 and 816 to 881 (+7.8%). Decode unchanged. Greedy output identical on three prompts per model and on LLaVA 1.6 Mistral with an image.

Acceptance Criteria

• Unit test in src/models/llama3_tests.rs shows all three entry points match the sliced full forward at the last and at an interior position
• Greedy output identical on real checkpoints (Llama 3.1 8B, Qwen2.5 7B, LLaVA 1.6 Mistral with an image)
• Benchmark record in docs/benchmark_results/
• Remaining families (Qwen3, Qwen3-MoE, Phi-3, Mixtral, Gemma 2/3, DeepSeek V3, Nemotron-H, Jamba) tracked as a follow-up issue

Verification

cargo test --release llama3_tests::
cargo clippy --all-targets -- -D warnings
cargo fmt --check