Context
Native Bumblebee.Text.generation (the CM7–CM13 expr-compiler work on feat/expr-compiler) compiles the whole generation graph — transformer forward + defn while decode loop + KV cache + sampling — to the single-NIF replay. Measured ~5.07× over the evaluator on Qwen3-0.6B greedy decode (~61 vs ~12 tok/s, byte-identical output; see bench/qwen3_tokens_per_sec.exs).
But that number is achieved with mx::compile fusion disabled. A while-containing program is a data-dependent, host-controlled loop that mx::compile can't trace, so eval_program degrades it to a plain sync replay (the guard in c_src/program.cpp's eval_program_nif). The elementwise runs the replay leaves separate (RMSNorm/softmax/SiLU gating/residual adds) are therefore unfused. CM6 measured ~1.5–1.6× from this fusion on a decode-shaped transformer block, so there's likely meaningful headroom on top of the 5×.
Idea
The outer loop can't be traced, but the loop body (the per-token forward sub-program, instr.subprograms[1] of the While instruction) is shape-stable — offset is a runtime s32 input, so one compiled program serves every position. Wrap just the body sub-program in mx::compile (cached per stream, like the existing CM6 Program::CompiledEntry machinery), replaying the fused callable each iteration. The condition sub-program could get the same treatment.
Where
c_src/program.cpp — the While arm in replay_program: instead of replay_program(body_p, state, s) raw, replay a per-stream-cached mx::compile'd version of the body.- Reuse the
Program::compiled per-stream cache pattern (program.hpp) — but keyed/stored on the child (body) program.
Caveats / acceptance
- Not bit-identical.
mx::compile reassociates f32 (last-few-ULP), so the fused decode won't match the evaluator bit-for-bit (same caveat as CM6). generation_native_test.exs asserts bit-identical, so this must be opt-in (e.g. an eval-mode/option flag) or its conformance gate must switch to an all-close tolerance + a greedy-token-prefix match rather than exact.
- Shape-stability holds (offset is a runtime input), so the compiled body should cache-hit across tokens — verify it does, not recompile per step.
- Acceptance:
bench/qwen3_tokens_per_sec.exs shows the fused-while native lane beats the current native lane on Qwen3-0.6B; document the speedup and the tolerance trade-off.
Follow-up to the CM13 native-generation work (PRs #156–#162).
Context
Native
Bumblebee.Text.generation(the CM7–CM13 expr-compiler work onfeat/expr-compiler) compiles the whole generation graph — transformer forward +defn whiledecode loop + KV cache + sampling — to the single-NIF replay. Measured ~5.07× over the evaluator on Qwen3-0.6B greedy decode (~61 vs ~12 tok/s, byte-identical output; seebench/qwen3_tokens_per_sec.exs).But that number is achieved with
mx::compilefusion disabled. Awhile-containing program is a data-dependent, host-controlled loop thatmx::compilecan't trace, soeval_programdegrades it to a plain sync replay (the guard inc_src/program.cpp'seval_program_nif). The elementwise runs the replay leaves separate (RMSNorm/softmax/SiLU gating/residual adds) are therefore unfused. CM6 measured ~1.5–1.6× from this fusion on a decode-shaped transformer block, so there's likely meaningful headroom on top of the 5×.Idea
The outer loop can't be traced, but the loop body (the per-token forward sub-program,
instr.subprograms[1]of theWhileinstruction) is shape-stable —offsetis a runtimes32input, so one compiled program serves every position. Wrap just the body sub-program inmx::compile(cached per stream, like the existing CM6Program::CompiledEntrymachinery), replaying the fused callable each iteration. The condition sub-program could get the same treatment.Where
c_src/program.cpp— theWhilearm inreplay_program: instead ofreplay_program(body_p, state, s)raw, replay a per-stream-cachedmx::compile'd version of the body.Program::compiledper-stream cache pattern (program.hpp) — but keyed/stored on the child (body) program.Caveats / acceptance
mx::compilereassociates f32 (last-few-ULP), so the fused decode won't match the evaluator bit-for-bit (same caveat as CM6).generation_native_test.exsasserts bit-identical, so this must be opt-in (e.g. an eval-mode/option flag) or its conformance gate must switch to an all-close tolerance + a greedy-token-prefix match rather than exact.bench/qwen3_tokens_per_sec.exsshows the fused-while native lane beats the current native lane on Qwen3-0.6B; document the speedup and the tolerance trade-off.Follow-up to the CM13 native-generation work (PRs #156–#162).