Skip to content

fix(loader): gate CUTLASS grouped NVFP4 MoE on architecture, not model_type string #260

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
rsafier wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix(loader): gate CUTLASS grouped NVFP4 MoE on architecture, not model_type string #260

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
34 changes: 20 additions & 14 deletions crates/spark-model/src/weight_loader/qwen35/load_layers.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -95,22 +95,25 @@ 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 {
QuantFormat::Nvfp4
};
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()
Expand Down Expand Up @@ -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")
})
}
Loading