metrale_model_arch/glm5next_dsa/layer/workspace.rs
1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-26: [`Glm5NextDsaWorkspace`], the per-layer scratch `decode_k` reuses, with the
4//! two switches that decide whether its batched-selector buffers exist and are used
5//! ([`dsa_select_rows_enabled`], [`batch_select_enabled`]).
6//!
7//! Owner: model-arch (GLM-5.3 DSA).
8//! Invariants:
9//! - The batched-selector buffers are NULL exactly when [`dsa_select_rows_enabled`] is off,
10//! and `bt`/`sl` exactly when `METRALE_GLM_DSA_ALLOC_PER_STEP=1`.
11
12use anyhow::Result;
13use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend};
14
15use super::super::Glm5NextDsaConfig;
16use super::super::select::DsaSelectScratch;
17
18/// 2026-09-25: Whether one selection pass covers all `k` rows (`select_rows_batched`) or each
19/// row selects on its own (`select_row`). All four terms must hold:
20///
21/// * `workspace_ready`: the `[max_rows]` selector buffers exist. [`dsa_select_rows_enabled`]
22/// off (`METRALE_DSA_SELECT_ROWS=0`) leaves them unallocated.
23/// * `is_prefill`: stated by the caller, because `ForwardContext` does not separate a
24/// prefill from a speculative verify: both have `decode_step` false, and `verify_a` sets
25/// `graph_capture` false.
26/// * `!graph_capture`: `select_rows_batched` copies `q_pos` host-to-device, which a
27/// capturing stream does not allow.
28/// * `k > 1`.
29///
30/// One pass over the group selects, for each row, the pools a single-row pass at that row's
31/// length would. `dsa_index_scores` makes pool `p` a candidate for row `r` only when the
32/// pool is complete and its last token is valid and at or before `q_pos[r]`, so the extra
33/// pools of the group's larger `P` score `-FLT_MAX` for row `r`; `dsa_topk_pools` orders by
34/// score, then pool index, so the top `select_k` is unique; and `dsa_expand_selection`
35/// clamps `select_k` to the row's own pool count, so the row's tail lands in the slots a
36/// single-row pass uses (the kernel's comment carries the measurement behind that clamp).
37pub(crate) fn batch_select_enabled(
38 workspace_ready: bool,
39 is_prefill: bool,
40 graph_capture: bool,
41 k: usize,
42) -> bool {
43 workspace_ready && is_prefill && !graph_capture && k > 1
44}
45
46/// 2026-09-25: `METRALE_DSA_SELECT_ROWS`: on unless set to `0`. When off,
47/// `Glm5NextDsaWorkspace::new` does not allocate the batched-selector buffers.
48pub(crate) fn dsa_select_rows_enabled() -> bool {
49 std::env::var("METRALE_DSA_SELECT_ROWS").as_deref() != Ok("0")
50}
51
52/// 2026-09-25: Scratch reused across `decode_k` calls; each layer owns one.
53pub struct Glm5NextDsaWorkspace {
54 pub(super) q_a: DevicePtr,
55 pub(super) q_resid: DevicePtr,
56 pub(super) q_abs: DevicePtr,
57 pub(super) kv_a: DevicePtr,
58 pub(super) q_idx: DevicePtr,
59 pub(super) head_weights: DevicePtr,
60 pub(super) q_pos: DevicePtr,
61 pub(super) q_mask: DevicePtr,
62 /// 2026-09-25: `[max_rows, ...]` copies of `q_idx` / `head_weights` / `q_pos` / `q_mask`
63 /// for the batched prefill selector. NULL when [`dsa_select_rows_enabled`] is off.
64 pub(super) q_idx_rows: DevicePtr,
65 pub(super) head_weights_rows: DevicePtr,
66 pub(super) q_pos_rows: DevicePtr,
67 pub(super) q_mask_rows: DevicePtr,
68 pub(super) slot: DevicePtr,
69 pub(super) attn_out: DevicePtr,
70 /// 2026-09-25: Block table for the paged gather, `bt_cap` i32 entries. NULL when
71 /// `METRALE_GLM_DSA_ALLOC_PER_STEP=1`; `decode_k` then allocates one per call.
72 pub(super) bt: DevicePtr,
73 /// 2026-09-25: `seq_lens` for the paged gather, `[max_rows]` i32; NULL under the same
74 /// setting as `bt`.
75 pub(super) sl: DevicePtr,
76 /// 2026-09-25: Capacity of `bt` in entries; `decode_k` refuses a larger upload.
77 pub(super) bt_cap: usize,
78 /// 2026-09-25: The largest `k` `decode_k` accepts.
79 pub(super) max_rows: usize,
80 /// 2026-09-25: `[index_head_dim]` BF16 staging rows at fixed addresses (`stage_k`,
81 /// `stage_gate`). On the replay-safe path the indexer projections write here and
82 /// `dsa_indexer_store` copies them to the row at a device-side position, which a
83 /// replayed graph reads live.
84 pub(super) stage_k: DevicePtr,
85 pub(super) stage_gate: DevicePtr,
86 /// 2026-09-25: `[5]` i32 selector geometry, written on the device by `dsa_write_geom`
87 /// for each row on the replay-safe path.
88 pub(super) geom_dev: DevicePtr,
89 pub(super) select: DsaSelectScratch,
90}
91
92impl Glm5NextDsaWorkspace {
93 /// 2026-09-25: `max_rows` (at least 1) is the largest `k` this workspace serves. Sized by
94 /// it: the projection outputs, `attn_out`, `sl`, the batched-selector buffers, and the
95 /// selection scratch, which is planned at [`super::super::state::max_dsa_context`] tokens and
96 /// `max_rows` query rows. The staging rows and `bt` do not depend on it.
97 pub fn new(gpu: &dyn GpuBackend, cfg: &Glm5NextDsaConfig, max_rows: usize) -> Result<Self> {
98 let rows = max_rows.max(1);
99 let geom = super::super::select::DsaSelectGeometry::plan(
100 cfg,
101 super::super::state::max_dsa_context(cfg),
102 rows,
103 )?;
104 let bt_cap = super::super::state::max_dsa_context(cfg).max(1);
105 let persist = std::env::var("METRALE_GLM_DSA_ALLOC_PER_STEP").as_deref() != Ok("1");
106 let batch_select = dsa_select_rows_enabled();
107 Ok(Self {
108 q_a: gpu.alloc(rows * (cfg.q_lora_rank * 2))?,
109 q_resid: gpu.alloc(rows * (cfg.q_lora_rank * 2))?,
110 q_abs: gpu.alloc(rows * (cfg.local_heads * cfg.kv_lora_rank * 2))?,
111 kv_a: gpu.alloc(rows * (cfg.kv_lora_rank * 2))?,
112 q_idx: gpu.alloc(cfg.index_heads * cfg.index_head_dim * 4)?,
113 head_weights: gpu.alloc(cfg.index_heads * 4)?,
114 q_pos: gpu.alloc(4)?,
115 q_idx_rows: if batch_select {
116 gpu.alloc(rows * cfg.index_heads * cfg.index_head_dim * 4)?
117 } else {
118 DevicePtr(0)
119 },
120 head_weights_rows: if batch_select {
121 gpu.alloc(rows * cfg.index_heads * 4)?
122 } else {
123 DevicePtr(0)
124 },
125 q_pos_rows: if batch_select {
126 gpu.alloc(rows * 4)?
127 } else {
128 DevicePtr(0)
129 },
130 q_mask_rows: if batch_select {
131 // 2026-09-25: All 1s (every row is a real query), set once here.
132 let p = gpu.alloc(rows)?;
133 gpu.memset_async(p, 1, rows, 0)?;
134 gpu.synchronize(0)?;
135 p
136 } else {
137 DevicePtr(0)
138 },
139 q_mask: {
140 // 2026-09-25: Always 1 (one real query), set once here so no step writes it.
141 let p = gpu.alloc(1)?;
142 gpu.memset_async(p, 1, 1, 0)?;
143 gpu.synchronize(0)?;
144 p
145 },
146 slot: gpu.alloc(8)?,
147 attn_out: gpu.alloc(rows * (cfg.local_heads * cfg.kv_lora_rank * 2))?,
148 bt: if persist {
149 gpu.alloc(bt_cap * 4)?
150 } else {
151 DevicePtr(0)
152 },
153 // 2026-09-25: `[rows]`: `attend_rows` launches all rows at once and
154 // `glm5next_dsa_mla_decode_fp8` reads `seq_lens[blockIdx.y]`.
155 sl: if persist {
156 gpu.alloc(rows * 4)?
157 } else {
158 DevicePtr(0)
159 },
160 bt_cap,
161 max_rows: rows,
162 stage_k: gpu.alloc(cfg.index_head_dim * 2)?,
163 stage_gate: gpu.alloc(cfg.index_head_dim * 2)?,
164 geom_dev: gpu.alloc(5 * 4)?,
165 select: DsaSelectScratch::alloc(gpu, cfg, &geom)?,
166 })
167 }
168}