spark_model/layers/ops/
wide_prefill.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Which kernel serves an HDIM>256 prefill, and the BR its grid must be built
4//! for. Split out of `prefill_attn_main_a.rs` for the repo's 500-line cap, and
5//! it belongs alone regardless: the NAME is chosen in `qwen3_attention::init`
6//! while the GRID is built in `prefill_attention`, two different files, and a
7//! mismatch between them does not fail — it silently computes the wrong q-tiles.
8
9use spark_runtime::gpu::{GpuBackend, KernelHandle};
10
11/// The HDIM>256 prefill kernel: its module/entry name and the BR its grid must
12/// be built for. ONE reader, because the name is chosen in `qwen3_attention::init`
13/// and the grid here, and a mismatch is silent.
14///
15/// Default is the tensor-core instantiation (`BR=32`). `ATLAS_ATTN_512_TC=0`
16/// selects the scalar reference (`BR=16`) — kept reachable because it is the
17/// oracle the TC path was validated against (cosine 0.999998, 64.7x faster on
18/// S=1024/4q/2kv/causal).
19pub fn wide_prefill_kernel(gpu: &dyn GpuBackend) -> (KernelHandle, u32) {
20    // ★ RESOLVE WITH A FALLBACK, NOT A FIXED NAME. Only targets that ship the
21    // HDIM=512 instantiation have `inferspark_prefill_512tc`; gemma-4-31b, for
22    // one, ships only the scalar `inferspark_prefill_512`. Returning the TC name
23    // unconditionally leaves those targets with KernelHandle(0), which makes the
24    // caller's `hd > 256 && handle != 0` guard go FALSE and quietly routes
25    // 512-wide heads into the 64-wide kernel — no error, wrong results. That is
26    // the PR #296 failure class (a kernel handle absent, a silent fallback, and
27    // both gates green), and this path very nearly reproduced it.
28    if std::env::var("ATLAS_ATTN_512_TC").ok().as_deref() != Some("0") {
29        let tc =
30            crate::layers::try_kernel(gpu, "inferspark_prefill_512tc", "inferspark_prefill_512tc");
31        if tc.0 != 0 {
32            return (tc, 32);
33        }
34        tracing::debug!(
35            "inferspark_prefill_512tc absent for this target; using the scalar \
36             HDIM=512 reference"
37        );
38    }
39    (
40        crate::layers::try_kernel(gpu, "inferspark_prefill_512", "inferspark_prefill_512"),
41        16,
42    )
43}
spark_model/layers/ops/wide_prefill.rs

spark_model/layers/ops/
wide_prefill.rs
