lablup / lablup/mlxcel

fix(drafter): validate quantization against SUPPORTED_AFFINE_BITS

Open
#1,790 0 comments 0 reactions 0 assignees View on GitHub
area:core priority:low status:ready type:bug
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.