Skip to content

perf(qwen3): add crossrow QMV verifier - #275

Draft
dusterbloom wants to merge 4 commits into
panbanda:mainfrom
dusterbloom:codex/qwen38-minimal-upstream
Draft

perf(qwen3): add crossrow QMV verifier#275
dusterbloom wants to merge 4 commits into
panbanda:mainfrom
dusterbloom:codex/qwen38-minimal-upstream

Conversation

@dusterbloom

Copy link
Copy Markdown
Contributor

Summary

Port the smallest validated Qwen3.8-27B performance path from the nightly work into current upstream/main.

  • Add the crossrow Metal QMV verifier for affine 4-bit, group-64 weights.
  • Enable only validated M2/M3/M5/M6/M8/M9 shapes; M4/M7 remain on stock QMV.
  • Fall back to stock QMV on opt-out, unsupported shapes, reshape failures, or kernel errors.
  • Pin the required MLX 0.31 fork in Cargo.toml/Cargo.lock.
  • Keep Qwen3-Next weight loading chunked/evaluated for the measured path.

Validation

  • cargo fmt --all -- --check
  • cargo check -p higgs-models
  • cargo test -p higgs-models --lib crossrow_qmv -- --nocapture --test-threads=1 (5 passed including M4/M7 fallback contract)
  • cargo build --release
  • omen 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.

Use MLX 0.31 for the Metal ABI and preserve stock QMV for unsupported row counts or kernel errors.
@coderabbitai

coderabbitai Bot commented Aug 18, 2026

Copy link
Copy Markdown
Contributor

Important

Draft PR not reviewed

Draft PRs are not automatically reviewed by default.

  • Trigger a manual review

To automatically review draft PRs, update your CodeRabbit configuration:

reviews:
  auto_review:
    drafts: true

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

Comment @coderabbitai help to get the list of available commands.

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.
@dusterbloom

Copy link
Copy Markdown
Contributor Author

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 reinterpret_cast<const device bfloat16_t*>(x), but nothing gates the input dtype. When a QLinear with affine-4/g64 weights receives an fp32 activation (here: the GDN in_proj_ba projection during short suffix prefills, m ∈ {2,3,5,6,8,9}), the buffer is misread — halved row stride plus garbage odd 16-bit words — producing corrupt values and NaN bit patterns. Downstream this surfaced as all-248320-logits NaN (verified with a fail-fast finiteness tripwire in sample()) and stable one-token doom loops.

Evidence (nightly + serving branch): same prompt, same build — crossrow enabled → !!!!!… repetition loop / NaN logits; HIGGS_CROSSROW_QMV=0 → clean; with the dtype gate → clean and token-identical to the disabled run.

Fix: eligibility now requires x.dtype() == Dtype::Bfloat16; everything else falls back to stock quantized matmul, which is the path this PR already documents as the correctness baseline. Build + crossrow bit-exact tests pass locally. (This branch predates the Escha loader, so the live repro itself can't run here — it was verified on the branches that carry the serving path.)

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.
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