From 6a482105cadfd79687e3386dc67bcc224c07948b Mon Sep 17 00:00:00 2001 From: Richard Safier Date: Sat, 4 Jul 2026 23:34:57 -0400 Subject: [PATCH] fix(loader): gate CUTLASS grouped NVFP4 MoE on architecture, not model_type string MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The fast CUTLASS grouped NVFP4 MoE path (opt-in via ATLAS_HOLO_LOW_MEMORY_MOE=1) was gated on `config.model_type == "holo3_1_moe"`, silently dropping every other NVFP4-MoE checkpoint to the slow `moe_w4a16` path. But holo3_1_moe / qwen3_5_moe / qwen3_6_moe are the SAME architecture at the weight level (factory:: loader_for_config maps all three to Qwen35WeightLoader); "holo3_1_moe" is a label fabricated at config/dispatch.rs purely for kernel-target directory lookup. The grouped path operates on the loaded NVFP4 experts regardless of source quant format (the inline comment already said so). Repro: Qwen-AgentWorld-35B-A3B-NVFP4 (model_type qwen3_5_moe, compressed-tensors NVFP4) fell back to moe_w4a16 even with LOW_MEMORY_MOE=1. Fix — gate on the architecture (MoE + NVFP4), not the label: - holo_nvfp4_moe -> nvfp4_moe = `config.num_experts > 0 && quant_format == Nvfp4` - is_holo_modelopt_mixed_precision -> is_modelopt_mixed_precision: drop the `model_type == "holo3_1_moe"` conjunct, keep the quant-config test. No kernels move; the config/dispatch.rs model_type rewrites (kernel-dir resolution) are untouched. holo3_1_moe still qualifies (num_experts=256, Nvfp4) so Holo-3.1-35B is unchanged; the mixed-precision path still fires for Holo and now also for any architecturally-identical sibling shipping the same quant. --- .../src/weight_loader/qwen35/load_layers.rs | 34 +++++++++++-------- 1 file changed, 20 insertions(+), 14 deletions(-) diff --git a/crates/spark-model/src/weight_loader/qwen35/load_layers.rs b/crates/spark-model/src/weight_loader/qwen35/load_layers.rs index c221fdd12..53fa16765 100644 --- a/crates/spark-model/src/weight_loader/qwen35/load_layers.rs +++ b/crates/spark-model/src/weight_loader/qwen35/load_layers.rs @@ -95,7 +95,7 @@ pub(super) fn load_layers( // Resolve runtime quantization format from the detected on-disk // variant. This determines which kernels are used for // decode/prefill/verify. - let modelopt_mixed_precision = is_holo_modelopt_mixed_precision(config); + let modelopt_mixed_precision = is_modelopt_mixed_precision(config); let quant_format = if variant == Nvfp4Variant::Fp8Dequanted { QuantFormat::Fp8 } else { @@ -103,14 +103,17 @@ pub(super) fn load_layers( }; let native_fp8 = quant_format == QuantFormat::Fp8; // FAST_MOE=full (transposed prefill tables + the CUTLASS grouped NVFP4 MoE) - // applies to any Holo-family (holo3_1_moe) NVFP4 MoE — both modelopt - // MIXED_PRECISION (Holo-3.1) and uniform NVFP4 (e.g. compressed-tensors, - // Ornith-35B). The transpose/grouped path operates on the loaded NVFP4 experts - // regardless of source quant format; the mixed-precision-specific handling - // (fp8 experts, native-modelopt SSM/attn) stays gated on - // `modelopt_mixed_precision` and simply doesn't fire for the uniform case. - let holo_nvfp4_moe = config.model_type == "holo3_1_moe" && quant_format == QuantFormat::Nvfp4; - let low_memory_modelopt_moe = (modelopt_mixed_precision || holo_nvfp4_moe) + // applies to ANY NVFP4 MoE: the transpose/grouped path operates on the loaded + // NVFP4 experts regardless of source quant format (modelopt MIXED_PRECISION, + // compressed-tensors, uniform NVFP4). Keyed on the ARCHITECTURE (MoE + NVFP4), + // not a model_type string — holo3_1_moe / qwen3_5_moe / qwen3_6_moe are the + // same arch at the weight level (factory::loader_for_config maps all three to + // Qwen35WeightLoader), so gating on the label silently dropped sibling and + // third-party NVFP4-MoE checkpoints to the slow moe_w4a16 path. The + // mixed-precision-specific handling (fp8 experts, native-modelopt SSM/attn) + // stays gated on `modelopt_mixed_precision` and simply doesn't fire otherwise. + let nvfp4_moe = config.num_experts > 0 && quant_format == QuantFormat::Nvfp4; + let low_memory_modelopt_moe = (modelopt_mixed_precision || nvfp4_moe) && std::env::var("ATLAS_HOLO_LOW_MEMORY_MOE").ok().as_deref() == Some("1"); let holo_fast_moe_mode = if low_memory_modelopt_moe { holo_fast_moe_mode() @@ -987,9 +990,12 @@ fn parse_layer_ranges(spec: &str) -> Vec<(usize, usize)> { ranges } -fn is_holo_modelopt_mixed_precision(config: &ModelConfig) -> bool { - config.model_type == "holo3_1_moe" - && config.quantization_config.as_ref().is_some_and(|qc| { - qc.quant_method == "modelopt" && qc.quant_algo.eq_ignore_ascii_case("MIXED_PRECISION") - }) +/// Modelopt mixed-precision NVFP4 (fp8 experts + NVFP4 FFN). Per-checkpoint quant +/// property, NOT tied to a model_type label — keyed purely on the quant config so +/// any Qwen3.5/3.6-family MoE shipping modelopt MIXED_PRECISION is recognized (was +/// previously also gated on model_type == "holo3_1_moe", excluding siblings). +fn is_modelopt_mixed_precision(config: &ModelConfig) -> bool { + config.quantization_config.as_ref().is_some_and(|qc| { + qc.quant_method == "modelopt" && qc.quant_algo.eq_ignore_ascii_case("MIXED_PRECISION") + }) }