perf(qwen3): add crossrow QMV verifier - #275
Conversation
Use MLX 0.31 for the Metal ABI and preserve stock QMV for unsupported row counts or kernel errors.
|
Important Draft PR not reviewedDraft PRs are not automatically reviewed by default.
To automatically review draft PRs, update your CodeRabbit configuration: reviews:
auto_review:
drafts: trueThanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
The crossrow Metal kernel hard-codes a bfloat16 view of the activation buffer. An fp32 input misreads it (halved row stride; odd 16-bit words become garbage activations with NaN/Inf bit patterns), so the projection returns corrupt values that propagate as non-finite logits and stable one-token repetition loops. Observed live on the Escha W2 checkpoint: short suffix prefills with 2/3/5/6/8/9 rows hit this on the GDN in_proj_ba projection. Gate eligibility on x.dtype() == BFloat16; anything else takes the stock quantized matmul.
|
Found a live correctness bug in this crossrow path while debugging repetition loops on the Escha W2 checkpoint — fixed in 201afa9. Bug: the Metal kernel hard-codes Evidence (nightly + serving branch): same prompt, same build — crossrow enabled → Fix: eligibility now requires The kernel could later be taught to handle f32 natively; until then the gate keeps the fast path bit-exact and correctness-neutral. |
crossrow_qmv.rs (new in this PR) violated the workspace restriction lints: as_conversions (try_from/unwrap_or instead; crossrow_group_count drops const, which no caller needs), format_push_string (write!), and shadow_reuse in the kernel builders. Synthetic test-data casts keep under a tests-module as_conversions allow. bonsai_q1.rs is pre-existing base-branch lint debt (chunks_exact -> as_chunks) that blocks the -D warnings job for every PR; included so this PR's CI can go green.
Run the repo's lints (clippy 1.98, RUSTFLAGS=-Dwarnings, matching CI) on the code this PR introduces or modifies: write_with_newline, doc_markdown, redundant_pub_crate, useless_conversion, cast lints, format_push_string, shadowing, needless borrows, question_mark, unnecessary_map_or, missing_const_for_fn. Also auto-fixes the same chunks_exact-to-as_chunks warning in turboquant tests (pre-existing base debt). Remaining -Dwarnings failures on this branch are the pre-existing mlx_rs::transforms::eval disallowed-method debt in files this PR does not touch (gemma2/3, siglip, llava_qwen2, lib.rs, qwen3_next internals, turboquant); those predate this PR and need the mlx_exec gate migration, not this PR.
Summary
Port the smallest validated Qwen3.8-27B performance path from the nightly work into current
upstream/main.Cargo.toml/Cargo.lock.Validation
cargo fmt --all -- --checkcargo check -p higgs-modelscargo test -p higgs-models --lib crossrow_qmv -- --nocapture --test-threads=1(5 passed including M4/M7 fallback contract)cargo build --releaseomen diff -t upstream/main: 5 files, 668 additions, 7 deletions; medium risk due to the single kernel implementation.The benchmark result motivating this port was Qwen3.8-27B depth-5 Low Power Mode validation at 7.39 tok/s median versus 5.85 tok/s stock, with the exact verifier trajectory preserved.