feat(loader): native U8 NVFP4 loading for pre-quantized checkpoints

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,

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,
             })
         }
     }

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<DenseWeight> {
                             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,13 +559,112 @@ 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
                     // a fp32 checkpoint and truncates to BF16 with logging;
                     // 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

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()
 }

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")?;

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,
     })
 }
 

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,
     })
 }
 

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,

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,
     })
 }

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<QuantizedWeight> {
+        // 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,
         })
     }
 

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,
         })
     };
 

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 {