metrale_gpu_runtime/buffers/
moe_fp8_scratch.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: One arena-owned slab that every grouped FP8 MoE layer reuses in turn.
4//! Quantizers overwrite their live activation/scale ranges. The worklist builder
5//! overwrites the live entries and `total_tiles` on EVERY invocation, including
6//! empty routes; consumers never read the stale tail. No host synchronization
7//! or allocation is needed inside CUDA capture, and addresses survive replay.
8//!
9//! Owner: gpu-runtime.
10//! Invariants:
11//! - The regions are disjoint and 256-byte aligned.
12//! - `Layout::new(c, rows).bytes <= Layout::new(c, max_batch_tokens).bytes` for
13//!   every `rows <= max_batch_tokens`.
14
15use super::BufferArena;
16use crate::gpu::DevicePtr;
17use metrale_config::ModelConfig;
18
19pub(super) struct Layout {
20    pub(super) scale: usize,
21    pub(super) worklist: usize,
22    pub(super) counter: usize,
23    pub(super) small_worklist: usize,
24    pub(super) small_counter: usize,
25    pub(super) small_bytes: usize,
26    pub(super) bytes: usize,
27}
28
29impl Layout {
30    pub(super) fn new(c: &ModelConfig, rows: usize) -> Self {
31        if c.num_experts == 0 {
32            return Self {
33                scale: 0,
34                worklist: 0,
35                counter: 0,
36                small_worklist: 0,
37                small_counter: 0,
38                small_bytes: 0,
39                bytes: 0,
40            };
41        }
42        let expanded = rows * c.num_experts_per_tok;
43        let shared_k = c.hidden_size.max(c.shared_expert_intermediate_size);
44        let activation = (rows * shared_k).max(expanded * c.moe_intermediate_size);
45        let scales = (rows * shared_k.div_ceil(128))
46            .max(expanded * c.moe_intermediate_size.div_ceil(128))
47            * 4;
48        // 2026-09-25: Same PM4 geometry as the grouped kernel: M=128, N=64. Per-expert
49        // rounding is bounded by +num_experts; +1 retains the launcher's slack.
50        let items = (expanded.div_ceil(128) + c.num_experts + 1)
51            * c.hidden_size.max(c.moe_intermediate_size).div_ceil(64);
52        let align = |n: usize| n.div_ceil(256) * 256;
53        let scale = align(activation);
54        let worklist = scale + align(scales);
55        let counter = worklist + align(items * 8);
56        // 2026-09-25: Only E256 is qualified for the adaptive Hopper dispatch. One small
57        // tile per expert, N=64; the capacity comes from the dimensions, never a fixed KiB.
58        let small_bytes = if c.num_experts == 256 {
59            small_list_bytes(c.num_experts, c.hidden_size.max(c.moe_intermediate_size))
60                .expect("adaptive FP8 small worklist size overflow")
61        } else {
62            0
63        };
64        let small_worklist = if small_bytes > 0 {
65            align(counter + 4)
66        } else {
67            0
68        };
69        let small_counter = small_worklist + align(small_bytes);
70        Self {
71            scale,
72            worklist,
73            counter,
74            small_worklist,
75            small_counter,
76            small_bytes,
77            bytes: if small_bytes > 0 {
78                small_counter + 4
79            } else {
80                counter + 4
81            },
82        }
83    }
84}
85
86fn small_list_bytes(experts: usize, width: usize) -> Option<usize> {
87    experts.checked_mul(width.div_ceil(64))?.checked_mul(8)
88}
89
90/// 2026-09-25: Disjoint regions; shared/gate/up/down phases reuse these on one stream.
91pub struct MoeFp8Scratch {
92    pub activation: DevicePtr,
93    pub scales: DevicePtr,
94    pub worklist: DevicePtr,
95    pub total_tiles: DevicePtr,
96    pub worklist_bytes: usize,
97    pub small_worklist: DevicePtr,
98    pub small_total_tiles: DevicePtr,
99    pub small_worklist_bytes: usize,
100}
101
102impl BufferArena {
103    /// 2026-09-25: Fixed offsets for a given shape, without allocating or mutating ownership.
104    ///
105    /// # Errors
106    /// When `rows` exceeds the arena's batch, the model has no experts, the shape
107    /// needs more than the slab holds, or the slab was released.
108    pub fn moe_fp8_scratch(&self, c: &ModelConfig, rows: usize) -> anyhow::Result<MoeFp8Scratch> {
109        let l = Layout::new(c, rows);
110        anyhow::ensure!(
111            rows <= self.max_batch_tokens
112                && l.bytes > 0
113                && l.bytes <= self.sizes.moe_fp8_scratch
114                && self.moe_fp8_scratch != DevicePtr::NULL,
115            "FP8 MoE scratch exceeds arena: rows={rows}, requested={} bytes, capacity={} bytes",
116            l.bytes,
117            self.sizes.moe_fp8_scratch
118        );
119        Ok(MoeFp8Scratch {
120            activation: self.moe_fp8_scratch,
121            scales: self.moe_fp8_scratch.offset(l.scale),
122            worklist: self.moe_fp8_scratch.offset(l.worklist),
123            total_tiles: self.moe_fp8_scratch.offset(l.counter),
124            worklist_bytes: l.counter - l.worklist,
125            small_worklist: if l.small_bytes > 0 {
126                self.moe_fp8_scratch.offset(l.small_worklist)
127            } else {
128                DevicePtr::NULL
129            },
130            small_total_tiles: if l.small_bytes > 0 {
131                self.moe_fp8_scratch.offset(l.small_counter)
132            } else {
133                DevicePtr::NULL
134            },
135            small_worklist_bytes: l.small_bytes,
136        })
137    }
138}
139
140#[cfg(test)]
141#[path = "moe_fp8_scratch_tests.rs"]
142mod tests;
metrale_gpu_runtime/buffers/moe_fp8_scratch.rs

metrale_gpu_runtime/buffers/
moe_fp8_scratch.rs
