diff --git a/crates/spark-model/src/mistral_loader/loader_impl/phase_assemble.rs b/crates/spark-model/src/mistral_loader/loader_impl/phase_assemble.rs index fb667f760..1b9bf1980 100644 --- a/crates/spark-model/src/mistral_loader/loader_impl/phase_assemble.rs +++ b/crates/spark-model/src/mistral_loader/loader_impl/phase_assemble.rs @@ -39,6 +39,7 @@ pub(super) fn assemble_layer( weight_scale: spark_runtime::gpu::DevicePtr::NULL, weight_scale_2: 0.0, input_scale: spark_runtime::gpu::DevicePtr::NULL, + weight_scale_2_vec: spark_runtime::gpu::DevicePtr::NULL, }; let attn = AttentionWeights { q_proj: null, diff --git a/crates/spark-model/src/tp_shard/quant_shard.rs b/crates/spark-model/src/tp_shard/quant_shard.rs index 7a0f37b74..84b0d20c4 100644 --- a/crates/spark-model/src/tp_shard/quant_shard.rs +++ b/crates/spark-model/src/tp_shard/quant_shard.rs @@ -85,6 +85,7 @@ pub fn shard_quantized_nvfp4( weight_scale: s_dst, weight_scale_2: src.weight_scale_2, input_scale: src.input_scale, + weight_scale_2_vec: src.weight_scale_2_vec, }) } TpShardKind::RowParallel => { @@ -132,6 +133,7 @@ pub fn shard_quantized_nvfp4( weight_scale: s_dst, weight_scale_2: src.weight_scale_2, input_scale: src.input_scale, + weight_scale_2_vec: src.weight_scale_2_vec, }) } } diff --git a/crates/spark-model/src/weight_loader/qwen35_dense.rs b/crates/spark-model/src/weight_loader/qwen35_dense.rs index e64d554d2..282981176 100644 --- a/crates/spark-model/src/weight_loader/qwen35_dense.rs +++ b/crates/spark-model/src/weight_loader/qwen35_dense.rs @@ -288,7 +288,37 @@ impl ModelWeightLoader for Qwen35DenseWeightLoader { DenseWeight, crate::weight_map::QuantizedWeight, )> { - let src = dense_auto(store, &format!("{p}.{name}.weight"), gpu)?; + // Pre-quantized Standard NVFP4 (e.g. sakamakismile): weight is U8 + // on disk. Load directly as QuantizedWeight without BF16 roundtrip. + // TP sharding of pre-quantized NVFP4 is not yet supported: enforce + // tp_size=1 explicitly, or every rank silently loads the full + // unsharded weight (duplicated, not sharded — wrong results with + // no error). + let weight_key = format!("{p}.{name}.weight"); + if matches!( + store.get(&weight_key).map(|w| w.dtype), + Ok(WeightDtype::UInt8) + ) { + anyhow::ensure!( + tp_size == 1, + "pre-quantized NVFP4 weight '{weight_key}' (U8 on disk) \ + cannot be loaded under tensor parallelism (tp_size={tp_size}): \ + TP sharding of pre-quantized NVFP4 checkpoints is not yet \ + implemented. Use tp_size=1, or dequantize this checkpoint to \ + BF16 first so it goes through the shard-then-requantize path." + ); + let null_dense = DenseWeight { + weight: spark_runtime::gpu::DevicePtr::NULL, + }; + let qw = quantized_auto( + store, + &format!("{p}.{name}"), + gpu, + Nvfp4Variant::Standard, + )?; + return Ok((null_dense, qw)); + } + let src = dense_auto(store, &weight_key, gpu)?; let (sharded_ptr, local_n, local_k) = shard_dense_bf16( src.weight, full_n, full_k, kind, tp_rank, tp_size, gpu, )?; @@ -425,6 +455,14 @@ impl ModelWeightLoader for Qwen35DenseWeightLoader { |name: &str, rows: usize, cols: usize| -> Result { if store.contains(&format!("{name}.weight_packed")) { dequant_nvfp4_to_bf16(store, name, rows, cols, gpu) + } else if matches!( + store.get(&format!("{name}.weight")).map(|w| w.dtype), + Ok(WeightDtype::UInt8) + ) { + // Standard-convention NVFP4 (packed bytes at + // `.weight`, not `.weight_packed`) — same dequant, + // different on-disk key. + dequant_nvfp4_to_bf16(store, name, rows, cols, gpu) } else { dense_auto(store, &format!("{name}.weight"), gpu) } @@ -508,14 +546,11 @@ impl ModelWeightLoader for Qwen35DenseWeightLoader { continue; } - let qkv_dense = load_ssm_proj(&format!("{la}.in_proj_qkv"), qkv_rows, h)?; - let z_dense = load_ssm_proj(&format!("{la}.in_proj_z"), z_rows, h)?; - let out_proj_dense = load_ssm_proj(&format!("{la}.out_proj"), h, value_dim)?; - - // A, B are always BF16 - let in_proj_a = dense(store, &format!("{la}.in_proj_a.weight"))?; - let in_proj_b = dense(store, &format!("{la}.in_proj_b.weight"))?; - let conv1d = dense(store, &format!("{la}.conv1d.weight"))?; + // A, B, conv1d, A_log, dt_bias, norm are independent of the + // qkv/z/out_proj on-disk format below — load them once up + // front so both the native-NVFP4 fast path and the legacy + // dequant/requant path can share them. + // // A_log and dt_bias MUST be FP32 — consumer kernels in // `ssm_preprocess.cu` and `mamba2_ssm_decode.cu` declare // them `const float*`. Loading via `dense()` kept BF16 @@ -524,6 +559,14 @@ impl ModelWeightLoader for Qwen35DenseWeightLoader { // error amplification through GDR recurrence at long // context. The MoE sister loader (`ssm_qwen35.rs`) // already promotes these; dense was missing the mirror. + // + // in_proj_a/b: route through `load_ssm_proj` (not the raw + // `dense()` byte-reinterpret) so a Standard-NVFP4 A/B + // (U8-packed) checkpoint dequants correctly instead of + // being read as BF16 garbage. + let in_proj_a = load_ssm_proj(&format!("{la}.in_proj_a"), nv, h)?; + let in_proj_b = load_ssm_proj(&format!("{la}.in_proj_b"), nv, h)?; + let conv1d = dense(store, &format!("{la}.conv1d.weight"))?; let a_log = dense_keep_f32(store, &format!("{la}.A_log"), gpu)?; let dt_bias = dense_keep_f32(store, &format!("{la}.dt_bias"), gpu)?; // norm.weight: use `dense_f32_safe` (FP32-aware: detects @@ -531,6 +574,97 @@ impl ModelWeightLoader for Qwen35DenseWeightLoader { // bf16 passes through). Mirrors `weight_map/ssm_qwen35.rs` // MoE sister loader (backported here 2026-05-20). let norm = dense_f32_safe(store, &format!("{la}.norm.weight"), gpu)?; + let ba_dense = interleave_ba(&in_proj_a, &in_proj_b, nv, nk, h, gpu)?; + let qkvz_size = config.ssm_qkvz_size(); + + // Native Standard-NVFP4 GDN (pre-quantized checkpoint, e.g. + // sakamakismile): in_proj_qkv / in_proj_z / out_proj ship + // U8-packed NVFP4 directly on disk (`.weight` dtype UInt8, + // not `.weight_packed` — that's the compressed-tensors + // convention `load_ssm_proj` already dequants above). Load + // them straight into `QuantizedWeight` and concat on GPU, + // skipping the BF16-dequant→re-quantize roundtrip entirely + // (that roundtrip is what the FP8/BF16-opt-in paths above + // exist to avoid for FP8-native and BF16-preferring + // checkpoints; here there's no lossy step to avoid in the + // first place — the data is already NVFP4). Requires all + // three projections to be U8; a partial-U8 checkpoint + // falls through to the legacy path below, where the + // `load_ssm_proj` UInt8 branch added above still dequants + // each U8 tensor correctly on its own. + let native_nvfp4 = matches!( + store + .get(&format!("{la}.in_proj_qkv.weight")) + .map(|w| w.dtype), + Ok(WeightDtype::UInt8) + ) && matches!( + store + .get(&format!("{la}.in_proj_z.weight")) + .map(|w| w.dtype), + Ok(WeightDtype::UInt8) + ) && matches!( + store.get(&format!("{la}.out_proj.weight")).map(|w| w.dtype), + Ok(WeightDtype::UInt8) + ); + if native_nvfp4 { + let qkv_qw = quantized_auto( + store, + &format!("{la}.in_proj_qkv"), + gpu, + Nvfp4Variant::Standard, + )?; + let z_qw = quantized_auto( + store, + &format!("{la}.in_proj_z"), + gpu, + Nvfp4Variant::Standard, + )?; + let qkvz_nvfp4 = qkv_qw.concat_rows(&z_qw, qkv_rows, z_rows, h, gpu)?; + let qkvz_nvfp4_t = qkvz_nvfp4.transpose_for_gemm(gpu, qkvz_size, h)?; + + let out_proj_nvfp4 = quantized_auto( + store, + &format!("{la}.out_proj"), + gpu, + Nvfp4Variant::Standard, + )?; + let out_proj_nvfp4_t = + out_proj_nvfp4.transpose_for_gemm(gpu, h, value_dim)?; + + let ssm = SsmWeights { + in_proj_qkvz: DenseWeight { + weight: spark_runtime::gpu::DevicePtr::NULL, + }, + in_proj_ba: ba_dense, + conv1d, + a_log, + dt_bias, + norm, + out_proj: out_proj_nvfp4, + }; + let mut layer = Qwen3SsmLayer::new_sequential( + input_norm, + ssm, + post_attn_norm, + ffn, + Some(qkvz_nvfp4), + Some(qkvz_nvfp4_t), + Some(out_proj_nvfp4_t), + config, + gpu, + )?; + layer.predequant_for_prefill(gpu, config, stream)?; + tracing::info!( + "SSM[{lp}] native NVFP4 GDN: qkvz+out_proj loaded pre-quantized \ + (U8-packed on disk; no BF16 dequant/requant roundtrip)" + ); + layers.push(Box::new(layer)); + continue; + } + + let qkv_dense = load_ssm_proj(&format!("{la}.in_proj_qkv"), qkv_rows, h)?; + let z_dense = load_ssm_proj(&format!("{la}.in_proj_z"), z_rows, h)?; + let out_proj_dense = load_ssm_proj(&format!("{la}.out_proj"), h, value_dim)?; let qkvz_dense = gpu_concat_rows(&qkv_dense, qkv_rows, &z_dense, z_rows, h, gpu)?; @@ -539,10 +673,6 @@ impl ModelWeightLoader for Qwen35DenseWeightLoader { gpu.free(qkv_dense.weight)?; gpu.free(z_dense.weight)?; - let ba_dense = interleave_ba(&in_proj_a, &in_proj_b, nv, nk, h, gpu)?; - - let qkvz_size = config.ssm_qkvz_size(); - // GDN ≥FP8 precision policy (2026-07-04). The nvidia // Qwen3.6-27B-NVFP4 checkpoint ships GDN in_proj_qkv / // out_proj as native F8_E4M3 (modelopt sensitivity diff --git a/crates/spark-model/src/weight_loader/step3p7.rs b/crates/spark-model/src/weight_loader/step3p7.rs index 5ae471af2..0d0bfd37a 100644 --- a/crates/spark-model/src/weight_loader/step3p7.rs +++ b/crates/spark-model/src/weight_loader/step3p7.rs @@ -98,6 +98,7 @@ fn slice_fused_experts( } else { fused_input_scale.offset(e * input_scale_bytes_per_expert) }, + weight_scale_2_vec: DevicePtr::NULL, }) .collect() } diff --git a/crates/spark-model/src/weight_loader/step3p7/load_layers.rs b/crates/spark-model/src/weight_loader/step3p7/load_layers.rs index bff476d4e..bd0443fa6 100644 --- a/crates/spark-model/src/weight_loader/step3p7/load_layers.rs +++ b/crates/spark-model/src/weight_loader/step3p7/load_layers.rs @@ -193,6 +193,7 @@ fn load_moe_ffn( weight_scale, weight_scale_2, input_scale, + weight_scale_2_vec: DevicePtr::NULL, }) }; let gate_proj = load_expert_proj("gate_proj")?; diff --git a/crates/spark-model/src/weight_map/loaders_fp8.rs b/crates/spark-model/src/weight_map/loaders_fp8.rs index a83031617..d39556411 100644 --- a/crates/spark-model/src/weight_map/loaders_fp8.rs +++ b/crates/spark-model/src/weight_map/loaders_fp8.rs @@ -212,6 +212,7 @@ pub(crate) fn quantize_to_nvfp4( weight_scale: scale_buf, weight_scale_2: scale2, input_scale: DevicePtr::NULL, + weight_scale_2_vec: DevicePtr::NULL, }) } diff --git a/crates/spark-model/src/weight_map/model_a.rs b/crates/spark-model/src/weight_map/model_a.rs index 0efd4c620..c2e26acf9 100644 --- a/crates/spark-model/src/weight_map/model_a.rs +++ b/crates/spark-model/src/weight_map/model_a.rs @@ -115,6 +115,7 @@ pub(crate) fn quantized( } else { DevicePtr::NULL }, + weight_scale_2_vec: DevicePtr::NULL, }) } diff --git a/crates/spark-model/src/weight_map/moe.rs b/crates/spark-model/src/weight_map/moe.rs index f9b6245b3..086099a2d 100644 --- a/crates/spark-model/src/weight_map/moe.rs +++ b/crates/spark-model/src/weight_map/moe.rs @@ -56,6 +56,7 @@ impl MoeWeights { weight_scale: DevicePtr::NULL, weight_scale_2: 1.0, input_scale: DevicePtr::NULL, + weight_scale_2_vec: DevicePtr::NULL, }; let null_expert = ExpertWeight { gate_proj: null_quant, diff --git a/crates/spark-model/src/weight_map/quant_helpers.rs b/crates/spark-model/src/weight_map/quant_helpers.rs index 6f575739e..693495760 100644 --- a/crates/spark-model/src/weight_map/quant_helpers.rs +++ b/crates/spark-model/src/weight_map/quant_helpers.rs @@ -361,5 +361,6 @@ pub(crate) fn quantized_v2( // dynamically and never reads this field. Absent ⇒ NULL (W4A4/W4A8 // checkpoints still load their `input_global_scale` unchanged). input_scale: ptr(store, &format!("{prefix}.input_global_scale")).unwrap_or(DevicePtr::NULL), + weight_scale_2_vec: DevicePtr::NULL, }) } diff --git a/crates/spark-model/src/weight_map/quantized.rs b/crates/spark-model/src/weight_map/quantized.rs index b8c7b5e10..a0e6c6e6d 100644 --- a/crates/spark-model/src/weight_map/quantized.rs +++ b/crates/spark-model/src/weight_map/quantized.rs @@ -77,6 +77,10 @@ pub struct QuantizedWeight { pub weight_scale_2: f32, /// Input activation scale (FP32 on device, for FP8 activation path). pub input_scale: DevicePtr, + /// Per-row FP32 scale2 on device (`[N]` floats). When set, the `w4a16_gemv_prs` + /// kernel reads scale2 per output row instead of the scalar `weight_scale_2`, + /// eliminating precision loss from per-tensor absmax on outlier rows. + pub weight_scale_2_vec: DevicePtr, } impl QuantizedWeight { @@ -87,14 +91,83 @@ impl QuantizedWeight { weight_scale: DevicePtr::NULL, weight_scale_2: 0.0, input_scale: DevicePtr::NULL, + weight_scale_2_vec: DevicePtr::NULL, } } + /// Whether this weight has per-row scale2 (for PRS GEMV dispatch). + pub fn has_per_row_scale2(&self) -> bool { + self.weight_scale_2_vec != DevicePtr::NULL + } + /// Whether this weight points to NULL (remote expert placeholder). pub fn is_null(&self) -> bool { self.weight == DevicePtr::NULL } + /// Concatenate two NVFP4 weights by rows: `[N1, K/2]` + `[N2, K/2]` → `[N1+N2, K/2]`. + /// + /// Both weights MUST share the same `K` (input dimension) and the same scalar + /// `weight_scale_2`. The packed weight bytes and FP8 block scales are concatenated + /// on-GPU via `cuMemcpy`. + pub fn concat_rows( + &self, + other: &QuantizedWeight, + n1: usize, + n2: usize, + k: usize, + gpu: &dyn GpuBackend, + ) -> anyhow::Result { + // The concatenated weight carries a single scalar scale2 (self's) for + // ALL rows — a mismatched `other` would silently dequantize its rows + // with the wrong per-tensor scale. This bit-exact equality only holds + // for `Nvfp4Variant::Standard` (NVIDIA ModelOpt) checkpoints, whose + // convention is a single global per-tensor `weight_scale_2` scalar + // shared across every row of the tensor — so two tensors quantized + // together by the same run share the identical f32 bit pattern. + // Other conventions (e.g. compressed-tensors) may carry independent + // per-tensor scales even for logically concatenable projections. + anyhow::ensure!( + self.weight_scale_2 == other.weight_scale_2, + "concat_rows: weight_scale_2 mismatch (self={}, other={}) — both NVFP4 \ + tensors must share the same per-tensor scale to be concatenated. \ + This is expected for ModelOpt/Standard NVFP4 checkpoints (single \ + global per-tensor scale2); re-quantize with the ModelOpt/Standard \ + quantizer, or report which checkpoint/quantizer produced independent \ + per-tensor scales for these projections", + self.weight_scale_2, + other.weight_scale_2, + ); + const GROUP_SIZE: usize = 16; + let half_k = k / 2; + let num_groups = k / GROUP_SIZE; + + let total_n = n1 + n2; + let packed_size = total_n * half_k; + let scale_size = total_n * num_groups; + + let new_weight = gpu.alloc(packed_size)?; + let new_scale = gpu.alloc(scale_size)?; + + gpu.copy_d2d(self.weight, new_weight, n1 * half_k)?; + gpu.copy_d2d(other.weight, new_weight.offset(n1 * half_k), n2 * half_k)?; + + gpu.copy_d2d(self.weight_scale, new_scale, n1 * num_groups)?; + gpu.copy_d2d( + other.weight_scale, + new_scale.offset(n1 * num_groups), + n2 * num_groups, + )?; + + Ok(QuantizedWeight { + weight: new_weight, + weight_scale: new_scale, + weight_scale_2: self.weight_scale_2, + input_scale: DevicePtr::NULL, + weight_scale_2_vec: DevicePtr::NULL, + }) + } + /// Transpose weight layout from [N, K/2] to [K/2, N] for coalesced GEMM reads. /// /// Also transposes scale from [N, K/GROUP_SIZE] to [K/GROUP_SIZE, N]. @@ -141,6 +214,7 @@ impl QuantizedWeight { weight_scale: new_scale, weight_scale_2: self.weight_scale_2, input_scale: self.input_scale, + weight_scale_2_vec: self.weight_scale_2_vec, }) } diff --git a/crates/spark-model/src/weight_map/ssm_qwen35.rs b/crates/spark-model/src/weight_map/ssm_qwen35.rs index 245e56d11..c6c408dc8 100644 --- a/crates/spark-model/src/weight_map/ssm_qwen35.rs +++ b/crates/spark-model/src/weight_map/ssm_qwen35.rs @@ -419,6 +419,7 @@ pub(crate) fn load_moe_no_shared( weight_scale: alloc_zero(scale_sz)?, weight_scale_2: 0.0, input_scale: DevicePtr::NULL, + weight_scale_2_vec: DevicePtr::NULL, }) }; diff --git a/crates/spark-model/src/weight_map/ssm_qwen35_more.rs b/crates/spark-model/src/weight_map/ssm_qwen35_more.rs index d9890ee40..8678191b7 100644 --- a/crates/spark-model/src/weight_map/ssm_qwen35_more.rs +++ b/crates/spark-model/src/weight_map/ssm_qwen35_more.rs @@ -108,6 +108,7 @@ pub(crate) fn load_moe_minimax( weight_scale: alloc_zero(scale_sz)?, weight_scale_2: 0.0, input_scale: DevicePtr::NULL, + weight_scale_2_vec: DevicePtr::NULL, }) }; let shared_expert = ExpertWeight { @@ -310,6 +311,7 @@ pub(crate) fn load_moe_gemma4( weight_scale: alloc_zero(s)?, weight_scale_2: 0.0, input_scale: DevicePtr::NULL, + weight_scale_2_vec: DevicePtr::NULL, }) }; let shared_expert = ExpertWeight {