Repository navigation
fix(drafter): validate quantization against SUPPORTED_AFFINE_BITS #1790
Copy link
Copy link
Open
Labels
area:coremlxcel-core: MLX FFI, primitives, KV cache, layersmlxcel-core: MLX FFI, primitives, KV cache, layerspriority:lowLow priorityLow prioritystatus:in-progressCurrently being worked onCurrently being worked ontype:bugBug fixes, error corrections, or issue resolutionsBug fixes, error corrections, or issue resolutions
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.
- The Qwen 3.5 MTP drafter quantizes its dense projections at load with
bitsandgroup_sizetaken straight from its config.quantize_weights_with_modeis a non-Resultbridge, so an unsupported value throws inside MLX and aborts the process, the failure PR fix(split-mtp): fix refusal text, --force shards, --q-bits validation #1778 fixed forsplit-mtp.group_size: 0panics earlier, on the%atmodel.rs:265. The drafter loads lazily in the batch worker at the first speculative request, so either takes the server down. - 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:243quantize_dense_projections: config values at:248,% group_sizeat:265,ffi::quantize_weights_with_mode(w, group_size, bits, "affine")at:270; called fromfrom_pathat:212.- The bridge returns
UniquePtr<MlxQuantizedWeights>, notResult(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 checkquantization.bitsorgroup_size(:85-95).src/lib/mlxcel-core/src/drafter/dflash/markov.rs:200(const SUPPORTED_BITS: [i64; 5] = [2, 3, 4, 6, 8];) invalidate_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_projectionsoncecandidatesis known to be non-empty (so before the%), returnResult, and map the error toDrafterError::Configinfrom_path. Validate there rather than invalidate_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 inlayers.rs(:1369,:1400,:1452-1454). The two validators' lists already omit their CLI caller,preflight_checksinsrc/commands/split_mtp.rs:181-184(added by PR fix(split-mtp): fix refusal text, --force shards, --q-bits validation #1778), so add it together with the drafter callers.
Acceptance criteria
- An unsupported drafter
bitsorgroup_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
cargo test --workspace --profile test-fast --features metal,accelerate
cargo clippy --workspace --all-targets -- -D warnings
cargo fmt --all -- --checkEach new test must fail with its fix reverted.
Reactions are currently unavailable
Metadata
Metadata
Assignees
Labels
area:coremlxcel-core: MLX FFI, primitives, KV cache, layersmlxcel-core: MLX FFI, primitives, KV cache, layerspriority:lowLow priorityLow prioritystatus:in-progressCurrently being worked onCurrently being worked ontype:bugBug fixes, error corrections, or issue resolutionsBug fixes, error corrections, or issue resolutions
Activity