onnx-runtime-ep-cuda 0.1.0-dev.5

//! `pkg.nxrt::IndexShare` v1: device-resident selected-token attention for the
//! GLM-5.2 IndexShare / DeepSeek DSA building block.
//!
//! The frozen CPU reference in
//! `crates/onnx-runtime-ep-cpu/src/kernels/index_share.rs` is the authoritative
//! numerical oracle. This kernel reproduces its math **on the device**: the
//! `past ⧺ current` KV cache concatenation, the per-row selected-token gather,
//! the scaled dot-product scores, a numerically-stable fp32 softmax over the
//! selected keys, and the probability·value reduction all run in NVRTC kernels.
//! Query, key, value, bias, the present KV cache, and the output tensor stay
//! resident on the device. In eager (non-capturing) execution the small
//! `selected_indices` tensor is copied D2H so the ONNX-required deterministic
//! index validation (strictly-increasing order, trailing `-1` padding, range,
//! not-all-`-1`) produces the same hard errors as the CPU oracle. During
//! CUDA-graph capture and replay that host round-trip is illegal (it would
//! synchronize the stream), so the identical checks run **on the device** in
//! [`validate_index_rows`], which latches any violation into the runtime's
//! persistent capture-error word (read back by the host at the per-step logits
//! sync, outside the captured region). The bulk attention row kernel additionally
//! clamps every gathered key index into range so a poisoned replay can never
//! issue an out-of-bounds load.
//!
//! ## Determinism / bit-parity
//!
//! Every reduction sums in the same fixed ascending order as the CPU reference:
//! each score's dot product accumulates over `head_size` in one thread, the
//! softmax max/exp/sum runs sequentially in the block's lead thread, and the
//! `probs·V` accumulation sums over the selected keys in ascending order in one
//! thread per output channel. `sqrt(scale)` is folded into each Q and K operand
//! (matching the reference's `(Q·√scale)·(K·√scale)`), so results are
//! byte-identical to the CPU oracle.
//!
//! ## Capture support
//!
//! After a warmed eager execution has sized the module-global-style pooled
//! scratch (present K/V staging and the per-row score scratch keep stable device
//! addresses across warmup → capture → replay), the launch path is legal to
//! record into a CUDA graph and replay with only device-buffer contents
//! changing:
//!
//!   * No `stream.synchronize()` on the capturing path (the initial input-upload
//!     wait and the trailing completion wait are both skipped while capturing;
//!     same-stream ordering guarantees inputs are ready).
//!   * No per-call `cudaMalloc`/`cudaFree`: scratch is reused from the pool that
//!     the warmup pass allocated. Growth (a fresh alloc) only ever happens in
//!     eager mode, so captured replay keeps fixed buffer addresses.
//!   * No `selected_indices` D2H copy: [`validate_index_rows`] performs the
//!     deterministic ONNX index validation on the device and latches violations
//!     into the runtime capture-error word.
//!
//! Capture stays gated off until such a warmup has run (mirroring
//! [`super::gather`]); until then [`capture_support`] reports the missing
//! precondition.
//!
//! ## Claim-time gating
//!
//! [`unsupported_reason`] (in the CUDA provider) delegates to the CPU oracle's
//! own `unsupported_reason`, so the two backends reject exactly the same
//! dtype/layout/arity/shape combinations at claim time rather than claiming a
//! node and falling back inside the kernel.

use std::borrow::Cow;
use std::ffi::c_void;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex};

use cudarc::driver::sys::CUdeviceptr;
use cudarc::driver::{LaunchConfig, PushKernelArg};
use onnx_runtime_ep_api::{
    CaptureSupport, EpError, Kernel, KernelFactory, Result, TensorMut, TensorView,
};
use onnx_runtime_ir::{DataType, Node};

use crate::error::driver_err;
use crate::runtime::{CudaRuntime, cuptr};

const OP: &str = "IndexShare";

/// Claim-time validation preserving the CPU oracle's structural ABI checks
/// while extending its f32-only execution oracle to CUDA's f16/bf16 storage
/// variants. The CPU validation receives an f32 dtype projection after this
/// method has enforced CUDA's homogeneous floating dtype contract.
pub(crate) fn unsupported_reason(
    node: &Node,
    shapes: &[onnx_runtime_ir::Shape],
    input_dtypes: &[DataType],
) -> Option<Cow<'static, str>> {
    let dtype_at = |index| {
        input_dtypes
            .get(index)
            .copied()
            .unwrap_or(DataType::Undefined)
    };
    let dtype = dtype_at(0);
    if !matches!(
        dtype,
        DataType::Float32 | DataType::Float16 | DataType::BFloat16
    ) {
        return Some(Cow::Owned(format!(
            "IndexShare: query dtype {dtype:?} unsupported on CUDA (expected f32, f16, or bf16)"
        )));
    }
    for index in [1, 2, 3, 4, 6] {
        let candidate = dtype_at(index);
        if candidate != DataType::Undefined && candidate != dtype {
            return Some(Cow::Borrowed(
                "IndexShare: query, key, value, past_key, past_value, and attention_bias must use the same floating dtype on CUDA",
            ));
        }
    }
    let projected: Vec<_> = input_dtypes
        .iter()
        .map(|&candidate| {
            if matches!(
                candidate,
                DataType::Float32 | DataType::Float16 | DataType::BFloat16
            ) {
                DataType::Float32
            } else {
                candidate
            }
        })
        .collect();
    onnx_runtime_ep_cpu::kernels::index_share::unsupported_reason(node, shapes, &projected)
}

/// Capture-error latch bit raised by [`validate_index_rows`] when a captured
/// replay observes a `selected_indices` row that fails the deterministic ONNX
/// validation (out of range, not strictly increasing, index after trailing
/// `-1` padding, or an all-`-1` row). The host reads the shared latch at the
/// per-step logits sync, so a poisoned replay is rejected before its token is
/// consumed.
pub const INDEX_SHARE_CAPTURE_ERROR_INDEX: u32 = 512;
const INPUT_NAMES: [&str; 7] = [
    "query",
    "key",
    "value",
    "past_key",
    "past_value",
    "selected_indices",
    "attention_bias",
];

/// Threads per block for the bulk-copy present-cache builder.
const BLOCK: u32 = 256;
/// Threads per block for `index_share_row` (one block services one output row).
const ROW_THREADS: u32 = 128;
const MODULE: &str = "index_share_f32_f16_bf16_v3";
const SOURCE: &str = r#"
#include <cuda_fp16.h>
#include <cuda_bf16.h>
#define NEG_INF __int_as_float(0xff800000)

// dtype is 0 for f32, 1 for f16, and 2 for bf16. Scores and reductions remain
// fp32; only externally visible tensors and the K/V cache use this storage type.
__device__ __forceinline__ float load_float(
    const void* data, unsigned long long index, int dtype) {
  if (dtype == 0) {
    return ((const float*)data)[index];
  }
  if (dtype == 1) {
    return __half2float(((const __half*)data)[index]);
  }
  return __bfloat162float(((const __nv_bfloat16*)data)[index]);
}

__device__ __forceinline__ void store_float(
    void* data, unsigned long long index, float value, int dtype) {
  if (dtype == 0) {
    ((float*)data)[index] = value;
  } else if (dtype == 1) {
    ((__half*)data)[index] = __float2half_rn(value);
  } else {
    ((__nv_bfloat16*)data)[index] = __float2bfloat16_rn(value);
  }
}

// Gather a K/V input plus an optional past cache into a contiguous
// [batch, kv_heads, total_seq, head_size] present buffer (past ++ current along
// the sequence axis). This is a pure copy, so the present outputs are
// bit-identical to the CPU reference's concatenation.
extern "C" __global__ void build_present(
    const void* past, const void* cur, void* out, int dtype, int has_past,
    unsigned long long batch, unsigned long long heads,
    unsigned long long past_seq, unsigned long long cur_seq,
    unsigned long long total_seq, unsigned long long dim,
    unsigned long long elements) {
  for (unsigned long long idx = (unsigned long long)blockIdx.x * blockDim.x + threadIdx.x;
       idx < elements; idx += (unsigned long long)gridDim.x * blockDim.x) {
    unsigned long long d = idx % dim;
    unsigned long long rem = idx / dim;
    unsigned long long t = rem % total_seq;
    rem /= total_seq;
    unsigned long long h = rem % heads;
    unsigned long long b = rem / heads;
    float val;
    if (has_past && t < past_seq) {
      val = load_float(past, ((b * heads + h) * past_seq + t) * dim + d, dtype);
    } else {
      unsigned long long c = has_past ? (t - past_seq) : t;
      val = load_float(cur, ((b * heads + h) * cur_seq + c) * dim + d, dtype);
    }
    store_float(out, idx, val, dtype);
  }
}

// Additive attention bias for logical index (b, h, q, k), broadcasting a
// rank<=4 bias right-aligned against [b, h, q, k]. Mirrors the CPU reference's
// Bias::at exactly (size-1 axes broadcast; no -inf padding for a short last
// dim, which the claim gate already forbids).
__device__ __forceinline__ float bias_at(
    const void* bias, int dtype, int rank,
    unsigned long long bd0, unsigned long long bd1,
    unsigned long long bd2, unsigned long long bd3,
    unsigned long long b, unsigned long long h,
    unsigned long long q, unsigned long long k) {
  unsigned long long logical[4] = {b, h, q, k};
  unsigned long long dims[4] = {bd0, bd1, bd2, bd3};
  unsigned long long off = 0;
  for (int axis = 4 - rank; axis < 4; ++axis) {
    unsigned long long dim = dims[axis];
    unsigned long long index = (dim == 1ULL) ? 0ULL : logical[axis];
    off = off * dim + index;
  }
  return load_float(bias, off, dtype);
}

// Recover the logical valid length that the fixed-capacity present drops from
// its shape (present aliases past at capacity, so its sequence extent no longer
// encodes the logical length) from the causal/padding bias frontier:
// valid_len = 1 + max{k : bias(b,.,.,k) finite}; write_pos = valid_len - current_seq.
// One thread owns one batch. Mirrors the CPU oracle's `capacity_valid_lens`
// (max over finite columns) and `build_capacity_present` write position exactly.
// Runs on the capturing path too -- the frontier is recomputed from the live
// bias every replay -- so it stays fully on-device with no host round-trip,
// which is what keeps the capacity present capture-safe.
extern "C" __global__ void capacity_write_pos(
    const void* bias, int dtype, int rank,
    unsigned long long bd0, unsigned long long bd1,
    unsigned long long bd2, unsigned long long bd3,
    unsigned long long batch, unsigned long long q_heads, unsigned long long q_seq,
    unsigned long long cache_seq, unsigned long long current_seq,
    long long* valid_len, long long* write_pos) {
  const unsigned long long b =
      (unsigned long long)blockIdx.x * blockDim.x + threadIdx.x;
  if (b >= batch) {
    return;
  }
  unsigned long long valid = 0;
  for (unsigned long long h = 0; h < q_heads; ++h) {
    for (unsigned long long qi = 0; qi < q_seq; ++qi) {
      for (unsigned long long k = 0; k < cache_seq; ++k) {
        const float v =
            bias_at(bias, dtype, rank, bd0, bd1, bd2, bd3, b, h, qi, k);
        if (isfinite(v) && k + 1 > valid) {
          valid = k + 1;
        }
      }
    }
  }
  valid_len[b] = (long long)valid;
  write_pos[b] = (long long)(valid >= current_seq ? valid - current_seq : 0);
}

// Build the fixed-capacity ("in-place") present that ALIASES past at
// `cache_seq` positions: every capacity row is copied from past, then the
// current token(s) overwrite `[write_pos, write_pos + cur_seq)`. Positions at or
// beyond valid_len are never gathered (the selected indices only name positions
// < valid_len), so attention over this layout is byte-identical to the growing
// concat present. `write_pos` is the per-batch device array produced by
// `capacity_write_pos`. Mirrors the CPU oracle's `build_capacity_present`.
extern "C" __global__ void build_present_capacity(
    const void* past, const void* cur, void* out, int dtype,
    const long long* write_pos,
    unsigned long long batch, unsigned long long heads,
    unsigned long long cache_seq, unsigned long long cur_seq,
    unsigned long long dim, unsigned long long elements) {
  for (unsigned long long idx = (unsigned long long)blockIdx.x * blockDim.x + threadIdx.x;
       idx < elements; idx += (unsigned long long)gridDim.x * blockDim.x) {
    unsigned long long d = idx % dim;
    unsigned long long rem = idx / dim;
    unsigned long long t = rem % cache_seq;
    rem /= cache_seq;
    unsigned long long h = rem % heads;
    unsigned long long b = rem / heads;
    const unsigned long long wp = (unsigned long long)write_pos[b];
    float val;
    if (t >= wp && t < wp + cur_seq) {
      const unsigned long long c = t - wp;
      val = load_float(cur, ((b * heads + h) * cur_seq + c) * dim + d, dtype);
    } else {
      val = load_float(past, ((b * heads + h) * cache_seq + t) * dim + d, dtype);
    }
    store_float(out, idx, val, dtype);
  }
}

__device__ __forceinline__ long long load_index(
    const void* indices, unsigned long long offset, int index_is_i64) {
  return index_is_i64
      ? ((const long long*)indices)[offset]
      : (long long)((const int*)indices)[offset];
}

// Device port of the CPU oracle's deterministic `selected_indices` validation.
// One thread owns one [batch, index_head, query] row and scans its
// `selected_width` columns, reproducing the exact rejection rules: an index
// below the -1 sentinel, a non-(-1) index after trailing -1 padding, an index
// outside [0, total_seq), a non-strictly-increasing (or duplicate) index, and
// an all-(-1) row. Any violation latches INDEX_SHARE_CAPTURE_ERROR_INDEX into
// the shared capture-error word via atomicOr; the host reads it back outside the
// captured region. This replaces the host D2H validation on the capture path.
extern "C" __global__ void validate_index_rows(
    const void* indices, unsigned int* capture_error,
    unsigned long long batch, unsigned long long index_heads,
    unsigned long long q_seq, unsigned long long selected_width,
    unsigned long long total_seq, int index_is_i64) {
  const unsigned long long rows = batch * index_heads * q_seq;
  for (unsigned long long row = (unsigned long long)blockIdx.x * blockDim.x + threadIdx.x;
       row < rows; row += (unsigned long long)gridDim.x * blockDim.x) {
    const unsigned long long base = row * selected_width;
    long long previous = 0;
    int seen_padding = 0;
    unsigned long long count = 0;
    int bad = 0;
    for (unsigned long long s = 0; s < selected_width; ++s) {
      const long long index = load_index(indices, base + s, index_is_i64);
      if (index == -1) {
        seen_padding = 1;
        continue;
      }
      if (index < -1) {
        bad = 1;
        break;
      }
      if (seen_padding) {
        bad = 1;
        break;
      }
      if ((unsigned long long)index >= total_seq) {
        bad = 1;
        break;
      }
      if (count > 0 && index <= previous) {
        bad = 1;
        break;
      }
      previous = index;
      count += 1;
    }
    if (count == 0) {
      bad = 1;
    }
    if (bad && capture_error) {
      atomicOr(capture_error, 512u);
    }
  }
}

// One block per (batch, q_head, query) output row. Gathers the selected keys,
// computes scaled QK scores (+ optional bias), a numerically-stable softmax
// over the valid selections, and the probability-weighted value sum.
extern "C" __global__ void index_share_row(
    const void* q, const void* present_k, const void* present_v,
    const void* indices, const void* bias, float* scores, void* y,
    unsigned long long batch, unsigned long long q_heads, unsigned long long kv_heads,
    unsigned long long q_seq, unsigned long long total_seq, unsigned long long head_size,
    unsigned long long index_heads, unsigned long long selected_width,
    unsigned long long group, float sqrt_scale,
    int dtype, int index_is_i64, int has_bias, int bias_rank,
    unsigned long long bd0, unsigned long long bd1,
    unsigned long long bd2, unsigned long long bd3) {
  const unsigned long long row = blockIdx.x;
  const unsigned long long total_rows = batch * q_heads * q_seq;
  if (row >= total_rows) {
    return;
  }
  const unsigned long long qi = row % q_seq;
  unsigned long long rem = row / q_seq;
  const unsigned long long qh = rem % q_heads;
  const unsigned long long b = rem / q_heads;
  const unsigned long long kvh = qh / group;
  const unsigned long long ih = (index_heads == 1ULL) ? 0ULL : qh;
  const unsigned long long index_row =
      ((b * index_heads + ih) * q_seq + qi) * selected_width;
  const unsigned long long score_row = row * selected_width;
  const int tid = threadIdx.x;
  const int nthreads = blockDim.x;

  // Valid count: entries before the trailing -1 padding. Eager host validation
  // and the capturing-path `validate_index_rows` both guarantee
  // strictly-increasing indices with only trailing -1 padding, so counting
  // non-(-1) entries reproduces the CPU take_while.
  __shared__ unsigned long long valid_sh;
  if (tid == 0) {
    unsigned long long valid = 0;
    for (unsigned long long s = 0; s < selected_width; ++s) {
      if (load_index(indices, index_row + s, index_is_i64) == -1) {
        break;
      }
      valid += 1;
    }
    valid_sh = valid;
  }
  __syncthreads();
  const unsigned long long valid = valid_sh;

  const unsigned long long qoff = ((b * q_heads + qh) * q_seq + qi) * head_size;

  // Stage 1: scaled QK score per selected key (sqrt(scale) folded into each
  // operand), plus optional additive bias.
  for (unsigned long long s = tid; s < valid; s += nthreads) {
    const long long raw_key = load_index(indices, index_row + s, index_is_i64);
    // Clamp into [0, total_seq): eager mode has already validated every index,
    // so this is a no-op there. On a poisoned captured replay (caught by
    // validate_index_rows) it keeps the gather in bounds; the tainted output is
    // discarded by the host once the capture-error latch is read.
    const unsigned long long key =
        (raw_key >= 0 && (unsigned long long)raw_key < total_seq)
            ? (unsigned long long)raw_key
            : 0ULL;
    const unsigned long long koff = ((b * kv_heads + kvh) * total_seq + key) * head_size;
    float acc = 0.0f;
    for (unsigned long long d = 0; d < head_size; ++d) {
      acc += (load_float(q, qoff + d, dtype) * sqrt_scale) *
             (load_float(present_k, koff + d, dtype) * sqrt_scale);
    }
    if (has_bias) {
      acc += bias_at(bias, dtype, bias_rank, bd0, bd1, bd2, bd3, b, qh, qi, key);
    }
    scores[score_row + s] = acc;
  }
  __syncthreads();

  // Stage 2: numerically-stable softmax over the valid scores. The lead thread
  // reduces in ascending order to match the CPU reference bit-for-bit.
  __shared__ float inv_sum_sh;
  __shared__ int all_masked_sh;
  if (tid == 0) {
    float m = NEG_INF;
    for (unsigned long long s = 0; s < valid; ++s) {
      m = fmaxf(m, scores[score_row + s]);
    }
    if (m == NEG_INF) {
      all_masked_sh = 1;
      inv_sum_sh = 0.0f;
    } else {
      all_masked_sh = 0;
      float sum = 0.0f;
      for (unsigned long long s = 0; s < valid; ++s) {
        const float e = expf(scores[score_row + s] - m);
        scores[score_row + s] = e;
        sum += e;
      }
      inv_sum_sh = 1.0f / sum;
    }
  }
  __syncthreads();

  const unsigned long long ybase = ((b * q_heads + qh) * q_seq + qi) * head_size;
  if (all_masked_sh) {
    for (unsigned long long d = tid; d < head_size; d += nthreads) {
      store_float(y, ybase + d, 0.0f, dtype);
    }
    return;
  }

  // Normalize probabilities in place (prob = exp * inv_sum), matching the CPU
  // reference which stores the normalized weights before the value reduction.
  const float inv = inv_sum_sh;
  for (unsigned long long s = tid; s < valid; s += nthreads) {
    scores[score_row + s] *= inv;
  }
  __syncthreads();

  // Stage 3: Y = sum_s prob[s] * V[key_s]. Each thread owns whole output
  // channels and sums over selected keys in ascending order.
  for (unsigned long long d = tid; d < head_size; d += nthreads) {
    float acc = 0.0f;
    for (unsigned long long s = 0; s < valid; ++s) {
      const long long raw_key = load_index(indices, index_row + s, index_is_i64);
      const unsigned long long key =
          (raw_key >= 0 && (unsigned long long)raw_key < total_seq)
              ? (unsigned long long)raw_key
              : 0ULL;
      const unsigned long long voff = ((b * kv_heads + kvh) * total_seq + key) * head_size + d;
      acc += scores[score_row + s] * load_float(present_v, voff, dtype);
    }
    store_float(y, ybase + d, acc, dtype);
  }
}
"#;

/// Resolved geometry shared between the CPU reference and this kernel.
#[derive(Clone, Copy)]
struct Dims {
    batch: usize,
    q_heads: usize,
    kv_heads: usize,
    q_seq: usize,
    current_seq: usize,
    past_seq: usize,
    total_seq: usize,
    head_size: usize,
    index_heads: usize,
    selected_width: usize,
    /// Present row stride: `total_seq` for the growing concat present, or the
    /// fixed `past_seq` capacity when the present aliases past in place.
    cache_seq: usize,
    /// The 3-output present aliases the fixed-capacity past bindings in place
    /// (no growing `past ++ current`); the valid length is carried by the bias.
    capacity_mode: bool,
}

/// Right-aligned broadcast metadata for the optional additive bias.
struct BiasMeta {
    ptr: CUdeviceptr,
    present: bool,
    rank: i32,
    dims: [u64; 4],
}

/// Pooled, stable-address device scratch reused across calls so the capturing
/// path never issues a per-call `cudaMalloc`/`cudaFree`. Each buffer grows (a
/// fresh allocation) only in eager mode; a warmed graph replay keeps the exact
/// addresses the warmup pass established. A zero pointer means "not yet
/// allocated"; `capacity` is the current allocation size in bytes.
#[derive(Debug, Default)]
struct ScratchPool {
    present_key: CUdeviceptr,
    present_key_capacity: usize,
    present_value: CUdeviceptr,
    present_value_capacity: usize,
    scores: CUdeviceptr,
    scores_capacity: usize,
    /// Per-batch `[valid_len; write_pos]` (2 * batch i64) recovered on-device
    /// from the bias frontier for the fixed-capacity present path.
    frontier: CUdeviceptr,
    frontier_capacity: usize,
}

impl ScratchPool {
    fn ensure_present_key(
        &mut self,
        runtime: &CudaRuntime,
        bytes: usize,
        capturing: bool,
    ) -> Result<CUdeviceptr> {
        ensure_scratch(
            runtime,
            &mut self.present_key,
            &mut self.present_key_capacity,
            bytes,
            capturing,
            "present_key",
        )
    }

    fn ensure_present_value(
        &mut self,
        runtime: &CudaRuntime,
        bytes: usize,
        capturing: bool,
    ) -> Result<CUdeviceptr> {
        ensure_scratch(
            runtime,
            &mut self.present_value,
            &mut self.present_value_capacity,
            bytes,
            capturing,
            "present_value",
        )
    }

    fn ensure_scores(
        &mut self,
        runtime: &CudaRuntime,
        bytes: usize,
        capturing: bool,
    ) -> Result<CUdeviceptr> {
        ensure_scratch(
            runtime,
            &mut self.scores,
            &mut self.scores_capacity,
            bytes,
            capturing,
            "scores",
        )
    }

    fn ensure_frontier(
        &mut self,
        runtime: &CudaRuntime,
        bytes: usize,
        capturing: bool,
    ) -> Result<CUdeviceptr> {
        ensure_scratch(
            runtime,
            &mut self.frontier,
            &mut self.frontier_capacity,
            bytes,
            capturing,
            "frontier",
        )
    }
}

/// Reuse (or, only in eager mode, grow) a pooled scratch buffer, keeping its
/// device address stable whenever the current capacity already suffices. A grow
/// allocates before freeing the previous block so a failed allocation leaves the
/// pool intact. Growing while capturing is impossible (`cudaMalloc` is illegal
/// during capture), so an under-sized buffer on the capturing path is a hard
/// error pointing at the missing warmup.
fn ensure_scratch(
    runtime: &CudaRuntime,
    ptr: &mut CUdeviceptr,
    capacity: &mut usize,
    bytes: usize,
    capturing: bool,
    what: &str,
) -> Result<CUdeviceptr> {
    let bytes = bytes.max(1);
    if *ptr != 0 && *capacity >= bytes {
        return Ok(*ptr);
    }
    if capturing {
        return Err(error(format!(
            "{what} scratch ({bytes} bytes) exceeds the warmed pool capacity ({} bytes); \
             a fixed-shape eager warmup must run before capture",
            *capacity
        )));
    }
    let fresh = runtime.alloc_raw(bytes)?;
    if *ptr != 0 {
        // SAFETY: the previous pointer came from this runtime's `alloc_raw` and
        // is freed exactly once here, after the replacement is secured.
        unsafe {
            let _ = runtime.free_raw(*ptr);
        }
    }
    *ptr = fresh;
    *capacity = bytes;
    Ok(fresh)
}

pub struct IndexShareFactory {
    pub runtime: Arc<CudaRuntime>,
}

impl KernelFactory for IndexShareFactory {
    fn create(&self, node: &Node, _input_shapes: &[Vec<usize>]) -> Result<Box<dyn Kernel>> {
        let num_heads = required_positive_int(node, "num_heads")?;
        let kv_num_heads = optional_positive_int(node, "kv_num_heads")?.unwrap_or(num_heads);
        if num_heads % kv_num_heads != 0 {
            return Err(error(format!(
                "num_heads {num_heads} must be a multiple of kv_num_heads {kv_num_heads}"
            )));
        }
        let scale = node
            .attr("scale")
            .map(|attribute| {
                attribute
                    .as_float()
                    .ok_or_else(|| error("attribute 'scale' must be a float"))
            })
            .transpose()?;
        if scale.is_some_and(|scale| !scale.is_finite() || scale <= 0.0) {
            return Err(error("attribute 'scale' must be finite and > 0"));
        }
        Ok(Box::new(IndexShareKernel {
            runtime: self.runtime.clone(),
            num_heads,
            kv_num_heads,
            scale,
            scratch: Mutex::new(ScratchPool::default()),
            warmed: AtomicBool::new(false),
        }))
    }
}

#[derive(Debug)]
pub struct IndexShareKernel {
    runtime: Arc<CudaRuntime>,
    num_heads: usize,
    kv_num_heads: usize,
    scale: Option<f32>,
    /// Pooled stable-address scratch (present K/V staging + per-row scores).
    scratch: Mutex<ScratchPool>,
    /// Set after a successful eager execution has compiled every NVRTC kernel
    /// and sized the scratch pool, which is the precondition for capturing this
    /// kernel into a CUDA graph.
    warmed: AtomicBool,
}

impl Drop for IndexShareKernel {
    fn drop(&mut self) {
        let pool = self
            .scratch
            .get_mut()
            .expect("cuda_ep IndexShare scratch pool poisoned");
        for ptr in [
            pool.present_key,
            pool.present_value,
            pool.scores,
            pool.frontier,
        ] {
            if ptr != 0 {
                // SAFETY: every non-zero pointer came from this runtime's
                // `alloc_raw` in `ScratchPool::ensure` and is freed exactly once.
                unsafe {
                    let _ = self.runtime.free_raw(ptr);
                }
            }
        }
    }
}

impl Kernel for IndexShareKernel {
    fn execute(&self, inputs: &[TensorView], outputs: &mut [TensorMut]) -> Result<()> {
        if !(6..=7).contains(&inputs.len()) {
            return Err(error(format!(
                "expected 6 or 7 inputs, got {}",
                inputs.len()
            )));
        }
        if !matches!(outputs.len(), 1 | 3) {
            return Err(error(format!(
                "expected 1 output or 3 outputs (paired present K/V), got {}",
                outputs.len()
            )));
        }
        // Whether the EP stream is recording into a CUDA graph. On the capturing
        // path a stream synchronize and any host D2H copy are illegal, so the
        // input-upload wait below and the deterministic index validation are
        // skipped; same-stream ordering guarantees the inputs are ready and the
        // device-side `validate_index_rows` latch covers correctness instead.
        let capturing = self.runtime.is_capturing()?;
        if !capturing {
            // Inputs may have been uploaded asynchronously on the EP stream.
            self.runtime.synchronize()?;
        }

        for &index in &[0, 1, 2, 5] {
            if inputs[index].is_absent() {
                return Err(error(format!(
                    "required input {index} ('{}') is absent",
                    INPUT_NAMES[index]
                )));
            }
        }
        let has_past_key = optional_input(inputs, 3).is_some();
        let has_past_value = optional_input(inputs, 4).is_some();
        if has_past_key != has_past_value {
            return Err(error("past_key and past_value must be provided together"));
        }
        let dtype = require_floating_dtype(&inputs[0], 0)?;
        for &index in &[1, 2] {
            if inputs[index].dtype != dtype {
                return Err(error(
                    "query, key, and value must use the same floating dtype",
                ));
            }
        }
        for index in [3, 4, 6] {
            if let Some(input) = optional_input(inputs, index)
                && input.dtype != dtype
            {
                return Err(error(
                    "query, key, value, past_key, past_value, and attention_bias must use the same floating dtype",
                ));
            }
        }
        if !matches!(inputs[5].dtype, DataType::Int32 | DataType::Int64) {
            return Err(error(format!(
                "input 5 ('selected_indices') dtype {:?} unsupported; expected Int32 or Int64",
                inputs[5].dtype
            )));
        }
        for (index, output) in outputs.iter().enumerate() {
            if output.dtype != dtype {
                return Err(error(format!(
                    "output {index} dtype {:?} must match query dtype {dtype:?}",
                    output.dtype,
                )));
            }
        }
        for &index in &[0, 1, 2, 5] {
            if !inputs[index].is_contiguous() {
                return Err(error(format!(
                    "input {index} ('{}') must be contiguous",
                    INPUT_NAMES[index]
                )));
            }
        }
        for index in [3, 4, 6] {
            if let Some(input) = optional_input(inputs, index)
                && !input.is_contiguous()
            {
                return Err(error(format!(
                    "input {index} ('{}') must be contiguous",
                    INPUT_NAMES[index]
                )));
            }
        }
        for output in outputs.iter() {
            if !output.is_contiguous() {
                return Err(error("outputs must be contiguous"));
            }
        }

        let dims = self.validate_shapes(inputs, outputs)?;

        // The capacity present aliases the fixed-capacity past bindings in place,
        // so its shape no longer encodes the logical length; the causal/padding
        // attention_bias frontier carries it (mirrors default-domain Attention at
        // fixed capacity). Reject the mode when that signal is absent.
        if dims.capacity_mode && optional_input(inputs, 6).is_none() {
            return Err(error(
                "capacity-mode IndexShare (present aliases fixed-capacity past) requires attention_bias to carry the valid length",
            ));
        }

        // Eager (non-capturing) execution mirrors the CPU oracle: copy the small
        // index tensor D2H and reject any malformed row synchronously. Capturing
        // execution cannot round-trip to the host, so this is skipped and the
        // device-side `validate_index_rows` latch (checked outside the captured
        // region) enforces the same contract instead. The capacity path validates
        // against the bias-derived per-batch valid length recovered on-device
        // below, so it defers this to inside the launch sequence.
        if !capturing && !dims.capacity_mode {
            let indices = self.read_indices(&inputs[5], dims)?;
            validate_indices(&indices, dims, &vec![dims.total_seq; dims.batch])?;
        }

        let bias = self.bias_meta(inputs, dims)?;

        let q_ptr = cuptr(inputs[0].data_ptr::<u8>() as *const c_void);
        let key_ptr = cuptr(inputs[1].data_ptr::<u8>() as *const c_void);
        let value_ptr = cuptr(inputs[2].data_ptr::<u8>() as *const c_void);
        let past_key_ptr = optional_input(inputs, 3)
            .map(|view| cuptr(view.data_ptr::<u8>() as *const c_void))
            .unwrap_or(0);
        let past_value_ptr = optional_input(inputs, 4)
            .map(|view| cuptr(view.data_ptr::<u8>() as *const c_void))
            .unwrap_or(0);
        let indices_ptr = cuptr(inputs[5].data_ptr::<u8>() as *const c_void);
        let index_is_i64 = i32::from(inputs[5].dtype == DataType::Int64);

        // Present K/V element counts and the output element count. `cache_seq`
        // is `total_seq` for the growing concat present and the fixed capacity
        // when the present aliases past in place.
        let present_elements = dims.batch * dims.kv_heads * dims.cache_seq * dims.head_size;
        let output_elements = dims.batch * dims.q_heads * dims.q_seq * dims.head_size;
        let scores_elements = dims.batch * dims.q_heads * dims.q_seq * dims.selected_width;

        // Present outputs are written directly into the caller's output slots
        // when requested (outputs 1 and 2); otherwise present K/V land in scratch
        // that only feeds the attention kernel.
        let want_present = outputs.len() == 3;
        let (output_head, output_tail) = outputs.split_at_mut(1);
        let y_ptr = cuptr(output_head[0].data_ptr_mut::<u8>() as *const c_void);
        let (present_key_out, present_value_out) = if want_present {
            (
                cuptr(output_tail[0].data_ptr_mut::<u8>() as *const c_void),
                cuptr(output_tail[1].data_ptr_mut::<u8>() as *const c_void),
            )
        } else {
            (0, 0)
        };

        let result = (|| -> Result<()> {
            let mut pool = self
                .scratch
                .lock()
                .expect("cuda_ep IndexShare scratch pool poisoned");

            // Present K/V land in the caller's output slots when requested;
            // otherwise they use pooled scratch that only feeds the row kernel.
            let present_key_ptr = if want_present {
                present_key_out
            } else {
                pool.ensure_present_key(
                    &self.runtime,
                    present_elements * dtype.storage_bytes(1),
                    capturing,
                )?
            };
            let present_value_ptr = if want_present {
                present_value_out
            } else {
                pool.ensure_present_value(
                    &self.runtime,
                    present_elements * dtype.storage_bytes(1),
                    capturing,
                )?
            };
            let scores_ptr = pool.ensure_scores(&self.runtime, scores_elements * 4, capturing)?;

            if dims.capacity_mode {
                // Recover the per-batch valid length + current-token write
                // position from the bias frontier on-device (capture-safe), then
                // build the fixed-capacity present that aliases past in place.
                let frontier_ptr = pool.ensure_frontier(
                    &self.runtime,
                    2 * dims.batch * std::mem::size_of::<i64>(),
                    capturing,
                )?;
                let valid_len_ptr = frontier_ptr;
                let write_pos_ptr = frontier_ptr + (dims.batch * std::mem::size_of::<i64>()) as u64;
                self.launch_capacity_write_pos(
                    &bias,
                    dims,
                    dtype_code(dtype)?,
                    valid_len_ptr,
                    write_pos_ptr,
                )?;
                self.build_present_capacity(
                    past_key_ptr,
                    key_ptr,
                    present_key_ptr,
                    write_pos_ptr,
                    dims,
                    dtype_code(dtype)?,
                )?;
                self.build_present_capacity(
                    past_value_ptr,
                    value_ptr,
                    present_value_ptr,
                    write_pos_ptr,
                    dims,
                    dtype_code(dtype)?,
                )?;
                // Eager: validate the selected indices against the bias-derived
                // per-batch valid length, matching the CPU oracle exactly.
                if !capturing {
                    self.runtime.synchronize()?;
                    let mut valid_raw = vec![0u8; dims.batch * std::mem::size_of::<i64>()];
                    // SAFETY: `valid_len_ptr` is a live pooled allocation of at
                    // least `batch` i64 written by `capacity_write_pos` above.
                    unsafe {
                        self.runtime.dtoh(&mut valid_raw, valid_len_ptr)?;
                    }
                    let valid_lens: Vec<usize> = valid_raw
                        .chunks_exact(std::mem::size_of::<i64>())
                        .map(|raw| i64::from_ne_bytes(raw.try_into().unwrap()).max(0) as usize)
                        .collect();
                    let indices = self.read_indices(&inputs[5], dims)?;
                    validate_indices(&indices, dims, &valid_lens)?;
                }
            } else {
                self.build_present(
                    past_key_ptr,
                    key_ptr,
                    present_key_ptr,
                    has_past_key,
                    dims,
                    dtype_code(dtype)?,
                )?;
                self.build_present(
                    past_value_ptr,
                    value_ptr,
                    present_value_ptr,
                    has_past_value,
                    dims,
                    dtype_code(dtype)?,
                )?;
            }

            // On the capturing path record the device-side index validation so a
            // poisoned replay latches the shared capture-error word (read by the
            // host at the per-step logits sync, outside the captured region). The
            // bound is the `cache_seq` capacity (a safe superset of the per-batch
            // valid length for the capacity path; identical to `total_seq` for
            // the concat path).
            if capturing {
                self.launch_index_validation(indices_ptr, dims, index_is_i64)?;
            }

            self.launch_rows(
                q_ptr,
                present_key_ptr,
                present_value_ptr,
                indices_ptr,
                &bias,
                scores_ptr,
                y_ptr,
                dims,
                dtype_code(dtype)?,
                index_is_i64,
                output_elements,
            )?;
            if capturing {
                Ok(())
            } else {
                self.runtime.synchronize()
            }
        })();

        if result.is_ok() && !capturing {
            // A warmed eager pass has compiled every kernel and sized the pooled
            // scratch: capture may now record this kernel with stable addresses.
            self.warmed.store(true, Ordering::Relaxed);
        }
        result
    }

    fn supports_strided_input(&self, _index: usize) -> bool {
        false
    }

    fn capture_support(&self) -> CaptureSupport {
        if self.warmed.load(Ordering::Relaxed) {
            CaptureSupport::Supported
        } else {
            CaptureSupport::unsupported(
                "requires a warmed fixed-shape eager IndexShare pass to size the pooled scratch and \
                 prime device-side selected_indices validation",
            )
        }
    }
}

impl IndexShareKernel {
    fn validate_shapes(&self, inputs: &[TensorView], outputs: &[TensorMut]) -> Result<Dims> {
        for &index in &[0, 1, 2, 5] {
            require_rank(index, inputs[index].shape)?;
        }
        for index in [3, 4] {
            if let Some(input) = optional_input(inputs, index) {
                require_rank(index, input.shape)?;
            }
        }
        let q = inputs[0].shape;
        let key = inputs[1].shape;
        let value = inputs[2].shape;
        let (batch, q_heads, q_seq, head_size) = (q[0], q[1], q[2], q[3]);
        if q_heads != self.num_heads {
            return Err(error(format!(
                "query head dimension {q_heads} must equal num_heads {}",
                self.num_heads
            )));
        }
        if key[0] != batch || value[0] != batch {
            return Err(error("query, key, and value batch dimensions must match"));
        }
        if key[1] != self.kv_num_heads || value[1] != self.kv_num_heads {
            return Err(error(format!(
                "key/value head dimensions must equal kv_num_heads {}",
                self.kv_num_heads
            )));
        }
        if key[2] != value[2] || key[3] != head_size || value[3] != head_size {
            return Err(error(
                "key/value sequence and head dimensions must match query/schema",
            ));
        }
        let current_seq = key[2];
        let mut past_seq = 0;
        if let (Some(past_key), Some(past_value)) =
            (optional_input(inputs, 3), optional_input(inputs, 4))
        {
            if past_key.shape != past_value.shape {
                return Err(error("past_key and past_value shapes must match"));
            }
            if past_key.shape[0] != batch
                || past_key.shape[1] != self.kv_num_heads
                || past_key.shape[3] != head_size
            {
                return Err(error(
                    "past key/value must have shape [B, kv_num_heads, S_past, H]",
                ));
            }
            past_seq = past_key.shape[2];
        }
        let total_seq = past_seq
            .checked_add(current_seq)
            .ok_or_else(|| error("total cache sequence length overflow"))?;
        let selected = inputs[5].shape;
        let index_heads = selected[1];
        if selected[0] != batch
            || (index_heads != 1 && index_heads != q_heads)
            || selected[2] != q_seq
        {
            return Err(error(format!(
                "selected_indices must have shape [B, 1|N, S_q, K], got {selected:?}"
            )));
        }
        if selected[3] == 0 {
            return Err(error("selected_indices K dimension must be nonzero"));
        }
        if outputs[0].shape != q {
            return Err(error(format!(
                "output shape {:?} must equal query shape {q:?}",
                outputs[0].shape
            )));
        }
        let mut cache_seq = total_seq;
        let mut capacity_mode = false;
        if outputs.len() == 3 {
            let concat = [batch, self.kv_num_heads, total_seq, head_size];
            let capacity = [batch, self.kv_num_heads, past_seq, head_size];
            // The present may either grow (`past ++ current`, sequence ==
            // total_seq) or alias the fixed-capacity `past` in place (sequence
            // == past_seq, requires a past cache). `past_seq < total_seq` always
            // (current_seq >= 1), so the two shapes are unambiguous and no
            // existing concat caller trips the capacity branch.
            if outputs[1].shape == concat && outputs[2].shape == concat {
                // Growing concat present.
            } else if past_seq > 0 && outputs[1].shape == capacity && outputs[2].shape == capacity {
                capacity_mode = true;
                cache_seq = past_seq;
            } else {
                return Err(error(format!(
                    "present_key and present_value shapes must be {concat:?} (growing) or {capacity:?} (fixed capacity)"
                )));
            }
        }
        // The bias spans the gathered cache: `total_seq` for the concat present,
        // or the fixed `cache_seq` capacity when the present aliases past.
        if let Some(bias) = optional_input(inputs, 6) {
            validate_bias_shape(bias.shape, [batch, q_heads, q_seq, cache_seq])?;
        }
        Ok(Dims {
            batch,
            q_heads,
            kv_heads: self.kv_num_heads,
            q_seq,
            current_seq,
            past_seq,
            total_seq,
            head_size,
            index_heads,
            selected_width: selected[3],
            cache_seq,
            capacity_mode,
        })
    }

    fn read_indices(&self, view: &TensorView, dims: Dims) -> Result<Vec<i64>> {
        let count = dims.batch * dims.index_heads * dims.q_seq * dims.selected_width;
        let byte_len = view.dtype.storage_bytes(count);
        let mut host = vec![0u8; byte_len];
        if !host.is_empty() {
            // SAFETY: `view` is a live contiguous device tensor and the host
            // buffer is exactly its fixed-width storage size.
            unsafe {
                self.runtime
                    .dtoh(&mut host, cuptr(view.data_ptr::<u8>() as *const c_void))?;
            }
        }
        Ok(host
            .chunks_exact(view.dtype.byte_size())
            .map(|raw| match view.dtype {
                DataType::Int32 => i32::from_ne_bytes(raw.try_into().unwrap()) as i64,
                DataType::Int64 => i64::from_ne_bytes(raw.try_into().unwrap()),
                _ => unreachable!("index dtype was validated"),
            })
            .collect())
    }

    fn bias_meta(&self, inputs: &[TensorView], dims: Dims) -> Result<BiasMeta> {
        match optional_input(inputs, 6) {
            Some(view) => {
                let rank = view.shape.len();
                if rank > 4 {
                    return Err(error(format!("attention_bias rank {rank} exceeds 4")));
                }
                let expected = view.numel();
                let actual = view.shape.iter().product::<usize>();
                if expected != actual {
                    return Err(error("attention_bias element count mismatch"));
                }
                let _ = dims;
                let mut broadcast = [1u64; 4];
                for (axis, &dim) in view.shape.iter().enumerate() {
                    broadcast[4 - rank + axis] = dim as u64;
                }
                Ok(BiasMeta {
                    ptr: cuptr(view.data_ptr::<u8>() as *const c_void),
                    present: true,
                    rank: rank as i32,
                    dims: broadcast,
                })
            }
            None => Ok(BiasMeta {
                ptr: 0,
                present: false,
                rank: 0,
                dims: [1u64; 4],
            }),
        }
    }

    /// Record the device-side deterministic `selected_indices` validation on the
    /// EP stream. One thread scans each `[batch, index_head, query]` row and
    /// latches [`INDEX_SHARE_CAPTURE_ERROR_INDEX`] into the runtime capture-error
    /// word on any violation. Used only on the capturing path, where the host
    /// D2H validation is illegal; the latch is read back outside the captured
    /// region so a poisoned replay is rejected before its token is consumed.
    fn launch_index_validation(
        &self,
        indices_ptr: CUdeviceptr,
        dims: Dims,
        index_is_i64: i32,
    ) -> Result<()> {
        let rows = (dims.batch * dims.index_heads * dims.q_seq) as u64;
        if rows == 0 {
            return Ok(());
        }
        let func = self
            .runtime
            .nvrtc_function(MODULE, SOURCE, "validate_index_rows")?;
        let capture_error = self.runtime.capture_error_ptr();
        let batch = dims.batch as u64;
        let index_heads = dims.index_heads as u64;
        let q_seq = dims.q_seq as u64;
        let selected_width = dims.selected_width as u64;
        let total_seq = dims.cache_seq as u64;
        let mut builder = self.runtime.stream().launch_builder(&func);
        builder
            .arg(&indices_ptr)
            .arg(&capture_error)
            .arg(&batch)
            .arg(&index_heads)
            .arg(&q_seq)
            .arg(&selected_width)
            .arg(&total_seq)
            .arg(&index_is_i64);
        // SAFETY: argument types/order match `validate_index_rows`; the index
        // tensor is a live contiguous device allocation and `capture_error` is
        // the runtime's persistent four-byte latch word.
        unsafe {
            builder.launch(LaunchConfig {
                grid_dim: (rows.div_ceil(BLOCK as u64).clamp(1, 65_535) as u32, 1, 1),
                block_dim: (BLOCK, 1, 1),
                shared_mem_bytes: 0,
            })
        }
        .map_err(|e| driver_err("launch validate_index_rows", e))
        .map(|_| ())
    }

    fn build_present(
        &self,
        past_ptr: CUdeviceptr,
        current_ptr: CUdeviceptr,
        out_ptr: CUdeviceptr,
        has_past: bool,
        dims: Dims,
        dtype: i32,
    ) -> Result<()> {
        let elements = (dims.batch * dims.kv_heads * dims.total_seq * dims.head_size) as u64;
        if elements == 0 {
            return Ok(());
        }
        let func = self
            .runtime
            .nvrtc_function(MODULE, SOURCE, "build_present")?;
        let has_past_i = i32::from(has_past);
        let batch = dims.batch as u64;
        let heads = dims.kv_heads as u64;
        let past_seq = dims.past_seq as u64;
        let cur_seq = dims.current_seq as u64;
        let total_seq = dims.total_seq as u64;
        let head_size = dims.head_size as u64;
        let mut builder = self.runtime.stream().launch_builder(&func);
        builder
            .arg(&past_ptr)
            .arg(&current_ptr)
            .arg(&out_ptr)
            .arg(&dtype)
            .arg(&has_past_i)
            .arg(&batch)
            .arg(&heads)
            .arg(&past_seq)
            .arg(&cur_seq)
            .arg(&total_seq)
            .arg(&head_size)
            .arg(&elements);
        // SAFETY: argument types/order match `build_present`; all pointers refer
        // to live contiguous device allocations validated above.
        unsafe {
            builder.launch(LaunchConfig {
                grid_dim: (
                    elements.div_ceil(BLOCK as u64).clamp(1, 65_535) as u32,
                    1,
                    1,
                ),
                block_dim: (BLOCK, 1, 1),
                shared_mem_bytes: 0,
            })
        }
        .map_err(|e| driver_err("launch build_present", e))
        .map(|_| ())
    }

    /// Recover the per-batch valid length and current-token write position from
    /// the bias frontier on-device (capture-safe: no host round-trip), writing
    /// `valid_len` and `write_pos` (each `batch` i64) into pooled scratch.
    fn launch_capacity_write_pos(
        &self,
        bias: &BiasMeta,
        dims: Dims,
        dtype: i32,
        valid_len_ptr: CUdeviceptr,
        write_pos_ptr: CUdeviceptr,
    ) -> Result<()> {
        if !bias.present {
            return Err(error(
                "capacity-mode IndexShare requires attention_bias to derive the valid length",
            ));
        }
        let batch = dims.batch as u64;
        if batch == 0 {
            return Ok(());
        }
        let func = self
            .runtime
            .nvrtc_function(MODULE, SOURCE, "capacity_write_pos")?;
        let rank = bias.rank;
        let (bd0, bd1, bd2, bd3) = (bias.dims[0], bias.dims[1], bias.dims[2], bias.dims[3]);
        let q_heads = dims.q_heads as u64;
        let q_seq = dims.q_seq as u64;
        let cache_seq = dims.cache_seq as u64;
        let current_seq = dims.current_seq as u64;
        let mut builder = self.runtime.stream().launch_builder(&func);
        builder
            .arg(&bias.ptr)
            .arg(&dtype)
            .arg(&rank)
            .arg(&bd0)
            .arg(&bd1)
            .arg(&bd2)
            .arg(&bd3)
            .arg(&batch)
            .arg(&q_heads)
            .arg(&q_seq)
            .arg(&cache_seq)
            .arg(&current_seq)
            .arg(&valid_len_ptr)
            .arg(&write_pos_ptr);
        // SAFETY: argument types/order match `capacity_write_pos`; the bias is a
        // live contiguous device allocation and the two output pointers are the
        // pooled `batch`-i64 halves of the frontier scratch.
        unsafe {
            builder.launch(LaunchConfig {
                grid_dim: (batch.div_ceil(BLOCK as u64).clamp(1, 65_535) as u32, 1, 1),
                block_dim: (BLOCK, 1, 1),
                shared_mem_bytes: 0,
            })
        }
        .map_err(|e| driver_err("launch capacity_write_pos", e))
        .map(|_| ())
    }

    /// Build the fixed-capacity present that aliases `past` at `cache_seq`
    /// positions, overwriting the current token(s) at the per-batch `write_pos`
    /// produced by [`Self::launch_capacity_write_pos`]. Byte-identical to the CPU
    /// oracle's `build_capacity_present`.
    fn build_present_capacity(
        &self,
        past_ptr: CUdeviceptr,
        current_ptr: CUdeviceptr,
        out_ptr: CUdeviceptr,
        write_pos_ptr: CUdeviceptr,
        dims: Dims,
        dtype: i32,
    ) -> Result<()> {
        let elements = (dims.batch * dims.kv_heads * dims.cache_seq * dims.head_size) as u64;
        if elements == 0 {
            return Ok(());
        }
        let func = self
            .runtime
            .nvrtc_function(MODULE, SOURCE, "build_present_capacity")?;
        let batch = dims.batch as u64;
        let heads = dims.kv_heads as u64;
        let cache_seq = dims.cache_seq as u64;
        let cur_seq = dims.current_seq as u64;
        let head_size = dims.head_size as u64;
        let mut builder = self.runtime.stream().launch_builder(&func);
        builder
            .arg(&past_ptr)
            .arg(&current_ptr)
            .arg(&out_ptr)
            .arg(&dtype)
            .arg(&write_pos_ptr)
            .arg(&batch)
            .arg(&heads)
            .arg(&cache_seq)
            .arg(&cur_seq)
            .arg(&head_size)
            .arg(&elements);
        // SAFETY: argument types/order match `build_present_capacity`; all
        // pointers refer to live contiguous device allocations and `write_pos`
        // holds `batch` i64 written by `capacity_write_pos`.
        unsafe {
            builder.launch(LaunchConfig {
                grid_dim: (
                    elements.div_ceil(BLOCK as u64).clamp(1, 65_535) as u32,
                    1,
                    1,
                ),
                block_dim: (BLOCK, 1, 1),
                shared_mem_bytes: 0,
            })
        }
        .map_err(|e| driver_err("launch build_present_capacity", e))
        .map(|_| ())
    }

    #[allow(clippy::too_many_arguments)]
    fn launch_rows(
        &self,
        q_ptr: CUdeviceptr,
        present_key_ptr: CUdeviceptr,
        present_value_ptr: CUdeviceptr,
        indices_ptr: CUdeviceptr,
        bias: &BiasMeta,
        scores_ptr: CUdeviceptr,
        y_ptr: CUdeviceptr,
        dims: Dims,
        dtype: i32,
        index_is_i64: i32,
        output_elements: usize,
    ) -> Result<()> {
        let total_rows = (dims.batch * dims.q_heads * dims.q_seq) as u64;
        if total_rows == 0 || output_elements == 0 {
            return Ok(());
        }
        let func = self
            .runtime
            .nvrtc_function(MODULE, SOURCE, "index_share_row")?;
        let scale = self
            .scale
            .unwrap_or_else(|| 1.0 / (dims.head_size as f32).sqrt());
        let sqrt_scale = scale.sqrt();
        let group = (dims.q_heads / dims.kv_heads) as u64;
        let batch = dims.batch as u64;
        let q_heads = dims.q_heads as u64;
        let kv_heads = dims.kv_heads as u64;
        let q_seq = dims.q_seq as u64;
        // Present row stride/gather bound: `cache_seq` (== total_seq for the
        // concat present, the fixed capacity when present aliases past).
        let total_seq = dims.cache_seq as u64;
        let head_size = dims.head_size as u64;
        let index_heads = dims.index_heads as u64;
        let selected_width = dims.selected_width as u64;
        let has_bias = i32::from(bias.present);
        let bias_rank = bias.rank;
        let (bd0, bd1, bd2, bd3) = (bias.dims[0], bias.dims[1], bias.dims[2], bias.dims[3]);
        let mut builder = self.runtime.stream().launch_builder(&func);
        builder
            .arg(&q_ptr)
            .arg(&present_key_ptr)
            .arg(&present_value_ptr)
            .arg(&indices_ptr)
            .arg(&bias.ptr)
            .arg(&scores_ptr)
            .arg(&y_ptr)
            .arg(&batch)
            .arg(&q_heads)
            .arg(&kv_heads)
            .arg(&q_seq)
            .arg(&total_seq)
            .arg(&head_size)
            .arg(&index_heads)
            .arg(&selected_width)
            .arg(&group)
            .arg(&sqrt_scale)
            .arg(&dtype)
            .arg(&index_is_i64)
            .arg(&has_bias)
            .arg(&bias_rank)
            .arg(&bd0)
            .arg(&bd1)
            .arg(&bd2)
            .arg(&bd3);
        // SAFETY: argument types/order match `index_share_row`; all pointers
        // refer to live contiguous device allocations, the scores scratch is
        // sized for `batch*q_heads*q_seq*selected_width` f32, and every index was
        // range-checked on the host.
        unsafe {
            builder.launch(LaunchConfig {
                grid_dim: (total_rows.min(u32::MAX as u64).max(1) as u32, 1, 1),
                block_dim: (ROW_THREADS, 1, 1),
                shared_mem_bytes: 0,
            })
        }
        .map_err(|e| driver_err("launch index_share_row", e))
        .map(|_| ())
    }
}

fn validate_indices(indices: &[i64], dims: Dims, per_batch_bound: &[usize]) -> Result<()> {
    for (b, &bound) in per_batch_bound.iter().enumerate() {
        for h in 0..dims.index_heads {
            for q in 0..dims.q_seq {
                let row = ((b * dims.index_heads + h) * dims.q_seq + q) * dims.selected_width;
                let mut previous = None;
                let mut padding = false;
                let mut count = 0;
                for (column, &index) in indices[row..row + dims.selected_width].iter().enumerate() {
                    if index == -1 {
                        padding = true;
                        continue;
                    }
                    if index < -1 {
                        return Err(index_error(
                            b,
                            h,
                            q,
                            column,
                            format!("invalid sentinel {index}"),
                        ));
                    }
                    if padding {
                        return Err(index_error(
                            b,
                            h,
                            q,
                            column,
                            format!("index {index} follows trailing -1 padding"),
                        ));
                    }
                    if index as usize >= bound {
                        return Err(index_error(
                            b,
                            h,
                            q,
                            column,
                            format!("index {index} is out of range for cache length {bound}"),
                        ));
                    }
                    if let Some(previous) = previous
                        && index <= previous
                    {
                        let reason = if index == previous {
                            format!("duplicate index {index}")
                        } else {
                            format!("indices are not strictly increasing: {previous} then {index}")
                        };
                        return Err(index_error(b, h, q, column, reason));
                    }
                    previous = Some(index);
                    count += 1;
                }
                if count == 0 {
                    return Err(error(format!(
                        "selected_indices row [batch={b}, head={h}, query={q}] is all -1"
                    )));
                }
            }
        }
    }
    Ok(())
}

fn validate_bias_shape(shape: &[usize], target: [usize; 4]) -> Result<()> {
    if shape.len() > 4 {
        return Err(error(format!(
            "attention_bias rank {} exceeds 4",
            shape.len()
        )));
    }
    for (axis, &dimension) in shape.iter().enumerate() {
        let expected = target[4 - shape.len() + axis];
        if dimension != 1 && dimension != expected {
            return Err(error(format!(
                "attention_bias dimension {dimension} is not broadcastable to {target:?}"
            )));
        }
    }
    Ok(())
}

fn index_error(batch: usize, head: usize, query: usize, column: usize, reason: String) -> EpError {
    error(format!(
        "selected_indices [batch={batch}, head={head}, query={query}, column={column}]: {reason}"
    ))
}

fn require_rank(index: usize, shape: &[usize]) -> Result<()> {
    if shape.len() != 4 {
        return Err(error(format!(
            "input {index} ('{}') rank {} unsupported; expected 4",
            INPUT_NAMES[index],
            shape.len()
        )));
    }
    Ok(())
}

fn require_floating_dtype(input: &TensorView, index: usize) -> Result<DataType> {
    if !matches!(
        input.dtype,
        DataType::Float32 | DataType::Float16 | DataType::BFloat16
    ) {
        return Err(error(format!(
            "input {index} ('{}') dtype {:?} unsupported; expected Float32, Float16, or BFloat16",
            INPUT_NAMES[index], input.dtype
        )));
    }
    Ok(input.dtype)
}

fn dtype_code(dtype: DataType) -> Result<i32> {
    match dtype {
        DataType::Float32 => Ok(0),
        DataType::Float16 => Ok(1),
        DataType::BFloat16 => Ok(2),
        _ => Err(error(format!("unsupported floating dtype {dtype:?}"))),
    }
}

fn required_positive_int(node: &Node, name: &str) -> Result<usize> {
    let value = node
        .attr(name)
        .ok_or_else(|| error(format!("missing required integer attribute '{name}'")))?
        .as_int()
        .ok_or_else(|| error(format!("attribute '{name}' must be an integer")))?;
    usize::try_from(value)
        .ok()
        .filter(|&value| value > 0)
        .ok_or_else(|| error(format!("attribute '{name}' must be > 0")))
}

fn optional_positive_int(node: &Node, name: &str) -> Result<Option<usize>> {
    node.attr(name)
        .map(|attribute| {
            let value = attribute
                .as_int()
                .ok_or_else(|| error(format!("attribute '{name}' must be an integer")))?;
            usize::try_from(value)
                .ok()
                .filter(|&value| value > 0)
                .ok_or_else(|| error(format!("attribute '{name}' must be > 0")))
        })
        .transpose()
}

fn optional_input<'a>(inputs: &'a [TensorView<'a>], index: usize) -> Option<&'a TensorView<'a>> {
    inputs.get(index).filter(|input| !input.is_absent())
}

fn error(message: impl Into<String>) -> EpError {
    EpError::KernelFailed(format!("cuda_ep {OP}: {}", message.into()))
}