spark_model/layers/qwen3_ssm/

init.rs

// SPDX-License-Identifier: AGPL-3.0-only

//! Qwen3SsmLayer constructors + setters.

use super::*;

impl Qwen3SsmLayer {
    pub fn new(
        input_norm: DenseWeight,
        ssm: SsmWeights,
        post_attn_norm: DenseWeight,
        ffn: FfnComponent,
        qkvz_nvfp4: Option<QuantizedWeight>,
        config: &atlas_core::config::ModelConfig,
        gpu: &dyn GpuBackend,
    ) -> Result<Self> {
        let nv = config.linear_num_value_heads;
        let vd = config.linear_value_head_dim;
        let nk = config.linear_num_key_heads;
        let kd = config.linear_key_head_dim;
        let d_conv = config.linear_conv_kernel_dim;

        // conv_dim = Q_flat + K_flat + V_flat = 2*key_dim + value_dim = 8192
        let conv_dim = nk * kd * 2 + nv * vd;

        // Resolved BEFORE the struct literal because two fields need it: the
        // tensor-core spine's handle IS the bit that decides which spine the
        // prefill launches, so the scalar spine's route line has to read it,
        // and a field initializer cannot read a sibling field.
        let gdn_tc_spine = gdn_prefill_tc_kernel(gpu);

        Ok(Self {
            // mHC is attached later by the loader, and only for models that
            // carry a hc_mult-wide residual highway. The handles are gated on
            // the same condition `ArchProbes` uses, so a plain GDN model
            // never issues the lookup.
            hc: None,
            ple: None,
            hc_pre_k: hc_kernel(config, gpu, "hc_pre"),
            hc_post_k: hc_kernel(config, gpu, "hc_post"),
            hc_expand_k: hc_kernel(config, gpu, "hc_expand"),
            input_norm,
            ssm,
            post_attn_norm,
            ffn,
            lora_out_proj: None,
            qkvz_nvfp4,
            qkvz_nvfp4_t: None,
            out_proj_nvfp4_t: None,
            out_proj_dense: None,
            qkvz_fp8w: None,
            out_proj_fp8w: None,
            qkvz_fp8w_rowwise: None,
            out_proj_fp8w_rowwise: None,
            qkvz_rowwise_bf16: std::sync::atomic::AtomicU64::new(0),
            out_proj_rowwise_bf16: std::sync::atomic::AtomicU64::new(0),
            qkvz_q2: None,
            q2_0_gemv_k: super::super::try_kernel(gpu, "q2_0_gemv_vec", "q2_0_gemv_vec"),
            dequant_q2_0_gn_k: super::super::try_kernel(
                gpu,
                "dequant_gguf_bf16",
                "dequant_q2_0_gn_to_bf16",
            ),
            // The keep-packed MMQ family ships only in targets that serve
            // GGUF Q2 checkpoints; probing here would fail the boot audit on
            // every other GDN target. `set_packed_q2_qkvz` resolves them —
            // the only path that installs weights their dispatch sites check.
            q2_0_mmq_nc_k: KernelHandle(0),
            q2_0_mmq_wc_k: KernelHandle(0),
            q4k_quant_act_k: KernelHandle(0),
            sequential_qkvz: false,
            // Resolved ONCE here from the driver, then carried on the layer:
            // the projection dispatch asks "does this grid still fill the
            // machine?" and that question has no portable answer.
            sm_count: gpu.sm_count()?,
            rms_norm_residual_k: gpu.kernel("norm", "rms_norm_residual")?,
            // `output_gate_type: "sigmoid"` (qwen4_exp) swaps the gated-norm
            // handles for the sigmoid twins ONCE, here, so no forward call
            // site branches on it. Every other model keeps the SiLU originals.
            gated_rms_norm_k: if config.gdn_norm_sigmoid {
                gpu.kernel("gated_norm_sigmoid", "gated_rms_norm_sigmoid")?
            } else {
                gpu.kernel("norm", "gated_rms_norm")?
            },
            gated_rms_norm_f32_k: if config.gdn_norm_sigmoid {
                super::super::try_kernel(
                    gpu,
                    "gated_norm_sigmoid",
                    "gated_rms_norm_f32_input_sigmoid",
                )
            } else {
                super::super::try_kernel(gpu, "norm", "gated_rms_norm_f32_input")
            },
            // Only the SiLU (non-sigmoid) family has a strided twin today; a
            // `gdn_norm_sigmoid` model gets KernelHandle(0) and keeps the
            // per-sequence loop, which is correct, just launch-heavy.
            gated_rms_norm_f32_strided_k: if config.gdn_norm_sigmoid {
                KernelHandle(0)
            } else {
                super::super::try_kernel(gpu, "norm", "gated_rms_norm_f32_input_strided")
            },
            dense_gemv_k: gpu.kernel("gemv", "dense_gemv_bf16")?,
            dense_gemv_batch2_k: gpu.kernel("dense_gemv_bf16_batch2", "dense_gemv_bf16_batch2")?,
            w4a16_gemv_k: gpu.kernel("w4a16_gemv", "w4a16_gemv")?,
            w4a16_gemv_sw_k: super::super::try_kernel(gpu, "w4a16_gemv", "w4a16_gemv_sw"),
            w8a16_gemv_k: gpu.kernel("w8a16_gemv", "w8a16_gemv")?,
            w4a16_gemv_qkvz_k: gpu.kernel("w4a16_gemv", "w4a16_gemv_qkvz")?,
            deinterleave_k: gpu.kernel("ssm_preprocess", "deinterleave_qkvz")?,
            conv1d_k: gpu.kernel("causal_conv1d", "causal_conv1d_update")?,
            conv1d_l2norm_k: gpu.kernel("causal_conv1d", "causal_conv1d_update_l2norm")?,
            // FP32 conv1d output prevents BF16 truncation in the recurrent
            // path from compounding past ~8k tokens. The Metal backend
            // (kernels/metal/common/causal_conv1d_update_l2norm.metal) only
            // ships the BF16 variant; on those targets we fall back to the
            // BF16 kernel via the `.0 != 0` gate at the use site
            // (ssm_forward.rs). Warn instead of error: missing-on-Metal is
            // expected, and a startup `error!` would page on benign cases.
            // Strided twin of `conv1d_l2norm_f32_k` for the batched multi-seq
            // decode path. Optional: absent on older kernel sets, where the
            // multi-seq conv stays a per-sequence loop.
            conv1d_l2norm_f32_strided_k: super::super::try_kernel(
                gpu,
                "causal_conv1d",
                "causal_conv1d_update_l2norm_f32_strided",
            ),
            conv1d_l2norm_f32_k: {
                let h = super::super::try_kernel(
                    gpu,
                    "causal_conv1d",
                    "causal_conv1d_update_l2norm_f32",
                );
                if h.0 == 0 {
                    tracing::warn!(
                        "FP32 conv1d kernel not loaded; SSM uses BF16 conv \
                         output. Expect long-context coherence drift past ~8k \
                         tokens on this backend."
                    );
                }
                h
            },
            gdn_k: gpu.kernel("gated_delta_rule", "gated_delta_rule_decode")?,
            gdn_f32_k: super::super::try_kernel(
                gpu,
                "gated_delta_rule",
                "gated_delta_rule_decode_f32",
            ),
            gdn_f32_norm_k: super::super::try_kernel(
                gpu,
                "gated_delta_rule",
                "gated_delta_rule_decode_f32_norm",
            ),
            gdn_f32_conv_norm_k: super::super::try_kernel(
                gpu,
                "gated_delta_rule",
                "gated_delta_rule_decode_f32_conv_norm",
            ),
            gdn_f32_strided_k: super::super::try_kernel(
                gpu,
                "gated_delta_rule",
                "gated_delta_rule_decode_f32_strided",
            ),
            gdn_f32_strided_norm_k: super::super::try_kernel(
                gpu,
                "gated_delta_rule",
                "gated_delta_rule_decode_f32_strided_norm",
            ),
            gdn_f32_strided_norm_half_k: super::super::try_kernel(
                gpu,
                "gated_delta_rule",
                "gated_delta_rule_decode_f32_strided_norm_half",
            ),
            gdn_f32_strided_norm_smem_k: super::super::try_kernel(
                gpu,
                "gated_delta_rule",
                "gated_delta_rule_decode_f32_strided_norm_smem",
            ),
            gdn_f16_strided_norm_half_k: super::super::try_kernel(
                gpu,
                "gated_delta_rule",
                "gated_delta_rule_decode_f16_strided_norm_half",
            ),
            gdn_f16_norm_k: super::super::try_kernel(
                gpu,
                "gated_delta_rule",
                "gated_delta_rule_decode_f16_norm",
            ),
            ssm_h_f16_to_f32_k: super::super::try_kernel(
                gpu,
                "ssm_h_dtype",
                "ssm_h_state_f16_to_f32",
            ),
            ssm_h_f32_to_f16_k: super::super::try_kernel(
                gpu,
                "ssm_h_dtype",
                "ssm_h_state_f32_to_f16",
            ),
            ba_gates_k: gpu.kernel("ssm_preprocess", "dense_gemv_ba_gates")?,
            residual_add_k: gpu.kernel("residual_add", "bf16_residual_add")?,
            l2_norm_k: gpu.kernel("norm", "l2_norm_bf16")?,
            residual_add_rms_norm_k: gpu.kernel("norm", "residual_add_rms_norm")?,
            residual_add_rms_norm_gatef32_k: crate::layers::try_kernel(
                gpu,
                "norm",
                "residual_add_rms_norm_gatef32",
            ),
            gated_rms_norm_prefill_k: if config.gdn_norm_sigmoid {
                gpu.kernel("gated_norm_sigmoid", "gated_rms_norm_prefill_sigmoid")?
            } else {
                gpu.kernel("norm", "gated_rms_norm_prefill")?
            },
            w4a16_gemm_k: gpu.kernel("w4a16", "w4a16_gemm")?,
            w4a16_gemm_t_k: crate::layers::tgemm_kernel(gpu),
            w4a16_gemm_t_k64_k: crate::layers::k64_kernel(gpu)?,
            w4a16_gemm_t_k64_n64_k: crate::layers::k64_n64_kernel(gpu),
            w4a16_gemm_t_m128_k: gpu.kernel("w4a16", "w4a16_gemm_t_m128")?,
            // 8-warp pipelined M128 (try_kernel: 0 when absent → falls back to m128/n128).
            w4a16_gemm_t_m128_v2_k: super::super::w4a16_v2_kernel(gpu),
            w4a16_gemv_batch2_k: gpu.kernel("w4a16_gemv", "w4a16_gemv_batch2")?,
            dense_gemm_k: gpu.kernel("gemm", "dense_gemm_bf16")?,
            // try_kernel: 0-handle if absent (gated at dispatch); the pipelined
            // BF16 GEMM lives in the same `gemm` module as dense_gemm_bf16.
            dense_gemm_pipelined_k: super::super::try_kernel(
                gpu,
                "gemm",
                "dense_gemm_bf16_pipelined",
            ),
            gdn_prefill_k: gpu.kernel("gated_delta_rule", "gated_delta_rule_prefill")?,
            gdn_prefill_split_k: gpu
                .kernel("gated_delta_rule", "gated_delta_rule_prefill_split")?,
            gdn_prefill_split4_k: gpu
                .kernel("gated_delta_rule", "gated_delta_rule_prefill_split4")?,
            gdn_prefill_persistent_k: super::super::try_kernel(
                gpu,
                "gated_delta_rule_persistent",
                "gated_delta_rule_prefill_persistent",
            ),
            gdn_prefill_persistent_wy4_k: super::super::try_kernel(
                gpu,
                "gated_delta_rule_persistent",
                "gated_delta_rule_prefill_persistent_wy4",
            ),
            gdn_prefill_regresident_k: super::super::try_kernel(
                gpu,
                "gated_delta_rule_regresident",
                "gated_delta_rule_prefill_regresident",
            ),
            gdn_prefill_fla_recompute_wu_k: super::super::try_kernel(
                gpu,
                "gated_delta_rule_fla",
                "gated_delta_rule_recompute_wu",
            ),
            gdn_prefill_fla_recompute_wu_hopper_k: init_kernels::prefill_wu_hopper_k(gpu),
            gdn_prefill_fla_chunk_fwd_o_hopper_k: init_kernels::prefill_fwd_o_hopper_k(gpu),
            gdn_prefill_fla_chunk_delta_h_k: super::super::try_kernel(
                gpu,
                "gated_delta_rule_fla",
                "gated_delta_rule_chunk_delta_h_ksplit",
            ),
            gdn_prefill_fla_chunk_delta_h_tc_vblock_k: super::super::try_kernel(
                gpu,
                "gated_delta_rule_fla",
                "gated_delta_rule_chunk_delta_h_tc_vblock",
            ),
            gdn_prefill_fla_chunk_delta_h_tcfuse_k: gdn_tc_spine,
            // ONE handle for the scalar fused GDN state spine, and the route
            // line naming whichever spine the prefill will launch — see
            // `init_kernels::fused_spine_kernel` for both.
            gdn_prefill_fla_chunk_delta_h_fused_k: fused_spine_kernel(gpu, gdn_tc_spine),
            gdn_prefill_fla_chunk_delta_h_tma_k: super::super::try_kernel(
                gpu,
                "gated_delta_rule_fla",
                "gated_delta_rule_chunk_delta_h_tma",
            ),
            gdn_prefill_fla_chunk_fwd_o_k: super::super::try_kernel(
                gpu,
                "gated_delta_rule_fla",
                "gated_delta_rule_chunk_fwd_o",
            ),
            gdn_prefill_wy32_k: super::super::try_kernel(
                gpu,
                "gated_delta_rule_wy64_prefill",
                "gated_delta_rule_prefill_wy64",
            ),
            // ── Q12 Phase 2b: batched GDN kernel handles ──
            gdn_prefill_wy32_batched_k: super::super::try_kernel(
                gpu,
                "gated_delta_rule_wy64_prefill",
                "gated_delta_rule_prefill_wy64_batched",
            ),
            gdn_prefill_persistent_batched_k: super::super::try_kernel(
                gpu,
                "gated_delta_rule_persistent",
                "gated_delta_rule_prefill_persistent_batched",
            ),
            gdn_prefill_persistent_wy4_batched_k: super::super::try_kernel(
                gpu,
                "gated_delta_rule_persistent",
                "gated_delta_rule_prefill_persistent_wy4_batched",
            ),
            gdn_prefill_split4_batched_k: super::super::try_kernel(
                gpu,
                "gated_delta_rule",
                "gated_delta_rule_prefill_split4_batched",
            ),
            compute_gdn_gates_k: gpu.kernel("ssm_preprocess", "compute_gdn_gates")?,
            ba_gates_prefill_k: gpu.kernel("ssm_preprocess", "dense_gemm_ba_gates_prefill")?,
            ba_gates_prefill_hopper_k: init_kernels::ba_gates_hopper_k(gpu),
            conv1d_prefill_k: gpu.kernel("causal_conv1d", "causal_conv1d_update_prefill")?,
            conv1d_prefill_tp_k: super::super::try_kernel(
                gpu,
                "causal_conv1d",
                "causal_conv1d_update_prefill_tp",
            ),
            gdn_chunk2_k: gpu.kernel("gated_delta_rule", "gated_delta_rule_chunk2")?,
            conv1d_chunk2_k: gpu.kernel("causal_conv1d", "causal_conv1d_update_chunk2")?,
            gdn_chunk3_k: gpu.kernel("gated_delta_rule", "gated_delta_rule_chunk3")?,
            w4a16_gemv_batch3_k: gpu.kernel("w4a16_gemv", "w4a16_gemv_batch3")?,
            gdn_wy2_k: gpu.kernel("gated_delta_rule_wy", "gated_delta_rule_wy2")?,
            // Register-resident wy2 twin (own gb10-common module so non-gb10
            // targets simply resolve 0 and keep the base wy2 — try_kernel
            // misses are a silent handle 0, so `wy2_kernel` logs the
            // resolution outcome once at first K=2 dispatch).
            gdn_wy2_resident_k: super::super::try_kernel(
                gpu,
                "gated_delta_rule_wy2_resident",
                "gated_delta_rule_wy2_resident",
            ),
            gdn_wy3_k: gpu.kernel("gated_delta_rule_wy3", "gated_delta_rule_wy3")?,
            // Register-resident wy3 twin (same pattern as wy2's above:
            // try_kernel so non-gb10 targets resolve 0 and keep base wy3;
            // `wy3_kernel` logs the resolution outcome at first K=3 dispatch).
            gdn_wy3_resident_k: super::super::try_kernel(
                gpu,
                "gated_delta_rule_wy3_resident",
                "gated_delta_rule_wy3_resident",
            ),
            gdn_wy4_k: gpu.kernel("gated_delta_rule_wy4", "gated_delta_rule_wy4")?,
            // ── ATLAS_SSM_H_FP16 stage 2: FP16 h-state twins of the MTP
            // verify WY kernels. try_kernel for the same reason as the
            // resident twins above — a miss is a silent handle 0, and the
            // selectors gate on `.0 != 0` before ever picking one. Without
            // these the flag and `--speculative` are mutually exclusive,
            // because every WY kernel above writes the state as FP32 and an
            // FP32 kernel over an FP16 pool produces fluent garbage, not an
            // error. Preflight refuses the combination unless the K values the
            // configured draft count can reach all have a twin here.
            gdn_wy2_f16_k: super::super::try_kernel(
                gpu,
                "gated_delta_rule_wy_f16",
                "gated_delta_rule_wy2_f16",
            ),
            gdn_wy2_resident_f16_k: super::super::try_kernel(
                gpu,
                "gated_delta_rule_wy2_resident_f16",
                "gated_delta_rule_wy2_resident_f16",
            ),
            gdn_wy3_f16_k: super::super::try_kernel(
                gpu,
                "gated_delta_rule_wy3_f16",
                "gated_delta_rule_wy3_f16",
            ),
            gdn_wy3_resident_f16_k: super::super::try_kernel(
                gpu,
                "gated_delta_rule_wy3_resident_f16",
                "gated_delta_rule_wy3_resident_f16",
            ),
            gdn_wy4_f16_k: super::super::try_kernel(
                gpu,
                "gated_delta_rule_wy4_f16",
                "gated_delta_rule_wy4_f16",
            ),
            // STAGE 1 fused K=2 verify epilogue. Only present in the gb10
            // common PTX module set; NULL on targets lacking the .cu, in which
            // case the num_tokens==2 arm keeps the per-token path even when
            // ATLAS_GDN_FUSED_VERIFY is set.
            gdn_verify_fused_conv_k2_k: super::super::try_kernel(
                gpu,
                "gdn_verify_fused_k2",
                "gdn_verify_fused_conv_k2",
            ),
            gdn_verify_fused_norm_k2_k: super::super::try_kernel(
                gpu,
                "gdn_verify_fused_k2",
                "gdn_verify_fused_norm_k2",
            ),
            // Generic-K fused verify conv (K=17 DFlash arm). gb10 common
            // module; NULL on targets lacking the .cu, in which case the
            // K=17 arm keeps its per-token conv loop.
            gdn_verify_fused_conv_kn_k: super::super::try_kernel(
                gpu,
                "gdn_verify_fused_conv_kn",
                "gdn_verify_fused_conv_kn",
            ),
            // Batched twin (gridDim.y = n_seq) for batched speculative decoding.
            gdn_verify_fused_conv_kn_batched_k: super::super::try_kernel(
                gpu,
                "gdn_verify_fused_conv_kn",
                "gdn_verify_fused_conv_kn_batched",
            ),
            // Exact-verify `_snap` twins (#435): model-shadow staged
            // (qwen3.6-27b/nvfp4), 0 elsewhere — the exact arm then uses the
            // parent kernel + copy_d2d snapshots (same bits, more launches).
            // Every other GDN target declares these three lookups in its
            // MODEL.toml [expected_absent] (#438) — the boot gate fails closed.
            gdn_f32_norm_snap_k: super::super::try_kernel(
                gpu,
                "gated_delta_rule_snap",
                "gated_delta_rule_decode_f32_norm_snap",
            ),
            gdn_f32_strided_norm_snap_k: super::super::try_kernel(
                gpu,
                "gated_delta_rule_snap",
                "gated_delta_rule_decode_f32_strided_norm_snap",
            ),
            gdn_verify_fused_conv_kn_f32_k: super::super::try_kernel(
                gpu,
                "gdn_verify_fused_conv_kn_f32",
                "gdn_verify_fused_conv_kn_f32",
            ),
            // wy17 ships only in qwen3.6-35b-a3b's and qwen3.6-27b's PTX sets;
            // NULL elsewhere (declared [expected_absent] in those MODEL.tomls).
            // decode_batched(K=17) checks for non-NULL before dispatching the fused path.
            gdn_wy17_k: super::super::try_kernel(
                gpu,
                "gated_delta_rule_wy17",
                "gated_delta_rule_wy17",
            ),
            gdn_wyn_k: init_kernels::wyn_kernels(gpu),
            gdn_wyn_f16_k: init_kernels::wyn_f16_kernels(gpu),
            h_state_bytes: nv * vd * kd * 4, // FP32 [nv, kd, vd] transposed for coalescing
            conv_state_bytes: conv_dim * d_conv * 4, // FP32 [conv_dim, d_conv]
            qkvz_fp8: None,
            out_proj_fp8: None,
            fp8_gemm_k: gpu.kernel("w4a16", "fp8_gemm_t")?,
            fp8_gemm_t_m128_k: gpu.kernel("w4a16", "fp8_gemm_t_m128")?,
            w8a16_gemm_k: super::super::try_kernel(gpu, "w8a16_gemm", "w8a16_gemm"),
            w8a16_gemm_pipelined_k: super::super::try_kernel(
                gpu,
                "w8a16_gemm_pipelined",
                "w8a16_gemm_pipelined",
            ),
            w8a16_gemv_batch4_k: super::super::try_kernel(
                gpu,
                "w8a16_gemv_batch4",
                "w8a16_gemv_batch4",
            ),
            w8a16_gemv_batch16_k: super::super::try_kernel(
                gpu,
                "w8a16_gemv_batch4",
                "w8a16_gemv_batch16",
            ),
            // NVFP4 batched decode GEMV (all entries live in the w4a16_gemv module).
            w4a16_batchm: crate::layers::w4a16_gemv_tiers::W4a16BatchmTiers::resolve(gpu),
            w4a16_gemv_batch16_k: super::super::try_kernel(gpu, "w4a16_gemv", "w4a16_gemv_batch16"),
            w8a16_gemm_t_k: super::super::try_kernel(gpu, "w8a16_gemm_t", "w8a16_gemm_t"),
            per_token_group_quant_fp8_k: ops::Fp8ActQuant::resolve(gpu),
            fp8_gemm_t_blockscaled_k: super::super::try_kernel(
                gpu,
                "fp8_gemm_t_blockscaled",
                "fp8_gemm_t_blockscaled",
            ),
            fp8_act_scale_kmajor_k: super::super::try_kernel(
                gpu,
                "fp8_scale_transpose",
                "fp8_act_scale_to_kmajor",
            ),
        })
    }

    // `new_sequential` moved to `init_sequential.rs` (≤500 LoC split).
}

#[path = "init_kernels.rs"]
mod init_kernels;
use init_kernels::{fused_spine_kernel, gdn_prefill_tc_kernel, hc_kernel};

#[path = "init_sequential.rs"]
mod init_sequential;