trial(translate/riva-4b-v2): Riva-Translate-4B CoreML probe — kernel floor findings + MLX baseline - #81
Open
Alex-Wengg wants to merge 2 commits into
Open
trial(translate/riva-4b-v2): Riva-Translate-4B CoreML probe — kernel floor findings + MLX baseline#81Alex-Wengg wants to merge 2 commits into
Alex-Wengg wants to merge 2 commits into
Conversation
…4B-Instruct-v2 Feasibility probe answering whether a 4B Mistral-class translation LLM is viable on-device via CoreML. Conversion recipe follows the qwen3-asr stateful-KV-cache pattern, traced fully in fp16 so it converts on a 24GB host (~70s), with embed_tokens host-side and lm_head as a separate model. Results (M-series 24GB): fp16/GPU decodes at 11.5 tok/s with exact 14/14 greedy parity vs fp16 torch (logits corr 0.999999). int4 per-block-32 cuts the decoder 7.0GB -> 2.0GB, loads in 3.8s, prefills 3x faster, and still produces a correct translation (different but valid greedy path). Decode is dispatch-overhead-bound, not bandwidth-bound (int4 did not speed it up). CPU_AND_NE is strictly worse (209ms/tok, corr 0.785) - same conclusion as the Qwen3-0.6B LLM-on-ANE trial: autoregressive decode of this shape gains nothing from the ANE. Verdict: converts cleanly, runs fine on Mac, marginal on iOS; MLX remains the better host for 4B-class decode. mlpackages/npy artifacts gitignored.
…seline Optimization pass on the Riva-4B CoreML probe. Decoder-only profiling across quant formats shows CoreML's quantized GEMV kernels hit a ~40ms/tok floor (int4 moves 4x less data than fp16 but decodes only 1.5x faster; palettized LUT is 2.3x SLOWER than fp16). RangeDim shape churn measured free; host-side gaps between predicts measured near-free (busy-wait probe); alternating between two MLModel instances costs ~24ms/step, fixed by a --fused variant (decoder + final norm + tied lm_head in one stateful graph, last-position logits) which lifts prefill to 228 tok/s. Sustained-load bimodality observed (~41ms vs ~68ms regimes, GPU clock management). MLX 4-bit baseline on the same machine: 106.7 tok/s decode, 3.5GB peak, same correct translation - 4.4x CoreML's best steady-state (24 tok/s). The gap lives in the kernels, not the host or model structure, so the verdict stands quantitatively: 4B decode belongs on MLX; CoreML keeps prefill/ encoder workloads. Speculative decoding and chunked pipelines intentionally not pursued (draft-model machinery / adds the handoff cost it would need to avoid).
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Feasibility probe of nvidia/Riva-Translate-4B-Instruct-v2 (Mistral-class 4.2B text NMT, 37 languages) on CoreML, answering: does a 4B translation LLM run acceptably on-device? Conversion works with exact parity; the transferable value is the profiling findings, which apply to any large autoregressive CoreML decoder.
Conversion recipe
Follows the
qwen3-asrstateful-KV template with two memory-driven changes for a 24GB host:embed_tokenshost-side, lm_head separate; a--fusedvariant puts decoder + final norm + tied lm_head in one graph (last-position logits)fp16 result: 14/14 greedy-token parity with PyTorch (logits corr 0.999999) on an en→de prompt.
Findings (M5 Pro, 24GB)
Verdict
4B-class decode belongs on MLX; CoreML wins prefill (228 tok/s vs MLX ~39 on a 36-token prompt), keeping the CoreML/ANE-encoder + MLX-decoder hybrid pattern relevant. Scripts + full writeup in
models/translate/riva-translate-4b-v2/coreml/README.md; large artifacts gitignored.🤖 Generated with Claude Code