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

//! CUDA implementation of `com.microsoft::GroupQueryAttention`.
//!
//! BSH query/key/value inputs are prepared into BNSH buffers with NVRTC kernels,
//! including cache append and optional RoPE. Multi-token prefill then uses the
//! shared tiled online-softmax flash kernel when its measured shape gate wins;
//! decode and unsupported/slower shapes retain the existing attention baseline.
//! Present key/value outputs remain BNSH and preserve a fixed cache capacity.

use std::ffi::c_void;
use std::sync::{Arc, Mutex};

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

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

use super::attention::{AttentionDtype, run_attention_phase2a};
use super::flash_attention;
use super::gqa_decode;
use super::gqa_decode_fp16;

const PREP_SRC: &str = r#"
extern "C" __global__ void gqa_transpose_bsh_to_bnsh(
    const float* src, float* dst, int batch, int seq, int heads, int dim)
{
    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
    const int count = batch * heads * seq * dim;
    if (idx >= count) return;
    int x = idx;
    const int d = x % dim; x /= dim;
    const int s = x % seq; x /= seq;
    const int h = x % heads; const int b = x / heads;
    dst[idx] = src[((b * seq + s) * heads + h) * dim + d];
}

extern "C" __global__ void gqa_split_packed_qkv(
    const float* packed, float* query, float* key, float* value,
    int batch, int seq, int q_heads, int kv_heads, int dim)
{
    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
    const int q_hidden = q_heads * dim;
    const int kv_hidden = kv_heads * dim;
    const int packed_hidden = q_hidden + 2 * kv_hidden;
    const int count = batch * seq * packed_hidden;
    if (idx >= count) return;
    const int feature = idx % packed_hidden;
    const int token = idx / packed_hidden;
    if (feature < q_hidden) {
        query[token * q_hidden + feature] = packed[idx];
    } else if (feature < q_hidden + kv_hidden) {
        key[token * kv_hidden + feature - q_hidden] = packed[idx];
    } else {
        value[token * kv_hidden + feature - q_hidden - kv_hidden] = packed[idx];
    }
}

extern "C" __global__ void gqa_prepare_metadata(
    const int* seqlens_k, int* total_lengths, int* past_lengths,
    int* query_starts, int batch, int current_key_length, int query_length,
    int past_capacity, int present_capacity, const long long* position_ids,
    int validate_positions, int cache_rows, int* error_flag)
{
    const int b = blockIdx.x * blockDim.x + threadIdx.x;
    if (b >= batch) return;
    // Latch-first poison propagation: if any prior captured kernel (an earlier
    // step, an earlier replay, or an earlier layer in this replay) already
    // recorded a bounds violation, force this step into the sentinel/skip state
    // deterministically. A later in-range step can therefore never resume writes
    // over a KV row that a poisoned step skipped, so no stale hole can form.
    // `atomicOr(flag, 0)` performs a coherent global read without mutating it.
    if (error_flag && atomicOr(error_flag, 0) != 0) {
        total_lengths[b] = -1;
        past_lengths[b] = -1;
        query_starts[b] = -1;
        return;
    }
    const long long total = (long long)seqlens_k[b] + 1;
    const long long past = total - current_key_length;
    const long long query_start = total - query_length;
    int error = 0;
    if (total > 2147483647LL) error |= 1;
    if (past < 0) error |= 2;
    if (query_start < 0) error |= 4;
    if (past > past_capacity) error |= 8;
    if (total > present_capacity) error |= 16;
    if (validate_positions) {
        for (int s = 0; s < query_length; ++s) {
            const long long position = position_ids
                ? position_ids[b * query_length + s]
                : query_start + s;
            if (position < 0 || position >= (long long)cache_rows) {
                error |= 32;
            }
        }
    }
    if (error) {
        if (error_flag) atomicOr(error_flag, error);
        total_lengths[b] = -1;
        past_lengths[b] = -1;
        query_starts[b] = -1;
        return;
    }
    total_lengths[b] = (int)total;
    past_lengths[b] = (int)past;
    query_starts[b] = (int)query_start;
}

__device__ __forceinline__ int gqa_prepare_metadata_batch1(
    const int* seqlens_k, int* total_lengths, int* past_lengths,
    int* query_starts, int past_capacity, int present_capacity,
    const long long* position_ids, int validate_positions, int cache_rows,
    int* error_flag, int write_metadata)
{
    if (error_flag && atomicOr(error_flag, 0) != 0) {
        if (write_metadata) {
            total_lengths[0] = -1;
            past_lengths[0] = -1;
            query_starts[0] = -1;
        }
        return -1;
    }
    const long long total = (long long)seqlens_k[0] + 1;
    const long long past = total - 1;
    const long long query_start = total - 1;
    int error = 0;
    if (total > 2147483647LL) error |= 1;
    if (past < 0) error |= 2;
    if (query_start < 0) error |= 4;
    if (past > past_capacity) error |= 8;
    if (total > present_capacity) error |= 16;
    if (validate_positions) {
        const long long position = position_ids ? position_ids[0] : past;
        if (position < 0 || position >= (long long)cache_rows) error |= 32;
    }
    if (error) {
        if (error_flag) atomicOr(error_flag, error);
        if (write_metadata) {
            total_lengths[0] = -1;
            past_lengths[0] = -1;
            query_starts[0] = -1;
        }
        return -1;
    }
    if (write_metadata) {
        total_lengths[0] = (int)total;
        past_lengths[0] = (int)past;
        query_starts[0] = (int)query_start;
    }
    return (int)past;
}

extern "C" __global__ void gqa_build_cache(
    const float* current, const float* past, float* present,
    const int* past_lengths, int batch, int seq, int heads, int dim,
    int past_capacity, int present_capacity)
{
    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
    const int count = batch * heads * present_capacity * dim;
    if (idx >= count) return;
    int x = idx;
    const int d = x % dim; x /= dim;
    const int s = x % present_capacity; x /= present_capacity;
    const int h = x % heads; const int b = x / heads;
    const int past_len = past_lengths[b];
    if (past_len < 0) return;
    float value = 0.0f;
    if (s < past_len && past) {
        value = past[((b * heads + h) * past_capacity + s) * dim + d];
    } else if (s >= past_len && s < past_len + seq) {
        const int current_s = s - past_len;
        value = current[((b * seq + current_s) * heads + h) * dim + d];
    }
    present[idx] = value;
}

extern "C" __global__ void gqa_append_cache(
    const float* current, float* present, const int* past_lengths,
    int batch, int seq, int heads, int dim, int present_capacity)
{
    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
    const int count = batch * heads * seq * dim;
    if (idx >= count) return;
    int x = idx;
    const int d = x % dim; x /= dim;
    const int s = x % seq; x /= seq;
    const int h = x % heads; const int b = x / heads;
    const int target_s = past_lengths[b] + s;
    if (target_s < 0 || target_s >= present_capacity) return;
    present[((b * heads + h) * present_capacity + target_s) * dim + d] =
        current[((b * seq + s) * heads + h) * dim + d];
}

extern "C" __global__ void gqa_rope_bnsh(
    float* tensor, const float* cos_cache, const float* sin_cache,
    const long long* position_ids, const int* past_lengths,
    int batch, int seq, int heads, int dim, int rotary_dim, int tensor_capacity,
    int current_offset, int cache_rows, int interleaved, int cache_is_half)
{
    (void)cache_is_half;
    const int half = rotary_dim / 2;
    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
    const int count = batch * heads * seq * half;
    if (idx >= count) return;
    int x = idx;
    const int k = x % half; x /= half;
    const int s = x % seq; x /= seq;
    const int h = x % heads; const int b = x / heads;
    const int past = past_lengths[b];
    if (past < 0) return;
    const long long position = position_ids
        ? position_ids[b * seq + s]
        : (long long)past + s;
    if (position < 0 || position >= (long long)cache_rows) return;
    const int pos = (int)position;
    const int d0 = interleaved ? 2 * k : k;
    const int d1 = interleaved ? 2 * k + 1 : k + half;
    const int tensor_s = current_offset ? past_lengths[b] + s : s;
    const size_t base = ((size_t)(b * heads + h) * tensor_capacity + tensor_s) * dim;
    const float x0 = tensor[base + d0];
    const float x1 = tensor[base + d1];
    const float c = cos_cache[pos * half + k];
    const float sn = sin_cache[pos * half + k];
    tensor[base + d0] =
        __fsub_rn(__fmul_rn(c, x0), __fmul_rn(sn, x1));
    tensor[base + d1] =
        __fadd_rn(__fmul_rn(sn, x0), __fmul_rn(c, x1));
}

extern "C" __global__ void gqa_transpose_bnsh_to_bsh(
    const float* src, float* dst, int batch, int seq, int heads, int dim)
{
    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
    const int count = batch * seq * heads * dim;
    if (idx >= count) return;
    int x = idx;
    const int d = x % dim; x /= dim;
    const int h = x % heads; x /= heads;
    const int s = x % seq; const int b = x / seq;
    dst[idx] = src[((b * heads + h) * seq + s) * dim + d];
}

extern "C" __global__ void gqa_attention_reference_f32(
    const float* query,
    const float* key,
    const float* value,
    float* output,
    float* scores,
    const int* total_lengths,
    const int batch,
    const int query_heads,
    const int kv_heads,
    const int query_seq,
    const int head_size,
    const int cache_capacity,
    const int group_size,
    const float scale,
    const int local_window,
    const float softcap)
{
    const int row = blockIdx.x;
    const int rows = batch * query_heads * query_seq;
    if (row >= rows) return;
    const int query_pos = row % query_seq;
    const int query_head = (row / query_seq) % query_heads;
    const int batch_index = row / (query_heads * query_seq);
    const int kv_head = query_head / group_size;
    const int total = total_lengths[batch_index];
    const int causal_limit = total - query_seq + query_pos;
    const int local_start =
        local_window > 0 && causal_limit + 1 > local_window
            ? causal_limit + 1 - local_window
            : 0;
    float* row_scores = scores + (long)row * cache_capacity;

    if (threadIdx.x == 0) {
        const float negative_infinity = __int_as_float(0xff800000);
        float maximum = negative_infinity;
        for (int key_pos = 0; key_pos < total; ++key_pos) {
            float score = negative_infinity;
            if (key_pos >= local_start && key_pos <= causal_limit) {
                score = 0.0f;
                const long q_base =
                    ((long)(batch_index * query_heads + query_head) * query_seq + query_pos)
                    * head_size;
                const long k_base =
                    ((long)(batch_index * kv_heads + kv_head) * cache_capacity + key_pos)
                    * head_size;
                for (int d = 0; d < head_size; ++d) {
                    score = __fadd_rn(
                        score,
                        __fmul_rn(query[q_base + d], key[k_base + d]));
                }
                score = __fmul_rn(score, scale);
                if (softcap != 0.0f) {
                    score = __fmul_rn(softcap, tanhf(score / softcap));
                }
            }
            row_scores[key_pos] = score;
            maximum = fmaxf(maximum, score);
        }
        for (int key_pos = 0; key_pos < total; ++key_pos) {
            float probability = isfinite(row_scores[key_pos])
                ? (float)exp((double)(row_scores[key_pos] - maximum))
                : 0.0f;
            row_scores[key_pos] = probability;
        }
    }
    __syncthreads();

    float sum = 0.0f;
    for (int key_pos = 0; key_pos < total; ++key_pos) {
        sum = __fadd_rn(sum, row_scores[key_pos]);
    }
    for (int d = threadIdx.x; d < head_size; d += blockDim.x) {
        float result = 0.0f;
        for (int key_pos = 0; key_pos < total; ++key_pos) {
            const long v_index =
                ((long)(batch_index * kv_heads + kv_head) * cache_capacity + key_pos)
                * head_size + d;
            const float weighted =
                __fmul_rn(row_scores[key_pos] / sum, value[v_index]);
            result = __fadd_rn(result, weighted);
        }
        output[
            ((long)(batch_index * query_heads + query_head) * query_seq + query_pos)
                * head_size + d] = result;
    }
}

// Fused single-token (Sq=1, Sk=1) decode prep: split (implicit, reads packed or
// unpacked source directly), BSH->BNSH transpose (identity for Sq=1), in-place
// KV cache append, and RoPE for Q and present-K -- all in one launch. Replaces
// the split+transpose_in+append(K,V)+rope(Q,K) chain on the aliased device-KV
// decode path. For batch 1, metadata is derived independently by thread 0 of
// every CTA and shared within that CTA; block 0 also writes the attention
// metadata arrays. This avoids a cross-CTA dependency while preserving the
// sticky error latch and sentinel-gated cache writes.
extern "C" __global__ void gqa_fuse_decode_prep(
    const float* q_src, const float* k_src, const float* v_src, int packed,
    float* q_bnsh, float* present_k, float* present_v,
    const int* seqlens_k, int* total_lengths, int* past_lengths,
    int* query_starts, int past_capacity, int* error_flag, int derive_metadata,
    const float* cos_cache, const float* sin_cache,
    const long long* position_ids, int batch, int q_heads, int kv_heads, int dim,
    int rotary_dim, int present_capacity, int cache_rows, int do_rotary,
    int interleaved, int cache_is_half)
{
    (void)cache_is_half;
    const int head_half = dim / 2;
    const int rotary_half = rotary_dim / 2;
    const int qN = q_heads * head_half;
    const int kvN = kv_heads * head_half;
    const int per_batch = qN + 2 * kvN;
    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
    __shared__ int batch1_past;
    if (derive_metadata) {
        if (threadIdx.x == 0) {
            batch1_past = gqa_prepare_metadata_batch1(
                seqlens_k, total_lengths, past_lengths, query_starts,
                past_capacity, present_capacity, position_ids, do_rotary,
                cache_rows, error_flag, blockIdx.x == 0);
        }
        __syncthreads();
    }
    if (idx >= batch * per_batch) return;
    const int b = idx / per_batch;
    int local = idx - b * per_batch;
    const int past = derive_metadata ? batch1_past : past_lengths[b];
    const int q_hidden = q_heads * dim;
    const int kv_hidden = kv_heads * dim;
    const int packed_hidden = q_hidden + 2 * kv_hidden;
    const long long position = position_ids ? position_ids[b] : (long long)past;
    const int rope_ok = do_rotary && position >= 0 && position < (long long)cache_rows;
    const int pos = rope_ok ? (int)position : 0;
    int region, h, k;
    if (local < qN) { region = 0; h = local / head_half; k = local % head_half; }
    else if (local < qN + kvN) { local -= qN; region = 1; h = local / head_half; k = local % head_half; }
    else { local -= qN + kvN; region = 2; h = local / head_half; k = local % head_half; }
    const int is_rotary = k < rotary_half;
    const int tail = rotary_dim + 2 * (k - rotary_half);
    const int d0 = is_rotary ? (interleaved ? 2 * k : k) : tail;
    const int d1 = is_rotary ? (interleaved ? 2 * k + 1 : k + rotary_half) : tail + 1;
    if (region == 0) {
        const long src = (long)b * (packed ? packed_hidden : q_hidden) + (long)h * dim;
        const long dst = (long)(b * q_heads + h) * dim;
        const float x0 = q_src[src + d0];
        const float x1 = q_src[src + d1];
        if (rope_ok && past >= 0 && is_rotary) {
            const float c = cos_cache[pos * rotary_half + k];
            const float sn = sin_cache[pos * rotary_half + k];
            q_bnsh[dst + d0] = __fsub_rn(__fmul_rn(c, x0), __fmul_rn(sn, x1));
            q_bnsh[dst + d1] = __fadd_rn(__fmul_rn(sn, x0), __fmul_rn(c, x1));
        } else {
            q_bnsh[dst + d0] = x0;
            q_bnsh[dst + d1] = x1;
        }
        return;
    }
    if (past < 0 || past >= present_capacity) return;
    const long dst = ((long)(b * kv_heads + h) * present_capacity + past) * dim;
    if (region == 1) {
        const long src = (long)b * (packed ? packed_hidden : kv_hidden)
                       + (packed ? q_hidden : 0) + (long)h * dim;
        const float x0 = k_src[src + d0];
        const float x1 = k_src[src + d1];
        if (rope_ok && is_rotary) {
            const float c = cos_cache[pos * rotary_half + k];
            const float sn = sin_cache[pos * rotary_half + k];
            present_k[dst + d0] = __fsub_rn(__fmul_rn(c, x0), __fmul_rn(sn, x1));
            present_k[dst + d1] = __fadd_rn(__fmul_rn(sn, x0), __fmul_rn(c, x1));
        } else {
            present_k[dst + d0] = x0;
            present_k[dst + d1] = x1;
        }
    } else {
        const long src = (long)b * (packed ? packed_hidden : kv_hidden)
                       + (packed ? (q_hidden + kv_hidden) : 0) + (long)h * dim;
        present_v[dst + d0] = v_src[src + d0];
        present_v[dst + d1] = v_src[src + d1];
    }
}
"#;

const PREP_HALF_SRC: &str = r#"
#include <cuda_fp16.h>
#include <cuda_bf16.h>

__device__ __forceinline__ int gqa_prepare_metadata_batch1(
    const int* seqlens_k, int* total_lengths, int* past_lengths,
    int* query_starts, int past_capacity, int present_capacity,
    const long long* position_ids, int validate_positions, int cache_rows,
    int* error_flag, int write_metadata)
{
    if (error_flag && atomicOr(error_flag, 0) != 0) {
        if (write_metadata) {
            total_lengths[0] = -1;
            past_lengths[0] = -1;
            query_starts[0] = -1;
        }
        return -1;
    }
    const long long total = (long long)seqlens_k[0] + 1;
    const long long past = total - 1;
    const long long query_start = total - 1;
    int error = 0;
    if (total > 2147483647LL) error |= 1;
    if (past < 0) error |= 2;
    if (query_start < 0) error |= 4;
    if (past > past_capacity) error |= 8;
    if (total > present_capacity) error |= 16;
    if (validate_positions) {
        const long long position = position_ids ? position_ids[0] : past;
        if (position < 0 || position >= (long long)cache_rows) error |= 32;
    }
    if (error) {
        if (error_flag) atomicOr(error_flag, error);
        if (write_metadata) {
            total_lengths[0] = -1;
            past_lengths[0] = -1;
            query_starts[0] = -1;
        }
        return -1;
    }
    if (write_metadata) {
        total_lengths[0] = (int)total;
        past_lengths[0] = (int)past;
        query_starts[0] = (int)query_start;
    }
    return (int)past;
}

template <typename T> __device__ __forceinline__ float gqa_load(T value);
template <> __device__ __forceinline__ float gqa_load<__half>(__half value) {
    return __half2float(value);
}
template <> __device__ __forceinline__ float gqa_load<__nv_bfloat16>(__nv_bfloat16 value) {
    return __bfloat162float(value);
}
template <typename T> __device__ __forceinline__ T gqa_store(float value);
template <> __device__ __forceinline__ __half gqa_store<__half>(float value) {
    return __float2half_rn(value);
}
template <> __device__ __forceinline__ __nv_bfloat16 gqa_store<__nv_bfloat16>(float value) {
    return __float2bfloat16_rn(value);
}

__device__ __forceinline__ float gqa_load_cache(
    const void* cache, int index, int cache_is_half) {
    return cache_is_half
        ? __half2float(reinterpret_cast<const __half*>(cache)[index])
        : reinterpret_cast<const float*>(cache)[index];
}

template <typename T>
__device__ void gqa_transpose_bsh_to_bnsh_body(
    const T* src, T* dst, int batch, int seq, int heads, int dim)
{
    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
    const int count = batch * heads * seq * dim;
    if (idx >= count) return;
    int x = idx;
    const int d = x % dim; x /= dim;
    const int s = x % seq; x /= seq;
    const int h = x % heads; const int b = x / heads;
    dst[idx] = src[((b * seq + s) * heads + h) * dim + d];
}

template <typename T>
__device__ void gqa_split_packed_qkv_body(
    const T* packed, T* query, T* key, T* value,
    int batch, int seq, int q_heads, int kv_heads, int dim)
{
    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
    const int q_hidden = q_heads * dim;
    const int kv_hidden = kv_heads * dim;
    const int packed_hidden = q_hidden + 2 * kv_hidden;
    const int count = batch * seq * packed_hidden;
    if (idx >= count) return;
    const int feature = idx % packed_hidden;
    const int token = idx / packed_hidden;
    if (feature < q_hidden) {
        query[token * q_hidden + feature] = packed[idx];
    } else if (feature < q_hidden + kv_hidden) {
        key[token * kv_hidden + feature - q_hidden] = packed[idx];
    } else {
        value[token * kv_hidden + feature - q_hidden - kv_hidden] = packed[idx];
    }
}

template <typename T>
__device__ void gqa_build_cache_body(
    const T* current, const T* past, T* present,
    const int* past_lengths, int batch, int seq, int heads, int dim,
    int past_capacity, int present_capacity)
{
    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
    const int count = batch * heads * present_capacity * dim;
    if (idx >= count) return;
    int x = idx;
    const int d = x % dim; x /= dim;
    const int s = x % present_capacity; x /= present_capacity;
    const int h = x % heads; const int b = x / heads;
    const int past_len = past_lengths[b];
    if (past_len < 0) return;
    T result = gqa_store<T>(0.0f);
    if (s < past_len && past) {
        result = past[((b * heads + h) * past_capacity + s) * dim + d];
    } else if (s >= past_len && s < past_len + seq) {
        const int current_s = s - past_len;
        result = current[((b * seq + current_s) * heads + h) * dim + d];
    }
    present[idx] = result;
}

template <typename T>
__device__ void gqa_append_cache_body(
    const T* current, T* present, const int* past_lengths,
    int batch, int seq, int heads, int dim, int present_capacity)
{
    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
    const int count = batch * heads * seq * dim;
    if (idx >= count) return;
    int x = idx;
    const int d = x % dim; x /= dim;
    const int s = x % seq; x /= seq;
    const int h = x % heads; const int b = x / heads;
    const int target_s = past_lengths[b] + s;
    if (target_s < 0 || target_s >= present_capacity) return;
    present[((b * heads + h) * present_capacity + target_s) * dim + d] =
        current[((b * seq + s) * heads + h) * dim + d];
}

template <typename T>
__device__ void gqa_rope_bnsh_body(
    T* tensor, const void* cos_cache, const void* sin_cache,
    const long long* position_ids, const int* past_lengths,
    int batch, int seq, int heads, int dim, int rotary_dim, int tensor_capacity,
    int current_offset, int cache_rows, int interleaved, int cache_is_half)
{
    const int half = rotary_dim / 2;
    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
    const int count = batch * heads * seq * half;
    if (idx >= count) return;
    int x = idx;
    const int k = x % half; x /= half;
    const int s = x % seq; x /= seq;
    const int h = x % heads; const int b = x / heads;
    const int past = past_lengths[b];
    if (past < 0) return;
    const long long position = position_ids
        ? position_ids[b * seq + s]
        : (long long)past + s;
    if (position < 0 || position >= (long long)cache_rows) return;
    const int pos = (int)position;
    const int d0 = interleaved ? 2 * k : k;
    const int d1 = interleaved ? 2 * k + 1 : k + half;
    const int tensor_s = current_offset ? past_lengths[b] + s : s;
    const size_t base = ((size_t)(b * heads + h) * tensor_capacity + tensor_s) * dim;
    const float x0 = gqa_load<T>(tensor[base + d0]);
    const float x1 = gqa_load<T>(tensor[base + d1]);
    const float c = gqa_load_cache(cos_cache, pos * half + k, cache_is_half);
    const float sn = gqa_load_cache(sin_cache, pos * half + k, cache_is_half);
    tensor[base + d0] = gqa_store<T>(c * x0 - sn * x1);
    tensor[base + d1] = gqa_store<T>(sn * x0 + c * x1);
}

template <typename T>
__device__ void gqa_transpose_bnsh_to_bsh_body(
    const T* src, T* dst, int batch, int seq, int heads, int dim)
{
    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
    const int count = batch * seq * heads * dim;
    if (idx >= count) return;
    int x = idx;
    const int d = x % dim; x /= dim;
    const int h = x % heads; x /= heads;
    const int s = x % seq; const int b = x / seq;
    dst[idx] = src[((b * heads + h) * seq + s) * dim + d];
}

// Half/bf16 counterpart of `gqa_fuse_decode_prep` (see the f32 source for the
// full contract). RoPE reads/writes go through the shared float load/store and
// cache helpers so the fused result is bit-identical to the unfused
// split+transpose+append+rope chain; raw (non-rotary) writes stay direct T
// copies to avoid any extra float round-trip.
template <typename T>
__device__ void gqa_fuse_decode_prep_body(
    const T* q_src, const T* k_src, const T* v_src, int packed,
    T* q_bnsh, T* present_k, T* present_v,
    const int* seqlens_k, int* total_lengths, int* past_lengths,
    int* query_starts, int past_capacity, int* error_flag, int derive_metadata,
    const void* cos_cache, const void* sin_cache,
    const long long* position_ids, int batch, int q_heads, int kv_heads, int dim,
    int rotary_dim, int present_capacity, int cache_rows, int do_rotary,
    int interleaved, int cache_is_half)
{
    const int head_half = dim / 2;
    const int rotary_half = rotary_dim / 2;
    const int qN = q_heads * head_half;
    const int kvN = kv_heads * head_half;
    const int per_batch = qN + 2 * kvN;
    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
    __shared__ int batch1_past;
    if (derive_metadata) {
        if (threadIdx.x == 0) {
            batch1_past = gqa_prepare_metadata_batch1(
                seqlens_k, total_lengths, past_lengths, query_starts,
                past_capacity, present_capacity, position_ids, do_rotary,
                cache_rows, error_flag, blockIdx.x == 0);
        }
        __syncthreads();
    }
    if (idx >= batch * per_batch) return;
    const int b = idx / per_batch;
    int local = idx - b * per_batch;
    const int past = derive_metadata ? batch1_past : past_lengths[b];
    const int q_hidden = q_heads * dim;
    const int kv_hidden = kv_heads * dim;
    const int packed_hidden = q_hidden + 2 * kv_hidden;
    const long long position = position_ids ? position_ids[b] : (long long)past;
    const int rope_ok = do_rotary && position >= 0 && position < (long long)cache_rows;
    const int pos = rope_ok ? (int)position : 0;
    int region, h, k;
    if (local < qN) { region = 0; h = local / head_half; k = local % head_half; }
    else if (local < qN + kvN) { local -= qN; region = 1; h = local / head_half; k = local % head_half; }
    else { local -= qN + kvN; region = 2; h = local / head_half; k = local % head_half; }
    const int is_rotary = k < rotary_half;
    const int tail = rotary_dim + 2 * (k - rotary_half);
    const int d0 = is_rotary ? (interleaved ? 2 * k : k) : tail;
    const int d1 = is_rotary ? (interleaved ? 2 * k + 1 : k + rotary_half) : tail + 1;
    if (region == 0) {
        const long src = (long)b * (packed ? packed_hidden : q_hidden) + (long)h * dim;
        const long dst = (long)(b * q_heads + h) * dim;
        if (rope_ok && past >= 0 && is_rotary) {
            const float x0 = gqa_load<T>(q_src[src + d0]);
            const float x1 = gqa_load<T>(q_src[src + d1]);
            const float c = gqa_load_cache(cos_cache, pos * rotary_half + k, cache_is_half);
            const float sn = gqa_load_cache(sin_cache, pos * rotary_half + k, cache_is_half);
            q_bnsh[dst + d0] = gqa_store<T>(c * x0 - sn * x1);
            q_bnsh[dst + d1] = gqa_store<T>(sn * x0 + c * x1);
        } else {
            q_bnsh[dst + d0] = q_src[src + d0];
            q_bnsh[dst + d1] = q_src[src + d1];
        }
        return;
    }
    if (past < 0 || past >= present_capacity) return;
    const long dst = ((long)(b * kv_heads + h) * present_capacity + past) * dim;
    if (region == 1) {
        const long src = (long)b * (packed ? packed_hidden : kv_hidden)
                       + (packed ? q_hidden : 0) + (long)h * dim;
        if (rope_ok && is_rotary) {
            const float x0 = gqa_load<T>(k_src[src + d0]);
            const float x1 = gqa_load<T>(k_src[src + d1]);
            const float c = gqa_load_cache(cos_cache, pos * rotary_half + k, cache_is_half);
            const float sn = gqa_load_cache(sin_cache, pos * rotary_half + k, cache_is_half);
            present_k[dst + d0] = gqa_store<T>(c * x0 - sn * x1);
            present_k[dst + d1] = gqa_store<T>(sn * x0 + c * x1);
        } else {
            present_k[dst + d0] = k_src[src + d0];
            present_k[dst + d1] = k_src[src + d1];
        }
    } else {
        const long src = (long)b * (packed ? packed_hidden : kv_hidden)
                       + (packed ? (q_hidden + kv_hidden) : 0) + (long)h * dim;
        present_v[dst + d0] = v_src[src + d0];
        present_v[dst + d1] = v_src[src + d1];
    }
}

#define DEFINE_GQA_HALF_KERNELS(TYPE, SUFFIX) \
extern "C" __global__ void gqa_transpose_bsh_to_bnsh_##SUFFIX( \
    const TYPE* src, TYPE* dst, int batch, int seq, int heads, int dim) { \
    gqa_transpose_bsh_to_bnsh_body<TYPE>(src, dst, batch, seq, heads, dim); \
} \
extern "C" __global__ void gqa_split_packed_qkv_##SUFFIX( \
    const TYPE* packed, TYPE* query, TYPE* key, TYPE* value, \
    int batch, int seq, int q_heads, int kv_heads, int dim) { \
    gqa_split_packed_qkv_body<TYPE>( \
        packed, query, key, value, batch, seq, q_heads, kv_heads, dim); \
} \
extern "C" __global__ void gqa_build_cache_##SUFFIX( \
    const TYPE* current, const TYPE* past, TYPE* present, \
    const int* past_lengths, int batch, int seq, int heads, int dim, \
    int past_capacity, int present_capacity) { \
    gqa_build_cache_body<TYPE>(current, past, present, past_lengths, batch, seq, heads, \
                               dim, past_capacity, present_capacity); \
} \
extern "C" __global__ void gqa_append_cache_##SUFFIX( \
    const TYPE* current, TYPE* present, const int* past_lengths, \
    int batch, int seq, int heads, int dim, int present_capacity) { \
    gqa_append_cache_body<TYPE>( \
        current, present, past_lengths, batch, seq, heads, dim, present_capacity); \
} \
extern "C" __global__ void gqa_rope_bnsh_##SUFFIX( \
    TYPE* tensor, const void* cos_cache, const void* sin_cache, \
    const long long* position_ids, const int* past_lengths, \
    int batch, int seq, int heads, int dim, int rotary_dim, int tensor_capacity, \
    int current_offset, int cache_rows, int interleaved, int cache_is_half) { \
    gqa_rope_bnsh_body<TYPE>(tensor, cos_cache, sin_cache, position_ids, past_lengths, \
                             batch, seq, heads, dim, rotary_dim, tensor_capacity, current_offset, \
                             cache_rows, interleaved, cache_is_half); \
} \
extern "C" __global__ void gqa_transpose_bnsh_to_bsh_##SUFFIX( \
    const TYPE* src, TYPE* dst, int batch, int seq, int heads, int dim) { \
    gqa_transpose_bnsh_to_bsh_body<TYPE>(src, dst, batch, seq, heads, dim); \
} \
extern "C" __global__ void gqa_fuse_decode_prep_##SUFFIX( \
    const TYPE* q_src, const TYPE* k_src, const TYPE* v_src, int packed, \
    TYPE* q_bnsh, TYPE* present_k, TYPE* present_v, \
    const int* seqlens_k, int* total_lengths, int* past_lengths, \
    int* query_starts, int past_capacity, int* error_flag, int derive_metadata, \
    const void* cos_cache, const void* sin_cache, \
    const long long* position_ids, int batch, int q_heads, int kv_heads, int dim, \
    int rotary_dim, int present_capacity, int cache_rows, int do_rotary, \
    int interleaved, int cache_is_half) { \
    gqa_fuse_decode_prep_body<TYPE>(q_src, k_src, v_src, packed, q_bnsh, present_k, \
        present_v, seqlens_k, total_lengths, past_lengths, query_starts, past_capacity, \
        error_flag, derive_metadata, cos_cache, sin_cache, position_ids, batch, \
        q_heads, kv_heads, dim, rotary_dim, present_capacity, cache_rows, do_rotary, \
        interleaved, cache_is_half); \
}

DEFINE_GQA_HALF_KERNELS(__half, f16)
DEFINE_GQA_HALF_KERNELS(__nv_bfloat16, bf16)
"#;

const PREP_MODULE: &str = "group_query_attention_prep";
const PREP_HALF_MODULE: &str = "group_query_attention_prep_half_v1";
const BLOCK: u32 = 256;
const WS_TOTALS: usize = 0;
const WS_PAST_LENGTHS: usize = 1;
const WS_QUERY_STARTS: usize = 2;
const WS_PACKED_Q: usize = 3;
const WS_PACKED_K: usize = 4;
const WS_PACKED_V: usize = 5;
const WS_Q_BNSH: usize = 6;
const WS_OUT_BNSH: usize = 7;
const WS_PRESENT_K: usize = 8;
const WS_PRESENT_V: usize = 9;
const WS_SCORES: usize = 10;
const WS_COUNT: usize = 11;

/// Bit flags a captured GQA prep kernel `atomicOr`s into the runtime's latching
/// capture-error word when a decode-metadata invariant is violated during graph
/// replay. Exposed so hosts (and tests) can identify which bound was breached.
pub const GQA_CAPTURE_ERROR_TOTAL_OVERFLOW: u32 = 1;
pub const GQA_CAPTURE_ERROR_PAST_NEGATIVE: u32 = 2;
pub const GQA_CAPTURE_ERROR_QUERY_NEGATIVE: u32 = 4;
pub const GQA_CAPTURE_ERROR_PAST_CAPACITY: u32 = 8;
pub const GQA_CAPTURE_ERROR_PRESENT_CAPACITY: u32 = 16;
pub const GQA_CAPTURE_ERROR_POSITION: u32 = 32;

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

impl KernelFactory for GroupQueryAttentionFactory {
    fn create(&self, node: &Node, _input_shapes: &[Vec<usize>]) -> Result<Box<dyn Kernel>> {
        let required_heads = |name: &str| -> Result<usize> {
            let value = node.attr(name).and_then(|a| a.as_int()).ok_or_else(|| {
                EpError::KernelFailed(format!(
                    "cuda_ep GroupQueryAttention: missing required `{name}` attribute"
                ))
            })?;
            usize::try_from(value)
                .ok()
                .filter(|&v| v > 0)
                .ok_or_else(|| {
                    EpError::KernelFailed(format!(
                        "cuda_ep GroupQueryAttention: `{name}` must be > 0"
                    ))
                })
        };
        let num_heads = required_heads("num_heads")?;
        let kv_num_heads = required_heads("kv_num_heads")?;
        if !num_heads.is_multiple_of(kv_num_heads) {
            return Err(EpError::KernelFailed(format!(
                "cuda_ep GroupQueryAttention: num_heads {num_heads} must be a multiple of kv_num_heads {kv_num_heads}"
            )));
        }
        for name in ["k_quant_type", "v_quant_type"] {
            if let Some(value) = node.attr(name)
                && value.as_str() != Some("NONE")
            {
                return Err(EpError::KernelFailed(format!(
                    "cuda_ep GroupQueryAttention: `{name}` other than NONE is not supported"
                )));
            }
        }
        for (name, message) in [
            ("kv_cache_bit_width", "quantized KV cache"),
            ("qk_output", "qk_output"),
            ("smooth_softmax", "smooth_softmax"),
        ] {
            if node.attr(name).and_then(|a| a.as_int()).unwrap_or(0) != 0 {
                return Err(EpError::KernelFailed(format!(
                    "cuda_ep GroupQueryAttention: {message} is not supported"
                )));
            }
        }
        let softcap = node
            .attr("softcap")
            .and_then(|a| a.as_float())
            .unwrap_or(0.0);
        if softcap < 0.0 {
            return Err(EpError::KernelFailed(
                "cuda_ep GroupQueryAttention: softcap must be non-negative".into(),
            ));
        }
        Ok(Box::new(GroupQueryAttentionKernel::new(
            self.runtime.clone(),
            num_heads,
            kv_num_heads,
            node.attr("scale").and_then(|a| a.as_float()),
            node.attr("do_rotary").and_then(|a| a.as_int()).unwrap_or(0) != 0,
            node.attr("rotary_interleaved")
                .and_then(|a| a.as_int())
                .unwrap_or(0)
                != 0,
            node.attr("local_window_size")
                .and_then(|a| a.as_int())
                .unwrap_or(-1),
            softcap,
        )?))
    }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum GroupQueryAttentionBackend {
    Auto,
    Fused,
    Phase2a,
}

#[derive(Debug)]
pub struct GroupQueryAttentionKernel {
    runtime: Arc<CudaRuntime>,
    num_heads: usize,
    kv_num_heads: usize,
    scale: Option<f32>,
    do_rotary: bool,
    rotary_interleaved: bool,
    local_window_size: i64,
    softcap: f32,
    backend: GroupQueryAttentionBackend,
    prep_fusion_disabled: bool,
    workspace: Mutex<GqaWorkspace>,
    last_capture_safe_signature: Mutex<Option<GqaCaptureSignature>>,
}

#[derive(Clone, Debug, PartialEq, Eq)]
struct GqaCaptureSignature {
    dtype: DataType,
    batch: usize,
    query_sequence_length: usize,
    key_sequence_length: usize,
    q_hidden: usize,
    k_hidden: usize,
    dim: usize,
    past_capacity: usize,
    present_capacity: usize,
    packed_qkv: bool,
    explicit_positions: bool,
    cache_rows: usize,
    input_shapes: Vec<Option<Vec<usize>>>,
    output_shapes: Vec<Vec<usize>>,
    backend: GroupQueryAttentionBackend,
}

#[derive(Clone, Copy, Debug, Default)]
struct WorkspaceSlot {
    ptr: CUdeviceptr,
    bytes: usize,
}

#[derive(Debug)]
struct GqaWorkspace {
    runtime: Arc<CudaRuntime>,
    slots: [WorkspaceSlot; WS_COUNT],
}

impl GqaWorkspace {
    fn new(runtime: Arc<CudaRuntime>) -> Self {
        Self {
            runtime,
            slots: [WorkspaceSlot::default(); WS_COUNT],
        }
    }

    fn reserve(&mut self, index: usize, bytes: usize) -> Result<CUdeviceptr> {
        let bytes = bytes.max(1);
        let slot = self.slots[index];
        if slot.bytes >= bytes {
            return Ok(slot.ptr);
        }
        if self.runtime.is_capturing()? {
            return Err(EpError::KernelFailed(format!(
                "cuda_ep GroupQueryAttention: workspace slot {index} requires {bytes} bytes during CUDA graph capture; warm the fixed decode shape before capture"
            )));
        }
        let ptr = self.runtime.alloc_raw(bytes)?;
        if slot.ptr != 0 {
            // Dynamic prefill/growing-cache shapes may outgrow a slot. Preserve
            // the fixed-capacity decode fast path (which never reaches here),
            // but wait before replacing storage that queued work may still use.
            if let Err(error) = self.runtime.synchronize() {
                // SAFETY: `ptr` was allocated above and has not escaped.
                let _ = unsafe { self.runtime.free_raw(ptr) };
                return Err(error);
            }
            // SAFETY: synchronization completed all prior users of `slot.ptr`,
            // which is still exclusively owned by this workspace.
            if let Err(error) = unsafe { self.runtime.free_raw(slot.ptr) } {
                // SAFETY: `ptr` was allocated above and has not escaped.
                let _ = unsafe { self.runtime.free_raw(ptr) };
                return Err(error);
            }
        }
        self.slots[index] = WorkspaceSlot { ptr, bytes };
        Ok(ptr)
    }
}

impl Drop for GqaWorkspace {
    fn drop(&mut self) {
        for slot in self.slots.iter_mut().rev() {
            if slot.ptr != 0 {
                // SAFETY: every live slot pointer was allocated by this runtime
                // and remains exclusively owned by this workspace.
                let _ = unsafe { self.runtime.free_raw(slot.ptr) };
                slot.ptr = 0;
            }
        }
    }
}

fn checked_i32(value: usize, name: &str) -> Result<i32> {
    i32::try_from(value).map_err(|_| {
        EpError::KernelFailed(format!(
            "cuda_ep GroupQueryAttention: {name} {value} exceeds i32"
        ))
    })
}

fn require_matching_capture_signature(
    runtime: &CudaRuntime,
    warmed: Option<&GqaCaptureSignature>,
    current: Option<&GqaCaptureSignature>,
) -> Result<()> {
    if runtime.is_capturing()? && (current.is_none() || warmed != current) {
        return Err(EpError::KernelFailed(
            "cuda_ep GroupQueryAttention: dtype, decode mode, or shape changed during CUDA graph capture; warm the exact one-token fixed device-KV signature before capture".into(),
        ));
    }
    Ok(())
}

/// Human-readable description of the invariant(s) a captured GQA decode step
/// latched into the runtime capture-error word, given its raw bitmask. Returns
/// `None` for a zero (un-poisoned) mask.
pub fn gqa_capture_error_description(error: u32) -> Option<String> {
    if error == 0 {
        return None;
    }
    let mut violations = Vec::new();
    if error & GQA_CAPTURE_ERROR_TOTAL_OVERFLOW != 0 {
        violations.push("seqlens_k + 1 overflows int32");
    }
    if error & GQA_CAPTURE_ERROR_PAST_NEGATIVE != 0 {
        violations.push("seqlens_k + 1 is shorter than current key sequence");
    }
    if error & GQA_CAPTURE_ERROR_QUERY_NEGATIVE != 0 {
        violations.push("seqlens_k + 1 is shorter than current query sequence");
    }
    if error & GQA_CAPTURE_ERROR_PAST_CAPACITY != 0 {
        violations.push("effective past length exceeds past cache extent");
    }
    if error & GQA_CAPTURE_ERROR_PRESENT_CAPACITY != 0 {
        violations.push("valid sequence length exceeds present cache capacity");
    }
    if error & GQA_CAPTURE_ERROR_POSITION != 0 {
        violations.push("position_ids or implicit rotary position exceeds cache rows");
    }
    if violations.is_empty() {
        violations.push("unrecognized capture-safety violation");
    }
    Some(violations.join("; "))
}

fn require_dense(view: &TensorView, name: &str, dtype: DataType) -> Result<()> {
    if view.dtype != dtype {
        return Err(EpError::KernelFailed(format!(
            "cuda_ep GroupQueryAttention: {name} must have dtype {dtype:?}, got {:?}",
            view.dtype
        )));
    }
    if !view.is_contiguous() {
        return Err(EpError::KernelFailed(format!(
            "cuda_ep GroupQueryAttention: non-contiguous {name} is not supported; materialise it first"
        )));
    }
    Ok(())
}

fn validate_sequence_lengths_shape(shape: &[usize], numel: usize, batch: usize) -> Result<()> {
    if shape == [batch] || shape == [batch, 1] {
        return Ok(());
    }

    let scalar = shape.is_empty() && numel == 1;
    if scalar {
        return if batch == 1 {
            Ok(())
        } else {
            Err(EpError::KernelFailed(format!(
                "cuda_ep GroupQueryAttention: scalar seqlens_k can only be promoted to [1] when batch_size is 1, got batch_size {batch}; provide contiguous int32 [batch_size] or [batch_size, 1] values for every row"
            )))
        };
    }

    Err(EpError::KernelFailed(format!(
        "cuda_ep GroupQueryAttention: seqlens_k must be non-negative contiguous int32 with shape [batch_size], [batch_size, 1], or a scalar for batch_size 1 (for batch_size {batch}: [{batch}] or [{batch}, 1]), got shape {shape:?}"
    )))
}

fn read_i32(runtime: &CudaRuntime, view: &TensorView, name: &str) -> Result<Vec<i32>> {
    require_dense(view, name, DataType::Int32)?;
    let mut bytes = vec![0u8; view.numel() * 4];
    // SAFETY: the source tensor has exactly `bytes.len()` bytes.
    unsafe {
        runtime.dtoh(&mut bytes, cuptr(view.data_ptr::<u8>() as *const c_void))?;
    }
    Ok(bytes
        .chunks_exact(4)
        .map(|x| i32::from_ne_bytes([x[0], x[1], x[2], x[3]]))
        .collect())
}

fn read_i64(runtime: &CudaRuntime, view: &TensorView, name: &str) -> Result<Vec<i64>> {
    require_dense(view, name, DataType::Int64)?;
    let mut bytes = vec![0u8; view.numel() * 8];
    // SAFETY: the source tensor has exactly `bytes.len()` bytes.
    unsafe {
        runtime.dtoh(&mut bytes, cuptr(view.data_ptr::<u8>() as *const c_void))?;
    }
    Ok(bytes
        .chunks_exact(8)
        .map(|x| i64::from_ne_bytes([x[0], x[1], x[2], x[3], x[4], x[5], x[6], x[7]]))
        .collect())
}

macro_rules! launch_1d {
    ($runtime:expr, $module:expr, $source:expr, $entry:expr, $count:expr, $builder:ident, $args:block) => {{
        let launch_count: usize = $count;
        if launch_count != 0 {
            let function = $runtime.nvrtc_function($module, $source, $entry)?;
            let grid =
                u32::try_from(launch_count.div_ceil(BLOCK as usize)).map_err(|_| {
                    EpError::KernelFailed(
                        "cuda_ep GroupQueryAttention: launch grid exceeds u32".into(),
                    )
                })?;
            let mut $builder = $runtime.stream().launch_builder(&function);
            $args
            // SAFETY: each invocation supplies the argument ABI for its entry point;
            // input/output buffers outlive execution, and workspace buffers remain
            // owned by the kernel while stream-ordered work is pending.
            unsafe {
                $builder.launch(LaunchConfig {
                    grid_dim: (grid, 1, 1),
                    block_dim: (BLOCK, 1, 1),
                    shared_mem_bytes: 0,
                })
            }
            .map_err(|e| driver_err(&format!("launch {}", $entry), e))?;
        }
    }};
}

impl GroupQueryAttentionKernel {
    #[allow(clippy::too_many_arguments)]
    pub fn new(
        runtime: Arc<CudaRuntime>,
        num_heads: usize,
        kv_num_heads: usize,
        scale: Option<f32>,
        do_rotary: bool,
        rotary_interleaved: bool,
        local_window_size: i64,
        softcap: f32,
    ) -> Result<Self> {
        if num_heads == 0
            || kv_num_heads == 0
            || !num_heads.is_multiple_of(kv_num_heads)
            || local_window_size == 0
            || local_window_size < -1
            || softcap < 0.0
        {
            return Err(EpError::KernelFailed(
                "cuda_ep GroupQueryAttention: invalid heads, local window, or softcap".into(),
            ));
        }
        Ok(Self {
            workspace: Mutex::new(GqaWorkspace::new(runtime.clone())),
            runtime,
            num_heads,
            kv_num_heads,
            scale,
            do_rotary,
            rotary_interleaved,
            local_window_size,
            softcap,
            backend: GroupQueryAttentionBackend::Auto,
            prep_fusion_disabled: false,
            last_capture_safe_signature: Mutex::new(None),
        })
    }

    pub fn with_backend(mut self, backend: GroupQueryAttentionBackend) -> Self {
        self.backend = backend;
        self
    }

    /// Forces the unfused per-op decode prep chain (split, transpose, append,
    /// RoPE, output transpose) instead of the single fused decode-prep launch.
    /// Exposed so tests can prove the fused kernel is bit-identical to the
    /// unfused reference path on the same inputs.
    pub fn with_prep_fusion_disabled(mut self, disabled: bool) -> Self {
        self.prep_fusion_disabled = disabled;
        self
    }

    /// Reads the internal metadata workspace for GPU parity tests.
    #[doc(hidden)]
    pub fn read_prepared_metadata_for_test(
        &self,
        batch: usize,
    ) -> Result<(Vec<i32>, Vec<i32>, Vec<i32>)> {
        let workspace = self.workspace.lock().map_err(|_| {
            EpError::KernelFailed("cuda_ep GroupQueryAttention: workspace lock poisoned".into())
        })?;
        let read_slot = |index: usize| -> Result<Vec<i32>> {
            let slot = workspace.slots[index];
            let bytes_len = batch
                .checked_mul(std::mem::size_of::<i32>())
                .ok_or_else(|| {
                    EpError::KernelFailed(
                        "cuda_ep GroupQueryAttention: test metadata size overflow".into(),
                    )
                })?;
            if slot.ptr == 0 || slot.bytes < bytes_len {
                return Err(EpError::KernelFailed(
                    "cuda_ep GroupQueryAttention: test metadata workspace is unavailable".into(),
                ));
            }
            let mut bytes = vec![0u8; bytes_len];
            // SAFETY: the workspace slot is live and was reserved for at least
            // `bytes_len` bytes by the most recent successful execution.
            unsafe {
                self.runtime.dtoh(&mut bytes, slot.ptr)?;
            }
            Ok(bytes
                .chunks_exact(4)
                .map(|x| i32::from_ne_bytes([x[0], x[1], x[2], x[3]]))
                .collect())
        };
        Ok((
            read_slot(WS_TOTALS)?,
            read_slot(WS_PAST_LENGTHS)?,
            read_slot(WS_QUERY_STARTS)?,
        ))
    }

    /// Resolves the configured backend using the same shape gate as execution.
    pub fn selected_backend_for_shape(
        &self,
        dtype: DataType,
        query_sequence_length: usize,
        valid_sequence_length: usize,
        head_size: usize,
    ) -> GroupQueryAttentionBackend {
        let fused_supported = flash_attention::supported(query_sequence_length, head_size);
        let measured_fused_win = valid_sequence_length <= 128
            || (dtype == DataType::Float16
                && head_size.is_multiple_of(16)
                && valid_sequence_length <= 512
                && self.runtime.capabilities().compute_capability().0 >= 7);
        if fused_supported
            && (self.backend == GroupQueryAttentionBackend::Fused
                || (self.backend == GroupQueryAttentionBackend::Auto && measured_fused_win))
        {
            GroupQueryAttentionBackend::Fused
        } else {
            GroupQueryAttentionBackend::Phase2a
        }
    }

    fn run(&self, inputs: &[TensorView], outputs: &mut [TensorMut]) -> Result<()> {
        let mut last_signature = self.last_capture_safe_signature.lock().map_err(|_| {
            EpError::KernelFailed(
                "cuda_ep GroupQueryAttention: capture signature lock poisoned".into(),
            )
        })?;
        let warmed_signature = last_signature.take();
        if !(7..=14).contains(&inputs.len()) || !(1..=3).contains(&outputs.len()) {
            return Err(EpError::KernelFailed(format!(
                "cuda_ep GroupQueryAttention: expected 7..14 inputs and 1..3 outputs, got {} and {}",
                inputs.len(),
                outputs.len()
            )));
        }
        let packed_qkv = inputs[1].is_absent() && inputs[2].is_absent();
        if inputs[1].is_absent() != inputs[2].is_absent() {
            return Err(EpError::KernelFailed(
                "cuda_ep GroupQueryAttention: key and value must both be present for unpacked Q/K/V or both absent for packed QKV".into(),
            ));
        }
        for (index, feature) in [
            (10, "attention_bias"),
            (11, "head_sink"),
            (12, "quantized-cache k_scale"),
            (13, "quantized-cache v_scale"),
        ] {
            if inputs.get(index).is_some_and(|v| !v.is_absent()) {
                return Err(EpError::KernelFailed(format!(
                    "cuda_ep GroupQueryAttention: {feature} is not supported"
                )));
            }
        }
        if self.local_window_size == 0 || self.local_window_size < -1 {
            return Err(EpError::KernelFailed(
                "cuda_ep GroupQueryAttention: local_window_size must be -1 or a positive integer"
                    .into(),
            ));
        }

        let q = &inputs[0];
        let dtype = AttentionDtype::from_onnx(q.dtype).map_err(|_| {
            EpError::KernelFailed(format!(
                "cuda_ep GroupQueryAttention: query dtype {:?} is not supported; expected Float32, Float16, or BFloat16",
                q.dtype
            ))
        })?;
        require_dense(q, "query", q.dtype)?;
        if q.dtype != DataType::Float32 {
            self.runtime
                .require_nvrtc_half_headers("GroupQueryAttention")?;
        }
        let element_size = dtype.element_size() as usize;
        let (
            prep_module,
            prep_src,
            split_entry,
            transpose_in_entry,
            build_entry,
            append_entry,
            rope_entry,
            transpose_out_entry,
            fuse_entry,
        ) = match q.dtype {
            DataType::Float32 => (
                PREP_MODULE,
                PREP_SRC,
                "gqa_split_packed_qkv",
                "gqa_transpose_bsh_to_bnsh",
                "gqa_build_cache",
                "gqa_append_cache",
                "gqa_rope_bnsh",
                "gqa_transpose_bnsh_to_bsh",
                "gqa_fuse_decode_prep",
            ),
            DataType::Float16 => (
                PREP_HALF_MODULE,
                PREP_HALF_SRC,
                "gqa_split_packed_qkv_f16",
                "gqa_transpose_bsh_to_bnsh_f16",
                "gqa_build_cache_f16",
                "gqa_append_cache_f16",
                "gqa_rope_bnsh_f16",
                "gqa_transpose_bnsh_to_bsh_f16",
                "gqa_fuse_decode_prep_f16",
            ),
            DataType::BFloat16 => (
                PREP_HALF_MODULE,
                PREP_HALF_SRC,
                "gqa_split_packed_qkv_bf16",
                "gqa_transpose_bsh_to_bnsh_bf16",
                "gqa_build_cache_bf16",
                "gqa_append_cache_bf16",
                "gqa_rope_bnsh_bf16",
                "gqa_transpose_bnsh_to_bsh_bf16",
                "gqa_fuse_decode_prep_bf16",
            ),
            _ => unreachable!(),
        };
        if q.shape.len() != 3 {
            return Err(EpError::KernelFailed(format!(
                "cuda_ep GroupQueryAttention: query must be rank 3 [B,S,H], got {:?}",
                q.shape
            )));
        }
        let (batch, q_seq, input_hidden) = (q.shape[0], q.shape[1], q.shape[2]);
        let (q_hidden, k_seq, k_hidden, dim) = if packed_qkv {
            let packed_heads = self.num_heads + 2 * self.kv_num_heads;
            if batch == 0
                || q_seq == 0
                || input_hidden == 0
                || !input_hidden.is_multiple_of(packed_heads)
            {
                return Err(EpError::KernelFailed(format!(
                    "cuda_ep GroupQueryAttention: packed query must be [B,S,(num_heads + 2*kv_num_heads)*head_size], got {:?}",
                    q.shape
                )));
            }
            let dim = input_hidden / packed_heads;
            (self.num_heads * dim, q_seq, self.kv_num_heads * dim, dim)
        } else {
            let (k, v) = (&inputs[1], &inputs[2]);
            for (view, name) in [(k, "key"), (v, "value")] {
                require_dense(view, name, q.dtype)?;
                if view.shape.len() != 3 {
                    return Err(EpError::KernelFailed(format!(
                        "cuda_ep GroupQueryAttention: unpacked {name} must be rank 3 [B,S,H*D], got {:?}",
                        view.shape
                    )));
                }
            }
            let (k_batch, k_seq, k_hidden) = (k.shape[0], k.shape[1], k.shape[2]);
            if batch == 0
                || q_seq == 0
                || input_hidden == 0
                || k_hidden == 0
                || !input_hidden.is_multiple_of(self.num_heads)
                || !k_hidden.is_multiple_of(self.kv_num_heads)
                || v.shape != [batch, k_seq, k_hidden]
                || k_batch != batch
            {
                return Err(EpError::KernelFailed(
                    "cuda_ep GroupQueryAttention: incompatible query/key/value batch, sequence, or hidden dimensions".into(),
                ));
            }
            let dim = input_hidden / self.num_heads;
            if k_hidden / self.kv_num_heads != dim {
                return Err(EpError::KernelFailed(
                    "cuda_ep GroupQueryAttention: query and key/value head sizes must match".into(),
                ));
            }
            (input_hidden, k_seq, k_hidden, dim)
        };
        if outputs[0].dtype != q.dtype
            || outputs[0].shape != [batch, q_seq, q_hidden]
            || !outputs[0].is_contiguous()
        {
            return Err(EpError::KernelFailed(format!(
                "cuda_ep GroupQueryAttention: output must be contiguous {:?} [B,S,H*D] = [{batch},{q_seq},{q_hidden}], got {:?}",
                q.dtype, outputs[0].shape
            )));
        }

        let has_past_key = !inputs[3].is_absent();
        let has_past_value = !inputs[4].is_absent();
        if has_past_key != has_past_value {
            return Err(EpError::KernelFailed(
                "cuda_ep GroupQueryAttention: past_key and past_value must be provided together"
                    .into(),
            ));
        }
        let past_capacity = if has_past_key {
            for (view, name) in [(&inputs[3], "past_key"), (&inputs[4], "past_value")] {
                require_dense(view, name, q.dtype)?;
                if view.shape.len() != 4
                    || view.shape[0] != batch
                    || view.shape[1] != self.kv_num_heads
                    || view.shape[3] != dim
                {
                    return Err(EpError::KernelFailed(format!(
                        "cuda_ep GroupQueryAttention: {name} must be BNSH [{batch},{},{},{}], got {:?}",
                        self.kv_num_heads, view.shape[2], dim, view.shape
                    )));
                }
            }
            if inputs[3].shape != inputs[4].shape {
                return Err(EpError::KernelFailed(
                    "cuda_ep GroupQueryAttention: past_key and past_value shapes must match".into(),
                ));
            }
            inputs[3].shape[2]
        } else {
            0
        };

        if inputs[5].dtype == DataType::Int32 && !inputs[5].is_contiguous() {
            return Err(EpError::KernelFailed(
                "cuda_ep GroupQueryAttention: non-contiguous seqlens_k was provided; expected non-negative contiguous int32 with shape [batch_size] or [batch_size, 1]"
                    .into(),
            ));
        }
        require_dense(&inputs[5], "seqlens_k", DataType::Int32)?;
        validate_sequence_lengths_shape(inputs[5].shape, inputs[5].numel(), batch)?;
        require_dense(&inputs[6], "total_sequence_length", DataType::Int32)?;
        if inputs[6].numel() != 1 {
            return Err(EpError::KernelFailed(
                "cuda_ep GroupQueryAttention: total_sequence_length must be one non-negative int32 scalar".into(),
            ));
        }
        let current_key_length = checked_i32(k_seq, "key sequence length")?;
        let query_length = checked_i32(q_seq, "query sequence length")?;
        let requested_present_capacity = outputs
            .get(1)
            .map(|output| output.shape.get(2).copied().unwrap_or(past_capacity));

        let explicit_positions = inputs.get(9).filter(|view| !view.is_absent());
        let (
            cos_ptr,
            sin_ptr,
            positions_ptr,
            cache_rows,
            cache_rows_usize,
            rotary_dim,
            rotary_dim_usize,
            rope_cache_is_half,
        ) = if self.do_rotary {
            if k_seq != 0 && q_seq != k_seq {
                return Err(EpError::KernelFailed(
                        "cuda_ep GroupQueryAttention: do_rotary requires equal query/key sequence lengths unless current key/value are empty".into(),
                    ));
            }
            let cos = inputs
                .get(7)
                .filter(|view| !view.is_absent())
                .ok_or_else(|| {
                    EpError::KernelFailed(
                        "cuda_ep GroupQueryAttention: do_rotary=1 requires cos_cache".into(),
                    )
                })?;
            let sin = inputs
                .get(8)
                .filter(|view| !view.is_absent())
                .ok_or_else(|| {
                    EpError::KernelFailed(
                        "cuda_ep GroupQueryAttention: do_rotary=1 requires sin_cache".into(),
                    )
                })?;
            require_dense(cos, "cos_cache", DataType::Float32)
                .or_else(|_| require_dense(cos, "cos_cache", DataType::Float16))?;
            let cache_dtype = cos.dtype;
            let cache_is_half = match cache_dtype {
                DataType::Float32 => 0i32,
                DataType::Float16 if matches!(q.dtype, DataType::Float16 | DataType::BFloat16) => {
                    1i32
                }
                other => {
                    return Err(EpError::KernelFailed(format!(
                        "cuda_ep GroupQueryAttention: cos_cache/sin_cache dtype {other:?} unsupported for query dtype {:?}; expected Float32, or Float16 with half-precision queries",
                        q.dtype
                    )));
                }
            };
            require_dense(sin, "sin_cache", cache_dtype)?;
            if cos.shape.len() != 2 {
                return Err(EpError::KernelFailed(format!(
                    "cuda_ep GroupQueryAttention: cos_cache must be rank 2 [max_sequence_length, rotary_dim/2], got shape {:?}",
                    cos.shape
                )));
            }
            if sin.shape != cos.shape {
                return Err(EpError::KernelFailed(format!(
                    "cuda_ep GroupQueryAttention: sin_cache shape {:?} must exactly match cos_cache shape {:?} so both caches describe the same rotary_dim",
                    sin.shape, cos.shape
                )));
            }
            let rotary_dim_usize = cos.shape[1].checked_mul(2).ok_or_else(|| {
                    EpError::KernelFailed(format!(
                        "cuda_ep GroupQueryAttention: rotary_dim derived from cos_cache width {} overflows; use a finite cache width with 1 <= width <= head_size/2",
                        cos.shape[1]
                    ))
                })?;
            if rotary_dim_usize < 2
                || !rotary_dim_usize.is_multiple_of(2)
                || rotary_dim_usize > dim
                || cos.shape[1] != rotary_dim_usize / 2
            {
                return Err(EpError::KernelFailed(format!(
                    "cuda_ep GroupQueryAttention: rotary_dim derived from cos_cache width {} is {}; it must be even and satisfy 2 <= rotary_dim <= head_size={} (cache width must equal rotary_dim/2)",
                    cos.shape[1], rotary_dim_usize, dim
                )));
            }
            let position_ptr = if let Some(position_ids) = explicit_positions {
                require_dense(position_ids, "position_ids", DataType::Int64)?;
                if position_ids.shape != [batch, q_seq] {
                    return Err(EpError::KernelFailed(
                        "cuda_ep GroupQueryAttention: position_ids must be valid non-negative int64 [batch_size, sequence_length]".into(),
                    ));
                }
                cuptr(position_ids.data_ptr::<u8>() as *const c_void)
            } else {
                0
            };
            (
                cuptr(cos.data_ptr::<u8>() as *const c_void),
                cuptr(sin.data_ptr::<u8>() as *const c_void),
                position_ptr,
                checked_i32(cos.shape[0], "rotary cache rows")?,
                cos.shape[0],
                checked_i32(rotary_dim_usize, "rotary dimension")?,
                rotary_dim_usize,
                cache_is_half,
            )
        } else {
            (0, 0, 0, 0, 0, 0, 0, 0)
        };

        let structurally_valid_outputs = requested_present_capacity.is_some_and(|capacity| {
            let expected = [batch, self.kv_num_heads, capacity, dim];
            outputs.len() == 3
                && outputs[1].dtype == q.dtype
                && outputs[1].shape == expected
                && outputs[1].is_contiguous()
                && outputs[2].dtype == q.dtype
                && outputs[2].shape == expected
                && outputs[2].is_contiguous()
        });
        let aliased_device_kv = structurally_valid_outputs
            && (outputs[1].data.0 as *mut u8)
                .wrapping_add(outputs[1].byte_offset)
                .cast_const()
                == inputs[3].data_ptr::<u8>()
            && (outputs[2].data.0 as *mut u8)
                .wrapping_add(outputs[2].byte_offset)
                .cast_const()
                == inputs[4].data_ptr::<u8>();
        let capture_candidate = requested_present_capacity
            .filter(|&present_capacity| {
                (q.dtype == DataType::Float32
                    || (q.dtype == DataType::Float16 && gqa_decode_fp16::supported(q_seq, dim)))
                    && q_seq == 1
                    && k_seq <= 1
                    && has_past_key
                    && present_capacity >= 1
                    && past_capacity == present_capacity
                    && aliased_device_kv
                    && self.selected_backend_for_shape(q.dtype, q_seq, present_capacity, dim)
                        == GroupQueryAttentionBackend::Phase2a
            })
            .map(|present_capacity| GqaCaptureSignature {
                dtype: q.dtype,
                batch,
                query_sequence_length: q_seq,
                key_sequence_length: k_seq,
                q_hidden,
                k_hidden,
                dim,
                past_capacity,
                present_capacity,
                packed_qkv,
                explicit_positions: explicit_positions.is_some(),
                cache_rows: cache_rows_usize,
                input_shapes: inputs
                    .iter()
                    .enumerate()
                    .map(|(index, input)| {
                        (!input.is_absent()).then(|| {
                            if index == 5 {
                                vec![batch]
                            } else {
                                input.shape.to_vec()
                            }
                        })
                    })
                    .collect(),
                output_shapes: outputs.iter().map(|output| output.shape.to_vec()).collect(),
                backend: GroupQueryAttentionBackend::Phase2a,
            });
        require_matching_capture_signature(
            &self.runtime,
            warmed_signature.as_ref(),
            capture_candidate.as_ref(),
        )?;
        let capture_safe_decode = capture_candidate
            .as_ref()
            .is_some_and(|candidate| warmed_signature.as_ref() == Some(candidate));

        let mut valid_sequence_length = None;
        let mut validated_query_starts = None;
        let total_sequence_length = if capture_safe_decode {
            None
        } else {
            let seqlens = read_i32(&self.runtime, &inputs[5], "seqlens_k")?;
            if seqlens.iter().any(|&length| length < 0) {
                return Err(EpError::KernelFailed(
                    "cuda_ep GroupQueryAttention: seqlens_k must be non-negative int32 [batch_size]"
                        .into(),
                ));
            }
            let total_scalar = read_i32(&self.runtime, &inputs[6], "total_sequence_length")?;
            if total_scalar[0] < 0 {
                return Err(EpError::KernelFailed(
                    "cuda_ep GroupQueryAttention: total_sequence_length must be one non-negative int32 scalar".into(),
                ));
            }
            let total_sequence_length = total_scalar[0] as usize;
            let totals: Vec<i32> = seqlens
                .iter()
                .map(|&length| length.checked_add(1))
                .collect::<Option<_>>()
                .ok_or_else(|| {
                    EpError::KernelFailed(
                        "cuda_ep GroupQueryAttention: seqlens_k + 1 overflows int32".into(),
                    )
                })?;
            let maximum = totals.iter().copied().max().unwrap_or(0) as usize;
            if maximum > total_sequence_length {
                return Err(EpError::KernelFailed(format!(
                    "cuda_ep GroupQueryAttention: valid sequence length {maximum} exceeds physical total_sequence_length capacity {total_sequence_length}"
                )));
            }
            let mut query_starts = Vec::with_capacity(batch);
            for &total in &totals {
                let past = total.checked_sub(current_key_length).ok_or_else(|| {
                    EpError::KernelFailed(
                        "cuda_ep GroupQueryAttention: seqlens_k + 1 is shorter than current key sequence"
                            .into(),
                    )
                })?;
                let query_start = total.checked_sub(query_length).ok_or_else(|| {
                    EpError::KernelFailed(
                        "cuda_ep GroupQueryAttention: seqlens_k + 1 is shorter than current query sequence"
                            .into(),
                    )
                })?;
                if past as usize > past_capacity {
                    return Err(EpError::KernelFailed(
                        "cuda_ep GroupQueryAttention: effective past length exceeds past cache extent"
                            .into(),
                    ));
                }
                query_starts.push(query_start);
            }
            valid_sequence_length = Some(maximum);
            validated_query_starts = Some(query_starts);
            Some(total_sequence_length)
        };

        if !capture_safe_decode && self.do_rotary {
            if let Some(position_ids) = explicit_positions {
                let ids = read_i64(&self.runtime, position_ids, "position_ids")?;
                let cache_rows_i64 = i64::from(cache_rows);
                if ids
                    .iter()
                    .any(|&position| position < 0 || position >= cache_rows_i64)
                {
                    return Err(EpError::KernelFailed(
                        "cuda_ep GroupQueryAttention: position_ids must be valid non-negative int64 [batch_size, sequence_length]".into(),
                    ));
                }
            } else if validated_query_starts.as_ref().is_some_and(|starts| {
                starts
                    .iter()
                    .any(|&start| start as usize + q_seq > cache_rows_usize)
            }) {
                return Err(EpError::KernelFailed(
                    "cuda_ep GroupQueryAttention: rotary position exceeds cache rows".into(),
                ));
            }
        }

        let minimum_present_capacity =
            past_capacity.max(total_sequence_length.unwrap_or(past_capacity));
        let present_capacity = requested_present_capacity.unwrap_or(minimum_present_capacity);
        if present_capacity < minimum_present_capacity {
            return Err(EpError::KernelFailed(format!(
                "cuda_ep GroupQueryAttention: present cache capacity {present_capacity} is smaller than required {minimum_present_capacity}"
            )));
        }
        let expected_cache_shape = [batch, self.kv_num_heads, present_capacity, dim];
        for (index, name) in [(1, "present_key"), (2, "present_value")] {
            if let Some(output) = outputs.get(index)
                && (output.dtype != q.dtype
                    || output.shape != expected_cache_shape
                    || !output.is_contiguous())
            {
                return Err(EpError::KernelFailed(format!(
                    "cuda_ep GroupQueryAttention: {name} must be contiguous {:?} BNSH {:?}, got {:?}",
                    q.dtype, expected_cache_shape, output.shape
                )));
            }
        }

        // Single-token decode lets the trailing BNSH->BSH output transpose
        // collapse into the attention write (identical layouts when Sq==1), and
        // when the KV cache is appended in place (aliased device-KV, matching
        // past/present capacity) the whole split+transpose+append+RoPE prep
        // chain fuses into one launch. Prefill/growing-cache/non-aliased shapes
        // keep the unfused kernels.
        let single_token = q_seq == 1;
        let fuse_prep = single_token
            && k_seq == 1
            && dim.is_multiple_of(2)
            && has_past_key
            && aliased_device_kv
            && past_capacity == present_capacity
            && !self.prep_fusion_disabled;
        let fuse_metadata = fuse_prep && batch == 1;

        let mut workspace = self.workspace.lock().map_err(|_| {
            EpError::KernelFailed("cuda_ep GroupQueryAttention: workspace lock poisoned".into())
        })?;
        let metadata_bytes = batch * std::mem::size_of::<i32>();
        let totals_gpu = workspace.reserve(WS_TOTALS, metadata_bytes)?;
        let past_lengths_gpu = workspace.reserve(WS_PAST_LENGTHS, metadata_bytes)?;
        let query_starts_gpu = workspace.reserve(WS_QUERY_STARTS, metadata_bytes)?;
        let metadata_error_gpu = if capture_candidate.is_some() {
            // Capture-safe decode steps latch any bounds violation into the
            // runtime-shared word so every captured GQA layer poisons the same
            // flag, and the host detects it once per step at the logits sync.
            self.runtime.capture_error_ptr()
        } else {
            0
        };
        let packed_q = (packed_qkv && !fuse_prep)
            .then(|| workspace.reserve(WS_PACKED_Q, batch * q_seq * q_hidden * element_size))
            .transpose()?;
        let packed_k = (packed_qkv && !fuse_prep)
            .then(|| workspace.reserve(WS_PACKED_K, batch * k_seq * k_hidden * element_size))
            .transpose()?;
        let packed_v = (packed_qkv && !fuse_prep)
            .then(|| workspace.reserve(WS_PACKED_V, batch * k_seq * k_hidden * element_size))
            .transpose()?;
        let q_bnsh = workspace.reserve(WS_Q_BNSH, batch * q_seq * q_hidden * element_size)?;
        // Sq==1 writes attention output straight into the BSH output tensor
        // (identical layout), so the BNSH scratch is only needed for Sq>1.
        let out_bnsh = (!single_token)
            .then(|| workspace.reserve(WS_OUT_BNSH, outputs[0].numel() * element_size))
            .transpose()?;
        let owned_present_k = (outputs.len() < 2)
            .then(|| {
                workspace.reserve(
                    WS_PRESENT_K,
                    expected_cache_shape.iter().product::<usize>() * element_size,
                )
            })
            .transpose()?;
        let owned_present_v = (outputs.len() < 3)
            .then(|| {
                workspace.reserve(
                    WS_PRESENT_V,
                    expected_cache_shape.iter().product::<usize>() * element_size,
                )
            })
            .transpose()?;
        let present_k_ptr = if let Some(output) = outputs.get_mut(1) {
            cuptr(output.data_ptr_mut::<u8>() as *const c_void)
        } else {
            *owned_present_k.as_ref().ok_or_else(|| {
                EpError::KernelFailed(
                    "cuda_ep GroupQueryAttention: internal present-key allocation missing".into(),
                )
            })?
        };
        let present_v_ptr = if let Some(output) = outputs.get_mut(2) {
            cuptr(output.data_ptr_mut::<u8>() as *const c_void)
        } else {
            *owned_present_v.as_ref().ok_or_else(|| {
                EpError::KernelFailed(
                    "cuda_ep GroupQueryAttention: internal present-value allocation missing".into(),
                )
            })?
        };
        let output_ptr = cuptr(outputs[0].data_ptr_mut::<u8>() as *const c_void);
        // Sq==1 makes BNSH and BSH layouts identical, so attention writes into
        // the real output tensor directly and the trailing transpose is skipped.
        let attention_out = if single_token {
            output_ptr
        } else {
            *out_bnsh.as_ref().ok_or_else(|| {
                EpError::KernelFailed(
                    "cuda_ep GroupQueryAttention: internal BNSH output allocation missing".into(),
                )
            })?
        };

        let batch_i = checked_i32(batch, "batch")?;
        let q_seq_i = checked_i32(q_seq, "query sequence length")?;
        let k_seq_i = checked_i32(k_seq, "key sequence length")?;
        let heads_i = checked_i32(self.num_heads, "num_heads")?;
        let kv_heads_i = checked_i32(self.kv_num_heads, "kv_num_heads")?;
        let dim_i = checked_i32(dim, "head_size")?;
        let past_capacity_i = checked_i32(past_capacity, "past capacity")?;
        let present_capacity_i = checked_i32(present_capacity, "present capacity")?;
        let local_window_i = i32::try_from(self.local_window_size.max(0)).map_err(|_| {
            EpError::KernelFailed(
                "cuda_ep GroupQueryAttention: local_window_size exceeds i32".into(),
            )
        })?;
        let seqlens_ptr = cuptr(inputs[5].data_ptr::<u8>() as *const c_void);
        let validate_positions_i: i32 = self.do_rotary.into();
        if fuse_metadata {
            onnx_runtime_ep_api::record_kernel_variant_stage!(
                "metadata",
                "gqa_prep_fused_with_metadata",
                "batch-1 fixed-capacity single-token decode derives past/total/query-start \
                 metadata inside fused prep with device-side bounds and sticky error latching"
            );
        } else {
            onnx_runtime_ep_api::record_kernel_variant_stage!(
                "metadata",
                "metadata_separate",
                "metadata remains a separate launch: batch={}, prep_fused={}; folding requires \
                 batch==1 and the eligible fixed-capacity single-token fused prep path",
                batch,
                fuse_prep
            );
            launch_1d!(
                self.runtime,
                PREP_MODULE,
                PREP_SRC,
                "gqa_prepare_metadata",
                batch,
                builder,
                {
                    builder
                        .arg(&seqlens_ptr)
                        .arg(&totals_gpu)
                        .arg(&past_lengths_gpu)
                        .arg(&query_starts_gpu)
                        .arg(&batch_i)
                        .arg(&current_key_length)
                        .arg(&query_length)
                        .arg(&past_capacity_i)
                        .arg(&present_capacity_i)
                        .arg(&positions_ptr)
                        .arg(&validate_positions_i)
                        .arg(&cache_rows)
                        .arg(&metadata_error_gpu);
                }
            );
        }
        let input_q_ptr = cuptr(q.data_ptr::<u8>() as *const c_void);
        if fuse_prep {
            let prep_variant = if fuse_metadata {
                "gqa_prep_fused_with_metadata"
            } else {
                "gqa_prep_fused"
            };
            onnx_runtime_ep_api::record_kernel_variant_stage!(
                "prep",
                prep_variant,
                "decode prep fused into one launch: Sq==1, k_seq==1, even head_dim={}, \
                 aliased device-KV, past_capacity==present_capacity={}, metadata_fused={}",
                dim,
                present_capacity,
                fuse_metadata
            );
            // Fused single-token decode prep. One launch subsumes the packed
            // split, BSH->BNSH query transpose, in-place K/V cache append, and
            // Q/present-K RoPE that the branch below performs separately. Batch
            // 1 also derives metadata per CTA, with block 0 writing the arrays;
            // larger batches consume the stream-ordered separate metadata.
            let (q_src, k_src, v_src, packed_flag) = if packed_qkv {
                (input_q_ptr, input_q_ptr, input_q_ptr, 1i32)
            } else {
                (
                    input_q_ptr,
                    cuptr(inputs[1].data_ptr::<u8>() as *const c_void),
                    cuptr(inputs[2].data_ptr::<u8>() as *const c_void),
                    0i32,
                )
            };
            let interleaved_i: i32 = self.rotary_interleaved.into();
            let do_rotary_i: i32 = self.do_rotary.into();
            let derive_metadata_i: i32 = fuse_metadata.into();
            let fused_count = batch * (self.num_heads + 2 * self.kv_num_heads) * (dim / 2);
            launch_1d!(
                self.runtime,
                prep_module,
                prep_src,
                fuse_entry,
                fused_count,
                builder,
                {
                    builder
                        .arg(&q_src)
                        .arg(&k_src)
                        .arg(&v_src)
                        .arg(&packed_flag)
                        .arg(&q_bnsh)
                        .arg(&present_k_ptr)
                        .arg(&present_v_ptr)
                        .arg(&seqlens_ptr)
                        .arg(&totals_gpu)
                        .arg(&past_lengths_gpu)
                        .arg(&query_starts_gpu)
                        .arg(&past_capacity_i)
                        .arg(&metadata_error_gpu)
                        .arg(&derive_metadata_i)
                        .arg(&cos_ptr)
                        .arg(&sin_ptr)
                        .arg(&positions_ptr)
                        .arg(&batch_i)
                        .arg(&heads_i)
                        .arg(&kv_heads_i)
                        .arg(&dim_i)
                        .arg(&rotary_dim)
                        .arg(&present_capacity_i)
                        .arg(&cache_rows)
                        .arg(&do_rotary_i)
                        .arg(&interleaved_i)
                        .arg(&rope_cache_is_half);
                }
            );
        } else {
            onnx_runtime_ep_api::record_kernel_variant_stage!(
                "prep",
                "gqa_prep_unfused",
                "decode prep unfused (split/transpose/append/rope as separate launches): \
                 single_token={}, k_seq={}, even_head_dim={}, has_past_key={}, \
                 aliased_device_kv={}, past==present_capacity={}, prep_fusion_disabled={}",
                single_token,
                k_seq,
                dim.is_multiple_of(2),
                has_past_key,
                aliased_device_kv,
                past_capacity == present_capacity,
                self.prep_fusion_disabled
            );
            let (q_ptr, k_ptr, v_ptr) = if packed_qkv {
                let q_scratch = packed_q.as_ref().ok_or_else(|| {
                    EpError::KernelFailed(
                        "cuda_ep GroupQueryAttention: internal packed-query allocation missing"
                            .into(),
                    )
                })?;
                let k_scratch = packed_k.as_ref().ok_or_else(|| {
                    EpError::KernelFailed(
                        "cuda_ep GroupQueryAttention: internal packed-key allocation missing"
                            .into(),
                    )
                })?;
                let v_scratch = packed_v.as_ref().ok_or_else(|| {
                    EpError::KernelFailed(
                        "cuda_ep GroupQueryAttention: internal packed-value allocation missing"
                            .into(),
                    )
                })?;
                let packed_count = q.numel();
                launch_1d!(
                    self.runtime,
                    prep_module,
                    prep_src,
                    split_entry,
                    packed_count,
                    builder,
                    {
                        builder
                            .arg(&input_q_ptr)
                            .arg(q_scratch)
                            .arg(k_scratch)
                            .arg(v_scratch)
                            .arg(&batch_i)
                            .arg(&q_seq_i)
                            .arg(&heads_i)
                            .arg(&kv_heads_i)
                            .arg(&dim_i);
                    }
                );
                (*q_scratch, *k_scratch, *v_scratch)
            } else {
                (
                    input_q_ptr,
                    cuptr(inputs[1].data_ptr::<u8>() as *const c_void),
                    cuptr(inputs[2].data_ptr::<u8>() as *const c_void),
                )
            };
            launch_1d!(
                self.runtime,
                prep_module,
                prep_src,
                transpose_in_entry,
                batch * q_seq * q_hidden,
                builder,
                {
                    builder
                        .arg(&q_ptr)
                        .arg(&q_bnsh)
                        .arg(&batch_i)
                        .arg(&q_seq_i)
                        .arg(&heads_i)
                        .arg(&dim_i);
                }
            );

            let past_k_ptr = if has_past_key {
                cuptr(inputs[3].data_ptr::<u8>() as *const c_void)
            } else {
                0
            };
            let past_v_ptr = if has_past_value {
                cuptr(inputs[4].data_ptr::<u8>() as *const c_void)
            } else {
                0
            };
            for (current, past, present) in [
                (k_ptr, past_k_ptr, present_k_ptr),
                (v_ptr, past_v_ptr, present_v_ptr),
            ] {
                if past != 0 && past == present && past_capacity == present_capacity {
                    launch_1d!(
                        self.runtime,
                        prep_module,
                        prep_src,
                        append_entry,
                        batch * self.kv_num_heads * k_seq * dim,
                        builder,
                        {
                            builder
                                .arg(&current)
                                .arg(&present)
                                .arg(&past_lengths_gpu)
                                .arg(&batch_i)
                                .arg(&k_seq_i)
                                .arg(&kv_heads_i)
                                .arg(&dim_i)
                                .arg(&present_capacity_i);
                        }
                    );
                } else {
                    launch_1d!(
                        self.runtime,
                        prep_module,
                        prep_src,
                        build_entry,
                        expected_cache_shape.iter().product::<usize>(),
                        builder,
                        {
                            builder
                                .arg(&current)
                                .arg(&past)
                                .arg(&present)
                                .arg(&past_lengths_gpu)
                                .arg(&batch_i)
                                .arg(&k_seq_i)
                                .arg(&kv_heads_i)
                                .arg(&dim_i)
                                .arg(&past_capacity_i)
                                .arg(&present_capacity_i);
                        }
                    );
                }
            }

            if self.do_rotary {
                let interleaved_i: i32 = self.rotary_interleaved.into();
                for (tensor, positions, seq_i, heads, capacity, current_offset) in [
                    (q_bnsh, query_starts_gpu, q_seq_i, heads_i, q_seq_i, 0i32),
                    (
                        present_k_ptr,
                        past_lengths_gpu,
                        k_seq_i,
                        kv_heads_i,
                        present_capacity_i,
                        1i32,
                    ),
                ] {
                    let count =
                        batch * (heads as usize) * (seq_i as usize) * (rotary_dim_usize / 2);
                    launch_1d!(
                        self.runtime,
                        prep_module,
                        prep_src,
                        rope_entry,
                        count,
                        builder,
                        {
                            builder
                                .arg(&tensor)
                                .arg(&cos_ptr)
                                .arg(&sin_ptr)
                                .arg(&positions_ptr)
                                .arg(&positions)
                                .arg(&batch_i)
                                .arg(&seq_i)
                                .arg(&heads)
                                .arg(&dim_i)
                                .arg(&rotary_dim)
                                .arg(&capacity)
                                .arg(&current_offset)
                                .arg(&cache_rows)
                                .arg(&interleaved_i)
                                .arg(&rope_cache_is_half);
                        }
                    );
                }
            }
        }

        let scale = self
            .scale
            .filter(|&scale| scale != 0.0)
            .unwrap_or_else(|| 1.0 / (dim as f32).sqrt());
        let attention_sequence_length = valid_sequence_length.unwrap_or(present_capacity);
        let selected_backend =
            self.selected_backend_for_shape(q.dtype, q_seq, attention_sequence_length, dim);
        crate::trace::record_kernel_metrics(inputs, outputs, || {
            let score_elements = (batch as u64)
                .saturating_mul(self.num_heads as u64)
                .saturating_mul(q_seq as u64)
                .saturating_mul(attention_sequence_length as u64);
            let qk_flops = score_elements.saturating_mul(dim as u64).saturating_mul(2);
            let pv_flops = score_elements.saturating_mul(dim as u64).saturating_mul(2);
            let softmax_flops = score_elements.saturating_mul(4).saturating_add(
                (batch as u64)
                    .saturating_mul(self.num_heads as u64)
                    .saturating_mul(q_seq as u64),
            );
            qk_flops
                .saturating_add(pv_flops)
                .saturating_add(softmax_flops)
        });
        let use_fused = selected_backend == GroupQueryAttentionBackend::Fused;
        if use_fused {
            onnx_runtime_ep_api::record_kernel_variant!(
                "attention_flash_fused",
                "fused flash attention: backend={:?}, dtype={:?}, q_seq={}, \
                 valid_seq_len={}, head_dim={} passed the fused-support + measured-win gates",
                self.backend,
                q.dtype,
                q_seq,
                attention_sequence_length,
                dim
            );
            flash_attention::run(
                &self.runtime,
                q.dtype,
                self.num_heads,
                self.kv_num_heads,
                true,
                batch,
                q_seq,
                attention_sequence_length,
                present_capacity,
                dim,
                self.num_heads / self.kv_num_heads,
                scale,
                q_bnsh,
                present_k_ptr,
                present_v_ptr,
                attention_out,
                0,
                0,
                totals_gpu,
                query_starts_gpu,
                local_window_i,
                self.softcap,
            )?;
        } else if q.dtype == DataType::Float32 && gqa_decode::supported(q_seq, dim) {
            onnx_runtime_ep_api::record_kernel_variant!(
                "attention_gqa_decode_f32_splitk",
                "capture-safe f32 split-K single-token decode: q_seq={}, head_dim={}; \
                 active split count (1/2/4/8/16, max {}) is chosen on-device; \
                 flash backend={:?} not selected",
                q_seq,
                dim,
                gqa_decode::MAX_SPLITS,
                selected_backend
            );
            // Capture-safe split-K single-token GQA decode. Reads the valid
            // length on-device and uses fixed module-global scratch, so both
            // partial and merge launches record/replay inside a CUDA graph.
            gqa_decode::run(
                &self.runtime,
                batch,
                self.num_heads,
                self.kv_num_heads,
                q_seq,
                dim,
                present_capacity,
                self.num_heads / self.kv_num_heads,
                scale,
                q_bnsh,
                present_k_ptr,
                present_v_ptr,
                attention_out,
                totals_gpu,
                local_window_i,
                self.softcap,
            )?;
        } else if q.dtype == DataType::Float16 && gqa_decode_fp16::supported(q_seq, dim) {
            onnx_runtime_ep_api::record_kernel_variant!(
                "attention_gqa_decode_fp16_splitk",
                "capture-safe fp16 split-K flash-decode: q_seq={}, even head_dim={} (<=256); \
                 active split count (up to {}) chosen on-device from the valid length \
                 and a host occupancy target that fills the multiprocessors",
                q_seq,
                dim,
                gqa_decode_fp16::MAX_SPLITS
            );
            // Capture-safe fp16 flash-decode sibling of the f32 `gqa_decode`
            // branch above. Same launcher signature/units; passes the fp16
            // device pointers for query/present-K/present-V/output. Reads the
            // valid length on-device from `totals_gpu` and allocates only
            // fixed-size dynamic shared memory, so it records/replays inside a
            // CUDA graph. Unsupported fp16 shapes (e.g. prefill Sq>1) still fall
            // through to the phase-2a path below.
            gqa_decode_fp16::run(
                &self.runtime,
                batch,
                self.num_heads,
                self.kv_num_heads,
                q_seq,
                dim,
                present_capacity,
                self.num_heads / self.kv_num_heads,
                scale,
                q_bnsh,
                present_k_ptr,
                present_v_ptr,
                attention_out,
                totals_gpu,
                local_window_i,
                self.softcap,
            )?;
        } else if q.dtype == DataType::Float32 {
            onnx_runtime_ep_api::record_kernel_variant!(
                "attention_reference_f32",
                "f32 reference attention: flash backend={:?} for q_seq={}, valid_seq_len={}, \
                 head_dim={}; capture-safe gqa_decode does not support this shape",
                selected_backend,
                q_seq,
                attention_sequence_length,
                dim
            );
            let attention_rows = batch
                .checked_mul(self.num_heads)
                .and_then(|rows| rows.checked_mul(q_seq))
                .ok_or_else(|| {
                    EpError::KernelFailed(
                        "cuda_ep GroupQueryAttention: attention row count overflow".into(),
                    )
                })?;
            let score_count = attention_rows
                .checked_mul(present_capacity)
                .ok_or_else(|| {
                    EpError::KernelFailed(
                        "cuda_ep GroupQueryAttention: score scratch size overflow".into(),
                    )
                })?;
            let score_scratch = workspace.reserve(WS_SCORES, score_count.max(1) * 4)?;
            let attention_rows_u32 = u32::try_from(attention_rows).map_err(|_| {
                EpError::KernelFailed(
                    "cuda_ep GroupQueryAttention: attention row count exceeds u32".into(),
                )
            })?;
            let kv_heads_i = checked_i32(self.kv_num_heads, "KV head count")?;
            let group_i = checked_i32(
                self.num_heads / self.kv_num_heads,
                "query-to-KV head group size",
            )?;
            let func = self.runtime.nvrtc_function(
                PREP_MODULE,
                PREP_SRC,
                "gqa_attention_reference_f32",
            )?;
            let mut builder = self.runtime.stream().launch_builder(&func);
            builder
                .arg(&q_bnsh)
                .arg(&present_k_ptr)
                .arg(&present_v_ptr)
                .arg(&attention_out)
                .arg(&score_scratch)
                .arg(&totals_gpu)
                .arg(&batch_i)
                .arg(&heads_i)
                .arg(&kv_heads_i)
                .arg(&q_seq_i)
                .arg(&dim_i)
                .arg(&present_capacity_i)
                .arg(&group_i)
                .arg(&scale)
                .arg(&local_window_i)
                .arg(&self.softcap);
            // SAFETY: the scratch and BNSH buffers are sized above, and the scalar
            // ABI matches `gqa_attention_reference_f32`.
            unsafe {
                builder.launch(LaunchConfig {
                    grid_dim: (attention_rows_u32, 1, 1),
                    block_dim: (BLOCK, 1, 1),
                    shared_mem_bytes: 0,
                })
            }
            .map_err(|error| driver_err("launch GQA reference attention", error))?;
        } else {
            onnx_runtime_ep_api::record_kernel_variant!(
                "attention_phase2a",
                "phase2a general attention: dtype={:?}, q_seq={}, valid_seq_len={}, head_dim={} \
                 not selected for fused flash or a capture-safe dtype-specific decode kernel",
                q.dtype,
                q_seq,
                attention_sequence_length,
                dim
            );
            run_attention_phase2a(
                &self.runtime,
                dtype,
                self.num_heads,
                self.kv_num_heads,
                true,
                batch,
                q_seq,
                attention_sequence_length,
                dim,
                present_capacity,
                self.num_heads / self.kv_num_heads,
                scale,
                q_bnsh,
                present_k_ptr,
                present_v_ptr,
                attention_out,
                0,
                0,
                totals_gpu,
                query_starts_gpu,
                local_window_i,
                self.softcap,
            )?;
        }

        // For Sq==1 attention already wrote the BSH output in place, so the
        // BNSH->BSH transpose is skipped. Sq>1 still materialises via `out_bnsh`.
        if !single_token {
            let out_bnsh_ptr = out_bnsh.ok_or_else(|| {
                EpError::KernelFailed(
                    "cuda_ep GroupQueryAttention: internal BNSH output allocation missing".into(),
                )
            })?;
            launch_1d!(
                self.runtime,
                prep_module,
                prep_src,
                transpose_out_entry,
                outputs[0].numel(),
                builder,
                {
                    builder
                        .arg(&out_bnsh_ptr)
                        .arg(&output_ptr)
                        .arg(&batch_i)
                        .arg(&q_seq_i)
                        .arg(&heads_i)
                        .arg(&dim_i);
                }
            );
        }
        *last_signature = capture_candidate;
        Ok(())
    }
}

impl Kernel for GroupQueryAttentionKernel {
    fn execute(&self, inputs: &[TensorView], outputs: &mut [TensorMut]) -> Result<()> {
        self.run(inputs, outputs)
    }

    fn supports_strided_input(&self, _input_idx: usize) -> bool {
        true
    }

    fn capture_support(&self) -> onnx_runtime_ep_api::CaptureSupport {
        // Eligibility is tied to the exact one-token, fixed-capacity, in-place
        // device-KV decode signature warmed by the most recent successful call.
        match self.last_capture_safe_signature.lock() {
            Ok(signature) if signature.is_some() => onnx_runtime_ep_api::CaptureSupport::Supported,
            Ok(_) => onnx_runtime_ep_api::CaptureSupport::unsupported(
                "requires a warmed f32/fp16 q_seq==1 k_seq<=1 fixed-capacity device-KV decode path; the current signature was not warmed as capture-safe",
            ),
            Err(_) => onnx_runtime_ep_api::CaptureSupport::unsupported(
                "GroupQueryAttention capture signature is unavailable because its state lock was poisoned",
            ),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn sequence_lengths_shape_accepts_canonical_per_batch_forms() {
        for batch in [1, 3] {
            validate_sequence_lengths_shape(&[batch], batch, batch).unwrap();
            validate_sequence_lengths_shape(&[batch, 1], batch, batch).unwrap();
        }
    }

    #[test]
    fn sequence_lengths_shape_rejects_noncanonical_singleton_layouts_actionably() {
        for shape in [vec![1, 3], vec![3, 1, 1]] {
            let error = validate_sequence_lengths_shape(&shape, 3, 3)
                .expect_err("noncanonical seqlens_k shape must fail");
            let message = format!("{error}");
            assert!(message.contains("[batch_size], [batch_size, 1]"));
            assert!(message.contains(&format!("got shape {shape:?}")));
        }
    }

    #[test]
    fn sequence_lengths_shape_promotes_unit_batch_scalar_and_rejects_larger_batches() {
        validate_sequence_lengths_shape(&[], 1, 1).unwrap();
        let error = validate_sequence_lengths_shape(&[], 1, 2)
            .expect_err("a scalar cannot represent a multi-row batch");
        assert!(format!("{error}").contains("batch_size 2"));
    }

    fn runtime() -> Option<Arc<CudaRuntime>> {
        let previous_hook = std::panic::take_hook();
        std::panic::set_hook(Box::new(|_| {}));
        let runtime = std::panic::catch_unwind(|| CudaRuntime::new(0).ok().map(Arc::new))
            .ok()
            .flatten();
        std::panic::set_hook(previous_hook);
        runtime
    }

    #[test]
    fn persistent_workspace_reuses_fixed_shape_allocation() {
        let Some(runtime) = runtime() else {
            eprintln!("skipping CUDA GQA workspace test: CUDA runtime unavailable");
            return;
        };
        let before = runtime.allocation_counts();
        let mut workspace = GqaWorkspace::new(runtime.clone());

        let first = workspace.reserve(WS_Q_BNSH, 4096).unwrap();
        let allocated = runtime.allocation_counts();
        assert_eq!(allocated.allocations, before.allocations + 1);
        assert_eq!(allocated.frees, before.frees);

        assert_eq!(workspace.reserve(WS_Q_BNSH, 4096).unwrap(), first);
        assert_eq!(workspace.reserve(WS_Q_BNSH, 2048).unwrap(), first);
        assert_eq!(runtime.allocation_counts(), allocated);

        let grown = workspace.reserve(WS_Q_BNSH, 8192).unwrap();
        assert_ne!(grown, first);
        let grown_counts = runtime.allocation_counts();
        assert_eq!(grown_counts.allocations, before.allocations + 2);
        assert_eq!(grown_counts.frees, before.frees + 1);

        drop(workspace);
        let after = runtime.allocation_counts();
        assert_eq!(after.allocations, before.allocations + 2);
        assert_eq!(after.frees, before.frees + 2);
    }
}