Skip to content

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
mainfrom
feat/riva-translate-4b-coreml
Open

trial(translate/riva-4b-v2): Riva-Translate-4B CoreML probe — kernel floor findings + MLX baseline#81
Alex-Wengg wants to merge 2 commits into
mainfrom
feat/riva-translate-4b-coreml

Conversation

@Alex-Wengg

Copy link
Copy Markdown
Member

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-asr stateful-KV template with two memory-driven changes for a 24GB host:

  • 34-layer decoder traced end-to-end in fp16 (no fp32 master copy) — converts in ~70s
  • embed_tokens host-side, lm_head separate; a --fused variant 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)

decoder size ms/tok effective BW
fp16 7.0GB 58.8 ~142GB/s
int8 per-channel 3.5GB 47.5 74GB/s
int4 per-block-32 2.0GB 41.3 49GB/s
int4 per-channel 1.8GB 39.1 46GB/s
4-bit palettized LUT 1.8GB 95.5 pathological
  1. CoreML quantized GEMV kernels floor at ~40ms/tok — int4 moves 4× less data than fp16 but decodes only 1.5× faster. Palettized LUT is 2.3× slower than fp16: avoid on GPU.
  2. Alternating two MLModel instances costs ~24ms/step (41.1 + 1.5ms measured 66.8ms interleaved) — fixed by the fused single-graph variant.
  3. RangeDim shape churn and host-side gaps are free (measured, not assumed).
  4. CPU_AND_NE strictly worse (209ms/tok, logits corr 0.785) — same conclusion as the Qwen3-0.6B LLM-on-ANE trial.
  5. MLX 4-bit on the same machine: 106.7 tok/s (3.5GB peak, same correct translation) — 4.4× CoreML's best; the gap lives in the kernels and no restructuring closes it.

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

…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).
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant