Porting Sakana AI's TRINITY Qwen-based model to Elixir/Bumblebee/Nx/Axon

I put in some elbow grease yesterday. Seems to be pretty solid so far. Using a modified, tiny 0.6B Qwen SLM for robust routing might turn out to be quite useful!

Sakana is using their version of TRINITY (obviously, theirs is not served from Elixir/Bumblebee/etc) in their forthcoming commercial AI framework. So, this is pretty cutting edge. Glad to stand on the shoulders of giants (both in the research community, as well as the ML community in Elixir)

I’m running this on a 5060 Ti 16GB, but it should work on most any GPU since it’s an SLM.

Right now it’s just saying what model it would call, so it’s not complete framework yet, just a proving ground.

$ XLA_TARGET=cuda12 mix run examples/qwen_router_prompt_eval.exs

=== QWEN ROUTER PROMPT EVAL ===

What this does
  Loads the local adapted Qwen router once, sends fixed transcripts through it,
  and checks whether the selected agent slot and role match expectations.

  No external LLM/provider calls are made.

Agent labels
  The agent names are labels from the original Sakana checkpoint.
  Example: agent 4 is labeled "google/gemma-3-27b-it".
  This eval does not call Gemma; it only reports that the router selected
  checkpoint slot 4.

Artifact
  priv/sakana_trinity/adapted_qwen3_0_6b_layer26

Model
  base router model: Qwen/Qwen3-0.6B
  router head shape: [10, 1024]
  assertion mode: strict

Native logs
  hidden in normal mode: tmp/examples/qwen_router_prompt_eval.native.log
  use --debug-native-logs to print XLA/CUDA compiler logs inline


[1/12] math_direct - PASS

Prompt sent to router:
  user: What is 17 + 25? Answer briefly.

Expected route:
  agent 4: google/gemma-3-27b-it
  role  2: Verifier

Router returned:
  agent 4: google/gemma-3-27b-it
  role  2: Verifier

Router input tokens: 15


[2/12] math_proof - PASS

Prompt sent to router:
  user: Prove that the sum of the first n odd positive integers is n squared. Route this to the best next role.

Expected route:
  agent 0: gpt-5
  role  0: Worker

Router returned:
  agent 0: gpt-5
  role  0: Worker

Router input tokens: 26


[3/12] code_debug - PASS

Prompt sent to router:
  user: A Python function mutates its default list argument across calls. Identify the bug and propose the smallest fix.

Expected route:
  agent 0: gpt-5
  role  0: Worker

Router returned:
  agent 0: gpt-5
  role  0: Worker

Router input tokens: 23


[4/12] security_review - PASS

Prompt sent to router:
  user: Review this login flow for security risks: passwords are hashed, sessions are cookies, and reset tokens never expire.

Expected route:
  agent 4: google/gemma-3-27b-it
  role  1: Thinker

Router returned:
  agent 4: google/gemma-3-27b-it
  role  1: Thinker

Router input tokens: 24


[5/12] planning - PASS

Prompt sent to router:
  user: Create a concise implementation plan for migrating a small Elixir service from in-memory state to Postgres.

Expected route:
  agent 4: google/gemma-3-27b-it
  role  0: Worker

Router returned:
  agent 4: google/gemma-3-27b-it
  role  0: Worker

Router input tokens: 22


[6/12] verification_after_worker - PASS

Prompt sent to router:
  user: Calculate 6 * 7 and verify the answer.
  assistant: Worker answer: 6 * 7 = 42.

Expected route:
  agent 4: google/gemma-3-27b-it
  role  2: Verifier

Router returned:
  agent 4: google/gemma-3-27b-it
  role  2: Verifier

Router input tokens: 28


[7/12] needs_revision - PASS

Prompt sent to router:
  user: Check whether the answer is correct: 19 + 24 = 41.
  assistant: Worker answer: 19 + 24 = 41.

Expected route:
  agent 4: google/gemma-3-27b-it
  role  2: Verifier

Router returned:
  agent 4: google/gemma-3-27b-it
  role  2: Verifier

Router input tokens: 38


[8/12] ambiguous_decomposition - PASS

Prompt sent to router:
  user: This problem has unclear requirements. We may need to split it into assumptions, risks, and a concrete next action.

Expected route:
  agent 0: gpt-5
  role  2: Verifier

Router returned:
  agent 0: gpt-5
  role  2: Verifier

Router input tokens: 25


[9/12] creative_but_constrained - PASS

Prompt sent to router:
  user: Draft a friendly but precise support reply explaining a billing correction. Keep it under 120 words.

Expected route:
  agent 4: google/gemma-3-27b-it
  role  2: Verifier

Router returned:
  agent 4: google/gemma-3-27b-it
  role  2: Verifier

Router input tokens: 23


[10/12] longer_context - PASS

Prompt sent to router:
  system: You are routing work inside a three-role TRINITY loop. Worker solves, Thinker plans or redirects, Verifier checks.
  user: Given a release checklist, identify the next best role. The feature compiles, unit tests pass, docs changed, but no smoke test has been run yet.

Expected route:
  agent 4: google/gemma-3-27b-it
  role  2: Verifier

Router returned:
  agent 4: google/gemma-3-27b-it
  role  2: Verifier

Router input tokens: 62


[11/12] provider_failure_triage - PASS

Prompt sent to router:
  user: The last provider call timed out after 30 seconds. Decide whether to retry, ask a thinker for a smaller plan, or verify the partial answer.

Expected route:
  agent 4: google/gemma-3-27b-it
  role  2: Verifier

Router returned:
  agent 4: google/gemma-3-27b-it
  role  2: Verifier

Router input tokens: 33


[12/12] final_answer_check - PASS

Prompt sent to router:
  user: Solve and then verify: the capital of France is Paris.
  assistant: Worker answer: Paris is the capital of France.
  assistant: Thinker note: this is a factual lookup and likely ready.

Expected route:
  agent 4: google/gemma-3-27b-it
  role  0: Worker

Router returned:
  agent 4: google/gemma-3-27b-it
  role  0: Worker

Router input tokens: 41


Summary
  passed: 12
  failed: 0
  roles selected: Worker=4, Thinker=1, Verifier=7
  agent slots selected: 0=3, 4=9

PASS qwen_router_prompt_eval

Just for fun, I ported this to work on my 24GB M4 MBP using the Emily backend. In doing so, during the export, I ran into a limitation of the current native mlx libraries - their svd functions have no ‘thin’ mode and always materialise the full matrix. For the Qwen embedder that turned out to be ~92GB of memory.

I updated Emily to support this mode (in specific cases) directly via the Gram matrix, could do the one-time export in ~2s and was able to run the qwen router example.

If you are interested, the changes I made are here

Nx provides a default implementation that doesn’t have any explicit evals, so that might be a way to avoid those materializations.

Ah, that makes sense. Thanks both.

@ausimian, the Emily result is extremely useful to know. The fact that native MLX SVD materialized the full matrix, but the Gram-matrix path made the one-time export work on a 24GB M4, is exactly the kind of backend detail I would not have known to check.

@polvalente, your point about Nx’s default implementation is really helpful. If the default implementation avoids explicit evals, then it may avoid the immediate materialization issue seen in the native MLX SVD path.

This helped me separate two concerns: what TRINITY needs from Nx today, and what backend-neutral code can safely assume. In TRINITY, the export path only needs the reduced SVD outputs for reconstruction. As an application author, though, it is tempting to read:

Nx.LinAlg.svd(tensor, full_matrices?: false)

as a “thin SVD” guarantee in both output shape and execution/memory behavior. The safer cross-backend interpretation seems to be:

1. full_matrices?: false is an output-shape contract.

2. Execution strategy and memory behavior are backend-dependent.

If that is the intended boundary, I have two practical questions:

1. If a backend’s native SVD is known to materialize full matrices anyway, should the preferred backend behavior be to notice full_matrices?: false and fall back to Nx’s default implementation when that is likely to behave better?

2. Would the Nx team be open to a small docs PR clarifying that full_matrices? controls SVD output shapes but does not guarantee a reduced-memory execution profile?

For the current TRINITY path, this does not force an immediate change: it is still fine to call Nx.LinAlg.svd/2 and let Nx/the selected backend choose native SVD, default Nx lowering, or a backend-specific workaround. But the next rework is to make backend/runtime selection a profile concern instead of something CUDA-shaped in the application code, so I want to make sure I am relying on the right Nx contract before generalizing it. I do not want application code branching on Emily vs EXLA vs Torchx unless there is no better option.

Side note: I did go down a rabbit hole thinking about how production code could ask a backend about these paths before starting a large SVD. I wrote that up here: Concrete proposal: SVD execution capabilities in Nx. It is probably too much architecture for this immediate issue, but I am leaving it as background in case backend-specific memory behavior becomes a larger Nx API discussion later.

@gtcode any backend-specific issues should be addressed in the specific backends. For instance, EMLX, which is the original MLX-based backend for Nx, does not fall back to the binary backend at all, and relies on these default implementations from Nx for the most part when things aren’t supported by the native backend.

I think it doesn’t really make sense to document this specific issue in Nx because this can happen with literally any option, as different backends are limited by different constraints. Maybe a general comment on this would be welcome.

Also, feel free to test this out with EMLX in the github main version, as I’m actively working on improvements. v0.3 should be out soon.

I gave a very quick pass through the proposal, and I think there’s something being overlooked there. It focuses on Nx.block, but the same concerns are pertinent to all Nx callbacks.

Also keep in mind that blocks and block structs can be defined by Nx users, not just Nx itself.

Additionally, even Nx.dot might have different execution profiles depending on the backend.

Finally, lowered vs eval’d is also something that’s not a concern for the backend, but for the compiler.

Do feel welcome to open an issue for us to discuss there and hopefully achieve a great feature, but I suggest rethinking the approach to something more general that includes all Nx callbacks.

@ausimian — this is great, thank you. The ~92 GB → ~2 s number on the embedder is very helpful – I would not have been able to diagnose that failure mode from the CUDA side, and your Emily Gram path is the right idea.

I went through your two commits and the Emily PR. The diagnosis and the backend-level fix both look mostly right. The one thing I’d note is that since you opened that PR, the trinity_coordinator has moved to a runtime-profile architecture (mix … --runtime-profile <name>), and the prompt-eval suite is now 37 fixed cases with snapshot + determinism + margin-floor checks rather than the original 12. So, I don’t want to merge the PR verbatim. config :nx, default_backend: Emily.Backend globally and making "emily" a valid XLA_TARGET would blast through a few things on our end. But, I do want to land an explicit :emily runtime profile that takes the same intent: Emily as the default backend when the operator chooses that lane, untouched otherwise.

The blocker for me is that I don’t have Apple hardware. Would you be up for one validation pass against current main so I can shape the upstream integration with confidence? I think it’s two commands:

# Pull current main, add Emily 0.4.0 (the Gram fix is already in there), then:

mix trinity.sakana.export_adapted \
  --force \
  --svd-compute-type f32 \
  --out tmp/emily_adapted_qwen3_0_6b_layer26

mix run examples/qwen_router_prompt_eval.exs \
  --artifact-dir tmp/emily_adapted_qwen3_0_6b_layer26 \
  --snapshot examples/fixtures/qwen_router_prompt_eval_logits.json \
  --determinism-runs 2

What I’d love to see in the output:

• export reaches status=complete with no OOM,
• per-tensor SVD time for the embedder,
• 37/37 prompt cases pass,
• if anything fails, whether it’s route_hash (likely backend-sensitive) or one of the decision-stable fields (agent_id / role_id / token_count / transcript_hash) — that distinction tells me whether Emily needs its own snapshot lane or whether something deeper drifted.

If you’d rather not run the full thing, even just the export half is useful. Either way, thanks again for doing this — turning it into a real Apple lane upstream is now basically a paperwork exercise on my side because of the work you did.

@gtcode How can I run this on EMLX as well?

I was able to adapt it locally. I hit the same 92GB allocation, so I’ll see if I can add a better upstream Nx implementation for Thin SVD so that EXLA can also take advantage of it.

I do want to encourage people to reach out on the Nx issues tracker when they hit these types of limitations.

I got a working version of thin SVD working here: refactor: better memory footprint for thin svd by polvalente · Pull Request #1753 · elixir-nx/nx · GitHub

It uses a smaller memory footprint which enabled EMLX to run the export without OOMs and without having to rewrite anything. This also benefits EXLA, which uses the same Nx base code.

PR should be merged soon, I just gotta fix some test assertions and documentation.

With just this change and replacing all EXLA.Backend cuda config with the equivalent EMLX, device: :gpu setup, I got the outputs:

mix trinity.sakana.export_adapted
–force
–svd-compute-type f32
–out tmp/emily_adapted_qwen3_0_6b_layer26

…

tensor_export_progress: path=language_modeling_head.output.kernel decompose_ms=32908 reconstruct_ms=0
tensor_export_finished: path=language_modeling_head.output.kernel
Export status: complete
Export complete flag: true
Tensor count: 9
Completed: 9 / 9

and

mix run examples/qwen_router_prompt_eval.exs \
–artifact-dir tmp/emily_adapted_qwen3_0_6b_layer26 \
–snapshot examples/fixtures/qwen_router_prompt_eval_logits.json \
–determinism-runs 2

Summary
passed: 37
failed: 0
roles selected: Worker=18, Thinker=4, Verifier=15
agent slots selected: 0=17, 1=1, 4=19

edit: note that I didn’t configure emlx_axon, so the EMLX.Fast rewrites were not used