fix(drafter): validate quantization against SUPPORTED_AFFINE_BITS
- Dominant language
- Rust
- Stars
- 467
- Forks
- 54
- Avg merge
- 4h 25m
- Merged PRs (30d)
- 310
Description
## Problem
Found during review of PR #1778, which added `SUPPORTED_AFFINE_BITS`, `SUPPORTED_AFFINE_GROUP_SIZES` and two validators to `mlxcel_core::layers` and applied them to `split-mtp`. Two drafter paths still bypass them.
1. The Qwen 3.5 MTP drafter quantizes its dense projections at load with `bits` and `group_size` taken straight from its config. `quantize_weights_with_mode` is a non-`Result` bridge, so an unsupported value throws inside MLX and aborts the process, the failure PR #1778 fixed for `split-mtp`. `group_size: 0` panics earlier, on the `%` at `model.rs:265`. The drafter loads lazily in the batch worker at the first speculative request, so either takes the server down.
2. The Markov (DFlash) head checks a quantized factor's packed width against its own list `[2, 3, 4, 6, 8]`, which omits 5, so a 5-bit factor is refused at load ("no supported bit depth") although MLX supports 5.
## Evidence
- `src/lib/mlxcel-core/src/drafter/qwen3_5_mtp/model.rs:243` `quantize_dense_projections`: config values at `:248`, `% group_size` at `:265`, `ffi::quantize_weights_with_mode(w, group_size, bits, "affine")` at `:270`; called from `from_path` at `:212`.
- The bridge returns `UniquePtr`, not `Result` (`src/lib/mlxcel-core/src/lib.rs:2754-2759`).
- `validate_bounds` (`src/lib/mlxcel-core/src/drafter/qwen3_5_mtp/config.rs:169-236`), the drafter's config gate, does not check `quantization.bits` or `group_size` (`:85-95`).
- `src/lib/mlxcel-core/src/drafter/dflash/markov.rs:200` (`const SUPPORTED_BITS: [i64; 5] = [2, 3, 4, 6, 8];`) in `validate_markov_factor` (`:162`).
- `src/lib/mlxcel-core/src/layers.rs`: `validate_affine_quantization_bits` (`:1370`), `validate_affine_quantization_group_size` (`:1401`), `SUPPORTED_AFFINE_BITS` (`:1455`), `SUPPORTED_AFFINE_GROUP_SIZES` (`:1464`).
## Proposed fix
- Call both shared validators in `quantize_dense_projections` once `candidates` is known to be non-empty (so before the `%`), return `Result`, and map the error to `DrafterError::Config` in `from_path`. Validate there rather than in `validate_bounds`: a pre-converted checkpoint whose tensors this function skips should not be refused over values it never uses.
- Derive Markov's list from `SUPPORTED_AFFINE_BITS`.
- Update the `Used by:` lists in `layers.rs` (`:1369`, `:1400`, `:1452-1454`). The two validators' lists already omit their CLI caller, `preflight_checks` in `src/commands/split_mtp.rs:181-184` (added by PR #1778), so add it together with the drafter callers.
## Acceptance criteria
- [ ] An unsupported drafter `bits` or `group_size` (including 0) makes the load return an error instead of aborting or panicking (test).
- [ ] A 5-bit quantized Markov factor passes `VanillaMarkovHead::from_weights`, and an unexplainable width is still refused (tests).
## Verification
```bash
cargo test --workspace --profile test-fast --features metal,accelerate
cargo clippy --workspace --all-targets -- -D warnings
cargo fmt --all -- --check
```
Each new test must fail with its fix reverted.
Contributor guide
Research direction
Start with quantize_dense_projections and from_path in src/lib/mlxcel-core/src/drafter/qwen3_5_mtp/model.rs, then read the shared validators in layers.rs and their error handling. Check validate_markov_factor in drafter/dflash/markov.rs and the existing drafter and split-mtp tests. Done means invalid bits or group sizes return errors, 5-bit Markov factors pass, invalid widths remain rejected, and the listed cargo test, clippy, and fmt checks pass.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- rust
- Domain
- backend, machine-learning
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 68/100