From 762c13c348a03881876a0ba9f7e00d5b911edfff Mon Sep 17 00:00:00 2001 From: justinchuby <223556219+Copilot@users.noreply.github.com> Date: Tue, 11 Aug 2026 15:43:13 -0700 Subject: [PATCH 1/2] Convert flash prefill KV reads to stride descriptors Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../src/kernels/attention.rs | 2 + .../src/kernels/flash_attention.rs | 31 +- .../src/kernels/gqa_decode_fp16.rs | 4 +- .../src/kernels/group_query_attention.rs | 159 +++---- .../src/kernels/kv_stride.rs | 150 +++++-- .../tests/gqa_seqmajor_parity_gpu.rs | 397 +++++++++++++++++- docs/MEMORY_ARCHITECTURE.md | 22 +- 7 files changed, 629 insertions(+), 136 deletions(-) diff --git a/crates/onnx-runtime-ep-cuda/src/kernels/attention.rs b/crates/onnx-runtime-ep-cuda/src/kernels/attention.rs index 499913086b..ca36f872c5 100644 --- a/crates/onnx-runtime-ep-cuda/src/kernels/attention.rs +++ b/crates/onnx-runtime-ep-cuda/src/kernels/attention.rs @@ -70,6 +70,7 @@ use onnx_runtime_memory_governor::MemoryRole; use crate::blas::{GemmDtype, GemmEx, WORKSPACE_BYTES, gemm_ex}; use crate::error::{driver_err, not_implemented}; +use crate::kernels::kv_stride::KvCacheStrides; use crate::runtime::{CudaRuntime, cuptr}; use super::flash_attention; @@ -733,6 +734,7 @@ impl AttentionKernel { 0, 0, 0.0, + &KvCacheStrides::head_major_bnsh(), ) } else { // Carve the governed scratch from the executor-prepared workspace diff --git a/crates/onnx-runtime-ep-cuda/src/kernels/flash_attention.rs b/crates/onnx-runtime-ep-cuda/src/kernels/flash_attention.rs index 608526a42e..c4dc884085 100644 --- a/crates/onnx-runtime-ep-cuda/src/kernels/flash_attention.rs +++ b/crates/onnx-runtime-ep-cuda/src/kernels/flash_attention.rs @@ -13,6 +13,7 @@ use onnx_runtime_ep_api::{EpError, Result}; use onnx_runtime_ir::DataType; use crate::error::{driver_err, not_implemented}; +use crate::kernels::kv_stride::{KvCachePath, KvCacheStrides}; use crate::runtime::CudaRuntime; const QUERY_TILE: usize = 8; @@ -82,8 +83,7 @@ __device__ __forceinline__ void flash_attention_body( const unsigned long long q_base = ((unsigned long long)(b * q_heads + h) * sq + qi) * dim; - const unsigned long long kv_base = - (unsigned long long)(b * kv_heads + kvh) * kv_capacity * dim; + const unsigned long long kv_base = FLASH_KV_BASE(b, kvh); for (int key0 = 0; key0 < sk; key0 += FLASH_K_TILE) { const int tile_keys = min(FLASH_K_TILE, sk - key0); @@ -95,7 +95,8 @@ __device__ __forceinline__ void flash_attention_body( float kval = 0.0f; float vval = 0.0f; if (j < sk) { - const unsigned long long offset = kv_base + (unsigned long long)j * dim + p; + const unsigned long long offset = + kv_base + (unsigned long long)j * FLASH_KV_STRIDE + p; kval = flash_load(key[offset]); vval = flash_load(value[offset]); } @@ -314,8 +315,7 @@ extern "C" __global__ void flash_attention_f16_tc( const int logical_sk = total_lengths ? total_lengths[b] : sk; const unsigned long long q_plane = (unsigned long long)(b * q_heads + h) * sq * dim; - const unsigned long long kv_plane = - (unsigned long long)(b * kv_heads + kvh) * kv_capacity * dim; + const unsigned long long kv_plane = FLASH_KV_BASE(b, kvh); for (int index = tid; index < FLASH_TC_Q * FLASH_TC_D; index += blockDim.x) { const int row = index / FLASH_TC_D; @@ -337,7 +337,8 @@ extern "C" __global__ void flash_attention_f16_tc( const int p = index - row * FLASH_TC_D; const int kj = key0 + row; if (kj < sk && p < dim) { - const unsigned long long offset = kv_plane + (unsigned long long)kj * dim + p; + const unsigned long long offset = + kv_plane + (unsigned long long)kj * FLASH_KV_STRIDE + p; k_tile[index] = key[offset]; v_tile[index] = value[offset]; } else { @@ -516,6 +517,7 @@ pub(super) fn run( past_lengths: CUdeviceptr, local_window: i32, softcap: f32, + kv_strides: &KvCacheStrides, ) -> Result<()> { if !supported(sq, head_dim) { return Err(not_implemented(format!( @@ -526,27 +528,28 @@ pub(super) fn run( let tensor_core_f16 = dtype == DataType::Float16 && head_dim.is_multiple_of(16) && runtime.capabilities().compute_capability().0 >= 7; - let (module, source, entry, query_tile) = match dtype { + kv_strides.require_converted_path_support(KvCachePath::FlashPrefillRead)?; + let (module, source_body, entry, query_tile) = match dtype { DataType::Float32 => ( - "flash_attention_f32_v1", + kv_strides.flash_f32_module_key()?, FLASH_F32_SOURCE.as_str(), "flash_attention_f32", QUERY_TILE, ), DataType::Float16 if tensor_core_f16 => ( - "flash_attention_half_v2", + kv_strides.flash_half_module_key()?, FLASH_HALF_SOURCE.as_str(), "flash_attention_f16_tc", 16, ), DataType::Float16 => ( - "flash_attention_half_v2", + kv_strides.flash_half_module_key()?, FLASH_HALF_SOURCE.as_str(), "flash_attention_f16", QUERY_TILE, ), DataType::BFloat16 => ( - "flash_attention_half_v2", + kv_strides.flash_half_module_key()?, FLASH_HALF_SOURCE.as_str(), "flash_attention_bf16", QUERY_TILE, @@ -562,9 +565,10 @@ pub(super) fn run( // crt/mma.h from a separately packaged wheel. Say so here rather than let // NVRTC report a missing file whose name identifies neither the kernel nor // the package that carries it. - if source == FLASH_HALF_SOURCE.as_str() { + if source_body == FLASH_HALF_SOURCE.as_str() { runtime.require_nvrtc_tensor_core_headers("fused Attention")?; } + let source = format!("{}{}", kv_strides.flash_prelude(), source_body); let as_i32 = |name: &str, value: usize| { i32::try_from(value).map_err(|_| { @@ -600,7 +604,7 @@ pub(super) fn run( }; } - let function = runtime.nvrtc_function(module, source, entry)?; + let function = runtime.nvrtc_function(module, &source, entry)?; let mut builder = runtime.stream().launch_builder(&function); builder .arg(&q) @@ -869,6 +873,7 @@ mod tests { 0, 0, 0.0, + &KvCacheStrides::head_major_bnsh(), ) .unwrap(); diff --git a/crates/onnx-runtime-ep-cuda/src/kernels/gqa_decode_fp16.rs b/crates/onnx-runtime-ep-cuda/src/kernels/gqa_decode_fp16.rs index e79c9c9266..200beed07b 100644 --- a/crates/onnx-runtime-ep-cuda/src/kernels/gqa_decode_fp16.rs +++ b/crates/onnx-runtime-ep-cuda/src/kernels/gqa_decode_fp16.rs @@ -52,7 +52,7 @@ use cudarc::driver::{LaunchConfig, PushKernelArg}; use onnx_runtime_ep_api::{EpError, Result}; use crate::error::driver_err; -use crate::kernels::kv_stride::KvCacheStrides; +use crate::kernels::kv_stride::{KvCachePath, KvCacheStrides}; use crate::runtime::CudaRuntime; const ENTRY: &str = "gqa_decode_attention_f16"; @@ -408,7 +408,7 @@ pub(super) fn run( kv_strides: &KvCacheStrides, ) -> Result<()> { runtime.require_nvrtc_half_headers("gqa_decode_attention_f16")?; - kv_strides.require_converted_path_support()?; + kv_strides.require_converted_path_support(KvCachePath::Fp16DecodeRead)?; let module_key = kv_strides.decode_module_key()?; let decode_src = format!("{}{}", kv_strides.decode_prelude(), DECODE_SRC); diff --git a/crates/onnx-runtime-ep-cuda/src/kernels/group_query_attention.rs b/crates/onnx-runtime-ep-cuda/src/kernels/group_query_attention.rs index a72cae3693..33bd2847a0 100644 --- a/crates/onnx-runtime-ep-cuda/src/kernels/group_query_attention.rs +++ b/crates/onnx-runtime-ep-cuda/src/kernels/group_query_attention.rs @@ -1,10 +1,11 @@ //! 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. +//! BSH query/key/value inputs are prepared into descriptor-selected KV 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. Declared present shapes remain BNSH-compatible +//! while the native backend may physically store converted paths as BSNH. use std::ffi::c_void; use std::sync::{Arc, Mutex}; @@ -15,7 +16,7 @@ use onnx_runtime_ep_api::{EpError, Kernel, KernelFactory, Result, TensorMut, Ten use onnx_runtime_ir::{DataType, Node}; use crate::error::driver_err; -use crate::kernels::kv_stride::KvCacheStrides; +use crate::kernels::kv_stride::{KvCachePath, KvCacheStrides}; use crate::runtime::{CudaRuntime, cuptr}; use super::attention::{AttentionDtype, run_attention_phase2a}; @@ -24,13 +25,15 @@ use super::gqa_decode; use super::gqa_decode_fp16; const PREP_SRC: &str = r#" -// Default (head-major BNSH) KV-cache append destination index. A non-default -// layout (seq-major BSNH) prepends its own `#define GQA_KV_DST` before this -// source, and this `#ifndef` guard yields to it, so the fused prep write indexes -// from the KV stride descriptor with no runtime layout branch. +// Default (head-major BNSH) KV-cache index. A non-default layout prepends its +// own descriptor-generated definition before this source. +#ifndef GQA_KV_INDEX +#define GQA_KV_INDEX(b, h, slot, heads, capacity, dim) \ + ( ((long)((b) * (heads) + (h)) * (capacity) + (slot)) * (dim) ) +#endif #ifndef GQA_KV_DST #define GQA_KV_DST(b, h, slot) \ - ( ((long)((b) * kv_heads + (h)) * present_capacity + (slot)) * dim ) + GQA_KV_INDEX(b, h, slot, kv_heads, present_capacity, dim) #endif extern "C" __global__ void gqa_transpose_bsh_to_bnsh( const float* src, float* dst, int batch, int seq, int heads, int dim) @@ -177,12 +180,12 @@ extern "C" __global__ void gqa_build_cache( if (past_len < 0) return; float value = 0.0f; if (s < past_len && past) { - value = past[((b * heads + h) * past_capacity + s) * dim + d]; + value = past[GQA_KV_INDEX(b, h, s, heads, past_capacity, 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; + present[GQA_KV_INDEX(b, h, s, heads, present_capacity, dim) + d] = value; } extern "C" __global__ void gqa_append_cache( @@ -198,7 +201,7 @@ extern "C" __global__ void gqa_append_cache( 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] = + present[GQA_KV_INDEX(b, h, target_s, heads, present_capacity, dim) + d] = current[((b * seq + s) * heads + h) * dim + d]; } @@ -227,7 +230,9 @@ extern "C" __global__ void gqa_rope_bnsh( 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 size_t base = current_offset + ? GQA_KV_INDEX(b, h, tensor_s, heads, tensor_capacity, dim) + : ((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]; @@ -437,11 +442,15 @@ const PREP_HALF_SRC: &str = r#" #include #include -// Default (head-major BNSH) KV-cache append destination index; a seq-major -// layout prepends its own `#define GQA_KV_DST` and this guard yields to it. +// Default (head-major BNSH) KV-cache index; a non-default layout prepends its +// descriptor-generated definition and these guards yield to it. +#ifndef GQA_KV_INDEX +#define GQA_KV_INDEX(b, h, slot, heads, capacity, dim) \ + ( ((long)((b) * (heads) + (h)) * (capacity) + (slot)) * (dim) ) +#endif #ifndef GQA_KV_DST #define GQA_KV_DST(b, h, slot) \ - ( ((long)((b) * kv_heads + (h)) * present_capacity + (slot)) * dim ) + GQA_KV_INDEX(b, h, slot, kv_heads, present_capacity, dim) #endif __device__ __forceinline__ int gqa_prepare_metadata_batch1( @@ -563,12 +572,12 @@ __device__ void gqa_build_cache_body( if (past_len < 0) return; T result = gqa_store(0.0f); if (s < past_len && past) { - result = past[((b * heads + h) * past_capacity + s) * dim + d]; + result = past[GQA_KV_INDEX(b, h, s, heads, past_capacity, 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; + present[GQA_KV_INDEX(b, h, s, heads, present_capacity, dim) + d] = result; } template @@ -585,7 +594,7 @@ __device__ void gqa_append_cache_body( 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] = + present[GQA_KV_INDEX(b, h, target_s, heads, present_capacity, dim) + d] = current[((b * seq + s) * heads + h) * dim + d]; } @@ -614,7 +623,9 @@ __device__ void gqa_rope_bnsh_body( 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 size_t base = current_offset + ? GQA_KV_INDEX(b, h, tensor_s, heads, tensor_capacity, dim) + : ((size_t)(b * heads + h) * tensor_capacity + tensor_s) * dim; const float x0 = gqa_load(tensor[base + d0]); const float x1 = gqa_load(tensor[base + d1]); const float c = gqa_load_cache(cos_cache, pos * half + k, cache_is_half); @@ -784,8 +795,8 @@ DEFINE_GQA_HALF_KERNELS(__half, f16) DEFINE_GQA_HALF_KERNELS(__nv_bfloat16, bf16) "#; -const PREP_MODULE: &str = "group_query_attention_prep_v3"; -const PREP_HALF_MODULE: &str = "group_query_attention_prep_half_v3"; +const PREP_MODULE: &str = "group_query_attention_prep_v4"; +const PREP_HALF_MODULE: &str = "group_query_attention_prep_half_v4"; const BLOCK: u32 = 256; const WS_TOTALS: usize = 0; const WS_PAST_LENGTHS: usize = 1; @@ -916,10 +927,9 @@ pub struct GroupQueryAttentionKernel { /// KV cache physical layout as a stride descriptor. The default is /// head-major BNSH (ORT-compatible), which every reader and writer honors. /// The native backend sets this to seq-major BSNH on GQA nodes whose KV it - /// exclusively owns; a non-default layout is honored solely on the fused - /// fp16 single-token decode path, and any other path rejects it. The - /// converted kernels generate their index arithmetic from this descriptor at - /// NVRTC module-build time, so there is no runtime layout branch. + /// exclusively owns. Converted flash-prefill and fp16 decode kernels + /// generate their index arithmetic from this descriptor at NVRTC + /// module-build time; every unconverted path rejects it. kv_strides: KvCacheStrides, workspace: Mutex, last_capture_safe_signature: Mutex>, @@ -1201,8 +1211,7 @@ impl GroupQueryAttentionKernel { /// Select the KV cache stride descriptor. The native backend sets a /// non-default (seq-major) descriptor on GQA nodes whose KV it exclusively - /// owns; a non-default layout is honored solely on the fused fp16 decode - /// path (see [`Self::run`]). + /// owns; converted paths are selected and hard-gated in [`Self::run`]. pub(crate) fn with_kv_strides(mut self, kv_strides: KvCacheStrides) -> Self { self.kv_strides = kv_strides; self @@ -1344,8 +1353,8 @@ impl GroupQueryAttentionKernel { } let element_size = dtype.element_size() as usize; let ( - prep_module, - prep_src, + base_prep_module, + base_prep_src, split_entry, transpose_in_entry, build_entry, @@ -1798,22 +1807,30 @@ impl GroupQueryAttentionKernel { && !self.prep_fusion_disabled; let fuse_metadata = fuse_prep && batch == 1; - // A non-default (seq-major BSNH) KV layout is a native-backend-only - // capability wired through exactly one path in this increment: the fused - // fp16 single-token decode prep write plus the fp16 split-K decode read. - // Every other prep path (unfused append/build/RoPE) and every other - // reader (flash / f32 decode / reference / phase2a) still assumes - // head-major BNSH, so reject a non-default cache that would reach them - // rather than silently corrupt it. - if !self.kv_strides.is_head_major() && !(fuse_prep && q.dtype == DataType::Float16) { - return Err(EpError::KernelFailed( - "cuda_ep GroupQueryAttention: seq-major (BSNH) KV requires the fused fp16 \ - single-token decode prep path (aliased in-place cache, past==present capacity, \ - even head_dim, q_seq==k_seq==1); the requested dtype/shape would route to \ - head-major-only append/build/RoPE kernels" - .into(), - )); - } + 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); + let use_fused = selected_backend == GroupQueryAttentionBackend::Fused; + let prep_path = if fuse_prep && q.dtype == DataType::Float16 { + KvCachePath::FusedDecodePrep + } else if q_seq > 1 && use_fused { + KvCachePath::FlashPrefillPrep + } else { + KvCachePath::UnfusedDecodePrep + }; + self.kv_strides.require_converted_path_support(prep_path)?; + let prep_src_owned; + let (prep_module, prep_src): (&str, &str) = if self.kv_strides.is_head_major() { + (base_prep_module, base_prep_src) + } else { + prep_src_owned = format!("{}{}", self.kv_strides.prep_prelude(), base_prep_src); + let module = match q.dtype { + DataType::Float32 => self.kv_strides.prep_f32_module_key()?, + DataType::Float16 | DataType::BFloat16 => self.kv_strides.prep_half_module_key()?, + _ => unreachable!(), + }; + (module, prep_src_owned.as_str()) + }; let mut workspace = self.workspace.lock().map_err(|_| { EpError::KernelFailed("cuda_ep GroupQueryAttention: workspace lock poisoned".into()) })?; @@ -1982,25 +1999,10 @@ impl GroupQueryAttentionKernel { 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); - // The fused write indexes the KV cache from the stride descriptor. - // Head-major reuses the shared prep module (its `#ifndef GQA_KV_DST` - // default write formula); a seq-major layout prepends a specialized - // `GQA_KV_DST` prelude and loads a distinct, per-layout module. - let fuse_src_owned: String; - let (fuse_module, fuse_src): (&str, &str) = if self.kv_strides.is_head_major() { - (prep_module, prep_src) - } else { - self.kv_strides.require_converted_path_support()?; - fuse_src_owned = format!("{}{}", self.kv_strides.prep_prelude(), prep_src); - ( - self.kv_strides.prep_half_module_key()?, - fuse_src_owned.as_str(), - ) - }; launch_1d!( self.runtime, - fuse_module, - fuse_src, + prep_module, + prep_src, fuse_entry, fused_count, builder, @@ -2224,9 +2226,6 @@ impl GroupQueryAttentionKernel { .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) @@ -2243,17 +2242,18 @@ impl GroupQueryAttentionKernel { .saturating_add(pv_flops) .saturating_add(softmax_flops) }); - let use_fused = selected_backend == GroupQueryAttentionBackend::Fused; - if !self.kv_strides.is_head_major() - && !(!use_fused && gqa_decode_fp16::supported(q_seq, dim)) - { - return Err(EpError::KernelFailed( - "cuda_ep GroupQueryAttention: seq-major (BSNH) KV must read through the \ - capture-safe fp16 split-K decode kernel; the selected reader (fused flash / \ - f32 decode / reference / phase2a) still expects head-major BNSH" - .into(), - )); - } + let read_path = if use_fused { + KvCachePath::FlashPrefillRead + } else if q.dtype == DataType::Float32 && gqa_decode::supported(q_seq, dim) { + KvCachePath::F32DecodeRead + } else if q.dtype == DataType::Float16 && gqa_decode_fp16::supported(q_seq, dim) { + KvCachePath::Fp16DecodeRead + } else if q.dtype == DataType::Float32 { + KvCachePath::ReferenceRead + } else { + KvCachePath::Phase2aRead + }; + self.kv_strides.require_converted_path_support(read_path)?; if use_fused { onnx_runtime_ep_api::record_kernel_variant!( "attention_flash_fused", @@ -2288,6 +2288,7 @@ impl GroupQueryAttentionKernel { query_starts_gpu, local_window_i, self.softcap, + &self.kv_strides, )?; } else if q.dtype == DataType::Float32 && gqa_decode::supported(q_seq, dim) { onnx_runtime_ep_api::record_kernel_variant!( diff --git a/crates/onnx-runtime-ep-cuda/src/kernels/kv_stride.rs b/crates/onnx-runtime-ep-cuda/src/kernels/kv_stride.rs index 7ac1303ba5..3411105d0a 100644 --- a/crates/onnx-runtime-ep-cuda/src/kernels/kv_stride.rs +++ b/crates/onnx-runtime-ep-cuda/src/kernels/kv_stride.rs @@ -1,5 +1,4 @@ -//! Descriptor-driven KV-cache stride indexing for the converted GQA decode -//! path. +//! Descriptor-driven KV-cache stride indexing for converted GQA paths. //! //! The CUDA EP used to select the KV cache physical layout with a two-valued //! `kv_layout` integer (`0` = head-major BNSH, `1` = seq-major BSNH) that every @@ -29,6 +28,41 @@ pub(crate) enum KvDim { HeadDim, } +/// A GQA path that may touch the physical KV-cache layout. +#[derive(Clone, Copy, PartialEq, Eq, Debug)] +pub(crate) enum KvCachePath { + FusedDecodePrep, + FlashPrefillPrep, + FlashPrefillRead, + Fp16DecodeRead, + UnfusedDecodePrep, + F32DecodeRead, + ReferenceRead, + Phase2aRead, +} + +impl KvCachePath { + fn name(self) -> &'static str { + match self { + Self::FusedDecodePrep => "fused fp16 decode prep", + Self::FlashPrefillPrep => "flash prefill prep", + Self::FlashPrefillRead => "flash prefill read", + Self::Fp16DecodeRead => "fp16 split-K decode read", + Self::UnfusedDecodePrep => "unfused decode prep", + Self::F32DecodeRead => "f32 split-K decode read", + Self::ReferenceRead => "reference attention read", + Self::Phase2aRead => "phase2a attention read", + } + } +} + +const CONVERTED_PATHS: &[KvCachePath] = &[ + KvCachePath::FusedDecodePrep, + KvCachePath::FlashPrefillPrep, + KvCachePath::FlashPrefillRead, + KvCachePath::Fp16DecodeRead, +]; + /// Symbolic per-axis strides of a KV-cache binding, general over layout. /// /// The concrete element stride of an axis is the product of the runtime @@ -98,18 +132,15 @@ impl KvCacheStrides { } /// Whether this is the head-major default layout that every reader and - /// writer honors. Non-default layouts are gated to the one converted path. + /// writer honors. Non-default layouts are gated to converted paths. pub(crate) fn is_head_major(&self) -> bool { self.named == "bnsh" } - /// Reject any descriptor the converted single-token decode path cannot - /// honor, rather than mis-indexing it. The path requires a contiguous - /// (unit-stride) `head_dim` innermost axis — the fp16 read vectorizes it as - /// `half2` and the fused write addresses it as `dst + d` — and it addresses - /// a whole-buffer binding, so a non-zero reservation offset or a - /// reservation-spanning seq extent (a token-major view) is not honored yet. - pub(crate) fn require_converted_path_support(&self) -> Result<()> { + /// Reject a descriptor/path pair that is not converted, rather than + /// silently mis-indexing it. Converted paths require a contiguous + /// (unit-stride) `head_dim` innermost axis and whole-buffer bindings. + pub(crate) fn require_converted_path_support(&self, path: KvCachePath) -> Result<()> { if !self.head_dim.is_empty() { return Err(EpError::KernelFailed( "cuda_ep GroupQueryAttention: KV descriptor has a non-unit head_dim stride; the \ @@ -132,6 +163,14 @@ impl KvCacheStrides { .into(), )); } + if !self.is_head_major() && !CONVERTED_PATHS.contains(&path) { + return Err(EpError::KernelFailed(format!( + "cuda_ep GroupQueryAttention: seq-major (BSNH) KV cannot use the {} path; \ + converted paths are fused fp16 decode prep, flash prefill prep/read, and fp16 \ + split-K decode read", + path.name() + ))); + } Ok(()) } @@ -144,13 +183,34 @@ impl KvCacheStrides { } } - /// The `'static` module-cache key for a non-default (seq-major) fp16/bf16 - /// prep module. Head-major reuses the shared prep module const (with its - /// `#ifndef` default write formula), so only the specialized layout needs a - /// distinct key here. + /// The `'static` module-cache key for a non-default f32 prep module. + pub(crate) fn prep_f32_module_key(&self) -> Result<&'static str> { + match self.named { + "bsnh" => Ok("group_query_attention_prep_v4_bsnh"), + _ => Err(self.no_cached_module_err()), + } + } + + /// The `'static` module-cache key for a non-default fp16/bf16 prep module. pub(crate) fn prep_half_module_key(&self) -> Result<&'static str> { match self.named { - "bsnh" => Ok("group_query_attention_prep_half_v3_bsnh"), + "bsnh" => Ok("group_query_attention_prep_half_v4_bsnh"), + _ => Err(self.no_cached_module_err()), + } + } + + pub(crate) fn flash_f32_module_key(&self) -> Result<&'static str> { + match self.named { + "bnsh" => Ok("flash_attention_f32_v2_bnsh"), + "bsnh" => Ok("flash_attention_f32_v2_bsnh"), + _ => Err(self.no_cached_module_err()), + } + } + + pub(crate) fn flash_half_module_key(&self) -> Result<&'static str> { + match self.named { + "bnsh" => Ok("flash_attention_half_v3_bnsh"), + "bsnh" => Ok("flash_attention_half_v3_bsnh"), _ => Err(self.no_cached_module_err()), } } @@ -177,18 +237,30 @@ impl KvCacheStrides { ) } - /// The CUDA `#define` prelude that specializes the fused decode-prep write - /// kernel's destination-index arithmetic for this descriptor. The macro - /// references the kernel's in-scope `kv_heads`, `present_capacity`, and - /// `dim` variables. Prepended only for a non-default layout; head-major uses - /// the `#ifndef` default baked into the prep source. + /// The CUDA `#define` prelude that specializes flash/prefill KV reads. + pub(crate) fn flash_prelude(&self) -> String { + let batch = product_expr(&self.batch, "kv_heads", "kv_capacity", "dim"); + let head = product_expr(&self.head, "kv_heads", "kv_capacity", "dim"); + let seq = product_expr(&self.seq, "kv_heads", "kv_capacity", "dim"); + format!( + "#define FLASH_KV_BASE(b, h) ( (long)(b) * {batch} + (long)(h) * {head} )\n\ + #define FLASH_KV_STRIDE ( {seq} )\n" + ) + } + + /// The CUDA prelude that specializes prep cache reads/writes. The generic + /// index macro accepts the runtime head count, capacity, and head dimension + /// so the build path can address past and present buffers with different + /// capacities. pub(crate) fn prep_prelude(&self) -> String { - let batch = product_expr(&self.batch, "kv_heads", "present_capacity", "dim"); - let head = product_expr(&self.head, "kv_heads", "present_capacity", "dim"); - let seq = product_expr(&self.seq, "kv_heads", "present_capacity", "dim"); + let batch = product_expr(&self.batch, "heads", "capacity", "dim"); + let head = product_expr(&self.head, "heads", "capacity", "dim"); + let seq = product_expr(&self.seq, "heads", "capacity", "dim"); format!( - "#define GQA_KV_DST(b, h, slot) ( (long)(b) * {batch} \ - + (long)(h) * {head} + (long)(slot) * {seq} )\n" + "#define GQA_KV_INDEX(b, h, slot, heads, capacity, dim) ( (long)(b) * {batch} \ + + (long)(h) * {head} + (long)(slot) * {seq} )\n\ + #define GQA_KV_DST(b, h, slot) \ + GQA_KV_INDEX(b, h, slot, kv_heads, present_capacity, dim)\n" ) } } @@ -225,14 +297,19 @@ mod tests { fn named_layouts_are_honored() { assert!( KvCacheStrides::head_major_bnsh() - .require_converted_path_support() + .require_converted_path_support(KvCachePath::ReferenceRead) .is_ok() ); assert!( KvCacheStrides::seq_major_bsnh() - .require_converted_path_support() + .require_converted_path_support(KvCachePath::FlashPrefillRead) .is_ok() ); + assert!( + KvCacheStrides::seq_major_bsnh() + .require_converted_path_support(KvCachePath::ReferenceRead) + .is_err() + ); } // A descriptor the converted path cannot honor must be rejected with an @@ -242,7 +319,9 @@ mod tests { let mut view = KvCacheStrides::seq_major_bsnh(); view.offset_elements = 4096; view.reservation_override = true; - let error = view.require_converted_path_support().unwrap_err(); + let error = view + .require_converted_path_support(KvCachePath::Fp16DecodeRead) + .unwrap_err(); assert!(matches!(error, EpError::KernelFailed(_))); assert!(view.decode_module_key().is_ok()); } @@ -251,7 +330,10 @@ mod tests { fn non_unit_head_dim_is_rejected() { let mut bad = KvCacheStrides::head_major_bnsh(); bad.head_dim = vec![KvDim::KvHeads]; - assert!(bad.require_converted_path_support().is_err()); + assert!( + bad.require_converted_path_support(KvCachePath::Fp16DecodeRead) + .is_err() + ); } // The generated arithmetic must reproduce the two hand-written layouts. @@ -276,6 +358,14 @@ mod tests { // distributed sum b*(N*cap*dim) + h*(cap*dim) + slot*dim. assert!(bnsh.contains("(long)(slot) * ((long)dim)")); let bsnh = KvCacheStrides::seq_major_bsnh().prep_prelude(); - assert!(bsnh.contains("(long)(slot) * ((long)kv_heads * (long)dim)")); + assert!(bsnh.contains("(long)(slot) * ((long)heads * (long)dim)")); + } + + #[test] + fn flash_prelude_matches_layout_formulae() { + let bnsh = KvCacheStrides::head_major_bnsh().flash_prelude(); + assert!(bnsh.contains("#define FLASH_KV_STRIDE ( ((long)dim) )")); + let bsnh = KvCacheStrides::seq_major_bsnh().flash_prelude(); + assert!(bsnh.contains("#define FLASH_KV_STRIDE ( ((long)kv_heads * (long)dim) )")); } } diff --git a/crates/onnx-runtime-ep-cuda/tests/gqa_seqmajor_parity_gpu.rs b/crates/onnx-runtime-ep-cuda/tests/gqa_seqmajor_parity_gpu.rs index 6ebae9d1fd..85bb7224fb 100644 --- a/crates/onnx-runtime-ep-cuda/tests/gqa_seqmajor_parity_gpu.rs +++ b/crates/onnx-runtime-ep-cuda/tests/gqa_seqmajor_parity_gpu.rs @@ -20,7 +20,9 @@ use onnx_runtime_ep_api::{ DeviceBuffer, DevicePtr, DevicePtrMut, ExecutionProvider, Kernel, TensorMut, TensorView, }; use onnx_runtime_ep_cuda::runtime::cuptr; -use onnx_runtime_ep_cuda::{CudaExecutionProvider, GroupQueryAttentionKernel}; +use onnx_runtime_ep_cuda::{ + CudaExecutionProvider, GroupQueryAttentionBackend, GroupQueryAttentionKernel, +}; use onnx_runtime_ir::{DataType, compute_contiguous_strides}; const BATCH: usize = 1; @@ -133,6 +135,45 @@ fn cpu_reference( out } +fn cpu_prefill_reference( + query: &[f16], + key: &[Vec>], + value: &[Vec>], + sequence_length: usize, + scale: f32, +) -> Vec { + let mut out = vec![0.0_f32; sequence_length * QUERY_HEADS * HEAD_DIM]; + for query_pos in 0..sequence_length { + for qh in 0..QUERY_HEADS { + let kvh = qh / GROUP; + let mut scores = vec![0.0_f32; query_pos + 1]; + let mut max_score = f32::NEG_INFINITY; + for key_pos in 0..=query_pos { + let mut dot = 0.0_f32; + for d in 0..HEAD_DIM { + let q_index = (query_pos * QUERY_HEADS + qh) * HEAD_DIM + d; + dot += query[q_index].to_f32() * key[kvh][key_pos][d].to_f32(); + } + scores[key_pos] = dot * scale; + max_score = max_score.max(scores[key_pos]); + } + let mut denominator = 0.0_f32; + for score in &mut scores { + *score = (*score - max_score).exp(); + denominator += *score; + } + for d in 0..HEAD_DIM { + let mut acc = 0.0_f32; + for key_pos in 0..=query_pos { + acc += scores[key_pos] / denominator * value[kvh][key_pos][d].to_f32(); + } + out[(query_pos * QUERY_HEADS + qh) * HEAD_DIM + d] = acc; + } + } + } + out +} + fn fp16_values(bytes: &[u8]) -> Vec { bytes .chunks_exact(2) @@ -394,3 +435,357 @@ fn seq_major_decode_is_bit_identical_to_head_major() { ep.deallocate(buffer).unwrap(); } } + +#[cfg_attr( + not(feature = "gpu-tests"), + ignore = "requires CUDA device; enable the gpu-tests feature on a CUDA runner" +)] +#[test] +fn seq_major_full_generation_prefill_and_decode_is_bit_identical() { + const PREFILL_LEN: usize = 10; + + let ep = require_cuda(); + let runtime = ep.runtime(); + let scale = 1.0_f32; + let device = ep.device_id(); + + let prefill_query: Vec = (0..PREFILL_LEN * QUERY_HEADS * HEAD_DIM) + .map(|i| f16::from_f32(((i * 17 % 97) as f32 - 48.0) / 256.0)) + .collect(); + let decode_query: Vec = (0..QUERY_HEADS * HEAD_DIM) + .map(|i| f16::from_f32(((i * 23 % 89) as f32 - 44.0) / 256.0)) + .collect(); + let mut key = vec![vec![vec![f16::ZERO; HEAD_DIM]; PREFILL_LEN + 1]; KV_HEADS]; + let mut value = vec![vec![vec![f16::ZERO; HEAD_DIM]; PREFILL_LEN + 1]; KV_HEADS]; + for h in 0..KV_HEADS { + for t in 0..=PREFILL_LEN { + for d in 0..HEAD_DIM { + key[h][t][d] = + f16::from_f32((((h * 131 + t * 13 + d * 7) % 101) as f32 - 50.0) / 256.0); + value[h][t][d] = + f16::from_f32((((h * 29 + t * 19 + d * 3) % 113) as f32 - 56.0) / 128.0); + } + } + } + let flatten_current = |position: usize| { + let mut values = vec![f16::ZERO; KV_HEADS * HEAD_DIM]; + for h in 0..KV_HEADS { + for d in 0..HEAD_DIM { + values[h * HEAD_DIM + d] = key[h][position][d]; + } + } + values + }; + let flatten_prefill = |logical: &[Vec>]| { + let mut values = vec![f16::ZERO; PREFILL_LEN * KV_HEADS * HEAD_DIM]; + for t in 0..PREFILL_LEN { + for h in 0..KV_HEADS { + for d in 0..HEAD_DIM { + values[(t * KV_HEADS + h) * HEAD_DIM + d] = logical[h][t][d]; + } + } + } + values + }; + let prefill_key = flatten_prefill(&key); + let prefill_value = flatten_prefill(&value); + let decode_key = flatten_current(PREFILL_LEN); + let decode_value = { + let mut values = vec![f16::ZERO; KV_HEADS * HEAD_DIM]; + for h in 0..KV_HEADS { + for d in 0..HEAD_DIM { + values[h * HEAD_DIM + d] = value[h][PREFILL_LEN][d]; + } + } + values + }; + + let prefill_query_shape = [BATCH, PREFILL_LEN, QUERY_HEADS * HEAD_DIM]; + let prefill_current_shape = [BATCH, PREFILL_LEN, KV_HEADS * HEAD_DIM]; + let decode_query_shape = [BATCH, 1, QUERY_HEADS * HEAD_DIM]; + let decode_current_shape = [BATCH, 1, KV_HEADS * HEAD_DIM]; + let cache_shape = [BATCH, KV_HEADS, CACHE_CAPACITY, HEAD_DIM]; + let seqlens_shape = [BATCH]; + let scalar_shape: [usize; 0] = []; + let prefill_query_strides = compute_contiguous_strides(&prefill_query_shape); + let prefill_current_strides = compute_contiguous_strides(&prefill_current_shape); + let cache_strides = compute_contiguous_strides(&cache_shape); + let seqlens_strides = compute_contiguous_strides(&seqlens_shape); + let scalar_strides = compute_contiguous_strides(&scalar_shape); + + let prefill_query_buffer = upload(&ep, typed_bytes(&prefill_query)).unwrap(); + let prefill_key_buffer = upload(&ep, typed_bytes(&prefill_key)).unwrap(); + let prefill_value_buffer = upload(&ep, typed_bytes(&prefill_value)).unwrap(); + let decode_query_buffer = upload(&ep, typed_bytes(&decode_query)).unwrap(); + let decode_key_buffer = upload(&ep, typed_bytes(&decode_key)).unwrap(); + let decode_value_buffer = upload(&ep, typed_bytes(&decode_value)).unwrap(); + let prefill_seqlens = [(PREFILL_LEN - 1) as i32]; + let decode_seqlens = [PREFILL_LEN as i32]; + let capacity = [CACHE_CAPACITY as i32]; + let prefill_seqlens_buffer = upload(&ep, typed_bytes(&prefill_seqlens)).unwrap(); + let decode_seqlens_buffer = upload(&ep, typed_bytes(&decode_seqlens)).unwrap(); + let capacity_buffer = upload(&ep, typed_bytes(&capacity)).unwrap(); + + let run_layout = |seq_major: bool, capture: bool| -> (Vec, Vec) { + let kernel = GroupQueryAttentionKernel::new( + runtime.clone(), + QUERY_HEADS, + KV_HEADS, + Some(scale), + false, + false, + -1, + 0.0, + ) + .unwrap() + .with_backend(GroupQueryAttentionBackend::Fused) + .with_kv_layout(if seq_major { 1 } else { 0 }); + let mut cache_key = ep + .allocate( + BATCH * KV_HEADS * CACHE_CAPACITY * HEAD_DIM * std::mem::size_of::(), + 256, + ) + .unwrap(); + let mut cache_value = ep + .allocate( + BATCH * KV_HEADS * CACHE_CAPACITY * HEAD_DIM * std::mem::size_of::(), + 256, + ) + .unwrap(); + let mut prefill_output = ep + .allocate( + PREFILL_LEN * QUERY_HEADS * HEAD_DIM * std::mem::size_of::(), + 256, + ) + .unwrap(); + let mut decode_output = ep + .allocate(QUERY_HEADS * HEAD_DIM * std::mem::size_of::(), 256) + .unwrap(); + + { + let prefill_inputs = [ + TensorView::new( + DevicePtr(prefill_query_buffer.as_ptr()), + DataType::Float16, + &prefill_query_shape, + &prefill_query_strides, + device, + ), + TensorView::new( + DevicePtr(prefill_key_buffer.as_ptr()), + DataType::Float16, + &prefill_current_shape, + &prefill_current_strides, + device, + ), + TensorView::new( + DevicePtr(prefill_value_buffer.as_ptr()), + DataType::Float16, + &prefill_current_shape, + &prefill_current_strides, + device, + ), + TensorView::absent(DataType::Float16), + TensorView::absent(DataType::Float16), + TensorView::new( + DevicePtr(prefill_seqlens_buffer.as_ptr()), + DataType::Int32, + &seqlens_shape, + &seqlens_strides, + device, + ), + TensorView::new( + DevicePtr(capacity_buffer.as_ptr()), + DataType::Int32, + &scalar_shape, + &scalar_strides, + device, + ), + ]; + let mut prefill_outputs = [ + TensorMut::new( + DevicePtrMut(prefill_output.as_mut_ptr()), + DataType::Float16, + &prefill_query_shape, + &prefill_query_strides, + device, + ), + TensorMut::new( + DevicePtrMut(cache_key.as_mut_ptr()), + DataType::Float16, + &cache_shape, + &cache_strides, + device, + ), + TensorMut::new( + DevicePtrMut(cache_value.as_mut_ptr()), + DataType::Float16, + &cache_shape, + &cache_strides, + device, + ), + ]; + kernel + .execute(&prefill_inputs, &mut prefill_outputs) + .unwrap(); + } + + { + let decode_query_strides = compute_contiguous_strides(&decode_query_shape); + let decode_current_strides = compute_contiguous_strides(&decode_current_shape); + let decode_inputs = [ + TensorView::new( + DevicePtr(decode_query_buffer.as_ptr()), + DataType::Float16, + &decode_query_shape, + &decode_query_strides, + device, + ), + TensorView::new( + DevicePtr(decode_key_buffer.as_ptr()), + DataType::Float16, + &decode_current_shape, + &decode_current_strides, + device, + ), + TensorView::new( + DevicePtr(decode_value_buffer.as_ptr()), + DataType::Float16, + &decode_current_shape, + &decode_current_strides, + device, + ), + TensorView::new( + DevicePtr(cache_key.as_ptr()), + DataType::Float16, + &cache_shape, + &cache_strides, + device, + ), + TensorView::new( + DevicePtr(cache_value.as_ptr()), + DataType::Float16, + &cache_shape, + &cache_strides, + device, + ), + TensorView::new( + DevicePtr(decode_seqlens_buffer.as_ptr()), + DataType::Int32, + &seqlens_shape, + &seqlens_strides, + device, + ), + TensorView::new( + DevicePtr(capacity_buffer.as_ptr()), + DataType::Int32, + &scalar_shape, + &scalar_strides, + device, + ), + ]; + let mut decode_outputs = [ + TensorMut::new( + DevicePtrMut(decode_output.as_mut_ptr()), + DataType::Float16, + &decode_query_shape, + &decode_query_strides, + device, + ), + TensorMut::new( + DevicePtrMut(cache_key.as_mut_ptr()), + DataType::Float16, + &cache_shape, + &cache_strides, + device, + ), + TensorMut::new( + DevicePtrMut(cache_value.as_mut_ptr()), + DataType::Float16, + &cache_shape, + &cache_strides, + device, + ), + ]; + kernel.execute(&decode_inputs, &mut decode_outputs).unwrap(); + if capture { + runtime.begin_graph_capture(&[]).unwrap(); + kernel.execute(&decode_inputs, &mut decode_outputs).unwrap(); + runtime.end_graph_capture().unwrap(); + assert_eq!( + runtime.graph_segment_count().unwrap(), + 1, + "full-generation decode must install one captured segment" + ); + runtime.replay_graph().unwrap(); + runtime.reset_graph().unwrap(); + } + } + + let prefill_bytes = read( + &ep, + &prefill_output, + PREFILL_LEN * QUERY_HEADS * HEAD_DIM * std::mem::size_of::(), + ) + .unwrap(); + let decode_bytes = read( + &ep, + &decode_output, + QUERY_HEADS * HEAD_DIM * std::mem::size_of::(), + ) + .unwrap(); + for buffer in [decode_output, prefill_output, cache_value, cache_key] { + ep.deallocate(buffer).unwrap(); + } + (prefill_bytes, decode_bytes) + }; + + let (bnsh_prefill, bnsh_decode) = run_layout(false, false); + let (bsnh_prefill, bsnh_decode) = run_layout(true, false); + let (bnsh_captured_prefill, bnsh_captured_decode) = run_layout(false, true); + let (bsnh_captured_prefill, bsnh_captured_decode) = run_layout(true, true); + assert_eq!(bnsh_prefill, bsnh_prefill); + assert_eq!(bnsh_decode, bsnh_decode); + assert_eq!(bnsh_prefill, bnsh_captured_prefill); + assert_eq!(bnsh_decode, bnsh_captured_decode); + assert_eq!(bsnh_prefill, bsnh_captured_prefill); + assert_eq!(bsnh_decode, bsnh_captured_decode); + + let prefill_expected = cpu_prefill_reference(&prefill_query, &key, &value, PREFILL_LEN, scale); + let prefill_got = fp16_values(&bnsh_prefill); + let prefill_max_err = prefill_got + .iter() + .zip(&prefill_expected) + .map(|(got, expected)| (got - expected).abs()) + .fold(0.0_f32, f32::max); + assert!( + prefill_max_err < 3e-3, + "head-major GPU prefill diverged from CPU oracle: max_abs={prefill_max_err:e}" + ); + let decode_expected = cpu_reference(&decode_query, &key, &value, scale); + let decode_got = fp16_values(&bnsh_decode); + let decode_max_err = decode_got + .iter() + .zip(&decode_expected) + .map(|(got, expected)| (got - expected).abs()) + .fold(0.0_f32, f32::max); + assert!( + decode_max_err < 3e-3, + "head-major GPU full-generation decode diverged from CPU oracle: \ + max_abs={decode_max_err:e}" + ); + + for buffer in [ + capacity_buffer, + decode_seqlens_buffer, + prefill_seqlens_buffer, + decode_value_buffer, + decode_key_buffer, + decode_query_buffer, + prefill_value_buffer, + prefill_key_buffer, + prefill_query_buffer, + ] { + ep.deallocate(buffer).unwrap(); + } +} diff --git a/docs/MEMORY_ARCHITECTURE.md b/docs/MEMORY_ARCHITECTURE.md index 83fab4b87a..d011fe60b2 100644 --- a/docs/MEMORY_ARCHITECTURE.md +++ b/docs/MEMORY_ARCHITECTURE.md @@ -154,20 +154,20 @@ than copying conclusions, following the anti-staleness rule established in **Status, as of 2026-08-11:** -- **Implemented:** head-major BNSH remains the default. PR #782 added a - per-backend `KvLayout` tag (default `HeadMajorBnsh`) and landed the - seq-major BSNH fused fp16 single-token decode append/read pair. Seq-major is - hard-gated to those converted paths; unsupported readers and writers fail - rather than silently mis-index. It is not yet an end-to-end layout for every - attention path. +- **Implemented:** head-major BNSH remains the default. PR #782 landed the + seq-major BSNH fused fp16 single-token decode append/read pair, and #792 + replaced its layout enum with a symbolic stride descriptor and static + per-layout NVRTC specialization. Flash prefill now uses the same descriptor + for cache preparation and reads, enabling full seq-major prefill + fp16 + decode generations when the fused-flash shape gate applies. Unsupported + readers and writers still fail rather than silently mis-index. - **Measured, not implemented:** token-major across all layers. Its residency floor and 192 KiB-stride read cost were measured in #787, but no production kernel or binding layout uses it. -- **Proposed, in progress:** replace the growing layout enum with a stride - descriptor and make layout selection a per-EP, per-platform capability - (#783). The descriptor, binding views, layout negotiation, prefix multi-map, - staggering, and commit-ahead policy described below are design intent, not - present-tense capabilities. +- **Proposed, in progress:** make layout selection a per-EP, per-platform + capability (#783). Binding views, layout negotiation, prefix multi-map, + staggering, and commit-ahead policy described below remain design intent, + not present-tense capabilities. #### Governing rule: layout controls residency From d63a86c21b48ad3791e5c98d249dccb579c2ddec Mon Sep 17 00:00:00 2001 From: justinchuby <223556219+Copilot@users.noreply.github.com> Date: Wed, 12 Aug 2026 04:18:40 -0700 Subject: [PATCH 2/2] docs(squad): record what the VMM/offload/streaming/batching push established A durable record of the load-bearing results, including the claims that were refuted. Three of this push's most useful outcomes were corrections: seq-major had not realised the 8x floor (bindings ignored the layout descriptor), the decoder was not declining capture (a cached env var in a long-lived test process), and a fixed KV stride does not remove re-capture (the engine invalidates unconditionally on growth). Also records the #736 audit's recurring finding -- four of five slices found over-reservation rather than an ungoverned allocation -- and the guidance that follows from it. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: d60eb808-7cc6-4abc-b48d-2a6dd3841624 --- .../inbox/copilot-vmm-push-summary.md | 97 +++++++++++++++++++ 1 file changed, 97 insertions(+) create mode 100644 .squad/decisions/inbox/copilot-vmm-push-summary.md diff --git a/.squad/decisions/inbox/copilot-vmm-push-summary.md b/.squad/decisions/inbox/copilot-vmm-push-summary.md new file mode 100644 index 0000000000..fd6c9689f1 --- /dev/null +++ b/.squad/decisions/inbox/copilot-vmm-push-summary.md @@ -0,0 +1,97 @@ +### 2026-08-12: What the VMM / offload / streaming / batching push established +**By:** Copilot (coordinator) +**What:** A record of the load-bearing results from this push, so the reasoning +survives in the tree rather than only in issue threads. + +Every claim below is backed by a merged, executable test. Where a claim was +*refuted*, that is recorded too — the refutations were often worth more than the +confirmations. + +#### The governing rule (measured, #772/#776/#787) + +`cuMemMap` maps whole granule-aligned windows onto whole physical granules, so +`committed bytes = granule x (windows containing at least one live byte)`. +**Layout therefore controls residency**, because layout decides where the live +bytes fall in virtual address space. The allocator cannot compact what the layout +scattered. + +`CU_MEM_ALLOC_GRANULARITY_MINIMUM == RECOMMENDED == 2 MiB` on this device, so the +floor cannot be fixed by shrinking the granule — only by layout. This is a +**CUDA-specific severity**: minimum mapping granularity spans ~500x across +platforms (Level Zero and Vulkan ~64 KiB, CPU `mmap` 4 KiB), which is why layout +must be a queried per-EP, per-platform capability (#783) rather than a constant. + +#### Confirmed + +- **Floor is layout-determined:** 768 granules (~1.5 GiB) head-major -> 96 (~192 MiB) + seq-major -> 1 per sequence (~2 MiB) token-major, the last measured at a **768x** + reduction (#787). +- **Strided reads are not the obstacle:** seq/head bandwidth ratio 0.80-1.02, and + a 192 KB token-major stride measured **1.000** at a 6 GiB working set — reads are + DRAM-bound independently of stride, because device memory is already 2 MiB-page + backed (#778/#787). +- **Offload and capture are no longer mutually exclusive** (#796). Weights page + under a stable VA, so page-in remaps physical granules instead of returning a new + pointer. This unblocked #755. +- **Managed no-spill VMM is the default**, with automatic weight streaming when a + model exceeds the budget (#798). A fitting model does **not** page: verified + `FullResident`, offload off, **0 page-ins**. +- **Prefix sharing is sound:** one handle maps into N=8 sequences under captured + replay; the ledger charges it **once**, keeps it alive until the last sharer, and + an additional sharer costs **0** incremental bytes (#793/#803). + +#### Refuted, and why the refutation mattered + +- **"seq-major landed means the 8x floor is realised."** It did not: #794 measured + head-major and seq-major committing **identical** physical bytes, because the + bindings did not consume the layout descriptor. Fixed in #797. I had stated the + 8x as achieved and had to correct the floor table before merging. +- **"The decoder structurally declines capture for this model."** It does not + (#804). `captures=0` came from a cached `ONNX_GENAI_CUDA_GRAPH=0` in a + long-lived test process. **Two PRs (#794, #801) misattributed a failed + measurement to this.** +- **"A fixed KV stride removes growth-triggered re-capture."** True in mechanism, + irrelevant in practice: seq-major moves **0 bytes** on growth versus head-major's + 688,576, and the graph is invalidated anyway because the engine invalidates + **unconditionally** on growth (#805). The real blocker was never the layout. +- **"Tokens per granule" as a KV cost model.** Wrong for head-major — publicly + retracted — yet exactly correct for token-major. Same granule, same model, same + VMM; only the layout differs. It is the clearest demonstration that layout is the + whole story. + +#### The audit's recurring finding (#736, six slices) + +Four of five completed slices found **over-reservation** — bytes charged against +the device authority on a path that never uses them — rather than the ungoverned +allocation the audit was framed to look for: + +| Slice | Finding | +|---|---| +| #751 IndexShare | staging the common 3-output decode path never touches | +| #795 GQA `WS_SCORES` | needed only on the f32 reference route (~128 MiB hot path) | +| #799 cuBLASLt GEMM | 32 MiB was a heuristic **ceiling**; measured use 0-96 bytes | +| #802 default-domain Attention scores | **genuinely needed** on every route | +| #806 GQA QKV staging | unpacked and fp16/bf16 fused-decode routes charge 0 | + +Guidance, now recorded in `MEMORY_ARCHITECTURE.md`: **start from use, not from +allocation.** Locating `alloc_raw` is the easy part; the valuable question is what +the bytes are used for on each dispatch route. And governing a bypass *without* +sizing it to use makes things worse — it converts invisible waste into charged +waste, tightening committed-granule admission (#745) and therefore reducing +admissible concurrency. + +#### Method notes that earned their keep + +- **Order-dependent test state cost two wrong conclusions this week** — a + process-frozen `RuntimeConfig` (#804) and a CUDA context warmed by an + alphabetically-earlier sibling (#797). On a measurement-driven line a wrong + number does not merely fail, it redirects design. #807 added a debug-only freeze + guard and a single-stream test helper, plus an inventory of the remaining + order-dependent tests. +- **Negative results were delivered as first-class outcomes** — the granule lever + being unavailable, quantization being a split result (tokens/granule improve but + the byte floor does not), and #802's "genuinely needed". Each closed a question + rather than leaving it to be re-investigated. +- **Never extrapolate an unmeasured number.** Where a large-model run was + impossible (`qwen14b-zp` lacks `inference_metadata.yaml` and is not + native-loadable, #384), that was reported as not measured with the reason.