metrale_model_layers/layers/ops/
wide_prefill.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Which kernel serves a head_dim > 256 prefill, and the BR its
4//! grid must be built for.
5//!
6//! Owner: model-layers ops.
7//! Invariants: none beyond the types.
8
9use metrale_gpu_runtime::gpu::{GpuBackend, KernelHandle};
10
11/// 2026-09-25: The head_dim > 256 prefill kernel and the BR its grid must be
12/// built for. `qwen3_attention::init` takes the handle from here and
13/// `prefill_attention` takes the BR from here, so the two cannot disagree; a
14/// grid built for the wrong BR would not fail, it would cover the wrong
15/// q-tiles.
16///
17/// The tensor-core `attn_prefill_512tc` (`BR=32`) when the target carries it;
18/// otherwise, or under `METRALE_ATTN_512_TC=0`, the scalar `attn_prefill_512`
19/// (`BR=16`), whose handle is zero when the target lacks it too.
20pub fn wide_prefill_kernel(gpu: &dyn GpuBackend) -> (KernelHandle, u32) {
21    // 2026-09-25: Resolve with a fallback, not a fixed name. Not every target
22    // ships `attn_prefill_512tc` (deepseek-v4-flash ships only the scalar
23    // `attn_prefill_512`). A zero handle would make the callers'
24    // `hd > 256 && prefill_attn_512_k != 0` guard false and send the wide heads
25    // down the general prefill path without an error.
26    if std::env::var("METRALE_ATTN_512_TC").ok().as_deref() != Some("0") {
27        let tc = crate::layers::try_kernel(gpu, "attn_prefill_512tc", "attn_prefill_512tc");
28        if tc.0 != 0 {
29            return (tc, 32);
30        }
31        tracing::debug!(
32            "attn_prefill_512tc absent for this target; using the scalar \
33             HDIM=512 reference"
34        );
35    }
36    (
37        crate::layers::try_kernel(gpu, "attn_prefill_512", "attn_prefill_512"),
38        16,
39    )
40}
metrale_model_layers/layers/ops/wide_prefill.rs

metrale_model_layers/layers/ops/
wide_prefill.rs
