gam-sae 0.3.155

//! GPU score-block kernel for the collapsed-linear-lane router (#1026).
//!
//! The collapsed linear lane ([`crate::sparse_dict`]) scales a linear SAE
//! dictionary to `K ≈ 32_000` atoms by routing each row against the WHOLE
//! dictionary one atom-tile at a time and keeping only the top-`s` atoms online.
//! That route step — `scores[r][a] = Σ_c x[r][c]·decoder[a][c]` over a
//! `rows × tile × P` block — is the dominant cost of a fit (a single fit at
//! `K ≈ 32k` is the measured 1e4–1e6× hardware gap the issue tracks) and is the
//! embarrassingly-parallel shape a GPU exists for.
//!
//! # What this offloads (and what it does NOT)
//!
//! This computes one atom-tile's `rows × tile` score block on the device, then
//! folds that tile into per-row top-`s` state on the device before moving to the
//! next tile. The minibatch router [`route_minibatch_required`] walks the whole
//! `K`-wide dictionary in atom-column tiles (each launch's block capped at
//! `GPU_ROUTE_TILE_ELEMS`), so peak device score memory is `rows × tile`,
//! **independent of `K`**. The host downloads only the final `rows × s`
//! `(atom, score)` shortlists instead of every score tile. The lane's no-`N×K`
//! memory discipline is preserved exactly as on the CPU; the GPU does both the
//! `O(rows·tile·P)` multiply-accumulate and the high-K top-`s` scan.
//!
//! # Bit-exact parity (the gate, not a tolerance)
//!
//! The CPU oracle accumulates `acc += x[c]·d[c]` as SEPARATE f32 multiply then
//! f32 add (Rust emits no fused multiply-add for `a*b+c` unless `f32::mul_add`
//! is called, and `-ffp-contract` is off), in ascending `c` order. NVRTC
//! defaults to `--fmad=true`, which contracts `a*b+c` into a single-rounding
//! FMA — a ~1 ULP difference that can flip a near-tie top-`s` selection and make
//! the routed support, and hence the whole fit, diverge from the CPU oracle.
//!
//! So the kernel forces SEPARATE rounding with `__fmul_rn` + `__fadd_rn`, in the
//! SAME ascending-`c` order, giving a score block that is **bit-for-bit**
//! identical to the CPU `score_row_tile` (every `f32` equal under `to_bits`).
//! Because the scores are identical and the device fold uses the same
//! `(|score| desc, atom asc)` ordering as [`super::scoring::TopSSelector`], the
//! routed support is IDENTICAL to the CPU oracle — parity is exact by
//! construction, not bounded by a tolerance.

#![cfg(target_os = "linux")]

use ndarray::ArrayView2;

/// The bit-exact-parity NVRTC kernel. A `BM × BN` output tile per CUDA block;
/// each thread owns one `(row, atom)` output and accumulates over `P` columns in
/// ascending order with separate-rounding f32 ops so the result matches the CPU
/// sequential `acc += x·d` to the bit.
///
/// The row and atom operands for the tile are cooperatively staged into fixed
/// `PK`-wide shared-memory chunks. This keeps per-block shared memory bounded by
/// `(BM + BN) * PK * sizeof(float)` even when the live T1 feature width is large
/// (`P=2048`), while every output still visits chunks and columns in ascending
/// `c`. The arithmetic is unchanged — each term is summed with
/// `__fmul_rn`/`__fadd_rn`; shared memory only holds exact copies of the same
/// operands, so the CPU-oracle parity gate is preserved by construction.
///
/// `PP` (the column count) is baked in as a `#define` so the inner loop is a
/// fixed trip count (matching the other NVRTC kernels in this repo, which
/// monomorphise their shape macros for a pure `compile_ptx`). `BM`/`BN` are the
/// output-tile dimensions; the host launch (`score_block_device`) must use a
/// `(BN, BM)` block and a `ceil(n_atoms/BN) × ceil(n_rows/BM)` grid.
pub const SCORE_BLOCK_KERNEL_SOURCE: &str = r#"
// Register-blocked score GEMM. A block computes a BM×BN output tile with a
// TM×TN thread block, each thread owning an RM×RN micro-tile of outputs
// (RM=BM/TM, RN=BN/TN). Every output still accumulates its own dot product in
// strictly ascending c with SEPARATE-rounding f32 ops (__fmul_rn/__fadd_rn, no
// FMA contraction), so each result is bit-identical to the CPU
// `acc += x[c]*d[c]` reference and to the earlier one-output-per-thread kernel.
// The micro-tile buys (a) RM*RN independent accumulator chains per thread to
// hide the serial __fadd_rn latency the old kernel was bound by, and (b) operand
// reuse — each staged row/atom column is consumed RN/RM times from registers
// instead of re-read from shared per output.
#define BM 64
#define BN 64
#define TM 16
#define TN 16
#define RM (BM / TM)
#define RN (BN / TN)
#define PK 32

static __device__ __forceinline__
void sparse_dict_score_block_impl(
    const float* __restrict__ rows,    // [n_rows * PP] row-major
    const float* __restrict__ atoms,   // [total_atoms * PP] row-major decoder
    int n_rows,
    int n_atoms,
    unsigned int atom_offset,          // decoder slice base (0 for the tile form)
    float* __restrict__ scores)        // [n_rows * n_atoms] row-major
{
  // Shared operand chunks: BM rows and BN atoms, each PK columns long. Chunking
  // keeps shared memory fixed-size for high-P jobs while preserving ascending-c
  // accumulation order (the acc registers persist across chunks).
  __shared__ float sr[BM][PK];
  __shared__ float sa[BN][PK];
  const int row0  = blockIdx.y * BM;
  const int atom0 = blockIdx.x * BN;
  const int tx = threadIdx.x;   // 0..TN-1
  const int ty = threadIdx.y;   // 0..TM-1
  const int lin = ty * TN + tx;
  const int nthreads = TM * TN;
  // This thread owns outputs rows [row0 + ty*RM, +RM) × atoms [atom0 + tx*RN, +RN).
  float acc[RM][RN];
  #pragma unroll
  for (int i = 0; i < RM; ++i)
    #pragma unroll
    for (int j = 0; j < RN; ++j) acc[i][j] = 0.0f;
  for (int c0 = 0; c0 < PP; c0 += PK) {
    const int chunk = (PP - c0 < PK) ? (PP - c0) : PK;
    // Cooperative, coalesced load of one P-chunk of row/atom operands
    // (zero-padded past ragged row/atom/chunk tails; +0.0 is an exact no-op in
    // the ascending-c accumulation so parity holds on padded lanes).
    for (int e = lin; e < BM * PK; e += nthreads) {
      int rr = e / PK, kc = e - rr * PK;
      int gr = row0 + rr;
      int cc = c0 + kc;
      sr[rr][kc] = (gr < n_rows && kc < chunk) ? rows[(long long)gr * PP + cc] : 0.0f;
    }
    for (int e = lin; e < BN * PK; e += nthreads) {
      int aa = e / PK, kc = e - aa * PK;
      int ga = atom0 + aa;
      int cc = c0 + kc;
      unsigned int global_atom = atom_offset + (unsigned int)ga;
      sa[aa][kc] = (ga < n_atoms && kc < chunk) ? atoms[(long long)global_atom * PP + cc] : 0.0f;
    }
    __syncthreads();
    for (int kc = 0; kc < chunk; ++kc) {
      // Stage this column's RM row-fragments and RN atom-fragments into
      // registers, then cross them: RM*RN separate-rounding MACs reusing 8 loads.
      float rf[RM];
      float af[RN];
      #pragma unroll
      for (int i = 0; i < RM; ++i) rf[i] = sr[ty * RM + i][kc];
      #pragma unroll
      for (int j = 0; j < RN; ++j) af[j] = sa[tx * RN + j][kc];
      #pragma unroll
      for (int i = 0; i < RM; ++i)
        #pragma unroll
        for (int j = 0; j < RN; ++j)
          acc[i][j] = __fadd_rn(acc[i][j], __fmul_rn(rf[i], af[j]));
    }
    __syncthreads();
  }
  #pragma unroll
  for (int i = 0; i < RM; ++i) {
    int r = row0 + ty * RM + i;
    if (r >= n_rows) continue;
    #pragma unroll
    for (int j = 0; j < RN; ++j) {
      int a = atom0 + tx * RN + j;
      if (a < n_atoms) scores[(long long)r * n_atoms + a] = acc[i][j];
    }
  }
}

extern "C" __global__
void sparse_dict_score_block(
    const float* __restrict__ rows,    // [n_rows * PP] row-major
    const float* __restrict__ atoms,   // [n_atoms * PP] row-major (decoder tile)
    int n_rows,
    int n_atoms,
    float* __restrict__ scores)        // [n_rows * n_atoms] row-major
{
  sparse_dict_score_block_impl(rows, atoms, n_rows, n_atoms, 0u, scores);
}

extern "C" __global__
void sparse_dict_score_block_offset(
    const float* __restrict__ rows,    // [n_rows * PP] row-major
    const float* __restrict__ atoms,   // [total_atoms * PP] row-major decoder
    int n_rows,
    int n_atoms,
    unsigned int atom_offset,
    float* __restrict__ scores)        // [n_rows * n_atoms] row-major tile
{
  sparse_dict_score_block_impl(rows, atoms, n_rows, n_atoms, atom_offset, scores);
}

#define EMPTY_TOP_ATOM 0xffffffffu

static __device__ __forceinline__
float sparse_dict_abs_f32(float v) {
  return (v < 0.0f) ? -v : v;
}

static __device__ __forceinline__
int sparse_dict_better(float mag, unsigned int atom,
                       float ref_mag, unsigned int ref_atom) {
  return (mag > ref_mag) || (mag == ref_mag && atom < ref_atom);
}

static __device__ __forceinline__
int sparse_dict_worse(float mag, unsigned int atom,
                      float ref_mag, unsigned int ref_atom) {
  return (mag < ref_mag) || (mag == ref_mag && atom > ref_atom);
}

static __device__ __forceinline__
void sparse_dict_recompute_worst(const unsigned int* atoms,
                                 const float* mags,
                                 int count,
                                 int* worst_idx) {
  int worst = 0;
  for (int j = 1; j < count; ++j) {
    if (sparse_dict_worse(mags[j], atoms[j], mags[worst], atoms[worst])) {
      worst = j;
    }
  }
  *worst_idx = worst;
}

static __device__ __forceinline__
void sparse_dict_offer_top_s(unsigned int* atoms,
                             float* scores,
                             float* mags,
                             int active,
                             unsigned int atom,
                             float score,
                             int* count,
                             int* worst_idx) {
  if (active <= 0) {
    return;
  }
  const float mag = sparse_dict_abs_f32(score);
  if (*count < active) {
    const int slot = *count;
    atoms[slot] = atom;
    scores[slot] = score;
    mags[slot] = mag;
    *count = slot + 1;
    if (*count == active) {
      sparse_dict_recompute_worst(atoms, mags, *count, worst_idx);
    }
    return;
  }
  const int worst = *worst_idx;
  if (sparse_dict_better(mag, atom, mags[worst], atoms[worst])) {
    atoms[worst] = atom;
    scores[worst] = score;
    mags[worst] = mag;
    sparse_dict_recompute_worst(atoms, mags, active, worst_idx);
  }
}

static __device__ __forceinline__
void sparse_dict_sort_top_s(unsigned int* atoms,
                            float* scores,
                            float* mags,
                            int active,
                            int count) {
  for (int i = 1; i < count; ++i) {
    const unsigned int atom = atoms[i];
    const float score = scores[i];
    const float mag = mags[i];
    int j = i;
    while (j > 0 && sparse_dict_better(mag, atom, mags[j - 1], atoms[j - 1])) {
      atoms[j] = atoms[j - 1];
      scores[j] = scores[j - 1];
      mags[j] = mags[j - 1];
      --j;
    }
    atoms[j] = atom;
    scores[j] = score;
    mags[j] = mag;
  }
  for (int j = count; j < active; ++j) {
    atoms[j] = EMPTY_TOP_ATOM;
    scores[j] = 0.0f;
    mags[j] = -1.0f;
  }
}

extern "C" __global__
void sparse_dict_fold_top_s(
    const float* __restrict__ scores,  // [n_rows * n_atoms] current tile
    int n_rows,
    int n_atoms,
    unsigned int atom_offset,
    int active,
    unsigned int* __restrict__ top_atoms, // [n_rows * active]
    float* __restrict__ top_scores,       // [n_rows * active]
    float* __restrict__ top_mags)         // [n_rows * active]
{
  const int row = blockIdx.x;
  if (row >= n_rows || active <= 0) {
    return;
  }
  const int tid = threadIdx.x;
  const int nthreads = blockDim.x;
  const int candidate_slots = nthreads * active;

  extern __shared__ unsigned char smem[];
  unsigned int* cand_atoms = (unsigned int*)smem;
  unsigned int* best_atoms = cand_atoms + candidate_slots;
  float* cand_scores = (float*)(best_atoms + active);
  float* best_scores = cand_scores + candidate_slots;
  float* cand_mags = best_scores + active;
  float* best_mags = cand_mags + candidate_slots;

  const int local_base = tid * active;
  for (int j = 0; j < active; ++j) {
    cand_atoms[local_base + j] = EMPTY_TOP_ATOM;
    cand_scores[local_base + j] = 0.0f;
    cand_mags[local_base + j] = -1.0f;
  }
  if (tid == 0) {
    for (int j = 0; j < active; ++j) {
      best_atoms[j] = EMPTY_TOP_ATOM;
      best_scores[j] = 0.0f;
      best_mags[j] = -1.0f;
    }
  }
  __syncthreads();

  int local_count = 0;
  int local_worst = 0;
  unsigned int* local_atoms = cand_atoms + local_base;
  float* local_scores = cand_scores + local_base;
  float* local_mags = cand_mags + local_base;
  const long long row_base = (long long)row * n_atoms;
  for (int atom = tid; atom < n_atoms; atom += nthreads) {
    const float score = scores[row_base + atom];
    sparse_dict_offer_top_s(
        local_atoms,
        local_scores,
        local_mags,
        active,
        atom_offset + (unsigned int)atom,
        score,
        &local_count,
        &local_worst);
  }
  __syncthreads();

  if (tid == 0) {
    int best_count = 0;
    int best_worst = 0;
    const long long out_base = (long long)row * active;
    if (atom_offset != 0u) {
      for (int j = 0; j < active; ++j) {
        const unsigned int atom = top_atoms[out_base + j];
        if (atom != EMPTY_TOP_ATOM) {
          sparse_dict_offer_top_s(
              best_atoms,
              best_scores,
              best_mags,
              active,
              atom,
              top_scores[out_base + j],
              &best_count,
              &best_worst);
        }
      }
    }
    for (int t = 0; t < nthreads; ++t) {
      const int base = t * active;
      for (int j = 0; j < active; ++j) {
        const unsigned int atom = cand_atoms[base + j];
        if (atom != EMPTY_TOP_ATOM) {
          sparse_dict_offer_top_s(
              best_atoms,
              best_scores,
              best_mags,
              active,
              atom,
              cand_scores[base + j],
              &best_count,
              &best_worst);
        }
      }
    }
    sparse_dict_sort_top_s(best_atoms, best_scores, best_mags, active, best_count);
    for (int j = 0; j < active; ++j) {
      top_atoms[out_base + j] = best_atoms[j];
      top_scores[out_base + j] = best_scores[j];
      top_mags[out_base + j] = best_mags[j];
    }
  }
}
"#;

/// Output-tile dimensions the [`SCORE_BLOCK_KERNEL_SOURCE`] kernel is written
/// for (`BM`/`BN`); the host grid uses them to tile the `n_rows × n_atoms`
/// output. Kept in sync with the `#define`s at the top of the kernel string.
pub const SCORE_BLOCK_TILE_M: u32 = 64;
pub const SCORE_BLOCK_TILE_N: u32 = 64;

/// Thread-block dimensions (`TM`/`TN`) for the register-blocked kernel: each
/// thread owns an `(BM/TM) × (BN/TN)` micro-tile of outputs, so the launch uses
/// a `TN × TM` thread block over the `BM × BN` output tile. Kept in sync with the
/// `#define`s at the top of the kernel string.
pub const SCORE_BLOCK_THREADS_M: u32 = 16;
pub const SCORE_BLOCK_THREADS_N: u32 = 16;

/// Prepend the `PP` shape macro so the NVRTC compile is a pure `compile_ptx`
/// (mirrors `sae_rowjet::softmax_kernel_source` / `arrow_schur_nvrtc`).
#[must_use]
pub fn score_block_kernel_source(p: usize) -> String {
    format!("#define PP {p}\n{SCORE_BLOCK_KERNEL_SOURCE}")
}

/// Minimum score-block element count (`n_rows · n_atoms`) below which the device
/// launch is not worth its fixed cost (probe + H2D + D2H). Below this the CPU
/// reference is used. Tuned to the same genus as the other SAE device floors
/// (the calibrated row-kernel crossover used by the complete SAE row jet).
pub const DEVICE_SCORE_BLOCK_MIN_ELEMS: usize = gam_gpu::DEFAULT_DICTIONARY_SCORE_MIN_ELEMS;

/// Which path produced a score block. Returned by the fail-loud entry point so
/// callers (and the parity test) can ASSERT the device engaged rather than
/// silently falling back — the #1026/#1551 'GPU 0%' failure mode.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ScoreBlockPath {
    /// The NVRTC `sparse_dict_score_block` kernel ran on the device.
    Device,
    /// The CPU `score_block_cpu` reference ran.
    Cpu,
}

/// Peak score elements per device launch for the tiled GPU router. The router
/// NEVER materialises the whole `m × K` block: it walks `K` in atom-column tiles
/// sized so each launch's `m × cols` block stays under this cap (~2M f32 ≈ 8 MB
/// device score buffer), then discards it after folding. This keeps peak score
/// memory bounded **independent of `K`** — the same discipline the CPU lane
/// ([`super::scoring::top_s_online`]) keeps with its `rows × tile` column tiles —
/// so a `K ≈ 32_000` fit does not balloon a `device alloc` linearly in `K`.
const GPU_ROUTE_TILE_ELEMS: usize = gam_gpu::DEFAULT_DICTIONARY_SCORE_TILE_ELEMS;

pub fn route_minibatch_required(
    rows: ArrayView2<'_, f32>,
    decoder: ArrayView2<'_, f32>,
    s: usize,
    tile: usize,
    mode: gam_gpu::GpuPolicy,
) -> Result<(Vec<Vec<(u32, f32)>>, ScoreBlockPath, usize), gam_gpu::GpuError> {
    use super::scoring::top_s_online;

    let m = rows.nrows();
    let k = decoder.nrows();
    let active = s.max(1).min(k.max(1));

    // CPU per-row path (bit-identical oracle), used for Off, below break-even,
    // or a genuine device absence under Auto.
    let cpu_route = || -> Vec<Vec<(u32, f32)>> {
        rows.outer_iter()
            .map(|row| top_s_online(row, decoder, s, tile))
            .collect()
    };

    if mode == gam_gpu::GpuPolicy::Off {
        return Ok((cpu_route(), ScoreBlockPath::Cpu, 0));
    }

    // Engagement is decided on the TOTAL work `m × K` (that is what justifies the
    // device's fixed launch cost), but the launches themselves are K-tiled so the
    // buffers never grow with K.
    let plan = gam_gpu::DictionaryScoreRoutePlan::with_limits(
        m,
        k,
        decoder.ncols(),
        DEVICE_SCORE_BLOCK_MIN_ELEMS,
        GPU_ROUTE_TILE_ELEMS,
    );
    if !plan.device_admitted {
        if mode == gam_gpu::GpuPolicy::Required {
            return Err(gam_gpu::gpu_err!(
                "route_minibatch GpuPolicy::Required: block of {m}×{k} = {} elems is below \
                 the device launch break-even (DEVICE_SCORE_BLOCK_MIN_ELEMS={DEVICE_SCORE_BLOCK_MIN_ELEMS}); \
                 refusing to silently run on the CPU",
                m.saturating_mul(k)
            ));
        }
        gam_gpu::engagement::note_route_engagement(
        "gam-sae sparse_dict score router",
        "falling back to CPU",
            false,
            &format!(
                "block {m}x{k} = {} elems below the device launch break-even \
                 (DEVICE_SCORE_BLOCK_MIN_ELEMS={DEVICE_SCORE_BLOCK_MIN_ELEMS})",
                m.saturating_mul(k)
            ),
        );
        return Ok((cpu_route(), ScoreBlockPath::Cpu, 0));
    }
    if m == 0 || k == 0 {
        return Ok((cpu_route(), ScoreBlockPath::Cpu, 0));
    }

    let runtime = if mode == gam_gpu::GpuPolicy::Required {
        Some(gam_gpu::GpuRuntime::require()?)
    } else {
        gam_gpu::GpuRuntime::resolve(mode)?
    };
    if runtime.is_none() {
        gam_gpu::engagement::note_route_engagement(
        "gam-sae sparse_dict score router",
        "falling back to CPU",
        false, "Auto admission found no CUDA device");
        return Ok((cpu_route(), ScoreBlockPath::Cpu, 0));
    }

    // Atom-columns per device launch: bound the per-launch block to
    // GPU_ROUTE_TILE_ELEMS, at least one column, never more than K.
    let tile_cols = plan.tile_items;

    let out = device::route_decoder_tiled_device(rows, decoder, active, tile_cols)?;
    gam_gpu::engagement::note_route_engagement(
        "gam-sae sparse_dict score router",
        "falling back to CPU",
        true,
        &format!("block {m}x{k}, tile_cols={tile_cols}, active={active}"),
    );
    Ok((
        out.selections,
        ScoreBlockPath::Device,
        out.device_dtoh_bytes,
    ))
}

mod device {
    use super::score_block_kernel_source;
    use gam_gpu::backend_probe::CachedBackend;
    use gam_gpu::gpu_error::{GpuError, GpuResultExt};
    use ndarray::ArrayView2;
    use std::sync::Arc;

    use cudarc::driver::{CudaModule, LaunchConfig, PushKernelArg};

    use super::super::score_router_backend::ScoreRouterBackend as Backend;

    static BACKEND: CachedBackend<Backend> = CachedBackend::new();

    fn backend() -> Result<&'static Backend, GpuError> {
        BACKEND.get_or_probe("sparse_dict_score_block", Backend::from_parts)
    }

    fn module_for(b: &Backend, p: usize) -> Result<Arc<CudaModule>, GpuError> {
        b.modules
            .get_or_compile(&b.ctx, p, "sparse_dict score-block", score_block_kernel_source)
    }

    const TOP_S_FOLD_THREADS: u32 = 32;

    /// Number of bounded-progress checkpoints across one tiled route's atom-tile
    /// walk (#2227). This is a telemetry/backlog cadence, not a numerical tuning
    /// knob: the tile loop synchronises `min(tile_count, this)` times so the
    /// in-flight async launch backlog is bounded to `ceil(tile_count/this)` tiles
    /// and any device fault or stall is attributed to the tile window that
    /// produced it, instead of surfacing (if at all) as one unattributed block in
    /// the terminal synchronize with no telemetry for the whole high-`K` route.
    const ROUTE_PROGRESS_CHECKPOINTS: usize = 16;

    pub(super) struct RouteDeviceOutput {
        pub(super) selections: Vec<Vec<(u32, f32)>>,
        pub(super) device_dtoh_bytes: usize,
    }

    fn fold_shared_bytes(
        active: usize,
        threads: u32,
        max_shared_mem_per_block: usize,
    ) -> Result<u32, GpuError> {
        let slots = (threads as usize)
            .checked_add(1)
            .and_then(|v| v.checked_mul(active))
            .ok_or_else(|| gam_gpu::gpu_err!("sparse_dict top-s fold shared-memory overflow"))?;
        let bytes = slots
            .checked_mul(
                std::mem::size_of::<u32>()
                    + std::mem::size_of::<f32>()
                    + std::mem::size_of::<f32>(),
            )
            .ok_or_else(|| gam_gpu::gpu_err!("sparse_dict top-s fold shared-memory overflow"))?;
        if max_shared_mem_per_block > 0 && bytes > max_shared_mem_per_block {
            return Err(gam_gpu::gpu_err!(
                "sparse_dict top-s fold requires {bytes} shared-memory bytes per row block \
                 (active={active}, threads={threads}) but the selected device reports \
                 max_shared_mem_per_block={max_shared_mem_per_block}"
            ));
        }
        u32::try_from(bytes)
            .map_err(|_| gam_gpu::gpu_err!("sparse_dict top-s fold shared-memory bytes overflow"))
    }

    /// Route-time score-block stream for a full decoder. Rows and the whole decoder
    /// stay resident for the route; one reusable score buffer rotates across K.
    /// Each score tile is folded into resident per-row top-`s` state by
    /// `sparse_dict_fold_top_s`, so the host downloads only the final `(atom,
    /// score)` shortlists instead of every `rows × tile` score.
    pub(super) fn route_decoder_tiled_device(
        rows: ArrayView2<'_, f32>,
        decoder: ArrayView2<'_, f32>,
        active: usize,
        tile_cols: usize,
    ) -> Result<RouteDeviceOutput, GpuError> {
        let n_rows = rows.nrows();
        let k = decoder.nrows();
        let p = rows.ncols();
        if p != decoder.ncols() {
            return Err(gam_gpu::gpu_err!(
                "sparse_dict tiled score: P mismatch rows={p} decoder={}",
                decoder.ncols()
            ));
        }
        if n_rows == 0 || k == 0 || p == 0 {
            return Ok(RouteDeviceOutput {
                selections: vec![Vec::new(); n_rows],
                device_dtoh_bytes: 0,
            });
        }
        let active = active.max(1).min(k);
        if k > u32::MAX as usize {
            return Err(gam_gpu::gpu_err!(
                "sparse_dict tiled route K={k} exceeds u32 atom-index storage"
            ));
        }

        let b = backend()?;
        let module = module_for(b, p)?;
        let score_func = module
            .load_function("sparse_dict_score_block_offset")
            .gpu_ctx("sparse_dict tiled score-offset load_function")?;
        let fold_func = module
            .load_function("sparse_dict_fold_top_s")
            .gpu_ctx("sparse_dict top-s fold load_function")?;
        let stream = b.stream.clone();

        let rows_storage: Vec<f32>;
        let rows_host: &[f32] = if let Some(slice) = rows.as_slice() {
            slice
        } else {
            rows_storage = rows.iter().copied().collect();
            rows_storage.as_slice()
        };
        assert_eq!(
            rows_host.len(),
            n_rows * p,
            "tiled score rows flatten length"
        );
        let rows_dev = stream
            .clone_htod(rows_host)
            .gpu_ctx("sparse_dict tiled score htod rows")?;

        let decoder_storage: Vec<f32>;
        let decoder_host: &[f32] = if let Some(slice) = decoder.as_slice() {
            slice
        } else {
            decoder_storage = decoder.iter().copied().collect();
            decoder_storage.as_slice()
        };
        assert_eq!(
            decoder_host.len(),
            k * p,
            "tiled score decoder flatten length"
        );

        let n_rows_i32 = i32::try_from(n_rows).map_err(|_| {
            gam_gpu::gpu_err!("sparse_dict tiled score n_rows={n_rows} overflows i32")
        })?;
        let active_i32 = i32::try_from(active).map_err(|_| {
            gam_gpu::gpu_err!("sparse_dict tiled score active={active} overflows i32")
        })?;
        let tile_m = super::SCORE_BLOCK_TILE_M;
        let tile_n = super::SCORE_BLOCK_TILE_N;
        let tile_cols = tile_cols.max(1);
        let max_tile_cols = tile_cols.min(k);
        let decoder_dev = stream
            .clone_htod(decoder_host)
            .gpu_ctx("sparse_dict tiled score htod decoder")?;
        let mut scores_dev = stream
            .alloc_zeros::<f32>(n_rows * max_tile_cols)
            .gpu_ctx("sparse_dict tiled score alloc scores")?;
        let mut top_atoms_dev = stream
            .alloc_zeros::<u32>(n_rows * active)
            .gpu_ctx("sparse_dict top-s alloc atoms")?;
        let mut top_scores_dev = stream
            .alloc_zeros::<f32>(n_rows * active)
            .gpu_ctx("sparse_dict top-s alloc scores")?;
        let mut top_mags_dev = stream
            .alloc_zeros::<f32>(n_rows * active)
            .gpu_ctx("sparse_dict top-s alloc mags")?;
        let fold_shared =
            fold_shared_bytes(active, TOP_S_FOLD_THREADS, b.max_shared_mem_per_block)?;

        // Bounded-progress checkpoints (#2227). The walk enqueues every score+fold
        // launch on one stream; without intermediate synchronisation a device fault
        // or stall in any tile surfaces only at the terminal synchronize, as a
        // single unattributed block with no telemetry for the whole high-`K` route.
        // Synchronise on a cadence derived from the tile count so the async backlog
        // is bounded and each fault is attributed to its tile window; the heartbeat
        // is `log::debug!` so an ordinary (info-level) per-minibatch run is not
        // flooded, while `RUST_LOG=debug` exposes intra-route progress.
        let tile_count = k.div_ceil(tile_cols);
        let checkpoint_stride = tile_count
            .div_ceil(ROUTE_PROGRESS_CHECKPOINTS.max(1))
            .max(1);
        let route_started = std::time::Instant::now();
        let mut tiles_done = 0usize;
        let mut checkpoint_lo = 0usize;
        let mut start = 0usize;
        while start < k {
            let end = (start + tile_cols).min(k);
            let n_atoms = end - start;
            let n_atoms_i32 = i32::try_from(n_atoms).map_err(|_| {
                gam_gpu::gpu_err!("sparse_dict tiled score n_atoms={n_atoms} overflows i32")
            })?;
            let atom_offset = u32::try_from(start).map_err(|_| {
                gam_gpu::gpu_err!("sparse_dict tiled score atom offset={start} overflows u32")
            })?;
            let grid_x: u32 = u32::try_from(n_atoms.div_ceil(tile_n as usize))
                .map_err(|_| gam_gpu::gpu_err!("sparse_dict tiled score grid_x overflow"))?;
            let grid_y: u32 = u32::try_from(n_rows.div_ceil(tile_m as usize))
                .map_err(|_| gam_gpu::gpu_err!("sparse_dict tiled score grid_y overflow"))?;
            let cfg = LaunchConfig {
                grid_dim: (grid_x, grid_y, 1),
                block_dim: (
                    super::SCORE_BLOCK_THREADS_N,
                    super::SCORE_BLOCK_THREADS_M,
                    1,
                ),
                shared_mem_bytes: 0,
            };
            let mut builder = stream.launch_builder(&score_func);
            builder
                .arg(&rows_dev)
                .arg(&decoder_dev)
                .arg(&n_rows_i32)
                .arg(&n_atoms_i32)
                .arg(&atom_offset)
                .arg(&mut scores_dev);
            // SAFETY: grid/block validated; device pointers are cudarc-checked
            // allocations on this stream. The kernel reads the resident rows and
            // resident decoder slice `[atom_offset, atom_offset + n_atoms)` and
            // writes exactly `n_rows * n_atoms` scores.
            unsafe { builder.launch(cfg) }.gpu_ctx("sparse_dict tiled score launch")?;

            let fold_cfg = LaunchConfig {
                grid_dim: (
                    u32::try_from(n_rows)
                        .map_err(|_| gam_gpu::gpu_err!("sparse_dict top-s fold grid overflow"))?,
                    1,
                    1,
                ),
                block_dim: (TOP_S_FOLD_THREADS, 1, 1),
                shared_mem_bytes: fold_shared,
            };
            let mut fold = stream.launch_builder(&fold_func);
            fold.arg(&scores_dev)
                .arg(&n_rows_i32)
                .arg(&n_atoms_i32)
                .arg(&atom_offset)
                .arg(&active_i32)
                .arg(&mut top_atoms_dev)
                .arg(&mut top_scores_dev)
                .arg(&mut top_mags_dev);
            // SAFETY: the fold kernel launches one block per row, reads the
            // score tile just written by the previous launch on this stream, and
            // updates exactly `n_rows * active` shortlist slots.
            unsafe { fold.launch(fold_cfg) }.gpu_ctx("sparse_dict top-s fold launch")?;
            start = end;
            tiles_done += 1;
            if tiles_done % checkpoint_stride == 0 || start >= k {
                stream.synchronize().gpu_ctx_with(|err| {
                    format!(
                        "sparse_dict tiled route progress checkpoint (tiles {checkpoint_lo}..{tiles_done} of {tile_count}, atoms 0..{start} of {k}): {err}"
                    )
                })?;
                log::debug!(
                    "[SAE score route] tiles {tiles_done}/{tile_count} atoms {start}/{k} \
                     elapsed {:.2}s",
                    route_started.elapsed().as_secs_f64(),
                );
                checkpoint_lo = tiles_done;
            }
        }

        let mut top_atoms = vec![0u32; n_rows * active];
        let mut top_scores = vec![0.0f32; n_rows * active];
        stream
            .memcpy_dtoh(&top_atoms_dev, &mut top_atoms)
            .gpu_ctx("sparse_dict top-s dtoh atoms")?;
        stream
            .memcpy_dtoh(&top_scores_dev, &mut top_scores)
            .gpu_ctx("sparse_dict top-s dtoh scores")?;
        stream
            .synchronize()
            .gpu_ctx("sparse_dict tiled route synchronize")?;

        let mut selections = Vec::with_capacity(n_rows);
        for r in 0..n_rows {
            let mut row = Vec::with_capacity(active);
            let base = r * active;
            for j in 0..active {
                let atom = top_atoms[base + j];
                if atom != u32::MAX {
                    row.push((atom, top_scores[base + j]));
                }
            }
            selections.push(row);
        }
        Ok(RouteDeviceOutput {
            selections,
            device_dtoh_bytes: n_rows
                .saturating_mul(active)
                .saturating_mul(std::mem::size_of::<u32>() + std::mem::size_of::<f32>()),
        })
    }
}