gam_gpu/
dictionary_score.rs

1//! Shape and memory planning for high-`K` dictionary score routing.
2//!
3//! Sparse SAE dictionary routers all have the same hot loop: score a minibatch
4//! of `n_rows` residual rows against `n_items` candidate atoms/blocks, keep a
5//! tiny online top-`s`, and never materialize the full `n_rows x n_items` score
6//! matrix. This module owns the reusable admission and tile-size invariants for
7//! that pattern. Domain crates still own their kernels and selection semantics.
8
9/// Minimum `n_rows * n_items` score elements before a cold device route is worth
10/// its launch and host/device transfer cost.
11pub const DEFAULT_DICTIONARY_SCORE_MIN_ELEMS: usize = 1 << 20;
12
13/// Maximum score elements per device launch. With `f32` scores this is 8 MiB,
14/// matching the library row-chunk target and keeping peak score memory bounded
15/// independent of dictionary width.
16pub const DEFAULT_DICTIONARY_SCORE_TILE_ELEMS: usize =
17    gam_runtime::resource::LIBRARY_ROW_CHUNK_TARGET_BYTES / std::mem::size_of::<f32>();
18
19/// Device admission and tile geometry for one minibatch-by-dictionary score
20/// route.
21#[derive(Clone, Copy, Debug, Eq, PartialEq)]
22pub struct DictionaryScoreRoutePlan {
23    /// Minibatch rows scored together.
24    pub n_rows: usize,
25    /// Candidate atoms/blocks scored for each row.
26    pub n_items: usize,
27    /// Dot-product width for one score.
28    pub feature_dim: usize,
29    /// Minimum `n_rows * n_items` elements required for device admission.
30    pub device_min_score_elems: usize,
31    /// Maximum `n_rows * tile_items` score elements held by one device launch.
32    pub max_tile_score_elems: usize,
33    /// Candidate items per launch tile.
34    pub tile_items: usize,
35    /// Number of candidate tiles covering `0..n_items`.
36    pub tile_count: usize,
37    /// True when the total route work is large enough to use the device.
38    pub device_admitted: bool,
39    /// Peak score-block bytes for a full tile.
40    pub peak_score_bytes: usize,
41    /// Lower-bound arithmetic for dispatch diagnostics: one multiply and one add
42    /// per `(row, item, feature)` score term.
43    pub dot_flops_lower_bound: u128,
44}
45
46impl DictionaryScoreRoutePlan {
47    /// Build a plan with explicit thresholds. The function is pure and
48    /// allocation-free so call sites can test routing decisions without a CUDA
49    /// runtime.
50    #[must_use]
51    pub fn with_limits(
52        n_rows: usize,
53        n_items: usize,
54        feature_dim: usize,
55        device_min_score_elems: usize,
56        max_tile_score_elems: usize,
57    ) -> Self {
58        let total_score_elems = n_rows.saturating_mul(n_items);
59        let nondegenerate = n_rows > 0 && n_items > 0 && feature_dim > 0;
60        let tile_items = if !nondegenerate {
61            0
62        } else {
63            (max_tile_score_elems / n_rows).clamp(1, n_items)
64        };
65        let tile_count = if tile_items == 0 {
66            0
67        } else {
68            n_items.div_ceil(tile_items)
69        };
70        let peak_tile_items = tile_items.min(n_items);
71        let peak_score_elems = n_rows.saturating_mul(peak_tile_items);
72        let dot_flops_lower_bound = 2u128
73            .saturating_mul(n_rows as u128)
74            .saturating_mul(n_items as u128)
75            .saturating_mul(feature_dim as u128);
76
77        Self {
78            n_rows,
79            n_items,
80            feature_dim,
81            device_min_score_elems,
82            max_tile_score_elems,
83            tile_items,
84            tile_count,
85            device_admitted: nondegenerate && total_score_elems >= device_min_score_elems,
86            peak_score_bytes: peak_score_elems.saturating_mul(std::mem::size_of::<f32>()),
87            dot_flops_lower_bound,
88        }
89    }
90
91    /// Build a plan with the library defaults used by sparse dictionary routers.
92    #[must_use]
93    pub fn default_for_shape(n_rows: usize, n_items: usize, feature_dim: usize) -> Self {
94        Self::with_limits(
95            n_rows,
96            n_items,
97            feature_dim,
98            DEFAULT_DICTIONARY_SCORE_MIN_ELEMS,
99            DEFAULT_DICTIONARY_SCORE_TILE_ELEMS,
100        )
101    }
102
103    /// True when the plan covers no route work.
104    #[must_use]
105    pub const fn is_degenerate(self) -> bool {
106        self.n_rows == 0 || self.n_items == 0 || self.feature_dim == 0
107    }
108}
109
110#[cfg(test)]
111mod tests {
112    use super::*;
113
114    #[test]
115    fn target_k32k_shape_is_admitted_and_memory_bounded() {
116        let plan = DictionaryScoreRoutePlan::default_for_shape(256, 32_768, 64);
117        assert!(plan.device_admitted);
118        assert_eq!(plan.tile_items, 8_192);
119        assert_eq!(plan.tile_count, 4);
120        // The tile is sized off the canonical library row-chunk target, so the
121        // expected peak is that target by name, not a transcription of it
122        // (#2704).
123        assert_eq!(
124            plan.peak_score_bytes,
125            gam_runtime::resource::LIBRARY_ROW_CHUNK_TARGET_BYTES
126        );
127        assert_eq!(
128            plan.dot_flops_lower_bound,
129            2u128 * 256u128 * 32_768u128 * 64u128
130        );
131    }
132
133    #[test]
134    fn peak_score_memory_does_not_grow_with_dictionary_width() {
135        let small = DictionaryScoreRoutePlan::default_for_shape(512, 4_096, 48);
136        let large = DictionaryScoreRoutePlan::default_for_shape(512, 131_072, 48);
137        assert_eq!(small.tile_items, large.tile_items);
138        assert_eq!(small.peak_score_bytes, large.peak_score_bytes);
139        assert!(large.tile_count > small.tile_count);
140    }
141
142    #[test]
143    fn sub_floor_and_degenerate_shapes_stay_on_host() {
144        let tiny = DictionaryScoreRoutePlan::default_for_shape(16, 1024, 64);
145        assert!(!tiny.device_admitted);
146        assert_eq!(tiny.tile_count, 1);
147
148        for plan in [
149            DictionaryScoreRoutePlan::default_for_shape(0, 1024, 64),
150            DictionaryScoreRoutePlan::default_for_shape(16, 0, 64),
151            DictionaryScoreRoutePlan::default_for_shape(16, 1024, 0),
152        ] {
153            assert!(plan.is_degenerate());
154            assert!(!plan.device_admitted);
155            assert_eq!(plan.tile_items, 0);
156            assert_eq!(plan.tile_count, 0);
157            assert_eq!(plan.peak_score_bytes, 0);
158        }
159    }
160
161    #[test]
162    fn tiny_tile_budget_still_makes_forward_progress() {
163        let plan = DictionaryScoreRoutePlan::with_limits(512, 1000, 32, 1, 7);
164        assert_eq!(plan.tile_items, 1);
165        assert_eq!(plan.tile_count, 1000);
166        assert_eq!(plan.peak_score_bytes, 512 * std::mem::size_of::<f32>());
167    }
168}
gam_gpu/dictionary_score.rs

gam_gpu/
dictionary_score.rs
