spark_model/layers/qwen3_ssm/
init_fp8.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! FP8 weight-install setters and the NVFP4→FP8 prefill pre-dequant for
4//! [`Qwen3SsmLayer`]. Split out of `init.rs` (500-LoC cap).
5
6use anyhow::Result;
7use spark_runtime::gpu::{DevicePtr, GpuBackend};
8
9use super::Qwen3SsmLayer;
10use crate::weight_map::Fp8Weight;
11
12impl Qwen3SsmLayer {
13    /// Install native FP8 block-scaled weights for the decode GEMV path.
14    ///
15    /// Inputs MUST be tagged `WeightQuantFormat::Fp8BlockScaled` — that is
16    /// the canonical input format for the `w8a16_gemv` kernel
17    /// (`out[n] = sum_k A[k] * E4M3_LUT[B[n,k]] * block_scale[n/BS, k/BS]`,
18    /// see `kernels/gb10/common/w8a16_gemv.cu`). The kernel reads the
19    /// scale buffer at `[N/BS, K/BS]` BF16 — exactly the shape produced
20    /// by `load_fp8_block_scaled_as_fp8weight`.
21    ///
22    /// This setter does NOT install the raw FP8 DevicePtr fields used by
23    /// the prefill `fp8_gemm_n128` kernel — that kernel takes no scale
24    /// argument and assumes single-scale FP8 (baked-in scale) produced
25    /// by `bf16_to_fp8`. Block-scaled bytes would silently produce wrong
26    /// outputs there. For prefill, call `set_fp8_prefill_only_weights`
27    /// separately with single-scale FP8 derived from a BF16 dequant.
28    pub fn set_fp8_decode_weights(&mut self, qkvz: Option<Fp8Weight>, out_proj: Option<Fp8Weight>) {
29        if let Some(ref w) = qkvz {
30            w.scale_format.expect(
31                crate::weight_map::WeightQuantFormat::Fp8BlockScaled,
32                "set_fp8_decode_weights::qkvz (w8a16_gemv expects [N/BS,K/BS] BF16 block scales)",
33            );
34        }
35        if let Some(ref w) = out_proj {
36            w.scale_format.expect(
37                crate::weight_map::WeightQuantFormat::Fp8BlockScaled,
38                "set_fp8_decode_weights::out_proj (w8a16_gemv expects [N/BS,K/BS] BF16 block scales)",
39            );
40        }
41        self.qkvz_fp8w = qkvz;
42        self.out_proj_fp8w = out_proj;
43    }
44
45    /// Install PER-ROW FP8 weights for the row-wise cuBLASLt PREFILL arm
46    /// (`ATLAS_FP8_ROWWISE=1`, mixed-precision compressed-tensors
47    /// checkpoints). Decode is untouched and keeps the NVFP4 copy.
48    ///
49    /// The `Fp8PerRow` assertion is the mirror of `set_fp8_decode_weights`'s
50    /// `Fp8BlockScaled` one: each setter refuses the other's layout, so the
51    /// two FP8 shapes cannot cross into each other's kernels. That crossing
52    /// does not fault — the smaller buffer is read in-bounds — so an assert
53    /// is the only thing that catches it.
54    pub fn set_fp8_rowwise_prefill_weights(
55        &mut self,
56        qkvz: Option<Fp8Weight>,
57        out_proj: Option<Fp8Weight>,
58    ) {
59        for (w, what) in [(&qkvz, "qkvz"), (&out_proj, "out_proj")] {
60            if let Some(w) = w {
61                w.scale_format.expect(
62                    crate::weight_map::WeightQuantFormat::Fp8PerRow,
63                    "set_fp8_rowwise_prefill_weights (cuBLASLt row-wise expects [N] f32)",
64                );
65                let _ = what;
66            }
67        }
68        self.qkvz_fp8w_rowwise = qkvz;
69        self.out_proj_fp8w_rowwise = out_proj;
70    }
71
72    /// Set raw FP8 DevicePtrs for the prefill GEMM path ONLY (no decode GEMV
73    /// scale fields). Used by the Qwen3.6-27B-FP8 native-FP8 SSM prefill path:
74    /// the FP8 buffer here is a single-scale FP8 (BF16 → FP8 truncation; values
75    /// already in FP8 range) suitable for `fp8_gemm_n128`. Decode falls back to
76    /// the NVFP4/BF16 paths via the existing `qkvz_nvfp4*` fields. PCND:
77    /// caller decides whether to install — never set implicitly.
78    pub fn set_fp8_prefill_only_weights(
79        &mut self,
80        qkvz_fp8: Option<DevicePtr>,
81        out_proj_fp8: Option<DevicePtr>,
82    ) {
83        if qkvz_fp8.is_some() {
84            self.qkvz_fp8 = qkvz_fp8;
85        }
86        if out_proj_fp8.is_some() {
87            self.out_proj_fp8 = out_proj_fp8;
88        }
89    }
90
91    /// Pre-dequant NVFP4 → FP8 for QKVZ and out_proj transposed weights.
92    /// Eliminates per-inference dequant overhead in prefill GEMMs.
93    pub fn predequant_for_prefill(
94        &mut self,
95        gpu: &dyn GpuBackend,
96        config: &atlas_core::config::ModelConfig,
97        stream: u64,
98    ) -> Result<()> {
99        let predequant_k = gpu.kernel("w4a16", "predequant_nvfp4_to_fp8")?;
100        let h = config.hidden_size;
101        let qkvz_size = config.ssm_qkvz_size();
102        let value_dim = config.linear_num_value_heads * config.linear_value_head_dim;
103
104        // QKVZ FP8 predequant: tested at ISL=1019, FP8 is ~50% slower (1900µs vs 1228µs)
105        // because weight matrix [12288, 2048] is bandwidth-dominated at M=1024 — the 2×
106        // larger FP8 weights (25 MB vs 12.6 MB NVFP4) cost more than the dequant saves.
107        let _ = qkvz_size; // suppress unused warning
108        // Use NON-transposed out_proj (ssm.out_proj is [N, K/2] layout).
109        // predequant_nvfp4_to_fp8 assumes [N, K/2] input layout.
110        if self.out_proj_nvfp4_t.is_some() {
111            self.out_proj_fp8 = Some(self.ssm.out_proj.predequant_to_fp8(
112                gpu,
113                predequant_k,
114                h,
115                value_dim,
116                stream,
117            )?);
118        }
119        Ok(())
120    }
121}
spark_model/layers/qwen3_ssm/init_fp8.rs

spark_model/layers/qwen3_ssm/
init_fp8.rs
