metrale_model_layers/layers/ops/n_tile.rs
1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-29: Grid sizing from the N tile a kernel publishes.
4//!
5//! Owner: model-layers ops.
6//! Invariants: a launcher whose grid covers N in tiles takes the tile from the resolved
7//! kernel (`GpuBackend::kernel_n_tile`), never from a width written at the call site.
8//! One entry name can carry different tiles in different target trees
9//! (`moe_w4a16_grouped_gemm_ptrtable_t` is 64 wide in `gb10/common` and 128 wide in the
10//! Qwen trees), and a grid sized for the wrong one leaves columns unwritten.
11
12use anyhow::Result;
13use metrale_gpu_runtime::gpu::{GpuBackend, KernelHandle};
14use metrale_gpu_runtime::kernel_args::div_ceil;
15
16/// 2026-09-29: `ceil(n / t)` blocks, `t` being `kernel`'s published N tile. An error when
17/// the kernel publishes none.
18pub fn n_tile_blocks(gpu: &dyn GpuBackend, kernel: KernelHandle, n: u32) -> Result<u32> {
19 Ok(div_ceil(n, gpu.kernel_n_tile(kernel)?))
20}
21
22#[cfg(all(test, feature = "cuda"))]
23#[path = "n_tile_gpu_tests.rs"]
24mod gpu_tests;
25#[cfg(test)]
26#[path = "n_tile_tests.rs"]
27mod tests;