Skip to content

refactor(rai): run Qwen3, Qwen3.5 and Gemma 4 as corelib's own models - #775

Closed
chizamd wants to merge 10 commits into
ROCm:mainfrom
chizamd:feature/rai-model-api
Closed

chizamd wants to merge 10 commits into
ROCm:mainfrom
chizamd:feature/rai-model-api

Conversation

@chizamd

@chizamd chizamd commented Oct 11, 2026

Copy link
Copy Markdown
Contributor

Stacked on #774 — review the 9 commits after abcfa3b4 (f4c49ab5..1a754611). Retarget or rebase once #774 lands.

Qwen3, Qwen3.5 and Gemma 4 on the rai backend stop driving corelib op by op and run as corelib's own models (ryzenai_corelib_model_*), the way Qwen3.6 already did. All four now share one engine and one backend; each model is a ~20-line spec. Phi-4 is unchanged (corelib's phi4mini takes ONNX only, and FastFlowLM ships a GGUF).

What changes

One engine: flm::rai::CorelibModel (include/rai/corelib_model.hpp, common/rai/corelib_model.cpp)

  • A causal_lm over a corelib model, generalized from the Qwen3.6 engine. ModelSpec states what differs between models:
    • Continuation — where a multi-row pass may start past position 0: anywhere (Qwen3.5/3.6), only at 0 (Qwen3, so a later prompt re-runs the conversation), or at 0 / a 4096-row page boundary (Gemma 4).
    • Rewind — kSeek for attention-only models (every pass states its position; rewind and restore just move it), kSnapshot for Qwen3.6 (model_state_save/restore), kReplay for Qwen3.5 (reset and replay).
  • Token ids and the context limit are checked before corelib is called; a failed forward, restore or reset poisons the engine.
  • clear_context on a kSeek model makes no corelib call: a pass never reads the cache past its position, and model_reset cost ~60 ms per request on Qwen3.

One backend: flm::rai::Factory(BackendSpec) (include/rai/corelib_model_backend.hpp)

  • Refuses preemption and a context outside the model's rule, reads the package's token ids, and resolves the GGUF before the runtime or weights are touched.
  • EOS ids are either given (Qwen3, Qwen3.5) or read from the package: tokenizer_config.json's eos_token plus config.json's eos_token_id (Qwen3.6, Gemma 4 → {1, 106}); Gemma 4's BOS comes from bos_token.

Removed

  • The per-op engines, GGUF readers, shape plans and host math of Qwen3, Qwen3.5 and Gemma 4; rai/gguf_file, rai/weight_cache, and the Q8/RMSNorm parts of rai/host_ops; 14 corelib entry points nothing calls any more.
  • There is no packed-weight cache on this path: corelib packs the weights on every load. Old <gguf stem>.rai-weights.bin/.json files can be deleted.

Tests

  • New src/test/rai_model/: 36 engine tests (every Rewind/Continuation combination, poisoning, capacity) and 14 backend tests against the fake corelib. The Qwen3.6 tests moved here.
  • Gemma 4's registry, catalog and server-gate tests moved verbatim to test_gemma4_registry.cpp, plus one test of the Gemma 4 spec. The op-level Gemma 4 and Qwen3 suites are gone.

Docs: src/common/models/README.md gains §2a ("a model corelib runs whole"); the Qwen page notes the dropped weight cache.

Results (xcomedusad-44, AIE4; old and new builds run interleaved)

Model Load, new vs old Long-prompt TTFT, new vs old Decode tok/s, new vs old Greedy text
qwen3-rai:0.6b 3.8–4.5 s vs 2.7–3.1 s 0.25–0.28 vs 0.26–0.39 s 69–79 vs 62–81 identical
qwen3.5-rai:9b 11–27 s vs 56 s 41.8 vs 49.3 s 13–20 vs 15.6 identical
qwen3.6-moe-rai:35b-a3b 87–88 vs 92 s 14.2 vs 14.75 s 18–20 vs 19 one near-tie token in 1 of 2 runs
gemma4-it-rai:e2b 7–15 s vs 3.2–4.9 s 1.17–1.39 vs 0.75–0.94 s 30.7 vs 40.4 one near-tie token (old flips it too)
gemma4-it-rai:e4b 12–17 s vs 3.5–5.6 s 1.95–2.41 vs 1.20–1.24 s 18.8 vs 26.6 identical
phi4-mini-it-rai:4b (unchanged) ~same 0.61–0.63 vs 0.605 s 35–39 vs 37 same as old
  • Gemma 4 E2B at 16k: passkey recalled; a short prompt after the full paged cache matches a fresh run.
  • Switching e2b → qwen3 → qwen3.5 → e4b → phi4 → e2b in one flm serve process: every answer matches a fresh run, no errors.
  • Fake-corelib suites: rai_model 2/2, phi4_rai 8/8, gemma4_rai 3/3 (real-package tokenizer test skipped).

Known regression: Gemma 4 is slower on corelib 0.26's built-in model. Prefill takes ~1.5× (E2B) / 1.6–2× (E4B), decode is 24–29% slower, and loads are slower. The cause is in corelib, not here: the host per-layer-embedding projection in Gemma4::BuildPle (FP32 weights, scalar, single-threaded on decode) and serial weight packing with no pre-packed input. A corelib fix is being prepared; FastFlowLM picks it up with a version bump. Qwen3 loads ~1 s slower without the packed-weight cache.

Not covered

  • Built and tested on Windows only; the Linux build of flm_rai was not compiled for this change.

Diff: 69 files; product code +801 / −8,423, tests +1,487 / −12,763.

chizamd and others added 10 commits October 10, 2026 17:03
…ma 4 at 16k

Moves the `--rai` backend from ryzenai-corelib 0.11 to 0.26 and builds on
what the new corelib offers.

corelib 0.26 (pinned exactly):
- Follow the ABI changes: `create_tensor_window` takes strides,
  `linear_attention_bf16` takes rows, and models are built with
  `model_weights_create_path` + `model_create`.
- A 3-D matmul output that writes the KV cache is a strided window
  `[max_seq*head, head, 1]` (Phi-4, Qwen3, Qwen3.5 V writes; Gemma 4 E4B
  staging window).
- Per-op engines open their stream through `corelib::OpenModelStream`,
  which enables the model's kernel bundle before any kernel query.
- corelib's per-model libraries are copied beside `flm.exe` on Windows
  and beside corelib by `home_install.sh` on Linux. Nothing needs
  `RYZENAI_CORELIB_DD_ROOT` or `RYZENAI_CORELIB_KERNELS_ROOT` any more.

Qwen3.6-35B-A3B (new `qwen3.6-moe-rai:35b-a3b`, text only):
- Runs on corelib's model API; context up to 262144 in 4096-row pages,
  a prompt is a single `model_forward`, and checkpoint/restore use
  corelib state snapshots.
- Fetch logits only for the row that gets sampled.

Gemma 4:
- Context of 1..4096 on flat attention, or a multiple of 4096 up to 16384
  on corelib paged attention; the aie_next catalog entries default to
  16384. A paged continuation starts from its own page.
- `gemma4-it-rai:e4b` is offered again.

Shared:
- `Sampler::sampler_topk_apply` selects top-k with a k-entry heap in one
  pass (same ranking and tie rule).

Tests: new `src/test/qwen36_rai`, and the Gemma 4, Phi-4 and Qwen3 rai
tests and fake corelib follow the 0.26 ABI and the 16k paged context.

Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
@chizamd

chizamd commented Oct 11, 2026

Copy link
Copy Markdown
Contributor Author

Moving this to the internal FLM-Medusa repo.

@chizamd chizamd closed this Oct 11, 2026
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