metrale_model_layers/layers/ops/
fp8_act_quant.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: The per-token FP8 activation quantizers a layer can launch, and
4//! the grid each one needs.
5//!
6//! Owner: model-layers ops.
7//! Invariants:
8//! - [`Fp8ActQuant::pick_with`] takes the kernel handle and the grid from the
9//!   same `twin` flag, so a pick never pairs one kernel with the other's grid.
10//!
11//! The shared kernel (`kernels/gb10/common/per_token_group_quant_fp8.cu`) runs
12//! one 128-thread CTA per 128-element K-group. The Hopper twin
13//! (`kernels/hopper/common/fp8_act_quant_hopper.cu`) gives each group 16
14//! threads that load one `uint4` each, and packs 8 groups into a CTA. Only
15//! `kernels/hopper` carries the twin; on other targets its handle is 0 and the
16//! shared kernel runs. `fp8_act_quant_floor.rs` decides when the twin runs.
17
18use metrale_gpu_runtime::gpu::{GpuBackend, KernelHandle};
19
20use super::Fp8QuantPick;
21
22/// 2026-09-25: Module of the shared quantizer; the file and the kernel share
23/// this name.
24pub const FP8_QUANT_MODULE: &str = "per_token_group_quant_fp8";
25/// 2026-09-25: Entry point of the shared quantizer.
26pub const FP8_QUANT_ENTRY: &str = "per_token_group_quant_fp8";
27/// 2026-09-25: Module (file stem) of the Hopper twin.
28pub const FP8_QUANT_HOPPER_MODULE: &str = "fp8_act_quant_hopper";
29/// 2026-09-25: Entry point of the Hopper twin. It differs from the shared
30/// entry point, so a Hopper image carries both kernels.
31pub const FP8_QUANT_HOPPER_ENTRY: &str = "per_token_group_quant_fp8_hopper";
32
33/// 2026-09-25: K-groups one Hopper CTA covers: 128 threads / 16 threads per
34/// group, as `HQ_GROUPS_PER_CTA` in the `.cu`. The kernel derives its group
35/// span from `gridDim.y`, so this value sets only the grid's Y extent.
36pub const FP8_QUANT_HOPPER_GROUPS_PER_CTA: u32 = 8;
37
38/// 2026-09-25: The FP8 activation quantizer a layer launches: the shared
39/// kernel, plus the Hopper twin when the image has one.
40#[derive(Clone, Copy, Debug)]
41pub struct Fp8ActQuant {
42    /// 2026-09-25: `per_token_group_quant_fp8`, or `KernelHandle(0)` when the
43    /// image lacks it.
44    pub shared: KernelHandle,
45    /// 2026-09-25: `per_token_group_quant_fp8_hopper`, or `KernelHandle(0)`
46    /// when the image lacks it.
47    pub hopper: KernelHandle,
48}
49
50impl Default for Fp8ActQuant {
51    /// 2026-09-25: No quantizer: both handles are 0.
52    fn default() -> Self {
53        Self {
54            shared: KernelHandle(0),
55            hopper: KernelHandle(0),
56        }
57    }
58}
59
60impl Fp8ActQuant {
61    /// 2026-09-25: Probe both kernels; a missing one resolves to
62    /// `KernelHandle(0)`.
63    pub fn resolve(gpu: &dyn GpuBackend) -> Self {
64        Self {
65            shared: crate::layers::try_kernel(gpu, FP8_QUANT_MODULE, FP8_QUANT_ENTRY),
66            hopper: crate::layers::try_target_kernel(
67                gpu,
68                FP8_QUANT_HOPPER_MODULE,
69                FP8_QUANT_HOPPER_ENTRY,
70            ),
71        }
72    }
73
74    /// 2026-09-25: A quantizer with only the shared kernel.
75    pub fn shared_only(shared: KernelHandle) -> Self {
76        Self {
77            shared,
78            hopper: KernelHandle(0),
79        }
80    }
81
82    /// 2026-09-25: Is there a quantizer to launch at all? True when either
83    /// handle is loaded, whatever the width; [`Self::pick`] is the per-launch
84    /// choice.
85    pub fn available(&self) -> bool {
86        self.shared.0 != 0 || self.hopper.0 != 0
87    }
88
89    /// 2026-09-25: Is the twin in this image? Whether a launch uses it is
90    /// [`Fp8QuantPick::twin`].
91    pub fn twin_present(&self) -> bool {
92        self.hopper.0 != 0
93    }
94
95    /// 2026-09-25: Which kernel this `(m, k)` launches, on which grid, and why
96    /// not the other, from the resolved `fp8_act_quant_hopper` lever and the
97    /// compiled target's SM count. [`Self::pick_with`] is the pure form.
98    pub fn pick(&self, m: u32, k: u32) -> Fp8QuantPick {
99        self.pick_with(
100            super::fp8_act_quant_hopper_enabled(),
101            m,
102            k,
103            metrale_kernels::TARGET_SM_COUNT,
104        )
105    }
106
107    /// 2026-09-25: [`Self::pick`] over an explicit lever value and SM count.
108    /// Pure, so CPU tests can grade any `(m, k)`.
109    pub fn pick_with(&self, requested: bool, m: u32, k: u32, sm_count: u32) -> Fp8QuantPick {
110        let reject = super::fp8_act_quant_hopper_reject(
111            requested,
112            self.twin_present(),
113            self.shared.0 != 0,
114            m,
115            k,
116            sm_count,
117        );
118        let twin = reject.is_none();
119        Fp8QuantPick {
120            kernel: if twin { self.hopper } else { self.shared },
121            grid: fp8_quant_grid(twin, m, k),
122            twin,
123            reject,
124            requested,
125        }
126    }
127
128    /// 2026-09-25: The handle this `(m, k)` launches.
129    pub fn kernel(&self, m: u32, k: u32) -> KernelHandle {
130        self.pick(m, k).kernel
131    }
132
133    /// 2026-09-25: The grid for `(m, k)`, from the same pick as
134    /// [`Self::kernel`].
135    pub fn grid(&self, m: u32, k: u32) -> [u32; 3] {
136        self.pick(m, k).grid
137    }
138}
139
140/// 2026-09-25: Grid for one quantizer launch. Pure, so
141/// `fp8_act_quant_tests.rs` checks both arms without a GPU.
142///
143/// `M` is on grid X in both arms: X allows 2^31-1 blocks, Y only 65535.
144///
145/// Shared arm: `(M, K/128, 1)`, one CTA per K-group; the kernel reads the
146/// group from `blockIdx.y`.
147///
148/// Hopper arm: `(M, max(1, ceil(K/128 / 8)), 1)`. The kernel derives its group
149/// span as `ceil(L / gridDim.y)`, so any Y in `1..=L` covers every group
150/// exactly once; the Y extent is a speed choice.
151pub fn fp8_quant_grid(hopper: bool, m: u32, k: u32) -> [u32; 3] {
152    let groups = k / 128;
153    if hopper {
154        [
155            m,
156            groups.div_ceil(FP8_QUANT_HOPPER_GROUPS_PER_CTA).max(1),
157            1,
158        ]
159    } else {
160        [m, groups, 1]
161    }
162}
163
164/// 2026-09-25: The half-open K-group range one Hopper CTA owns, computed as
165/// the `.cu` computes `gpc`/`g0`/`g_end`. `fp8_act_quant_tests.rs` uses it to
166/// check on the host that the grid's Y extent and the kernel's span partition
167/// the groups.
168pub fn fp8_quant_hopper_span(groups: u32, grid_y: u32, block_y: u32) -> (u32, u32) {
169    let gpc = groups.div_ceil(grid_y);
170    let g0 = block_y * gpc;
171    if g0 >= groups {
172        return (groups, groups);
173    }
174    (g0, (g0 + gpc).min(groups))
175}
176
177/// 2026-09-25: The group slot thread `tid` of a Hopper CTA works on, and the
178/// half-open range `[start, end)` of 8 elements it owns within that
179/// 128-element group, as `sub`/`lane` in the `.cu`. `None` when the slot is
180/// past the CTA's span.
181pub fn fp8_quant_hopper_lane(tid: u32, span: u32) -> Option<(u32, u32, u32)> {
182    const LANES: u32 = 16;
183    const ELEMS: u32 = 8;
184    let sub = tid / LANES;
185    let lane = tid % LANES;
186    if sub >= span {
187        return None;
188    }
189    Some((sub, lane * ELEMS, lane * ELEMS + ELEMS))
190}
191
192#[cfg(test)]
193#[path = "fp8_act_quant_tests.rs"]
194mod fp8_act_quant_tests;
metrale_model_layers/layers/ops/fp8_act_quant.rs

metrale_model_layers/layers/ops/
fp8_act_quant.rs
