spark_model/layers/glm5next_dsa/
attend.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! GLM-5.3 DSA attention — the launcher for the selected-index NoPE MLA paged decode.
4//!
5//! Consumes what [`super::select`] produced: a `[q_rows, out_width]` i32 row of token
6//! ids with `-1` holes. Gathers exactly those tokens from the paged FP8 latent cache.
7//!
8//! # Why this is not the masked kernel
9//!
10//! `dsa_mla_masked_attn` is an **oracle** (see [`super::MASKED_ATTN_MAX_KEYS`]). The
11//! production path gathers per row, which is what HF says the reference *cannot* do
12//! (`_supports_flash_attn = False`, "cannot be mapped to FA without a custom kernel that
13//! can select on a per indices bases per row") and what vLLM ships. The gather is
14//! **exactly equivalent**, not an approximation: the reference mask is pure set
15//! membership, with duplicates collapsed and no additive weighting.
16//!
17//! # 🪤 NoPE, and why no `common/` kernel would do
18//!
19//! GLM-5.3 has `qk_rope_head_dim == 0`: the latent **is** the whole cache token.
20//! `common/mla_paged_decode.cu` declares `kv_cache_dim` and never reads it (its strides
21//! come from `#define ROPE_DIM 64`); `common/mla_paged_decode_fp8.cu` uses the runtime
22//! stride but then overwrites dims 448–511 with rope taken from the *next* token. Both
23//! fail silently. Hence a GLM-target kernel with no rope arm at all.
24
25use anyhow::{Result, bail};
26use spark_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
27use spark_runtime::kernel_args::KernelLaunch;
28
29use super::{Glm5NextDsaConfig, select::DsaSelectGeometry};
30
31/// Module name the DSA decode kernel resolves from — an unlisted `.cu` takes its file
32/// stem, and this one lives in the `glm-5.3-flash` target, not `common/`.
33pub const DSA_DECODE_MODULE: &str = "glm5next_dsa_mla_decode";
34
35/// Threads per block: `NUM_WARPS * WARP_SIZE` in the kernel.
36const DECODE_BLOCK: u32 = 256;
37
38/// The selected-index MLA decode entry point.
39#[derive(Clone, Copy)]
40pub struct Glm5NextDsaDecodeKernel(KernelHandle);
41
42impl Glm5NextDsaDecodeKernel {
43    /// Resolved with `kernel()`, never `try_kernel`: a missing sparse decode entry point
44    /// must be a hard error. Falling back to a dense path would be a correctness bug
45    /// wearing a performance bug's clothes.
46    pub fn resolve(gpu: &dyn GpuBackend) -> Result<Self> {
47        Ok(Self(gpu.kernel(
48            DSA_DECODE_MODULE,
49            "glm5next_dsa_mla_decode_fp8",
50        )?))
51    }
52}
53
54/// Everything the decode reads, all caller-owned.
55#[derive(Debug, Clone, Copy)]
56pub struct DsaDecodeInputs {
57    /// `[num_q_heads * kv_lora_rank]` BF16 — this rank's absorbed queries.
58    pub q: DevicePtr,
59    /// FP8 paged latent cache. In absorbed NoPE MLA K and V are the **same** buffer;
60    /// both are taken so a caller that splits them is not forced to lie.
61    pub k_cache: DevicePtr,
62    pub v_cache: DevicePtr,
63    /// `[num_q_heads * kv_lora_rank]` BF16 output.
64    pub out: DevicePtr,
65    /// `[num_seqs, max_blocks_per_seq]` i32.
66    pub block_tables: DevicePtr,
67    /// `[num_seqs]` i32.
68    pub seq_lens: DevicePtr,
69    /// `[num_seqs, sel_width]` i32 — [`super::select::DsaSelectScratch::tokens`].
70    pub sel_indices: DevicePtr,
71    pub k_scale: f32,
72    pub v_scale: f32,
73}
74
75/// Paging geometry the decode needs and the selection does not.
76#[derive(Debug, Clone, Copy, PartialEq, Eq)]
77pub struct DsaDecodePaging {
78    pub num_seqs: usize,
79    pub num_q_heads: usize,
80    pub num_kv_heads: usize,
81    pub max_blocks_per_seq: usize,
82    pub block_size: usize,
83    pub cache_stride_bytes: u64,
84}
85
86impl DsaDecodePaging {
87    /// 🪤 `block_size % index_kpool == 0` is required by the *selector*, not this kernel:
88    /// pools are built over absolute positions, so a block that straddles a pool boundary
89    /// makes a pool's tokens span two pages. The gather itself is per-token and would not
90    /// notice — which is exactly why the check belongs here rather than nowhere.
91    pub fn validate(&self, cfg: &Glm5NextDsaConfig) -> Result<()> {
92        if self.num_seqs == 0 || self.num_q_heads == 0 {
93            bail!(
94                "DSA decode: degenerate launch ({} seqs, {} heads)",
95                self.num_seqs,
96                self.num_q_heads
97            );
98        }
99        if self.block_size == 0 {
100            bail!("DSA decode: block_size must be > 0");
101        }
102        if !self.block_size.is_multiple_of(cfg.index_kpool) {
103            bail!(
104                "DSA decode: block_size {} is not a multiple of index_kpool {} — a pool \
105                 would straddle a page boundary",
106                self.block_size,
107                cfg.index_kpool
108            );
109        }
110        if self.num_kv_heads != 1 {
111            bail!(
112                "DSA decode: MLA carries a single latent KV head, got {}",
113                self.num_kv_heads
114            );
115        }
116        Ok(())
117    }
118}
119
120/// Launch the selected-index decode. Enqueued on `stream`, not synchronised.
121///
122/// `q_rows` in `geom` must equal `paging.num_seqs`: this is the decode path, one query
123/// row per sequence. A mismatch would index the selection rows with the wrong stride,
124/// so it is refused rather than trusted.
125pub fn decode_attention(
126    gpu: &dyn GpuBackend,
127    kernel: Glm5NextDsaDecodeKernel,
128    cfg: &Glm5NextDsaConfig,
129    geom: &DsaSelectGeometry,
130    paging: &DsaDecodePaging,
131    inputs: &DsaDecodeInputs,
132    stream: u64,
133) -> Result<()> {
134    paging.validate(cfg)?;
135    if geom.q_rows != paging.num_seqs {
136        bail!(
137            "DSA decode: selection has {} query rows but {} sequences are being decoded; \
138             the selection row stride would be wrong",
139            geom.q_rows,
140            paging.num_seqs
141        );
142    }
143
144    // The kernel tiles 512 latent dims across 32 lanes at 16 each. The host guard for
145    // this is `Glm5NextDsaConfig::validate` (KERNEL_KV_LORA_DIM); restated at the launch
146    // because a mismatch here is silent memory corruption, not an error.
147    if cfg.kv_lora_rank != super::KERNEL_KV_LORA_DIM {
148        bail!(
149            "DSA decode: kv_lora_rank {} != kernel tiling {}",
150            cfg.kv_lora_rank,
151            super::KERNEL_KV_LORA_DIM
152        );
153    }
154
155    KernelLaunch::new(gpu, kernel.0)
156        .grid([paging.num_q_heads as u32, paging.num_seqs as u32, 1])
157        .block([DECODE_BLOCK, 1, 1])
158        .arg_ptr(inputs.q)
159        .arg_ptr(inputs.k_cache)
160        .arg_ptr(inputs.v_cache)
161        .arg_ptr(inputs.out)
162        .arg_ptr(inputs.block_tables)
163        .arg_ptr(inputs.seq_lens)
164        .arg_ptr(inputs.sel_indices)
165        .arg_u32(geom.out_width as u32)
166        .arg_u32(paging.max_blocks_per_seq as u32)
167        .arg_u32(paging.num_q_heads as u32)
168        .arg_u32(paging.num_kv_heads as u32)
169        .arg_u32(cfg.kv_lora_rank as u32)
170        .arg_u32(paging.block_size as u32)
171        // 🪤 NoPE: the score scale is over the latent width, which IS the whole cache
172        // token. DeepSeek-V4 divides by sqrt(576) because its token carries a rope tail.
173        .arg_f32((cfg.kv_lora_rank as f32).powf(-0.5))
174        .arg_f32(inputs.k_scale)
175        .arg_f32(inputs.v_scale)
176        .arg_u64(paging.cache_stride_bytes)
177        .launch(stream)?;
178    Ok(())
179}
180
181#[cfg(test)]
182mod tests;
spark_model/layers/glm5next_dsa/attend.rs

spark_model/layers/glm5next_dsa/
attend.rs
