diff --git a/.agents/skills/ort-release-notes/SKILL.md b/.agents/skills/ort-release-notes/SKILL.md new file mode 100644 index 0000000000000..9b7587057ee39 --- /dev/null +++ b/.agents/skills/ort-release-notes/SKILL.md @@ -0,0 +1,151 @@ +--- +name: ort-release-notes +description: Draft ONNX Runtime release notes using preset configurations for full ORT or scoped component releases. Use when generating highlights with PR links, compiling human contributor acknowledgments from compile_contributors.py output, and applying preset path filtering. +argument-hint: "preset base_ref target_ref [version] [output_dir]" +--- + +# ONNX Runtime Release Notes + +Use this skill to produce a consistent release-note draft from commit history and contributor metadata. + +## When To Use + +Use this skill when you need to: + +- Draft release notes for a full ONNX Runtime release +- Draft release notes for a scoped component (e.g., in-tree plugin EP) release +- Select a release profile by preset name instead of manually supplying path/version files +- Add PR links to highlight bullets +- Build a human-only contributor acknowledgment list from contributor metadata + +## Required Inputs + +Collect these inputs from the user or infer from context: + +1. `preset`: release profile name (for example, `ort`, `webgpu-plugin-ep`, `cuda-plugin-ep`) +2. `base_ref`: previous release tag +3. `target_ref`: release commit/tag/branch tip + +Optional inputs: + +- `version` override +- `output_dir` override + +## Presets + +Read preset definitions from [presets.json](./presets.json). + +The config defines shared output defaults: + +1. `outputDirPattern` +2. `draftFileName` + +Each preset defines: + +1. `displayName`: reader-facing product or component name +2. `versionFile` +3. `pathsFile` (nullable): file of git pathspecs to filter to, one per line. Use `:(top)` to anchor an entry at repo root. + +Example presets: + +1. `ort` (full ONNX Runtime) +2. `webgpu-plugin-ep` (scoped WebGPU Plugin EP) +3. `cuda-plugin-ep` (scoped CUDA Plugin EP) + +### CUDA Plugin EP Scope + +The `cuda-plugin-ep` preset uses the pathspecs in `plugin-ep-cuda/paths.txt` to scope release-note changes. + +## Workflow + +1. Determine release mode. + - Select preset and load configuration from [presets.json](./presets.json). + - Use the preset's `displayName` whenever the release-note content names the product or component. The preset + key is internal and must not appear in published content. + - If preset has `pathsFile`, run in scoped mode. Otherwise run full mode. +2. Resolve version, in this order: + 1. explicit `version` input + 2. value from preset `versionFile` +3. Resolve output directory, referred to as `resolved_output_dir` after this step. + The output directory contains contributor artifacts and the release notes draft. + Resolve it in this order: + 1. explicit `output_dir` input + 2. shared `outputDirPattern` rendered with preset name and resolved version +4. Gather metadata. + - If the output directory is missing or lacks contributor artifacts, generate them with + `tools/python/compile_contributors.py`. + - Generation can take a while because it scans commit history and fetches PR metadata. + - Use `--paths-file` only when preset has a `pathsFile`. + - If existing contributor artifacts are reused, verify `resolved_output_dir/logs.txt` matches base/target before trusting them. +5. Read `resolved_output_dir/detail.csv` as the primary source for PR numbers, titles, authors, target commits, and cherry-pick mapping. + - Use `resolved_output_dir/logs.txt` for contributor summary context and base/target verification. + - Use `git log` only as a fallback sanity check when artifacts are present but incomplete or suspect. + - Check PRs with unexpectedly large author lists for rebased history that imported unrelated commits. Replace those + authors with the actual PR author or authors before building contributor acknowledgments; do not credit authors + solely because they authored an unrelated imported commit. For example, PR #28299, the rebased history contains unrelated commits and co-author metadata. +6. Build highlight categories. + - Full ORT example categories: performance, model/operator support, execution providers, API/languages, + reliability/security, build/packaging/tooling, docs/dev workflow. + - Scoped mode: narrow categories to the component domain. +7. Draft markdown. + - Write the release-note draft to `resolved_output_dir/`. + - Contents: + - Intro sentence + - `## Highlights` + - Inline PR links on every highlight bullet + - `## Contributors` + - Optional scope note for scoped-component releases + - Use the preset's reader-facing `displayName`, not the internal preset key. + - Describe the scope in reader-facing terms, such as "commits affecting WebGPU Plugin EP code and + packaging." + - AI disclaimer if AI drafted + - Do not mention presets, `pathsFile`, configuration files, or other release-note-generation implementation + details in the release-note content. + - Do not refer to the release notes as a "draft" in their content. "Draft" is only an internal workflow + and file-naming concept. +8. Build contributors section. + - Start from `detail.csv` output + - Include humans only + - Exclude bots/agents (for example: `github-actions[bot]`, `app/copilot-swe-agent`, `claude`) + - Sort alphabetically +9. Validate draft quality. + - Every highlight bullet has at least one PR link + - PRs are traceable to metadata or git history + - Contributor list is human-only and alphabetical + - Scope is correct for full vs component release + +## PowerShell Command Patterns + +### compile_contributors.py + +Full ORT metadata: + +```powershell +python .\tools\python\compile_contributors.py \ + --base \ + --target \ + --dir +``` + +Scoped metadata: + +```powershell +python .\tools\python\compile_contributors.py \ + --base \ + --target \ + --dir \ + --paths-file +``` + +## Style and Policy + +Default policy unless release owners override: + +1. Treat the range as changes since the previous release. +2. Keep PR links inline with highlight claims. +3. Keep contributor acknowledgments human-only and best effort. +4. Include an AI disclaimer when highlights are AI drafted. +5. Use GitHub Releases pages as preferred style references. + E.g., [ORT 1.28 release page](https://github.com/microsoft/onnxruntime/releases/tag/v1.28.0). +6. Prefer preset-driven configuration over ad-hoc path/version arguments. +7. Use a single `output_dir` for contributor artifacts, logs, and the release-note draft. diff --git a/.agents/skills/ort-release-notes/presets.json b/.agents/skills/ort-release-notes/presets.json new file mode 100644 index 0000000000000..fb2d1d0201512 --- /dev/null +++ b/.agents/skills/ort-release-notes/presets.json @@ -0,0 +1,22 @@ +{ + "defaultPreset": "ort", + "outputDirPattern": "release-notes-output.{preset}.{version}", + "draftFileName": "release-notes-draft.md", + "presets": { + "ort": { + "displayName": "ONNX Runtime", + "versionFile": "VERSION_NUMBER", + "pathsFile": null + }, + "webgpu-plugin-ep": { + "displayName": "ONNX Runtime WebGPU Plugin EP", + "versionFile": "plugin-ep-webgpu/VERSION_NUMBER", + "pathsFile": "plugin-ep-webgpu/paths.txt" + }, + "cuda-plugin-ep": { + "displayName": "ONNX Runtime CUDA Plugin EP", + "versionFile": "plugin-ep-cuda/VERSION_NUMBER", + "pathsFile": "plugin-ep-cuda/paths.txt" + } + } +} diff --git a/.github/workflows/windows_cuda_no_cudnn.yml b/.github/workflows/windows_cuda_no_cudnn.yml index ac4f9c8d09f80..d2866713f1d29 100644 --- a/.github/workflows/windows_cuda_no_cudnn.yml +++ b/.github/workflows/windows_cuda_no_cudnn.yml @@ -121,6 +121,7 @@ jobs: --use_vcpkg ` --use_vcpkg_ms_internal_asset_cache ` --enable_cuda_profiling ` + --cmake_extra_defines onnxruntime_USE_TRT_FUSED_ATTENTION=OFF ` --cmake_extra_defines onnxruntime_QUICK_BUILD=ON ` --cmake_extra_defines CMAKE_CUDA_ARCHITECTURES=86 ` --cmake_extra_defines onnxruntime_BUILD_CUDA_EP_AS_PLUGIN=ON diff --git a/cmake/CMakeLists.txt b/cmake/CMakeLists.txt index b9b63c0595c45..a3a15416dd675 100644 --- a/cmake/CMakeLists.txt +++ b/cmake/CMakeLists.txt @@ -123,12 +123,18 @@ option(onnxruntime_USE_VSINPU "Build with VSINPU support" OFF) cmake_dependent_option(onnxruntime_USE_FLASH_ATTENTION "Build flash attention kernel for scaled dot product attention" ON "onnxruntime_USE_CUDA" OFF) option(onnxruntime_USE_LEAN_ATTENTION "Build lean attention kernel for scaled dot product attention" OFF) cmake_dependent_option(onnxruntime_USE_MEMORY_EFFICIENT_ATTENTION "Build memory efficient attention kernel for scaled dot product attention" ON "onnxruntime_USE_CUDA" OFF) +# The TensorRT fused MHA kernels are prebuilt cubins for sm70-sm89 that mainly benefit BERT-style +# encoder models. Turning this OFF drops ~14MB of embedded cubin data; attention falls back to +# flash / memory efficient / cuDNN / unfused kernels. +cmake_dependent_option(onnxruntime_USE_TRT_FUSED_ATTENTION "Build TensorRT fused multi-head attention cubin kernels" ON "onnxruntime_USE_CUDA" OFF) option(onnxruntime_USE_FP4_QMOE "Build CUDA QMoE FP4 kernels" OFF) option(onnxruntime_USE_FP8_QMOE "Build CUDA QMoE FP8 kernels" OFF) cmake_dependent_option(onnxruntime_USE_FPA_INTB_GEMM "Build FpA IntB gemm cuda kernels" ON "onnxruntime_USE_CUDA" OFF) option(onnxruntime_USE_INT4_KV_CACHE "Build cuda kernels for int4 kv cache" OFF) option(onnxruntime_USE_FP8_KV_CACHE "Build cuda kernels for fp8 kv cache" ON) option(onnxruntime_QUICK_BUILD "Speed up build by skipping some kernels for faster development" OFF) +# Raises the minimum driver to the CUDA 12.4 level (Linux >= 550.54.14, Windows >= 551.61); always on for CUDA >= 13.0. +cmake_dependent_option(onnxruntime_CUDA_FATBIN_COMPRESS_SIZE "Compress CUDA fatbins with -compress-mode=size" OFF "onnxruntime_USE_CUDA" OFF) option(onnxruntime_BUILD_FOR_NATIVE_MACHINE "Enable this option for turning on optimization specific to this machine" OFF) option(onnxruntime_USE_AVX "Use AVX instructions" OFF) @@ -788,6 +794,7 @@ else() set(onnxruntime_USE_FLASH_ATTENTION OFF) set(onnxruntime_USE_LEAN_ATTENTION OFF) set(onnxruntime_USE_MEMORY_EFFICIENT_ATTENTION OFF) + set(onnxruntime_USE_TRT_FUSED_ATTENTION OFF) endif() if (onnxruntime_USE_CUDA) @@ -809,6 +816,11 @@ if (onnxruntime_USE_CUDA) list(APPEND ORT_PROVIDER_FLAGS -DUSE_MEMORY_EFFICIENT_ATTENTION=1) endif() + if (onnxruntime_USE_TRT_FUSED_ATTENTION) + message( STATUS "Enable TensorRT fused multi-head attention for CUDA EP") + list(APPEND ORT_PROVIDER_FLAGS -DUSE_TRT_FUSED_ATTENTION=1) + endif() + if (onnxruntime_USE_FPA_INTB_GEMM) message( STATUS "Enable FpA IntB Gemm for CUDA EP") list(APPEND ORT_PROVIDER_FLAGS -DUSE_FPA_INTB_GEMM=1) @@ -1526,7 +1538,12 @@ if (onnxruntime_USE_CUDA) message(FATAL_ERROR "onnxruntime_USE_FP4_QMOE requires CUDA Toolkit version 12.8 or newer") endif() - if(CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL "13.0") + if(onnxruntime_CUDA_FATBIN_COMPRESS_SIZE AND CMAKE_CUDA_COMPILER_VERSION VERSION_LESS "12.8") + message(FATAL_ERROR "onnxruntime_CUDA_FATBIN_COMPRESS_SIZE requires CUDA Toolkit version 12.8 or newer") + endif() + + if(CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL "13.0" OR onnxruntime_CUDA_FATBIN_COMPRESS_SIZE) + message(STATUS "Compressing CUDA fatbins with -compress-mode=size (requires a CUDA 12.4 or newer driver)") set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} -Xfatbin=-compress-all -compress-mode=size") else() set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} -Xfatbin=-compress-all") diff --git a/cmake/external/cuda_configuration.cmake b/cmake/external/cuda_configuration.cmake index 9d7bfb0e5ad61..ff8b1a00aaea1 100644 --- a/cmake/external/cuda_configuration.cmake +++ b/cmake/external/cuda_configuration.cmake @@ -164,6 +164,7 @@ macro(setup_cuda_architectures) unset(ORT_HAS_SM80_OR_LATER) unset(ORT_HAS_SM90_OR_LATER) unset(ORT_HAS_SM100_OR_LATER) + unset(ORT_HAS_SM120_OR_LATER) foreach(CUDA_ARCH IN LISTS CMAKE_CUDA_ARCHITECTURES_ORIG) if(CUDA_ARCH MATCHES "^([0-9]+)") if(CMAKE_MATCH_1 GREATER_EQUAL 80) @@ -175,6 +176,9 @@ macro(setup_cuda_architectures) if(CMAKE_MATCH_1 GREATER_EQUAL 100) set(ORT_HAS_SM100_OR_LATER ON) endif() + if(CMAKE_MATCH_1 GREATER_EQUAL 120) + set(ORT_HAS_SM120_OR_LATER ON) + endif() endif() endforeach() if(ORT_HAS_SM80_OR_LATER) @@ -186,6 +190,9 @@ macro(setup_cuda_architectures) if(ORT_HAS_SM100_OR_LATER) add_definitions("-DHAS_SM100_OR_LATER") endif() + if(ORT_HAS_SM120_OR_LATER) + add_definitions("-DHAS_SM120_OR_LATER") + endif() set(ARCHITECTURES_WITH_KERNELS "80" "86" "89" "90" "100" "110" "120") foreach(CUDA_ARCH IN LISTS ARCHITECTURES_WITH_KERNELS) @@ -203,7 +210,7 @@ macro(setup_cuda_architectures) endforeach() # Enable accelerated features (like WGMMA, TMA and setmaxnreg) for SM >= 90. - set(ARCHITECTURES_WITH_ACCEL "90" "100" "101" "110" "120") + set(ARCHITECTURES_WITH_ACCEL "90" "100" "103" "110" "120" "121") unset(CMAKE_CUDA_ARCHITECTURES_NORMALIZED) foreach(CUDA_ARCH IN LISTS CMAKE_CUDA_ARCHITECTURES) if(CUDA_ARCH MATCHES "^([0-9]+)f$") diff --git a/cmake/onnxruntime_providers_cuda.cmake b/cmake/onnxruntime_providers_cuda.cmake index 367d65c1ddaa0..535b79bae8839 100644 --- a/cmake/onnxruntime_providers_cuda.cmake +++ b/cmake/onnxruntime_providers_cuda.cmake @@ -67,6 +67,14 @@ include(onnxruntime_cuda_source_filters.cmake) onnxruntime_filter_cuda_cu_sources(onnxruntime_cuda_contrib_ops_cu_srcs) + + if (NOT onnxruntime_USE_TRT_FUSED_ATTENTION) + # Drop the prebuilt TensorRT fused MHA cubin blobs. cudaDriverWrapper is kept because + # sparse attention depends on it. + list(FILTER onnxruntime_cuda_contrib_ops_cc_srcs EXCLUDE REGEX + ".*/bert/tensorrt_fused_multihead_attention/.*(\\.cubin\\.cc|_kernel\\.sm[0-9]+\\.cc)$") + endif() + onnxruntime_extract_sm_specific_cuda_sources(onnxruntime_cuda_contrib_ops_cu_srcs SM90_SOURCES onnxruntime_cuda_sm90_tma_srcs SM120_SOURCES onnxruntime_cuda_sm120_tma_srcs diff --git a/cmake/onnxruntime_providers_cuda_plugin.cmake b/cmake/onnxruntime_providers_cuda_plugin.cmake index ffb783f1a7cc5..dbe2284c9f367 100644 --- a/cmake/onnxruntime_providers_cuda_plugin.cmake +++ b/cmake/onnxruntime_providers_cuda_plugin.cmake @@ -44,6 +44,13 @@ list(FILTER CUDA_PLUGIN_EP_CU_SRCS EXCLUDE REGEX "onnxruntime/contrib_ops/cuda/c list(FILTER CUDA_PLUGIN_EP_CC_SRCS EXCLUDE REGEX "onnxruntime/contrib_ops/cuda/aten_ops/.*") list(FILTER CUDA_PLUGIN_EP_CC_SRCS EXCLUDE REGEX "onnxruntime/contrib_ops/cuda/collective/.*") +if (NOT onnxruntime_USE_TRT_FUSED_ATTENTION) + # Drop the prebuilt TensorRT fused MHA cubin blobs. cudaDriverWrapper.cc is kept because + # sparse attention depends on it. + list(FILTER CUDA_PLUGIN_EP_CC_SRCS EXCLUDE REGEX + ".*/bert/tensorrt_fused_multihead_attention/.*(\\.cubin\\.cc|_kernel\\.sm[0-9]+\\.cc)$") +endif() + # Exclude files that include cuda_execution_provider.h (directly or transitively), # which conflicts with the adapter shim CUDAExecutionProvider class. list(FILTER CUDA_PLUGIN_EP_CC_SRCS EXCLUDE REGEX ".*/cuda_execution_provider\\.cc$") @@ -123,6 +130,17 @@ onnxruntime_add_shared_library_module(onnxruntime_providers_cuda_plugin ${CUDA_PLUGIN_EP_CU_SRCS} ) +if(WIN32) + # Add version information to the packaged plugin DLL. + target_sources(onnxruntime_providers_cuda_plugin PRIVATE + "${ONNXRUNTIME_ROOT}/core/providers/cuda/onnxruntime_providers_cuda.rc") + target_compile_definitions(onnxruntime_providers_cuda_plugin PRIVATE + FILE_NAME=\"onnxruntime_providers_cuda.dll\") +elseif(UNIX AND NOT APPLE) + # The build output is packaged directly, so do not embed the build machine's CUDA path. + set_target_properties(onnxruntime_providers_cuda_plugin PROPERTIES SKIP_BUILD_RPATH TRUE) +endif() + # Mirror directory structure in the Visual Studio solution tree under "onnxruntime". source_group(TREE ${ONNXRUNTIME_ROOT} PREFIX "onnxruntime" FILES ${CUDA_EP_CC_SRCS} ${CUDA_EP_CU_SRCS}) source_group(TREE ${ONNXRUNTIME_ROOT} PREFIX "onnxruntime" FILES ${CUDA_CONTRIB_OPS_CC_SRCS} ${CUDA_CONTRIB_OPS_CU_SRCS}) @@ -346,6 +364,12 @@ if(NOT onnxruntime_DISABLE_CONTRIB_OPS) if(_cuda_plugin_llm_srcs) if(MSVC AND NOT onnxruntime_USE_FP4_QMOE) onnxruntime_filter_cuda_archs(_plugin_llm_cuda_architectures MIN_SM 75 EXCLUDE_SM120_REAL) + # A native-only Windows ARM64 build has no lower architecture left after the + # MSVC SM120 exclusion. Emit PTX privately for this object library so its host + # launchers and device kernels are still linked into the plugin. + if(NOT _plugin_llm_cuda_architectures AND ORT_HAS_SM120_OR_LATER) + set(_plugin_llm_cuda_architectures "120-virtual") + endif() else() onnxruntime_filter_cuda_archs(_plugin_llm_cuda_architectures MIN_SM 75) endif() diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index 23cdbac30a042..2331780e8e932 100644 --- a/docs/ContribOperators.md +++ b/docs/ContribOperators.md @@ -4234,21 +4234,37 @@ This version of the operator has been available since version 1 of the 'com.micr
do_rotary : int
Whether to use rotary position embedding. Default value is 0.
+
k_cache_dtype : string
+
Logical element type stored in 'key_cache', named after the ONNX element type it denotes: '' (the default) means the cache tensor's own element type is also the logical type. 'float16', 'bfloat16', 'int8' and 'float8e4m3fn' name that same type explicitly and must agree with the tensor. 'int4' and 'float4e2m1' name sub-byte types packed two per byte into a uint8 cache, where the last cache dimension holds (head_size + 1) / 2 bytes and logical element 2*i occupies the low-order bits of byte i. Every value is a signed, zero-symmetric type: quantization uses a scale with no zero point, so unsigned logical types are not expressible.
+
k_quant_type : string
+
Quantization granularity of the key cache: 'NONE', 'PER_TENSOR' or 'PER_CHANNEL'. Must be non-'NONE' exactly when 'key_cache' has a quantized element type, and then 'k_scale' is required. Default value is 'NONE'.
+
kv_cache_layout : string
+
Physical layout of the KV cache: 'SEPARATE' or 'LATENT'. 'SEPARATE' (the default) uses distinct 'key_cache' and 'value_cache' tensors. 'LATENT' selects absorbed Multi-head Latent Attention: there is a single cache, 'value' and 'value_cache' must be absent, 'kv_num_heads' must be 1, and V for every head is the leading 'v_head_size' channels of the same 'key_cache' row that supplies K. Default value is 'SEPARATE'.
kv_num_heads : int (required)
Number of attention heads for k and v
local_window_size : int
left_window_size for local attention (like Mistral). Default value is -1 meaning unused.
num_heads : int (required)
Number of attention heads for q
+
qk_norm_epsilon : float
+
Epsilon used by the Q/K RMSNorm when 'q_norm_weight' and 'k_norm_weight' are provided. Default value is 1e-6.
rotary_interleaved : int
Rotate using interleaved pattern. Default value is 0 (False).
+
rotary_offset : int
+
First channel within head_size covered by rotary embedding, so RoPE is applied to [rotary_offset, rotary_offset + rotary_dim) and channels outside that range are copied through. Must be a multiple of 8. MLA sets this to kv_lora_rank so that RoPE only touches the positional suffix of the latent row. Default value is 0.
scale : float
Custom scale will be used if specified. Default value is 1/sqrt(head_size)
softcap : float
Softcap value for attention weights. Default value is 0.
+
v_cache_dtype : string
+
Logical element type stored in 'value_cache', with the same values and packing rule as 'k_cache_dtype'. Default value is '' (use the cache tensor's element type).
+
v_head_size : int
+
Width of the value head, which may be narrower than head_size. Only valid when 'kv_cache_layout' is 'LATENT' (DeepSeek-V3 uses head_size=576 and v_head_size=512). When v_head_size differs from head_size the 'scale' attribute is required, because the 1/sqrt(head_size) default no longer matches the pre-absorption head width. Default value is 0, meaning the same as head_size.
+
v_quant_type : string
+
Quantization granularity of the value cache: 'NONE', 'PER_TENSOR' or 'PER_CHANNEL'. Must be non-'NONE' exactly when 'value_cache' has a quantized element type, and then 'v_scale' is required. Default value is 'NONE'.
-#### Inputs (8 - 10) +#### Inputs (8 - 17)
query : T
@@ -4256,11 +4272,11 @@ This version of the operator has been available since version 1 of the 'com.micr
key (optional) : T
Key with shape (num_tokens, kv_hidden_size)
value (optional) : T
-
Value with shape (num_tokens, kv_hidden_size)
-
key_cache : T
-
Block-based key cache with shape (num_blocks, block_size, kv_num_heads, head_size). This is updated in place within the op.
-
value_cache : T
-
Block-based value cache with shape (num_blocks, block_size, kv_num_heads, head_size). This is updated in place within the op. This should be the same shape as key_cache.
+
Value with shape (num_tokens, kv_hidden_size). Must be absent when 'kv_cache_layout' is 'LATENT'.
+
key_cache : T_CACHE
+
Block-based key cache with shape (num_blocks, block_size, kv_num_heads, head_size). This is updated in place within the op. When 'kv_cache_layout' is 'LATENT' this is the only cache, and V is read from its leading v_head_size channels.
+
value_cache (optional) : T_CACHE
+
Block-based value cache with shape (num_blocks, block_size, kv_num_heads, head_size). This is updated in place within the op. This should be the same shape as key_cache. Must be absent when 'kv_cache_layout' is 'LATENT'.
cumulative_sequence_length : S
A tensor with shape (batch_size + 1). It specifies the cumulative sequence lengths between the packed entries in Q/K/V.
past_seqlens : S
@@ -4271,17 +4287,31 @@ This version of the operator has been available since version 1 of the 'com.micr
2D tensor with shape (max total seqlen, head_size / 2).
sin_cache (optional) : T
2D tensor with shape (max total seqlen, head_size / 2).
+
slot_mapping (optional) : S
+
1D tensor with shape (num_tokens). For each query token, the flat slot index (block_id * block_size + offset_in_block) at which its key/value is written into the KV cache. A value of -1 skips the cache write for that token, which lets a scheduler suppress stores for prefix-cache hits or rejected speculative tokens. When absent, slots are derived from 'past_seqlens', 'cumulative_sequence_length' and 'block_table' as before. 'block_table' is still required, because it defines the read path.
+
head_sink (optional) : T
+
1D tensor with shape (num_heads). Each head has a learnable sink logit that participates in the softmax denominator but contributes no value, so attention can 'do nothing'.
+
q_norm_weight (optional) : T
+
1D tensor with shape (head_size). RMSNorm gain applied to each query head before rotary embedding. Must be provided together with 'k_norm_weight'.
+
k_norm_weight (optional) : T
+
1D tensor with shape (head_size). RMSNorm gain applied to each key head before rotary embedding and before the key is written to the KV cache. Must be provided together with 'q_norm_weight'.
+
k_scale (optional) : T_KV_SCALE
+
Dequantization scale of the key cache. Shape is (1) when 'k_quant_type' is 'PER_TENSOR' and (kv_num_heads, 1, head_size) when it is 'PER_CHANNEL'. Quantization is symmetric (no zero point).
+
v_scale (optional) : T_KV_SCALE
+
Dequantization scale of the value cache. Shape is (1) when 'v_quant_type' is 'PER_TENSOR' and (kv_num_heads, 1, head_size) when it is 'PER_CHANNEL'. Quantization is symmetric (no zero point).
+
attention_metadata (optional) : S
+
1D tensor with shape (2) holding [max_query_len_bound, max_kv_len_bound] in CPU memory. max_query_len_bound is an upper bound on the number of new tokens any one sequence contributes; max_kv_len_bound is an upper bound on past_seqlens[i] + query_len[i]. Both are replay-wide upper bounds, never exact per-step values: they must hold for every step this node -- or a CUDA Graph capturing it -- will serve, and 0 means 'unknown'. They may only select the backend and size launch dimensions and workspaces; they never enter a mask comparison, so over-estimating only costs empty work. The op can otherwise obtain these only by copying 'cumulative_sequence_length' and 'past_seqlens' back from the device and synchronizing the stream on every call, which stalls the pipeline once per node per step and makes the op impossible to capture into a CUDA Graph. Schedulers already track these bounds on the host, so supplying them is normally free. When absent, the op falls back to the device readback. The values are trusted: an under-sized bound violates the contract and may omit attention work.
#### Outputs (1 - 3)
output : T
-
3D output tensor with shape (num_tokens, hidden_size)
-
key_cache_out (optional) : T
+
2D output tensor with shape (num_tokens, num_heads * v_head_size), which is (num_tokens, hidden_size) unless 'kv_cache_layout' is 'LATENT' with a narrower v_head_size.
+
key_cache_out (optional) : T_CACHE
Block-based key cache with shape (num_blocks, block_size, kv_num_heads, head_size). This is always the same tensor as key_cache.
-
value_cache_out (optional) : T
-
Block-based value cache with shape (num_blocks, block_size, kv_num_heads, head_size). This is always the same tensor as value_cache.
+
value_cache_out (optional) : T_CACHE
+
Block-based value cache with shape (num_blocks, block_size, kv_num_heads, head_size). This is always the same tensor as value_cache. Must be absent when 'kv_cache_layout' is 'LATENT'.
#### Type Constraints @@ -4289,6 +4319,10 @@ This version of the operator has been available since version 1 of the 'com.micr
T : tensor(float16), tensor(bfloat16)
Constrain input and output to float tensors.
+
T_CACHE : tensor(float16), tensor(bfloat16), tensor(int8), tensor(float8e4m3fn)
+
Constrain the KV cache to float or quantized tensors.
+
T_KV_SCALE : tensor(float)
+
Constrain KV cache scales to float tensors.
S : tensor(int32)
Constrain Positional inputs to int tensor.
diff --git a/docs/JSEP_Deprecation.md b/docs/JSEP_Deprecation.md new file mode 100644 index 0000000000000..cc74749e3b8bc --- /dev/null +++ b/docs/JSEP_Deprecation.md @@ -0,0 +1,77 @@ +# JSEP deprecation + +**Status: deprecated.** JSEP — the JavaScript/TypeScript WebGPU execution path in `onnxruntime-web` — is being +replaced by the native WebGPU execution provider. Removal is planned. **The timeline is not yet fixed**; it will be +set once we understand real-world JSEP usage, and announced before anything is removed. + +This page is the authoritative statement of the deprecation and of the contribution policy below. Link to it when +redirecting JSEP work. + +## Contribution policy + +JSEP is in maintenance mode: **bug fixes and security fixes only.** + +| Change | Where it goes | +|---|---| +| Correctness or security fix in an existing JSEP kernel | JSEP — accepted | +| New operator | The native WebGPU EP — `onnxruntime/core/providers/webgpu/` or `onnxruntime/contrib_ops/webgpu/` | +| New feature, or performance work | The native WebGPU EP | + +Note that [`js/web/docs/webgpu-operators.md`](../js/web/docs/webgpu-operators.md) lists **JSEP** operators despite +its name, so it cannot be used to check what the native WebGPU EP already covers. + +If a model works on JSEP but not on the native WebGPU EP, that is a gap worth reporting — please open an issue +describing the model and the failure or add support in the native WebGPU EP. + +Adding to JSEP now means the work is deleted later and has to be written a second time against the native EP. + +## What JSEP is + +JSEP implements WebGPU compute in TypeScript, driven from C++ through an Asyncify-compiled WASM core. It spans multiple +areas of the repository: + +| Path | Contents | +|---|---| +| `js/web/lib/wasm/jsep/` | The TypeScript WebGPU backend and its kernel implementations | +| `onnxruntime/core/providers/js/` | The native "JS EP" — kernel stubs that dispatch back into JavaScript | +| `onnxruntime/contrib_ops/js/` | Contrib operator registrations for the same EP | +| `onnxruntime/wasm/pre-jsep.js` | Emscripten glue | +| `cmake/onnxruntime_providers_js.cmake`, `js/build_jsep.bat` | Build plumbing | + +Built with `--use_jsep` (`USE_JSEP`), and registered under the EP name `JsExecutionProvider`. + +This is **not** the same thing as the native WebGPU EP (`onnxruntime/core/providers/webgpu/`), which is a +conventional C++ execution provider (compiled to WASM for onnxruntime-web). Both register under the `webgpu` +backend key in JavaScript. Which one runs is a *build-time* choice, not a runtime one. + +## The replacement: native WebGPU EP + +`onnxruntime-web/webgpu` and `onnxruntime-web/jspi` are built against the native WebGPU EP **today**. No code +changes and no pending work are required to use them: + +```js +import * as ort from 'onnxruntime-web/webgpu'; +``` + +The default `onnxruntime-web` import still selects JSEP; the plan is to change that to the native WebGPU EP. A few +JSEP-only `env.webgpu` settings also have no native equivalent. Both are covered in the +[migration design doc](design/onnxruntime_web_jsep_to_webgpu_ep_migration.md). + +## Telling the two apart + +Useful when triaging a bug report, since both are called "WebGPU": + +| Import | WebGPU implementation | WASM artifact | +|---|---|---| +| `onnxruntime-web` (default) | JSEP | `ort-wasm-simd-threaded.jsep.wasm` | +| `onnxruntime-web/all` | JSEP | `ort-wasm-simd-threaded.jsep.wasm` | +| `onnxruntime-web/webgpu` | native WebGPU EP | `ort-wasm-simd-threaded.asyncify.wasm` | +| `onnxruntime-web/jspi` | native WebGPU EP | `ort-wasm-simd-threaded.jspi.wasm` | + +The `.jsep` infix in the WASM filename is the reliable signal. A report that does not identify the import or the +artifact is ambiguous and should be clarified before triage. + +## Related documents + +- [Migrate onnxruntime-web from JSEP to the native WebGPU EP](design/onnxruntime_web_jsep_to_webgpu_ep_migration.md) + — the migration design and phasing. diff --git a/docs/OperatorKernels.md b/docs/OperatorKernels.md index e2b01333249db..d359aa0e2b523 100644 --- a/docs/OperatorKernels.md +++ b/docs/OperatorKernels.md @@ -1109,7 +1109,7 @@ The **OpSet Version** column uses the following notation: |NhwcConv|*in* X:**T**
*in* W:**T**
*in* B:**T**
*out* Y:**T**|1+|**T** = tensor(float), tensor(float16)| |PackedAttention|*in* input:**T**
*in* weights:**T**
*in* bias:**T**
*in* token_offset:**M**
*in* cumulative_sequence_length:**M**
*in* attention_bias:**T**
*out* output:**T**|1+|**T** = tensor(float), tensor(float16)| |PackedMultiHeadAttention|*in* query:**T**
*in* key:**T**
*in* value:**T**
*in* bias:**T**
*in* token_offset:**M**
*in* cumulative_sequence_length:**M**
*in* attention_bias:**T**
*out* output:**T**|1+|**T** = tensor(float), tensor(float16)| -|PagedAttention|*in* query:**T**
*in* key:**T**
*in* value:**T**
*in* key_cache:**T**
*in* value_cache:**T**
*in* cumulative_sequence_length:**S**
*in* past_seqlens:**S**
*in* block_table:**S**
*in* cos_cache:**T**
*in* sin_cache:**T**
*out* output:**T**
*out* key_cache_out:**T**
*out* value_cache_out:**T**|1+|**S** = tensor(int32)
**T** = tensor(bfloat16), tensor(float16)| +|PagedAttention|*in* query:**T**
*in* key:**T**
*in* value:**T**
*in* key_cache:**T_CACHE**
*in* value_cache:**T_CACHE**
*in* cumulative_sequence_length:**S**
*in* past_seqlens:**S**
*in* block_table:**S**
*in* cos_cache:**T**
*in* sin_cache:**T**
*in* slot_mapping:**S**
*in* head_sink:**T**
*in* q_norm_weight:**T**
*in* k_norm_weight:**T**
*in* k_scale:**T_KV_SCALE**
*in* v_scale:**T_KV_SCALE**
*in* attention_metadata:**S**
*out* output:**T**
*out* key_cache_out:**T_CACHE**
*out* value_cache_out:**T_CACHE**|1+|**S** = tensor(int32)
**T** = tensor(bfloat16), tensor(float16)
**T_CACHE** = tensor(bfloat16), tensor(float16), tensor(float8e4m3fn), tensor(int8)
**T_KV_SCALE** = tensor(float)| |QAttention|*in* input:**T1**
*in* weight:**T2**
*in* bias:**T3**
*in* input_scale:**T3**
*in* weight_scale:**T3**
*in* mask_index:**T4**
*in* input_zero_point:**T1**
*in* weight_zero_point:**T2**
*in* past:**T3**
*out* output:**T3**
*out* present:**T3**|1+|**T1** = tensor(int8)
**T2** = tensor(int8)
**T3** = tensor(float), tensor(float16)
**T4** = tensor(int32)| |QMoE|*in* input:**T**
*in* router_probs:**T**
*in* fc1_experts_weights:**T1**
*in* fc1_scales:**T2**
*in* fc1_experts_bias:**T**
*in* fc2_experts_weights:**T1**
*in* fc2_scales:**T2**
*in* fc2_experts_bias:**T**
*in* fc3_experts_weights:**T1**
*in* fc3_scales:**T2**
*in* fc3_experts_bias:**T**
*in* fc1_zero_points:**T1**
*in* fc2_zero_points:**T1**
*in* fc3_zero_points:**T1**
*in* router_weights:**T**
*in* fc1_global_scale:**T4**
*in* fc2_global_scale:**T4**
*in* fc1_act_scale:**T4**
*in* fc2_act_scale:**T4**
*in* fc1_act_block_scale:**T2**
*in* fc2_act_block_scale:**T2**
*out* output:**T**|1+|**T** = tensor(bfloat16), tensor(float16)
**T1** = tensor(float8e4m3fn), tensor(uint8)
**T2** = tensor(bfloat16), tensor(float16), tensor(float8e4m3fn), tensor(float8e8m0)
**T4** = tensor(float)| |QOrderedAttention|*in* input:**Q**
*in* scale_input:**S**
*in* scale_Q_gemm:**S**
*in* scale_K_gemm:**S**
*in* scale_V_gemm:**S**
*in* Q_weight:**Q**
*in* K_weight:**Q**
*in* V_weight:**Q**
*in* scale_Q_weight:**S**
*in* scale_K_weight:**S**
*in* scale_V_weight:**S**
*in* Q_bias:**S**
*in* K_bias:**S**
*in* V_bias:**S**
*in* scale_QKT_gemm:**S**
*in* scale_QKT_softmax:**S**
*in* scale_values_gemm:**S**
*in* mask_index:**G**
*in* past:**Q**
*in* attention_bias:**S**
*out* output:**Q**|1+|**G** = tensor(int32)
**Q** = tensor(int8)
**S** = tensor(float)| diff --git a/docs/ReleaseNotesWorkflow.md b/docs/ReleaseNotesWorkflow.md new file mode 100644 index 0000000000000..672201929657d --- /dev/null +++ b/docs/ReleaseNotesWorkflow.md @@ -0,0 +1,35 @@ +# ONNX Runtime Release Notes Workflow + +ONNX Runtime release notes are drafted using the `ort-release-notes` agent skill: + +- [`.agents/skills/ort-release-notes/SKILL.md`](../.agents/skills/ort-release-notes/SKILL.md) + +The skill is the source of truth for the release-note workflow, including inputs, artifact discovery, contributor +handling, draft structure, validation, and PowerShell command patterns. Keep procedural details there so this document +does not drift out of sync. + +## Supported Presets + +Preset definitions live in: + +- [`.agents/skills/ort-release-notes/presets.json`](../.agents/skills/ort-release-notes/presets.json) + +For example, there is a preset for the core ONNX Runtime release and one for the WebGPU plugin EP release. + +Refer to that file for the available preset names and their version/path configuration. + +## Maintainer Usage + +Ask your favorite AI agent to draft release notes using a preset name, base ref, and target ref. For example: + +```text +Draft release notes using the ort-release-notes skill for preset webgpu-plugin-ep from +plugin-ep-webgpu/v0.1.0 to . +``` + +## Adding A Preset + +To add release-note support for another component, add a preset to +[`presets.json`](../.agents/skills/ort-release-notes/presets.json). Component presets should define a version file and, +when the release should be scoped, a path filter file similar to +[`plugin-ep-webgpu/paths.txt`](../plugin-ep-webgpu/paths.txt). diff --git a/docs/contrib_ops/cuda/paged_attention.md b/docs/contrib_ops/cuda/paged_attention.md new file mode 100644 index 0000000000000..2d5baa3d1813e --- /dev/null +++ b/docs/contrib_ops/cuda/paged_attention.md @@ -0,0 +1,1828 @@ +# PagedAttention — CUDA Design Document + +Status: **Draft / proposal** +Scope: `com.microsoft::PagedAttention`, CUDA Execution Provider +Related: [gqa.md](gqa.md) (`com.microsoft::GroupQueryAttention`) + +> **Compatibility decision.** `PagedAttention` is already serialized as `com.microsoft::PagedAttention` +> opset 1. [§4](#4-schema-the-compatible-contract) is the single normative contract for that opset and +> evolves it **additively only**. [§21](#21-deferred-breaking-contract) records the breaking +> cache-layout ideas that were considered and rejected for opset 1; they are deferred to a separately +> versioned schema and must not replace the shipped contract in place. + +## Table of Contents + +- [1. Purpose and Scope](#1-purpose-and-scope) +- [2. Current State](#2-current-state) +- [3. Design Principles](#3-design-principles) +- [4. Schema — The Compatible Contract](#4-schema--the-compatible-contract) +- [5. Feature: `slot_mapping`](#5-feature-slot_mapping) +- [6. Feature: Attention Sink (`head_sink`)](#6-feature-attention-sink-head_sink) +- [7. Feature: Fused QK-Norm](#7-feature-fused-qk-norm) +- [8. Feature: Quantized Paged KV Cache](#8-feature-quantized-paged-kv-cache) +- [9. Feature: Sliding Window Attention](#9-feature-sliding-window-attention) +- [10. Feature: `attention_bias`](#10-feature-attention_bias) +- [11. Feature: `output_qk`](#11-feature-output_qk) +- [12. Feature: Multi-head Latent Attention (MLA)](#12-feature-multi-head-latent-attention-mla) +- [13. Kernel Dispatch and Backend Plan](#13-kernel-dispatch-and-backend-plan) +- [14. Shared Code with GroupQueryAttention](#14-shared-code-with-groupqueryattention) +- [15. Validation Rules](#15-validation-rules) +- [16. Shape Inference and Tooling](#16-shape-inference-and-tooling) +- [17. Testing Plan](#17-testing-plan) +- [18. Known Defects to Fix First](#18-known-defects-to-fix-first) +- [19. Phasing](#19-phasing) +- [20. Open Questions](#20-open-questions) +- [21. Deferred Breaking Contract](#21-deferred-breaking-contract) + +--- + +## 1. Purpose and Scope + +`PagedAttention` is the **server / continuous-batching** attention operator for ONNX Runtime. +`GroupQueryAttention` (GQA) remains the **padded-batch** operator used by edge and single-stream +scenarios. The two ops are deliberately separate: + +| | `GroupQueryAttention` | `PagedAttention` | +|---|---|---| +| Query layout | `(batch_size, sequence_length, hidden_size)` — padded | `(token_count, hidden_size)` — packed varlen | +| KV cache | `(batch, kv_num_heads, max_seq, head_size)` BNSH, one contiguous buffer per sequence | `(num_blocks, block_size, kv_num_heads, head_size)` + `block_table` | +| Batching | Static, padded | Continuous / in-flight, ragged | +| Target | Edge, on-device, single stream | Multi-tenant serving, high throughput | + +Rank and semantics of `query` differ between the two, and the KV aliasing contract differs +(GQA shares one past/present buffer per sequence; PagedAttention mutates a global block pool in +place). Overloading GQA with a `block_table` input is rejected as a design: ORT kernel matching is +by op type and type constraints, not by optional-input presence, so every EP that registers GQA +would claim a paged node and fail at run time instead of at partition time. + +**This document specifies the feature work required to bring `PagedAttention` to parity with GQA +for popular LLMs** (GPT-OSS, Qwen3, Gemma 2/3, DeepSeek, Llama, Mistral, Phi), plus the paging +primitives (`slot_mapping`) that serving frameworks require, plus **Multi-head Latent Attention** +(§12) — which is not a GQA feature at all, but is native to this operator because MLA's entire +value proposition is a smaller paged KV cache. + +## 2. Current State + +Implemented today in `onnxruntime/contrib_ops/cuda/bert/paged_attention{.cc,.h,_impl.cu,_impl.h,_helper.h}`: + +| Capability | State | +|---|---| +| Packed varlen Q/K/V and packed QKV | Done | +| Block KV cache, in-place update (`ReshapeAndCache`) | Done | +| `block_table` gather for reads | Done | +| RoPE (`do_rotary`, `rotary_interleaved`, `cos_cache`/`sin_cache`) | Done | +| `scale`, `softcap` | Done | +| `local_window_size` | Wired to Flash varlen (`local_window_size - 1`) and MEA; not validated end-to-end | +| Backends | FlashAttention varlen → Memory-Efficient Attention (CUTLASS fMHA) fallback | +| dtypes | `float16`, `bfloat16` only | +| EPs | CUDA only | + +Not implemented: `slot_mapping` (a `slot_mappings` field exists in `PagedAttentionData` but is +never populated or consumed), `head_sink` / smooth softmax, QK-Norm, quantized cache, +`attention_bias`, `output_qk`, a dedicated paged decode kernel. + +Structural limits in the current implementation: + +- `batch_size <= 256` — `LaunchGetCumulativeSeqlensKV` uses a per-block `cub::BlockScan` with 256 + threads and independent blocks. +- `block_size % 256 == 0` — see [§18](#18-known-defects-to-fix-first); this is almost certainly a bug. + *(Partly true. It is a genuine FlashAttention tiling constraint, but a head-size-dependent one. See + the implementation note in §18.1: validation now accepts any power-of-two `block_size >= 16` and + the constraint has been moved into Flash backend eligibility, with fallback to the + memory-efficient backend.)* +- MEA fallback materializes a **tight gathered, GQA-expanded** KV buffer of + `[total_kv_tokens, num_heads, head_size]`; memory scales with context × query heads. +- A device→host sync per step to obtain `max_query_len` (and `total_kv_tokens` for MEA). + *(Fixed. Dispatch and workspace sizing now use static shapes and upper bounds; see §4.7. The sync + remains only as a prefill-side fallback that no capturable step reaches.)* + +## 3. Design Principles + +1. **Additive schema only.** `PagedAttention` is already shipped as contrib opset 1. All new inputs + are appended at the end as `OpSchema::Optional`, all new attributes have defaults matching current + behavior, and existing type constraints are only *widened*. No existing model breaks. This is + binding: [§4.2](#42-compatibility-invariant) states the invariant, and the breaking alternatives + that cannot satisfy it are deferred to [§21](#21-deferred-breaking-contract). +2. **Semantic parity with GQA, not schema parity.** A feature that exists in both ops must have + identical math, identical attribute names, identical defaults, and identical scale-tensor + layouts. Inputs that only make sense for padded batching (`past_key`, `past_value`, + `total_sequence_length`, padded `position_ids`, `present_*`) are **never** mirrored. +3. **Share the math, not the schema.** New features are implemented once in a shared kernel layer + parameterized by a KV *accessor* (contiguous BNSH vs. block table), so GQA and PagedAttention + cannot drift. See [§14](#14-shared-code-with-groupqueryattention). +4. **Fail at partition time, not run time.** No stub kernels that return `NOT_IMPLEMENTED`. An EP + either registers `PagedAttention` or does not. +5. **Prefer exact epilogues over new fused kernels** where an existing kernel already returns the + quantities needed (see the `head_sink` LSE rescale in [§6](#6-feature-attention-sink-head_sink)). + +## 4. Schema — The Compatible Contract + +This section is the **single normative index table** for `com.microsoft::PagedAttention` opset 1. +Every other section in this document refers to these indices. + +### 4.1 Decision + +Evolve the operator additively. Preserve inputs 0–9, outputs 0–2, and every existing attribute with +its current meaning. Add model coverage and serving features through trailing optional inputs, +optional attributes whose defaults reproduce current behavior, a widened cache type constraint, and +`value_cache` becoming optional only for an explicitly selected latent-cache mode. + +Do **not** merge `key_cache` and `value_cache`, replace the existing cache outputs, remove +`kv_num_heads`, rename `local_window_size`, or reinterpret an existing input combination. Those +options are recorded and deferred in [§21](#21-deferred-breaking-contract). + +### 4.2 Compatibility Invariant + +Every model valid under the shipped opset-1 schema remains valid and behaves identically when all +new inputs and attributes are absent. For such a model: + +```text +T_CACHE == T +value_cache is present +key and value are either both present or both absent +kv_cache_layout == "SEPARATE" +v_head_size == 0 +all quantization attributes select NONE +all trailing optional inputs are absent +``` + +The baseline this evolves from is the schema on ONNX Runtime `main` +(`onnxruntime/core/graph/contrib_ops/bert_defs.cc`): inputs 0–9 with `T`-typed caches, outputs 0–2 +under a both-or-neither rule, and the seven attributes `num_heads`, `kv_num_heads`, `scale`, +`softcap`, `local_window_size`, `do_rotary`, `rotary_interleaved`. + +### 4.3 Inputs + +Indices 0–9 are unchanged from the shipped schema. Indices 10–18 are new and optional. The order +matches the landing order in [§19](#19-phasing), so the schema grows monotonically per phase. + +| Idx | Name | Type | Shape | Status | +|-----|------|------|-------|--------| +| 0 | `query` | `T` | `(token_count, hidden_size)` or packed `(token_count, (num_heads + 2*kv_num_heads)*head_size)` | existing | +| 1 | `key` | `T` (opt) | `(token_count, kv_num_heads * head_size)` | existing — required in `LATENT` (§12) | +| 2 | `value` | `T` (opt) | `(token_count, kv_num_heads * head_size)` | existing — absent in `LATENT` (§12) | +| 3 | `key_cache` | **`T_CACHE`** | `(num_blocks, block_size, kv_num_heads, head_size)` | **type widened — §8** | +| 4 | `value_cache` | **`T_CACHE`** (opt) | `(num_blocks, block_size, kv_num_heads, head_size)` | **type widened; absent in `LATENT` — §12** | +| 5 | `cumulative_sequence_length` | `S` | `(batch_size + 1,)` | existing | +| 6 | `past_seqlens` | `S` | `(batch_size,)` | existing | +| 7 | `block_table` | `S` | `(batch_size, max_num_blocks_per_seq)` | existing — `-1` = unmapped (§9.2) | +| 8 | `cos_cache` | `T` (opt) | `(max_seq_len, rotary_dim/2)` | existing | +| 9 | `sin_cache` | `T` (opt) | `(max_seq_len, rotary_dim/2)` | existing | +| 10 | `slot_mapping` | `S` (opt) | `(token_count,)` | **new — §5** | +| 11 | `head_sink` | `T` (opt) | `(num_heads,)` | **new — §6** | +| 12 | `q_norm_weight` | `T` (opt) | `(head_size,)` | **new — §7** | +| 13 | `k_norm_weight` | `T` (opt) | `(head_size,)` | **new — §7** | +| 14 | `k_scale` | `T_KV_SCALE` (opt) | `(1,)` or `(kv_num_heads, 1, head_size)` | **new — §8** | +| 15 | `v_scale` | `T_KV_SCALE` (opt) | `(1,)` or `(kv_num_heads, 1, head_size)` | **new — §8**; absent in `LATENT` | +| 16 | `attention_metadata` | `S` (opt, **CPU**) | `(2,)` | **new — trusted bounds only, §4.7** | +| 17 | `query_positions` | `S` (opt) | `(token_count,)` | **new — §4.8** | +| 18 | `attention_bias` | `T` (opt) | `(batch_size or 1, num_heads or 1, query_length_capacity, context_length_capacity)` | **new — §10** | + +`max_context_len` is the largest per-sequence total KV length in the batch, bounded above by +`block_table.shape[1] * block_size`. + +`attention_bias` is deliberately last: it is a correctness fallback with potentially large memory +cost that disqualifies every fused backend. The rank and broadcast dimensions match GQA; only the +two sequence extents differ because PagedAttention is packed and ragged (§10). + +**Why this table says `head_size` and not `v_head_size`.** `v_head_size` may differ from `head_size` +**only** in `"LATENT"` mode (§12.3), and a `"LATENT"` node has no `value`, no `value_cache`, and no +`v_scale` — V is a *view* of the leading `v_head_size` channels of `key_cache`, so there is nothing +left to give a separate width to. Every V-carrying tensor above therefore exists only in +`"SEPARATE"` mode, where `effective_v_head_size == head_size` by definition. The document spells +`effective_v_head_size` only in the two places where the width genuinely varies: output 0 (§4.4) and +the V view of `key_cache` (§12.3). + +### 4.4 Outputs + +| Idx | Name | Type | Shape | Status | +|-----|------|------|-------|--------| +| 0 | `output` | `T` | `(token_count, num_heads * effective_v_head_size)` | existing (shape generalized by §12) | +| 1 | `key_cache_out` | `T_CACHE` (opt) | aliases `key_cache` | existing | +| 2 | `value_cache_out` | `T_CACHE` (opt) | aliases `value_cache` | existing | +| 3 | `output_qk` | `QK` (opt) | `(num_heads, token_count, max_context_len)` | **new — §11** | + +In `SEPARATE` mode outputs 1 and 2 are both present or both absent, preserving the shipped rule. In +`LATENT` mode output 1 may be present and output 2 must be absent. + +Requesting `output_qk` requires a four-entry ONNX output list. Unused optional output positions are +represented by an **empty output name**; indices are never compacted. A `LATENT` node requesting +`output_qk` therefore writes `[output, key_cache_out-or-empty, "", output_qk]`. + +`output_qk` uses its own `QK` type constraint rather than `T`, matching GQA, so the QK matrix can be +emitted in `float32` independently of the activation type. + +**Aliasing.** The shipped implementation requires the cache outputs to be the *same buffer* as the +corresponding inputs, but enforces it only as a runtime pointer-equality check in +`paged_attention.cc` that returns `INVALID_ARGUMENT`. There is no `.Alias()` in the kernel definition, +so a graph in which the allocation planner does not happen to reuse the input buffer fails at run +time rather than at partition time — a violation of design principle 4 (§3). Registering +`.Alias(3, 1).Alias(4, 2)` on the kernel def is fully compatible and belongs in P0 (§19). A +functional-output contract that stays correct when the planner cannot reuse the buffer is a +different, breaking contract and is deferred (§21). + +### 4.5 Attributes + +Names, defaults, and value sets follow GQA so a graph rewriter can move an attribute between the two +ops without translation. + +| Name | Type | Default | Status | +|---|---|---|---| +| `num_heads` | INT | required | existing | +| `kv_num_heads` | INT | required | existing | +| `scale` | FLOAT | `1/sqrt(head_size)` | existing — mandatory in `LATENT` (§12.6) | +| `softcap` | FLOAT | `0.0` | existing | +| `local_window_size` | INT | `-1` | existing — §9 | +| `do_rotary` | INT | `0` | existing | +| `rotary_interleaved` | INT | `0` | existing | +| `qk_norm_epsilon` | FLOAT | `1e-6` | **new — §7** | +| `k_quant_type` | STRING | `"NONE"` | **new — §8**; `NONE` \| `PER_TENSOR` \| `PER_CHANNEL` | +| `v_quant_type` | STRING | `"NONE"` | **new — §8**; same value set | +| `k_cache_dtype` | STRING | `""` | **new — §8**; `""` = the K cache tensor's own element type | +| `v_cache_dtype` | STRING | `""` | **new — §8**; `""` = the V cache tensor's own element type | +| `v_head_size` | INT | `0` | **new — §12**; `0` = `head_size`. Non-zero and `!= head_size` is legal **only** in `LATENT` | +| `rotary_offset` | INT | `0` | **new — §12.5** | +| `kv_cache_layout` | STRING | `"SEPARATE"` | **new — §12**; `SEPARATE` \| `LATENT` | +| `qk_output` | INT | `0` | **new — §11**; `0` none, `1` pre-softmax, `2` post-softmax | + +`k_cache_dtype` and `v_cache_dtype` name the *logical* element type of each cache. Every value is +spelled as the ONNX element type it denotes. `""` — the default — means the cache tensor's own +element type is also the logical type; `"float16"`, `"bfloat16"`, `"int8"` and `"float8e4m3fn"` name +that same type explicitly and must agree with the tensor. The reserved values `"int4"` and +`"float4e2m1"` describe sub-byte types packed two per byte into a `uint8` cache (§21.4), which +no ONNX tensor type can express here; they are rejected until a sub-byte backend exists. Every +value is a signed, zero-symmetric type — there is no zero-point input, so `uint4` / `uint8` are +deliberately not in the vocabulary (§8.3.1). + +A *dtype*, not a bit width, is the right unit here. Bit width alone cannot distinguish `int4` from +`float4e2m1` — same width, unrelated decode math — and for every non-packed format it is pure +redundancy against the cache tensor's element type, i.e. a second source of truth that can only ever +agree or disagree. K and V stay independent so future formats may use different precisions for each +(for example, TurboQuant-style mixed-precision caches). The current schema still binds both cache +tensors to `T_CACHE`; a future format whose K and V tensors have different ONNX element types must +additionally split that type constraint into `T_K_CACHE` and `T_V_CACHE`. Byte-backed sub-byte +formats can differ in logical width while both tensors remain `uint8`. + +Do **not** add `window_size_left` / `window_size_right` to opset 1. `local_window_size = W` has the +established meaning of admitting `W` positions *including* the current token; its Flash parameters +are `window_size_left = W - 1`, `window_size_right = 0`. + +### 4.6 Input-Mode State Machine + +Input mode is determined by `kv_cache_layout` and Q/K/V presence. No separate format attribute is +needed. + +| `kv_cache_layout` | `query` | `key` | `value` | `value_cache` | Mode | +|---|---|---|---|---|---| +| `SEPARATE` | Q | present | present | required | separate Q/K/V | +| `SEPARATE` | packed QKV | absent | absent | required | packed QKV | +| `LATENT` | absorbed Q | present (latent row) | absent | absent | absorbed MLA | + +Every other presence pattern is `INVALID_ARGUMENT`. In particular, `SEPARATE` preserves the shipped +rule exactly: K and V are packed inside `query` iff both `key` and `value` are absent. `LATENT` is +explicitly selected by `kv_cache_layout`, so its key-present/value-absent pattern cannot be confused +with packed QKV. + +Shape inference and runtime validation must inspect `kv_cache_layout` **before** applying the shipped +`value absent => packed QKV` branch. In `LATENT`, `query.shape[1] / num_heads` determines +`head_size`; in packed `SEPARATE`, the existing +`query.shape[1] / (num_heads + 2 * kv_num_heads)` formula remains unchanged. + +### 4.7 CUDA Graph Contract and `attention_metadata` + +Decode is launch-bound, so CUDA graph replay is the main performance lever and the schema must not +obstruct it. Two properties of the ORT implementation determine the whole contract: + +1. `cudaStreamSynchronize` on a capturing stream is illegal, so the unconditional per-step D→H sync + in `paged_attention.cc` (§18.4) makes the operator **uncapturable** today. +2. On replay, `InferenceSession::Run` short-circuits to + `cached_execution_provider_for_graph_replay_.ReplayGraph(...)` and never builds an + `ExecutionFrame`. **No kernel's `ComputeInternal` runs again.** Backend selection, grid + dimensions and workspace extents are all frozen at capture. + +Property 2 is the one that constrains the schema. A CPU-resident input is read exactly once, at +capture. `max_kv_len` and `total_kv_tokens` grow with every decode step, so a host input carrying +their *exact* values would leave a captured graph attending over the capture-step's KV length for the +rest of the sequence — silently wrong, and undetectable by the producer. "Recapture when the metadata +changes" is not a mitigation for decode, where it changes every token. + +> **Replay-invariance rule.** No host-visible value may determine a loop trip count, a mask boundary, +> or a memory extent that varies per step. Host values may only *select* the kernel and *size* the +> launch, and must stay valid for every step the captured graph will serve. Every per-step quantity +> is read on device from `cumulative_sequence_length`, `past_seqlens` and `block_table`. + +Three derivations satisfy the rule, none of which needs a synchronization: + +| Quantity | Source | Replay-safe because | +|---|---|---| +| decode vs. prefill dispatch | static shapes: `query.shape[0] == cumulative_sequence_length.shape[0] - 1` | shapes are fixed for a captured graph | +| grid size, split count, gather/workspace extents | static capacity bound `max_kv_len_bound = block_table.shape[1] * block_size` | independent of step | +| per-sequence KV length, causal and window masking, gather trip counts | device `past_seqlens` / `cumulative_sequence_length` | re-read from device memory on every replay | + +The shape test is a **performance heuristic only**. `token_count <= batch_size` does not prove that +every sequence contributes exactly one query token — one sequence may contribute two while another +contributes none. The paged decode kernel must therefore derive each token's sequence and position +from `cumulative_sequence_length` on device and stay correct for ragged input. Correctness never +depends on the heuristic being right, only speed. + +This removes the D→H sync **unconditionally** and without any new input, which is a stronger result +than making the sync conditional on a host hint. Sizing by the capacity bound instead of the exact +length costs empty split blocks that exit after a single device read — far cheaper than a per-layer, +per-step sync. + +`attention_metadata` is consequently demoted to optional **replay-wide bounds**: + +```text +attention_metadata : (2,) int32, OrtMemTypeCPUInput + [0] max_query_len_bound # 0 = unknown. Replay-wide upper bound on tokens from any one sequence. + [1] max_kv_len_bound # 0 = unknown. Replay-wide upper bound on total KV length of any sequence. +``` + +- Both entries are **upper bounds, never exact values**, and must hold for *every* step the node — + or the captured graph containing it — will serve. +- `0` means "no bound"; the implementation falls back to `token_count` for query length and to + `block_table.shape[1] * block_size` for KV length. +- The implementation clamps an over-large bound to the corresponding static limit. A non-zero + bound smaller than an actual per-step value violates the input contract and may omit work and + produce an incorrect result. The kernel cannot detect that violation without reading the device + length tensors back to the host. +- A valid bound may only shrink launch dimensions and workspace sizes. It must not enter a mask + comparison. Device loops use the current device lengths, additionally bounded by the trusted + launch/workspace extent. +- Even for an invalid bound, every device read must remain memory-safe: device lengths and static + tensor capacities guard all accesses. This safety property does not imply a correct result when + the producer violates the upper-bound contract. +- It must be a **graph input fed from the host**, never the output of an in-graph node. An in-graph + producer would be placed on the CPU EP and trip the `AreAllNodesInMainGraphAssignedToOneEp` check + that gates graph capture. + +`total_kv_tokens` is **not** part of the input. It was only ever used to size the MEA gather buffer, +which must now be sized by `batch_size * max_kv_len_bound` so that the allocation is replay-invariant. + +> **Implementation note (implemented).** `attention_metadata` exists as input 16 with the semantics +> above, and the D→H readback is gone from every configuration a CUDA graph can reach. +> +> Backend selection is now the static shape test of the table above: `token_count <= batch_size` +> selects the paged decode backend (when the cache is quantized or FlashAttention is ineligible), +> with no host knowledge of the actual per-sequence query lengths. That is safe because +> `PagedDecodeSplitKV` is indexed by **global query token** — it resolves each token's sequence and +> position from `cumulative_seqlens_q` on device and applies per-token causal and window masks — so +> a wrong heuristic costs speed, never correctness. The fused QK-Norm / rotary prologue was made +> token-indexed for the same reason, which removed its dependence on `max_query_len` entirely. +> +> Everything the host still needs is an upper bound, and each bound has a static fallback used when +> no metadata is supplied: +> +> | Host quantity | Consumer | Bound used | +> |---|---|---| +> | `max_query_len` | `params.seqlen_q` (Flash), `p.sequence_length` (MEA) — grid extent only | `max_query_len_bound`, else `token_count` | +> | `max_kv_len` | quantized-Flash `max_seqlen_k`, decode split count | `max_kv_len_bound`, else `block_table.shape[1] * block_size` | +> | `total_kv_tokens` | gather staging buffer extent | `batch_size * max_kv_len_bound` | +> +> Two narrow cases still take the readback, and only when the caller supplied **no** metadata at all: +> a gather backend (MEA, or FlashAttention on a quantized cache), because with no bound the +> capacity-based staging allocation would be a large over-allocation for a short prefill; and XQA, +> whose one-output-row-per-batch-index layout needs *proof* of exactly one token per sequence rather +> than the shape heuristic. Neither case can occur on a capturable step — a captured step is +> decode-shaped over a paged cache, so no gather runs, and a producer that captures must supply +> `attention_metadata` anyway, since replay-wide bounds are the only replay-safe host input. Both +> cases are prefill-side, and prefill is never captured. +> +> Measured with nsys on a single-node decode graph, this removes one `cudaStreamSynchronize` and two +> D→H `cudaMemcpyAsync` per node per `Run`, leaving zero of either — for the **unquantized** cache as +> well as the quantized ones, which is what makes `enable_cuda_graph=true` work for every KV type. + +Producers additionally owe the usual CUDA-graph obligations, which are outside this operator's +contract: fixed device addresses for `key_cache`, `value_cache`, `block_table`, `past_seqlens` and +`cumulative_sequence_length` across replays, and separate captures for prefill and decode. + +### 4.8 `query_positions` + +When absent, token `j` of sequence `b` has position `past_seqlens[b] + j` — the shipped behavior. +When present, element `j` supplies the linear logical position used for in-op RoPE and for backends +that support arbitrary-position causal and window masking. + +This covers linear position overrides and chunked prefill. It does **not** by itself encode tree +ancestry: Medusa and general tree attention additionally require a branch/ancestor mask, which is +deferred (§21). Backends that derive positions solely from packed row alignment are ineligible +whenever the supplied positions differ from the legacy sequence. + +Cached K is stored **after** RoPE, so reusing a cached block at a different logical position is +invalid unless the producer guarantees the block was rotated for that position, or the backend +re-rotates on read. `query_positions` does not make an already-rotated prefix relocatable. + +`query_positions` is `int32` and 1-D over `token_count`, unlike GQA's `position_ids`, which is +`int64` and 2-D over `(batch_size, sequence_length)`. The divergence is deliberate — the packed +varlen layout has no `(batch, seq)` grid — but it does mean a GQA↔PagedAttention rewriter needs a +`Cast` plus a reshape here. + +### 4.9 Type Constraints + +| Name | Allowed | Change | +|---|---|---| +| `T` | `float16`, `bfloat16` | unchanged | +| `T_CACHE` | `float16`, `bfloat16`, `int8`, `float8e4m3fn` | **new** (split out of `T`) | +| `T_KV_SCALE` | `float` | **new** | +| `QK` | `float`, `float16`, `bfloat16` | **new** — §11 | +| `S` | `int32` | unchanged | + +Splitting `T_CACHE` out of `T` is backward compatible: every previously valid model has +`T_CACHE == T`. The constraint name is `T_KV_SCALE`, matching GQA and the registration already in +`paged_attention.cc`. + +`uint8` is **intentionally** omitted from `T_CACHE`, even though GQA's `T_CACHE` already admits it +for packed INT4. There is no unsigned or sub-byte logical cache format specified for this operator +yet (§21), and widening a type constraint later is itself a compatible change, so nothing is lost by +waiting. This is a deliberate divergence from GQA, not an oversight. + +## 5. Feature: `slot_mapping` + +### 5.1 Problem + +Today the write slot for each token is *derived* inside `ReshapeAndCache`: + +``` +batch_id = binary_search(cumulative_seqlens_q, token_id) // which sequence owns this token +token_offset = token_id - cumulative_seqlens_q[batch_id] +position = past_seqlens[batch_id] + token_offset +block_id = block_table[batch_id * max_num_blocks_per_seq + position / block_size] +slot = block_id * block_size + position % block_size +``` + +This hard-codes "append `n_b` contiguous tokens at the end of sequence `b`". It cannot express: + +- **Prefix caching** — a prefill whose first *k* tokens hit an already-populated shared block must + *not* rewrite those slots (and must not have its KV recomputed into them). +- **Speculative / tree decoding** — draft tokens are written speculatively and rejected tokens must + be discarded; positions are not a contiguous run. +- **Chunked prefill with out-of-order scheduling**, where the scheduler owns slot assignment. +- **Cache-aware scheduling / block migration**, where the runtime, not the kernel, decides placement. + +It also concentrates a whole class of out-of-bounds bugs in the kernel (see [§18](#18-known-defects-to-fix-first)). + +### 5.2 Design + +Add optional input 10, `slot_mapping`, `int32`, shape `(token_count,)`. + +Each element is a **flat slot index** into the cache viewed as +`[num_blocks * block_size, kv_num_heads, head_size]`: + +``` +slot_mapping[t] = block_id * block_size + offset_in_block // 0 <= slot < num_blocks * block_size +slot_mapping[t] = -1 // do not write this token's K/V +``` + +Semantics: + +- **When present**, `slot_mapping` is authoritative for the *write* path. `ReshapeAndCache` performs + no binary search and no derivation; it reads `slot_mapping[token_id]` directly. Tokens with `-1` + skip the K/V store entirely (their Q still participates in attention). This is the prefix-cache + and rejected-speculative-token case. +- **When absent**, behavior is exactly today's derivation — full backward compatibility. +- `block_table` remains **required** in both cases: it drives the *read* path (which blocks make up + each sequence's context). `slot_mapping` only controls writes. +- `past_seqlens` and `cumulative_sequence_length` remain required; they still define the causal mask + and context length per sequence. + +### 5.3 Implementation + +`ReshapeAndCache` becomes templated on a slot resolver: + +```cpp +struct DerivedSlotResolver { /* current binary-search + block_table derivation */ }; +struct ExplicitSlotResolver { const int* slot_mapping; + __device__ int operator()(int token_id) const { return slot_mapping[token_id]; } }; +``` + +The explicit resolver removes the per-thread binary search over `cumulative_seqlens_q`, which is a +measurable win for large batches and eliminates the OOB failure mode. The kernel guards +`slot < 0 → return` and (debug builds) `slot < num_blocks * block_size`. + +`PagedAttentionData::slot_mappings` already exists as an unused field; it is renamed to +`slot_mapping` and wired. + +### 5.4 Validation + +- Rank 1, `shape[0] == token_count`. +- Element type `int32`. +- Range checking on device is a debug-build assertion only; host-side range checking would require a + D→H copy per step. The op documents that out-of-range values are undefined behavior, consistent + with `block_table` today. + +### 5.5 Why not derive-only + +Deriving slots keeps the graph smaller but pushes scheduling policy into the kernel. Every serving +framework that matters (vLLM, TRT-LLM, SGLang) passes an explicit slot mapping precisely because +the scheduler — not the attention kernel — owns block allocation. Keeping the derived path as the +default preserves the simple single-turn case; the explicit path unlocks serving. + +## 6. Feature: Attention Sink (`head_sink`) + +### 6.1 Math + +Identical to GQA. With per-head sink value $s_h$ over $T$ attended positions: + +$$ +\text{softmax}_i = \frac{e^{x_i - m}}{e^{s_h - m} + \sum_{j} e^{x_j - m}}, \qquad m = \max\!\left(s_h, \max_j x_j\right) +$$ + +Equivalent to appending one extra logit $s_h$ that contributes to the denominator only. Used by +GPT-OSS. The internal sink-enabled softmax path is selected whenever `head_sink` is present; there +is no standalone attribute for an unlearned zero-valued sink. + +### 6.2 Design: exact LSE epilogue (preferred) + +FlashAttention varlen already produces `softmax_lse` and the buffer is already allocated in +`paged_attention.cc`. For query token $t$ and head $h$, Flash returns + +$$ +\ell_{t,h} = \log\!\sum_j e^{x_j}, \qquad o_{t,h} = \frac{\sum_j e^{x_j} v_j}{\sum_j e^{x_j}} +$$ + +Adding the sink changes only the denominator, so the corrected output is an **exact** elementwise +rescale: + +$$ +o^{\text{sink}}_{t,h} = o_{t,h} \cdot \frac{e^{\ell_{t,h}}}{e^{\ell_{t,h}} + e^{s_h}} + = o_{t,h} \cdot \frac{1}{1 + e^{\,s_h - \ell_{t,h}}} +$$ + +This is a single `token_count × num_heads × head_size` pass over the output, computed in FP32 and +cast back to `T`. It requires **no change to the Flash kernel**, composes with sliding window, +softcap, GQA grouping, and packed QKV, and is numerically exact (the $1/(1+e^{s-\ell})$ form is +stable for both signs of $s - \ell$). + +Implementation: `LaunchApplyHeadSink(output, softmax_lse, head_sink, token_count, num_heads, head_size, stream)` +in `paged_attention_impl.cu`, invoked after `mha_varlen_fwd` when `data.head_sink != nullptr`. + +### 6.3 MEA fallback + +The CUTLASS fMHA path does not expose an LSE output in the current ORT wrapper. Options, in order of +preference: + +1. Enable the fMHA LSE output (`output_accum` / `logsumexp` variant) and reuse the same epilogue. +2. Reject `head_sink` on the MEA path and require Flash (SM80+, FP16/BF16). Since PagedAttention is + a server op and already requires Flash *or* MEA, and GPT-OSS deployments are SM80+, this is an + acceptable Phase-1 restriction — but it must be a clear `INVALID_ARGUMENT`, not silent wrong math. + +Phase 1 ships option 2 with option 1 as follow-up. + +### 6.4 Validation + +- Rank 1, `shape[0] == num_heads`, element type `T`. + +## 7. Feature: Fused QK-Norm + +### 7.1 Math + +Identical to GQA §3. Per head, over `head_size` channels, applied to Q and K **before** RoPE: + +$$ +x_\text{norm}[c] = x[c] \cdot \frac{1}{\sqrt{\frac{1}{H}\sum_{j} x[j]^2 + \epsilon}} \cdot w[c] +$$ + +`H = head_size`; `w` is `q_norm_weight` for Q and `k_norm_weight` for K, both 1D `(head_size,)` of +type `T`, shared across heads; `epsilon = qk_norm_epsilon` (default `1e-6`). Sum of squares is +reduced in FP32. + +Required by **Qwen3, Gemma 2/3, OLMo2, SmolLM3**. + +### 7.2 Design + +PagedAttention already has a prologue kernel that unpacks packed QKV, applies RoPE, and calls +`ReshapeAndCache`. QK-Norm is fused into that prologue: + +``` +[packed QKV split] → [QK-Norm on Q and K] → [RoPE on Q and K] → [ReshapeAndCache writes K,V] +``` + +Critically, **the normalized-and-RoPE'd K is what lands in the block cache**, matching GQA's +`UnpackRoPEAppend` ordering. Cached K is therefore directly consumable by attention on later steps — +QK-Norm is never re-applied to cached K. + +### 7.3 Workspace implication + +Today `workspace_buffer` is allocated only when `do_rotary_` or `is_packed_qkv`. QK-Norm requires a +writable Q (and K) buffer even when neither holds. The allocation condition becomes: + +```cpp +const bool needs_prologue = do_rotary_ || parameters.is_packed_qkv || has_qk_norm; +if (needs_prologue) { + workspace_buffer_bytes = sizeof(T) * token_count * (hidden_size + kv_hidden_size); +} +``` + +(When only QK-Norm is active, K must also be staged because its normalized form is what is cached.) + +### 7.4 Validation + +- `q_norm_weight` and `k_norm_weight` must be provided **together**; one alone is `INVALID_ARGUMENT`. +- Rank 1, `shape[0] == head_size`, element type `T`. +- `qk_norm_epsilon > 0`. + +## 8. Feature: Quantized Paged KV Cache + +### 8.1 Goal + +Store the block cache in INT8 or FP8 E4M3 while `query` remains FP16/BF16, halving (or better) the +dominant memory consumer in a serving deployment and proportionally reducing HBM traffic on the +decode path. Scope for this phase: **`PER_TENSOR` and `PER_CHANNEL`**, with `k_cache_dtype` and +`v_cache_dtype` left at `""` (or naming the cache tensor's own element type). +INT4 is deferred ([§19](#19-phasing)). + +### 8.2 Schema + +- `key_cache` / `value_cache` move from `T` to `T_CACHE ∈ {float16, bfloat16, int8, float8e4m3fn}`. + `uint8` is intentionally excluded until a sub-byte format is specified (§4.9, §21). +- `k_scale` / `v_scale` (inputs 14, 15), type `T_KV_SCALE` = **always FP32**, matching GQA. +- Attributes `k_quant_type`, `v_quant_type` ∈ `{"NONE", "PER_TENSOR", "PER_CHANNEL"}`, plus + independent `k_cache_dtype` and `v_cache_dtype` attributes, which stay `""` while every logical + type is expressible as an ONNX element type. +- Kernel becomes `PagedAttention`, registered for the same combinations GQA uses: + `{MLFloat16, BFloat16} × {same as T, int8_t, Float8E4M3FN}` (plus `uint8_t` if and when INT4 lands). + +### 8.3 Scale layout under the block layout + +The block cache is `(num_blocks, block_size, kv_num_heads, head_size)`. A `PER_CHANNEL` scale of +shape `(kv_num_heads, 1, head_size)` — the same shape GQA uses — broadcasts naturally over the +leading `(num_blocks, block_size)` dims. **This is why the GQA scale shape is reused verbatim**: the +quantize/dequantize helpers index only on `(kv_head, channel)` and are layout-agnostic. + +| Mode | Scale shape | Indexing | +|---|---|---| +| `PER_TENSOR` | `(1,)` | scalar | +| `PER_CHANNEL` | `(kv_num_heads, 1, head_size)` | `scale[h * head_size + c]` | + +`v_scale` uses the same last dimension, `head_size`: it exists only where a `value_cache` exists, +which is `"SEPARATE"` mode, and `effective_v_head_size == head_size` there (§12.3). Under +`"LATENT"` there is one physical cache and `k_scale` alone describes it, so `v_scale` is rejected +(§8.7, §12.9); the V dequant of that cache indexes `k_scale` with the `head_size` stride and reads +only its leading `v_head_size` entries. + +Symmetric quantization, same formulas as GQA: + +| Type | Range | Quantize | +|---|---|---| +| INT8 | `[-128, 127]` | `q = clamp(round(x / scale), -128, 127)` | +| FP8 E4M3 | `[-448, 448]` | `q = clamp(x / scale, -448, 448)` | +| INT4 (deferred) | `[-8, 7]`, 2/byte | last cache dim becomes `(head_size + 1) / 2` | + +#### 8.3.1 Zero point: always 0, and why the vocabulary is signed-only + +There is **no zero-point input and no zero-point attribute**. Dequantization is exactly +`x = q * scale`, with an implied zero point of `0`. Consequently every value `k_cache_dtype` / +`v_cache_dtype` may name is a *signed, zero-symmetric* type: `int8`, `float8e4m3fn`, `int4`, +`float4e2m1`, plus the unquantized floats. `uint4` and `uint8` are deliberately **absent**: an +unsigned logical type under symmetric quantization implies an offset of `2^(bits-1)` (128 for uint8, +8 for uint4) that nothing in the contract carries, so admitting the name would let two +implementations disagree about whether the stored code is biased. Unsigned logical types must arrive +together with the optional zero-point inputs sketched in §21.3, not before. + +> **Storage bias is not a zero point.** INT4 is *stored* in an unsigned nibble biased by `+8` +> (`store = q + 8`, `load = nibble - 8`), which is how ORT's MLAS KV path already packs its `S4` +> modes (`kInt4Bias` in `mlas/lib/qkv_quant_common.h`). That bias is a storage encoding removed +> before the value is used; the *logical* value stays signed in `[-8, 7]` and the dequantization is +> still `q_signed * scale`. Naming that format `uint4` would be wrong. + +**The kernels depend on this, not merely comply with it.** The §8.5 Phase 3 decode kernel folds +`k_scale` into Q at load time and `v_scale` into the reduce epilogue. Those foldings are exact only +because the zero point is 0: + +$$\sum_c q_c \cdot (k_c \cdot s_k) = \sum_c (q_c \cdot s_k) \cdot k_c$$ + +With a non-zero zero point $z$ the score becomes +$\sum_c q_c (k_c - z) s_k = s_k \sum_c q_c k_c - z\,s_k \sum_c q_c$, so a per-(token, head) +correction term appears in the QK product and a second one in the PV product. That is a structural +change to the kernel, not an extra subtraction — which is the concrete reason zero points are a +versioned-successor topic rather than a late addition. + +### 8.4 Write path + +`ReshapeAndCache` gains the quantization step. After QK-Norm and RoPE, K/V are quantized as they are +scattered to their slots. This is a natural fit: the kernel is already elementwise over +`(token, kv_head, channel)`, which is exactly the `PER_CHANNEL` scale index. + +### 8.5 Read path + +Phase 2 (correctness first): **dequantize-on-gather**. The MEA fallback already materializes a +gathered `[total_kv_tokens, num_heads, head_size]` KV buffer via `GatherAndExpandPagedKVCache`; that +kernel is extended to dequantize while gathering, and the Flash varlen path is routed through the +same gather when the cache is quantized. Cost: one extra pass over the live context per step, and +the gather buffer is FP16/BF16-sized. Correct, low-risk, and reuses existing code. + +Phase 3 (performance): a **paged decode kernel with in-kernel dequantization**, following the GQA +XQA approach — for `PER_TENSOR`, fold `k_scale` into the QK score scale and `v_scale` into the output +accumulator; for `PER_CHANNEL`, fold the K scale into Q before attention and apply the V scale to the +output afterward. Both are `O(num_heads * head_size)` passes and avoid touching the full cache. This +is the path that makes a quantized cache actually pay off; the gather-based Phase 2 mostly buys +memory capacity, not bandwidth. + +### 8.6 Build gating + +Mirror GQA: `onnxruntime_USE_FP8_KV_CACHE` (default ON), `onnxruntime_USE_INT4_KV_CACHE` +(default OFF). INT8 always built. + +### 8.7 Validation + +- `k_quant_type != "NONE"` requires `k_scale`; likewise for V. Conversely, a `k_scale` with + `k_quant_type == "NONE"` is `INVALID_ARGUMENT`. +- `T_CACHE != T` requires a non-`NONE` quant type; `T_CACHE == T` requires both to be `NONE` and both + scales to be absent. +- `k_cache_dtype` and `v_cache_dtype` are `""` for every cache this operator stores, quantized or + not: the cache tensor's element type is the logical element type. Naming that type explicitly + (`"float16"`, `"bfloat16"`, `"int8"`, `"float8e4m3fn"`) is accepted but must agree with the tensor. + `"int4"` and `"float4e2m1"` are reserved for a `uint8` packed cache and are rejected + until one exists. Unsigned logical types (`uint4`, `uint8`) are rejected outright: quantization + here has no zero point (§8.3.1). +- In `"LATENT"` mode only K storage exists: `k_quant_type` and `k_scale` describe the latent row, + `v_quant_type` and `v_cache_dtype` must be unset, and `v_scale` must be + absent because V is a view of K. +- FP8 is available when ORT is built with `onnxruntime_USE_FP8_KV_CACHE`; no additional runtime + architecture gate is required for the conversion path used by this operator. +- `PER_CHANNEL` scale shape must be exactly `(kv_num_heads, 1, head_size)` for both K and V. There + is no `v_head_size`-shaped scale: `v_scale` only exists alongside a `value_cache`, i.e. in + `"SEPARATE"` mode, where `effective_v_head_size == head_size`. + +> **Implementation note (P3, implemented).** Delivered: `int8` and `float8e4m3fn` caches with +> `PER_TENSOR` / `PER_CHANNEL` granularity, independently selectable for K and V, via +> `PagedAttention` registered for `{MLFloat16, BFloat16} × {int8_t, Float8E4M3FN}` in +> addition to the unquantized pairs. Deviations from the text above: +> +> - **Read path is §8.5 Phase 2 only.** `GatherAndExpandPagedKVCache` dequantizes while gathering, +> and the Flash varlen path is routed through the same gather when the cache is quantized (using +> `num_heads = kv_num_heads` so no GQA expansion happens, which keeps Flash's grouped layout). The +> §8.5 Phase 3 paged decode kernel with in-kernel dequantization is **not** implemented. +> - Because a quantized cache never reaches Flash's *paged* kernel, the `block_size` tiling +> constraint of §18.1 does not apply to it; Flash eligibility skips that check when the cache is +> quantized. Any power-of-two `block_size >= 16` works with a quantized cache on either backend. +> - **`uint8` / INT4 not added.** `T_CACHE` is `{float16, bfloat16, int8, float8e4m3fn}`, so +> `k_cache_dtype` and `v_cache_dtype` must be `""` or name the cache tensor's own element type. +> - **No SM89/SM90 gate for FP8.** `Float8E4M3FN`'s converting constructor uses +> `__nv_cvt_float_to_fp8`, which is available on every architecture ORT builds for from CUDA 11.8 +> onward, so the arch check in §8.7 would reject working configurations. FP8 remains gated at +> *build* time by `onnxruntime_USE_FP8_KV_CACHE`. +> - Parity tests compare the updated cache at one quantization step of slack. Rotary and RMSNorm are +> computed ~1 fp16 ULP differently on the host, which is enough to move a value across a rounding +> boundary and flip the stored code by one LSB. + +> **Implementation note (P5, implemented).** §8.5 Phase 3 is now in place as a purpose-built +> flash-decoding kernel (`PagedDecodeSplitKV` + `PagedDecodeReduce` in `paged_attention_impl.cu`), +> not by reusing the vendored XQA kernel. XQA does have paged-KV support, but its page list is +> `[batch][beam][2][max_pages]` over a *single* K/V pool whereas PagedAttention has two pools and one +> shared `block_table`, and it restricts `head_size ∈ {64, 128, 256}` and `group_size ∈ {4, 8, 16, +> 32}` while multiplying instantiations across page size × group size × head size × dtype × quant +> type. A ~250-line kernel covers the whole schema instead. +> +> - **Both scale foldings are exact and granularity-agnostic.** K folds into Q at load time +> (`q_sh[c] = float(q[c]) * GetCacheScale(k_scale, kv_head * head_size + c, k_per_channel)`), so +> `PER_TENSOR` is just the `per_channel == false` branch of the same expression rather than a +> separate "fold into the softmax scale" path. V folds into the epilogue: `v_scale_c` does not +> depend on the KV position, so it factors out of the accumulation entirely and never enters the +> softmax denominator. +> - The kernel reads pages in place at their stored width, so a decode step touches the KV cache once +> at `int8`/`fp8` bandwidth instead of gathering and dequantizing the whole live context. +> - `softcap` matches FlashAttention bit-for-bit: `softcap * tanh(qk_raw * scale / softcap)`, which is +> what `flash_api.cc` produces from `params.softcap = softmax_scale / softcap` and +> `params.scale_softmax = softcap`. +> - The attention sink enters as one extra `exp(sink - m_final)` term in the final denominator, which +> is algebraically identical to §6.2's `factor = 1 / (1 + exp(sink - lse))` epilogue but does not +> need FlashAttention's log-sum-exp output. §6.3's MEA restriction therefore applies only to MEA. +> - Sliding window uses the token's own causal position `q_pos = kv_len - 1`, admitting +> `t ∈ [kv_len - local_window_size, kv_len)`, matching Flash's `window_size_left = local_window_size - 1`. +> - **Backend gating.** The kernel is selected by the static shape test of §4.7 +> (`token_count <= batch_size`) when the cache is quantized *or* FlashAttention is unavailable; +> unquantized FlashAttention-eligible shapes keep using FlashAttention. `sdpa_kernel = 512` +> (`AttentionBackend::DECODER_ATTENTION`) forces it, which is how the unquantized path is tested. +> The shape test is a heuristic: one CTA owns one **global query token**, resolves its sequence and +> in-sequence position from `cumulative_seqlens_q` on device, and masks against +> `past_seqlens[b] + q_index + 1`, so the kernel is correct for arbitrary ragged input (including +> full prefill) and a wrong heuristic only costs speed. That is what removes the D→H sync. +> - Split-KV: `ComputePagedDecodeSplits` splits the KV range across up to 32 CTAs only when +> `token_count * num_heads` would leave the device under-occupied. `max_kv_len` may be an upper +> bound. Empty splits publish `(max = -FLT_MAX, denom = 0)` and the reduce kernel skips them, so +> their accumulator slice is never read. +> - The `FlashAttention` / `EfficientAttention` prologue (packed-QKV unpack, fused QK-Norm + rotary, +> `ReshapeAndCache`) was factored into a shared `PrepareQueryAndCache`, which the decode backend +> reuses. It lives outside the `USE_FLASH_ATTENTION` / `USE_MEMORY_EFFICIENT_ATTENTION` guards +> because the decode backend needs neither. +> - **Still deferred from P5:** the fused MLA decode backend (§12.7). + +## 9. Feature: Sliding Window Attention + +### 9.1 State + +`local_window_size` is already an attribute and is already passed to Flash varlen as +`local_window_size - 1` (Flash's `window_size_left` excludes the current token; ORT's convention +includes it) and to the MEA params. Required by **Mistral, Gemma 2/3, Phi-3, GPT-OSS** (which +alternates full and sliding layers). + +### 9.2 Work items + +1. **Convention verification.** Assert that PagedAttention's window semantics match GQA's exactly — + "the window includes the new token and only extends to the left". Add a direct GQA↔PagedAttention + parity test at several window sizes, including `window >= context` (must equal full attention) and + `window == 1`. +2. **Block-table pruning.** With a window of $W$ and context length $L_b$ for sequence $b$, all KV + blocks whose highest position is below $L_b - W$ are unreachable. The read path should start at + + ``` + first_block = max(0, (L_b - W) / block_size) + ``` + + and only walk `block_table[b][first_block ...]`. For a 128K context with a 4K window this reduces + both the gather volume and the number of block-table indirections by ~30×. This is the single + largest win available for long-context sliding-window models and is *only* expressible in the + paged layout. +3. **Freeing pruned blocks is a runtime concern**, not a kernel concern; the kernel must tolerate a + `block_table` whose leading entries are stale or `-1`. Define `-1` in `block_table` as "block not + mapped, treat as masked out", consistent with `slot_mapping`'s `-1`. +4. **Composition with `head_sink`.** GPT-OSS uses sinks *and* sliding windows in the same layer. The + LSE epilogue in [§6](#6-feature-attention-sink-head_sink-and-smooth-softmax) composes trivially + because Flash's LSE already reflects the window mask. +5. **Composition with `softcap`.** Both are already handled inside Flash varlen; add a combined test. + +### 9.3 Bounded-capacity (rolling) sliding-window cache + +§9.2 is about not *reading* out-of-window KV. This sub-section is about not *storing* it: when +`local_window_size = W > 0`, a sequence of context length $L$ only ever needs $O(W)$ cached tokens, +not $O(L)$. For a 128K context with a 4K window that is a ~32× reduction in KV memory, which is what +decides how many sequences fit on the device. + +**Yes — it is expressed through `slot_mapping`, but not through `slot_mapping` alone.** The +capability splits exactly along the write/read line of [§5](#5-feature-slot_mapping): + +| Concern | Mechanism | Owner | +|---|---|---| +| Reuse a physical block for a later token | `slot_mapping` (§5) — the *write* path | Runtime / scheduler | +| Stop reading an evicted block | `block_table[b][i] = -1` — the *read* path | Runtime / scheduler | +| Not attending out-of-window positions | `local_window_size` mask inside the backend | Kernel | +| Not *fetching* out-of-window blocks | window-clamped block walk (§9.2 item 2) | Kernel | + +The derived slot resolver cannot do this: it computes +`slot = block_table[b][position / block_size] * block_size + position % block_size`, so a block is +implicitly owned by one absolute position range forever and the allocation grows with $L$. An +explicit `slot_mapping` lets the same physical block be rewritten by a token $W$ positions later, +which is the whole trick. + +#### 9.3.1 Invariant: the block table stays indexed by absolute position + +`block_table[b][i]` must keep meaning "the block holding positions +`[i * block_size, (i+1) * block_size)` of sequence `b`". Do **not** compact or shift the row so that +entry 0 becomes the window start. Every backend derives its mask from absolute positions — +FlashAttention from `seqlen_k` (`n_block_min` is computed from `actual_seqlen_k`), the paged decode +kernel from `kv_len - 1`, the latent kernel from `past_seqlens[b] + s` — and RoPE has already been +applied at the absolute position when the row was written. Rotating the row would silently decouple +position from mask and from RoPE phase. + +So an evicting runtime presents a **sparse row**: entries below the window are `-1`, entries inside +it point into a small recycled pool. `-1` is defined as *"block not mapped — treat every position in +it as masked out"*, consistent with `slot_mapping`'s `-1`, and is already honoured by the slot +resolver, the gather kernel, the paged decode kernel and the latent kernel. + +#### 9.3.2 Allocation policies + +Let $N = \lceil W / \text{block\_size} \rceil + 1$. The `+1` absorbs the partially filled boundary +block, so the oldest still-in-window position is never evicted early. + +1. **Ring buffer** (TRT-LLM's *cyclic* KV cache, vLLM's sliding-window allocator). The runtime holds + $N$ blocks `ring[0 .. N-1]` per sequence and fills, for absolute position $p$ with + $i = \lfloor p / \text{block\_size} \rfloor$: + + ``` + slot_mapping[t] = ring[i % N] * block_size + (p % block_size) + block_table[b][i] = ring[i % N] // i inside the live range + block_table[b][j] = -1 // j below the live range + ``` + + Steady-state KV memory per sequence is exactly `N * block_size` tokens, independent of $L$. +2. **Free-list eviction.** After a step, return every block whose highest position falls below the + window start to the global allocator and write `-1` into its entry. Same asymptotic memory, more + allocator traffic, but it interleaves with prefix caching and with layers that are *not* sliding + (GPT-OSS alternates), which a per-sequence ring does not. + +**Eviction condition.** The earliest query position a sequence will use in the current step is +$p_{\min} = \texttt{past\_seqlens}[b]$, and its window admits positions $\ge p_{\min} - W + 1$. So +after the step, blocks with + +$$i < \left\lfloor \frac{\max(0,\ \texttt{past\_seqlens}[b] - W + 1)}{\text{block\_size}} \right\rfloor$$ + +are unreachable forever (query positions only move forward), and may be recycled. Eviction is +therefore a *post-step* runtime action; the kernel never frees anything. + +#### 9.3.3 Backend obligations + +| Backend | Behaviour on an out-of-window `-1` entry | +|---|---| +| Paged decode | `kv_begin = max(kv_begin, kv_len - W)` already skips the range; `block_id < 0` additionally forces `-FLT_MAX`. Safe. | +| Latent / MLA | `kv_begin = max(0, kv_end - W)`. Safe. | +| FlashAttention varlen (paged) | `n_block_min = max(0, (m_block * kBlockM + seqlen_k - seqlen_q - window_size_left) / kBlockN)` — out-of-window pages are never dereferenced. Requires `block_size % kBlockN == 0`, which is already a Flash *eligibility* condition (§13). Safe. | +| Gather path (MEA, quantized cache) | The gather **zero-fills** unmapped rows. A zero key yields logit $0$, not $-\infty$, so correctness relies on MEA's own window mask covering exactly the same positions. It does, because both derive from the same $W$ — but this makes the invariant below load-bearing. | + +> **Invariant (load-bearing).** A `-1` entry must never overlap a position the mask admits. The +> runtime may only unmap blocks strictly below the window start. Violating it is silently wrong on +> the gather path (zeros contribute weight) rather than loudly wrong. Add a debug-build assertion in +> the gather kernel: `block_id >= 0 || pos < kv_len - W`. + +#### 9.3.4 Follow-on kernel work this exposes + +1. **Window-clamped gather.** `GatherAndExpandPagedKVCache` still materializes $L$ tokens of + workspace even when only $W$ are readable. Starting the gather at + `first_block = max(0, (L_b - W) / block_size)` shrinks both the workspace and the gather traffic + to $O(W)$ — this is §9.2 item 2 applied to the gather path, and it is the difference between the + quantized/MEA backends being $O(L)$ and $O(W)$ per step. +2. **Split-KV layout.** `ComputePagedDecodeSplits` lays splits out over the full `kv_len` and each + CTA then clamps to `kv_len - W`, so with $W \ll L$ most CTAs launch only to exit immediately. The + split range should be computed over the clamped interval `[max(0, kv_len - W), kv_len)`. +3. **No schema change is required.** The block-table *row* stays $O(L / \text{block\_size})$, but at + 4 bytes per `block_size` tokens versus `2 * kv_num_heads * head_size * sizeof(T)` bytes per token + of KV, the row is ~0.01% of what it indexes — not worth a breaking change. Shrinking the row too + would need a per-sequence window origin (a rotated block table plus a `kv_start_positions` input); + that is listed with the other deferred contract changes in [§21](#21-deferred-breaking-contract). + +#### 9.3.5 Validation + +- Ring-buffer run (only $N$ blocks allocated per sequence, `slot_mapping` recycling them, leading + `block_table` entries `-1`) must be **bit-identical** to a full-cache run with the same window, for + every decode step past $L > W$. +- `W >= L` with recycling enabled must still equal full attention (nothing is ever evicted). +- Mixed batch: some sequences past the window, some not, in the same step. +- Composition: ring buffer × quantized cache (the gather path is the risky one) × `head_sink`. + +## 10. Feature: `attention_bias` + +### 10.1 GQA-compatible shape + +GQA defines `attention_bias` as +`(batch_size or 1, num_heads or 1, sequence_length, total_sequence_length)`. PagedAttention keeps +the same rank, dimension roles, and independent batch/head broadcasting. Its packed ragged batch has +no single query or context length, so the two sequence dimensions use batch maxima: + +```text +attention_bias : (batch_size or 1, num_heads or 1, query_length_capacity, context_length_capacity) +``` + +The last two dimensions are replay-wide capacities satisfying +`query_length_capacity >= max_b q_len_b` and +`context_length_capacity >= max_b (past_seqlens[b] + q_len_b)` for every step the node or captured +graph will serve. They may equal the exact maxima in an uncaptured run. The tensor may broadcast +across all sequences, all heads, or both, exactly as GQA does. + +### 10.2 Design + +Input 18, `attention_bias`, type `T`. For packed token `t` belonging to sequence `b`, let +`j = t - cumulative_sequence_length[b]` be its sequence-local query offset. Query head `h` and +logical KV position `k` use: + +```text +attention_bias[b_or_0, h_or_0, j, k] +``` + +The first two indices select `0` when that dimension broadcasts. Entries with `j >= q_len_b` or +`k >= past_seqlens[b] + q_len_b` are padding and are ignored. This is the direct packed equivalent +of GQA's indexing contract; flattening batch and query into `(num_heads, token_count, ...)` is +rejected because it loses GQA's batch-broadcast dimension and prevents schema-level parity. + +The bias capacities are host-visible shapes and therefore frozen during CUDA graph replay. Device +code must bounds-check `j` and `k` against them before reading the bias. An under-sized capacity is +a producer contract violation that may omit attention work, but it must not cause an out-of-bounds +read. This mirrors the trusted-bound rule for `attention_metadata` (§4.7). + +- Rejected on the Flash varlen, paged-decode, and current MEA paths. Initially supported only on the + unfused fallback, matching GQA. MEA may become eligible after its wrapper supports the rank-4 + batch/head broadcasting and ragged sequence-local row strides. +- Bias is applied after `scale` and before `softcap`, matching GQA. +- Emit an explicit `INVALID_ARGUMENT` when `attention_bias` is combined with a backend that cannot + serve it, naming the backend and the reason. + +### 10.3 No dedicated ALiBi input + +Do not add `alibi_slopes` in this schema revision. GQA has no corresponding input, and the general +`attention_bias` contract already expresses ALiBi when an exporter is willing to materialize it. +A dedicated fused representation remains an additive future extension if a concrete model or +performance requirement justifies adding it to both attention operators. + +## 11. Feature: `output_qk` + +### 11.1 Design + +- Attribute `qk_output`: `0` = no output (default), `1` = pre-softmax scores, `2` = post-softmax + probabilities. Same encoding as GQA's `QKOutputType`. +- Output 3, `output_qk`, type `QK`, shape `(num_heads, token_count, max_context_len)`. `QK` is a + separate type constraint (`float`, `float16`, `bfloat16`) so the score matrix can be emitted in + `float32` regardless of the activation type, exactly as in GQA. +- Use cases: interpretability, speculative-decode scoring, kernel debugging and parity triage. + +### 11.2 Constraints + +- Supported on the **unfused / MEA** paths only. Fused Flash kernels never materialize the score + matrix; requesting `output_qk` forces a fallback backend and is documented as such. +- Memory is `O(num_heads × token_count × max_context_len)` and can dwarf the KV cache itself. + The op must reject configurations where the output tensor would exceed a configurable byte + threshold rather than attempting the allocation. +- `qk_output != 0` with fewer than 4 outputs, or 4 outputs with `qk_output == 0`, is + `INVALID_ARGUMENT`. Shape inference must only touch output 3 when `ctx.getNumOutputs() > 3`. + +## 12. Feature: Multi-head Latent Attention (MLA) + +### 12.1 What MLA is + +DeepSeek-V2 / V3 / R1 replace the per-head K/V cache with a single low-rank **latent** vector per +token: + +- `compressed_kv` of width `kv_lora_rank` (512 in V3), shared by all heads; +- `k_pe` of width `qk_rope_head_dim` (64), shared by all heads (MQA-style), carrying RoPE. + +The cached footprint is `512 + 64 = 576` elements per token, against +`2 × num_heads × head_size = 2 × 128 × 128 = 32768` for a comparable MHA model — a ~57× reduction. +That reduction is the entire point of MLA, and it is a **paged-serving** feature: the win lands in +decode, where the KV cache dominates both capacity and bandwidth. MLA therefore belongs in +`PagedAttention` rather than in GQA. + +MLA has two mathematically equivalent evaluation forms: + +| Form | What attention sees as K/V | `head_size` | `v_head_size` | `kv_num_heads` | Used for | +|---|---|---|---|---|---| +| **Non-absorbed** | per-head `k = [k_nope(128); k_pe(64)]`, `v(128)`, produced by applying `W_UK` / `W_UV` to the latent | 192 | 128 | `num_heads` | Prefill with no cached prefix | +| **Absorbed** | the latent itself: `k = [compressed_kv(512); k_pe(64)]`, `v = compressed_kv(512)` | 576 | 512 | 1 | Decode, chunked prefill, prefix caching | + +Absorption folds `W_UK` into Q (`q_nope' = q_nope @ W_UKᵀ`, 128 → 512) and `W_UV` into the output +projection. Both are ordinary **MatMuls in the graph**; the attention op never sees a projection +weight. This keeps the operator an attention operator. + +### 12.2 Design: MLA is absorbed-form MQA with `v_head_size < head_size` + +Absorbed MLA is **already** the operator `PagedAttention` is, except for three properties: + +1. `v_head_size != head_size` (512 vs. 576). +2. V is not a separate tensor — it is the **leading `v_head_size` channels of K**. There is one + cache, not two. +3. RoPE applies to a **suffix** of the head dimension (channels 512–575), not a prefix. + +Everything else — packed varlen Q, `block_table`, `slot_mapping`, in-place cache update, causal +masking over ragged sequences, quantized cache — is unchanged. MLA is therefore added as a **mode of +`PagedAttention`**, expressed through three MLA-specific attributes, and not as a new operator. This is +the same conclusion FlashMLA, vLLM's MLA backend, and SGLang reached: absorbed MLA decode is MQA with +a wide head and a shared K/V buffer. + +### 12.3 Schema additions + +MLA is selected **explicitly** by `kv_cache_layout="LATENT"`, never inferred from input presence. + +| Addition | Meaning | +|---|---| +| Attribute `kv_cache_layout` (STRING, default `"SEPARATE"`) | `"LATENT"` selects absorbed MLA: one physical cache, V aliasing K. | +| Attribute `v_head_size` (INT, default `0`) | Head width of V and of each output head. `0` means "same as `head_size`". A value differing from `head_size` is legal **only** in `"LATENT"` mode. | +| Attribute `rotary_offset` (INT, default `0`) | First channel within `head_size` covered by RoPE (§12.5). | +| Input 4 `value_cache` becomes `Optional` | Absent in `"LATENT"`, where V is the leading `v_head_size` channels of `key_cache`. Still required in `"SEPARATE"`. | +| Input 2 `value` absent while `key` is present becomes legal | Only under `"LATENT"` (§4.6). In `SEPARATE`, the shipped rule stands: `key` **and** `value` absent means packed QKV. | +| Output 0 shape generalized to `(token_count, num_heads * effective_v_head_size)` | Identical to today whenever `v_head_size == 0`. | + +All six are backward compatible: existing models set none of the attributes, supply `value_cache`, +and get byte-identical behavior. + +**Why `v_head_size != head_size` is confined to `LATENT`.** Asymmetric K/V widths in `SEPARATE` mode +would be a second, independent feature: it needs a `value_cache` whose last dimension differs from +`key_cache`, which the shipped implementation contradicts (it builds `value_cache_out_shape[3]` from +`head_size` unconditionally and validates the two cache shapes as identical), and no available +backend supports it — ORT's Flash and CUTLASS fMHA both require `v_head_size == head_size`. Allowing +it in the schema without backend coverage would only create a validation surface with nothing behind +it. `effective_v_head_size` is therefore defined as: + +```text +effective_v_head_size = (kv_cache_layout == "LATENT" && v_head_size != 0) ? v_head_size : head_size +``` + +and `v_head_size != 0 && v_head_size != head_size` in `SEPARATE` mode is `INVALID_ARGUMENT`. + +A consequence worth stating explicitly, because it is why the rest of this document writes +`head_size` rather than `v_head_size` almost everywhere: the only tensors whose last dimension can +ever be `effective_v_head_size` are **output 0** and the **V view of `key_cache`**. `value`, +`value_cache`, and `v_scale` are all absent in `"LATENT"` (§4.6, §12.9), and in `"SEPARATE"` their +width is `head_size` by the rule above. Asymmetric K/V widths for a real `value_cache` arrive only +with the `"KV_CONCAT"` layout deferred in §21. + +### 12.4 Concrete node contract (DeepSeek-V3, absorbed) + +| Tensor / attribute | Value | +|---|---| +| `query` | `(token_count, 128 * 576)` — per head `[q_nope @ W_UKᵀ (512); q_rope (64)]` | +| `key` | `(token_count, 1 * 576)` — `[compressed_kv (512); k_pe (64)]` after `kv_a_layernorm` | +| `value` | absent | +| `key_cache` | `(num_blocks, block_size, 1, 576)` | +| `value_cache` | absent (aliases `key_cache`) | +| `block_table`, `slot_mapping`, `cumulative_sequence_length`, `past_seqlens` | as for any paged node | +| `output` | `(token_count, 128 * 512)`, consumed by the `W_UV`-absorbed output projection | +| `num_heads` / `kv_num_heads` | `128` / `1` | +| `head_size` (derived) | `576` | +| `kv_cache_layout` | `"LATENT"` | +| `v_head_size` | `512` | +| `rotary_offset` | `512` | +| `scale` | **must be set explicitly** (§12.6) | + +### 12.5 Offset (partial) RoPE + +`rotary_dim` is derived from `cos_cache` as today (`2 × cos_cache.shape[1]`). `rotary_offset` selects +where it starts: RoPE covers channels `[rotary_offset, rotary_offset + rotary_dim)` of each head and +channels outside that range are copied through unchanged. Default `0` reproduces current behavior +exactly. For absorbed MLA, `rotary_offset = kv_lora_rank = 512` and `rotary_dim = 64`. + +RoPE must be applied to K **before** the latent is written to the cache, so that the cache holds the +already-rotated `k_pe`. This matches the DeepSeek reference implementation and means RoPE is never +re-applied on cache reads. Exporters may equivalently apply RoPE in the graph and set `do_rotary=0`; +both spellings must produce identical results and both are covered by tests (§17). + +### 12.6 The softmax-scale trap + +DeepSeek's softmax scale is derived from the **pre-absorption** head width: +`scale = mscale² / sqrt(qk_nope_head_dim + qk_rope_head_dim)` = based on `192`, with an additional +YaRN `mscale` factor. It is *not* `1/sqrt(576)`. + +The operator's default (`scale == 0 → 1/sqrt(head_size)`) would therefore silently compute +`1/sqrt(576)` and produce plausible-but-wrong logits — the worst possible failure mode. Rule: +**when `v_head_size` is set and differs from `head_size`, an explicit `scale` attribute is +required**; omitting it is `INVALID_ARGUMENT`. The op refuses to guess. + +### 12.7 Kernel backend + +MLA is the one feature in this document that **cannot** be served by the existing backends: + +- ORT's vendored FlashAttention caps `head_size` at 256 and requires `v_head_size == head_size`, so + it cannot run 576/512. +- The CUTLASS fMHA (MEA) wrapper in ORT is constrained the same way. + +An MLA-capable backend is required. Candidates in order of preference: + +1. **FlashMLA** (DeepSeek, SM90) — purpose-built paged MLA decode with 576/512 and a shared K/V + buffer; it consumes precisely the cache layout proposed in §12.4. +2. **FlashInfer MLA** (SM80+) — wider architecture coverage, also paged. +3. **TensorRT-LLM MLA / XQA-MLA** kernels. +4. **cuDNN SDPA**, where it exposes asymmetric head dimensions over a paged cache. +5. **Unfused reference** — needed regardless, as the correctness oracle and for architectures with no + fused kernel. + +Recommended sequencing: land the unfused reference first (correct on any architecture, unblocks the +test matrix), then integrate exactly one fused backend chosen by target hardware. Adding a second +fused backend before the reference exists makes parity failures undebuggable. + +### 12.8 Non-absorbed prefill + +A prefill with **no cached prefix** does not read the paged cache at all: K and V come entirely from +the current chunk. Running it in absorbed form is correct but inflates QK FLOPs ~4× (512-wide instead +of 128-wide `q_nope`), so the recommended graph shape there is ordinary varlen attention +(`head_size = 192`, `v_head_size = 128`, `kv_num_heads = num_heads`) over the decompressed tensors, +plus a cache write of the **latent** form. + +Because the attended tensors and the cached tensor differ in that case, the write must be decoupled +from attention. Two supported spellings, neither of which needs a new operator: + +1. **Graph-level `ScatterND`** — view the cache as `(num_blocks * block_size, 1, 576)` and scatter + the latent rows using `slot_mapping` as indices. Pure standard ONNX. +2. **`PagedAttention` with `slot_mapping` entirely `-1`** (writes suppressed, §5) alongside the + scatter of option 1. + +Chunked prefill and prefix caching *do* read cached history and must therefore use the absorbed form, +exactly as decode does. The operator itself does not distinguish the three cases — the graph does. + +### 12.9 Interaction with the other features + +| Feature | With MLA | Rationale | +|---|---|---| +| `slot_mapping` (§5) | **Supported** | Orthogonal — write path only, and the latent is written exactly like any K/V row. | +| Quantized cache (§8) | **Supported** | An FP8 latent cache is standard in DeepSeek deployments. `PER_CHANNEL` scale shape becomes `(1, 1, head_size)` since `kv_num_heads == 1`. Because V *is* K, the same bytes are written once with `k_scale` and read back as both, so `k_scale` alone describes the cache and `v_scale` / `v_quant_type` must be unset — a second scale for the same bytes could only disagree. A `PER_CHANNEL` V dequant therefore indexes the scale with the `head_size` stride, reading only its leading `v_head_size` entries. | +| RoPE | **Supported** via `rotary_offset` (§12.5) | | +| `softcap` | Allowed, unused by DeepSeek | | +| Sliding window (§9) | Allowed but untested | No MLA model uses it; semantics are well defined (window is over positions, not channels). | +| `head_sink` (§6) | **Rejected** | No MLA model uses sinks. The LSE epilogue is valid math here, but shipping an untested combination invites silent errors. Revisit if a model needs it. | +| QK-Norm (§7) | **Rejected** | DeepSeek's `q_a_layernorm` / `kv_a_layernorm` act on the *latent* projections in the graph, before absorption. A `head_size`-wide RMSNorm in absorbed space is a different operation; accepting it would let an exporter produce silently wrong math. | +| `attention_bias` / `output_qk` (§10, §11) | Supported on the unfused path | `output_qk` shape is unchanged: `(num_heads, token_count, max_context_len)`. | + +### 12.10 Validation + +- `v_head_size != 0 && v_head_size != head_size` requires `kv_cache_layout == "LATENT"`; otherwise + `INVALID_ARGUMENT`. In `"SEPARATE"` mode `effective_v_head_size == head_size` always. +- `kv_cache_layout == "LATENT"` requires `key` present and `value` and `value_cache` absent (§4.6). +- `v_head_size ∈ [1, head_size]` when set; `0` means "equal to `head_size`". +- `v_head_size != head_size` requires an explicit `scale` attribute (§12.6). +- Initially require `kv_num_heads == 1`; widen only alongside a backend and tests for grouped latent + heads. `kv_num_heads` is `1` for DeepSeek, and any divisor of `num_heads` is accepted once a + backend exists. +- `rotary_offset >= 0`, `rotary_offset % 8 == 0`, `rotary_offset + rotary_dim <= head_size`. +- In `"LATENT"` mode `head_size` may exceed 256, but the selected backend must accept it; otherwise + return `INVALID_ARGUMENT` naming the backend and the supported head widths. +- `head_sink` or QK-Norm weights combined with `"LATENT"` ⇒ `INVALID_ARGUMENT` (§12.9). +- `v_scale` and a non-`"NONE"` `v_quant_type` combined with `"LATENT"` ⇒ `INVALID_ARGUMENT`: there is + one physical cache and `k_scale` already describes it (§12.9). +- Dispatch only to an MLA-capable backend or the unfused reference; never silently ignore an input. + +## 13. Kernel Dispatch and Backend Plan + +Target dispatch order once all phases land, first eligible wins: + +| Priority | Backend | Eligible when | +|---|---|---| +| 0 | **MLA backend** (new, §12.7) | `kv_cache_layout == "LATENT"` — FlashMLA / FlashInfer MLA / TRT-LLM MLA where available, unfused MLA reference otherwise | +| 1 | **Paged decode kernel** (new, Phase 4) | Decode-shaped batch per the static shape test in §4.7 (`token_count <= batch_size`); non-quantized or `PER_TENSOR`/`PER_CHANNEL` INT8/FP8; sliding window and `head_sink` supported | +| 2 | **FlashAttention varlen** | FP16/BF16, SM80+, non-quantized cache (or quantized via dequant-gather); supports sliding window, softcap, packed QKV, `head_sink` via LSE epilogue | +| 3 | **Memory-Efficient Attention (CUTLASS fMHA)** | Fallback for supported combinations and pre-SM80 | +| 4 | **Unfused** | Last resort — arbitrary `head_size`, `attention_bias`, `output_qk` | + +Implementation status: priority 0 is currently served by `PagedLatentAttentionKernel`, the unfused +MLA reference in `paged_attention_impl.cu`. It is selected unconditionally when +`kv_cache_layout == "LATENT"` — the other three backends cannot express `v_head_size != head_size` +over a single aliased cache, so there is no eligibility test to run. One CTA owns one +(query token, query head) pair and streams the KV range in tiles with an online-softmax state, +which makes prefill, chunked prefill and decode take the same path. A fused MLA backend (P5) will +slot in ahead of it under the same priority. + +Backend selection must depend only on values that are constant for a captured CUDA graph (§4.7): +shapes, attributes, and type constraints. It must never depend on a device-resident sequence length. + +Feature × backend matrix (target state): + +| Feature | Paged decode | Flash varlen | MEA | Unfused | +|---|---|---|---|---| +| Sliding window | Yes | Yes | Yes | Yes | +| `softcap` | Planned | Yes | Yes | Yes | +| `head_sink` | Yes (native) | Yes (LSE epilogue) | After fMHA LSE | Yes | +| QK-Norm | Yes (prologue) | Yes (prologue) | Yes (prologue) | Yes | +| Quantized cache | Yes (in-kernel) | Via dequant-gather | Via dequant-gather | Via dequant-gather | +| `query_positions` (§4.8) | Backend work required | Fallback unless legacy-equivalent | Backend work required | Yes | +| `attention_bias` | No | No | No initially | Yes | +| `output_qk` | No | No | Yes | Yes | +| `slot_mapping` | Yes (write path — backend independent) | Yes | Yes | Yes | +| `LATENT` MLA | No | No | No | Yes — plus the dedicated MLA backend (§12.7) | + +Feature acceptance and backend eligibility are separate concerns. Schema validation (§15) decides +whether the *request* is well formed; dispatch then selects a backend that implements the requested +combination. An unsupported combination must produce an `INVALID_ARGUMENT` naming both the feature +and the backend limitation — never a silently ignored input or attribute. + +QK-Norm, RoPE, offset RoPE, packed-QKV unpacking, quantized writes, and `slot_mapping` all live in +the **prologue** and are therefore backend independent. Only the attention math itself varies by +backend. + +The selected backend must be reported through `AttentionKernelDebugInfo` (`SdpaKernel=...`) exactly +as GQA does, so that `ORT_ENABLE_ATTENTION_KERNEL_DEBUG_INFO=1` works uniformly across both ops. + +> **Implementation note (P5, implemented).** Priority 1 is in place and reports +> `SdpaKernel=DECODER_ATTENTION`. Actual eligibility is broader than the table: any `head_size` whose +> working set fits `sharedMemPerBlock` (roughly `2 * head_size + 256` floats plus 512 B, so every +> head size the op supports on current hardware), any `block_size`, any GQA ratio, `softcap`, and +> arbitrary ragged query lengths are all supported. It is only *preferred* over FlashAttention when +> the cache is quantized or FlashAttention is ineligible, since Flash's tensor-core decode path is +> faster on an unquantized cache. Priority 0 (MLA) is still unimplemented. +> +> The gate is the static shape test of §4.7 (`token_count <= batch_size`), so no device-resident +> length reaches the host and the D→H sync is gone. The XQA fast path inside this backend keeps a +> stricter gate — it needs `token_count == batch_size` *and* `max_query_len == 1`, the latter from +> `attention_metadata` or, failing that, from the readback — because its output layout is one row +> per batch index rather than per query token. + +## 14. Shared Code with GroupQueryAttention + +Feature drift between the two ops is the principal long-term risk of keeping them separate. The +mitigation is structural, not procedural. + +1. **Common parameter base.** `PagedAttentionParameters` and `GroupQueryAttentionParameters` already + derive from `AttentionParameters` in `contrib_ops/cpu/bert/attention_parameters.h`. Move the + fields that are genuinely shared — `local_window_size`, `softcap`, + `qk_norm_epsilon`, `k_quant_type`, `v_quant_type`, + `rotary_interleaved` — + into the base so a feature has one definition. `v_head_size` and `rotary_offset` (§12) start in + `PagedAttentionParameters`; promote them if and when a second op needs asymmetric head widths. + +2. **Shared validation helpers.** `paged_attention_helper.h` already calls + `group_query_attention_helper::CheckRotaryCaches`. Extend the same pattern for + `CheckQKNormWeights`, `CheckHeadSink`, and `CheckKVQuantScales`, so both ops enforce identical + shapes and identical error text. + +3. **KV accessor abstraction (the important one).** Introduce a device-side accessor concept: + + ```cpp + struct ContiguousKVAccessor { /* BNSH, (b, h, s, d) -> offset */ }; + struct PagedKVAccessor { /* block_table + block_size, (b, h, s, d) -> offset */ }; + ``` + + Prologue kernels (QK-Norm → RoPE → quantize → store) and any future fused attention kernel are + templated on the accessor. A new feature is then written once and both ops receive it. This is the + only mechanism that makes "separate ops, shared math" hold over time. + +4. **Shared epilogues.** `LaunchApplyHeadSink`, per-channel dequant scaling, and softcap helpers live + in a common `attention_epilogues.cuh` used by both. + +## 15. Validation Rules + +Consolidated, to be implemented in `paged_attention_helper::CheckInputs`. Every violation returns +`INVALID_ARGUMENT` with a message naming the offending tensor and the expected value. + +**Existing (retain):** +- `num_heads % kv_num_heads == 0`; `num_heads <= max_threads_per_block`. +- `query` rank 2; `head_size % 8 == 0`. +- `key`/`value` both present or both absent; `key`/`value` rank 2 with dim 0 == `token_count`. +- Packed QKV: `hidden_size % (num_heads + 2 * kv_num_heads) == 0`. +- `key_cache`/`value_cache` rank 4, identical shapes, dim 2 == `kv_num_heads`, dim 3 == `head_size`. +- `cumulative_sequence_length` rank 1 with `dim0 >= 2`; `past_seqlens` rank 1 with `dim0 == batch_size`. +- `block_table` rank 2 with `dim0 == batch_size`. +- `cos_cache`/`sin_cache` both present or both absent; required when `do_rotary == 1`. +- `key_cache_out` must alias `key_cache`; same for value. +- `batch_size <= 256` (BlockScan limitation — to be lifted, see §18). + +**Corrected:** +- `block_size` must be a power of two in `{16, 32, 64, 128, 256}` (**replaces** `block_size % 256 == 0`; see §18). + *Implemented as: any power of two `>= 16`. Values that FlashAttention cannot address for the given + `head_size` select the memory-efficient backend rather than failing — see the note in §18.1.* + +**New:** +- `slot_mapping`: rank 1, `dim0 == token_count`, `int32`. +- `head_sink`: rank 1, `dim0 == num_heads`, type `T`. +- `q_norm_weight` / `k_norm_weight`: both-or-neither; rank 1, `dim0 == head_size`, type `T`. +- `k_scale` / `v_scale`: FP32; `(1,)` for `PER_TENSOR`, `(kv_num_heads, 1, head_size)` for + `PER_CHANNEL` (both K and V — `v_scale` only exists alongside a `value_cache`, so its last + dimension is always `head_size`); present iff the corresponding quant type is not `NONE`. +- `T_CACHE != T` iff a quant type is not `NONE`. +- `k_cache_dtype` and `v_cache_dtype` must be `""` or name the cache tensor's own element type: + every logical element type this operator stores is expressible as an ONNX element type. The + reserved sub-byte values are rejected until a `uint8` packed cache exists. FP8 availability is + controlled by `onnxruntime_USE_FP8_KV_CACHE`, without an additional runtime architecture gate. +- `attention_metadata`: rank 1, `dim0 == 2`, `int32`, CPU-resident; entries `>= 0`; each non-zero + entry is clamped to its static limit before use and may only size launch dimensions and workspace + (§4.7). It must never enter a mask comparison. Each value is a trusted upper bound for every step + served by the node or captured graph. +- `query_positions`: rank 1, `dim0 == token_count`, `int32`, entries `>= 0`. +- `attention_bias`: rank 4; `dim0 ∈ {1, batch_size}`, `dim1 ∈ {1, num_heads}`, + `dim2 == query_length_capacity`, and `dim3 == context_length_capacity`; both capacities must cover + every actual length served by the node or captured graph. Indexing uses the sequence-local query + offset and logical KV position, matching GQA (§10.2). Device code bounds-checks both indices; + fused backends reject the input and the initial implementation uses the unfused fallback. +- `qk_output != 0` iff the node has 4 outputs; `output_qk` rejected on Flash. +- `kv_cache_layout` must be one of its documented values, and the input presence pattern must match + a row in [§4.6](#46-input-mode-state-machine). +- MLA (§12.10): `kv_cache_layout == "LATENT"` requires `key` present, `value` and `value_cache` + absent, an explicit `scale`, and no `head_sink` or QK-Norm weights. + `v_head_size ∈ [1, head_size]`; `v_head_size != head_size` outside `"LATENT"` is + `INVALID_ARGUMENT`. `rotary_offset % 8 == 0` and `rotary_offset + rotary_dim <= head_size`. +- Backend-incompatible feature combinations must be reported at `Compute` entry with the reason, + never silently ignored. + +**Trust boundary.** `attention_metadata` is host-supplied and cannot be checked against the device +tensors without reintroducing the synchronization it exists to remove. It is therefore a *trusted +bound*: an under-sized value violates the contract and may omit valid attention work. Independently, +the kernel must keep every device read within the static tensor capacities even when the bound is +invalid. No other host-supplied value is permitted to bound a launch or workspace extent. + +## 16. Shape Inference and Tooling + +- `PagedAttentionTypeAndShapeInference` in `onnxruntime/core/graph/contrib_ops/bert_defs.cc`: + - Output 0 dim 1 becomes `num_heads * effective_v_head_size`. **Only compute the product when + `effective_v_head_size != head_size`** (i.e. in `"LATENT"` mode). In every `SEPARATE`-mode model + keep the shipped `propagateShapeFromInputToOutput(ctx, 0, 0)`, which needs no numeric dimension; + deriving `head_size = query.shape[1] / num_heads` unconditionally would make a model with a + symbolic hidden dimension — which infers fine today — start failing. Both the unpacked and + packed-QKV branches must apply the generalization when it does apply, otherwise every MLA graph + gets a wrong — and silently propagated — output width. + - Read `kv_cache_layout` before branching on Q/K/V presence. `LATENT` uses the unpacked-query + formula even though `value` is absent; only `SEPARATE` with both K and V absent uses the packed + formula. Unknown symbolic dimensions must stay symbolic. + - Outputs 1/2 must propagate **type from inputs 3/4** (not from input 0) once `T_CACHE` can differ + from `T`. The current code propagates elem type from input 0 to outputs 1/2 first and then from + 3/4 — the first propagation becomes wrong under quantization and must be removed. Element-type + propagation must run *before* shape propagation for the same output, or ONNX reports + `Mismatch between inferred and declared type`. + - `value_cache` (input 4) is now optional: guard `getInputShape(ctx, 4)` and the output-2 + propagation on `ctx.hasInput(4)` before any write. + - The shipped `getNumOutputs() > 1 ⇒ getNumOutputs() == 3` rule must be relaxed to admit a + four-entry output list with empty names at unused optional cache positions (§4.4), while keeping + the both-or-neither rule for `SEPARATE` mode. + - Output 3 (`output_qk`) is written only when `ctx.getNumOutputs() > 3`, guarded before any write, + consistent with the contrib-op shape-inference memory-safety rules. +- Kernel definition: register `.Alias(3, 1).Alias(4, 2)` and `.InputMemoryType(OrtMemTypeCPUInput, 16)` + (§4.4, §4.7). +- `onnxruntime/python/tools/symbolic_shape_infer.py` — `_infer_PagedAttention` must handle the new + optional inputs and the fourth output. +- Regenerate `docs/ContribOperators.md` and `docs/OperatorKernels.md`. +- Update `onnxruntime/python/tools/transformers/fusion_options.py` and any fusion that emits + `PagedAttention` so new attributes get explicit values. + +## 17. Testing Plan + +### 17.1 GQA ↔ PagedAttention cross-parity harness + +The highest-value test. For a given configuration, construct a block table that maps each sequence's +blocks contiguously and in order, build the equivalent padded GQA inputs, and assert +`PagedAttention` output == `GroupQueryAttention` output within tolerance. This makes GQA the +reference oracle for every shared feature and catches drift automatically. Run it across: + +- `head_sink` on/off × sliding window on/off × `softcap` on/off +- QK-Norm on/off × RoPE on/off × interleaved on/off +- packed QKV vs. unpacked +- quantized cache: `NONE`/`PER_TENSOR`/`PER_CHANNEL` × INT8/FP8 + +### 17.2 Paging-specific tests (no GQA equivalent) + +- `slot_mapping`: explicit mapping equals derived mapping when the mapping is the derived one. +- `slot_mapping` with `-1` entries: skipped tokens leave the cache byte-identical, and attention + output matches a run where those positions were pre-populated (prefix-cache simulation). +- Non-contiguous / shuffled / shared block tables — including two sequences sharing prefix blocks. +- `block_table` entries of `-1` (unmapped) are masked out. +- Ragged batches including sequences with **zero** new tokens. +- `token_count == 0` early-out. +- `batch_size` at and just above the BlockScan limit (must error, not corrupt). +- Sliding-window block pruning: pruned run bit-matches the unpruned run. +- Rolling sliding-window cache (§9.3): a run holding only `ceil(W / block_size) + 1` blocks per + sequence, recycled through `slot_mapping` with the evicted `block_table` entries set to `-1`, + bit-matches the full-cache run at every decode step. + +### 17.3 MLA tests + +GQA cannot be the oracle here — it has no MLA mode — so MLA needs its own reference chain: + +- **Absorbed ↔ non-absorbed equivalence.** For a random `W_UK` / `W_UV`, assert that absorbed-form + `PagedAttention` (576/512, `kv_num_heads=1`) matches a non-absorbed reference (192/128, + `kv_num_heads=num_heads`) built from the decompressed latent. This is the single test that proves + the whole design and must run before any fused MLA kernel is integrated. +- **HuggingFace parity.** Compare one DeepSeek-V2-Lite decoder layer against the reference PyTorch + implementation, including `mscale`/YaRN scaling, over prefill and several decode steps. +- **`rotary_offset`.** In-op offset RoPE (`do_rotary=1`, `rotary_offset=512`) must bit-match the + graph-applied spelling (`do_rotary=0`, RoPE fused into the producer). +- **Scale guard.** Omitting `scale` while `v_head_size != head_size` must fail with a clear error — + a regression here produces plausible-looking but wrong logits (§12.6). +- **V-aliases-K.** With `value_cache` absent, verify that the leading `v_head_size` channels of + `key_cache` supply V. Any present `value_cache` — including the same tensor as `key_cache` — must + be rejected in `LATENT` mode. +- **MLA × paging.** Shuffled block tables, `slot_mapping` with `-1`, and an FP8 latent cache, each + combined with MLA. +- **Rejected combinations.** MLA + `head_sink` and MLA + QK-Norm must fail with the documented + message, not silently compute something. + +### 17.4 Reference implementations + +Extend `onnxruntime/test/python/transformers/test_paged_attention_cuda.py` with a PyTorch reference +covering attention sinks, QK-Norm (RMSNorm-before-RoPE), per-channel dequantization, +windowing, and MLA absorption — reusing the GQA test helpers rather than duplicating them. + +### 17.5 Negative tests + +One test per validation rule in [§15](#15-validation-rules), asserting the error *message*, not just +failure — so that backend-incompatible combinations cannot regress into silent wrong results. + +### 17.6 Compatibility and CUDA graph regression + +These two guard the contract itself rather than a feature. + +- **Opset-1 compatibility.** Serialize a model using *only* the shipped contract — inputs 0–9, + outputs 0–2, the seven original attributes — and run it against the extended kernel **without + rewriting the graph**. Assert byte-identical output against a run on the pre-extension build. This + must be re-run in every phase (§19), not just the phase that adds an input. +- **CUDA graph decode replay.** Capture a decode graph, then replay it for `N` steps while the KV + length grows, and assert the result matches an uncaptured run step for step. This is the test that + would have caught the frozen-`max_kv_len` failure mode in §4.7: a naive implementation passes step + 1 and diverges from step 2 onward, so the test must check **every** step, not just the last. +- **Metadata-bound invariance.** The same decode sequence must produce identical results with + `attention_metadata` absent, with the tightest replay-wide valid bounds, and with deliberately + loose (over-large) bounds. + Any difference means a host value affected device-side logical attention bounds rather than only + launch or workspace capacity (§4.7). +- **Under-sized bound safety.** A deliberately too-small `max_kv_len_bound` violates the contract, + so its numerical result is unspecified, but it must not read out of bounds; run it under + `compute-sanitizer`. + +## 18. Known Defects to Fix First + +These block the feature work and should land ahead of it. + +1. **`block_size % 256 == 0` is wrong.** `paged_attention_helper.h::CheckKVCache` requires the block + size *in tokens* to be a multiple of 256. Conventional paged block sizes are 16/32/64/128 tokens; + 256-token blocks defeat the fine-grained allocation that paging exists to provide, and a 16-token + block table (the vLLM default) is rejected outright. The in-code `TODO(aciddelgado): block size + multiple of 8` suggests the intent was an alignment constraint on the *innermost* dimension. + Replace with `block_size ∈ {16, 32, 64, 128, 256}` and, if a byte-alignment constraint is truly + needed, express it on `block_size * head_size * sizeof(T_CACHE)`. + + > **Implementation note (correction).** The constraint is *not* purely a validation bug. The + > vendored FlashAttention split-KV kernel builds each `gK`/`gV` tile as a single contiguous + > `kBlockN × head_size` region addressed by one `(block_table_idx, block_table_offset)` pair, so a + > tile must never straddle a page. That requires `block_size % kBlockN == 0`, where + > `kBlockN = head_size <= 64 ? 256 : (head_size <= 128 ? 128 : 64)` + > (`flash_fwd_launch_template.h::run_mha_fwd_splitkv_dispatch`). Relaxing the *validation* alone + > would produce silent garbage on the Flash path. + > + > What was implemented instead: `CheckKVCache` accepts any power-of-two `block_size >= 16` (a + > superset of `{16, 32, 64, 128, 256}` that also keeps the existing `block_size = 512` test case + > valid), and `paged_attention.cc` treats `block_size % kBlockN == 0` as part of Flash backend + > *eligibility*. When a model uses a smaller page than Flash can address, the op transparently + > falls back to the memory-efficient backend, which gathers pages into a dense buffer first and + > therefore accepts any block size. The op only errors when neither backend is eligible. + > Lifting this properly requires teaching the Flash paged loader to split a tile across pages. +2. **Out-of-bounds binary search.** The binary search over `cumulative_seqlens_q` in + `ReshapeAndCache` and `GatherAndExpandPagedKVCache` can yield `batch_id == batch_size` when + `token_id >= cumulative_seqlens_q[batch_size]`, producing OOB reads of `past_seqlens` and + `block_table`. Guard with an early `return` when + `token_id >= cumulative_seqlens_q[batch_size]`. (`slot_mapping` removes the search entirely on the + write path, but the gather path still needs the fix.) +3. **`batch_size <= 256`.** Replace the per-block `cub::BlockScan` with a grid-wide scan (or a single + 256-thread block doing a strided serial scan) so continuous batching is not capped at 256 + concurrent sequences — a real limit for a serving op. +4. **Per-step D→H synchronization — and it blocks CUDA graph capture.** ~~`max_query_len` (and + `total_kv_tokens` for MEA) are obtained via `cudaStreamSynchronize` every step, once per layer. + This is not only a throughput bug for the op's primary use case: `cudaStreamSynchronize` on a + capturing stream is **illegal**, so the operator cannot be captured into a CUDA graph at all — + precisely the optimization decode needs most.~~ **Fixed.** Dispatch now derives from static + shapes, extents from the static capacity bound `block_table.shape[1] * block_size` or the + optional `attention_metadata` bounds, and every per-step quantity from device memory, as + specified in [§4.7](#47-cuda-graph-contract-and-attention_metadata). The sync survives only as a + prefill-side fallback for a dense gather or for XQA when no metadata is supplied, neither of + which a capturable step can reach. + +## 19. Phasing + +| Phase | Contents | Schema delta | +|---|---|---| +| **P0 — Foundation** | §18 defect fixes; `.Alias(3,1).Alias(4,2)` on the kernel def (§4.4); sliding-window semantic verification and GQA parity harness (§17.1) | none | +| **P1 — Paging primitives** | `slot_mapping` (§5); sliding-window block pruning (§9.2); rolling sliding-window cache (§9.3) | input 10 | +| **P2 — Model coverage** | `head_sink` via LSE epilogue (§6); fused QK-Norm (§7) | inputs 11–13, attr `qk_norm_epsilon` | +| **P3 — Memory** | Quantized cache INT8/FP8, `PER_TENSOR` + `PER_CHANNEL`, dequant-on-gather read path (§8) | `T_CACHE`, `T_KV_SCALE`, inputs 14–15, 4 attrs | +| **P4 — MLA (correctness)** | `kv_cache_layout="LATENT"`, `v_head_size`, `rotary_offset`, V-aliases-K, optional `value_cache`, unfused MLA reference kernel, absorbed↔non-absorbed equivalence tests (§12) | attrs `kv_cache_layout`, `v_head_size`, `rotary_offset`; input 4 optional | +| **P5 — Performance** | Paged decode kernel with in-kernel dequant; fused MLA backend (FlashMLA / FlashInfer MLA, §12.7); `softcap` on decode; **remove the D→H sync and make the op CUDA-graph-capturable (§4.7)**; optional `attention_metadata` replay-wide bounds | input 16 | +| **P6 — Completeness** | `query_positions` (§4.8); `attention_bias` (§10); `output_qk` (§11) | inputs 17–18, output 3, attr `qk_output` | +| **Later** | INT4 cache; MLA quantized latent cache tuning; non-CUDA EPs | — | + +Status: P0–P4 are implemented, except the `.Alias` registration. P5 is partially +implemented — the paged decode kernel with in-kernel dequantization (including `softcap`, sliding +window and `head_sink`) has landed (§8.5, §13), as has the sync removal and CUDA graph capture +(§4.7); the fused MLA backend has not. + +Every phase must re-run the old-model tests with all new inputs absent, in addition to the +feature-specific tests. The compatibility regression test should serialize a model using only the +shipped opset-1 contract and run it against the extended kernel **without rewriting the graph** +(§4.2). + +P0–P2 cover GPT-OSS, Qwen3, Gemma 2/3, Llama, Mistral and Phi. P3 and P5 are what make the op +competitive for throughput-oriented serving. P4–P5 add DeepSeek-V2/V3/R1. + +P4 is deliberately split from P5: the schema work and the unfused reference are small, low-risk, and +unblock the test matrix, whereas integrating a fused MLA kernel is a large dependency decision +(§12.7, §20). Shipping the reference first means the fused kernel arrives with an oracle already in +place. + +Note that MLA is folded into `PagedAttention` rather than given its own operator precisely because +the *only* differences are two attributes and a K/V aliasing rule (§12.2) — the layout, batching +model, and cache management are identical. That is the opposite of the GQA-vs-PagedAttention case in +§1, where the query rank and the cache contract genuinely differ. + +## 20. Open Questions + +1. ~~**Host-scalar inputs for `max_query_len` / `total_kv_tokens` (§18.4).** Adding them as optional + inputs removes a per-step sync but leaks scheduler state into the graph. Is that acceptable?~~ + **Resolved in [§4.7](#47-cuda-graph-contract-and-attention_metadata):** no host input is needed to + remove the sync, and an *exact* one would be actively unsafe — under CUDA graph replay no kernel + `Compute` runs, so a host value is frozen at capture while `max_kv_len` grows every step. Dispatch + comes from static shapes, extents from the static capacity bound, and every per-step quantity from + device memory. `attention_metadata` survives only as optional replay-wide bounds: valid bounds + preserve correctness, while an under-sized bound is a producer contract violation. +2. **`block_table == -1` semantics (§9.2).** Confirm "unmapped, masked out" rather than "invalid, + error" — the former is required for sliding-window block eviction. +3. **Non-CUDA EPs.** Is `PagedAttention` a server-GPU-only op (CUDA, ROCm, TensorRT), or is a CPU + reference implementation required? A CPU reference has real value as a test oracle even if it is + never used in production. Stub kernels returning `NOT_IMPLEMENTED` are not an acceptable middle + ground (§3.4). +4. **ORT-GenAI commitment.** Does the continuous-batching path in `onnxruntime-genai` commit to + emitting `PagedAttention` (as opposed to a `block_table`-extended GQA)? This determines the + priority of everything above. +5. **Ownership of the parity matrix (§13).** Who keeps the feature × backend table current, and is + the KV-accessor refactor (§14.3) in scope for P2 or a follow-up? +6. **Which fused MLA backend (§12.7)?** FlashMLA is the closest fit but is SM90-only and adds a + third-party dependency; FlashInfer MLA covers SM80+; TRT-LLM has its own. The choice determines + the hardware matrix ORT can serve DeepSeek on and should be made before P5 starts. +7. **Does MLA need a non-absorbed *paged* path (§12.8)?** The proposal handles prefill-with-prefix by + absorbing. If a workload shows the absorbed prefill FLOP inflation to be material, an in-op + decompression path would require passing `W_UK` / `W_UV` into the operator — which this design + deliberately avoids. Needs a measurement before it is reconsidered. +8. **MLA + quantized latent cache (§12.9).** FP8 on a 576-wide latent shared by 128 query heads has + a different error profile from FP8 on a per-head cache. Accuracy validation is required before + the combination is recommended, even though the schema supports it on day one. + +## 21. Deferred Breaking Contract + +Status: **deferred, not adopted.** [§4](#4-schema--the-compatible-contract) is the only normative +contract. This section records the breaking cache-format ideas that were evaluated, states why each +was rejected for opset 1, and preserves the analysis so it does not have to be redone if a merged or +sub-byte cache ever becomes a concrete requirement. + +### 21.1 Why they are deferred + +The original argument was that three of these changes cannot be expressed additively at any later +date, so they must land before the operator acquires an external emitter or never: + +| Change | Why it cannot be additive | +|---|---| +| Merge `key_cache` + `value_cache` → `kv_cache` | Changes input arity and the aliasing contract | +| `kv_cache_out` becomes required | Changes output arity | +| Positions no longer implicit in `past_seqlens[b] + j` | Changes the meaning of an existing graph | + +The first two hold. The third does not: [§4.8](#48-query_positions) shows that an optional +`query_positions` input expresses explicit positions additively, because absence keeps the legacy +derivation. That removes the only *model-coverage* item from the "now or never" list, and what +remains is physical cache-format cleanup with no feature behind it. + +Weighed against that, the operator is already serialized as `com.microsoft::PagedAttention` opset 1, +and P0–P3 plus half of P5 are implemented and tested against that contract. Breaking it buys +expressibility for formats no ORT model uses today, at the cost of invalidating every already +serialized graph and every test. The decision is therefore: + +> Treat the separate-cache representation as **permanent** for `com.microsoft::PagedAttention` +> opset 1. If a merged or sub-byte cache becomes a real requirement, introduce a separately versioned +> schema or a new operator name with a migration tool — do not change the meaning of inputs, outputs +> or attributes in place. + +The complete deferred list: + +- one merged `kv_cache` input containing K and V (§21.2); +- one required functional `kv_cache_out` instead of two optional aliasing outputs; +- removal of `kv_num_heads` in favor of `kv_cache.shape[2]`; +- quantization granularity inferred from scale shape, and zero points (§21.3); +- sub-byte logical types stored in `uint8` tensors — the `k_cache_dtype` / `v_cache_dtype` attributes + that name them are adopted in §4.5, but no backend decodes a packed cache yet (§21.4); +- inline scales or zero points packed into cache rows (§21.3, note); +- a physical `HND` cache layout (§21.6); +- renaming `local_window_size` to `window_size_left` / `window_size_right`, and the + lookahead window (`window_size_right > 0`) that tree speculative decoding needs; +- a tree-attention ancestry mask, which `query_positions` deliberately does **not** provide (§4.8). + +Everything else that was in the v2 proposal — `attention_metadata`, `query_positions`, +`kv_cache_layout`, `v_head_size`, `rotary_offset`, explicit quantization attributes — is expressible +additively and has been folded into §4. + +--- + +### 21.2 Single `kv_cache` with a layout attribute + +``` +kv_cache : (num_blocks, block_size, kv_num_heads, kv_pack_dim) +``` + +| `kv_cache_layout` | `kv_num_heads` | `kv_pack_dim` | K slice | V slice | +|---|---|---|---|---| +| `"KV_CONCAT"` | `H` | `head_size + v_head_size` | `[0, head_size)` | `[head_size, +v_head_size)` | +| `"LATENT"` (MLA) | `1` | `kv_lora_rank + qk_rope_head_dim` (576) | `[0, kv_pack_dim)` | **`[0, v_head_size)`** | + +`"KV_CONCAT"` covers MHA, GQA, and the asymmetric-head-size case (`v_head_size != head_size`) in one +rule. `"LATENT"` is the MLA case where the V view *overlaps* the K view rather than following it. + +This is the layout vLLM's FlashAttention, FlashInfer and Triton backends converged on +(`(num_blocks, num_kv_heads, block_size, 2 * head_size)` in their head-major ordering), and +TensorRT-LLM parameterizes the same axis as `kv_factor ∈ {1, 2}` +(`kv_cache_manager_v2.py`, `CacheType::kSELFKONLY`). A hard-coded `2` axis — the shape this document +previously implied — cannot express either MLA or asymmetric K/V and is rejected. + +**Derived, not attributes.** `k_head_size` is always `query.shape[-1] / num_heads`, including for +MLA where `query` is the 576-wide nope‖rope concatenation. `v_head_size` is derivable as +`kv_pack_dim - k_head_size` under `"KV_CONCAT"` but **must** be given under `"LATENT"` (it is +`kv_lora_rank`; 512 cannot be recovered from `kv_pack_dim = 576` and `k_head_size = 576`). + +**No explicit `k_offset` / `v_offset`.** They add nothing a named layout does not already determine, +and they cannot express the part that actually matters: the *write* path differs between the two +layouts. Under `"KV_CONCAT"` both `key` and `value` are supplied and scattered into disjoint regions +of the row. Under `"LATENT"` there is no separate V to write — `value` is **absent** and the V region +aliases bytes already written by the K store. With raw offsets an implementation would have to infer +"the ranges overlap, therefore skip the V write," which is implicit, easy to get wrong, and gives +shape inference no way to reject `value` being present. Named layouts also require no range/overlap +validation, and can grow new members later (e.g. an alignment-padded or scale-interleaved variant) +without having locked the operator into a byte-layout contract. + +This makes §12.2's "V aliases K" a layout selection rather than a schema fork. §4 achieves the same +effect additively with `kv_cache_layout="LATENT"` over the *separate* caches, at the cost of not +covering the asymmetric-K/V and packed-quant cases — which is why those are restricted to `LATENT` +in §12.3. + +**Performance is not the motivation.** There is no meaningful kernel-level gain from merging the two +pools for non-MLA models — vLLM splits the merged tensor back into two strided views before every +kernel call. The merge buys *expressibility* (MLA, asymmetric K/V, packed quant formats) and one +fewer aliased graph edge. The KV-transfer and allocator-packing benefits that motivate it in vLLM +and TensorRT-LLM do not apply to ORT, which does not do disaggregated prefill and does not own the +block pool. That thin margin is what makes the merge deferrable. + +### 21.3 Quantization: granularity from scale shape + +`k_quant_type` / `v_quant_type` would be removed; granularity read off the scale tensor instead: + +| `k_scale` shape | Granularity | +|---|---| +| absent | not quantized | +| `()` or `(1,)` | per-tensor | +| `(kv_num_heads, head_size)` | per-channel | +| `(num_blocks, block_size, kv_num_heads)` | per-token | +| `(num_blocks, kv_num_heads)` | per-block | + +One source of truth instead of two that can disagree, and it extends to per-token/per-block without +new enum values. This also drops the vestigial middle `1` in v1's `(kv_num_heads, 1, head_size)`. +K and V may use different granularities independently. + +`k_zero_point` / `v_zero_point` are new and optional, same shape as the corresponding scale, for +asymmetric integer quantization. Absent ⇒ symmetric (the current behavior). These are the inputs an +**unsigned** logical cache type would need: §8.3.1 keeps `k_cache_dtype` / `v_cache_dtype` restricted +to signed, zero-symmetric types precisely because they do not exist yet, and §8.3.1 also shows that +adding them is not free — a non-zero zero point breaks the decode kernel's scale folding and +introduces correction terms in both the QK and PV products. + +> Note: vLLM's DeepSeek V3.2/V4 formats store scales **inline in the cache row** rather than in a +> separate tensor (`fp8_ds_mla` = 512 NoPE + 16 scale + 128 RoPE = 656 B; DeepSeek-V4 = 448 + 128 + +> 8 = 584 B), so one page read fetches data and scale together. The shape-derived rule above covers +> the separate-tensor case only. If an inline format is ever needed it should be added as a +> `kv_cache_layout` member, not by overloading the scale inputs. + +### 21.4 `k_cache_dtype` / `v_cache_dtype` for sub-byte caches + +**The attributes themselves are adopted in §4.5**; only their sub-byte *values* are deferred, because +no backend decodes a packed cache yet. They were adopted rather than deferred because the obvious +alternative — a `k_cache_bit_width` / `v_cache_bit_width` pair — is redundant against the cache +tensor's element type for every format that exists today and still insufficient for the format it +was meant to describe. + +For `int8` and `float8e4m3fn` each cache tensor's own element type is the logical type and its +corresponding cache-dtype attribute stays `""`. Sub-byte needs more: + +- ORT's CUDA EP has no usable 4-bit tensor element type, so storage must be `uint8`. +- `kv_pack_dim` then counts **storage bytes**, and the logical head width is unrecoverable. +- Bit width alone is insufficient: `int4` and `float4e2m1` are both 4 bits with entirely different + decode math. Keeping K and V independent supports formats where they use different precisions. + +| `k_cache_dtype` / `v_cache_dtype` | Storage elem type | Packed logical width | +|---|---|---| +| `""` | corresponding cache tensor type | unchanged | +| `"float16"`, `"bfloat16"`, `"int8"`, `"float8e4m3fn"` | the same type, named explicitly | unchanged | +| `"int4"`, `"float4e2m1"` | `uint8` | logical width / 2 | +| `"int2"` | `uint8` | logical width / 4 | + +where `E = head_size + v_head_size` under `"KV_CONCAT"`, or `kv_pack_dim`'s logical width under +`"LATENT"`. Packing order must be specified or implementations will diverge: **logical element `2i` +occupies the low-order bits of byte `i`**, element `2i+1` the high-order bits. The storage type is +`uint8` for all of these, but the *logical* values stay signed: a 4-bit code is written as `q + 8` +and read back as `nibble - 8`, matching MLAS's `S4` packing, so the dequantization remains +`q_signed * scale` with no zero point (§8.3.1). `"uint4"` / `"uint2"` are not members and cannot be +until `k_zero_point` / `v_zero_point` (§21.3) exist. + +### 21.5 Tightened / clarified + +- **`kv_cache_out` is required.** As an optional output it is a dead-code-elimination hazard: the + cache mutation is a side effect, so a graph in which the output is unconsumed is legal but the node + is not removable. Making it required removes the ambiguity. §4 keeps the optional outputs and + instead recommends registering the alias on the kernel def (§4.4). +- **`window_size_left` / `window_size_right`** replace `local_window_size`, matching FlashAttention's + own parameter names and removing the current off-by-one ambiguity (`local_window_size - 1` is what + actually reaches the kernel). `window_size_right = 0` expresses causal; `> 0` expresses the + bidirectional-lookahead window that speculative decoding needs. §4.5 keeps `local_window_size` + with its established "W positions including the current token" semantic; the lookahead window is + deferred with the rename. +- **`block_table == -1`** means unmapped and fully masked, not an error (§9.2, §20.3). + *Adopted in §4.3, not breaking.* + +### 21.6 `kv_layout` + +`"NHD"` (default) is `(num_blocks, block_size, kv_num_heads, kv_pack_dim)`; `"HND"` is +`(num_blocks, kv_num_heads, block_size, kv_pack_dim)`. + +NHD is the default and the only value implemented, because ORT's Flash backend requires it: +`flash_api.h` documents the paged `kcache`/`vcache` as +`num_blocks x page_block_size x num_heads_k x head_size`, and `paged_attention_impl.cu` hands the +cache straight to `flash::mha_varlen_fwd`. NHD also keeps the per-channel scale index trivial (the +innermost dimension is the channel). + +HND has real but modest advantages on the read path — a head's page tile is one contiguous run, +which helps alignment/vectorization when `head_size * sizeof(T)` is not a multiple of 16, and suits +1-D TMA bulk copies. It is worse on the write path, where NHD gives one contiguous store per token. +Some kernels require one or the other (vLLM's AITER assembly path needs independently contiguous K +and V; trtllm-gen backends report `get_required_kv_cache_layout() == "HND"`), which is why the +attribute exists at all. **For MLA the two are identical**, since `kv_num_heads == 1`. + +### 21.7 Migration sketch, if a versioned successor is ever needed + +| opset 1 (§4) | Successor | +|---|---| +| `key_cache`, `value_cache` (inputs 3, 4) | single `kv_cache` (input 3) with `kv_pack_dim = 2 * head_size` | +| `key_cache_out`, `value_cache_out` | single required `kv_cache_out` | +| `kv_num_heads` attr | drop — read `kv_cache.shape[2]` | +| `local_window_size = w` | `window_size_left = w`, `window_size_right = 0` | +| `k_quant_type` / `v_quant_type` | drop — infer from `k_scale` / `v_scale` shape | +| `k_scale` shape `(kv_num_heads, 1, head_size)` | `(kv_num_heads, head_size)` | + +Implementation impact is concentrated in `ReshapeAndCache` (one packed row per token instead of two +scattered writes), the cache-read stride arithmetic in the P5 decode kernel (innermost stride becomes +`kv_pack_dim`, and the V base gains `k_head_size`), and the two K/V pointer derivations handed to +Flash and MEA, which become strided views over one tensor. + +Every row above is mechanical and scriptable, which is the other reason none of it is urgent: a +migration tool over serialized graphs is cheap compared with breaking a shipped contract in place. + +### 21.8 Disposition + +| Change | Compatible extension? | Disposition | +|---|---|---| +| §21.2 merged `kv_cache` | **No** | Deferred to a versioned successor | +| §21.5 required `kv_cache_out` | **No** | Deferred; §4.4 registers the alias instead | +| §21.3 scale-shape granularity, zero points | Yes | Deferred — explicit attributes in §4.5 are preferred while only two granularities exist | +| §21.4 `k_cache_dtype` / `v_cache_dtype` | Yes | **Adopted** — §4.5; only the sub-byte *values* wait for a packed-cache backend | +| §21.6 `kv_layout` | Yes | Deferred until a backend requires `HND` | +| §21.5 `window_size_*` rename | Yes, with deprecated aliases | Deferred with the lookahead window | +| `attention_metadata` | Yes | **Adopted, redesigned** — §4.7 | +| `query_positions` | Yes | **Adopted** — §4.8 | +| `kv_cache_layout`, `v_head_size`, `rotary_offset` | Yes | **Adopted** — §4.5, §12 | diff --git a/docs/contrib_ops/cuda/qmoe_gemv_experiments.md b/docs/contrib_ops/cuda/qmoe_gemv_experiments.md index 39e3020c75232..3c10d34f4161c 100644 --- a/docs/contrib_ops/cuda/qmoe_gemv_experiments.md +++ b/docs/contrib_ops/cuda/qmoe_gemv_experiments.md @@ -1399,3 +1399,239 @@ cd ~/onnxruntime/build/cu130/Release ``` Result: graph transformer tests `3 passed`; provider test `1 passed`. + +## 2026-06-21: NVFP4 GEMV Packed E2M1 Dequantize (`prmt` Quad Decode) + +### Change Under Test + +- Scope: `Fp4I2FConverter` in + `onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemv/details.h`, i.e. every MXFP4 / + NVFP4 QMoE GEMV kernel instantiated from `moe_gemv_fp4.cu` (`moe_gemv_kernel` + and `moe_gemv_interleaved_swiglu_kernel`). INT4/INT8 use `I2FConverter` and are + untouched. +- `Fp4I2FConverter::convert()` previously decoded one E2M1 code at a time. + Even though `decode()` was already branchless (a `prmt.b32` byte-select into a + packed magnitude table), the surrounding per-element mask / shift / or / pack + sequence dominated: ~53 ALU instructions per 8 codes. +- New packed path (`decode_quad`) taken whenever `N % 8 == 0`, which covers both + live FP4 configurations (`ColumnMajor` `StepK=8` used by NVFP4, and the opt-in + `ColumnMajorInterleaved` `StepK=32` used by MXFP4): + - Load the weights a whole 32-bit word (eight codes) at a time. + - `mag = w & 0x77777777` keeps the three magnitude bits per nibble and clears + bit 3 so `prmt` stays in byte-select mode instead of sign-replicate mode. + - `sgn = (w >> 3) & 0x11111111` becomes `0x00`/`0x80` sign bytes with one + `prmt`. + - One further `prmt` performs **four** magnitude table lookups at once, and one + or two more expand the bytes into two packed `half2` / `bfloat16x2` words. +- Numerically this is a pure instruction-count change: the magnitude constants + are the same exact `half`/`bf16` encodings of `{0, 0.5, 1, 1.5, 2, 3, 4, 6}`, + and nibble `i` still maps to output element `i`, so `pack_to_vec2`/`mma` need + no changes. There is no env var; the packed path is unconditional on device. + +### Bit-Exactness + +A standalone harness compared `decode_quad` against the per-element reference for +both `half` and `__nv_bfloat16`. The reachable input space is only 2^16 wide (a +`decode_quad` call is fully described by its 4x3-bit magnitude selector plus its +4x1-bit sign selector), so the sweep of 256x256x64 randomized 32-bit words covers +every reachable pattern -- this is an exhaustive check, not a sample: + +``` +cuda=no error mismatches=0 +``` + +### SASS Instruction Count (sm_90, CUDA 13.0) + +Same object (`moe_gemv_fp4.cu.o`), built with and without the change: + +| Kernel (fp16, `ColumnMajor`, CtaN=8, Threads=128, GroupSize=16) | Before | After | Delta | +|---|---:|---:|---:| +| `moe_gemv_interleaved_swiglu_kernel` (fc1) | 1088 | 760 | -30.1% | +| `moe_gemv_kernel` (fc2) | 960 | 640 | -33.3% | +| whole object, all FP4 instantiations | 300552 | 180112 | -40.1% | + +### Repro Notes + +Model: `qwen3.6_nvfp4_fp8dense_statscale_mtp_int8head_fp8kv_v10_w4` +(hidden 2048, inter 512, 256 experts, top_k 8, NVFP4 block_size 16, 40 layers), +H200 SXM, single GPU. + +```bash +source ~/git/venv/bin/activate +export CUDA_HOME=/home/tianlei/cuda13.0 CUDNN_HOME=/home/tianlei/cudnn_9.23_cuda13 +export LD_LIBRARY_PATH=/home/tianlei/ort_home_cu130_fp4_bench/lib:$CUDA_HOME/lib64:$CUDNN_HOME/lib:$LD_LIBRARY_PATH +export CUDA_VISIBLE_DEVICES=0 ORT_ENABLE_FP4_GEMV=1 ORT_MTP_DIRECT_ARENA_COMMIT=1 ORT_FP4_GEMV_AUTOTUNE=0 +# per-kernel timing: CUDA graph OFF so nsys reports real per-launch durations +nsys profile -t cuda --capture-range=cudaProfilerApi --capture-range-end=stop \ + -o /tmp/qmoe --force-overwrite true --export=sqlite \ + python scripts/h200_18/profile_mtp_decode.py $MODEL $HEAD off 3 20 10 +nsys stats --report cuda_gpu_kern_sum --format csv /tmp/qmoe.sqlite | grep moe_gemv +# end-to-end: CUDA graph ON +python scripts/h200_18/profile_mtp_decode.py $MODEL $HEAD on 3 200 50 +``` + +Caveat that cost time here: this workspace has **three** copies of +`libonnxruntime_providers_cuda.so` (`onnxruntime/capi`, `onnxruntime_genai`, and +the `ort_home` lib dir). A/B swaps must replace all three, otherwise both arms +silently load the same library and report identical numbers. + +### Nsight Systems Per-Kernel Results (CUDA graph OFF, 800 launches each) + +Interleaved base/new reps, `ORT_FP4_GEMV_AUTOTUNE=0` (shipping default tiling): + +| Kernel | Before (us) | After (us) | Delta | +|---|---:|---:|---:| +| `moe_gemv_interleaved_swiglu_kernel` (fc1) | 33.14 / 33.24 | 26.34 / 26.13 | **-20.9%** | +| `moe_gemv_kernel` (fc2) | 30.21 / 30.19 | 22.22 / 22.20 | **-26.5%** | + +40 launches of each per decode step, so about `-0.60 ms/step` of GPU time. + +### Model-Level Decode Benchmark (CUDA graph ON, 200 steps, 50 warmup) + +Interleaved reps; `ms/step` is the clean metric because `tok/s` also moves with +the MTP acceptance rate. + +| Rep | Before ms/step | After ms/step | +|---|---:|---:| +| 1 | 11.663 | 11.061 | +| 2 | 11.572 | 11.032 | +| 3 | 11.594 | 11.035 | +| mean | 11.610 | 11.043 | + +**-4.9% step time (+5.1% step rate).** Every "after" rep beat every "before" rep. + +### Validation + +```bash +cd onnxruntime/test/python/transformers +ORT_ENABLE_FP4_GEMV=1 python -m pytest -q test_qmoe_nvfp4_cuda.py # 22 passed +ORT_ENABLE_FP4_GEMV=0 python -m pytest -q test_qmoe_nvfp4_cuda.py # 22 passed +``` + +### Decision + +- Keep the change. It is bit-exact, has no env gate, and is the largest single + QMoE NVFP4 GEMV win measured so far. +- After this change the kernels are no longer dominated by dequantize. The next + FP4 GEMV lever is **tiling**, not the converter. + +### Follow-Up: Tiling Default Is Still Wrong For These Shapes + +`moe_gemv_fp4.cu` defaults to `Threads=128` with `StepK=8`, so `CtaK = 1024`. +For fc2 (`k = 512`) half the threads in every block do no work at all. Enabling +the existing autotuner picks `CtaN=8, Threads=64` for **both** GEMVs: + +| Config | fc1 (us) | fc2 (us) | ms/step | +|---|---:|---:|---:| +| `ORT_FP4_GEMV_AUTOTUNE=0` (default) | 26.2 | 22.2 | 11.028 / 11.004 | +| `ORT_FP4_GEMV_AUTOTUNE=1` | 23.9 | 17.5 | 10.787 / 10.751 | + +Another `-0.24 ms/step` (`-2.2%`) is available. `ORT_FP4_GEMV_AUTOTUNE` is off by +default because it synchronizes the inference stream, and it is skipped during +CUDA-graph capture, so the fix should be an analytic default (at minimum, drop to +`Threads=64` when `StepK * Threads > k`) rather than relying on the autotuner. +Resolved by the next section. + +## 2026-07-28: Analytic Default Tiling For The FP4 GEMV + +### Change Under Test + +Follow-up to the previous section, which showed the shipping FP4 GEMV tiling +(`CtaN=8, Threads=128`) is not the best choice for the NVFP4 decode shapes but +that the only way to get the better one was `ORT_FP4_GEMV_AUTOTUNE=1`. The +autotuner is off by default because it synchronizes the inference stream, and it +is skipped during CUDA-graph capture, so in practice it never runs in the +configuration that matters. + +New `gemv::Fp4MoeGemvDefaultConfig(expanded_num_rows, n, k)` in +`moe_gemv_fp4.cu` derives the tiling from the shape instead. `moe_quantization.cc` +now seeds `fc1_config` / `fc2_config` from it rather than from `kDefault`; the +autotuner, when enabled, still overrides it and the per-shape cache is unchanged. +Only `Threads` is derived. `CtaN` stays at `kDefaultCtaN`, so the analytic config +never changes which shapes `is_moe_gemv_fp4_supported` accepts. Note that +`Threads` is a tiling knob but **not** a bit-exact one: a block walks K in strides +of `CtaK = StepK * Threads` and the epilogue reduces across `Threads / 32` warps, +so changing it changes the floating-point summation order and the low bits of the +result can move. + +Two clauses, both about the fact that the grid is `(expanded_num_rows, n / CtaN)` +and therefore does **not** depend on `Threads`: + +- **(a) Idle threads.** A block walks K in strides of `CtaK = StepK * Threads` + (`StepK = 8` for the ColumnMajor FP4 layout). When `CtaK > k` the tail of every + block never enters the K loop. NVFP4 fc2 has `k = inter_size = 512` against + `CtaK = 1024`, i.e. half of every block does nothing. `k < 1024 -> Threads=64`. +- **(b) Epilogue width.** MAC work per block is fixed by `(CtaN, k)` regardless of + `Threads`, but the epilogue reduces partials across `Threads/32` warps through + shared memory. A narrower block does the same math with half the barriers and a + shallower reduction tree. This is gated on the grid being large enough that the + SM still fills: these kernels use ~72 registers/thread on sm_90 (~900 resident + threads/SM), which a 128-thread block reaches with ~7 resident blocks and a + 64-thread block only with ~14, so the threshold is **16 blocks/SM** + (`expanded_num_rows * (n / CtaN) >= 16 * multiProcessorCount`). + +`ORT_FP4_GEMV_DEFAULT_TILING=0` restores the fixed `kDefault` tiling. + +### Selected Configs + +Qwen3.6-35B-A3B-NVFP4 MTP decode (hidden 2048, inter 512, 256 experts, top_k 8, +expanded 32), H200, 132 SMs: + +| GEMV | n | k | blocks | Clause | Chosen | +|---|---:|---:|---:|---|---| +| fc1 (swiglu) | 1024 | 2048 | 4096 (31/SM) | (b) | `Threads=64` | +| fc2 | 2048 | 512 | 8192 (62/SM) | (a) | `Threads=64` | + +Single-token decode (expanded 8) keeps `Threads=128` for fc1 (1024 blocks, 7.8/SM, +below the clause (b) threshold) and still takes `Threads=64` for fc2 via clause +(a), which is the intended behavior. + +The analytic choice reproduces the autotuner's pick exactly: + +| Kernel | `AUTOTUNE=1` (us) | Analytic default (us) | +|---|---:|---:| +| `moe_gemv_interleaved_swiglu_kernel` (fc1) | 23.93 | 23.89 / 24.05 | +| `moe_gemv_kernel` (fc2) | 17.46 | 17.48 / 17.52 | + +### Model-Level Decode Benchmark (CUDA graph ON, 200 steps, 50 warmup) + +Same binary for both arms, toggled with `ORT_FP4_GEMV_DEFAULT_TILING`, interleaved +reps. Final threshold (16 blocks/SM): + +| Rep | Fixed default ms/step | Analytic default ms/step | +|---|---:|---:| +| 1 | 10.986 | 10.739 | +| 2 | 11.002 | 10.725 | +| 3 | 11.061 | 10.742 | +| mean | 11.016 | 10.735 | + +**-2.6% step time.** Every analytic rep beat every fixed-default rep. An earlier +build with the threshold at 8 blocks/SM measured -2.3% on the same shapes; the +16 blocks/SM threshold is the more conservative choice and loses nothing here. + +Combined with the packed E2M1 decode from the previous section, decode goes from +11.610 ms/step to 10.735 ms/step, **-7.5%**. + +### Validation + +```bash +cd onnxruntime/test/python/transformers +ORT_ENABLE_FP4_GEMV=1 python -m pytest -q test_qmoe_nvfp4_cuda.py # 22 passed +ORT_ENABLE_FP4_GEMV=0 python -m pytest -q test_qmoe_nvfp4_cuda.py # 22 passed +ORT_ENABLE_FP4_GEMV=1 ORT_FP4_GEMV_DEFAULT_TILING=0 \ + python -m pytest -q test_qmoe_nvfp4_cuda.py # 22 passed +``` + +### Decision + +- Keep. Every config computes the same dot products with the same accumulation + dtype, so this is purely about picking the faster launch shape, and it now + happens without a stream sync and works under CUDA graph capture. The summation + order does depend on `Threads`, so the choice is not bit-neutral; the tests + cover both clauses of the heuristic and the `=0` opt-out. +- Clause (a) is unconditional and carries most of the win (fc2: 22.2 -> 17.5 us). + Clause (b) adds the fc1 win (26.2 -> 23.9 us) and is the part that generalizes + least, hence the deliberately conservative occupancy-derived threshold and the + `ORT_FP4_GEMV_DEFAULT_TILING=0` opt-out. +- The autotuner is still the right tool for shapes the heuristic gets wrong; it + now starts from a better default and overrides it only when it measures a win. diff --git a/docs/design/onnxruntime_web_jsep_to_webgpu_ep_migration.md b/docs/design/onnxruntime_web_jsep_to_webgpu_ep_migration.md index 089b95d4ca9c2..5c5e4a51c7ca3 100644 --- a/docs/design/onnxruntime_web_jsep_to_webgpu_ep_migration.md +++ b/docs/design/onnxruntime_web_jsep_to_webgpu_ep_migration.md @@ -1,6 +1,13 @@ # Design: Migrate onnxruntime-web from JSEP to the native WebGPU EP -**Scope:** `onnxruntime-web` JavaScript/TypeScript package — WebGPU backend; WebNN initialization glue +**Scope:** the `onnxruntime-web` JavaScript/TypeScript package (WebGPU backend; WebNN initialization glue) **and** +the native JS execution provider behind it — `onnxruntime/core/providers/js/`, `onnxruntime/contrib_ops/js/` and +their build plumbing — which Phase 2 also removes (§10.2). Roughly 130 C++/CMake files that a JS-only reading of +this document would hide. + +**Deprecation notice:** [docs/JSEP_Deprecation.md](../JSEP_Deprecation.md) is the contributor-facing statement of +the freeze policy (bug and security fixes only), published as the first step of Phase 0 (§8). This document is the +implementation plan behind it. **Related work:** [Remove the WebGL (onnxjs) backend from onnxruntime-web](onnxruntime_web_remove_webgl_backend.md) — independent, but shares the `onnxruntime-web/all` bundle and the deprecation-warning utility. @@ -22,9 +29,16 @@ is unaffected: it is already the native C++ WebNN EP in both builds (both link ` selects the `WEBNN` EP either way). Removing JSEP only drops the shared JS init glue (`jsepInit('webnn', …)` → `webnnInit(…)`, same `WebNNBackend`), not the WebNN EP itself. -- **Phase 1 (this release):** flip the default (`.`) and `./all` bundles to the native WebGPU EP, add a temporary - `onnxruntime-web/jsep` escape-hatch export for one release, and ship deprecation warnings + docs. -- **Phase 2 (subsequent release):** delete JSEP and remove the temporary `/jsep` export and build flags. +- **Phase 0 (now):** announce the deprecation and freeze JSEP to bug and security fixes, then close the build + configuration, parity and CI gaps that block the flip. +- **Phase 1 (a future release):** flip the default (`.`) and `./all` bundles to the native WebGPU EP, add a temporary + `onnxruntime-web/jsep` escape-hatch export, and ship deprecation warnings + docs. +- **Phase 2 (a subsequent release after Phase 1):** delete JSEP — both the TypeScript backend and the native JS + EP — and remove the temporary `/jsep` export and build flags. + +Only Phase 0 is scheduled. Phase 1 begins when its release gate is met (§9), not at a target release. Phase 2 +begins when a review of reported parity gaps and JSEP usage concludes that the temporary `/jsep` export can be +withdrawn (§10) — at minimum one release after Phase 1. --- @@ -44,8 +58,8 @@ selects the `WEBNN` EP either way). Removing JSEP only drops the shared JS init - Make the native WebGPU EP the default WebGPU backend, transparently for existing consumers (same import, `webgpu` key, and public API). -- Give JSEP consumers a low-effort migration path plus a one-release safety net. -- Remove the JSEP TypeScript compute path and its build variants. +- Give JSEP consumers a low-effort migration path plus a safety net for the duration of the deprecation window. +- Remove the JSEP TypeScript compute path, the native JS execution provider, and their build variants. ### Non-goals @@ -79,12 +93,13 @@ which sets `DISABLE_JSEP = !!USE_WEBGPU_EP` and `DISABLE_WEBGPU = !USE_WEBGPU_EP - **Transparent default swap.** Both JSEP and the native WebGPU EP register under the `webgpu` key, so flipping the default bundle to the native WebGPU EP requires no consumer source changes. -- **Escape hatch.** Phase 1 adds a temporary `onnxruntime-web/jsep` export (built `USE_WEBGPU_EP=false`) that pins - JSEP for one release. It is deprecated and warns once, doubling as a parity-bug funnel. +- **Escape hatch.** Phase 1 adds a temporary `onnxruntime-web/jsep` export, built with JSEP selected, that pins + JSEP for the duration of the deprecation window — at least one release, closed at an explicit checkpoint rather + than on a fixed date (§10). It is deprecated and warns once, doubling as a parity-bug funnel. - **`/all` bundle.** `/all` bundles the WebGPU/WebNN backend and WebGL — two independent things changed by two efforts (this one flips WebGPU to native; the WebGL effort drops WebGL). It is kept to avoid breaking imports. Once both land, `/all` becomes a real alias of the webgpu/default artifact (`.asyncify.wasm`). The two - efforts may land in either order; whichever lands second performs the repoint (§9, WebGL doc §8). + efforts may land in either order; whichever lands second performs the repoint (§10, WebGL doc §8). --- @@ -106,14 +121,60 @@ Optional follow-up: a typed `enableInt64?: boolean` option. Graph capture forces int64 on (`enable_int64_{enable_graph_capture || enable_int64}` in `webgpu_execution_provider.cc`) to keep the captured region all-GPU. So `enableGraphCapture = true` silently moves int64 arithmetic to the GPU and truncates genuine `> 2³¹` values — an exception to default parity, flagged in the -migration guide (§11). +migration guide (§12). + +--- + +## 7. Build configuration parity + +The JSEP and native-WebGPU WASM artifacts are not built with the same operator and type surface today, and the +difference is invisible from the JavaScript layer. + +`.github/workflows/linux-wasm-ci-build-and-test-workflow.yml` defines a `reduced_size_build_args` set — +`--disable_ml_ops --disable_generation_ops --disable_types string float4 float8 optional sparsetensor +--include_ops_by_config onnxruntime/wasm/reduced_types.config --enable_reduced_operator_type_support` — and +applies it to the `.asyncify` and `.jspi` (native WebGPU EP) legs but **not** to the `.jsep` leg. + +`onnxruntime/wasm/reduced_types.config` keeps every operator (`!no_ops_specified_means_all_ops_are_required`) and +restricts only types: the globally-allowed set is `bool, int8_t, uint8_t, int32_t, uint32_t, int64_t, uint64_t, +float, MLFloat16`, so `double`, `int16_t` and `uint16_t` are dropped. The `--disable_*` flags additionally remove +the `ai.onnx.ml` operators, the generation operators, and the string / float4 / float8 / optional / sparse-tensor +types. + +**This is not a WebGPU-only concern.** The default bundle's `wasm` (CPU) backend runs on the same binary, so +flipping the default from `.jsep.wasm` to `.asyncify.wasm` *as currently built* would also narrow the operator and +type surface for CPU inference — a much broader change than swapping the WebGPU implementation, and one that fails +at load time rather than falling back. + +Three builds settle it: + +| Build | Artifact | Reduced-size args | +|---|---|---| +| A | `.jsep.wasm` | no — today's default bundle | +| B | `.asyncify.wasm` | yes — today's `/webgpu` bundle | +| C | `.asyncify.wasm` | no | + +`C − B` is the download cost of restoring full op/type parity; `B − A` and `C − A` are the user-visible size +change. `A` vs `C` is confounded and should not be read as the cost of the migration: JSEP ships no C++ WebGPU +kernels, while the native build compiles the full WebGPU kernel set plus Dawn. + +**Decision:** either drop the reduced-size args from the default artifact (parity, larger download) or apply them +uniformly and document the narrowed surface as a migration-visible change. Resolve from the measurement before the +flip; record the outcome here and in the migration guide (§12). --- -## 7. Items to validate before the flip +## 8. Phase 0 — Announce and close gaps (before the flip) -Parity checks to close before Phase 1. Items 1–3 need a real browser + GPU/WebNN run; item 4 is a known wiring gap -to fix. +Phase 0 is the pre-work, in three parts: + +- **Announce.** Publish [docs/JSEP_Deprecation.md](../JSEP_Deprecation.md) and the in-tree README pointers in + `onnxruntime/core/providers/js/` and `js/web/lib/wasm/jsep/`, so JSEP contributions can be redirected + immediately. This is deliberately first: it costs nothing, changes no behavior, and stops new JSEP work + accumulating while the rest of the plan is executed. +- **Resolve the build-configuration question** (§7), which gates the flip. +- **Close the parity and coverage gaps below.** Items 1–3 need a real browser + GPU/WebNN run; items 4–6 are known + wiring and coverage gaps to fix. 1. **Proxy-worker (`wasm.proxy = true`).** `./webgpu` already ships this exact path (native WebGPU EP + Asyncify + proxy over the EP-agnostic `proxy-wrapper.ts`), so this is a coverage check, not new wiring. CI today runs `--wasm.proxy` @@ -126,76 +187,167 @@ to fix. 3. **WebNN.** WebNN is already the native C++ WebNN EP in both builds (`session-options.ts` selects `WEBNN` either way); the flip only swaps the JS init bridge (`jsepInit('webnn', …)` → `webnnInit(…)`, same `WebNNBackend`). Confirm the native `webnnInit` path behaves the same in a `navigator.ml`-capable environment. -4. **Global `env.webgpu.*` settings (wiring gap).** `wasm-core-impl.ts` requests the adapter on the JS side but - calls `webgpuInit()` without forwarding it to the native WebGPU EP, and `startProfiling()` is a TODO on the native - path — so `adapter` / `powerPreference` / `forceFallbackAdapter` / `profiling` are silently dropped. Wire these - into the native path (or document the change) before the flip. +4. **Global `env.webgpu.*` settings (wiring gap).** `wasm-core-impl.ts` requests a `GPUAdapter` on the JS side and + then calls `webgpuInit()` without forwarding it, so `adapter`, `powerPreference` and `forceFallbackAdapter` are + silently dropped on the native path. Resolution: + - **`powerPreference`** — forward it to the existing `kPowerPreference` EP option + (`ep.webgpuexecutionprovider.powerPreference` in `webgpu_provider_options.h`), set in `session-options.ts` + alongside the other `webgpu` options. Note that the two differ when it is *unset*: JSEP passes `undefined` + to `requestAdapter()` and lets the browser choose, while `WebGpuContextConfig::power_preference` + (`webgpu_context.h`) defaults to `WGPUPowerPreference_HighPerformance`. Left alone, the flip would move + users who never expressed a preference onto a discrete GPU, so the wiring needs a behavior-preserving + default. + - **`adapter` / `forceFallbackAdapter`** — both are `@deprecated` in `js/common/lib/env.ts`, which names + `env.webgpu.device` as the replacement. That pointer is wrong: `env.webgpu.device` is output-only — both + paths write it after init and neither reads it. Custom devices go through the per-session option instead, + `executionProviders: [{ name: 'webgpu', device }]`, which `session-options.ts` already registers via + `webgpuRegisterDevice`. Document `adapter` / `forceFallbackAdapter` as no-ops on the native path, and fix + the `env.webgpu.device` doc comment. + - Drop the now-dead `navigator.gpu.requestAdapter()` call on the native path once nothing consumes its result. +5. **Profiling.** `env.webgpu.profiling.ondata` is a JSEP-only mechanism: a per-dispatch JavaScript callback + receiving `kernelId` / `kernelType` / `kernelName` / `programName`, timestamps and input/output tensor metadata + (`WebGpuProfilingDataV1` in `js/common/lib/env.ts`), falling back to `console.log` when unset. The native + WebGPU EP has no equivalent hook. It has a richer ORT-framework profiler (`webgpu_profiler.{h,cc}`, gated on + the session's `enableProfiling`), but its JSON trace is unreachable from JavaScript today: + `session-handler-inference.ts` leaves `startProfiling()` as a TODO on **both** paths, and `endProfiling()` in + `wasm-core-impl.ts` frees the filename returned by `_OrtEndProfiling` without surfacing it. + + **Decision:** `ondata` is dropped and documented as JSEP-only. Mapping `env.webgpu.profiling.mode` (and the + deprecated `profilingMode`) onto the session's `enableProfiling`, implementing `startProfiling()`, and + surfacing the JSON trace are a **Phase 2 prerequisite** (§10.1), not a flip gate — the `/jsep` escape hatch + preserves today's behavior for the whole deprecation window. +6. **CI coverage gaps.** Beyond items 1–3, the pipelines themselves need work: + - `tools/ci_build/github/azure-pipelines/templates/win-wasm-ci.yml` builds only the JSEP variants + (`wasm_simd_jsep`, `wasm_simd_threads_jsep`) and has no `BuildWebGPU` parameter, so once JSEP is gone the + Windows/ADO side produces no GPU-capable WASM at all. + - No CI leg passes `--jspi`, so the `/jspi` bundle is built but never exercised. + - GitHub Actions leg 07 runs `suite1` while the ADO leg 07 runs the default `suite0` — two nominally identical + legs testing different things. + - `--webgpu-ep` rebuilds `dist/ort.all[.min].js` in place and `karma.conf.js` always loads that path, so + consecutive legs overwrite each other's bundle. Per-session typed options are already a superset of JSEP's, confirmed by source audit. --- -## 8. Phase 1 — Flip + deprecate (this release) +## 9. Phase 1 — Flip + deprecate (a future release) -1. **Flip the default.** Build `.` and `./all` with the native WebGPU EP (`USE_WEBGPU_EP=true`). The `webgpu` key and - public API are unchanged. `/all` keeps WebGL until the WebGL effort removes it. -2. **Escape hatch.** Add the temporary, deprecated `onnxruntime-web/jsep` export (`USE_WEBGPU_EP=false`). +1. **Rename the build flag.** `USE_WEBGPU_EP` / `--webgpu-ep` inverts in meaning once the native EP is the + default. Replace it with a `--jsep` opt-in (`--no-jsep` selecting native). Land the rename while JSEP is still + the default so it carries no behavior change. +2. **Escape hatch.** Add the temporary, deprecated `onnxruntime-web/jsep` export (built with JSEP selected). 3. **Warn once.** The `/jsep` build warns once (respecting `env.logLevel`) via a shared deprecation-warning utility, also used by the WebGL effort; the native default emits nothing. -4. **Publish the tracking issue** ahead of the release as the early-warning channel. -5. **Docs.** Deprecation banners + migration guidance in `js/web/README` (hand-authored) and a release-notes entry - covering the `.jsep.wasm` → `.asyncify.wasm` filename change, `/jsep` pinning, and the int64 `extra` opt-in. -6. **Retarget the operator-doc generator.** `generate-webgpu-operator-md.ts` parses the JSEP registrations +4. **Retarget the operator-doc generator.** `generate-webgpu-operator-md.ts` parses the JSEP registrations (`js_execution_provider.cc`, `js_contrib_kernels.cc`), so post-flip `webgpu-operators.md` describes the wrong EP; point it at the native WebGPU EP registrations (those JSEP sources are removed in Phase 2). +5. **Flip the default.** Build `.` and `./all` with the native WebGPU EP. The `webgpu` key and public API are + unchanged. `/all` keeps WebGL until the WebGL effort removes it. With steps 1–4 already landed this is a + one-line default change — keep it an isolated, independently revertible commit. +6. **Open the tracking issue.** The public early-warning and parity-report channel: what changed, how to pin + `/jsep`, and where to report native-WebGPU-EP gaps. Link it from the deprecation notice and the release notes. +7. **Docs.** Deprecation banners + migration guidance in `js/web/README` (hand-authored) and a release-notes entry + covering the `.jsep.wasm` → `.asyncify.wasm` filename change, `/jsep` pinning, and the int64 `extra` opt-in. + This is the first user-facing announcement; capture the JSEP usage baseline (§11) before it ships. **Release gate.** The flip is gated on both: (a) a **blocking** CI job running the native WebGPU EP — the current native-WebGPU-EP legs are advisory (`continue-on-error` in `windows-web-ci-workflow.yml`, `continueOnError` in `win-web-ci.yml`) and must be promoted to blocking. The default `.` bundle and `./webgpu` are the same build, so a green `/webgpu` run covers the flip only if it exercises proxy/IO-binding/WebNN — which the current op-parity leg -does not; and (b) the §7 items resolved with targeted coverage (op-parity tests don't exercise those paths). +does not; and (b) the §8 items resolved with targeted coverage (op-parity tests don't exercise those paths). --- -## 9. Phase 2 — Removal (subsequent release) +## 10. Phase 2 — Removal (a subsequent release after Phase 1) + +The `/jsep` hatch stays available for **at least** one release. The window is not fixed in advance: it closes at an +explicit checkpoint that reviews parity regressions reported against the tracking issue, the warn-once funnel, and +the JSEP usage signal (per-file CDN statistics for `.jsep.wasm` versus `.asyncify.wasm` and `.jspi.wasm`, which is +the only breakdown npm download totals cannot give). Nothing in the code or the published package hard-codes a +removal release, so extending the window costs only keeping the JSEP CI legs alive. + +### 10.1 Prerequisites + +Both are non-destructive and must land before any deletion, so that the deletion diffs stay reviewable: + +- **Profiling.** Map `env.webgpu.profiling.mode` (and the deprecated `profilingMode`) onto the session's + `enableProfiling`, implement `startProfiling()`, and stop discarding the `_OrtEndProfiling` trace filename in + `endProfiling()`. Without this, removing JSEP removes the only working WebGPU profiling in the package (§8.5). +- **Relocate WebNN** out of `js/web/lib/wasm/jsep/` (`backend-webnn.ts`, `webnn/`) to a neutral path, repointing + importers including `test/test-runner.ts` and `test/unittests/pool-output-shape.ts`. -Timed one release after Phase 1, keeping the `/jsep` hatch available for exactly one release. Extend only if -native-WebGPU-EP parity regressions surface via real `/jsep` usage. +### 10.2 Web package -1. Drop JSEP WASM artifacts; update `build.ts` / `package.json`; remove the `USE_WEBGPU_EP` flag. Repoint `/all` +1. Drop JSEP WASM artifacts; update `build.ts` / `package.json`; remove the `--jsep` build flag. Repoint `/all` to the webgpu/default artifact — only once WebGL has also been dropped (WebGL doc §8); otherwise `/all` stays a distinct bundle until then. 2. Delete `BUILD_DEFS.DISABLE_JSEP` and the code it gates; simplify `index.ts`. -3. Relocate WebNN out of `jsep/` (`backend-webnn.ts`, `webnn/`) to a neutral path. -4. Remove `pre-jsep.js` glue; confirm `post-webgpu.js` / `post-webnn.js` cover initialization. -5. Remove the `onnxruntime-web/jsep` export. A lingering `import 'onnxruntime-web/jsep'` then fails with the +3. Remove `pre-jsep.js` glue; confirm `post-webgpu.js` / `post-webnn.js` cover initialization. +4. Remove the `onnxruntime-web/jsep` export. A lingering `import 'onnxruntime-web/jsep'` then fails with the native bundler error — an acceptable build-time failure on this temporary surface. The default `.` import is unaffected. +5. Drop the JSEP WASM build legs and the `build_jsep` / `BuildJsep` pipeline parameters. + +### 10.3 Native JS EP removal + +Phase 2 also removes the C++ half — roughly 130 files, almost entirely deletion: + +- `onnxruntime/core/providers/js/` (~90 files) and `onnxruntime/contrib_ops/js/` (~30 files). +- `cmake/onnxruntime_providers_js.cmake`, plus the `USE_JSEP` plumbing in `cmake/CMakeLists.txt`, + `onnxruntime_providers.cmake`, `onnxruntime_providers_cpu.cmake` (the unguarded + `onnxruntime_js_contrib_ops_cc_srcs` glob), `onnxruntime_unittests.cmake` and `onnxruntime_webassembly.cmake`; + and the `--use_jsep` argument in `tools/ci_build/build_args.py` / `build.py`. +- `onnxruntime/wasm/pre-jsep.js` and `js/build_jsep.bat`. +- The `USE_JSEP` guards in `onnxruntime/test/optimizer/graph_transform_test.cc` and + `group_query_attention_pre_norm_fusion_test.cc`. + +Two items need care: + +- **`kJsExecutionProvider` is public C API surface.** It is declared in + `include/onnxruntime/core/graph/constants.h`, named in `onnxruntime_c_api.h`, and referenced from + `provider_registration.cc`, `get_execution_providers.cc`, `provider_factory_creators.h`, `session_state.cc`, + `conv_activation_fusion.cc`, `graph_transformer_utils.cc` and `external_data_loader.h`. Removing it is an API + change and needs an explicit release-note entry; no tombstone is planned. +- **`post-webnn.js` is suppressed under JSEP** in `cmake/onnxruntime_webassembly.cmake`, so removing JSEP changes + which glue a WebNN build links. This needs a WebNN smoke test, not just a successful compile. + +**Ordering.** CI must stop passing `--use_jsep` (§10.2) *before* this code is deleted, or the pipelines break on +the deletion commit. --- -## 10. Risks and mitigations +## 11. Risks and mitigations | Risk | Likelihood | Mitigation | |---|---|---| -| Undiscovered native-WebGPU-EP parity gap vs. JSEP | Medium | One-release `/jsep` escape hatch + warn-once funnel; differential tests | +| Undiscovered native-WebGPU-EP parity gap vs. JSEP | Medium | `/jsep` escape hatch + warn-once funnel; differential tests | +| Default bundle silently narrows its operator/type surface (reduced-size build args) | High if unaddressed | Measure builds A/B/C and resolve before the flip (§7) | | int64 behavior change vs. JSEP | Low | None by default (native-off matches JSEP); `enableInt64 = 1` is an opt-in tradeoff | -| Proxy / IO-binding / WebNN unexercised in the native-WebGPU build in CI | Unknown | Validate before flip (§7) | -| Global `env.webgpu.*` settings dropped on native | Medium | Wire into the native path or document — release gate (§7, §8) | +| Proxy / IO-binding / WebNN unexercised in the native-WebGPU build in CI | Unknown | Validate before flip (§8) | +| Global `env.webgpu.*` settings dropped on native | Medium | `powerPreference` wired with a behavior-preserving default when unset; `adapter` / `forceFallbackAdapter` documented as no-ops, custom devices directed to the per-session `device` option — release gate (§8.4, §9) | +| `env.webgpu.profiling.ondata` has no native equivalent | Low | Documented as JSEP-only; native profiling wiring is a Phase 2 prerequisite (§8.5, §10.1) | +| Windows/ADO produces no WebGPU-EP WASM build | Medium | Add a `BuildWebGPU` leg to `win-wasm-ci.yml` before Phase 2 (§8.6) | +| `post-webnn.js` linkage changes for WebNN once JSEP is removed | Medium | WebNN smoke test on the Phase 2 build, not just a compile check (§10.3) | | WASM filename change breaks `wasmPaths` | Medium | Document `.jsep.wasm` → `.asyncify.wasm` | -| No telemetry on JSEP adoption | Medium | Warn-once funnel + one-release escape hatch | +| No telemetry on JSEP adoption | Medium | Per-file CDN statistics for `.jsep.wasm` vs `.asyncify.wasm` / `.jspi.wasm` (baseline captured before the user-facing announcement), plus the warn-once funnel and the escape hatch | --- -## 11. Migration guide +## 12. Migration guide - **Default import (`onnxruntime-web`) and `onnxruntime-web/all`:** no source change if you use default, package-managed asset resolution; the `webgpu` backend keeps working and converges to the native WebGPU EP. Consumers that pin WASM artifacts must still update the path (see `wasmPaths` below). - **Lower-overhead build (`onnxruntime-web/jspi`):** on JSPI-capable browsers, prefer `./jspi` — the same native WebGPU EP with a smaller WASM binary and lower per-call overhead than Asyncify (the universal fallback). -- **Need JSEP for one more release:** import `onnxruntime-web/jsep` (temporary, deprecated). File an issue if the - native WebGPU EP does not work for your model. +- **Need JSEP for now:** import `onnxruntime-web/jsep` (temporary, deprecated). File an issue if the + native WebGPU EP does not work for your model — those reports are what set the removal timeline. +- **WebGPU profiling:** `env.webgpu.profiling.ondata` is JSEP-only and has no native equivalent. Code that + registers an `ondata` callback keeps working on `/jsep` but silently receives nothing on the default bundle + (§8.5). +- **Operator and type coverage:** the default bundle's backing artifact changes, and with it the operator/type + surface available to *all* backends in that bundle, including `wasm` (CPU) — see §7 for the resolution and the + exact set involved. - **int64-heavy models:** no change needed — the default matches JSEP. Exceptions that move int64 arithmetic to the GPU (lossy for genuine `> 2³¹` values): `extra: { 'ep.webgpuexecutionprovider.enableInt64': '1' }`, and `enableGraphCapture = true` (forces int64 on; see §6). diff --git a/js/web/lib/wasm/jsep/README.md b/js/web/lib/wasm/jsep/README.md new file mode 100644 index 0000000000000..60c6f59919fd8 --- /dev/null +++ b/js/web/lib/wasm/jsep/README.md @@ -0,0 +1,14 @@ +# JSEP — deprecated + +This directory holds the TypeScript half of **JSEP**, the JavaScript WebGPU compute path. It is **deprecated** and +will be removed. The replacement is the native WebGPU execution provider, already shipping as +`onnxruntime-web/webgpu` and `onnxruntime-web/jspi`. + +**Bug fixes and security fixes only.** New operators, new features and performance work belong in the native +WebGPU EP (`onnxruntime/core/providers/webgpu/`). + +**WebNN is not deprecated.** `backend-webnn.ts` and `webnn/` live here only because they share JSEP's +initialization glue. WebNN is already the native C++ WebNN EP in every build, and this code will be relocated to a +neutral path rather than removed. + +See [docs/JSEP_Deprecation.md](../../../../../docs/JSEP_Deprecation.md) for more details. diff --git a/onnxruntime/contrib_ops/cpu/bert/attention_common.h b/onnxruntime/contrib_ops/cpu/bert/attention_common.h index e46075a86f811..c33d3b1c72bfd 100644 --- a/onnxruntime/contrib_ops/cpu/bert/attention_common.h +++ b/onnxruntime/contrib_ops/cpu/bert/attention_common.h @@ -85,6 +85,86 @@ inline KVQuantizationType StringToKVQuantizationType(std::string s) { "'. Valid values are: NONE, PER_TENSOR, PER_CHANNEL."); } +// Logical element type of a KV cache. Members are named after the ONNX element type they denote. +// DEFAULT means "whatever the cache tensor's own element type is" and is the only value a model +// needs when that type is expressible in ONNX, i.e. for float16 / bfloat16 / int8 / float8e4m3fn; +// naming one of those explicitly is allowed but must agree with the tensor. +// The sub-byte members have no ONNX tensor type here: they are packed two per byte into a uint8 +// cache, so the last cache dimension holds (head_size + 1) / 2 bytes and logical element 2*i +// occupies the low-order bits of byte i. +// Every member is a *signed*, zero-symmetric type. Quantization here has a scale but no zero point +// (there are no zero-point inputs), so an unsigned logical type such as uint4 or uint8 would need +// an implied offset of 2^(bits-1) that nothing in the contract can express. INT4 is still *stored* +// in an unsigned nibble biased by +8, but that is a storage encoding removed on read, not a +// quantization zero point. Unsigned types must arrive together with zero-point inputs. +enum class KVCacheDataType : int { + DEFAULT = 0, + FLOAT16 = 1, + BFLOAT16 = 2, + INT8 = 3, + FLOAT8E4M3FN = 4, + INT4 = 5, + FLOAT4E2M1 = 6, +}; + +// True for the packed sub-byte members, which are stored in a uint8 cache. +inline bool IsSubByteKVCacheDataType(KVCacheDataType t) { + return t == KVCacheDataType::INT4 || t == KVCacheDataType::FLOAT4E2M1; +} + +// True for the members that require a scale on read/write. +inline bool IsQuantizedKVCacheDataType(KVCacheDataType t) { + return t != KVCacheDataType::DEFAULT && t != KVCacheDataType::FLOAT16 && t != KVCacheDataType::BFLOAT16; +} + +inline const char* KVCacheDataTypeToString(KVCacheDataType t) { + switch (t) { + case KVCacheDataType::FLOAT16: + return "float16"; + case KVCacheDataType::BFLOAT16: + return "bfloat16"; + case KVCacheDataType::INT8: + return "int8"; + case KVCacheDataType::FLOAT8E4M3FN: + return "float8e4m3fn"; + case KVCacheDataType::INT4: + return "int4"; + case KVCacheDataType::FLOAT4E2M1: + return "float4e2m1"; + default: + return ""; + } +} + +inline KVCacheDataType StringToKVCacheDataType(std::string s) { + std::transform(s.begin(), s.end(), s.begin(), [](unsigned char c) { return std::tolower(c); }); + if (s.empty()) { + return KVCacheDataType::DEFAULT; + } + if (s == "float16") { + return KVCacheDataType::FLOAT16; + } + if (s == "bfloat16") { + return KVCacheDataType::BFLOAT16; + } + if (s == "int8") { + return KVCacheDataType::INT8; + } + if (s == "float8e4m3fn") { + return KVCacheDataType::FLOAT8E4M3FN; + } + if (s == "int4") { + return KVCacheDataType::INT4; + } + if (s == "float4e2m1") { + return KVCacheDataType::FLOAT4E2M1; + } + ORT_THROW("Invalid KV cache data type: '", s, + "'. Valid values are: '' (use the cache tensor's element type), float16, bfloat16, int8, " + "float8e4m3fn, int4, float4e2m1. Unsigned types are excluded because quantization here is " + "symmetric with no zero point."); +} + constexpr bool LAYOUT_BSNH = false; constexpr bool LAYOUT_BNSH = true; diff --git a/onnxruntime/contrib_ops/cpu/bert/attention_parameters.h b/onnxruntime/contrib_ops/cpu/bert/attention_parameters.h index 7b3c38cc883e8..c856528bbb9f1 100644 --- a/onnxruntime/contrib_ops/cpu/bert/attention_parameters.h +++ b/onnxruntime/contrib_ops/cpu/bert/attention_parameters.h @@ -132,6 +132,26 @@ struct PagedAttentionParameters : AttentionParameters { int local_window_size; // The window size includes new token. It only includes tokens on the left side. bool rotary_interleaved; float softcap; + // Internal attention-sink path, enabled when head_sink (input 11) is provided. + bool use_smooth_softmax = false; + // Per-head Q/K RMSNorm (QK-Norm) prologue applied before RoPE (inputs 12/13). + bool use_qk_norm = false; + float qk_norm_epsilon = 1e-6f; + // Quantized paged KV cache. Scales are inputs 14/15 and are always FP32, as in + // GroupQueryAttention. The storage element type is carried by the kernel's TCACHE specialization; + // the k_cache_dtype / v_cache_dtype attributes only override it for sub-byte formats packed into + // uint8, which no backend supports yet. + KVQuantizationType k_quant_type = KVQuantizationType::NONE; + KVQuantizationType v_quant_type = KVQuantizationType::NONE; + // Multi-head Latent Attention (kv_cache_layout == "LATENT"). There is a single physical cache: + // V of every head is the leading v_head_size channels of the same key_cache row, so 'value' and + // 'value_cache' are absent. The inherited v_head_size / v_hidden_size hold the effective V width + // and the output width; in SEPARATE mode they equal head_size / hidden_size. + bool is_latent_kv = false; + // First channel within head_size covered by rotary embedding. RoPE covers + // [rotary_offset, rotary_offset + rotary_dim); channels outside are copied through. Default 0 + // reproduces the original prefix-RoPE behavior. MLA uses rotary_offset == kv_lora_rank. + int rotary_offset = 0; }; // Parameters for sparse attention. diff --git a/onnxruntime/contrib_ops/cpu/bert/multihead_attention_helper.h b/onnxruntime/contrib_ops/cpu/bert/multihead_attention_helper.h index 66d8998ffe8da..e195fed8868b2 100644 --- a/onnxruntime/contrib_ops/cpu/bert/multihead_attention_helper.h +++ b/onnxruntime/contrib_ops/cpu/bert/multihead_attention_helper.h @@ -3,6 +3,8 @@ #pragma once +#include + #include "core/common/common.h" #include "core/providers/common.h" #include "contrib_ops/cpu/bert/attention_common.h" @@ -248,16 +250,22 @@ inline Status CheckCacheIndirection( "Input 'cache_indirection' is expected to have 3 dimensions, got ", cache_indir_dims.size()); } - num_beams = static_cast(cache_indir_dims[1]); - if (cache_indir_dims[1] == 0) { + if (cache_indir_dims[1] <= 0 || + cache_indir_dims[1] > static_cast(std::numeric_limits::max())) { return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 'cache_indirection' dimension 1 should be num_beams, got ", cache_indir_dims[1]); } - if (cache_indir_dims[0] != static_cast(batch_beam_size / num_beams)) { + num_beams = static_cast(cache_indir_dims[1]); + // Require num_beams to evenly divide batch_beam_size, and dim 0 to be exactly batch_beam_size / num_beams. + // Comparing dim 0 against the exact quotient (rather than multiplying dim 0 by num_beams) keeps the + // relation intact while avoiding int64_t overflow on the multiplication for arbitrary shape inputs. + if (batch_beam_size % num_beams != 0 || + cache_indir_dims[0] != static_cast(batch_beam_size / num_beams)) { return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "Input 'cache_indirection' dimension 0 should be batch_size, got ", - cache_indir_dims[0]); + "Input 'cache_indirection' dimension 0 (", cache_indir_dims[0], + ") times dimension 1 (num_beams=", num_beams, + ") must equal batch_beam_size (", batch_beam_size, ")"); } if (max_sequence_length > 0 && cache_indir_dims[2] != static_cast(max_sequence_length)) { // First condition is to avoid this check for cross attention layers where diff --git a/onnxruntime/contrib_ops/cuda/bert/attention_data.h b/onnxruntime/contrib_ops/cuda/bert/attention_data.h index 0b87ea831b745..be7d97d7be0a0 100644 --- a/onnxruntime/contrib_ops/cuda/bert/attention_data.h +++ b/onnxruntime/contrib_ops/cuda/bert/attention_data.h @@ -254,35 +254,58 @@ struct GroupQueryAttentionData { void* cudnn_handle = nullptr; }; -template +// TCACHE is the element type of the paged key/value cache. It equals T for an unquantized cache and +// is int8_t / Float8E4M3FN when the cache is quantized (see PagedAttentionParameters::k_quant_type). +template struct PagedAttentionData { // Input Tensors const T* query = nullptr; const T* key = nullptr; const T* value = nullptr; - T* key_cache = nullptr; - T* value_cache = nullptr; + TCACHE* key_cache = nullptr; + TCACHE* value_cache = nullptr; + // FP32 quantization scales for the paged cache: (1,) for PER_TENSOR and + // (kv_num_heads, 1, head_size) for PER_CHANNEL. nullptr when the cache is not quantized. + const float* k_scale = nullptr; + const float* v_scale = nullptr; const int* cumulative_seqlens_q = nullptr; const int* past_seqlens = nullptr; const int* block_table = nullptr; - const int* slot_mappings = nullptr; + // Optional explicit write slots, one per query token, into the cache viewed as + // [num_blocks * block_size, kv_num_heads, head_size]. A value of -1 suppresses the K/V store + // for that token (prefix cache hit / rejected speculative token). nullptr keeps the legacy + // derived mapping (past_seqlens + position within the sequence). + const int* slot_mapping = nullptr; const T* cos_cache = nullptr; const T* sin_cache = nullptr; + // Per-head attention sink (num_heads,). nullptr with use_smooth_softmax means a sink value of 0. + const T* head_sink = nullptr; + // QK-Norm weights (head_size,), shared across heads. Both are set or neither is. + const T* q_norm_weight = nullptr; + const T* k_norm_weight = nullptr; - // Flash buffers - T* softmax_lse = nullptr; + // Flash buffers. FlashAttention always emits FP32 log-sum-exp regardless of T; with + // params.num_splits <= 1 (which mha_varlen_fwd never overrides) the varlen layout is + // [num_heads, token_count]. + float* softmax_lse = nullptr; int* cumulative_seqlens_kv = nullptr; // Flash api takes cumulative sequence length for kv-cache // Fused op buffers T* workspace_buffer = nullptr; - // Memory-efficient attention (CUTLASS fMHA) buffers for the unfused fallback path - // taken when FlashAttention is unavailable (SM<80 or ORT_DISABLE_FLASH_ATTENTION). - T* gathered_key = nullptr; // [total_kv_tokens, num_heads, head_size], packed varlen (GQA-expanded) - T* gathered_value = nullptr; // [total_kv_tokens, num_heads, head_size], packed varlen (GQA-expanded) - T* fmha_buffer = nullptr; // CUTLASS fMHA output-accumulator workspace + // Dense KV staging buffers. Always used by the memory-efficient (CUTLASS fMHA) fallback, which + // needs a packed-varlen [total_kv_tokens, num_heads, head_size] GQA-expanded view of the cache. + // The FlashAttention path also uses them when the cache is quantized: Flash cannot read a + // quantized page directly, so the cache is dequantized into [total_kv_tokens, kv_num_heads, + // head_size] (no GQA expansion) and fed to the non-paged varlen entry point. + T* gathered_key = nullptr; + T* gathered_value = nullptr; + T* fmha_buffer = nullptr; // CUTLASS fMHA output-accumulator workspace // Populated by the caller after a D->H sync on cumulative_seqlens_kv[batch_size]. int total_kv_tokens = 0; + // Max per-batch total KV length. Only needed when the gathered (non-paged) Flash path is used, + // where it becomes mha_varlen_fwd's max_seqlen_k. + int max_kv_len = 0; // Actual max of per-batch new-query lengths (cumulative_seqlens_q[i+1] - cumulative_seqlens_q[i]). // Populated by the caller via the same D->H sync so the MEA path's rotary grid and MEA's @@ -291,12 +314,40 @@ struct PagedAttentionData { // producing silent per-token dropout in MEA and rotary. int max_query_len = 0; + // Paged decode (flash-decoding style) split-KV workspaces. Only allocated when the paged decode + // backend is selected. Layouts are [num_splits, token_count, num_heads, head_size] for the + // accumulator and [num_splits, token_count, num_heads] for the running max / denominator. + float* decode_partial_out = nullptr; + float* decode_partial_max = nullptr; + float* decode_partial_sum = nullptr; + int num_splits = 1; + + // Paged XQA decode workspaces. Only allocated when the XQA decode backend is selected + // (quantized cache, one new token per sequence -- see use_xqa_decode). + // xqa_workspace : XQA semaphores + multi-block scratch (GetXQAScratchSize bytes). + // xqa_page_table : block_table expanded from PagedAttention blocks to XQA's fixed 128-token + // pages, shape [batch_size, max_num_blocks_per_seq * pages_per_block]. + // xqa_query : scratch for Q pre-scaled by a PER_CHANNEL k_scale; unused otherwise. + // xqa_head_sink : head_sink converted to fp32, which is what XQA consumes. + void* xqa_workspace = nullptr; + size_t xqa_workspace_size = 0; + int* xqa_page_table = nullptr; + T* xqa_query = nullptr; + float* xqa_head_sink = nullptr; + // Output Tensors T* output = nullptr; // Kernel Flags bool use_flash_attention = false; bool use_memory_efficient_attention = false; + // Paged decode kernel: reads the paged cache in place and dequantizes inside the kernel, so it + // needs neither the dense staging buffers nor FlashAttention's page-alignment constraint. + bool use_paged_decode = false; + // XQA paged decode kernel: same in-place paged read, but tensor-core based and an order of + // magnitude faster than the generic decode kernel on a quantized cache. Takes precedence over + // use_paged_decode when set. + bool use_xqa_decode = false; }; } // namespace cuda diff --git a/onnxruntime/contrib_ops/cuda/bert/flash_attention/flash_api.cc b/onnxruntime/contrib_ops/cuda/bert/flash_attention/flash_api.cc index efda3f48b9cfc..d2424ea526a5b 100644 --- a/onnxruntime/contrib_ops/cuda/bert/flash_attention/flash_api.cc +++ b/onnxruntime/contrib_ops/cuda/bert/flash_attention/flash_api.cc @@ -294,14 +294,13 @@ Status mha_fwd(const cudaDeviceProp& dprops, bool kv_bsnh, int local_window_size, void* cache_batch_idx, - void* leftpad_k) { + void* leftpad_k, + void* head_sink) { auto round_multiple = [](int x, int m) { return (x + m - 1) / m * m; }; const int head_size_rounded = round_multiple(head_size, 32); const int seqlen_q_rounded = round_multiple(seqlen_q, 128); const int seqlen_k_rounded = round_multiple(seqlen_k, 128); - constexpr void* head_sink = nullptr; - Flash_fwd_params params; set_params_fprop(params, batch_size, diff --git a/onnxruntime/contrib_ops/cuda/bert/flash_attention/flash_api.h b/onnxruntime/contrib_ops/cuda/bert/flash_attention/flash_api.h index 7aed9fe10afbd..32ea803c998d8 100644 --- a/onnxruntime/contrib_ops/cuda/bert/flash_attention/flash_api.h +++ b/onnxruntime/contrib_ops/cuda/bert/flash_attention/flash_api.h @@ -63,7 +63,8 @@ Status mha_fwd(const cudaDeviceProp& dprops, bool kv_bsnh = true, int local_window_size = -1, void* cache_batch_idx = nullptr, - void* leftpad_k = nullptr); + void* leftpad_k = nullptr, + void* head_sink = nullptr); // num_heads Status mha_varlen_fwd(const cudaDeviceProp& dprops, cudaStream_t stream, diff --git a/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu index 89144f5e50787..4f075faae439f 100644 --- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu @@ -1323,17 +1323,12 @@ Status DequantizeFlashAttentionFallback( // (max_length sized) capacity on every decode step is pure memory traffic. bool is_bsnh = (parameters.past_kv_format == AttentionQkvFormat::Q_K_V_BSNH); - ORT_RETURN_IF_ERROR((LaunchDequantizeKV( - stream, k_dequant, reinterpret_cast(data.present_key), data.k_scale, - nullptr, parameters.batch_size, parameters.kv_num_heads, parameters.seqlen_present_kv_cache, - parameters.head_size, parameters.kv_cache_bit_width, parameters.k_quant_type, is_bsnh, - data.total_seq_lens))); - - ORT_RETURN_IF_ERROR((LaunchDequantizeKV( - stream, v_dequant, reinterpret_cast(data.present_value), data.v_scale, - nullptr, parameters.batch_size, parameters.kv_num_heads, parameters.seqlen_present_kv_cache, - parameters.head_size, parameters.kv_cache_bit_width, parameters.v_quant_type, is_bsnh, - data.total_seq_lens))); + ORT_RETURN_IF_ERROR((LaunchDequantizeKVPair( + stream, k_dequant, v_dequant, + reinterpret_cast(data.present_key), reinterpret_cast(data.present_value), + data.k_scale, data.v_scale, parameters.batch_size, parameters.kv_num_heads, + parameters.seqlen_present_kv_cache, parameters.head_size, parameters.kv_cache_bit_width, + parameters.k_quant_type, parameters.v_quant_type, is_bsnh, data.total_seq_lens))); // Step 3: Run Flash Attention on dequantized k/v bool is_causal = parameters.is_unidirectional; @@ -1425,7 +1420,12 @@ Status FlashAttentionAndQuantizeKV( reinterpret_cast(data.softmax_lse_accum), reinterpret_cast(data.out_accum), true, // kv_bsnh = true (BSNH) - local_window_size)); + local_window_size, + /*cache_batch_idx*/ nullptr, /*leftpad_k*/ nullptr, + // head_sink must be forwarded here as well. This is the only prompt path taken when the KV + // cache is quantized, so dropping it silently disables attention sinks for the whole prompt + // while the unquantized prompt path (FlashAttention) keeps them. + reinterpret_cast(const_cast(data.head_sink)))); if (parameters.k_quant_type != KVQuantizationType::NONE) { ORT_RETURN_IF_ERROR((LaunchQuantizeKV( diff --git a/onnxruntime/contrib_ops/cuda/bert/group_query_attention_qdq.cuh b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_qdq.cuh index 6655ccd467574..14c3191190adb 100644 --- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention_qdq.cuh +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_qdq.cuh @@ -6,6 +6,7 @@ #define KV_QUANT_SUPPORTED 1 #include +#include #include #include @@ -395,6 +396,173 @@ Status LaunchDequantizeKV(cudaStream_t stream, T* dequantized_data, return CUDA_CALL(cudaGetLastError()); } +// ============================================================================ +// Fused K + V dequantization. +// +// The decode fallback always dequantizes both caches back to back, and each launch only exposes +// batch_size * kv_num_heads independent (batch, head) slices -- two for a batch-1 GQA model with +// two KV heads. That leaves the vast majority of the device idle. Doing K and V in one launch +// doubles the CTA count, halves the number of launches, and lets both caches share the memory +// pipeline. blockIdx.z selects the tensor; everything else matches DequantizeKVVectorizedKernel. +// +// kVecSize is deliberately smaller here than in the single-tensor kernel: at decode sequence +// lengths the kernel is latency bound rather than instruction bound, so more threads with fewer +// elements each is faster than wide per-thread vectors. +// ============================================================================ +template +__global__ void DequantizeKVPairKernel(T* __restrict__ k_dequantized, T* __restrict__ v_dequantized, + const T_QUANT* __restrict__ k_quantized, + const T_QUANT* __restrict__ v_quantized, + const T_SCALE* __restrict__ k_scale, + const T_SCALE* __restrict__ v_scale, + const int* __restrict__ valid_seq_lens, + int num_heads, + int cache_sequence_length, + int head_size, + KVQuantizationType k_quant_type, + KVQuantizationType v_quant_type, + bool is_input_bsnh) { + static_assert(sizeof(T_QUANT) == 1, "Vectorized dequantization only supports 8-bit caches."); + static_assert(sizeof(T) == 2, "Vectorized dequantization only supports 16-bit outputs."); + + using LoadVec = typename DequantVecType::type; + constexpr int kStoreBytes = (kVecSize * 2 >= 16) ? 16 : kVecSize * 2; + using StoreVec = typename DequantVecType::type; + constexpr int kStoresPerVec = (kVecSize * 2) / kStoreBytes; + constexpr int kElemsPerStore = kVecSize / kStoresPerVec; + + const bool is_value = (blockIdx.z != 0); + T* dequantized_data = is_value ? v_dequantized : k_dequantized; + const T_QUANT* quantized_data = is_value ? v_quantized : k_quantized; + const T_SCALE* scale = is_value ? v_scale : k_scale; + const KVQuantizationType quant_type = is_value ? v_quant_type : k_quant_type; + + const int vecs_per_row = head_size / kVecSize; + const int rows_per_block = kBlockThreads / vecs_per_row; + const int vec_in_row = threadIdx.x % vecs_per_row; + const int row_in_block = threadIdx.x / vecs_per_row; + if (row_in_block >= rows_per_block) { + return; // kBlockThreads is not an exact multiple of vecs_per_row + } + + const int bn = blockIdx.y; // b * num_heads + n + const int n = bn % num_heads; + const int b = bn / num_heads; + + const int limit = (valid_seq_lens == nullptr) + ? cache_sequence_length + : ValidRowLimit(valid_seq_lens[b], cache_sequence_length); + + const int h0 = vec_in_row * kVecSize; + + float scales[kVecSize]; + if (quant_type == KVQuantizationType::PER_TENSOR) { + const float s0 = static_cast(scale[0]); +#pragma unroll + for (int i = 0; i < kVecSize; i++) { + scales[i] = s0; + } + } else { + const T_SCALE* channel_scale = scale + static_cast(n) * head_size + h0; +#pragma unroll + for (int i = 0; i < kVecSize; i++) { + scales[i] = static_cast(channel_scale[i]); + } + } + + // The output is always BNSH; only the input row stride differs between BNSH and BSNH. + const int64_t out_base = static_cast(bn) * cache_sequence_length * head_size + h0; + const int64_t in_base = is_input_bsnh + ? (static_cast(b) * cache_sequence_length * num_heads * head_size + + static_cast(n) * head_size + h0) + : (static_cast(bn) * cache_sequence_length * head_size + h0); + const int in_row_stride = is_input_bsnh ? (num_heads * head_size) : head_size; + + const int row_step = rows_per_block * gridDim.x; + for (int s = blockIdx.x * rows_per_block + row_in_block; s < limit; s += row_step) { + const LoadVec packed = *reinterpret_cast( + quantized_data + in_base + static_cast(s) * in_row_stride); + const auto* raw = reinterpret_cast(&packed); + + alignas(16) T out[kVecSize]; +#pragma unroll + for (int i = 0; i < kVecSize; i++) { + float value; +#ifdef USE_FP8_KV_CACHE + if constexpr (std::is_same::value) { + value = static_cast(raw[i]); + } else +#endif + { + value = static_cast(reinterpret_cast(raw)[i]); + } + out[i] = static_cast(value * scales[i]); + } + + StoreVec* dst = reinterpret_cast(dequantized_data + out_base + + static_cast(s) * head_size); +#pragma unroll + for (int i = 0; i < kStoresPerVec; i++) { + dst[i] = *reinterpret_cast(&out[i * kElemsPerStore]); + } + } +} + +// Dequantizes the key and the value cache in a single launch. Falls back to two independent +// LaunchDequantizeKV calls for cache formats the fused kernel does not cover (INT4, non-16-bit +// output, head sizes that do not tile into the fused block shape). +template +Status LaunchDequantizeKVPair(cudaStream_t stream, T* k_dequantized, T* v_dequantized, + const T_QUANT* k_quantized, const T_QUANT* v_quantized, + const T_SCALE* k_scale, const T_SCALE* v_scale, + int batch_size, int num_heads, int cache_sequence_length, + int head_size, int bit_width, + KVQuantizationType k_quant_type, KVQuantizationType v_quant_type, + bool is_input_bsnh, const int* valid_seq_lens) { + if (cache_sequence_length == 0) return Status::OK(); + + constexpr int kFusedVecSize = 8; + constexpr int kFusedBlockThreads = 128; + + // Same-binary opt-out so the fused kernel can be A/B'd against the two-launch path. + static const bool fused_disabled = [] { + const char* v = std::getenv("ORT_DISABLE_FUSED_KV_DEQUANT"); + return v != nullptr && v[0] != '\0' && v[0] != '0'; + }(); + + if constexpr (sizeof(T_QUANT) == 1 && sizeof(T) == 2) { + if (!fused_disabled && bit_width == 8 && head_size % kFusedVecSize == 0 && + (head_size / kFusedVecSize) <= kFusedBlockThreads) { + const int rows_per_block = kFusedBlockThreads / (head_size / kFusedVecSize); + const int row_chunks = (cache_sequence_length + rows_per_block - 1) / rows_per_block; + + // Cap the grid so a short sequence in a large cache does not launch a wave of blocks that + // immediately exit; the kernel loops over the remaining chunks. The cap only depends on + // shapes, so the launch configuration stays constant across decode steps (CUDA graph safe). + const int grid_yz = batch_size * num_heads * 2; + const int max_chunks = std::max(1, 2048 / grid_yz); + const dim3 grid(static_cast(std::min(row_chunks, max_chunks)), + static_cast(batch_size * num_heads), 2u); + + DequantizeKVPairKernel + <<>>( + k_dequantized, v_dequantized, k_quantized, v_quantized, k_scale, v_scale, + valid_seq_lens, num_heads, cache_sequence_length, head_size, + k_quant_type, v_quant_type, is_input_bsnh); + + return CUDA_CALL(cudaGetLastError()); + } + } + + ORT_RETURN_IF_ERROR((LaunchDequantizeKV( + stream, k_dequantized, k_quantized, k_scale, nullptr, batch_size, num_heads, + cache_sequence_length, head_size, bit_width, k_quant_type, is_input_bsnh, valid_seq_lens))); + + return LaunchDequantizeKV( + stream, v_dequantized, v_quantized, v_scale, nullptr, batch_size, num_heads, + cache_sequence_length, head_size, bit_width, v_quant_type, is_input_bsnh, valid_seq_lens); +} + // ============================================================================ // Folding per-channel KV dequantization scales into Q and into the attention output. // diff --git a/onnxruntime/contrib_ops/cuda/bert/paged_attention.cc b/onnxruntime/contrib_ops/cuda/bert/paged_attention.cc index 18832597a64ea..7cf64296f6d7e 100644 --- a/onnxruntime/contrib_ops/cuda/bert/paged_attention.cc +++ b/onnxruntime/contrib_ops/cuda/bert/paged_attention.cc @@ -1,6 +1,10 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. +#include +#include +#include + #include "core/providers/cuda/cuda_common.h" #include "core/platform/env_var_utils.h" #include "contrib_ops/cpu/utils/dump_tensor.h" @@ -10,6 +14,8 @@ #include "contrib_ops/cuda/bert/paged_attention_helper.h" #include "contrib_ops/cuda/bert/flash_attention/flash_api.h" #include "contrib_ops/cuda/bert/cutlass_fmha/memory_efficient_attention.h" +#include "contrib_ops/cuda/bert/xqa/xqa_paged_loader.h" +#include "contrib_ops/cuda/llm/common/cuda_runtime_utils.h" using namespace onnxruntime::cuda; using namespace ::onnxruntime::common; @@ -19,23 +25,75 @@ namespace onnxruntime { namespace contrib { namespace cuda { -#define REGISTER_KERNEL_TYPED(T) \ - ONNX_OPERATOR_TYPED_KERNEL_EX( \ - PagedAttention, \ - kMSDomain, \ - 1, \ - T, \ - kCudaExecutionProvider, \ - (*KernelDefBuilder::Create()) \ - .TypeConstraint("T", DataTypeImpl::GetTensorType()) \ - .TypeConstraint("S", DataTypeImpl::GetTensorType()), \ - PagedAttention); - -REGISTER_KERNEL_TYPED(MLFloat16) -REGISTER_KERNEL_TYPED(BFloat16) - -template -PagedAttention::PagedAttention(const OpKernelInfo& info) +#define REGISTER_KERNEL_TYPED(T, TCACHE) \ + ONNX_OPERATOR_TYPED_KERNEL_EX( \ + PagedAttention, \ + kMSDomain, \ + 1, \ + T##_##TCACHE, \ + kCudaExecutionProvider, \ + (*KernelDefBuilder::Create()) \ + .TypeConstraint("T", DataTypeImpl::GetTensorType()) \ + .TypeConstraint("T_CACHE", DataTypeImpl::GetTensorType()) \ + .TypeConstraint("T_KV_SCALE", DataTypeImpl::GetTensorType()) \ + .TypeConstraint("S", DataTypeImpl::GetTensorType()) \ + .InputMemoryType(OrtMemTypeCPUInput, 16), \ + PagedAttention); + +REGISTER_KERNEL_TYPED(MLFloat16, MLFloat16) +REGISTER_KERNEL_TYPED(BFloat16, BFloat16) +REGISTER_KERNEL_TYPED(MLFloat16, int8_t) +REGISTER_KERNEL_TYPED(BFloat16, int8_t) +#if defined(USE_FP8_KV_CACHE) && !defined(DISABLE_FLOAT8_TYPES) +REGISTER_KERNEL_TYPED(MLFloat16, Float8E4M3FN) +REGISTER_KERNEL_TYPED(BFloat16, Float8E4M3FN) +#endif + +// True when TCACHE stores quantized values that need a scale on read/write. +template +constexpr bool IsQuantizedCacheType() { + if constexpr (std::is_same::value) { + return true; +#if defined(USE_FP8_KV_CACHE) && !defined(DISABLE_FLOAT8_TYPES) + } else if constexpr (std::is_same::value) { + return true; +#endif + } else { + return false; + } +} + +// True when TCACHE is the FP8 cache element type. Split out from IsQuantizedCacheType because the +// XQA backend needs to tell the two quantized formats apart, and Float8E4M3FN is not necessarily a +// usable type in every build. +template +constexpr bool IsFp8CacheType() { +#if defined(USE_FP8_KV_CACHE) && !defined(DISABLE_FLOAT8_TYPES) + return std::is_same::value; +#else + return false; +#endif +} + +// The element type TCACHE stores, named as a KVCacheDataType so that an explicit k_cache_dtype / +// v_cache_dtype attribute can be checked against it. +template +constexpr KVCacheDataType CacheStorageDataType() { + if constexpr (std::is_same::value) { + return KVCacheDataType::INT8; +#if defined(USE_FP8_KV_CACHE) && !defined(DISABLE_FLOAT8_TYPES) + } else if constexpr (std::is_same::value) { + return KVCacheDataType::FLOAT8E4M3FN; +#endif + } else if constexpr (std::is_same::value) { + return KVCacheDataType::BFLOAT16; + } else { + return KVCacheDataType::FLOAT16; + } +} + +template +PagedAttention::PagedAttention(const OpKernelInfo& info) : CudaKernel(info) { int64_t num_heads = 0; int64_t kv_num_heads = 0; @@ -47,15 +105,46 @@ PagedAttention::PagedAttention(const OpKernelInfo& info) do_rotary_ = info.GetAttrOrDefault("do_rotary", 0) == 1; rotary_interleaved_ = info.GetAttrOrDefault("rotary_interleaved", 0) == 1; scale_ = info.GetAttrOrDefault("scale", 0.0f); + // scale == 0 selects the 1/sqrt(head_size) default. MLA derives its softmax scale from the + // pre-absorption head width, so that default is silently wrong there and validation requires a + // real value (see docs/contrib_ops/cuda/paged_attention.md §12.6). + has_explicit_scale_ = scale_ != 0.0f; softcap_ = info.GetAttrOrDefault("softcap", 0.0f); + qk_norm_epsilon_ = info.GetAttrOrDefault("qk_norm_epsilon", 1e-6f); + ORT_ENFORCE(std::isfinite(qk_norm_epsilon_) && qk_norm_epsilon_ > 0.0f, + "qk_norm_epsilon must be a positive finite number"); + k_quant_type_ = StringToKVQuantizationType(info.GetAttrOrDefault("k_quant_type", "NONE")); + v_quant_type_ = StringToKVQuantizationType(info.GetAttrOrDefault("v_quant_type", "NONE")); + // Empty (the default) means the cache tensor's own element type is the logical type, which covers + // every format this operator stores today. A non-empty value names a sub-byte logical type packed + // into a uint8 cache, which no build supports yet and is rejected during validation. The string is + // parsed once here; everything downstream compares the enum. + k_cache_dtype_ = StringToKVCacheDataType(info.GetAttrOrDefault("k_cache_dtype", "")); + v_cache_dtype_ = StringToKVCacheDataType(info.GetAttrOrDefault("v_cache_dtype", "")); + + // Multi-head Latent Attention. "SEPARATE" (the default) is the shipped two-cache layout; + // "LATENT" makes value/value_cache absent and aliases V onto the leading v_head_size channels of + // key_cache. Anything else is rejected here rather than silently treated as SEPARATE. + const std::string kv_cache_layout = info.GetAttrOrDefault("kv_cache_layout", "SEPARATE"); + ORT_ENFORCE(kv_cache_layout == "SEPARATE" || kv_cache_layout == "LATENT", + "'kv_cache_layout' must be 'SEPARATE' or 'LATENT', got '", kv_cache_layout, "'"); + is_latent_kv_ = kv_cache_layout == "LATENT"; + v_head_size_ = static_cast(info.GetAttrOrDefault("v_head_size", 0)); + ORT_ENFORCE(v_head_size_ >= 0, "'v_head_size' must be non-negative, got ", v_head_size_); + rotary_offset_ = static_cast(info.GetAttrOrDefault("rotary_offset", 0)); + ORT_ENFORCE(rotary_offset_ >= 0, "'rotary_offset' must be non-negative, got ", rotary_offset_); kernel_options_ = this->GetAttentionKernelOptions(); disable_flash_attention_ = sizeof(T) != 2 || !kernel_options_->UseFlashAttention(); disable_memory_efficient_attention_ = sizeof(T) != 2 || !kernel_options_->UseEfficientAttention(); + disable_paged_decode_ = sizeof(T) != 2 || !kernel_options_->UseDecoderAttention(); + // XQA defaults on for fp16/bf16 activations, matching GroupQueryAttention; ORT_ENABLE_XQA=0 + // disables it explicitly. + enable_xqa_ = sizeof(T) == 2 && (ParseEnvironmentVariableWithDefault("ORT_ENABLE_XQA", 1) != 0); } -template -Status PagedAttention::ComputeInternal(OpKernelContext* context) const { +template +Status PagedAttention::ComputeInternal(OpKernelContext* context) const { auto ort_stream = GetOrtStream(context); const Tensor* query = context->Input(0); @@ -68,11 +157,21 @@ Status PagedAttention::ComputeInternal(OpKernelContext* context) const { const Tensor* block_table = context->Input(7); const Tensor* cos_cache = context->Input(8); const Tensor* sin_cache = context->Input(9); + const Tensor* slot_mapping = context->Input(10); + const Tensor* head_sink = context->Input(11); + const Tensor* q_norm_weight = context->Input(12); + const Tensor* k_norm_weight = context->Input(13); + const Tensor* k_scale = context->Input(14); + const Tensor* v_scale = context->Input(15); + // Resident in CPU memory (see the kernel def's InputMemoryType above). + const Tensor* attention_metadata = context->Input(16); auto& device_prop = GetDeviceProp(); PagedAttentionParameters parameters; typedef typename ToCudaType::MappedType CudaT; - PagedAttentionData data; + typedef typename ToCudaType::MappedType CudaTCache; + constexpr bool kIsQuantizedCache = IsQuantizedCacheType(); + PagedAttentionData data; // Check shapes of inputs to op and set parameters ORT_RETURN_IF_ERROR(paged_attention_helper::CheckInputs(query, @@ -85,11 +184,28 @@ Status PagedAttention::ComputeInternal(OpKernelContext* context) const { block_table, cos_cache, sin_cache, + slot_mapping, + head_sink, + q_norm_weight, + k_norm_weight, + k_scale, + v_scale, + attention_metadata, ¶meters, num_heads_, kv_num_heads_, scale_, softcap_, + qk_norm_epsilon_, + k_quant_type_, + v_quant_type_, + k_cache_dtype_, + v_cache_dtype_, + CacheStorageDataType(), + is_latent_kv_, + v_head_size_, + rotary_offset_, + has_explicit_scale_, device_prop.maxThreadsPerBlock)); parameters.local_window_size = local_window_size_; parameters.do_rotary = do_rotary_; @@ -103,6 +219,9 @@ Status PagedAttention::ComputeInternal(OpKernelContext* context) const { DUMP_STRING("Q num heads = ", parameters.num_heads); DUMP_STRING("KV num heads = ", parameters.kv_num_heads); DUMP_STRING("Head size = ", parameters.head_size); + DUMP_STRING("V head size = ", parameters.v_head_size); + DUMP_STRING("Latent (MLA) KV layout = ", parameters.is_latent_kv); + DUMP_STRING("Rotary offset = ", parameters.rotary_offset); DUMP_STRING("Num blocks = ", parameters.num_blocks); DUMP_STRING("Block size = ", parameters.block_size); DUMP_STRING("Max num blocks per sequence = ", parameters.max_num_blocks_per_seq); @@ -115,10 +234,12 @@ Status PagedAttention::ComputeInternal(OpKernelContext* context) const { "cos_cache and sin_cache must be passed to PagedAttention when do_rotary = 1"); } - // Set output tensor shapes + // Set output tensor shapes. In LATENT mode the output head width is v_head_size, so it is + // narrower than the query head width (512 vs. 576 for DeepSeek-V3). v_hidden_size equals + // hidden_size in every SEPARATE-mode model. TensorShapeVector output_shape(2); output_shape[0] = static_cast(parameters.token_count); - output_shape[1] = static_cast(parameters.hidden_size); + output_shape[1] = static_cast(parameters.v_hidden_size); Tensor* output = context->Output(0, output_shape); TensorShapeVector key_cache_out_shape(4); @@ -128,17 +249,21 @@ Status PagedAttention::ComputeInternal(OpKernelContext* context) const { key_cache_out_shape[3] = static_cast(parameters.head_size); Tensor* key_cache_out = context->Output(1, key_cache_out_shape); - TensorShapeVector value_cache_out_shape(4); - value_cache_out_shape[0] = static_cast(parameters.num_blocks); - value_cache_out_shape[1] = static_cast(parameters.block_size); - value_cache_out_shape[2] = static_cast(parameters.kv_num_heads); - value_cache_out_shape[3] = static_cast(parameters.head_size); - Tensor* value_cache_out = context->Output(2, value_cache_out_shape); + // LATENT has a single physical cache, so there is no value_cache_out to produce. + Tensor* value_cache_out = nullptr; + if (!parameters.is_latent_kv) { + TensorShapeVector value_cache_out_shape(4); + value_cache_out_shape[0] = static_cast(parameters.num_blocks); + value_cache_out_shape[1] = static_cast(parameters.block_size); + value_cache_out_shape[2] = static_cast(parameters.kv_num_heads); + value_cache_out_shape[3] = static_cast(parameters.head_size); + value_cache_out = context->Output(2, value_cache_out_shape); + } - if (key_cache_out != nullptr && key_cache->Data() != key_cache_out->MutableData()) { + if (key_cache_out != nullptr && key_cache->Data() != key_cache_out->MutableData()) { return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "key_cache and key_cache_out must be the same buffer"); - } else if (value_cache_out != nullptr && value_cache->Data() != value_cache_out->MutableData()) { + } else if (value_cache_out != nullptr && value_cache->Data() != value_cache_out->MutableData()) { return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "value_cache and value_cache_out must be the same buffer"); } @@ -149,76 +274,93 @@ Status PagedAttention::ComputeInternal(OpKernelContext* context) const { return Status::OK(); } - // Kernel backend selection — FlashAttention preferred, fall back to MemoryEfficientAttention. + // Kernel backend selection. The choice depends only on static shapes and on the optional + // 'attention_metadata' bounds, never on a device-to-host readback, so it is identical on every + // replay of a captured CUDA Graph (docs/contrib_ops/cuda/paged_attention.md section 4.7). + // + // * FlashAttention (preferred for prefill / mixed batches). + // * PagedDecode: a flash-decoding style kernel that reads the paged cache in place and + // dequantizes in registers. + // * MemoryEfficientAttention: the general fallback; gathers pages into a dense buffer. + // + // The vendored FlashAttention paged kernel loads a whole kBlockN x head_size K/V tile using a + // single (page, offset) pair, so a tile must never straddle a page boundary: block_size has to be + // a multiple of kBlockN. kBlockN is fixed by head_size in run_mha_fwd_splitkv_dispatch. When the + // model uses a smaller page than that, we fall back to another backend (both of which accept any + // block_size) rather than rejecting the model. + // + // A quantized cache is exempt: FlashAttention cannot read a quantized page at all, so that path + // dequantizes the live context into a dense buffer and uses the non-paged varlen entry point, + // which has no page-alignment requirement. + const int flash_min_block_size = + parameters.head_size <= 64 ? 256 : (parameters.head_size <= 128 ? 128 : 64); + const bool flash_block_size_ok = kIsQuantizedCache || (parameters.block_size % flash_min_block_size) == 0; + + // LATENT (absorbed MLA) has exactly one eligible backend: neither FlashAttention nor the CUTLASS + // fMHA wrapper supports v_head_size != head_size or a head_size of 576, and the paged decode + // kernel assumes a separate value cache. See docs/contrib_ops/cuda/paged_attention.md §12.7. + const bool use_latent_attention = parameters.is_latent_kv; + if (use_latent_attention) { + const size_t latent_smem = GetPagedLatentSharedMemoryBytes(parameters.head_size, parameters.v_head_size); + if (latent_smem > static_cast(device_prop.sharedMemPerBlock)) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "PagedAttention: the unfused MLA backend needs ", latent_smem, + " bytes of shared memory for head_size=", parameters.head_size, + " and v_head_size=", parameters.v_head_size, ", but the device provides ", + device_prop.sharedMemPerBlock, " bytes per block."); + } + } + #if USE_FLASH_ATTENTION - bool use_flash_attention = !disable_flash_attention_ && - onnxruntime::flash::is_supported(device_prop, - parameters.head_size, - parameters.num_heads, - parameters.kv_num_heads); + const bool flash_eligible = !use_latent_attention && + !disable_flash_attention_ && + flash_block_size_ok && + onnxruntime::flash::is_supported(device_prop, + parameters.head_size, + parameters.num_heads, + parameters.kv_num_heads); #else - constexpr bool use_flash_attention = false; + const bool flash_eligible = false; #endif #if USE_MEMORY_EFFICIENT_ATTENTION const int sm = device_prop.major * 10 + device_prop.minor; const bool is_half = std::is_same::value; const bool is_bf16 = std::is_same::value; - bool use_memory_efficient_attention = - !use_flash_attention && + const bool mea_eligible = + !use_latent_attention && + !flash_eligible && !disable_memory_efficient_attention_ && has_memory_efficient_attention(sm, is_half, is_bf16, parameters.head_size, parameters.head_size); #else - constexpr bool use_memory_efficient_attention = false; + const bool mea_eligible = false; #endif - if (!use_flash_attention && !use_memory_efficient_attention) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "PagedAttention requires FlashAttention (sm>=80, fp16/bf16) or " - "MemoryEfficientAttention (fp16 sm>=53, bf16 sm>=80, head_size<=1024 and %8==0) " - "to be available. Check ORT_DISABLE_FLASH_ATTENTION / " - "ORT_DISABLE_MEMORY_EFFICIENT_ATTENTION env vars and dtype/head_size."); - } - - // Scratch buffers common to both backends. - size_t softmax_lse_bytes = 0; -#if USE_FLASH_ATTENTION - if (use_flash_attention) { - softmax_lse_bytes = onnxruntime::flash::get_softmax_lse_size(parameters.token_count, - parameters.num_heads); - } -#endif - auto softmax_lse_buffer = GetScratchBuffer(softmax_lse_bytes, GetComputeStream(context)); + // The decode kernel keeps Q, the running accumulator and one KV tile in shared memory, which + // bounds the head size it can serve. + const bool decode_eligible = + !use_latent_attention && + !disable_paged_decode_ && + GetPagedDecodeSharedMemoryBytes(parameters.head_size) <= static_cast(device_prop.sharedMemPerBlock); size_t cumulative_seqlens_kv_bytes = sizeof(int) * (parameters.batch_size + 1); auto cumulative_seqlens_kv_buffer = GetScratchBuffer(cumulative_seqlens_kv_bytes, GetComputeStream(context)); int* cumulative_seqlens_kv_ptr = reinterpret_cast(cumulative_seqlens_kv_buffer.get()); + // The fused prologue (QK-Norm and/or rotary) writes densified Q and K into the workspace, so it + // needs room for both. Plain packed-QKV only needs to densify Q. + const bool needs_qk_prologue = do_rotary_ || parameters.use_qk_norm; size_t workspace_buffer_bytes = 0; - if (do_rotary_) { + if (needs_qk_prologue) { workspace_buffer_bytes = sizeof(T) * parameters.token_count * (parameters.hidden_size + parameters.kv_hidden_size); } else if (parameters.is_packed_qkv) { workspace_buffer_bytes = sizeof(T) * parameters.token_count * parameters.hidden_size; } auto workspace_buffer = GetScratchBuffer(workspace_buffer_bytes, GetComputeStream(context)); - // Populate cumulative_seqlens_kv for both backends. The MEA path additionally needs - // the last element on the host to size the tight gather buffers, so we D->H sync below. - // - // LaunchGetCumulativeSeqlensKV uses a per-block cub::BlockScan with a block size of 256 - // and launches (batch_size + 255) / 256 blocks, so blocks scan independently. Enforce - // batch_size <= 256 so the cumulative sum is correct; a larger batch would silently - // produce wrong KV offsets. (A future grid-wide scan could lift this limit.) - constexpr int kMaxBatchSizeForCumulativeSeqlensKV = 256; - if (parameters.batch_size > kMaxBatchSizeForCumulativeSeqlensKV) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "PagedAttention currently supports batch_size <= ", - kMaxBatchSizeForCumulativeSeqlensKV, - " (LaunchGetCumulativeSeqlensKV limitation); got batch_size=", - parameters.batch_size, "."); - } - + // Populate cumulative_seqlens_kv for all backends. Every kernel that needs a per-sequence KV + // length reads it from here on device; the host only ever uses upper bounds. cudaStream_t cuda_stream = static_cast(ort_stream.get()->GetHandle()); ORT_RETURN_IF_ERROR(LaunchGetCumulativeSeqlensKV( cumulative_seqlens_kv_ptr, @@ -228,53 +370,273 @@ Status PagedAttention::ComputeInternal(OpKernelContext* context) const { int total_kv_tokens = 0; int max_query_len = 0; + int max_kv_len = 0; IAllocatorUniquePtr gathered_key_buffer; IAllocatorUniquePtr gathered_value_buffer; IAllocatorUniquePtr fmha_buffer; - // Compute max_query_len on the host for both FA and MEA. The previous - // `token_count - batch_size + 1` heuristic underestimates (or goes - // non-positive) when batches have zero new tokens. MEA additionally needs - // total_kv_tokens to size gather buffers. - if (use_flash_attention || use_memory_efficient_attention) { + // 'attention_metadata' supplies replay-wide *upper bounds* on the per-sequence query and KV + // lengths (docs/contrib_ops/cuda/paged_attention.md section 4.7). Bounds are all the backends + // need from the host: they only select the kernel, size launch dimensions and size workspaces. + // Every per-sequence length that enters a mask is re-read from device memory by the kernel + // itself, which is what keeps a captured graph correct as the sequences grow. + // + // When no metadata is supplied the bounds degrade to the static capacities: at most token_count + // query tokens can belong to a single sequence, and a sequence can address at most + // block_table.shape[1] * block_size cached tokens. Those are valid but loose, so they cost some + // empty thread blocks rather than correctness. + const int max_kv_len_capacity = parameters.max_num_blocks_per_seq * parameters.block_size; + const bool has_metadata_bounds = attention_metadata != nullptr; + int max_query_len_bound = parameters.token_count; + int max_kv_len_bound = max_kv_len_capacity; + if (has_metadata_bounds) { + const int* metadata = attention_metadata->Data(); + const int metadata_query_bound = metadata[0]; + const int metadata_kv_bound = metadata[1]; + if (metadata_query_bound < 0 || metadata_kv_bound < 0) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "PagedAttention: 'attention_metadata' entries must be non-negative, got [", + metadata_query_bound, ", ", metadata_kv_bound, "]. Use 0 for 'unknown'."); + } + // Clamp each bound to the static limit it can never exceed, so an over-large (or unknown) + // bound degrades to the same sizing we would use with no metadata at all. + if (metadata_query_bound > 0 && metadata_query_bound < max_query_len_bound) { + max_query_len_bound = metadata_query_bound; + } + if (metadata_kv_bound > 0 && metadata_kv_bound < max_kv_len_bound) { + max_kv_len_bound = metadata_kv_bound; + } + } + + // Backend selection from static shapes alone. + // + // token_count <= batch_size is a *heuristic* for "at most one new token per sequence": it does not + // prove that any particular sequence contributes at most one, only that they do on average. That + // is fine because PagedDecodeSplitKV is indexed by global query token and resolves each token's + // sequence and position from cumulative_seqlens_q on device, so it is correct for arbitrary ragged + // input. The inequality rather than equality matters for continuous batching, where a scheduled + // sequence may contribute no token on a given step. + // Prefer it only where it is a clear win: a quantized cache (it avoids materializing and + // dequantizing the whole live context) or when FlashAttention is unavailable (it beats the + // dense-gather MemoryEfficientAttention fallback). + const bool decode_shaped = parameters.token_count <= parameters.batch_size; + const bool use_paged_decode = decode_eligible && decode_shaped && (kIsQuantizedCache || !flash_eligible); + const bool use_flash_attention = flash_eligible && !use_paged_decode; + const bool use_memory_efficient_attention = mea_eligible && !use_paged_decode; + + // Both gather-based backends need a dense KV staging buffer when the cache is quantized + // (FlashAttention cannot read a quantized page, and the CUTLASS kernel is not paged at all). + const bool needs_dense_kv = use_memory_efficient_attention || (use_flash_attention && kIsQuantizedCache); + // The dense buffer keeps the grouped layout for FlashAttention (it does GQA internally) and is + // GQA-expanded for the CUTLASS kernel. + const int gathered_num_heads = use_memory_efficient_attention ? parameters.num_heads : parameters.kv_num_heads; + + // Within the decode-on-a-quantized-cache case, prefer XQA: it is the same tensor-core kernel + // GroupQueryAttention uses, reading the paged cache in place, and is roughly an order of + // magnitude faster than the portable PagedDecodeSplitKV kernel. Constraints come from the + // compiled instantiations (head_size, query/KV group size), from XQA itself (no softcap) and + // from the page remap (a PagedAttention block must split into whole 128-token XQA pages). + // XQA additionally lays its output out as one row per batch index, so unlike PagedDecodeSplitKV + // it needs *proof* that every sequence contributes exactly one token, not just the shape + // heuristic. token_count == batch_size rules out a sequence contributing none; max_query_len == 1, + // from the metadata bound or from the readback below, then rules out any contributing two. + bool xqa_candidate = false; + if (use_paged_decode && enable_xqa_ && kIsQuantizedCache && + parameters.token_count == parameters.batch_size) { + const int group_size = parameters.num_heads / parameters.kv_num_heads; + const bool is_fp8_cache = IsFp8CacheType(); + const auto is_supported_quant_type = [](KVQuantizationType t) { + return t == KVQuantizationType::PER_TENSOR || t == KVQuantizationType::PER_CHANNEL; + }; + xqa_candidate = + device_prop.major >= 8 && + parameters.softcap == 0.0f && + (parameters.head_size == 64 || parameters.head_size == 128) && + (group_size == 4 || group_size == 8 || group_size == 16 || group_size == 32) && + (parameters.block_size % kXqaTokensPerPage) == 0 && + is_supported_quant_type(k_quant_type_) && is_supported_quant_type(v_quant_type_) && + // FP8 arithmetic in the XQA kernel needs Ada (sm_89) or Hopper+. + (!is_fp8_cache || device_prop.major >= 9 || (device_prop.major == 8 && device_prop.minor == 9)); + } + + // Obtaining the exact lengths from the device means copying the two cumulative arrays back and + // blocking the host until they land, which drains everything already queued on the compute + // stream -- once per PagedAttention node, so once per layer per decoded token. It also makes the + // node impossible to capture into a CUDA Graph, since a stream synchronization is not a + // capturable operation. + // + // Nothing above needed it, and every remaining consumer is happy with an upper bound, with two + // exceptions that only arise when the caller supplied no bounds at all: + // + // 1. The gather backends stage the live context into a buffer indexed by the *exact* + // total_kv_tokens. Sizing it by batch_size * max_kv_len_bound is replay-invariant and + // correct, but with no metadata that bound is the full block-table capacity, which would be + // a large over-allocation for a short prefill. Read the exact value back instead. + // 2. XQA needs the one-token-per-sequence proof described above. + // + // Neither case can occur on a capturable step: a captured step is decode-shaped on a paged cache + // (so no gather runs), and a producer that captures must supply 'attention_metadata' anyway -- + // its bounds are the only replay-safe source of per-step information. The synchronization is + // therefore gone for every configuration CUDA Graphs can reach, including an unquantized cache. + const bool needs_readback = !has_metadata_bounds && (needs_dense_kv || xqa_candidate); + + if (!needs_readback) { + max_query_len = max_query_len_bound; + max_kv_len = max_kv_len_bound; + // Upper bound: no sequence holds more than max_kv_len_bound cached tokens. Read only by the + // gather backends, to size and to launch the staging buffer; the gather kernel derives each + // token's sequence from cumulative_seqlens_kv on device and skips indices past the real end. + total_kv_tokens = static_cast(std::min( + std::numeric_limits::max(), + static_cast(parameters.batch_size) * max_kv_len_bound)); + } else { const int kCumulativeCount = parameters.batch_size + 1; auto cum_q_pinned = this->AllocateBufferOnCPUPinned(kCumulativeCount); - IAllocatorUniquePtr cum_kv_pinned; + auto cum_kv_pinned = this->AllocateBufferOnCPUPinned(kCumulativeCount); CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(cum_q_pinned.get(), reinterpret_cast(cumulative_seqlens_q->Data()), sizeof(int) * kCumulativeCount, cudaMemcpyDeviceToHost, cuda_stream)); - if (use_memory_efficient_attention) { - cum_kv_pinned = this->AllocateBufferOnCPUPinned(kCumulativeCount); - CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(cum_kv_pinned.get(), cumulative_seqlens_kv_ptr, - sizeof(int) * kCumulativeCount, cudaMemcpyDeviceToHost, cuda_stream)); - } + CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(cum_kv_pinned.get(), cumulative_seqlens_kv_ptr, + sizeof(int) * kCumulativeCount, cudaMemcpyDeviceToHost, cuda_stream)); CUDA_RETURN_IF_ERROR(cudaStreamSynchronize(cuda_stream)); for (int i = 0; i < parameters.batch_size; ++i) { const int q_len_i = cum_q_pinned.get()[i + 1] - cum_q_pinned.get()[i]; if (q_len_i > max_query_len) { max_query_len = q_len_i; } - } - if (use_memory_efficient_attention) { - total_kv_tokens = cum_kv_pinned.get()[parameters.batch_size]; - if (total_kv_tokens == 0) { - return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, - "PagedAttention MEA fallback: total_kv_tokens is zero for non-empty input."); + const int kv_len_i = cum_kv_pinned.get()[i + 1] - cum_kv_pinned.get()[i]; + if (kv_len_i > max_kv_len) { + max_kv_len = kv_len_i; } - if (total_kv_tokens < 0) { - return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, - "PagedAttention MEA fallback: total_kv_tokens is negative (", total_kv_tokens, ")."); + } + total_kv_tokens = cum_kv_pinned.get()[parameters.batch_size]; + if (total_kv_tokens <= 0) { + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, + "PagedAttention: total_kv_tokens is not positive (", total_kv_tokens, + ") for non-empty input."); + } + } + + bool use_xqa_decode = xqa_candidate && max_query_len == 1; + if (use_xqa_decode) { + // The kernel's dynamic shared-memory request is fixed at compile time for its target SM and + // can exceed the opt-in limit of the device actually running it (e.g. a kernel JIT-compiled + // from sm_90 PTX onto consumer Blackwell). Query it once per node -- it depends only on + // head_size and the group size -- and fall back when it does not fit. The query is a + // cudaMemcpyFromSymbol, which synchronizes and is therefore illegal during graph capture, so + // skip XQA for that run and leave the result unresolved for a later non-capturing run. + int xqa_smem_ok = xqa_shared_memory_ok_.load(std::memory_order_relaxed); + if (xqa_smem_ok < 0) { + if (!onnxruntime::llm::common::isCapturing(cuda_stream)) { + const size_t required_smem = GetXQAPagedRequiredSharedMemoryBytes( + device_prop, parameters.head_size, parameters.num_heads, parameters.kv_num_heads, + IsFp8CacheType() ? XqaQuantType::kFp8 : XqaQuantType::kInt8, + std::is_same::value); + xqa_smem_ok = (required_smem == 0 || required_smem <= device_prop.sharedMemPerBlockOptin) ? 1 : 0; + xqa_shared_memory_ok_.store(xqa_smem_ok, std::memory_order_relaxed); + } else { + xqa_smem_ok = 0; } } + use_xqa_decode = (xqa_smem_ok != 0); } + DUMP_STRING("Backend = ", use_latent_attention ? "latent" + : use_xqa_decode ? "paged decode (XQA)" + : use_paged_decode ? "paged decode" + : use_flash_attention ? "flash attention" + : "memory efficient attention"); -#if USE_MEMORY_EFFICIENT_ATTENTION - if (use_memory_efficient_attention) { + if (!use_latent_attention && !use_paged_decode && !use_flash_attention && !use_memory_efficient_attention) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "PagedAttention requires FlashAttention (sm>=80, fp16/bf16, block_size a multiple of ", + flash_min_block_size, " for head_size ", parameters.head_size, + "), MemoryEfficientAttention (fp16 sm>=53, bf16 sm>=80, head_size<=1024 and %8==0), " + "or the paged decode kernel (fp16/bf16, decode-shaped batch, i.e. " + "token_count <= batch_size; this step has token_count=", + parameters.token_count, " and batch_size=", parameters.batch_size, + ") to be available. Check ORT_DISABLE_FLASH_ATTENTION / " + "ORT_DISABLE_MEMORY_EFFICIENT_ATTENTION / ORT_DISABLE_DECODER_ATTENTION env vars and " + "dtype/head_size/block_size."); + } + + // The attention-sink epilogue rescales the output using FlashAttention's log-sum-exp, which the + // CUTLASS memory-efficient kernel does not expose. The decode kernel folds the sink straight into + // its softmax denominator, so only the MEA path has to fail loudly here. + if (parameters.use_smooth_softmax && use_memory_efficient_attention) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "PagedAttention: 'head_sink' is only supported by the FlashAttention " + "and paged decode backends, but the MemoryEfficientAttention backend was selected " + "(head_size=", + parameters.head_size, ", block_size=", parameters.block_size, + "). FlashAttention requires sm>=80, fp16/bf16 and block_size a multiple of ", + flash_min_block_size, "."); + } + + size_t softmax_lse_bytes = 0; +#if USE_FLASH_ATTENTION + if (use_flash_attention) { + softmax_lse_bytes = onnxruntime::flash::get_softmax_lse_size(parameters.token_count, + parameters.num_heads); + } +#endif + auto softmax_lse_buffer = GetScratchBuffer(softmax_lse_bytes, GetComputeStream(context)); + + if (needs_dense_kv) { const size_t gather_elems = static_cast(total_kv_tokens) * - parameters.num_heads * parameters.head_size; + gathered_num_heads * parameters.head_size; gathered_key_buffer = GetScratchBuffer(sizeof(T) * gather_elems, GetComputeStream(context)); gathered_value_buffer = GetScratchBuffer(sizeof(T) * gather_elems, GetComputeStream(context)); + } + + // Split-KV workspaces for the decode kernel: one partial (accumulator, max, denominator) per + // split. Splitting only pays off when there are too few (token, head) pairs to fill the GPU. + int num_splits = 1; + IAllocatorUniquePtr decode_partial_out_buffer; + IAllocatorUniquePtr decode_partial_max_buffer; + IAllocatorUniquePtr decode_partial_sum_buffer; + if (use_paged_decode && !use_xqa_decode) { + num_splits = ComputePagedDecodeSplits(parameters.token_count, parameters.num_heads, max_kv_len, + device_prop.multiProcessorCount); + const size_t rows = static_cast(num_splits) * parameters.token_count * parameters.num_heads; + decode_partial_out_buffer = + GetScratchBuffer(sizeof(float) * rows * parameters.head_size, GetComputeStream(context)); + decode_partial_max_buffer = GetScratchBuffer(sizeof(float) * rows, GetComputeStream(context)); + decode_partial_sum_buffer = GetScratchBuffer(sizeof(float) * rows, GetComputeStream(context)); + } + // XQA scratch: semaphores + the multi-block (Flash Decoding) partials, plus the expanded page + // table, the optional pre-scaled Q copy and the fp32 attention sinks. + IAllocatorUniquePtr xqa_workspace_buffer; + IAllocatorUniquePtr xqa_page_table_buffer; + IAllocatorUniquePtr xqa_query_buffer; + IAllocatorUniquePtr xqa_head_sink_buffer; + size_t xqa_workspace_bytes = 0; + int xqa_max_pages_per_seq = 0; + if (use_xqa_decode) { + const int pages_per_block = parameters.block_size / kXqaTokensPerPage; + xqa_max_pages_per_seq = parameters.max_num_blocks_per_seq * pages_per_block; + xqa_workspace_bytes = GetXQAScratchSize( + device_prop, parameters.batch_size, parameters.num_heads, parameters.kv_num_heads, + parameters.head_size, xqa_max_pages_per_seq * kXqaTokensPerPage, + IsFp8CacheType() ? XqaQuantType::kFp8 : XqaQuantType::kInt8, + std::is_same::value); + xqa_workspace_buffer = GetScratchBuffer(xqa_workspace_bytes, GetComputeStream(context)); + xqa_page_table_buffer = GetScratchBuffer( + sizeof(int) * static_cast(parameters.batch_size) * xqa_max_pages_per_seq, + GetComputeStream(context)); + if (k_quant_type_ == KVQuantizationType::PER_CHANNEL) { + xqa_query_buffer = GetScratchBuffer( + sizeof(T) * static_cast(parameters.batch_size) * parameters.num_heads * parameters.head_size, + GetComputeStream(context)); + } + if (parameters.use_smooth_softmax && head_sink != nullptr) { + xqa_head_sink_buffer = GetScratchBuffer(sizeof(float) * parameters.num_heads, + GetComputeStream(context)); + } + } + +#if USE_MEMORY_EFFICIENT_ATTENTION + if (use_memory_efficient_attention) { if (MemoryEfficientAttentionParams::need_workspace(parameters.head_size, sizeof(T) == sizeof(float))) { // MEA output accumulator is float32 regardless of input dtype (see GQA pattern at // group_query_attention.cc:482); use sizeof(float), not sizeof(T). @@ -290,6 +652,7 @@ Status PagedAttention::ComputeInternal(OpKernelContext* context) const { AttentionKernelDebugInfo debug_info; debug_info.use_flash_attention = use_flash_attention; debug_info.use_efficient_attention = use_memory_efficient_attention; + debug_info.use_decoder_attention = use_paged_decode; debug_info.Print("PagedAttention", this->Node().Name(), @@ -301,17 +664,29 @@ Status PagedAttention::ComputeInternal(OpKernelContext* context) const { data.query = reinterpret_cast(query->Data()); data.key = key == nullptr ? nullptr : reinterpret_cast(key->Data()); data.value = value == nullptr ? nullptr : reinterpret_cast(value->Data()); - data.key_cache = reinterpret_cast(const_cast(key_cache->Data())); - data.value_cache = reinterpret_cast(const_cast(value_cache->Data())); + data.key_cache = reinterpret_cast(const_cast(key_cache->Data())); + // Absent in LATENT mode, where V is a slice of key_cache and ReshapeAndCache skips the V store. + data.value_cache = value_cache == nullptr + ? nullptr + : reinterpret_cast(const_cast(value_cache->Data())); + data.k_scale = k_scale == nullptr ? nullptr : k_scale->Data(); + data.v_scale = v_scale == nullptr ? nullptr : v_scale->Data(); data.cumulative_seqlens_q = reinterpret_cast(cumulative_seqlens_q->Data()); data.past_seqlens = reinterpret_cast(past_seqlens->Data()); data.cumulative_seqlens_kv = cumulative_seqlens_kv_ptr; data.block_table = reinterpret_cast(block_table->Data()); + data.slot_mapping = slot_mapping == nullptr ? nullptr : reinterpret_cast(slot_mapping->Data()); + data.head_sink = head_sink == nullptr ? nullptr : reinterpret_cast(head_sink->Data()); + data.q_norm_weight = q_norm_weight == nullptr ? nullptr : reinterpret_cast(q_norm_weight->Data()); + data.k_norm_weight = k_norm_weight == nullptr ? nullptr : reinterpret_cast(k_norm_weight->Data()); data.output = reinterpret_cast(output->MutableData()); data.use_flash_attention = use_flash_attention; data.use_memory_efficient_attention = use_memory_efficient_attention; + data.use_paged_decode = use_paged_decode; + data.use_xqa_decode = use_xqa_decode; if (softmax_lse_buffer != nullptr) { - data.softmax_lse = reinterpret_cast(softmax_lse_buffer.get()); + // FlashAttention always writes fp32 log-sum-exp, independent of T. + data.softmax_lse = reinterpret_cast(softmax_lse_buffer.get()); } if (workspace_buffer != nullptr) { data.workspace_buffer = reinterpret_cast(workspace_buffer.get()); @@ -320,19 +695,33 @@ Status PagedAttention::ComputeInternal(OpKernelContext* context) const { data.cos_cache = reinterpret_cast(cos_cache->Data()); data.sin_cache = reinterpret_cast(sin_cache->Data()); } - data.max_query_len = max_query_len; // consumed by both FA and MEA - if (use_memory_efficient_attention) { + data.max_query_len = max_query_len; // consumed by all backends + if (needs_dense_kv) { data.gathered_key = reinterpret_cast(gathered_key_buffer.get()); data.gathered_value = reinterpret_cast(gathered_value_buffer.get()); - if (fmha_buffer != nullptr) { - data.fmha_buffer = reinterpret_cast(fmha_buffer.get()); - } data.total_kv_tokens = total_kv_tokens; + data.max_kv_len = max_kv_len; + } + if (use_paged_decode && !use_xqa_decode) { + data.decode_partial_out = reinterpret_cast(decode_partial_out_buffer.get()); + data.decode_partial_max = reinterpret_cast(decode_partial_max_buffer.get()); + data.decode_partial_sum = reinterpret_cast(decode_partial_sum_buffer.get()); + data.num_splits = num_splits; + } + if (use_xqa_decode) { + data.xqa_workspace = xqa_workspace_buffer.get(); + data.xqa_workspace_size = xqa_workspace_bytes; + data.xqa_page_table = reinterpret_cast(xqa_page_table_buffer.get()); + data.xqa_query = reinterpret_cast(xqa_query_buffer.get()); + data.xqa_head_sink = reinterpret_cast(xqa_head_sink_buffer.get()); + } + if (use_memory_efficient_attention && fmha_buffer != nullptr) { + data.fmha_buffer = reinterpret_cast(fmha_buffer.get()); } cublasHandle_t cublas = GetCublasHandle(context); - return QkvToContext( + return QkvToContext( device_prop, cublas, ort_stream.get(), parameters, data); } diff --git a/onnxruntime/contrib_ops/cuda/bert/paged_attention.h b/onnxruntime/contrib_ops/cuda/bert/paged_attention.h index 027141f02b9ae..c537a8fb151e4 100644 --- a/onnxruntime/contrib_ops/cuda/bert/paged_attention.h +++ b/onnxruntime/contrib_ops/cuda/bert/paged_attention.h @@ -3,6 +3,7 @@ #pragma once +#include #include #include "core/providers/cuda/cuda_kernel.h" #include "contrib_ops/cuda/bert/paged_attention_impl.h" @@ -14,7 +15,9 @@ namespace cuda { using namespace onnxruntime::cuda; -template +// T is the activation type (float16 / bfloat16). TCACHE is the paged KV cache element type: +// the same as T for an unquantized cache, or int8_t / Float8E4M3FN for a quantized one. +template class PagedAttention final : public CudaKernel { public: PagedAttention(const OpKernelInfo& info); @@ -28,8 +31,29 @@ class PagedAttention final : public CudaKernel { bool rotary_interleaved_; float scale_; float softcap_; + float qk_norm_epsilon_; + KVQuantizationType k_quant_type_; + KVQuantizationType v_quant_type_; + // Logical element type stored in the cache when it cannot be expressed by the cache tensor's own + // element type (sub-byte formats packed into uint8). DEFAULT uses the tensor's element type. The + // attribute string is parsed once here so the hot path only compares enums. + KVCacheDataType k_cache_dtype_; + KVCacheDataType v_cache_dtype_; + // Multi-head Latent Attention (docs/contrib_ops/cuda/paged_attention.md §12). is_latent_kv_ comes + // from kv_cache_layout == "LATENT"; v_head_size_ == 0 means "same as head_size". + bool is_latent_kv_; + int v_head_size_; + int rotary_offset_; + bool has_explicit_scale_; bool disable_flash_attention_; bool disable_memory_efficient_attention_; + bool disable_paged_decode_; + // Tensor-core XQA decode kernel for a quantized paged cache. Defaults on; ORT_ENABLE_XQA=0 + // disables it and falls back to the portable PagedDecodeSplitKV kernel. + bool enable_xqa_; + // -1 = not yet resolved, 0 = the kernel needs more shared memory than this device allows, + // 1 = it fits. Resolved once per node because it only depends on head_size / group size. + mutable std::atomic xqa_shared_memory_ok_{-1}; const AttentionKernelOptions* kernel_options_; }; diff --git a/onnxruntime/contrib_ops/cuda/bert/paged_attention_helper.h b/onnxruntime/contrib_ops/cuda/bert/paged_attention_helper.h index 6fb8969aa9d0a..456ca9a3661ec 100644 --- a/onnxruntime/contrib_ops/cuda/bert/paged_attention_helper.h +++ b/onnxruntime/contrib_ops/cuda/bert/paged_attention_helper.h @@ -64,6 +64,58 @@ Status Check_Q_K_V(const T* query, const T* key, const T* value, const int num_h return Status::OK(); } +// LATENT (absorbed MLA) mode: 'query' is the absorbed query, 'key' is the latent row +// [compressed_kv; k_pe] shared by all heads, and 'value' is absent because V is the leading +// v_head_size channels of the same latent row. See docs/contrib_ops/cuda/paged_attention.md §12. +template +Status Check_Q_K_Latent(const T* query, const T* key, const T* value, const int num_heads, const int kv_num_heads, + int& token_count, int& q_hidden_size, int& kv_hidden_size, int& head_size) { + if (key == nullptr) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Input 'key' is required when 'kv_cache_layout' is 'LATENT'."); + } + if (value != nullptr) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Input 'value' must be absent when 'kv_cache_layout' is 'LATENT': the value of every " + "head is the leading 'v_head_size' channels of the latent key."); + } + + const auto& query_dims = query->Shape().GetDims(); + if (query_dims.size() != 2) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 'query' is expected to have 2 dimensions, got ", + query_dims.size()); + } + token_count = static_cast(query_dims[0]); + q_hidden_size = static_cast(query_dims[1]); + if (q_hidden_size % num_heads != 0) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Input 'query' hidden size must be a multiple of num_heads. Got ", q_hidden_size, + " % ", num_heads, " == ", q_hidden_size % num_heads); + } + head_size = q_hidden_size / num_heads; + if (head_size % 8 != 0) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "head_size must be a multiple of 8. Got head_size % 8 == ", head_size % 8); + } + + const auto& key_dims = key->Shape().GetDims(); + if (key_dims.size() != 2) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 'key' is expected to have 2 dimensions, got ", + key_dims.size()); + } + if (token_count != key_dims[0]) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Input 'query' and 'key' shall have same dim 0 (token count)"); + } + kv_hidden_size = static_cast(key_dims[1]); + if (kv_hidden_size != kv_num_heads * head_size) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Input 'key' is expected to have hidden size kv_num_heads * head_size = ", + kv_num_heads * head_size, " in 'LATENT' mode, got ", kv_hidden_size); + } + return Status::OK(); +} + template Status Check_QKV(const T* packed_qkv, const T* value, const int num_heads, const int kv_num_heads, int& token_count, int& q_hidden_size, int& kv_hidden_size, int& head_size) { @@ -88,29 +140,51 @@ Status Check_QKV(const T* packed_qkv, const T* value, const int num_heads, const return Status::OK(); } +// `value_cache` is null only in LATENT mode, where V aliases the leading channels of `key_cache` +// and there is no second physical cache to validate. template Status CheckKVCache(const T* key_cache, const T* value_cache, const int kv_num_heads, const int head_size, int& num_blocks, int& block_size) { const auto& key_cache_dims = key_cache->Shape().GetDims(); - const auto& value_cache_dims = value_cache->Shape().GetDims(); if (key_cache_dims.size() != 4) { return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 'key_cache' is expected to have 4 dimensions, got ", key_cache_dims.size()); } - if (value_cache_dims.size() != 4) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "Input 'value_cache' is expected to have 4 dimensions, got ", - value_cache_dims.size()); - } num_blocks = static_cast(key_cache_dims[0]); block_size = static_cast(key_cache_dims[1]); - // TODO(aciddelgado): block size multiple of 8 - if (block_size % 256 != 0) { + // The op itself only needs the block size to be a power of two >= 16 (the granularity every + // serving framework uses). The vendored FlashAttention paged kernel has a stricter, + // head-size-dependent requirement (a kBlockN tile must not straddle a page); that is enforced at + // backend-selection time in paged_attention.cc, which falls back to the gather-based + // memory-efficient path instead of rejecting the model. See docs/contrib_ops/cuda/paged_attention.md §18. + if (block_size < 16 || (block_size & (block_size - 1)) != 0) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "block_size must be a power of two and at least 16. Got block_size == ", + block_size); + } + + if (key_cache_dims[2] != kv_num_heads) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Input 'key_cache' shall have kv_num_heads, got ", + key_cache_dims[2]); + } + if (key_cache_dims[3] != head_size) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Input 'key_cache' dimension 3 should be same as head_size, got ", + key_cache_dims[3]); + } + + if (value_cache == nullptr) { + return Status::OK(); + } + + const auto& value_cache_dims = value_cache->Shape().GetDims(); + if (value_cache_dims.size() != 4) { return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "block_size must be a multiple of 256. Got block_size % 256 == ", - block_size % 256); + "Input 'value_cache' is expected to have 4 dimensions, got ", + value_cache_dims.size()); } if (value_cache_dims[0] != num_blocks) { return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, @@ -127,11 +201,6 @@ Status CheckKVCache(const T* key_cache, const T* value_cache, const int kv_num_h "Input 'key_cache' and 'value_cache' dimension 2 (kv num heads) should be the same, got ", key_cache_dims[2], " and ", value_cache_dims[2]); } - if (key_cache_dims[2] != kv_num_heads) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "Input 'key_cache' shall have kv_num_heads, got ", - key_cache_dims[2]); - } if (value_cache_dims[2] != kv_num_heads) { return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 'value_cache' shall have kv_num_heads, got ", @@ -143,11 +212,6 @@ Status CheckKVCache(const T* key_cache, const T* value_cache, const int kv_num_h "Input 'key_cache' and 'value_cache' dimension 3 (head size) should be the same, got ", key_cache_dims[3], " and ", value_cache_dims[3]); } - if (key_cache_dims[3] != head_size) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "Input 'key_cache' dimension 3 should be same as head_size, got ", - key_cache_dims[3]); - } if (value_cache_dims[3] != head_size) { return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 'past_value' dimension 3 should be same as head_size, got ", @@ -188,6 +252,138 @@ Status CheckBlockTable(const T* block_table, const int batch_size, int& max_num_ return Status::OK(); } +// slot_mapping (input 10) is the scheduler-owned write map: one flat slot index per query token +// into the cache viewed as [num_blocks * block_size, kv_num_heads, head_size], or -1 to skip the +// K/V store for that token. Element range is not validated on the host: that would require a +// device-to-host copy every step. Out-of-range values are undefined behavior, exactly as for +// block_table today. +template +Status CheckSlotMapping(const T* slot_mapping, const int token_count) { + const auto& dims = slot_mapping->Shape().GetDims(); + if (dims.size() != 1) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Input 'slot_mapping' is expected to have 1 dimension, got ", dims.size()); + } + if (dims[0] != token_count) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Input 'slot_mapping' dimension 0 should be token_count (", token_count, + "), got ", dims[0]); + } + return Status::OK(); +} + +template +Status CheckHeadSink(const T* head_sink, const int num_heads) { + const auto& dims = head_sink->Shape().GetDims(); + if (dims.size() != 1) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "head_sink must be a 1D tensor"); + } + if (dims[0] != num_heads) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "head_sink dimension 0 must be equal to the num heads, got ", dims[0]); + } + return Status::OK(); +} + +template +Status CheckQKNormWeights(const T* q_norm_weight, const T* k_norm_weight, const int head_size) { + if ((q_norm_weight != nullptr) != (k_norm_weight != nullptr)) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Input 'q_norm_weight' and 'k_norm_weight' must be provided together."); + } + if (q_norm_weight == nullptr) { + return Status::OK(); + } + const auto& q_dims = q_norm_weight->Shape().GetDims(); + const auto& k_dims = k_norm_weight->Shape().GetDims(); + if (q_dims.size() != 1 || q_dims[0] != head_size) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Input 'q_norm_weight' must be a 1D tensor of shape (head_size) = (", head_size, ")."); + } + if (k_dims.size() != 1 || k_dims[0] != head_size) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Input 'k_norm_weight' must be a 1D tensor of shape (head_size) = (", head_size, ")."); + } + return Status::OK(); +} + +// Validates one side (K or V) of the quantized paged KV cache contract. `is_quantized_cache` +// reflects the element type the kernel was instantiated for, so a mismatch between the cache dtype +// and the quant-type attribute is reported instead of silently producing garbage. +template +Status CheckKVCacheQuantization(const T* scale, const char* scale_name, const char* quant_type_name, + const KVQuantizationType quant_type, const bool is_quantized_cache, + const int kv_num_heads, const int head_size) { + if (quant_type == KVQuantizationType::NONE) { + if (scale != nullptr) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Input '", scale_name, "' must not be provided when '", quant_type_name, "' is 'NONE'."); + } + if (is_quantized_cache) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "The KV cache has a quantized element type, so '", quant_type_name, + "' must be 'PER_TENSOR' or 'PER_CHANNEL'."); + } + return Status::OK(); + } + + if (!is_quantized_cache) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "'", quant_type_name, + "' is set, but the KV cache element type is not quantized. " + "Use an int8 or float8e4m3fn cache, or set '", + quant_type_name, "' to 'NONE'."); + } + if (scale == nullptr) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Input '", scale_name, "' is required when '", quant_type_name, "' is not 'NONE'."); + } + + const auto& dims = scale->Shape().GetDims(); + const int64_t count = scale->Shape().Size(); + if (quant_type == KVQuantizationType::PER_TENSOR) { + if (count != 1) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Input '", scale_name, "' must have exactly 1 element for PER_TENSOR quantization, got ", + count); + } + return Status::OK(); + } + + // PER_CHANNEL. The canonical shape is (kv_num_heads, 1, head_size), matching GroupQueryAttention; + // any shape with the same element count and a trailing head_size is accepted so that callers may + // pass (kv_num_heads, head_size) directly. + if (count != static_cast(kv_num_heads) * head_size || dims.empty() || dims.back() != head_size) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Input '", scale_name, "' must have shape (kv_num_heads, 1, head_size) = (", + kv_num_heads, ", 1, ", head_size, ") for PER_CHANNEL quantization, got ", + scale->Shape().ToString()); + } + return Status::OK(); +} + +// Validates one side (K or V) of the `k_cache_dtype` / `v_cache_dtype` contract against +// `storage_dtype`, the element type the kernel was instantiated for. DEFAULT means "the cache +// tensor's element type is also the logical type" and always passes; naming that same type +// explicitly is allowed but must agree. The sub-byte members describe a logical type packed two per +// byte into a uint8 cache; the schema reserves them, but no backend decodes them yet, so they are +// rejected here instead of being silently mis-read. See docs/contrib_ops/cuda/paged_attention.md §8. +inline Status CheckKVCacheDataType(const KVCacheDataType cache_dtype, const KVCacheDataType storage_dtype, + const char* attr_name) { + if (cache_dtype == KVCacheDataType::DEFAULT || cache_dtype == storage_dtype) { + return Status::OK(); + } + if (IsSubByteKVCacheDataType(cache_dtype)) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "'", attr_name, "' == '", KVCacheDataTypeToString(cache_dtype), + "' requires a uint8 packed cache, which is not enabled in this build."); + } + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "'", attr_name, "' is '", KVCacheDataTypeToString(cache_dtype), + "', but the cache tensor's element type is '", KVCacheDataTypeToString(storage_dtype), + "'. Leave the attribute at '' to use the tensor's element type."); +} + template Status CheckInputs(const T* query, const T* key, @@ -199,12 +395,30 @@ Status CheckInputs(const T* query, const T* block_table, const T* cos_cache, const T* sin_cache, + const T* slot_mapping, + const T* head_sink, + const T* q_norm_weight, + const T* k_norm_weight, + const T* k_scale, + const T* v_scale, + const T* attention_metadata, void* parameters, int num_heads, int kv_num_heads, float scale, float softcap, + float qk_norm_epsilon, + KVQuantizationType k_quant_type, + KVQuantizationType v_quant_type, + KVCacheDataType k_cache_dtype, + KVCacheDataType v_cache_dtype, + KVCacheDataType cache_storage_dtype, + bool is_latent_kv, + int v_head_size_attr, + int rotary_offset, + bool has_explicit_scale, int max_threads_per_block) { + const bool is_quantized_cache = IsQuantizedKVCacheDataType(cache_storage_dtype); if (max_threads_per_block > 0 && num_heads > max_threads_per_block) { return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "num_heads should be no larger than ", max_threads_per_block); } @@ -214,13 +428,18 @@ Status CheckInputs(const T* query, num_heads % kv_num_heads); } - // Check query, key, and value + // Check query, key, and value. `kv_cache_layout` is inspected before the presence pattern, + // because LATENT's "key present, value absent" pattern would otherwise be indistinguishable from + // an ill-formed SEPARATE node. See docs/contrib_ops/cuda/paged_attention.md §4.6. int token_count = 0; int q_hidden_size = 0; int kv_hidden_size = 0; int head_size = 0; - const bool is_packed_qkv = key == nullptr; - if (!is_packed_qkv) { + const bool is_packed_qkv = !is_latent_kv && key == nullptr; + if (is_latent_kv) { + ORT_RETURN_IF_ERROR(Check_Q_K_Latent(query, key, value, num_heads, kv_num_heads, token_count, q_hidden_size, + kv_hidden_size, head_size)); + } else if (!is_packed_qkv) { ORT_RETURN_IF_ERROR(Check_Q_K_V(query, key, value, num_heads, kv_num_heads, token_count, q_hidden_size, kv_hidden_size, head_size)); } else { @@ -228,6 +447,67 @@ Status CheckInputs(const T* query, head_size)); } + // Effective V head size (§12.2). A V width that differs from head_size is only meaningful when V + // is a slice of the latent key, so it is confined to LATENT mode: no SEPARATE-mode backend + // supports asymmetric K/V widths, and value_cache's last dimension is head_size by construction. + int v_head_size = head_size; + if (v_head_size_attr != 0) { + if (v_head_size_attr < 1 || v_head_size_attr > head_size) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "'v_head_size' must be 0 (meaning head_size) or in [1, head_size] = [1, ", head_size, + "], got ", v_head_size_attr); + } + if (v_head_size_attr != head_size && !is_latent_kv) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "'v_head_size' (", v_head_size_attr, ") may only differ from head_size (", head_size, + ") when 'kv_cache_layout' is 'LATENT'."); + } + v_head_size = v_head_size_attr; + } + // The softmax-scale trap (§12.6): DeepSeek derives its scale from the pre-absorption head width, + // so the 1/sqrt(head_size) default would silently produce plausible-but-wrong logits. + if (v_head_size != head_size && !has_explicit_scale) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "An explicit 'scale' attribute is required when 'v_head_size' (", v_head_size, + ") differs from head_size (", head_size, + "): the default 1/sqrt(head_size) is not the intended scale for absorbed MLA."); + } + + if (is_latent_kv) { + if (value_cache != nullptr) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Input 'value_cache' must be absent when 'kv_cache_layout' is 'LATENT': the value cache " + "is the leading 'v_head_size' channels of 'key_cache'."); + } + if (kv_num_heads != 1) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "'kv_num_heads' must be 1 when 'kv_cache_layout' is 'LATENT', got ", kv_num_heads); + } + // §12.9: both combinations are well defined mathematically but no MLA model uses them, and an + // untested silent result is worse than a rejection. + if (head_sink != nullptr) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Input 'head_sink' is not supported when 'kv_cache_layout' is 'LATENT'."); + } + if (q_norm_weight != nullptr || k_norm_weight != nullptr) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Inputs 'q_norm_weight' / 'k_norm_weight' are not supported when 'kv_cache_layout' is " + "'LATENT': DeepSeek normalizes the latent projections in the graph, before absorption."); + } + // There is one physical cache, written once with k_scale, so a second scale for the same bytes + // could only disagree with it. V is dequantized with k_scale. + if (v_scale != nullptr || v_quant_type != KVQuantizationType::NONE || + v_cache_dtype != KVCacheDataType::DEFAULT) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Input 'v_scale' and attributes 'v_quant_type' / 'v_cache_dtype' must be unset when " + "'kv_cache_layout' is 'LATENT': the value elements are the key elements, so 'k_scale' " + "and 'k_cache_dtype' describe both."); + } + } else if (value_cache == nullptr) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Input 'value_cache' is required unless 'kv_cache_layout' is 'LATENT'."); + } + // Check KV-Cache int num_blocks = 0; int block_size = 0; @@ -240,6 +520,19 @@ Status CheckInputs(const T* query, // Check block table and slot mappings int max_num_blocks_per_seq = 0; ORT_RETURN_IF_ERROR(CheckBlockTable(block_table, batch_size, max_num_blocks_per_seq)); + if (slot_mapping != nullptr) { + ORT_RETURN_IF_ERROR(CheckSlotMapping(slot_mapping, token_count)); + } + + // Check attention sink and QK-Norm weights + if (head_sink != nullptr) { + ORT_RETURN_IF_ERROR(CheckHeadSink(head_sink, num_heads)); + } + ORT_RETURN_IF_ERROR(CheckQKNormWeights(q_norm_weight, k_norm_weight, head_size)); + if (q_norm_weight != nullptr && !(qk_norm_epsilon > 0.0f)) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "qk_norm_epsilon must be positive, got ", qk_norm_epsilon); + } // Check rotary cache int rotary_dim = 0; @@ -252,6 +545,42 @@ Status CheckInputs(const T* query, "Input 'cos_cache' and 'sin_cache' shall be both present or both absent."); } + // Offset (partial) rotary (§12.5). RoPE covers [rotary_offset, rotary_offset + rotary_dim) of + // each head; MLA rotates only the k_pe suffix. Default 0 is the shipped prefix behavior. + if (rotary_offset < 0 || rotary_offset % 8 != 0) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "'rotary_offset' must be non-negative and a multiple of 8, got ", rotary_offset); + } + if (rotary_offset + rotary_dim > head_size) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "'rotary_offset' + rotary_dim must not exceed head_size. Got ", rotary_offset, " + ", + rotary_dim, " > ", head_size); + } + + // Check quantized KV cache. LATENT has no value cache to describe, and the block above already + // required v_scale / v_quant_type to be unset there. + ORT_RETURN_IF_ERROR(CheckKVCacheQuantization(k_scale, "k_scale", "k_quant_type", k_quant_type, + is_quantized_cache, kv_num_heads, head_size)); + if (!is_latent_kv) { + ORT_RETURN_IF_ERROR(CheckKVCacheQuantization(v_scale, "v_scale", "v_quant_type", v_quant_type, + is_quantized_cache, kv_num_heads, v_head_size)); + } + ORT_RETURN_IF_ERROR(CheckKVCacheDataType(k_cache_dtype, cache_storage_dtype, "k_cache_dtype")); + ORT_RETURN_IF_ERROR(CheckKVCacheDataType(v_cache_dtype, cache_storage_dtype, "v_cache_dtype")); + + // Optional host-side [max_query_len_bound, max_kv_len_bound]. Only the shape is checked here. + // The entries are *trusted upper bounds* and cannot be cross-checked against the device tensors + // they bound without the readback this input exists to remove; see the trust boundary in + // docs/contrib_ops/cuda/paged_attention.md section 4.7. + if (attention_metadata != nullptr) { + const auto& metadata_dims = attention_metadata->Shape().GetDims(); + if (metadata_dims.size() != 1 || metadata_dims[0] != 2) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Input 'attention_metadata' must have shape (2), got ", + attention_metadata->Shape().ToString()); + } + } + if (parameters != nullptr) { PagedAttentionParameters* output_parameters = reinterpret_cast(parameters); output_parameters->batch_size = batch_size; @@ -261,6 +590,10 @@ Status CheckInputs(const T* query, output_parameters->num_heads = num_heads; output_parameters->kv_num_heads = kv_num_heads; output_parameters->head_size = head_size; + output_parameters->v_head_size = v_head_size; + output_parameters->v_hidden_size = num_heads * v_head_size; + output_parameters->is_latent_kv = is_latent_kv; + output_parameters->rotary_offset = rotary_offset; output_parameters->block_size = block_size; output_parameters->max_num_blocks_per_seq = max_num_blocks_per_seq; output_parameters->num_blocks = num_blocks; @@ -268,6 +601,11 @@ Status CheckInputs(const T* query, output_parameters->is_packed_qkv = is_packed_qkv; output_parameters->scale = scale; output_parameters->softcap = softcap; + output_parameters->use_smooth_softmax = head_sink != nullptr; + output_parameters->use_qk_norm = q_norm_weight != nullptr; + output_parameters->qk_norm_epsilon = qk_norm_epsilon; + output_parameters->k_quant_type = k_quant_type; + output_parameters->v_quant_type = v_quant_type; } return Status::OK(); diff --git a/onnxruntime/contrib_ops/cuda/bert/paged_attention_impl.cu b/onnxruntime/contrib_ops/cuda/bert/paged_attention_impl.cu index f0bb7fd81cf99..355dc5b6bbacf 100644 --- a/onnxruntime/contrib_ops/cuda/bert/paged_attention_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/paged_attention_impl.cu @@ -2,7 +2,9 @@ // Licensed under the MIT License. #include +#include // FLT_MAX #include +#include #include "core/providers/cuda/cu_inc/common.cuh" #include "core/providers/cuda/cuda_common.h" #include "core/providers/cuda/shared_inc/fpgeneric.h" @@ -11,6 +13,7 @@ #include "contrib_ops/cuda/bert/flash_attention/flash_api.h" #include "contrib_ops/cuda/bert/cutlass_fmha/memory_efficient_attention.h" #include "contrib_ops/cuda/bert/paged_attention_impl.h" +#include "contrib_ops/cuda/bert/xqa/xqa_paged_loader.h" #include "core/providers/cuda/shared_inc/cuda_call.h" #include "contrib_ops/cuda/bert/rotary_embedding_impl.h" #include @@ -21,6 +24,69 @@ namespace onnxruntime { namespace contrib { namespace cuda { +////////// Quantized paged KV cache helpers +// +// Symmetric, zero-point-free quantization with the same numerics GroupQueryAttention uses +// (group_query_attention_qdq.cuh): INT8 rounds to nearest and clamps to [-128, 127]; FP8 E4M3 +// clamps to +/-448 and lets the hardware convert. The paged layout +// [num_blocks, block_size, kv_num_heads, head_size] makes the scale index trivial: the innermost +// (kv_head, channel) pair is exactly the PER_CHANNEL scale index, and it is layout-independent, +// which is why the (kv_num_heads, 1, head_size) scale shape can be reused verbatim from GQA. + +constexpr int kPagedInt8Min = -128; +constexpr int kPagedInt8Max = 127; +constexpr float kPagedFp8E4M3Max = 448.0f; + +// True when the cache element type stores a quantized value that must be scaled on read/write. +template +struct IsQuantizedCache : std::false_type {}; +template <> +struct IsQuantizedCache : std::true_type {}; +#if defined(USE_FP8_KV_CACHE) && !defined(DISABLE_FLOAT8_TYPES) +template <> +struct IsQuantizedCache : std::true_type {}; +#endif + +// PER_CHANNEL scales are indexed by (kv_head * head_size + channel); PER_TENSOR uses scale[0]. +// `channel_index` is that flattened kv-hidden offset. +__device__ __forceinline__ float GetCacheScale(const float* __restrict__ scale, const int channel_index, + const bool per_channel) { + if (scale == nullptr) { + return 1.0f; + } + return per_channel ? scale[channel_index] : scale[0]; +} + +template +__device__ __forceinline__ TCACHE QuantizeToCache(const T value, const float scale) { + if constexpr (std::is_same::value) { + const float inv_scale = (scale == 0.0f) ? 0.0f : (1.0f / scale); + const int32_t q = static_cast(rintf(static_cast(value) * inv_scale)); + return static_cast(max(kPagedInt8Min, min(kPagedInt8Max, q))); +#if defined(USE_FP8_KV_CACHE) && !defined(DISABLE_FLOAT8_TYPES) + } else if constexpr (std::is_same::value) { + const float inv_scale = (scale == 0.0f) ? 0.0f : (1.0f / scale); + const float v = static_cast(value) * inv_scale; + return Float8E4M3FN(fmaxf(-kPagedFp8E4M3Max, fminf(kPagedFp8E4M3Max, v))); +#endif + } else { + return static_cast(value); + } +} + +template +__device__ __forceinline__ T DequantizeFromCache(const TCACHE value, const float scale) { + if constexpr (std::is_same::value) { + return static_cast(static_cast(value) * scale); +#if defined(USE_FP8_KV_CACHE) && !defined(DISABLE_FLOAT8_TYPES) + } else if constexpr (std::is_same::value) { + return static_cast(value.ToFloat() * scale); +#endif + } else { + return static_cast(value); + } +} + ////////// Auxiliary Kernels template @@ -86,36 +152,119 @@ Status LaunchUnpackCumulative(const T* input, T* output, const int token_count, return CUDA_CALL(cudaGetLastError()); } +// Fused per-head RMSNorm (QK-Norm) + rotary embedding over the unpadded TxNxH layout. +// +// Both transformations are pure prologue work on Q and K, so fusing them keeps a single read of the +// (possibly packed) input and a single write of the workspace, and it keeps the whole prologue +// independent of which attention backend runs afterwards. Order matters: QK-Norm is applied to the +// raw projection *before* RoPE, matching the reference implementations (and GroupQueryAttention's +// UnpackRoPEAppend), so the value that lands in the paged KV cache is normalized-then-rotated K. +// +// Set norm_weight == nullptr to skip normalization, and rotary_embedding_dim == 0 to skip rotary +// (the kernel then degenerates to an unpack/copy). At least one of the two must be enabled, +// otherwise the caller should not launch this kernel at all. +// +// `rotary_offset` is the first channel of the head that RoPE covers: the rotated span is +// [rotary_offset, rotary_offset + rotary_embedding_dim) and every channel outside it is copied +// through. rotary_offset == 0 is the shipped prefix-RoPE behavior; absorbed MLA rotates only the +// k_pe suffix and passes rotary_offset == kv_lora_rank. +// +// blockDim.x is the smallest power of two >= head_size so the reduction tree is exact; threads with +// h >= head_size participate in the reduction (contributing 0) but perform no global access. +// +// The grid is indexed by *global token*, not by (sequence position, sequence): the owning sequence +// is recovered from cumulative_seqlens_q with a binary search. That keeps the launch exactly +// token_count * num_heads blocks for any raggedness -- there is no per-sequence padding to skip -- +// and, more importantly, it removes the last dependence of the prologue on a host-computed +// max_query_len, which could only be obtained with a device-to-host synchronization +// (docs/contrib_ops/cuda/paged_attention.md section 4.7). template -__global__ void RotaryEmbeddingTNH(T* output, // TxNxH - const T* input, // TxNxH - const T* cos_cache, // Mx(H/2) - const T* sin_cache, // Mx(H/2) - const int32_t* past_seqlens, // B - const int32_t* cumulative_seqlens_q, // B+1 - const int head_size, - const int rotary_embedding_dim, - const bool interleaved, - const int3 in_strides, // TxNxH - const int3 out_strides) { // TxNxH +__global__ void QkNormRotaryTNH(T* output, // TxNxH + const T* input, // TxNxH + const T* cos_cache, // Mx(H/2) + const T* sin_cache, // Mx(H/2) + const int32_t* past_seqlens, // B + const int32_t* cumulative_seqlens_q, // B+1 + const T* norm_weight, // H, or nullptr + const float epsilon, + const int batch_size, + const int head_size, + const int rotary_embedding_dim, + const int rotary_offset, + const bool interleaved, + const int3 in_strides, // TxNxH + const int3 out_strides) { // TxNxH // Use .x in innermost loop to access global memory efficiently - const int b = blockIdx.y; - const int s = blockIdx.x; - const int n = blockIdx.z; + const int t = blockIdx.x; // index of the token in the unpadded input/output + const int n = blockIdx.y; const int h = threadIdx.x; - const int sequence_length = cumulative_seqlens_q[b + 1] - cumulative_seqlens_q[b]; - if (h >= head_size || s >= sequence_length) { + // cumulative_seqlens_q is non-decreasing, so this finds the first sequence whose exclusive end + // exceeds t -- which is the owning sequence, and correctly skips sequences with no new token. + int left = 0; + int right = batch_size; + while (left < right) { + const int mid = left + (right - left) / 2; + if (t < cumulative_seqlens_q[mid + 1]) { + right = mid; + } else { + left = mid + 1; + } + } + const int b = left; + // Defensive: a malformed cumulative_seqlens_q whose total disagrees with token_count would + // otherwise read past the end of past_seqlens. Uniform across the block. + if (b >= batch_size) { return; } + const int s = t - cumulative_seqlens_q[b]; // position of the token within its own sequence + + // Layout: blockDim.x floats for the reduction tree, then head_size elements of T holding the + // (optionally normalized) head so the rotary step can read its partner lane without a second + // global load. The float array comes first to keep both regions naturally aligned. + extern __shared__ char smem[]; + float* reduce_buffer = reinterpret_cast(smem); + T* head_values = reinterpret_cast(reduce_buffer + blockDim.x); - const int t = cumulative_seqlens_q[b] + s; // t is the index of the token in the unpadded input/output const T* input_data = input + t * in_strides.x + n * in_strides.y; T* output_data = output + t * out_strides.x + n * out_strides.y; - if (h >= rotary_embedding_dim) { - output_data[h] = input_data[h]; + const bool valid = h < head_size; + float value = valid ? static_cast(input_data[h]) : 0.0f; + + if (norm_weight != nullptr) { + // RMSNorm across head_size, accumulated in fp32 regardless of T. + reduce_buffer[h] = value * value; + __syncthreads(); + for (unsigned int stride = blockDim.x >> 1; stride > 0; stride >>= 1) { + if (h < static_cast(stride)) { + reduce_buffer[h] += reduce_buffer[h + stride]; + } + __syncthreads(); + } + const float inv_rms = rsqrtf(reduce_buffer[0] / static_cast(head_size) + epsilon); + if (valid) { + value = value * inv_rms * static_cast(norm_weight[h]); + } + } + + if (valid) { + head_values[h] = static_cast(value); + } + __syncthreads(); + + if (!valid) { + return; + } + + // A channel outside [rotary_offset, rotary_offset + rotary_embedding_dim) is copied through. + // rotary_embedding_dim == 0 makes every lane take this branch, which is the pure + // normalize-and-copy (or plain unpack) path. past_seqlens / cos_cache / sin_cache are then + // never dereferenced and may be null. + const int hr = h - rotary_offset; + if (hr < 0 || hr >= rotary_embedding_dim) { + output_data[h] = head_values[h]; return; } @@ -130,123 +279,289 @@ __global__ void RotaryEmbeddingTNH(T* output, // TxNx T sign = 0; int j = 0; if (interleaved) { - cache_idx = (h / 2) % half_rotary_embedding_dim; - sign = (h % 2 == 0) ? -1 : 1; - j = (h % 2 == 0) ? h + 1 : h - 1; // i - sign + cache_idx = (hr / 2) % half_rotary_embedding_dim; + sign = (hr % 2 == 0) ? -1 : 1; + j = (hr % 2 == 0) ? hr + 1 : hr - 1; // i - sign } else { - cache_idx = h % half_rotary_embedding_dim; - sign = (h < half_rotary_embedding_dim) ? -1 : 1; - j = (h + half_rotary_embedding_dim) % rotary_embedding_dim; + cache_idx = hr % half_rotary_embedding_dim; + sign = (hr < half_rotary_embedding_dim) ? -1 : 1; + j = (hr + half_rotary_embedding_dim) % rotary_embedding_dim; } - output_data[h] = input_data[h] * cos_data[cache_idx] + sign * input_data[j] * sin_data[cache_idx]; + output_data[h] = head_values[h] * cos_data[cache_idx] + sign * head_values[rotary_offset + j] * sin_data[cache_idx]; } +// Launches the fused QK-Norm / rotary prologue. Pass norm_weight == nullptr to disable QK-Norm and +// rotary_embedding_dim == 0 to disable rotary; the caller is responsible for not invoking this when +// both are disabled and the input is not packed. template -Status LaunchRotaryEmbeddingKernel(cudaStream_t stream, T* output, const T* input, const int32_t* past_seqlens, - const int32_t* cumulative_seqlens_q, const T* cos_cache, const T* sin_cache, - const int batch_size, const int max_seqlen_q, const int num_heads, - const int head_size, const int rotary_embedding_dim, const bool interleaved, - const int in_seq_stride, const int max_threads_per_block) { - ORT_ENFORCE(head_size <= max_threads_per_block, "Rotary embedding dim must be <= max_threads_per_block"); +Status LaunchQkNormRotaryKernel(cudaStream_t stream, T* output, const T* input, const int32_t* past_seqlens, + const int32_t* cumulative_seqlens_q, const T* cos_cache, const T* sin_cache, + const T* norm_weight, const float epsilon, const int batch_size, + const int token_count, const int num_heads, const int head_size, + const int rotary_embedding_dim, const int rotary_offset, const bool interleaved, + const int in_seq_stride, const int max_threads_per_block) { + if (batch_size == 0 || token_count == 0 || num_heads == 0) { + return Status::OK(); + } int3 in_strides = {in_seq_stride <= 0 ? num_heads * head_size : in_seq_stride, head_size, 1}; int3 out_strides = {num_heads * head_size, head_size, 1}; - int tpb = (head_size + 31) / 32 * 32; + // Round up to a power of two so the reduction tree halves exactly. + int tpb = 32; + while (tpb < head_size) { + tpb <<= 1; + } + ORT_ENFORCE(tpb <= max_threads_per_block, + "PagedAttention prologue requires head_size rounded up to a power of two (", tpb, + ") to be <= max_threads_per_block (", max_threads_per_block, ")"); - const dim3 grid(max_seqlen_q, batch_size, num_heads); + const size_t shared_bytes = static_cast(tpb) * sizeof(float) + static_cast(head_size) * sizeof(T); + const dim3 grid(token_count, num_heads); const dim3 block(tpb); - RotaryEmbeddingTNH<<>>( - output, input, cos_cache, sin_cache, past_seqlens, cumulative_seqlens_q, head_size, rotary_embedding_dim, - interleaved, in_strides, out_strides); + QkNormRotaryTNH<<>>( + output, input, cos_cache, sin_cache, past_seqlens, cumulative_seqlens_q, norm_weight, epsilon, batch_size, + head_size, rotary_embedding_dim, rotary_offset, interleaved, in_strides, out_strides); return CUDA_CALL(cudaGetLastError()); } +// Single-block inclusive scan over the per-sequence KV lengths. One block loops over the batch in +// kBlockSize-sized tiles carrying a running total, so there is no cap on batch_size (the previous +// implementation launched independent blocks whose cub::BlockScan did not compose, which silently +// produced wrong offsets past 256 concurrent sequences). template __global__ void GetCumulativeSeqlensKV(int32_t* cumulative_seqlens_kv, const int32_t* cumulative_seqlens_q, const int32_t* past_seqlens, const int batch_size) { - int id = blockDim.x * blockIdx.x + threadIdx.x; + typedef cub::BlockScan BlockScan; + __shared__ typename BlockScan::TempStorage temp_storage; + __shared__ int running_total; - if (id == 0) { + if (threadIdx.x == 0) { cumulative_seqlens_kv[0] = 0; + running_total = 0; } + __syncthreads(); - typedef cub::BlockScan BlockScan; - __shared__ typename BlockScan::TempStorage temp_storage; - - // Sum past_seqlens to new sequence length (which we get by subtracting cumulative_seqlens_q). - // Then do an inclusive sum across present sequence lengths to get the cumulative sequence length - if (id < batch_size) { - cumulative_seqlens_kv[id + 1] = past_seqlens[id] + cumulative_seqlens_q[id + 1] - cumulative_seqlens_q[id]; - int length = cumulative_seqlens_kv[id + 1]; - BlockScan(temp_storage).InclusiveSum(length, length); - cumulative_seqlens_kv[id + 1] = length; + for (int base = 0; base < batch_size; base += kBlockSize) { + const int id = base + static_cast(threadIdx.x); + // Sum past_seqlens to the new sequence length (which we get by subtracting cumulative_seqlens_q), + // then inclusive-scan across present sequence lengths. + const int length = (id < batch_size) + ? past_seqlens[id] + cumulative_seqlens_q[id + 1] - cumulative_seqlens_q[id] + : 0; + int prefix = 0; + int aggregate = 0; + BlockScan(temp_storage).InclusiveSum(length, prefix, aggregate); + if (id < batch_size) { + cumulative_seqlens_kv[id + 1] = running_total + prefix; + } + __syncthreads(); // all reads of running_total and of temp_storage are done + if (threadIdx.x == 0) { + running_total += aggregate; + } + __syncthreads(); // running_total visible, temp_storage safe to reuse } } Status LaunchGetCumulativeSeqlensKV(int32_t* cumulative_seqlens_kv, const int32_t* cumulative_seqlens_q, const int32_t* past_seqlens, const int batch_size, cudaStream_t stream) { - const int threads = 256; - const int blocks = (batch_size + threads - 1) / threads; - GetCumulativeSeqlensKV<256><<>>(cumulative_seqlens_kv, cumulative_seqlens_q, past_seqlens, - batch_size); + constexpr int kThreads = 256; + GetCumulativeSeqlensKV<<<1, kThreads, 0, stream>>>(cumulative_seqlens_kv, cumulative_seqlens_q, + past_seqlens, batch_size); return CUDA_CALL(cudaGetLastError()); } -template -__global__ void ReshapeAndCache(const T* __restrict__ key, const T* __restrict__ value, T* __restrict__ key_cache, - T* __restrict__ value_cache, const int* __restrict__ block_table, - const int* __restrict__ past_seqlens, const int* __restrict__ cumulative_seqlens_q, - const int batch_size, const int max_num_blocks_per_seq, const int token_count, - const int kv_hidden_size, const int block_size, const int key_stride, - const int value_stride) { - const int tid = threadIdx.x + blockIdx.x * blockDim.x; - if (tid >= token_count * kv_hidden_size) { - return; +// Resolves the flat cache slot that a query token's K/V is written to, in the cache viewed as +// [num_blocks * block_size, kv_num_heads, head_size]. A negative result suppresses the store. +// +// DerivedSlotResolver reproduces the legacy behavior: append the token at +// past_seqlens[b] + (token_id - cumulative_seqlens_q[b]) of its own sequence. The binary search is +// guarded against token_id >= cumulative_seqlens_q[batch_size], which previously walked off the end +// of past_seqlens / block_table. +struct DerivedSlotResolver { + const int* __restrict__ block_table; + const int* __restrict__ past_seqlens; + const int* __restrict__ cumulative_seqlens_q; + int batch_size; + int max_num_blocks_per_seq; + int block_size; + + __device__ __forceinline__ int operator()(int token_id) const { + if (token_id < 0 || token_id >= cumulative_seqlens_q[batch_size]) { + return -1; + } + // cumulative_seqlens_q is a non-decreasing prefix sum, so binary search finds the owning + // sequence in log2(batch_size) steps instead of the previous O(batch_size) scan. + int left = 0; + int right = batch_size - 1; + while (left < right) { + const int mid = left + (right - left) / 2; + if (token_id < cumulative_seqlens_q[mid + 1]) { + right = mid; + } else { + left = mid + 1; + } + } + const int batch_id = left; + const int position = past_seqlens[batch_id] + (token_id - cumulative_seqlens_q[batch_id]); + const int block_idx_in_seq = position / block_size; + if (block_idx_in_seq >= max_num_blocks_per_seq) { + return -1; + } + const int block_id = block_table[batch_id * max_num_blocks_per_seq + block_idx_in_seq]; + if (block_id < 0) { // unmapped block + return -1; + } + return block_id * block_size + position % block_size; + } +}; + +// ExplicitSlotResolver consumes the scheduler-provided slot_mapping (input 10) directly. This is +// what prefix caching, chunked prefill and speculative decoding need: the scheduler, not the +// kernel, owns block placement. It also removes the per-thread binary search entirely. +struct ExplicitSlotResolver { + const int* __restrict__ slot_mapping; + + __device__ __forceinline__ int operator()(int token_id) const { + return slot_mapping[token_id]; } - const int token_id = tid / kv_hidden_size; - const int hidden_offset = tid % kv_hidden_size; - int batch_id = 0; - for (int i = 0; i < batch_size; ++i) { - if (token_id < cumulative_seqlens_q[i + 1]) { - batch_id = i; - break; +}; + +template +__global__ void ReshapeAndCache(const T* __restrict__ key, const T* __restrict__ value, + TCACHE* __restrict__ key_cache, TCACHE* __restrict__ value_cache, + const float* __restrict__ k_scale, const float* __restrict__ v_scale, + const bool k_per_channel, const bool v_per_channel, + const SlotResolver resolver, const int64_t total_elems, + const int kv_hidden_size, const int key_stride, const int value_stride, + const int64_t num_slots) { + const int64_t stride = static_cast(gridDim.x) * blockDim.x; + for (int64_t tid = threadIdx.x + static_cast(blockIdx.x) * blockDim.x; + tid < total_elems; + tid += stride) { + const int token_id = static_cast(tid / kv_hidden_size); + const int hidden_offset = static_cast(tid % kv_hidden_size); + + const int slot = resolver(token_id); + // slot < 0 means "do not write this token" (unmapped block, or an explicit -1 in slot_mapping + // for a prefix-cache hit / rejected speculative token). The Q of such a token still attends. + if (slot < 0 || slot >= num_slots) { + continue; + } + + const int64_t key_id = static_cast(token_id) * key_stride + hidden_offset; + const int64_t dst_id = static_cast(slot) * kv_hidden_size + hidden_offset; + // hidden_offset is (kv_head * head_size + channel), which is exactly the PER_CHANNEL scale + // index. For an unquantized cache the scale pointers are null and this compiles to a copy. + key_cache[dst_id] = + QuantizeToCache(key[key_id], GetCacheScale(k_scale, hidden_offset, k_per_channel)); + // In LATENT (MLA) mode there is no separate value tensor or value cache: V is the leading + // v_head_size channels of the latent row that was just written above. + if (value_cache != nullptr) { + const int64_t value_id = static_cast(token_id) * value_stride + hidden_offset; + value_cache[dst_id] = + QuantizeToCache(value[value_id], GetCacheScale(v_scale, hidden_offset, v_per_channel)); } } - const int token_offset = token_id - cumulative_seqlens_q[batch_id]; - const int past_length = past_seqlens[batch_id]; - const int block_id = block_table[batch_id * max_num_blocks_per_seq + (past_length + token_offset) / block_size]; - const int block_offset = (past_length + token_offset) % block_size; +} + +template +Status LaunchReshapeAndCacheImpl(const T* key, const T* value, TCACHE* key_cache, TCACHE* value_cache, + const float* k_scale, const float* v_scale, const bool k_per_channel, + const bool v_per_channel, const SlotResolver& resolver, const int token_count, + const int kv_hidden_size, const int key_stride, const int value_stride, + const int64_t num_slots, cudaStream_t stream, const int max_threads_per_block) { + const int64_t total_elems = static_cast(token_count) * kv_hidden_size; + if (total_elems == 0) { + return Status::OK(); + } + const int threads = static_cast(std::min(max_threads_per_block, total_elems)); + const int blocks = static_cast(std::min((total_elems + threads - 1) / threads, 65535)); + ReshapeAndCache<<>>( + key, value, key_cache, value_cache, k_scale, v_scale, k_per_channel, v_per_channel, resolver, total_elems, + kv_hidden_size, key_stride, value_stride, num_slots); + return CUDA_CALL(cudaGetLastError()); +} - const int key_id = token_id * key_stride + hidden_offset; - const int value_id = token_id * value_stride + hidden_offset; - const int dst_id = block_id * block_size * kv_hidden_size + block_offset * kv_hidden_size + hidden_offset; - key_cache[dst_id] = key[key_id]; - value_cache[dst_id] = value[value_id]; +template +Status LaunchReshapeAndCache(const T* key, const T* value, TCACHE* key_cache, TCACHE* value_cache, + const float* k_scale, const float* v_scale, const bool k_per_channel, + const bool v_per_channel, const int* block_table, + const int* past_seqlens, const int* cumulative_seqlens_q, const int* slot_mapping, + const int batch_size, const int max_num_blocks_per_seq, const int token_count, + const int kv_hidden_size, const int block_size, const int num_blocks, + const int key_stride, const int value_stride, cudaStream_t stream, + const int max_threads_per_block) { + const int64_t num_slots = static_cast(num_blocks) * block_size; + if (slot_mapping != nullptr) { + ExplicitSlotResolver resolver{slot_mapping}; + return LaunchReshapeAndCacheImpl( + key, value, key_cache, value_cache, k_scale, v_scale, k_per_channel, v_per_channel, resolver, + token_count, kv_hidden_size, key_stride, value_stride, num_slots, stream, max_threads_per_block); + } + DerivedSlotResolver resolver{block_table, past_seqlens, cumulative_seqlens_q, batch_size, + max_num_blocks_per_seq, block_size}; + return LaunchReshapeAndCacheImpl( + key, value, key_cache, value_cache, k_scale, v_scale, k_per_channel, v_per_channel, resolver, + token_count, kv_hidden_size, key_stride, value_stride, num_slots, stream, max_threads_per_block); } +// Exact attention-sink epilogue. FlashAttention returns +// lse[t,h] = log(sum_j exp(x_j)) and o[t,h] = (sum_j exp(x_j) v_j) / sum_j exp(x_j), +// and a sink only adds the extra logit s_h to the denominator, so the corrected output is the +// elementwise rescale o *= exp(lse) / (exp(lse) + exp(s_h)) = 1 / (1 + exp(s_h - lse)). +// This is numerically stable for both signs of (s_h - lse), needs no change to the Flash kernel, +// and composes with sliding window, softcap and GQA grouping because lse already reflects the mask. template -Status LaunchReshapeAndCache(const T* key, const T* value, T* key_cache, T* value_cache, const int* block_table, - const int* past_seqlens, const int* cumulative_seqlens_q, const int batch_size, - const int max_num_blocks_per_seq, const int token_count, const int kv_hidden_size, - const int block_size, const int key_stride, const int value_stride, cudaStream_t stream, - const int max_threads_per_block) { - const int total_size = token_count * kv_hidden_size; - const int threads(std::min(total_size, max_threads_per_block)); - const int blocks((total_size + threads - 1) / threads); - ReshapeAndCache<<>>(key, value, key_cache, value_cache, block_table, past_seqlens, - cumulative_seqlens_q, batch_size, max_num_blocks_per_seq, - token_count, kv_hidden_size, block_size, key_stride, value_stride); +__global__ void ApplyHeadSink(T* __restrict__ output, const float* __restrict__ softmax_lse, + const T* __restrict__ head_sink, const int token_count, const int num_heads, + const int head_size, const int64_t total_elems) { + const int64_t stride = static_cast(gridDim.x) * blockDim.x; + const int64_t num_heads_times_head = static_cast(num_heads) * head_size; + for (int64_t tid = threadIdx.x + static_cast(blockIdx.x) * blockDim.x; + tid < total_elems; + tid += stride) { + const int head_id = static_cast((tid / head_size) % num_heads); + const int token_id = static_cast(tid / num_heads_times_head); + // Varlen LSE layout is [num_heads, token_count] (Flash's unpadded_lse with num_splits <= 1). + const float lse = softmax_lse[static_cast(head_id) * token_count + token_id]; + const float sink = (head_sink == nullptr) ? 0.0f : static_cast(head_sink[head_id]); + // lse == +inf marks a fully masked row (output already zero) -> factor 1; lse == -inf gives + // factor 0. expf saturates correctly in both cases, so no special casing is needed. + const float factor = 1.0f / (1.0f + expf(sink - lse)); + output[tid] = static_cast(static_cast(output[tid]) * factor); + } +} + +template +Status LaunchApplyHeadSink(T* output, const float* softmax_lse, const T* head_sink, const int token_count, + const int num_heads, const int head_size, cudaStream_t stream, + const int max_threads_per_block) { + const int64_t total_elems = static_cast(token_count) * num_heads * head_size; + if (total_elems == 0) { + return Status::OK(); + } + const int threads = static_cast(std::min(max_threads_per_block, total_elems)); + const int blocks = static_cast(std::min((total_elems + threads - 1) / threads, 65535)); + ApplyHeadSink<<>>(output, softmax_lse, head_sink, token_count, num_heads, + head_size, total_elems); return CUDA_CALL(cudaGetLastError()); } -// Gather paged KV into packed-varlen [total_kv_tokens, num_heads, head_size], expanding GQA heads. -// total_elems = total_kv_tokens * num_heads * head_size can exceed INT32_MAX for realistic +// Gather paged KV into packed-varlen [total_kv_tokens, out_num_heads, head_size], dequantizing on +// the fly when the cache is quantized. out_num_heads == num_heads expands GQA groups (what the +// CUTLASS memory-efficient kernel needs); out_num_heads == kv_num_heads keeps the grouped layout +// (what FlashAttention's non-paged varlen entry point needs). +// total_elems = total_kv_tokens * out_num_heads * head_size can exceed INT32_MAX for realistic // large-context GQA configs (e.g., 2M tokens * 64 * 128 = 16.4B), so the linear index is int64_t // and the kernel uses a grid-stride loop instead of a single (tid >= total_elems) early-exit. -template -__global__ void GatherAndExpandPagedKVCache(const T* __restrict__ key_cache, - const T* __restrict__ value_cache, +template +__global__ void GatherAndExpandPagedKVCache(const TCACHE* __restrict__ key_cache, + const TCACHE* __restrict__ value_cache, T* __restrict__ gathered_key, T* __restrict__ gathered_value, + const float* __restrict__ k_scale, + const float* __restrict__ v_scale, + const bool k_per_channel, + const bool v_per_channel, const int* __restrict__ block_table, const int* __restrict__ cumulative_seqlens_kv, const int batch_size, @@ -286,28 +601,47 @@ __global__ void GatherAndExpandPagedKVCache(const T* __restrict__ key_cache, } const int batch_id = left; + // Defensive: a malformed cumulative_seqlens_kv (or a total_kv_tokens that disagrees with it) + // would leave batch_id == batch_size and walk off the end of block_table. + if (batch_id >= batch_size) { + continue; + } + const int pos = token_id - cumulative_seqlens_kv[batch_id]; const int block_idx_in_seq = pos / block_size; const int block_offset = pos % block_size; + if (block_idx_in_seq >= max_num_blocks_per_seq) { + continue; + } const int block_id = block_table[batch_id * max_num_blocks_per_seq + block_idx_in_seq]; + if (block_id < 0) { + gathered_key[tid] = static_cast(0.f); + gathered_value[tid] = static_cast(0.f); + continue; + } // GQA expansion: each output head maps to kv_head_id = head_id / (num_heads / kv_num_heads). // For MHA (num_heads == kv_num_heads) this is the identity. const int kv_head_id = head_id / q_kv_head_ratio; + const int channel_index = kv_head_id * head_size + h; const int64_t paged_idx = static_cast(block_id) * page_stride + static_cast(block_offset) * kv_num_heads * head_size + kv_head_id * head_size + h; - gathered_key[tid] = key_cache[paged_idx]; - gathered_value[tid] = value_cache[paged_idx]; + gathered_key[tid] = + DequantizeFromCache(key_cache[paged_idx], GetCacheScale(k_scale, channel_index, k_per_channel)); + gathered_value[tid] = + DequantizeFromCache(value_cache[paged_idx], GetCacheScale(v_scale, channel_index, v_per_channel)); } } -template -Status LaunchGatherAndExpandPagedKVCache(const T* key_cache, const T* value_cache, +template +Status LaunchGatherAndExpandPagedKVCache(const TCACHE* key_cache, const TCACHE* value_cache, T* gathered_key, T* gathered_value, + const float* k_scale, const float* v_scale, + const bool k_per_channel, const bool v_per_channel, const int* block_table, const int* cumulative_seqlens_kv, const int batch_size, const int num_heads, const int kv_num_heads, const int head_size, @@ -318,33 +652,679 @@ Status LaunchGatherAndExpandPagedKVCache(const T* key_cache, const T* value_cach if (total_elems == 0) { return Status::OK(); } - // With the op's batch_size <= 256 precondition (paged_attention.cc) and MEA's - // head_size <= 1024 cap, blocks_needed = ceil(total_elems / threads) stays comfortably - // within int range for any realistic input, so no explicit clamp is needed. The kernel - // uses a grid-stride loop so launching fewer blocks than total_elems / threads would - // also be correct — we don't need an artificial "keep SMs busy" cap. + // The kernel uses a grid-stride loop, so the block count is clamped rather than allowed to + // overflow int for very large contexts. const int threads = static_cast(std::min(max_threads_per_block, total_elems)); - const int blocks = static_cast((total_elems + threads - 1) / threads); - GatherAndExpandPagedKVCache<<>>( - key_cache, value_cache, gathered_key, gathered_value, + const int blocks = static_cast(std::min((total_elems + threads - 1) / threads, 65535)); + GatherAndExpandPagedKVCache<<>>( + key_cache, value_cache, gathered_key, gathered_value, k_scale, v_scale, k_per_channel, v_per_channel, block_table, cumulative_seqlens_kv, batch_size, num_heads, kv_num_heads, head_size, block_size, max_num_blocks_per_seq, total_elems); return CUDA_CALL(cudaGetLastError()); } -////////// Launch Kernels +////////// Paged decode attention (flash-decoding style, reads the paged cache in place) +// +// Selected when the static shapes say the step is decode-shaped (token_count == batch_size). Unlike +// the gather-based path it never materializes a dense FP16 copy of the live context: K and V are +// read straight out of their pages and dequantized in registers, so a decode step touches the KV +// cache exactly once at its stored precision. +// +// The shape test is a heuristic, not a proof: token_count == batch_size does not guarantee that +// every sequence contributes exactly one query token (one may contribute two while another +// contributes none). The kernel is therefore indexed by *global query token* and derives both the +// owning sequence and the token's position inside it from cumulative_seqlens_q on device, so it +// stays correct for arbitrary ragged input -- including full prefill. Correctness never depends on +// the host heuristic being right, only speed. See +// docs/contrib_ops/cuda/paged_attention.md section 4.7. +// +// Both scales are folded instead of applied per element, which is exact and free: +// * K: q'_c = q_c * k_scale_c, so dot(q', k_raw) == dot(q, dequant(k)). PER_TENSOR degenerates to +// a uniform pre-scale of Q, which is why no separate code path is needed. +// * V: out_c = (sum_t p_t * v_raw[t][c]) * v_scale_c -- the scale does not depend on t, so it +// factors out of the accumulation entirely and is applied once in the reduce kernel. It also +// never enters the softmax denominator. +// +// The KV range of a sequence is split across `num_splits` CTAs; each emits a partial +// (max, denominator, unnormalized accumulator) triple that PagedDecodeReduce combines. -#if USE_FLASH_ATTENTION +constexpr int kPagedDecodeThreads = 128; +constexpr int kPagedDecodeTile = 128; // KV tokens scored per iteration +constexpr int kPagedDecodeMaxSplits = 32; + +// Cache element -> float. Identical to DequantizeFromCache with a scale of 1, kept separate because +// the decode kernel folds the scales into Q and into the output instead of applying them per read. +template +__device__ __forceinline__ float CacheToFloat(const TCACHE value) { +#if defined(USE_FP8_KV_CACHE) && !defined(DISABLE_FLOAT8_TYPES) + if constexpr (std::is_same::value) { + return value.ToFloat(); + } else // NOLINT(readability/braces) +#endif + { + return static_cast(value); + } +} + +// Number of channel groups the PV accumulation is split into. When head_size < blockDim, several +// groups of `head_size` threads each walk a disjoint subset of the tile's tokens and their partial +// accumulators are summed at the end; this keeps every thread busy and keeps the V reads +// warp-contiguous in the channel dimension. +__host__ __device__ __forceinline__ int PagedDecodeChannelGroups(const int head_size) { + return head_size >= kPagedDecodeThreads ? 1 : (kPagedDecodeThreads / head_size); +} + +template +__global__ void PagedDecodeSplitKV(const T* __restrict__ query, + const TCACHE* __restrict__ key_cache, + const TCACHE* __restrict__ value_cache, + const float* __restrict__ k_scale, + const int* __restrict__ cumulative_seqlens_q, + const int* __restrict__ cumulative_seqlens_kv, + const int* __restrict__ block_table, + float* __restrict__ partial_out, + float* __restrict__ partial_max, + float* __restrict__ partial_sum, + const int batch_size, + const int num_heads, + const int kv_num_heads, + const int head_size, + const int block_size, + const int max_num_blocks_per_seq, + const int token_count, + const int num_splits, + const float scale, + const float softcap, + const int local_window_size, + const bool k_per_channel) { + extern __shared__ float paged_decode_smem[]; + const int channel_groups = PagedDecodeChannelGroups(head_size); + const int acc_elems = channel_groups * head_size; + // Carve the dynamic block into: Q (head_size), the tile's logits (kPagedDecodeTile), the output + // accumulator (acc_elems), a block-reduction scratchpad (kPagedDecodeThreads) and the tile's + // resolved page ids (kPagedDecodeTile ints). + float* q_sh = paged_decode_smem; + float* logits_sh = q_sh + head_size; + float* acc_sh = logits_sh + kPagedDecodeTile; + float* red_sh = acc_sh + acc_elems; + int* block_id_sh = reinterpret_cast(red_sh + kPagedDecodeThreads); + + const int head_id = blockIdx.x; + const int token_id = blockIdx.y; + const int split_id = blockIdx.z; + const int tid = threadIdx.x; + + // One CTA owns one (query token, query head) pair. cumulative_seqlens_q is non-decreasing, so + // this binary search returns the first sequence whose exclusive end exceeds token_id -- the + // owning sequence -- and skips sequences that contribute no token at all. + int left = 0; + int right = batch_size; + while (left < right) { + const int mid = left + (right - left) / 2; + if (token_id < cumulative_seqlens_q[mid + 1]) { + right = mid; + } else { + left = mid + 1; + } + } + const int batch_id = left; + // Defensive: a cumulative_seqlens_q whose total disagrees with token_count would otherwise walk + // off the end of cumulative_seqlens_kv and block_table. Uniform across the block. + if (batch_id >= batch_size) { + return; + } + + const int64_t partial_head_index = + (static_cast(split_id) * token_count + token_id) * num_heads + head_id; + + // Causality is resolved per query token instead of assuming one new token per sequence: + // kv_len - q_len is past_seqlens[batch_id], so the token at offset q_index inside its sequence + // attends to cached positions [0, past + q_index]. For the decode case (q_len == 1) this reduces + // to the whole live context, as before. + const int q_index = token_id - cumulative_seqlens_q[batch_id]; + const int q_len = cumulative_seqlens_q[batch_id + 1] - cumulative_seqlens_q[batch_id]; + const int seq_kv_len = cumulative_seqlens_kv[batch_id + 1] - cumulative_seqlens_kv[batch_id]; + const int kv_len = seq_kv_len - q_len + q_index + 1; + + // Sliding window matches FlashAttention's window_size_left = local_window_size - 1 convention at + // query position kv_len - 1, i.e. positions in [kv_len - local_window_size, kv_len). + const int tokens_per_split = (kv_len + num_splits - 1) / num_splits; + int kv_begin = split_id * tokens_per_split; + const int kv_end = min(kv_len, kv_begin + tokens_per_split); + if (local_window_size > 0) { + kv_begin = max(kv_begin, kv_len - local_window_size); + } + + if (kv_begin >= kv_end) { + if (tid == 0) { + partial_max[partial_head_index] = -FLT_MAX; + partial_sum[partial_head_index] = 0.0f; + } + return; + } + + const int kv_head_id = head_id / (num_heads / kv_num_heads); + const int64_t head_offset_in_page = static_cast(kv_head_id) * head_size; + const int64_t token_stride_in_page = static_cast(kv_num_heads) * head_size; + + const T* q_ptr = query + (static_cast(token_id) * num_heads + head_id) * head_size; + for (int c = tid; c < head_size; c += kPagedDecodeThreads) { + q_sh[c] = static_cast(q_ptr[c]) * GetCacheScale(k_scale, kv_head_id * head_size + c, k_per_channel); + } + for (int c = tid; c < acc_elems; c += kPagedDecodeThreads) { + acc_sh[c] = 0.0f; + } + __syncthreads(); + + constexpr int kNumWarps = kPagedDecodeThreads / 32; + const int warp_id = tid / 32; + const int lane_id = tid % 32; + // FlashAttention computes softcap as scale_softmax * tanh(qk_raw * softmax_scale / softcap) with + // scale_softmax == softcap (flash_api.cc), so the effective logit is + // softcap * tanh(qk * scale / softcap). Match that exactly. + const float softcap_scale = softcap > 0.0f ? (scale / softcap) : 0.0f; + + float m_state = -FLT_MAX; + float l_state = 0.0f; + + for (int tile_begin = kv_begin; tile_begin < kv_end; tile_begin += kPagedDecodeTile) { + const int tile_len = min(kPagedDecodeTile, kv_end - tile_begin); + + // ---- QK: one warp per KV token, 32 lanes cooperating on the head-size dot product ---- + for (int t = warp_id; t < tile_len; t += kNumWarps) { + const int pos = tile_begin + t; + const int block_index = pos / block_size; + const int block_id = block_index < max_num_blocks_per_seq + ? block_table[batch_id * max_num_blocks_per_seq + block_index] + : -1; + float dot = 0.0f; + if (block_id >= 0) { + const TCACHE* k_ptr = key_cache + + (static_cast(block_id) * block_size + (pos % block_size)) * token_stride_in_page + + head_offset_in_page; + for (int c = lane_id; c < head_size; c += 32) { + dot += q_sh[c] * CacheToFloat(k_ptr[c]); + } + } +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + dot += __shfl_xor_sync(0xFFFFFFFFU, dot, offset); + } + if (lane_id == 0) { + block_id_sh[t] = block_id; + logits_sh[t] = block_id < 0 ? -FLT_MAX + : (softcap > 0.0f ? softcap * tanhf(dot * softcap_scale) : dot * scale); + } + } + __syncthreads(); + + // ---- tile max ---- + float local_max = -FLT_MAX; + for (int t = tid; t < tile_len; t += kPagedDecodeThreads) { + local_max = fmaxf(local_max, logits_sh[t]); + } + red_sh[tid] = local_max; + __syncthreads(); + for (int stride = kPagedDecodeThreads / 2; stride > 0; stride >>= 1) { + if (tid < stride) { + red_sh[tid] = fmaxf(red_sh[tid], red_sh[tid + stride]); + } + __syncthreads(); + } + const float m_tile = red_sh[0]; + __syncthreads(); + + // A tile whose pages are all unmapped contributes nothing and would otherwise make the + // rescale factor exp(-FLT_MAX - -FLT_MAX) == 1 for masked entries. + if (m_tile == -FLT_MAX) { + continue; + } + + const float m_new = fmaxf(m_state, m_tile); + const float alpha = __expf(m_state - m_new); // 0 on the first tile (m_state == -FLT_MAX) + + float local_sum = 0.0f; + for (int t = tid; t < tile_len; t += kPagedDecodeThreads) { + const float p = __expf(logits_sh[t] - m_new); + logits_sh[t] = p; + local_sum += p; + } + red_sh[tid] = local_sum; + __syncthreads(); + for (int stride = kPagedDecodeThreads / 2; stride > 0; stride >>= 1) { + if (tid < stride) { + red_sh[tid] += red_sh[tid + stride]; + } + __syncthreads(); + } + const float sum_tile = red_sh[0]; + __syncthreads(); + + l_state = l_state * alpha + sum_tile; + m_state = m_new; + + // ---- PV: consecutive threads own consecutive channels so the V reads stay coalesced ---- + if (channel_groups == 1) { + for (int c = tid; c < head_size; c += kPagedDecodeThreads) { + float acc = acc_sh[c] * alpha; + for (int t = 0; t < tile_len; ++t) { + const int block_id = block_id_sh[t]; + if (block_id < 0) { + continue; + } + const int pos = tile_begin + t; + const TCACHE* v_ptr = value_cache + + (static_cast(block_id) * block_size + (pos % block_size)) * token_stride_in_page + + head_offset_in_page; + acc += logits_sh[t] * CacheToFloat(v_ptr[c]); + } + acc_sh[c] = acc; + } + } else if (tid < acc_elems) { + const int group = tid / head_size; + const int c = tid - group * head_size; + float acc = acc_sh[tid] * alpha; + for (int t = group; t < tile_len; t += channel_groups) { + const int block_id = block_id_sh[t]; + if (block_id < 0) { + continue; + } + const int pos = tile_begin + t; + const TCACHE* v_ptr = value_cache + + (static_cast(block_id) * block_size + (pos % block_size)) * token_stride_in_page + + head_offset_in_page; + acc += logits_sh[t] * CacheToFloat(v_ptr[c]); + } + acc_sh[tid] = acc; + } + __syncthreads(); + } + + const int64_t out_base = + ((static_cast(split_id) * token_count + token_id) * num_heads + head_id) * head_size; + for (int c = tid; c < head_size; c += kPagedDecodeThreads) { + float acc = 0.0f; + for (int group = 0; group < channel_groups; ++group) { + acc += acc_sh[group * head_size + c]; + } + partial_out[out_base + c] = acc; + } + if (tid == 0) { + partial_max[partial_head_index] = m_state; + partial_sum[partial_head_index] = l_state; + } +} + +// Combine the per-split partials, apply the (folded) V scale and close the softmax. The attention +// sink enters here as one extra logit in the denominator, which is exactly what the FlashAttention +// path's ApplyHeadSink epilogue computes from the log-sum-exp. template -Status FlashAttention( - const cudaDeviceProp& device_prop, - cudaStream_t stream, - contrib::PagedAttentionParameters& parameters, - PagedAttentionData& data, - float scale) { - // Get parameters - const int max_threads_per_block = device_prop.maxThreadsPerBlock; +__global__ void PagedDecodeReduce(T* __restrict__ output, + const float* __restrict__ partial_out, + const float* __restrict__ partial_max, + const float* __restrict__ partial_sum, + const float* __restrict__ v_scale, + const T* __restrict__ head_sink, + const int num_heads, + const int kv_num_heads, + const int head_size, + const int token_count, + const int num_splits, + const bool v_per_channel, + const bool use_smooth_softmax) { + __shared__ float weight_sh[kPagedDecodeMaxSplits]; + __shared__ float max_sh[kPagedDecodeMaxSplits]; + __shared__ float sum_sh[kPagedDecodeMaxSplits]; + + const int head_id = blockIdx.x; + const int token_id = blockIdx.y; + const int tid = threadIdx.x; + const int64_t row_base = (static_cast(token_id) * num_heads + head_id) * head_size; + + if (tid < num_splits) { + const int64_t index = (static_cast(tid) * token_count + token_id) * num_heads + head_id; + max_sh[tid] = partial_max[index]; + sum_sh[tid] = partial_sum[index]; + } + __syncthreads(); + + float m_final = -FLT_MAX; + for (int s = 0; s < num_splits; ++s) { + if (sum_sh[s] > 0.0f) { + m_final = fmaxf(m_final, max_sh[s]); + } + } + + if (m_final == -FLT_MAX) { + for (int c = tid; c < head_size; c += blockDim.x) { + output[row_base + c] = static_cast(0.0f); + } + return; + } + + float l_final = 0.0f; + for (int s = 0; s < num_splits; ++s) { + const float w = sum_sh[s] > 0.0f ? __expf(max_sh[s] - m_final) : 0.0f; + if (tid == 0) { + weight_sh[s] = w; + } + l_final += sum_sh[s] * w; + } + if (use_smooth_softmax) { + const float sink = head_sink == nullptr ? 0.0f : static_cast(head_sink[head_id]); + l_final += __expf(sink - m_final); + } + __syncthreads(); + + const float inv_l = l_final > 0.0f ? (1.0f / l_final) : 0.0f; + const int kv_head_id = head_id / (num_heads / kv_num_heads); + for (int c = tid; c < head_size; c += blockDim.x) { + float acc = 0.0f; + for (int s = 0; s < num_splits; ++s) { + if (weight_sh[s] > 0.0f) { + acc += partial_out[((static_cast(s) * token_count + token_id) * num_heads + head_id) * head_size + c] * + weight_sh[s]; + } + } + output[row_base + c] = static_cast(acc * inv_l * + GetCacheScale(v_scale, kv_head_id * head_size + c, v_per_channel)); + } +} + +// Splits are only worth it when there are not enough (query token, head) pairs to fill the device. +// max_kv_len only sizes the launch, so a replay-invariant upper bound is a valid argument: an +// over-estimate costs empty splits that exit after a single device read. +int ComputePagedDecodeSplits(const int token_count, const int num_heads, const int max_kv_len, + const int multi_processor_count) { + const int base_ctas = token_count * num_heads; + if (base_ctas <= 0 || base_ctas >= 2 * multi_processor_count) { + return 1; + } + const int by_occupancy = (2 * multi_processor_count + base_ctas - 1) / base_ctas; + const int by_length = (max_kv_len + kPagedDecodeTile - 1) / kPagedDecodeTile; + return std::max(1, std::min(std::min(by_occupancy, by_length), kPagedDecodeMaxSplits)); +} + +size_t GetPagedDecodeSharedMemoryBytes(const int head_size) { + const size_t float_elems = static_cast(head_size) + kPagedDecodeTile + + static_cast(PagedDecodeChannelGroups(head_size)) * head_size + + kPagedDecodeThreads; + return float_elems * sizeof(float) + static_cast(kPagedDecodeTile) * sizeof(int); +} + +template +Status LaunchPagedDecodeAttention(const T* query, const TCACHE* key_cache, const TCACHE* value_cache, + const float* k_scale, const float* v_scale, + const bool k_per_channel, const bool v_per_channel, + const int* cumulative_seqlens_q, const int* cumulative_seqlens_kv, + const int* block_table, const T* head_sink, T* output, + float* partial_out, float* partial_max, float* partial_sum, + const int batch_size, const int num_heads, const int kv_num_heads, + const int head_size, const int block_size, const int max_num_blocks_per_seq, + const int token_count, const int num_splits, const float scale, + const float softcap, const int local_window_size, + const bool use_smooth_softmax, cudaStream_t stream) { + const size_t smem_bytes = GetPagedDecodeSharedMemoryBytes(head_size); + const dim3 grid(num_heads, token_count, num_splits); + PagedDecodeSplitKV<<>>( + query, key_cache, value_cache, k_scale, cumulative_seqlens_q, cumulative_seqlens_kv, block_table, + partial_out, partial_max, partial_sum, batch_size, num_heads, kv_num_heads, head_size, block_size, + max_num_blocks_per_seq, token_count, num_splits, scale, softcap, local_window_size, k_per_channel); + CUDA_RETURN_IF_ERROR(cudaGetLastError()); + + const dim3 reduce_grid(num_heads, token_count); + PagedDecodeReduce<<>>( + output, partial_out, partial_max, partial_sum, v_scale, head_sink, num_heads, kv_num_heads, + head_size, token_count, num_splits, v_per_channel, use_smooth_softmax); + return CUDA_CALL(cudaGetLastError()); +} + +////////// Unfused paged latent attention (absorbed MLA reference) +// +// Absorbed MLA is MQA with a wide key head and a narrower value head that *aliases* the leading +// v_head_size channels of the same cached row (docs/contrib_ops/cuda/paged_attention.md §12.2): +// +// key = [compressed_kv (kv_lora_rank) ; k_pe (qk_rope_head_dim)] -> head_size (576 in V3) +// value = compressed_kv -> v_head_size (512 in V3) +// +// Neither FlashAttention nor the CUTLASS fMHA wrapper can express that: both cap head_size at 256 +// and require v_head_size == head_size. This kernel is the correctness oracle called for in §12.7, +// and it is the only backend for LATENT until a fused MLA kernel lands. It handles arbitrary query +// lengths (prefill, chunked prefill and decode all take the same path), so causality is resolved +// per query token instead of assuming one new token per sequence. +// +// One CTA owns one (query token, query head) pair and streams the whole KV range in tiles, keeping +// an online-softmax (m, l) state and a v_head_size-wide fp32 accumulator in shared memory. The +// quantization scales are folded exactly as in the decode kernel: k_scale into Q at load time, and +// v_scale into the epilogue, so neither ever enters the softmax denominator. + +constexpr int kLatentThreads = 256; +constexpr int kLatentTile = 64; // KV tokens scored per iteration + +// Shared memory: Q (head_size) + output accumulator (v_head_size) + tile logits + block reduction +// scratchpad, all fp32, followed by the tile's resolved page ids. +size_t GetPagedLatentSharedMemoryBytes(const int head_size, const int v_head_size) { + const size_t float_elems = static_cast(head_size) + static_cast(v_head_size) + + kLatentTile + kLatentThreads; + return float_elems * sizeof(float) + static_cast(kLatentTile) * sizeof(int); +} + +template +__global__ void PagedLatentAttentionKernel(const T* __restrict__ query, // [token, N, head_size] + const TCACHE* __restrict__ key_cache, // paged, head_size wide + const TCACHE* __restrict__ value_cache, // == key_cache in LATENT + const float* __restrict__ k_scale, + const float* __restrict__ v_scale, + const int* __restrict__ cumulative_seqlens_q, + const int* __restrict__ past_seqlens, + const int* __restrict__ block_table, + T* __restrict__ output, // [token, N, v_head_size] + const int batch_size, + const int num_heads, + const int kv_num_heads, + const int head_size, + const int v_head_size, + const int block_size, + const int max_num_blocks_per_seq, + const float scale, + const float softcap, + const int local_window_size, + const bool k_per_channel, + const bool v_per_channel) { + extern __shared__ float paged_latent_smem[]; + float* q_sh = paged_latent_smem; + float* acc_sh = q_sh + head_size; + float* logits_sh = acc_sh + v_head_size; + float* red_sh = logits_sh + kLatentTile; + int* block_id_sh = reinterpret_cast(red_sh + kLatentThreads); + + const int token_id = blockIdx.x; + const int head_id = blockIdx.y; + const int tid = threadIdx.x; + + // Locate the sequence this packed token belongs to and its position inside that sequence. + // cumulative_seqlens_q is a non-decreasing prefix sum, so a binary search is exact for ragged + // batches, including sequences that contribute zero new tokens. + int left = 0; + int right = batch_size - 1; + while (left < right) { + const int mid = left + (right - left) / 2; + if (token_id < cumulative_seqlens_q[mid + 1]) { + right = mid; + } else { + left = mid + 1; + } + } + const int batch_id = left; + const int s = token_id - cumulative_seqlens_q[batch_id]; + + // Causality: this token's logical position is past_seqlens[b] + s, and it attends every cached + // position up to and including its own. That is exactly FlashAttention's bottom-right-aligned + // causal convention for seqlen_k = past + seqlen_q. + const int kv_end = past_seqlens[batch_id] + s + 1; + int kv_begin = 0; + if (local_window_size > 0) { + // local_window_size counts the current token, matching mha_varlen_fwd's window_size_left = W-1. + kv_begin = max(0, kv_end - local_window_size); + } + + const int kv_head_id = head_id / (num_heads / kv_num_heads); + const int64_t token_stride_in_page = static_cast(kv_num_heads) * head_size; + const int64_t head_offset_in_page = static_cast(kv_head_id) * head_size; + + // Fold k_scale into Q: dot(q * k_scale, k_raw) == dot(q, dequant(k)), exactly. + const T* q_ptr = query + (static_cast(token_id) * num_heads + head_id) * head_size; + for (int c = tid; c < head_size; c += kLatentThreads) { + q_sh[c] = static_cast(q_ptr[c]) * GetCacheScale(k_scale, kv_head_id * head_size + c, k_per_channel); + } + for (int c = tid; c < v_head_size; c += kLatentThreads) { + acc_sh[c] = 0.0f; + } + __syncthreads(); + + constexpr int kNumWarps = kLatentThreads / 32; + const int warp_id = tid / 32; + const int lane_id = tid % 32; + // Match FlashAttention's softcap spelling: softcap * tanh(qk * scale / softcap). + const float softcap_scale = softcap > 0.0f ? (scale / softcap) : 0.0f; + + float m_state = -FLT_MAX; + float l_state = 0.0f; + + for (int tile_begin = kv_begin; tile_begin < kv_end; tile_begin += kLatentTile) { + const int tile_len = min(kLatentTile, kv_end - tile_begin); + + // ---- QK over the full head_size (compressed_kv and k_pe together) ---- + for (int t = warp_id; t < tile_len; t += kNumWarps) { + const int pos = tile_begin + t; + const int block_index = pos / block_size; + const int block_id = block_index < max_num_blocks_per_seq + ? block_table[batch_id * max_num_blocks_per_seq + block_index] + : -1; + float dot = 0.0f; + if (block_id >= 0) { + const TCACHE* k_ptr = key_cache + + (static_cast(block_id) * block_size + (pos % block_size)) * token_stride_in_page + + head_offset_in_page; + for (int c = lane_id; c < head_size; c += 32) { + dot += q_sh[c] * CacheToFloat(k_ptr[c]); + } + } +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + dot += __shfl_xor_sync(0xFFFFFFFFU, dot, offset); + } + if (lane_id == 0) { + block_id_sh[t] = block_id; + logits_sh[t] = block_id < 0 ? -FLT_MAX + : (softcap > 0.0f ? softcap * tanhf(dot * softcap_scale) : dot * scale); + } + } + __syncthreads(); + + // ---- tile max ---- + float local_max = -FLT_MAX; + for (int t = tid; t < tile_len; t += kLatentThreads) { + local_max = fmaxf(local_max, logits_sh[t]); + } + red_sh[tid] = local_max; + __syncthreads(); + for (int stride = kLatentThreads / 2; stride > 0; stride >>= 1) { + if (tid < stride) { + red_sh[tid] = fmaxf(red_sh[tid], red_sh[tid + stride]); + } + __syncthreads(); + } + const float m_tile = red_sh[0]; + __syncthreads(); + + // A tile whose pages are all unmapped contributes nothing; skipping it also avoids the + // degenerate rescale exp(-FLT_MAX - -FLT_MAX) == 1. + if (m_tile == -FLT_MAX) { + continue; + } + + const float m_new = fmaxf(m_state, m_tile); + const float alpha = __expf(m_state - m_new); // 0 on the first contributing tile + + float local_sum = 0.0f; + for (int t = tid; t < tile_len; t += kLatentThreads) { + const float p = __expf(logits_sh[t] - m_new); + logits_sh[t] = p; + local_sum += p; + } + red_sh[tid] = local_sum; + __syncthreads(); + for (int stride = kLatentThreads / 2; stride > 0; stride >>= 1) { + if (tid < stride) { + red_sh[tid] += red_sh[tid + stride]; + } + __syncthreads(); + } + const float sum_tile = red_sh[0]; + __syncthreads(); + + l_state = l_state * alpha + sum_tile; + m_state = m_new; + + // ---- PV over the leading v_head_size channels of the same cached rows ---- + for (int c = tid; c < v_head_size; c += kLatentThreads) { + float acc = acc_sh[c] * alpha; + for (int t = 0; t < tile_len; ++t) { + const int block_id = block_id_sh[t]; + if (block_id < 0) { + continue; + } + const int pos = tile_begin + t; + const TCACHE* v_ptr = value_cache + + (static_cast(block_id) * block_size + (pos % block_size)) * token_stride_in_page + + head_offset_in_page; + acc += logits_sh[t] * CacheToFloat(v_ptr[c]); + } + acc_sh[c] = acc; + } + __syncthreads(); + } + + const int64_t out_base = (static_cast(token_id) * num_heads + head_id) * v_head_size; + const float inv_l = (m_state == -FLT_MAX || l_state <= 0.0f) ? 0.0f : (1.0f / l_state); + for (int c = tid; c < v_head_size; c += kLatentThreads) { + // V channels are a prefix of the head_size-wide cached row, so a PER_CHANNEL scale is indexed + // with the head_size stride even though only v_head_size of those channels are read. + output[out_base + c] = static_cast(acc_sh[c] * inv_l * + GetCacheScale(v_scale, kv_head_id * head_size + c, v_per_channel)); + } +} + +template +Status LaunchPagedLatentAttention(const T* query, const TCACHE* key_cache, const TCACHE* value_cache, + const float* k_scale, const float* v_scale, const bool k_per_channel, + const bool v_per_channel, const int* cumulative_seqlens_q, + const int* past_seqlens, const int* block_table, T* output, + const int batch_size, const int num_heads, const int kv_num_heads, + const int head_size, const int v_head_size, const int block_size, + const int max_num_blocks_per_seq, const int token_count, const float scale, + const float softcap, const int local_window_size, cudaStream_t stream) { + const size_t smem_bytes = GetPagedLatentSharedMemoryBytes(head_size, v_head_size); + const dim3 grid(token_count, num_heads); + PagedLatentAttentionKernel<<>>( + query, key_cache, value_cache, k_scale, v_scale, cumulative_seqlens_q, past_seqlens, block_table, output, + batch_size, num_heads, kv_num_heads, head_size, v_head_size, block_size, max_num_blocks_per_seq, scale, + softcap, local_window_size, k_per_channel, v_per_channel); + return CUDA_CALL(cudaGetLastError()); +} + +////////// Launch Kernels + +// Prologue shared by every backend: unpack packed QKV, run the fused QK-Norm + rotary kernel when +// requested, and scatter K/V into the paged cache (quantizing on the way in when the cache is +// quantized). None of this depends on which attention backend runs afterwards. On return +// *query_out points at the densified, post-prologue Q. +template +Status PrepareQueryAndCache(cudaStream_t stream, contrib::PagedAttentionParameters& parameters, + PagedAttentionData& data, const int max_threads_per_block, + T** query_out) { const int batch_size = parameters.batch_size; const int token_count = parameters.token_count; const int q_hidden_size = parameters.hidden_size; @@ -352,15 +1332,6 @@ Status FlashAttention( const int num_heads = parameters.num_heads; const int kv_num_heads = parameters.kv_num_heads; const int head_size = parameters.head_size; - const float softcap = parameters.softcap; - bool is_bf16 = std::is_same::value; - const int local_window_size = parameters.local_window_size; - const int max_num_blocks_per_seq = parameters.max_num_blocks_per_seq; - const int block_size = parameters.block_size; - // Host-computed actual max from paged_attention.cc. Used as both - // `params.seqlen_q` for mha_varlen_fwd and grid.x for the rotary kernel. - const int max_query_len = data.max_query_len; - const int max_seq_len = parameters.max_num_blocks_per_seq * parameters.block_size; T* query = const_cast(data.query); T* key; @@ -373,24 +1344,24 @@ Status FlashAttention( value = reinterpret_cast(key) + static_cast(kv_num_heads * head_size); } - // cumulative_seqlens_kv is populated by the caller (paged_attention.cc) before QkvToContext; - // shared across FA and MEA dispatch paths so the host can also read total_kv_tokens. int* cumulative_seqlens_q = const_cast(data.cumulative_seqlens_q); int* past_seqlens = const_cast(data.past_seqlens); - int* cumulative_seqlens_kv = data.cumulative_seqlens_kv; - if (parameters.do_rotary) { - // Will unpack Q and K in case of packed_qkv + if (parameters.do_rotary || parameters.use_qk_norm) { + // Fused QK-Norm + rotary prologue. Also unpacks Q and K in case of packed_qkv. auto q_buffer = data.workspace_buffer; auto k_buffer = data.workspace_buffer + token_count * num_heads * head_size; const int packed_seq_stride = parameters.is_packed_qkv ? (num_heads + 2 * kv_num_heads) * head_size : -1; - ORT_RETURN_IF_ERROR(LaunchRotaryEmbeddingKernel( - stream, q_buffer, query, past_seqlens, cumulative_seqlens_q, data.cos_cache, data.sin_cache, batch_size, - max_query_len, num_heads, head_size, parameters.rotary_dim, parameters.rotary_interleaved, packed_seq_stride, + const int rotary_dim = parameters.do_rotary ? parameters.rotary_dim : 0; + ORT_RETURN_IF_ERROR(LaunchQkNormRotaryKernel( + stream, q_buffer, query, past_seqlens, cumulative_seqlens_q, data.cos_cache, data.sin_cache, + data.q_norm_weight, parameters.qk_norm_epsilon, batch_size, token_count, num_heads, head_size, + rotary_dim, parameters.rotary_offset, parameters.rotary_interleaved, packed_seq_stride, max_threads_per_block)); - ORT_RETURN_IF_ERROR(LaunchRotaryEmbeddingKernel( - stream, k_buffer, key, past_seqlens, cumulative_seqlens_q, data.cos_cache, data.sin_cache, batch_size, - max_query_len, kv_num_heads, head_size, parameters.rotary_dim, parameters.rotary_interleaved, packed_seq_stride, + ORT_RETURN_IF_ERROR(LaunchQkNormRotaryKernel( + stream, k_buffer, key, past_seqlens, cumulative_seqlens_q, data.cos_cache, data.sin_cache, + data.k_norm_weight, parameters.qk_norm_epsilon, batch_size, token_count, kv_num_heads, head_size, + rotary_dim, parameters.rotary_offset, parameters.rotary_interleaved, packed_seq_stride, max_threads_per_block)); query = q_buffer; key = k_buffer; @@ -403,26 +1374,317 @@ Status FlashAttention( query = q_buffer; } - // Insert key and value into block-based KV cache - int* block_table = const_cast(data.block_table); - const int key_stride = parameters.is_packed_qkv && !parameters.do_rotary ? q_hidden_size + 2 * kv_hidden_size : kv_hidden_size; + // Insert key and value into block-based KV cache. The prologue (if it ran) already densified K + // into the workspace, so only the "no prologue" packed-QKV case still needs the packed stride. + const bool k_is_packed = parameters.is_packed_qkv && !(parameters.do_rotary || parameters.use_qk_norm); + const int key_stride = k_is_packed ? q_hidden_size + 2 * kv_hidden_size : kv_hidden_size; const int value_stride = parameters.is_packed_qkv ? q_hidden_size + 2 * kv_hidden_size : kv_hidden_size; - ORT_RETURN_IF_ERROR(LaunchReshapeAndCache(key, value, data.key_cache, data.value_cache, block_table, past_seqlens, - cumulative_seqlens_q, batch_size, max_num_blocks_per_seq, token_count, - kv_hidden_size, block_size, key_stride, value_stride, stream, - max_threads_per_block)); + const bool k_per_channel = parameters.k_quant_type == KVQuantizationType::PER_CHANNEL; + const bool v_per_channel = parameters.v_quant_type == KVQuantizationType::PER_CHANNEL; + ORT_RETURN_IF_ERROR((LaunchReshapeAndCache( + key, value, data.key_cache, data.value_cache, data.k_scale, data.v_scale, k_per_channel, v_per_channel, + const_cast(data.block_table), past_seqlens, cumulative_seqlens_q, data.slot_mapping, batch_size, + parameters.max_num_blocks_per_seq, token_count, kv_hidden_size, parameters.block_size, + parameters.num_blocks, key_stride, value_stride, stream, max_threads_per_block))); + + *query_out = query; + return Status::OK(); +} + +// LATENT (absorbed MLA) backend. The latent row is written to the single physical cache by the +// shared prologue (which also applies offset RoPE to the k_pe suffix), then the unfused latent +// kernel reads K and V out of that same cache. +template +Status LatentAttention( + const cudaDeviceProp& device_prop, + cudaStream_t stream, + contrib::PagedAttentionParameters& parameters, + PagedAttentionData& data, + float scale) { + T* query = nullptr; + ORT_RETURN_IF_ERROR((PrepareQueryAndCache(stream, parameters, data, + device_prop.maxThreadsPerBlock, &query))); + + // V shares the physical elements of K, so it must be dequantized with the scale those elements + // were stored with: k_scale, not v_scale (which validation requires to be absent in LATENT). + ORT_RETURN_IF_ERROR((LaunchPagedLatentAttention( + query, data.key_cache, /*value_cache*/ data.key_cache, data.k_scale, /*v_scale*/ data.k_scale, + parameters.k_quant_type == KVQuantizationType::PER_CHANNEL, + parameters.k_quant_type == KVQuantizationType::PER_CHANNEL, + data.cumulative_seqlens_q, data.past_seqlens, data.block_table, data.output, + parameters.batch_size, parameters.num_heads, parameters.kv_num_heads, parameters.head_size, + parameters.v_head_size, parameters.block_size, parameters.max_num_blocks_per_seq, + parameters.token_count, scale, parameters.softcap, parameters.local_window_size, stream))); + + DUMP_TENSOR_INIT(); + DUMP_TENSOR("latent (MLA) paged attention output", data.output, parameters.token_count, parameters.num_heads, + parameters.v_head_size); + + return Status::OK(); +} + +template +Status PagedDecodeAttention( + const cudaDeviceProp& device_prop, + cudaStream_t stream, + contrib::PagedAttentionParameters& parameters, + PagedAttentionData& data, + float scale) { + T* query = nullptr; + ORT_RETURN_IF_ERROR((PrepareQueryAndCache(stream, parameters, data, + device_prop.maxThreadsPerBlock, &query))); + + ORT_RETURN_IF_ERROR((LaunchPagedDecodeAttention( + query, data.key_cache, data.value_cache, data.k_scale, data.v_scale, + parameters.k_quant_type == KVQuantizationType::PER_CHANNEL, + parameters.v_quant_type == KVQuantizationType::PER_CHANNEL, + data.cumulative_seqlens_q, data.cumulative_seqlens_kv, data.block_table, data.head_sink, data.output, + data.decode_partial_out, data.decode_partial_max, data.decode_partial_sum, + parameters.batch_size, parameters.num_heads, parameters.kv_num_heads, parameters.head_size, + parameters.block_size, parameters.max_num_blocks_per_seq, parameters.token_count, data.num_splits, + scale, parameters.softcap, parameters.local_window_size, parameters.use_smooth_softmax, stream))); + + DUMP_TENSOR_INIT(); + DUMP_TENSOR("paged decode attention output", data.output, parameters.token_count, parameters.num_heads, + parameters.head_size); + + return Status::OK(); +} + +////////// Paged XQA decode backend +// +// PagedDecodeSplitKV above is a portable scalar kernel: it dequantizes one cache element per +// thread into fp32 and reduces through shared memory, which sustains only a small fraction of +// HBM bandwidth. XQA is the tensor-core decode kernel already used by GroupQueryAttention, and +// TensorRT-LLM's copy of it (contrib_ops/cuda/bert/xqa) supports a paged cache directly. Three +// things have to be reconciled to use it here: +// +// 1. Page size. XQA requires tokensPerPage to divide its CTA tile in the sequence dimension, so +// the kernels are compiled for kXqaTokensPerPage (128) tokens. PagedAttention's block_size is +// a graph attribute (256 by default). Because the KV pool is contiguous -- +// [num_blocks, block_size, kv_num_heads, head_size] -- and XQA's PAGED_KV_CACHE_LAYOUT == 1 +// page is exactly [tokens_per_page, kv_num_heads, head_size], block b is bit-for-bit the +// concatenation of pages [b * pages_per_block, (b + 1) * pages_per_block). ExpandBlockTable- +// ToPages rewrites the block table accordingly; it costs O(batch * max_num_blocks_per_seq). +// 2. Per-channel scales. XQA only accepts a scalar dequantization scale per cache. A PER_CHANNEL +// scale is folded out exactly the same way GroupQueryAttention does it (see the derivation +// next to LaunchScaleHeadsByChannelScale in group_query_attention_qdq.cuh): k_scale into Q +// (it multiplies the QK contraction dim) and v_scale into the attention output (it is a free +// dim of the PV accumulation, so it never touches the softmax denominator). +// 3. Attention sinks. XQA consumes them as fp32, laid out [kv_head][group] -- which is ORT's +// [num_heads] order -- so only a dtype conversion is needed. + +// block_table [batch, max_num_blocks_per_seq] -> page_table [batch, max_num_blocks_per_seq * +// pages_per_block]. An unmapped block (-1) expands to unmapped pages. +__global__ void ExpandBlockTableToPages(const int* __restrict__ block_table, + int* __restrict__ page_table, + const int max_num_blocks_per_seq, + const int pages_per_block, + const int total_pages) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= total_pages) { + return; + } + const int pages_per_seq = max_num_blocks_per_seq * pages_per_block; + const int seq = i / pages_per_seq; + const int page_in_seq = i - seq * pages_per_seq; + const int block_id = block_table[seq * max_num_blocks_per_seq + page_in_seq / pages_per_block]; + page_table[i] = block_id < 0 ? -1 : block_id * pages_per_block + (page_in_seq % pages_per_block); +} + +// Multiply every head vector by a PER_CHANNEL scale indexed [kv_head, channel]. Used to fold +// k_scale into Q before XQA and v_scale into XQA's output afterwards. dst may alias src (the +// output scaling is done in place), so neither pointer is marked __restrict__. +template +__global__ void PagedFoldChannelScaleKernel(T* dst, + const T* src, + const float* __restrict__ channel_scale, + const int num_heads, const int head_size, + const int group_size, const int64_t total_elements) { + const int64_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i >= total_elements) { + return; + } + const int h = static_cast(i / head_size) % num_heads; + const int c = static_cast(i % head_size); + dst[i] = static_cast(static_cast(src[i]) * channel_scale[(h / group_size) * head_size + c]); +} + +template +__global__ void PagedConvertHeadSinkToFloatKernel(float* __restrict__ dst, const T* __restrict__ src, + const int count) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < count) { + dst[i] = static_cast(src[i]); + } +} + +template +Status PagedXqaDecodeAttention( + const cudaDeviceProp& device_prop, + cudaStream_t stream, + contrib::PagedAttentionParameters& parameters, + PagedAttentionData& data, + float scale) { + const int max_threads_per_block = device_prop.maxThreadsPerBlock; + const int batch_size = parameters.batch_size; + const int num_heads = parameters.num_heads; + const int kv_num_heads = parameters.kv_num_heads; + const int head_size = parameters.head_size; + + T* query = nullptr; + ORT_RETURN_IF_ERROR((PrepareQueryAndCache(stream, parameters, data, max_threads_per_block, &query))); + + const int pages_per_block = parameters.block_size / onnxruntime::contrib::cuda::kXqaTokensPerPage; + const int max_pages_per_seq = parameters.max_num_blocks_per_seq * pages_per_block; + { + const int total_pages = batch_size * max_pages_per_seq; + const int blocks = (total_pages + max_threads_per_block - 1) / max_threads_per_block; + ExpandBlockTableToPages<<>>( + data.block_table, data.xqa_page_table, parameters.max_num_blocks_per_seq, pages_per_block, total_pages); + CUDA_RETURN_IF_ERROR(cudaGetLastError()); + } + + const bool k_per_channel = parameters.k_quant_type == KVQuantizationType::PER_CHANNEL; + const bool v_per_channel = parameters.v_quant_type == KVQuantizationType::PER_CHANNEL; + const int64_t q_elements = static_cast(batch_size) * num_heads * head_size; + + if (k_per_channel) { + // Q may point straight at the (const) graph input when there is no packed-QKV / rotary + // prologue, so the scaled copy always goes to a dedicated scratch buffer. + const int blocks = static_cast((q_elements + max_threads_per_block - 1) / max_threads_per_block); + PagedFoldChannelScaleKernel<<>>( + data.xqa_query, query, data.k_scale, num_heads, head_size, num_heads / kv_num_heads, q_elements); + CUDA_RETURN_IF_ERROR(cudaGetLastError()); + query = data.xqa_query; + } + + const float* attention_sinks = nullptr; + if (parameters.use_smooth_softmax && data.head_sink != nullptr) { + const int blocks = (num_heads + max_threads_per_block - 1) / max_threads_per_block; + PagedConvertHeadSinkToFloatKernel<<>>( + data.xqa_head_sink, data.head_sink, num_heads); + CUDA_RETURN_IF_ERROR(cudaGetLastError()); + attention_sinks = data.xqa_head_sink; + } + + constexpr bool kIsFp8Cache = +#if defined(USE_FP8_KV_CACHE) && !defined(DISABLE_FLOAT8_TYPES) + std::is_same::value; +#else + false; +#endif + const XqaQuantType kv_quant_type = kIsFp8Cache ? XqaQuantType::kFp8 : XqaQuantType::kInt8; + ORT_RETURN_IF_ERROR(LaunchXQAPagedKernel( + device_prop, stream, + reinterpret_cast(query), + reinterpret_cast(data.key_cache), + reinterpret_cast(data.value_cache), + reinterpret_cast(data.output), + data.xqa_page_table, + batch_size, num_heads, kv_num_heads, head_size, max_pages_per_seq, + scale, parameters.local_window_size, data.past_seqlens, attention_sinks, + // A PER_CHANNEL scale has already been folded into Q / will be applied to the output, so the + // kernel must use a scale of 1 (which it does when the pointer is null). + k_per_channel ? nullptr : data.k_scale, + v_per_channel ? nullptr : data.v_scale, + kv_quant_type, std::is_same::value, + data.xqa_workspace, data.xqa_workspace_size)); + + if (v_per_channel) { + const int blocks = static_cast((q_elements + max_threads_per_block - 1) / max_threads_per_block); + PagedFoldChannelScaleKernel<<>>( + data.output, data.output, data.v_scale, num_heads, head_size, num_heads / kv_num_heads, q_elements); + CUDA_RETURN_IF_ERROR(cudaGetLastError()); + } + + DUMP_TENSOR_INIT(); + DUMP_TENSOR("paged xqa decode attention output", data.output, parameters.token_count, parameters.num_heads, + parameters.head_size); + + return Status::OK(); +} + +#if USE_FLASH_ATTENTION +template +Status FlashAttention( + const cudaDeviceProp& device_prop, + cudaStream_t stream, + contrib::PagedAttentionParameters& parameters, + PagedAttentionData& data, + float scale) { + // Get parameters + const int max_threads_per_block = device_prop.maxThreadsPerBlock; + const int batch_size = parameters.batch_size; + const int token_count = parameters.token_count; + const int num_heads = parameters.num_heads; + const int kv_num_heads = parameters.kv_num_heads; + const int head_size = parameters.head_size; + const float softcap = parameters.softcap; + bool is_bf16 = std::is_same::value; + const int local_window_size = parameters.local_window_size; + const int max_num_blocks_per_seq = parameters.max_num_blocks_per_seq; + const int block_size = parameters.block_size; + // Upper bound on the number of query tokens any one sequence contributes, from paged_attention.cc. + // mha_varlen_fwd only uses it as `params.seqlen_q` to size the grid in the query dimension; every + // actual per-sequence length is re-read from cu_seqlens inside the kernel, and an m-block past a + // sequence's real length exits immediately. An over-estimate therefore costs empty blocks, not + // correctness (docs/contrib_ops/cuda/paged_attention.md section 4.7). + const int max_query_len = data.max_query_len; + + // cumulative_seqlens_kv is populated by the caller (paged_attention.cc) before QkvToContext; + // shared across the FA and MEA dispatch paths. + int* cumulative_seqlens_q = const_cast(data.cumulative_seqlens_q); + int* cumulative_seqlens_kv = data.cumulative_seqlens_kv; + int* block_table = const_cast(data.block_table); + + T* query = nullptr; + ORT_RETURN_IF_ERROR((PrepareQueryAndCache(stream, parameters, data, max_threads_per_block, &query))); + const bool k_per_channel = parameters.k_quant_type == KVQuantizationType::PER_CHANNEL; + const bool v_per_channel = parameters.v_quant_type == KVQuantizationType::PER_CHANNEL; // Launch kernel void* q = reinterpret_cast(query); - void* key_cache = reinterpret_cast(data.key_cache); - void* value_cache = reinterpret_cast(data.value_cache); void* output = reinterpret_cast(data.output); void* softmax_lse = reinterpret_cast(data.softmax_lse); - ORT_RETURN_IF_ERROR(onnxruntime::flash::mha_varlen_fwd( - device_prop, stream, q, key_cache, value_cache, output, cumulative_seqlens_q, cumulative_seqlens_kv, - /*seqused_k*/ nullptr, block_table, softmax_lse, batch_size, num_heads, kv_num_heads, head_size, - max_query_len, max_seq_len, token_count, scale, softcap, /*is_causal*/ true, is_bf16, local_window_size - 1, - max_num_blocks_per_seq, block_size)); + + if constexpr (IsQuantizedCache::value) { + // FlashAttention cannot read a quantized page, so dequantize the live context into a dense + // packed-varlen [total_kv_tokens, kv_num_heads, head_size] buffer (no GQA expansion — Flash + // does the grouping itself) and use the non-paged varlen entry point. That path leaves + // params.num_splits at 0 exactly like the paged one, so the fp32 [num_heads, token_count] + // softmax_lse layout the head-sink epilogue relies on is unchanged. + ORT_RETURN_IF_ERROR((LaunchGatherAndExpandPagedKVCache( + data.key_cache, data.value_cache, data.gathered_key, data.gathered_value, + data.k_scale, data.v_scale, k_per_channel, v_per_channel, + block_table, cumulative_seqlens_kv, batch_size, /*num_heads*/ kv_num_heads, kv_num_heads, + head_size, block_size, max_num_blocks_per_seq, data.total_kv_tokens, stream, max_threads_per_block))); + + ORT_RETURN_IF_ERROR(onnxruntime::flash::mha_varlen_fwd( + device_prop, stream, q, reinterpret_cast(data.gathered_key), + reinterpret_cast(data.gathered_value), output, cumulative_seqlens_q, cumulative_seqlens_kv, + /*seqused_k*/ nullptr, /*block_table*/ nullptr, softmax_lse, batch_size, num_heads, kv_num_heads, head_size, + max_query_len, data.max_kv_len, token_count, scale, softcap, /*is_causal*/ true, is_bf16, + local_window_size - 1)); + } else { + void* key_cache = reinterpret_cast(data.key_cache); + void* value_cache = reinterpret_cast(data.value_cache); + const int max_seq_len = max_num_blocks_per_seq * block_size; + ORT_RETURN_IF_ERROR(onnxruntime::flash::mha_varlen_fwd( + device_prop, stream, q, key_cache, value_cache, output, cumulative_seqlens_q, cumulative_seqlens_kv, + /*seqused_k*/ nullptr, block_table, softmax_lse, batch_size, num_heads, kv_num_heads, head_size, + max_query_len, max_seq_len, token_count, scale, softcap, /*is_causal*/ true, is_bf16, local_window_size - 1, + max_num_blocks_per_seq, block_size)); + } + + if (parameters.use_smooth_softmax) { + // Rescale by the softmax denominator that the sink logit adds. mha_varlen_fwd leaves + // params.num_splits at 0, so the split-combine kernel never runs and softmax_lse carries the + // unpadded [num_heads, token_count] fp32 layout this epilogue expects. If varlen ever enables + // num_splits > 1, both the layout and this epilogue must be revisited. + ORT_RETURN_IF_ERROR(LaunchApplyHeadSink(data.output, data.softmax_lse, data.head_sink, token_count, + num_heads, head_size, stream, max_threads_per_block)); + } DUMP_TENSOR_INIT(); DUMP_TENSOR("flash attention output", data.output, token_count, num_heads, head_size); @@ -437,82 +1699,46 @@ Status FlashAttention( // the paged KV cache into a packed-varlen [total_kv_tokens, num_heads, head_size] buffer and // dispatches to CUTLASS memory-efficient attention via its seqstart_q / seqstart_k varlen ABI. // Caller must populate data.gathered_key / data.gathered_value / data.total_kv_tokens. -template +template Status EfficientAttention( const cudaDeviceProp& device_prop, cudaStream_t stream, contrib::PagedAttentionParameters& parameters, - PagedAttentionData& data, + PagedAttentionData& data, float scale) { const int max_threads_per_block = device_prop.maxThreadsPerBlock; const int batch_size = parameters.batch_size; const int token_count = parameters.token_count; - const int q_hidden_size = parameters.hidden_size; - const int kv_hidden_size = parameters.kv_hidden_size; const int num_heads = parameters.num_heads; const int kv_num_heads = parameters.kv_num_heads; const int head_size = parameters.head_size; const int block_size = parameters.block_size; const int max_num_blocks_per_seq = parameters.max_num_blocks_per_seq; const int local_window_size = parameters.local_window_size; + // Upper bounds from paged_attention.cc, not exact values. total_kv_tokens only sizes the gather + // (the gather kernel skips indices past the real end of the packed layout) and max_query_len only + // sizes MEA's `grid_x = ceil_div(sequence_length, kQueriesPerBlock)`; in varlen mode the CUTLASS + // kernel re-reads num_queries / num_keys from seqstart_q / seqstart_k on device, so a block past + // a sequence's real length returns without doing any work. const int total_kv_tokens = data.total_kv_tokens; - // Use the caller-computed actual max of per-batch new-query lengths, not the - // `token_count - batch_size + 1` heuristic: the heuristic assumes >=1 new token per batch - // and underestimates otherwise, which would silently drop query tokens from the - // rotary grid and from MEA's `grid_x = ceil_div(sequence_length, kQueriesPerBlock)`. const int max_query_len = data.max_query_len; - T* query = const_cast(data.query); - T* key; - T* value; - if (!parameters.is_packed_qkv) { - key = const_cast(data.key); - value = const_cast(data.value); - } else { - key = reinterpret_cast(query) + static_cast(num_heads * head_size); - value = reinterpret_cast(key) + static_cast(kv_num_heads * head_size); - } - // cumulative_seqlens_kv is populated by the caller (paged_attention.cc) before QkvToContext; // shared across FA and MEA dispatch paths. int* cumulative_seqlens_q = const_cast(data.cumulative_seqlens_q); - int* past_seqlens = const_cast(data.past_seqlens); int* cumulative_seqlens_kv = data.cumulative_seqlens_kv; - - if (parameters.do_rotary) { - auto q_buffer = data.workspace_buffer; - auto k_buffer = data.workspace_buffer + token_count * num_heads * head_size; - const int packed_seq_stride = parameters.is_packed_qkv ? (num_heads + 2 * kv_num_heads) * head_size : -1; - ORT_RETURN_IF_ERROR(LaunchRotaryEmbeddingKernel( - stream, q_buffer, query, past_seqlens, cumulative_seqlens_q, data.cos_cache, data.sin_cache, batch_size, - max_query_len, num_heads, head_size, parameters.rotary_dim, parameters.rotary_interleaved, packed_seq_stride, - max_threads_per_block)); - ORT_RETURN_IF_ERROR(LaunchRotaryEmbeddingKernel( - stream, k_buffer, key, past_seqlens, cumulative_seqlens_q, data.cos_cache, data.sin_cache, batch_size, - max_query_len, kv_num_heads, head_size, parameters.rotary_dim, parameters.rotary_interleaved, packed_seq_stride, - max_threads_per_block)); - query = q_buffer; - key = k_buffer; - } else if (parameters.is_packed_qkv) { - auto q_buffer = data.workspace_buffer; - const int packed_seq_stride = q_hidden_size + 2 * kv_hidden_size; - ORT_RETURN_IF_ERROR(LaunchUnpackCumulative( - query, q_buffer, token_count, q_hidden_size, packed_seq_stride, stream, max_threads_per_block)); - query = q_buffer; - } - int* block_table = const_cast(data.block_table); - const int key_stride = parameters.is_packed_qkv && !parameters.do_rotary ? q_hidden_size + 2 * kv_hidden_size : kv_hidden_size; - const int value_stride = parameters.is_packed_qkv ? q_hidden_size + 2 * kv_hidden_size : kv_hidden_size; - ORT_RETURN_IF_ERROR(LaunchReshapeAndCache(key, value, data.key_cache, data.value_cache, block_table, past_seqlens, - cumulative_seqlens_q, batch_size, max_num_blocks_per_seq, token_count, - kv_hidden_size, block_size, key_stride, value_stride, stream, - max_threads_per_block)); - ORT_RETURN_IF_ERROR(LaunchGatherAndExpandPagedKVCache( + T* query = nullptr; + ORT_RETURN_IF_ERROR((PrepareQueryAndCache(stream, parameters, data, max_threads_per_block, &query))); + const bool k_per_channel = parameters.k_quant_type == KVQuantizationType::PER_CHANNEL; + const bool v_per_channel = parameters.v_quant_type == KVQuantizationType::PER_CHANNEL; + + ORT_RETURN_IF_ERROR((LaunchGatherAndExpandPagedKVCache( data.key_cache, data.value_cache, data.gathered_key, data.gathered_value, + data.k_scale, data.v_scale, k_per_channel, v_per_channel, block_table, cumulative_seqlens_kv, batch_size, num_heads, kv_num_heads, - head_size, block_size, max_num_blocks_per_seq, total_kv_tokens, stream, max_threads_per_block)); + head_size, block_size, max_num_blocks_per_seq, total_kv_tokens, stream, max_threads_per_block))); MemoryEfficientAttentionParams p; p.sm = device_prop.major * 10 + device_prop.minor; @@ -554,16 +1780,31 @@ Status EfficientAttention( ////////// API Functions -template +template Status QkvToContext( const cudaDeviceProp& device_prop, cublasHandle_t& /*cublas*/, Stream* ort_stream, contrib::PagedAttentionParameters& parameters, - PagedAttentionData& data) { + PagedAttentionData& data) { auto stream = static_cast(ort_stream->GetHandle()); const float scale = parameters.scale == 0.0f ? 1.f / sqrt(static_cast(parameters.head_size)) : parameters.scale; + // LATENT (MLA) has its own backend: no other kernel can serve v_head_size != head_size over a + // single aliased cache. Validation guarantees an explicit scale here, so the default above is + // never the one used. + if (parameters.is_latent_kv) { + return LatentAttention(device_prop, stream, parameters, data, scale); + } + + if (data.use_xqa_decode) { + return PagedXqaDecodeAttention(device_prop, stream, parameters, data, scale); + } + + if (data.use_paged_decode) { + return PagedDecodeAttention(device_prop, stream, parameters, data, scale); + } + #if USE_FLASH_ATTENTION if (data.use_flash_attention) { return FlashAttention(device_prop, stream, parameters, data, scale); @@ -579,21 +1820,25 @@ Status QkvToContext( return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "No PagedAttention kernel available for the current configuration."); } -template struct PagedAttentionData; -template Status QkvToContext( - const cudaDeviceProp& device_prop, - cublasHandle_t& cublas, - Stream* ort_stream, - contrib::PagedAttentionParameters& parameters, - PagedAttentionData& data); +#define INSTANTIATE_PAGED_ATTENTION(T, TCACHE) \ + template struct PagedAttentionData; \ + template Status QkvToContext( \ + const cudaDeviceProp& device_prop, \ + cublasHandle_t& cublas, \ + Stream* ort_stream, \ + contrib::PagedAttentionParameters& parameters, \ + PagedAttentionData& data); -template struct PagedAttentionData; -template Status QkvToContext( - const cudaDeviceProp& device_prop, - cublasHandle_t& cublas, - Stream* ort_stream, - contrib::PagedAttentionParameters& parameters, - PagedAttentionData& data); +INSTANTIATE_PAGED_ATTENTION(half, half) +INSTANTIATE_PAGED_ATTENTION(BFloat16, BFloat16) +INSTANTIATE_PAGED_ATTENTION(half, int8_t) +INSTANTIATE_PAGED_ATTENTION(BFloat16, int8_t) +#if defined(USE_FP8_KV_CACHE) && !defined(DISABLE_FLOAT8_TYPES) +INSTANTIATE_PAGED_ATTENTION(half, Float8E4M3FN) +INSTANTIATE_PAGED_ATTENTION(BFloat16, Float8E4M3FN) +#endif + +#undef INSTANTIATE_PAGED_ATTENTION } // namespace cuda } // namespace contrib diff --git a/onnxruntime/contrib_ops/cuda/bert/paged_attention_impl.h b/onnxruntime/contrib_ops/cuda/bert/paged_attention_impl.h index 22f9793be0af6..2ab7f1fd5727f 100644 --- a/onnxruntime/contrib_ops/cuda/bert/paged_attention_impl.h +++ b/onnxruntime/contrib_ops/cuda/bert/paged_attention_impl.h @@ -14,13 +14,13 @@ namespace onnxruntime { namespace contrib { namespace cuda { -template +template Status QkvToContext( const cudaDeviceProp& device_prop, cublasHandle_t& cublas, Stream* stream, contrib::PagedAttentionParameters& parameters, - PagedAttentionData& data); + PagedAttentionData& data); template Status LaunchUnpackQKVCumulative(const T* packed_qkv, T* unpacked_q, T* unpacked_k, T* unpacked_v, const int num_heads, @@ -32,6 +32,17 @@ Status LaunchUnpackQKVCumulative(const T* packed_qkv, T* unpacked_q, T* unpacked Status LaunchGetCumulativeSeqlensKV(int32_t* cumulative_seqlens_kv, const int32_t* cumulative_seqlens_q, const int32_t* past_seqlens, const int batch_size, cudaStream_t stream); +// Paged decode backend sizing helpers, used by paged_attention.cc to test eligibility (the kernel +// needs more dynamic shared memory than the device provides for very wide heads) and to size the +// split-KV workspaces. +size_t GetPagedDecodeSharedMemoryBytes(const int head_size); +int ComputePagedDecodeSplits(const int token_count, const int num_heads, const int max_kv_len, + const int multi_processor_count); + +// Shared memory required by the unfused latent (absorbed MLA) kernel. Used by paged_attention.cc to +// reject a latent configuration the device cannot hold instead of failing at launch time. +size_t GetPagedLatentSharedMemoryBytes(const int head_size, const int v_head_size); + } // namespace cuda } // namespace contrib } // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/bert/tensorrt_fused_multihead_attention/cross_attention/fmha_cross_attention.h b/onnxruntime/contrib_ops/cuda/bert/tensorrt_fused_multihead_attention/cross_attention/fmha_cross_attention.h index 7f363339eeef5..d80e1e037a5fb 100644 --- a/onnxruntime/contrib_ops/cuda/bert/tensorrt_fused_multihead_attention/cross_attention/fmha_cross_attention.h +++ b/onnxruntime/contrib_ops/cuda/bert/tensorrt_fused_multihead_attention/cross_attention/fmha_cross_attention.h @@ -116,6 +116,7 @@ struct Fused_multihead_attention_params_mhca { } }; +#if USE_TRT_FUSED_ATTENTION extern const unsigned char cubin_fmha_mhca_fp16_128_64_sm75_cu_cubin[]; extern const unsigned char cubin_fmha_mhca_fp16_128_64_sm80_cu_cubin[]; extern const unsigned char cubin_fmha_mhca_fp16_128_64_sm86_cu_cubin[]; @@ -141,8 +142,9 @@ extern const uint32_t cubin_fmha_mhca_fp16_128_128_sm89_cu_cubin_len; extern const uint32_t cubin_fmha_mhca_fp16_128_256_sm80_cu_cubin_len; extern const uint32_t cubin_fmha_mhca_fp16_128_256_sm86_cu_cubin_len; extern const uint32_t cubin_fmha_mhca_fp16_128_256_sm89_cu_cubin_len; +#endif // USE_TRT_FUSED_ATTENTION -static const struct FusedMultiHeadCrossAttentionKernelMetaInfoV2 { +struct FusedMultiHeadCrossAttentionKernelMetaInfoV2 { Data_type mDataType; int32_t mS; int32_t mD; @@ -154,7 +156,10 @@ static const struct FusedMultiHeadCrossAttentionKernelMetaInfoV2 { int32_t mThreadsPerCTA; int32_t mUnrollStep; bool mInterleaved; -} sMhaKernelMetaInfos[] = { +}; + +#if USE_TRT_FUSED_ATTENTION +static const FusedMultiHeadCrossAttentionKernelMetaInfoV2 sMhaKernelMetaInfos[] = { {DATA_TYPE_FP16, 128, 64, kSM_75, cubin_fmha_mhca_fp16_128_64_sm75_cu_cubin, cubin_fmha_mhca_fp16_128_64_sm75_cu_cubin_len, "fmha_mhca_fp16_128_64_sm75_kernel", 40960, 128, 0, false}, {DATA_TYPE_FP16, 128, 64, kSM_75, cubin_fmha_mhca_fp16_128_64_sm75_cu_cubin, cubin_fmha_mhca_fp16_128_64_sm75_cu_cubin_len, "fmha_mhca_fp16_128_64_sm75_kernel_nl", 36864, 128, 32, false}, @@ -178,6 +183,7 @@ static const struct FusedMultiHeadCrossAttentionKernelMetaInfoV2 { {DATA_TYPE_FP16, 128, 128, kSM_89, cubin_fmha_mhca_fp16_128_128_sm89_cu_cubin, cubin_fmha_mhca_fp16_128_128_sm89_cu_cubin_len, "fmha_mhca_fp16_128_128_sm89_kernel_nl", 81920, 128, 32, false}, {DATA_TYPE_FP16, 128, 256, kSM_89, cubin_fmha_mhca_fp16_128_256_sm89_cu_cubin, cubin_fmha_mhca_fp16_128_256_sm89_cu_cubin_len, "fmha_mhca_fp16_128_256_sm89_kernel", 163840, 256, 0, false}, {DATA_TYPE_FP16, 128, 256, kSM_89, cubin_fmha_mhca_fp16_128_256_sm89_cu_cubin, cubin_fmha_mhca_fp16_128_256_sm89_cu_cubin_len, "fmha_mhca_fp16_128_256_sm89_kernel_nl", 81920, 256, 16, false}}; +#endif // USE_TRT_FUSED_ATTENTION static Fused_multihead_attention_params_mhca getMHCAParams( // sizes @@ -273,16 +279,28 @@ using FusedMHACrossKernelFactory = TSharedCubinKernelFactory min_head_size) && (head_size <= max_head_size) && (kv_sequence_length <= 128); // TODO: shall we remove this constraint on kv_sequence_length? +#else + ORT_UNUSED_PARAMETER(sm); + ORT_UNUSED_PARAMETER(head_size); + ORT_UNUSED_PARAMETER(kv_sequence_length); + return false; +#endif } inline FusedMultiHeadCrossAttentionKernel const* get_fused_cross_attention_kernels(int32_t sm) { +#if USE_TRT_FUSED_ATTENTION return FusedMHACrossKernelFactory::Get().getCubinKernels( sMhaKernelMetaInfos, sizeof(sMhaKernelMetaInfos) / sizeof(sMhaKernelMetaInfos[0]), DATA_TYPE_FP16, sm); +#else + ORT_UNUSED_PARAMETER(sm); + return nullptr; +#endif } inline void run_fused_cross_attention( diff --git a/onnxruntime/contrib_ops/cuda/bert/tensorrt_fused_multihead_attention/flash_attention/fmha_flash_attention.h b/onnxruntime/contrib_ops/cuda/bert/tensorrt_fused_multihead_attention/flash_attention/fmha_flash_attention.h index 98d109dc35a49..7b648023bb757 100644 --- a/onnxruntime/contrib_ops/cuda/bert/tensorrt_fused_multihead_attention/flash_attention/fmha_flash_attention.h +++ b/onnxruntime/contrib_ops/cuda/bert/tensorrt_fused_multihead_attention/flash_attention/fmha_flash_attention.h @@ -25,6 +25,7 @@ namespace onnxruntime { namespace contrib { namespace cuda { +#if USE_TRT_FUSED_ATTENTION extern const unsigned char cubin_fmha_v2_flash_attention_fp16_64_64_S_64_sm70_cu_cubin[]; extern const unsigned char cubin_fmha_v2_flash_attention_fp16_64_64_S_16_sm75_cu_cubin[]; extern const unsigned char cubin_fmha_v2_flash_attention_fp16_64_64_S_32_sm75_cu_cubin[]; @@ -129,9 +130,11 @@ extern const uint32_t cubin_fmha_v2_flash_attention_fp16_128_32_S_128_sm89_cu_cu extern const uint32_t cubin_fmha_v2_flash_attention_fp16_64_16_S_160_sm89_cu_cubin_len; extern const uint32_t cubin_fmha_v2_flash_attention_fp16_64_16_S_256_sm89_cu_cubin_len; +#endif // USE_TRT_FUSED_ATTENTION + constexpr int32_t S{0}; -static const struct FusedMultiHeadFlashAttentionKernelMetaInfoV2 { +struct FusedMultiHeadFlashAttentionKernelMetaInfoV2 { Data_type mDataType; int32_t mS; int32_t mQStep; @@ -145,7 +148,10 @@ static const struct FusedMultiHeadFlashAttentionKernelMetaInfoV2 { int32_t mThreadsPerCTA; int32_t mUnrollStep; bool mInterleaved; -} sMhaKernelMetaInfos[] = { +}; + +#if USE_TRT_FUSED_ATTENTION +static const FusedMultiHeadFlashAttentionKernelMetaInfoV2 sMhaKernelMetaInfos[] = { // SM70 kernel is from FasterTransformer {DATA_TYPE_FP16, S, 64, 64, 64, kSM_70, cubin_fmha_v2_flash_attention_fp16_64_64_S_64_sm70_cu_cubin, cubin_fmha_v2_flash_attention_fp16_64_64_S_64_sm70_cu_cubin_len, "fmha_v2_flash_attention_fp16_0_64_sm70_kernel", 24576, 128, 0, false}, {DATA_TYPE_FP16, S, 64, 64, 64, kSM_70, cubin_fmha_v2_flash_attention_fp16_64_64_S_64_sm70_cu_cubin, cubin_fmha_v2_flash_attention_fp16_64_64_S_64_sm70_cu_cubin_len, "fmha_v2_flash_attention_fp16_0_64_sm70_kernel_nl", 24576, 128, 64, false}, @@ -253,6 +259,7 @@ static const struct FusedMultiHeadFlashAttentionKernelMetaInfoV2 { {DATA_TYPE_FP16, S, 64, 16, 160, kSM_89, cubin_fmha_v2_flash_attention_fp16_64_16_S_160_sm89_cu_cubin, cubin_fmha_v2_flash_attention_fp16_64_16_S_160_sm89_cu_cubin_len, "fmha_v2_flash_attention_fp16_64_16_S_160_sm89_kernel_nl", 98304, 128, 64, false}, {DATA_TYPE_FP16, S, 64, 16, 256, kSM_89, cubin_fmha_v2_flash_attention_fp16_64_16_S_256_sm89_cu_cubin, cubin_fmha_v2_flash_attention_fp16_64_16_S_256_sm89_cu_cubin_len, "fmha_v2_flash_attention_fp16_64_16_S_256_sm89_kernel", 98304, 128, 0, false}, {DATA_TYPE_FP16, S, 64, 16, 256, kSM_89, cubin_fmha_v2_flash_attention_fp16_64_16_S_256_sm89_cu_cubin, cubin_fmha_v2_flash_attention_fp16_64_16_S_256_sm89_cu_cubin_len, "fmha_v2_flash_attention_fp16_64_16_S_256_sm89_kernel_nl", 98304, 128, 64, false}}; +#endif // USE_TRT_FUSED_ATTENTION //////////////////////////////////////////////////////////////////////////////////////////////////// class FusedMultiHeadFlashAttentionKernel @@ -328,16 +335,28 @@ class FusedMultiHeadFlashAttentionKernel using FusedMHAFlashKernelFactory = TSharedCubinKernelFactory; inline FusedMultiHeadFlashAttentionKernel const* get_flash_attention_kernels(Data_type type, int32_t sm) { +#if USE_TRT_FUSED_ATTENTION return FusedMHAFlashKernelFactory::Get().getCubinKernels( sMhaKernelMetaInfos, sizeof(sMhaKernelMetaInfos) / sizeof(sMhaKernelMetaInfos[0]), type, sm); +#else + ORT_UNUSED_PARAMETER(type); + ORT_UNUSED_PARAMETER(sm); + return nullptr; +#endif } inline bool has_flash_attention_kernel(int sm, int head_size) { +#if USE_TRT_FUSED_ATTENTION return (sm == 70 && head_size == 64) || ((sm == 75 || sm == 80 || sm == 86 || sm == 89) && (head_size == 16 || head_size == 32 || head_size == 40 || head_size == 64 || head_size == 80 || head_size == 128 || head_size == 160 || head_size == 256)); +#else + ORT_UNUSED_PARAMETER(sm); + ORT_UNUSED_PARAMETER(head_size); + return false; +#endif } } // namespace cuda diff --git a/onnxruntime/contrib_ops/cuda/bert/tensorrt_fused_multihead_attention/fused_multihead_attention_v2.h b/onnxruntime/contrib_ops/cuda/bert/tensorrt_fused_multihead_attention/fused_multihead_attention_v2.h index e66ce6103e247..0167aedb85674 100644 --- a/onnxruntime/contrib_ops/cuda/bert/tensorrt_fused_multihead_attention/fused_multihead_attention_v2.h +++ b/onnxruntime/contrib_ops/cuda/bert/tensorrt_fused_multihead_attention/fused_multihead_attention_v2.h @@ -78,6 +78,23 @@ struct Fused_multihead_attention_params_v2 { }; //////////////////////////////////////////////////////////////////////////////////////////////////// +struct FusedMultiHeadAttentionKernelMetaInfoV2 { + Data_type mDataType; + unsigned int mS; + unsigned int mD; + unsigned int mSM; + const unsigned char* mCubin; + unsigned int mCubinSize; + const char* mFuncName; + unsigned int mSharedMemBytes; + unsigned int mThreadsPerCTA; + unsigned int mUnrollStep; + bool mInterleaved; + bool mWithRelativePositionBias = false; + bool mFlashAttention = false; +}; + +#if USE_TRT_FUSED_ATTENTION extern const unsigned char cubin_fmha_v2_fp16_384_32_sm80_cu_cubin[]; extern const unsigned char cubin_fmha_v2_fp16_256_32_sm80_cu_cubin[]; extern const unsigned char cubin_fmha_v2_fp16_192_32_sm80_cu_cubin[]; @@ -154,21 +171,7 @@ extern const unsigned int fused_multihead_attention_v2_fp16_128_64_kernel_sm70_c extern const unsigned int fused_multihead_attention_v2_fp16_256_64_kernel_sm70_cubin_len; extern const unsigned int fused_multihead_attention_v2_fp16_384_64_kernel_sm70_cubin_len; -static const struct FusedMultiHeadAttentionKernelMetaInfoV2 { - Data_type mDataType; - unsigned int mS; - unsigned int mD; - unsigned int mSM; - const unsigned char* mCubin; - unsigned int mCubinSize; - const char* mFuncName; - unsigned int mSharedMemBytes; - unsigned int mThreadsPerCTA; - unsigned int mUnrollStep; - bool mInterleaved; - bool mWithRelativePositionBias = false; - bool mFlashAttention = false; -} sMhaKernelMetaInfosV2[] = { +static const FusedMultiHeadAttentionKernelMetaInfoV2 sMhaKernelMetaInfosV2[] = { // Volta {DATA_TYPE_FP16, 64, @@ -1094,6 +1097,7 @@ static const struct FusedMultiHeadAttentionKernelMetaInfoV2 { false}, #endif }; +#endif // USE_TRT_FUSED_ATTENTION class FusedMultiHeadAttentionXMMAKernelV2 : public TFusedMultiHeadAttentionXMMAKernel { @@ -1230,8 +1234,14 @@ class FusedMultiHeadAttentionXMMAKernelV2 : public TFusedMultiHeadAttentionXMMAK using FusedMHAKernelFactoryV2 = TFusedMHAKernelFactory; inline const FusedMultiHeadAttentionXMMAKernelV2* getXMMAKernelsV2(Data_type type, unsigned int sm) { +#if USE_TRT_FUSED_ATTENTION return FusedMHAKernelFactoryV2::Get().getXMMAKernels( sMhaKernelMetaInfosV2, sizeof(sMhaKernelMetaInfosV2) / sizeof(sMhaKernelMetaInfosV2[0]), type, sm); +#else + ORT_UNUSED_PARAMETER(type); + ORT_UNUSED_PARAMETER(sm); + return nullptr; +#endif } } // namespace cuda diff --git a/onnxruntime/contrib_ops/cuda/bert/tensorrt_fused_multihead_attention/mha_runner.cu b/onnxruntime/contrib_ops/cuda/bert/tensorrt_fused_multihead_attention/mha_runner.cu index 7ec5fd47dfca8..ff141dc55a3b6 100644 --- a/onnxruntime/contrib_ops/cuda/bert/tensorrt_fused_multihead_attention/mha_runner.cu +++ b/onnxruntime/contrib_ops/cuda/bert/tensorrt_fused_multihead_attention/mha_runner.cu @@ -199,6 +199,13 @@ FusedMHARunnerFP16v2::FusedMHARunnerFP16v2(int num_heads, bool FusedMHARunnerFP16v2::IsSupported(int sm, int head_size, int sequence_length, bool enable_flash_attention) { +#if !USE_TRT_FUSED_ATTENTION + ORT_UNUSED_PARAMETER(sm); + ORT_UNUSED_PARAMETER(head_size); + ORT_UNUSED_PARAMETER(sequence_length); + ORT_UNUSED_PARAMETER(enable_flash_attention); + return false; +#else bool use_flash = enable_flash_attention && sequence_length >= kMinSequenceLengthFlashAttention; if (use_flash && has_flash_attention_kernel(sm, head_size)) { return true; @@ -219,6 +226,7 @@ bool FusedMHARunnerFP16v2::IsSupported(int sm, int head_size, int sequence_lengt // Normal (not flash) fused kernel supports sequence length up to 384. constexpr int max_sequence_length = 384; return sequence_length <= max_sequence_length; +#endif } void FusedMHARunnerFP16v2::Run(int batch_size, diff --git a/onnxruntime/contrib_ops/cuda/bert/xqa/mhaUtils.cuh b/onnxruntime/contrib_ops/cuda/bert/xqa/mhaUtils.cuh index 4fb427c1df8c5..dbf6374d51768 100644 --- a/onnxruntime/contrib_ops/cuda/bert/xqa/mhaUtils.cuh +++ b/onnxruntime/contrib_ops/cuda/bert/xqa/mhaUtils.cuh @@ -234,6 +234,10 @@ struct KVCacheList { const KVCachePageIndex* kvCachePageList; // shape: KVCachePageIndex[batchSize][beamWidth][2][maxNbPagesPerSeq]. const SeqLenDataType* seqLenList; // shape: [batchSize][beamWidth] (for compatibility) uint32_t maxNbPagesPerSeq; + // Added to the value read from seqLenList. ORT callers pass the *past* sequence length (the + // length before the token being decoded) plus extraSeqLen == 1, matching the contiguous + // KVCacheList specialization below and keeping getCacheSeqLen() uniform over both. + uint32_t extraSeqLen; }; template <> diff --git a/onnxruntime/contrib_ops/cuda/bert/xqa/mha_impl.cuh b/onnxruntime/contrib_ops/cuda/bert/xqa/mha_impl.cuh index 5fad67aba946d..104fec9159a62 100644 --- a/onnxruntime/contrib_ops/cuda/bert/xqa/mha_impl.cuh +++ b/onnxruntime/contrib_ops/cuda/bert/xqa/mha_impl.cuh @@ -1906,9 +1906,16 @@ CUBIN_EXPORT __global__ #endif #if USE_PAGED_KV_CACHE constexpr uint32_t xIterSeqStride = cacheVTileSeqStride * nbVItersPerXIter; + // `if constexpr` inside a non-template function still type-checks the discarded branch, so + // both divisors below must stay non-zero for every instantiation even though only one branch + // is ever live (ORT builds XQA with -Werror all-warnings, which turns a constant-folded + // "right operand of % is zero" in the dead branch into a build failure). + constexpr uint32_t nbXItersPerPage = + (xIterSeqStride <= tokensPerPage ? exactDiv(tokensPerPage, xIterSeqStride) : 1U); + constexpr uint32_t nbPagesPerXIter = + (xIterSeqStride <= tokensPerPage ? 1U : exactDiv(xIterSeqStride, tokensPerPage)); if constexpr (xIterSeqStride <= tokensPerPage) { - const uint32_t nbXItersPerPage = exactDiv(tokensPerPage, xIterSeqStride); - assert(nbXItersPerPage <= nbXItersPerCtaTile); + static_assert(nbXItersPerPage <= nbXItersPerCtaTile); if (xIter % nbXItersPerPage == nbXItersPerPage - 1 && vIter == nbVItersPerXIter - 1 && (idxBeam == beamWidth - 1 || isConvergedTile(seqIter))) { const auto step = 1; // cacheVTileSeqLen * gemm1NbWarpGrps / tokensPerPage; idxPageBeg += (idxPageBeg % nbPagesPerCtaTile == nbPagesPerCtaTile - 1 @@ -1920,7 +1927,7 @@ CUBIN_EXPORT __global__ } else { assert(nbVItersPerXIter == 1); if ((idxBeam == beamWidth - 1 || isConvergedTile(seqIter)) && vIter == nbVItersPerXIter - 1) { - const auto step = exactDiv(xIterSeqStride, tokensPerPage); + const auto step = nbPagesPerXIter; idxPageBeg += (idxPageBeg % nbPagesPerCtaTile + step >= nbPagesPerCtaTile ? nbPagesPerCtaTile * (nbSubSeqPerSeq - 1) + step : step); @@ -2542,9 +2549,9 @@ void launchMHA(const cudaDeviceProp& prop, uint32_t nbKHeads, #if USE_PAGED_KV_CACHE const uint32_t maxNbPagesPerSeq = exactDiv(maxSeqLen, tokensPerPage); #if PAGED_KV_CACHE_LAYOUT == 1 - const KVCacheList cacheList{kCacheVLLM, vCacheVLLM, kvCachePageList, seqLen, maxNbPagesPerSeq}; + const KVCacheList cacheList{kCacheVLLM, vCacheVLLM, kvCachePageList, seqLen, maxNbPagesPerSeq, 1}; #else - const KVCacheList cacheList{pool, kvCachePageList, seqLen, maxNbPagesPerSeq}; + const KVCacheList cacheList{pool, kvCachePageList, seqLen, maxNbPagesPerSeq, 1}; #endif cudaLaunchKernelEx(&launchCfg, kernel_mha, #if SPEC_DEC diff --git a/onnxruntime/contrib_ops/cuda/bert/xqa/utils.cuh b/onnxruntime/contrib_ops/cuda/bert/xqa/utils.cuh index ca21a705db4fb..cca2042ef3587 100644 --- a/onnxruntime/contrib_ops/cuda/bert/xqa/utils.cuh +++ b/onnxruntime/contrib_ops/cuda/bert/xqa/utils.cuh @@ -47,7 +47,9 @@ inline constexpr float log2e = 1.4426950408889634f; // std::log2(M_E) */ // this reason, don't set safeInitRowMax with a huge absolute value. // #define SAFE_INIT_ROW_MAX (-1e+5F) // moved to defines.h -inline constexpr int32_t kBAD_PAGE_INDEX = -1; +// Marked __constant__ (like kE4M3_MAX below) because the paged-KV kernels odr-use it from device +// code (Vec::filled takes a const reference), which a plain host constexpr variable cannot satisfy. +__constant__ constexpr int32_t kBAD_PAGE_INDEX = -1; __constant__ constexpr float kE4M3_MAX = 448.F; #ifdef __CUDA_ARCH__ @@ -668,11 +670,48 @@ __device__ __host__ inline void assertClose([[maybe_unused]] half a, [[maybe_unu assertClose(__half2float(a), __half2float(b), threshold); } +// Converts four packed signed int8 values into four half values without using I2F. +// +// A half whose bit pattern is 0x64XX is exactly 1024 + XX: the exponent field of 0x6400 +// selects 2^10, where the mantissa ULP is 1, so the low 8 mantissa bits hold the byte +// verbatim. Splicing a byte in with prmt and subtracting a magic constant therefore +// performs the conversion using only full-rate integer/half2 ALU instructions. The +// generic path instead costs one quarter-rate I2F per element plus a sign-extension +// sequence and a repack, which is the bulk of the int8-vs-fp8 gap in this kernel. +// +// Source bytes are biased first (XOR 0x80 maps two's-complement s to unsigned s + 128), +// so the spliced value is 1152 + s and the constant to subtract is half(1152) == 0x6480. +// 1152 + s lies in [1024, 1279] and the result s lies in [-128, 127]; both are exactly +// representable in half, so this is bit-identical to the I2F path. +// +// selector0 / selector1 are prmt byte selectors choosing which source byte feeds the low +// and high half of each output word, which lets callers fold a byte permutation in. +template +__device__ inline Vec cvtS8x4ToF16x4(uint32_t i8data) { + static constexpr uint32_t kSignFlip = 0x80808080U; // two's complement -> biased unsigned + static constexpr uint32_t kExpBytes = 0x64646464U; // half exponent byte for 1024.0 + static constexpr uint32_t kMagic = 0x64806480U; // half2(1152.0, 1152.0) + uint32_t const biased = i8data ^ kSignFlip; + Vec ret; + asm("prmt.b32 %0, %1, %2, %3;\n" : "=r"(ret.data[0]) : "r"(biased), "n"(kExpBytes), "n"(selector0)); + asm("prmt.b32 %0, %1, %2, %3;\n" : "=r"(ret.data[1]) : "r"(biased), "n"(kExpBytes), "n"(selector1)); + // kMagic must be a register operand: f16x2 immediates are not encodable. + asm("sub.f16x2 %0, %1, %2;\n" : "=r"(ret.data[0]) : "r"(ret.data[0]), "r"(kMagic)); + asm("sub.f16x2 %0, %1, %2;\n" : "=r"(ret.data[1]) : "r"(ret.data[1]), "r"(kMagic)); + return ret; +} + template __device__ inline Vec convertKCacheWordToF16(uint32_t i8data) { static_assert(mha::is_same_v || mha::is_same_v, "not implemented"); static_assert(sizeof(CacheElem) == 1); Vec ret; +#if (defined __CUDA_ARCH__) + if constexpr (mha::is_same_v && mha::is_same_v) { + // dst[i] = src[i], so byte i of the input feeds output half i. + return cvtS8x4ToF16x4<0x5150, 0x5352>(i8data); + } +#endif #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 890) if constexpr (mha::is_same_v && mha::is_same_v) { uint16_t (&src)[2] = reinterpret_cast(i8data); @@ -700,6 +739,13 @@ __device__ inline Vec convertVCacheWordToF16(uint32_t i8data) { static_assert(mha::is_same_v || mha::is_same_v, "not implemented"); static_assert(sizeof(CacheElem) == 1); Vec ret; +#if (defined __CUDA_ARCH__) + if constexpr (mha::is_same_v && mha::is_same_v) { + // dst[i][j] = src[j][i], i.e. the 2x2 byte transpose {b0,b1,b2,b3} -> {b0,b2,b1,b3}. + // Folded into the prmt selectors, so it costs nothing extra here. + return cvtS8x4ToF16x4<0x5250, 0x5351>(i8data); + } +#endif #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 890) if constexpr (mha::is_same_v && mha::is_same_v) { uint32_t (&dst)[2] = reinterpret_cast(ret); diff --git a/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_paged_bf16_fp8_128.cu b/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_paged_bf16_fp8_128.cu new file mode 100644 index 0000000000000..255f49ab67ec7 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_paged_bf16_fp8_128.cu @@ -0,0 +1,14 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#define HEAD_ELEMS 128 +#define HEAD_DIM_NAMESPACE H128 +#define XQA_PAGED_CACHE_ELEM 2 +#define XQA_PAGED_INPUT_FP16 0 +#define XQA_PAGED_QUERY_T __nv_bfloat16 +#define XQA_PAGED_FAMILY bf16_fp8 +#define XQA_PAGED_LAUNCH_FN LaunchXQAPagedFp8KernelBF16 + +#ifdef USE_FP8_KV_CACHE +#include "xqa_paged_loader_impl.cuh" +#endif diff --git a/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_paged_bf16_fp8_64.cu b/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_paged_bf16_fp8_64.cu new file mode 100644 index 0000000000000..e4e9893540b62 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_paged_bf16_fp8_64.cu @@ -0,0 +1,14 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#define HEAD_ELEMS 64 +#define HEAD_DIM_NAMESPACE H64 +#define XQA_PAGED_CACHE_ELEM 2 +#define XQA_PAGED_INPUT_FP16 0 +#define XQA_PAGED_QUERY_T __nv_bfloat16 +#define XQA_PAGED_FAMILY bf16_fp8 +#define XQA_PAGED_LAUNCH_FN LaunchXQAPagedFp8KernelBF16 + +#ifdef USE_FP8_KV_CACHE +#include "xqa_paged_loader_impl.cuh" +#endif diff --git a/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_paged_bf16_int8_128.cu b/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_paged_bf16_int8_128.cu new file mode 100644 index 0000000000000..1ce2956c7b2c0 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_paged_bf16_int8_128.cu @@ -0,0 +1,12 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#define HEAD_ELEMS 128 +#define HEAD_DIM_NAMESPACE H128 +#define XQA_PAGED_CACHE_ELEM 1 +#define XQA_PAGED_INPUT_FP16 0 +#define XQA_PAGED_QUERY_T __nv_bfloat16 +#define XQA_PAGED_FAMILY bf16_int8 +#define XQA_PAGED_LAUNCH_FN LaunchXQAPagedInt8KernelBF16 + +#include "xqa_paged_loader_impl.cuh" diff --git a/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_paged_bf16_int8_64.cu b/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_paged_bf16_int8_64.cu new file mode 100644 index 0000000000000..759e41b350f76 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_paged_bf16_int8_64.cu @@ -0,0 +1,12 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#define HEAD_ELEMS 64 +#define HEAD_DIM_NAMESPACE H64 +#define XQA_PAGED_CACHE_ELEM 1 +#define XQA_PAGED_INPUT_FP16 0 +#define XQA_PAGED_QUERY_T __nv_bfloat16 +#define XQA_PAGED_FAMILY bf16_int8 +#define XQA_PAGED_LAUNCH_FN LaunchXQAPagedInt8KernelBF16 + +#include "xqa_paged_loader_impl.cuh" diff --git a/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_paged_fp16_fp8_128.cu b/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_paged_fp16_fp8_128.cu new file mode 100644 index 0000000000000..81a8cf899c285 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_paged_fp16_fp8_128.cu @@ -0,0 +1,14 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#define HEAD_ELEMS 128 +#define HEAD_DIM_NAMESPACE H128 +#define XQA_PAGED_CACHE_ELEM 2 +#define XQA_PAGED_INPUT_FP16 1 +#define XQA_PAGED_QUERY_T half +#define XQA_PAGED_FAMILY fp16_fp8 +#define XQA_PAGED_LAUNCH_FN LaunchXQAPagedFp8Kernel + +#ifdef USE_FP8_KV_CACHE +#include "xqa_paged_loader_impl.cuh" +#endif diff --git a/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_paged_fp16_fp8_64.cu b/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_paged_fp16_fp8_64.cu new file mode 100644 index 0000000000000..ebd658d90949c --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_paged_fp16_fp8_64.cu @@ -0,0 +1,14 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#define HEAD_ELEMS 64 +#define HEAD_DIM_NAMESPACE H64 +#define XQA_PAGED_CACHE_ELEM 2 +#define XQA_PAGED_INPUT_FP16 1 +#define XQA_PAGED_QUERY_T half +#define XQA_PAGED_FAMILY fp16_fp8 +#define XQA_PAGED_LAUNCH_FN LaunchXQAPagedFp8Kernel + +#ifdef USE_FP8_KV_CACHE +#include "xqa_paged_loader_impl.cuh" +#endif diff --git a/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_paged_fp16_int8_128.cu b/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_paged_fp16_int8_128.cu new file mode 100644 index 0000000000000..ef6c2de918304 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_paged_fp16_int8_128.cu @@ -0,0 +1,12 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#define HEAD_ELEMS 128 +#define HEAD_DIM_NAMESPACE H128 +#define XQA_PAGED_CACHE_ELEM 1 +#define XQA_PAGED_INPUT_FP16 1 +#define XQA_PAGED_QUERY_T half +#define XQA_PAGED_FAMILY fp16_int8 +#define XQA_PAGED_LAUNCH_FN LaunchXQAPagedInt8Kernel + +#include "xqa_paged_loader_impl.cuh" diff --git a/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_paged_fp16_int8_64.cu b/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_paged_fp16_int8_64.cu new file mode 100644 index 0000000000000..0195042ac419f --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_paged_fp16_int8_64.cu @@ -0,0 +1,12 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#define HEAD_ELEMS 64 +#define HEAD_DIM_NAMESPACE H64 +#define XQA_PAGED_CACHE_ELEM 1 +#define XQA_PAGED_INPUT_FP16 1 +#define XQA_PAGED_QUERY_T half +#define XQA_PAGED_FAMILY fp16_int8 +#define XQA_PAGED_LAUNCH_FN LaunchXQAPagedInt8Kernel + +#include "xqa_paged_loader_impl.cuh" diff --git a/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_paged_impl_gen.cuh b/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_paged_impl_gen.cuh new file mode 100644 index 0000000000000..8baf9c2576018 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_paged_impl_gen.cuh @@ -0,0 +1,142 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +// Template for the paged-KV XQA kernel instantiation. Mirrors xqa_impl_gen.cuh but binds the +// paged (PAGED_KV_CACHE_LAYOUT == 1 / vLLM-style) entry point of launchMHA. +// +// Expected macros: +// NAMESPACE_NAME: name of the namespace (e.g. grp8_int8_paged) +// GRP_SIZE: integer value for HEAD_GRP_SIZE + +namespace NAMESPACE_NAME { +// Undefine dependent guard to allow header re-processing +#undef MHA_H_DEPENDENT + +#define HEAD_GRP_SIZE GRP_SIZE + +// See xqa_impl_gen.cuh for why the SM80 guard is written this way. +#undef XQA_HAS_SM80_TARGET +#ifdef __CUDA_ARCH__ +#if __CUDA_ARCH__ >= 800 +#define XQA_HAS_SM80_TARGET 1 +#endif +#elif defined(HAS_SM80_OR_LATER) || !defined(__CUDACC__) +#define XQA_HAS_SM80_TARGET 1 +#endif + +#ifdef XQA_HAS_SM80_TARGET +#include "mha_impl.cuh" +#endif + +#undef HEAD_GRP_SIZE + +template +inline Status Launch( + [[maybe_unused]] const cudaDeviceProp& device_prop, + [[maybe_unused]] cudaStream_t stream, + [[maybe_unused]] const void* query, + [[maybe_unused]] const void* key_cache, + [[maybe_unused]] const void* value_cache, + [[maybe_unused]] void* output, + [[maybe_unused]] const int* page_table, + [[maybe_unused]] const int batch_size, + [[maybe_unused]] const int num_heads, + [[maybe_unused]] const int kv_num_heads, + [[maybe_unused]] const int head_size, + [[maybe_unused]] const int max_pages_per_seq, + [[maybe_unused]] const float scale, + [[maybe_unused]] const int local_window_size, + [[maybe_unused]] const int* past_seq_lens, + [[maybe_unused]] const float* attention_sinks, + [[maybe_unused]] const float* k_cache_scale, + [[maybe_unused]] const float* v_cache_scale, + [[maybe_unused]] void* workspace, + [[maybe_unused]] size_t workspace_size) { +#ifdef XQA_HAS_SM80_TARGET + const InputHead* q_ptr = reinterpret_cast(query); + GMemCacheHead* k_ptr = reinterpret_cast(const_cast(key_cache)); + GMemCacheHead* v_ptr = reinterpret_cast(const_cast(value_cache)); + OutputHead* out_ptr = reinterpret_cast(output); + + // maxSeqLen must be a whole number of pages: launchMHA derives maxNbPagesPerSeq (the page-table + // row stride) as exactDiv(maxSeqLen, tokensPerPage). + const uint32_t max_seq_len = static_cast(max_pages_per_seq) * tokensPerPage; + + uint32_t* semaphores = nullptr; + void* scratch = nullptr; + + if (workspace != nullptr) { + uint32_t nbSeq = static_cast(batch_size * kv_num_heads); + size_t semaphore_size = nbSeq * sizeof(uint32_t); + size_t padded_sem_size = roundUp(semaphore_size, 128); + + uint32_t nbSubSeqPerSeq = computeNbSubSeqPerSeqMHA( + device_prop, + static_cast(batch_size), + static_cast(kv_num_heads), + max_seq_len); + size_t required_scratch_size = NAMESPACE_NAME::GetScratchSize(nbSeq, nbSubSeqPerSeq); + size_t total_required = padded_sem_size + required_scratch_size; + + if (workspace_size < total_required) { + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Paged XQA workspace size is too small. Expected at least ", + total_required, ", but got ", workspace_size); + } + semaphores = reinterpret_cast(workspace); + scratch = reinterpret_cast(workspace) + padded_sem_size; + + cudaMemsetAsync(semaphores, 0, semaphore_size, stream); + } + +#if SLIDING_WINDOW + // See xqa_impl_gen.cuh: -1 (global) maps to a window >= max_seq_len so no masking work is done. + uint32_t const sliding_win_size = (local_window_size > 0) + ? static_cast(local_window_size) + : max_seq_len; +#endif + + launchMHA( + device_prop, + static_cast(kv_num_heads), +#if SLIDING_WINDOW + sliding_win_size, +#endif + scale, + out_ptr, + q_ptr, + attention_sinks, + k_ptr, + v_ptr, + reinterpret_cast(page_table), + max_seq_len, + reinterpret_cast(past_seq_lens), + static_cast(batch_size), + k_cache_scale, + v_cache_scale, + semaphores, + scratch, + stream); + return Status::OK(); +#else + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "XQA is only supported on Ampere (SM80) or newer GPUs."); +#endif +} + +#ifndef GENERATE_CUBIN +// See xqa_impl_gen.cuh::GetSmemSize. +inline size_t GetSmemSize() { +#ifdef XQA_HAS_SM80_TARGET + uint32_t size = 0; + if (cudaMemcpyFromSymbol(&size, smemSize, sizeof(smemSize)) != cudaSuccess) { + (void)cudaGetLastError(); + return 0; + } + return static_cast(size); +#else + return 0; +#endif +} +#endif // GENERATE_CUBIN + +#undef XQA_HAS_SM80_TARGET +} // namespace NAMESPACE_NAME diff --git a/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_paged_loader.cu b/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_paged_loader.cu new file mode 100644 index 0000000000000..f45b00b3c75f5 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_paged_loader.cu @@ -0,0 +1,157 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +// +// Head-size / dtype / cache-dtype dispatcher for the paged-KV XQA decode kernels. The kernels +// themselves live in xqa_paged___.cu (each of which instantiates the four +// supported query/KV group sizes through xqa_paged_loader_impl.cuh). + +#include "contrib_ops/cuda/bert/xqa/xqa_paged_loader.h" + +#include +#include + +namespace onnxruntime { +namespace contrib { +namespace cuda { + +// Signature shared by every per-TU entry point. +#define XQA_PAGED_DECL(fn) \ + Status fn( \ + const cudaDeviceProp& device_prop, \ + cudaStream_t stream, \ + const void* query, \ + const void* key_cache, \ + const void* value_cache, \ + void* output, \ + const int* page_table, \ + const int batch_size, \ + const int num_heads, \ + const int kv_num_heads, \ + const int head_size, \ + const int max_pages_per_seq, \ + const float scale, \ + const int local_window_size, \ + const int* past_seq_lens, \ + const float* attention_sinks, \ + const float* k_cache_scale, \ + const float* v_cache_scale, \ + void* workspace, \ + size_t workspace_size); \ + size_t fn##_SmemSize(const int num_heads, const int kv_num_heads) + +#define XQA_PAGED_ARGS \ + device_prop, stream, query, key_cache, value_cache, output, page_table, batch_size, num_heads, \ + kv_num_heads, head_size, max_pages_per_seq, scale, local_window_size, past_seq_lens, \ + attention_sinks, k_cache_scale, v_cache_scale, workspace, workspace_size + +namespace H64 { +XQA_PAGED_DECL(LaunchXQAPagedInt8Kernel); +XQA_PAGED_DECL(LaunchXQAPagedInt8KernelBF16); +#ifdef USE_FP8_KV_CACHE +XQA_PAGED_DECL(LaunchXQAPagedFp8Kernel); +XQA_PAGED_DECL(LaunchXQAPagedFp8KernelBF16); +#endif +} // namespace H64 + +namespace H128 { +XQA_PAGED_DECL(LaunchXQAPagedInt8Kernel); +XQA_PAGED_DECL(LaunchXQAPagedInt8KernelBF16); +#ifdef USE_FP8_KV_CACHE +XQA_PAGED_DECL(LaunchXQAPagedFp8Kernel); +XQA_PAGED_DECL(LaunchXQAPagedFp8KernelBF16); +#endif +} // namespace H128 + +Status LaunchXQAPagedKernel( + const cudaDeviceProp& device_prop, + cudaStream_t stream, + const void* query, + const void* key_cache, + const void* value_cache, + void* output, + const int* page_table, + const int batch_size, + const int num_heads, + const int kv_num_heads, + const int head_size, + const int max_pages_per_seq, + const float scale, + const int local_window_size, + const int* past_seq_lens, + const float* attention_sinks, + const float* k_cache_scale, + const float* v_cache_scale, + const XqaQuantType kv_quant_type, + const bool is_bf16, + void* workspace, + size_t workspace_size) { + if (device_prop.major < 8) { + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "XQA is only supported on Ampere (SM80) or newer GPUs."); + } + + if (kv_quant_type == XqaQuantType::kInt8) { + if (head_size == 64) { + return is_bf16 ? H64::LaunchXQAPagedInt8KernelBF16(XQA_PAGED_ARGS) + : H64::LaunchXQAPagedInt8Kernel(XQA_PAGED_ARGS); + } else if (head_size == 128) { + return is_bf16 ? H128::LaunchXQAPagedInt8KernelBF16(XQA_PAGED_ARGS) + : H128::LaunchXQAPagedInt8Kernel(XQA_PAGED_ARGS); + } + } else if (kv_quant_type == XqaQuantType::kFp8) { +#ifdef USE_FP8_KV_CACHE + if (head_size == 64) { + return is_bf16 ? H64::LaunchXQAPagedFp8KernelBF16(XQA_PAGED_ARGS) + : H64::LaunchXQAPagedFp8Kernel(XQA_PAGED_ARGS); + } else if (head_size == 128) { + return is_bf16 ? H128::LaunchXQAPagedFp8KernelBF16(XQA_PAGED_ARGS) + : H128::LaunchXQAPagedFp8Kernel(XQA_PAGED_ARGS); + } +#else + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Paged XQA was built without FP8 KV cache support."); +#endif + } else { + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, + "Paged XQA is only compiled for INT8/FP8 KV caches; the FP16/BF16 cache " + "uses the FlashAttention paged path."); + } + + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Paged XQA only supports head_size 64 or 128. Input has ", + head_size); +} + +size_t GetXQAPagedRequiredSharedMemoryBytes( + const cudaDeviceProp& device_prop, + int head_size, + int num_heads, + int kv_num_heads, + XqaQuantType kv_quant_type, + [[maybe_unused]] bool is_bf16) { + if (device_prop.major < 8 || kv_num_heads <= 0) { + return 0; + } + // FP16 and BF16 kernels have identical shared-memory footprints (both 2-byte elements), so the + // FP16 instantiation is queried for both. + if (kv_quant_type == XqaQuantType::kInt8) { + if (head_size == 64) { + return H64::LaunchXQAPagedInt8Kernel_SmemSize(num_heads, kv_num_heads); + } else if (head_size == 128) { + return H128::LaunchXQAPagedInt8Kernel_SmemSize(num_heads, kv_num_heads); + } + } else if (kv_quant_type == XqaQuantType::kFp8) { +#ifdef USE_FP8_KV_CACHE + if (head_size == 64) { + return H64::LaunchXQAPagedFp8Kernel_SmemSize(num_heads, kv_num_heads); + } else if (head_size == 128) { + return H128::LaunchXQAPagedFp8Kernel_SmemSize(num_heads, kv_num_heads); + } +#endif + } + return 0; +} + +#undef XQA_PAGED_ARGS +#undef XQA_PAGED_DECL + +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_paged_loader.h b/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_paged_loader.h new file mode 100644 index 0000000000000..f58c6c4b7f9bd --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_paged_loader.h @@ -0,0 +1,71 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include +#include +#include "core/providers/cuda/cuda_common.h" +#include "contrib_ops/cuda/bert/xqa/xqa_loader.h" + +namespace onnxruntime { +namespace contrib { +namespace cuda { + +// Number of tokens the paged XQA kernels are compiled for. +// +// XQA requires tokensPerPage to divide the kernel's CTA tile in the sequence dimension +// (mha_impl.cuh: `nbPagesPerCtaTile = exactDiv(ctaTile.x, tokensPerPage)`), which caps it at 128. +// PagedAttention's block_size is independent of this: a block of `block_size` tokens is presented +// to XQA as `block_size / kXqaTokensPerPage` consecutive pages. That remap is exact because the KV +// pool is contiguous -- [num_blocks, block_size, kv_num_heads, head_size] -- and XQA's +// PAGED_KV_CACHE_LAYOUT == 1 page is exactly [tokens_per_page, kv_num_heads, head_size]. So block b +// covers pages [b * kPagesPerBlock, (b + 1) * kPagesPerBlock). +constexpr int kXqaTokensPerPage = 128; + +// Paged-KV XQA decode launcher. Unlike LaunchXQAKernel (contiguous per-request cache) this reads +// K and V from a shared block pool addressed through a page table. +// +// Preconditions: one query token per sequence, head_size in {64, 128}, group_size in +// {4, 8, 16, 32}, quantized (INT8/FP8) cache, block_size % kXqaTokensPerPage == 0. +Status LaunchXQAPagedKernel( + const cudaDeviceProp& device_prop, + cudaStream_t stream, + const void* query, // [batch_size, num_heads, head_size] + const void* key_cache, // [num_blocks, block_size, kv_num_heads, head_size] + const void* value_cache, // [num_blocks, block_size, kv_num_heads, head_size] + void* output, // [batch_size, num_heads, head_size] + const int* page_table, // [batch_size, max_pages_per_seq], in units of kXqaTokensPerPage + const int batch_size, + const int num_heads, + const int kv_num_heads, + const int head_size, + const int max_pages_per_seq, // page-table stride; max_seq_len = max_pages_per_seq * kXqaTokensPerPage + const float scale, // softmax scale applied to Q*K.T + const int local_window_size, // -1 => global attention + const int* past_seq_lens, // [batch_size]; the kernel attends to past_seq_lens[i] + 1 tokens + const float* attention_sinks, // [num_heads] fp32, nullptr if unused + const float* k_cache_scale, // per-tensor dequant scale; nullptr means "1" (folded into Q) + const float* v_cache_scale, // per-tensor dequant scale; nullptr means "1" (applied to output) + const XqaQuantType kv_quant_type, + const bool is_bf16, // dtype of query and output + void* workspace, + size_t workspace_size); + +// Workspace bytes required by LaunchXQAPagedKernel (semaphores + multi-block scratch). The paged +// and contiguous kernels share the CTA tile and the scratch layout, so this is GetXQAScratchSize +// called with max_seq_len = max_pages_per_seq * kXqaTokensPerPage. + +// Dynamic shared memory the paged kernel requests, read from the loaded module. Returns 0 when it +// cannot be determined. Callers must skip XQA when this exceeds device_prop.sharedMemPerBlockOptin. +size_t GetXQAPagedRequiredSharedMemoryBytes( + const cudaDeviceProp& device_prop, + int head_size, + int num_heads, + int kv_num_heads, + XqaQuantType kv_quant_type, + bool is_bf16); + +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_paged_loader_impl.cuh b/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_paged_loader_impl.cuh new file mode 100644 index 0000000000000..772499cc8fae1 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_paged_loader_impl.cuh @@ -0,0 +1,179 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +// +// Shared body of the paged-KV XQA translation units. Each xqa_paged___.cu +// defines the macros below and includes this file; the macros pick the KV element type, the query +// element type, the head size and the exported entry-point name. +// +// Expected macros: +// HEAD_ELEMS head size (64 / 128 / 256) +// HEAD_DIM_NAMESPACE namespace for the head size (H64 / H128 / H256) +// XQA_PAGED_CACHE_ELEM 1 = INT8 KV cache, 2 = FP8 KV cache +// XQA_PAGED_INPUT_FP16 1 = FP16 query/output, 0 = BF16 +// XQA_PAGED_QUERY_T query element type (half / __nv_bfloat16) +// XQA_PAGED_FAMILY token used to make the per-TU namespaces unique (e.g. fp16_int8) +// XQA_PAGED_LAUNCH_FN name of the exported dispatcher for this (query, cache) pair + +#pragma once +#include "xqa_paged_loader.h" +#include + +#ifndef HEAD_ELEMS +#error "HEAD_ELEMS must be defined before including xqa_paged_loader_impl.cuh" +#endif +#ifndef HEAD_DIM_NAMESPACE +#error "HEAD_DIM_NAMESPACE must be defined before including xqa_paged_loader_impl.cuh" +#endif +#ifndef XQA_PAGED_CACHE_ELEM +#error "XQA_PAGED_CACHE_ELEM must be defined before including xqa_paged_loader_impl.cuh" +#endif +#ifndef XQA_PAGED_INPUT_FP16 +#error "XQA_PAGED_INPUT_FP16 must be defined before including xqa_paged_loader_impl.cuh" +#endif +#ifndef XQA_PAGED_QUERY_T +#error "XQA_PAGED_QUERY_T must be defined before including xqa_paged_loader_impl.cuh" +#endif +#ifndef XQA_PAGED_FAMILY +#error "XQA_PAGED_FAMILY must be defined before including xqa_paged_loader_impl.cuh" +#endif +#ifndef XQA_PAGED_LAUNCH_FN +#error "XQA_PAGED_LAUNCH_FN must be defined before including xqa_paged_loader_impl.cuh" +#endif + +#define CACHE_ELEM_ENUM XQA_PAGED_CACHE_ELEM +#define INPUT_FP16 XQA_PAGED_INPUT_FP16 +// Paged KV cache, vLLM/SGLang layout: separate K and V pools, each page laid out as +// [tokens_per_page, kv_num_heads, head_size]. This matches PagedAttention's block pool exactly. +#define TOKENS_PER_PAGE 128 +#define USE_PAGED_KV_CACHE 1 +#define PAGED_KV_CACHE_LAYOUT 1 +#define ALLOW_MULTI_BLOCK_MODE 1 +// Compiled with sliding-window support so one kernel serves both global attention +// (local_window_size == -1, mapped to a window >= max_seq_len) and sliding-window models. +#define SLIDING_WINDOW 1 + +#pragma nv_diag_suppress 177 +#pragma nv_diag_suppress 20012 + +#include "cuda_hint.cuh" +#include "mha.h" +#include "ldgsts.cuh" +#include "mhaUtils.cuh" +#include "mha_components.cuh" +#include "mma.cuh" +#include "utils.cuh" +#include "hostUtils.h" + +#undef HEAD_GRP_SIZE +#undef M_TILESIZE + +// Token-pasting helpers so a single TU body can produce unique namespace names per +// (query dtype, cache dtype) family. +#define XQA_PAGED_CAT_(a, b, c) a##b##c +#define XQA_PAGED_CAT(a, b, c) XQA_PAGED_CAT_(a, b, c) +#define XQA_PAGED_NS(grp) XQA_PAGED_CAT(grp, XQA_PAGED_FAMILY, _paged) + +namespace onnxruntime { +namespace contrib { +namespace cuda { +namespace HEAD_DIM_NAMESPACE { + +#define NAMESPACE_NAME XQA_PAGED_NS(grp4_) +#define GRP_SIZE 4 +#define M_TILESIZE 8 +#include "xqa_paged_impl_gen.cuh" +#undef NAMESPACE_NAME +#undef GRP_SIZE +#undef M_TILESIZE + +#define NAMESPACE_NAME XQA_PAGED_NS(grp8_) +#define GRP_SIZE 8 +#define M_TILESIZE 8 +#include "xqa_paged_impl_gen.cuh" +#undef NAMESPACE_NAME +#undef GRP_SIZE +#undef M_TILESIZE + +#define NAMESPACE_NAME XQA_PAGED_NS(grp16_) +#define GRP_SIZE 16 +#define M_TILESIZE 16 +#include "xqa_paged_impl_gen.cuh" +#undef NAMESPACE_NAME +#undef GRP_SIZE +#undef M_TILESIZE + +#define NAMESPACE_NAME XQA_PAGED_NS(grp32_) +#define GRP_SIZE 32 +#define M_TILESIZE 32 +#include "xqa_paged_impl_gen.cuh" +#undef NAMESPACE_NAME +#undef GRP_SIZE +#undef M_TILESIZE + +#define XQA_PAGED_DISPATCH(grp) \ + return XQA_PAGED_NS(grp)::Launch( \ + device_prop, stream, query, key_cache, value_cache, output, page_table, batch_size, num_heads, \ + kv_num_heads, head_size, max_pages_per_seq, scale, local_window_size, past_seq_lens, \ + attention_sinks, k_cache_scale, v_cache_scale, workspace, workspace_size) + +Status XQA_PAGED_LAUNCH_FN( + const cudaDeviceProp& device_prop, + cudaStream_t stream, + const void* query, + const void* key_cache, + const void* value_cache, + void* output, + const int* page_table, + const int batch_size, + const int num_heads, + const int kv_num_heads, + const int head_size, + const int max_pages_per_seq, + const float scale, + const int local_window_size, + const int* past_seq_lens, + const float* attention_sinks, + const float* k_cache_scale, + const float* v_cache_scale, + void* workspace, + size_t workspace_size) { + const int group_size = num_heads / kv_num_heads; + switch (group_size) { + case 4: + XQA_PAGED_DISPATCH(grp4_); + case 8: + XQA_PAGED_DISPATCH(grp8_); + case 16: + XQA_PAGED_DISPATCH(grp16_); + case 32: + XQA_PAGED_DISPATCH(grp32_); + default: + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, + "Paged XQA only supports group_size 4, 8, 16, 32. Input has ", group_size); + } +} + +// Shared-memory requirement of the instantiation that would actually run, so the caller can fall +// back when it exceeds the device's per-block opt-in limit. See xqa_impl_gen.cuh::GetSmemSize. +size_t XQA_PAGED_CAT(XQA_PAGED_LAUNCH_FN, _, SmemSize)(const int num_heads, const int kv_num_heads) { + const int group_size = num_heads / kv_num_heads; + switch (group_size) { + case 4: + return XQA_PAGED_NS(grp4_)::GetSmemSize(); + case 8: + return XQA_PAGED_NS(grp8_)::GetSmemSize(); + case 16: + return XQA_PAGED_NS(grp16_)::GetSmemSize(); + case 32: + return XQA_PAGED_NS(grp32_)::GetSmemSize(); + default: + return 0; + } +} + +#undef XQA_PAGED_DISPATCH + +} // namespace HEAD_DIM_NAMESPACE +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc b/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc index 1c898ea5912a8..446d71f457052 100644 --- a/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc +++ b/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc @@ -121,8 +121,14 @@ class CUDA_MS_OP_TYPED_CLASS_NAME(1, BFloat16_uint8_t, GroupQueryAttention); class CUDA_MS_OP_TYPED_CLASS_NAME(1, MLFloat16_Float8E4M3FN, GroupQueryAttention); class CUDA_MS_OP_TYPED_CLASS_NAME(1, BFloat16_Float8E4M3FN, GroupQueryAttention); #endif -class CUDA_MS_OP_TYPED_CLASS_NAME(1, MLFloat16, PagedAttention); -class CUDA_MS_OP_TYPED_CLASS_NAME(1, BFloat16, PagedAttention); +class CUDA_MS_OP_TYPED_CLASS_NAME(1, MLFloat16_MLFloat16, PagedAttention); +class CUDA_MS_OP_TYPED_CLASS_NAME(1, BFloat16_BFloat16, PagedAttention); +class CUDA_MS_OP_TYPED_CLASS_NAME(1, MLFloat16_int8_t, PagedAttention); +class CUDA_MS_OP_TYPED_CLASS_NAME(1, BFloat16_int8_t, PagedAttention); +#if defined(USE_FP8_KV_CACHE) && !defined(DISABLE_FLOAT8_TYPES) +class CUDA_MS_OP_TYPED_CLASS_NAME(1, MLFloat16_Float8E4M3FN, PagedAttention); +class CUDA_MS_OP_TYPED_CLASS_NAME(1, BFloat16_Float8E4M3FN, PagedAttention); +#endif class CUDA_MS_OP_TYPED_CLASS_NAME(1, float, DecoderAttention); class CUDA_MS_OP_TYPED_CLASS_NAME(1, MLFloat16, DecoderAttention); class CUDA_ONNX_OP_TYPED_CLASS_NAME(1, int32_t, DynamicSlice); @@ -386,8 +392,14 @@ Status RegisterCudaContribKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, #endif - BuildKernelCreateInfo, - BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, +#if defined(USE_FP8_KV_CACHE) && !defined(DISABLE_FLOAT8_TYPES) + BuildKernelCreateInfo, + BuildKernelCreateInfo, +#endif BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, diff --git a/onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemv/details.h b/onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemv/details.h index 45c0089218144..b4dc5858ad5eb 100644 --- a/onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemv/details.h +++ b/onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemv/details.h @@ -74,6 +74,47 @@ class GMemIterator { int stride_; }; +// Access shape for the CtaN per-column scales a GEMV block loads for one K tile. +// +// The scales of the CtaN columns owned by a block sit `Interleave` elements apart, so for the +// non-interleaved ColumnMajor layout (Interleave == 1) the whole CtaN-wide scale vector is +// contiguous and can be fetched with one wide access instead of CtaN scalar ones. +// +// This matters far more than the byte count suggests. With a groupwise scale (e.g. NVFP4's +// GroupSize = 16) and StepK = 8, a warp's 32 lanes cover 16 distinct scale rows that are `n` +// elements apart, so *every* scale load touches 16 different sectors. Issuing CtaN of them costs +// CtaN * 16 sectors while using only 2 bytes out of each 32-byte sector; a single vector load +// costs 16 sectors and uses CtaN * 2 bytes of each. On the Qwen3.6 NVFP4 MoE decode shape +// (CtaN = 8) that is 128 of the 176 L1 sectors a warp requests per K iteration. +// +// The vector form needs the scale row base to be 16-byte aligned. Callers guarantee +// n % CtaN == 0, and kVectorized implies CtaN * sizeof(TypeA) % 16 == 0, so both the row stride +// (a multiple of n) and the column offset (a multiple of CtaN) are 16-byte aligned. +template +struct ScalesAccess { + static constexpr int kBytes = CtaN * static_cast(sizeof(T)); + static constexpr bool kVectorized = (Interleave == 1) && (kBytes % 16 == 0); + // GMemIterator<..., TVec, Strided, Continuous, T>: the vector form loads all CtaN scales in + // `Continuous` 16-byte chunks at ii = 0, the scalar form keeps one element per `ii`. + using TVec = std::conditional_t; + static constexpr int kStrided = kVectorized ? 1 : CtaN; + static constexpr int kContinuous = kVectorized ? kBytes / 16 : 1; +}; + +// Loads the CtaN scales for K tile `iter` into `vec_scale`, using the access shape `Access` +// describes. `scales_iterator` must have been declared with that same shape. +template +__device__ __forceinline__ void load_scales(Iterator& scales_iterator, T* vec_scale, int iter) { + if constexpr (Access::kVectorized) { + scales_iterator.load(vec_scale, iter); + } else { +#pragma unroll + for (int i = 0; i < CtaN; ++i) { + scales_iterator.load(vec_scale + i, iter, i); + } + } +} + struct FP16DetailsA { using Type = half; using Type2 = half2; @@ -282,19 +323,95 @@ struct Fp4I2FConverter { #endif } +#if defined(__CUDA_ARCH__) + // Decodes four consecutive E2M1 codes into two packed AType2 words. + // + // `mag_sel` holds the four 3-bit magnitudes as the four low nibbles (bit 3 cleared so prmt + // stays in byte-select mode rather than sign-replicate mode) and `sgn_sel` holds the four sign + // bits as 0/1 nibbles. One prmt then performs *four* magnitude table lookups at once, which is + // where this beats the per-element `decode()` above: prmt selects four bytes per instruction, + // so the whole 4-element lookup costs one instruction instead of four. + // + // Selector notation below: `prmt.b32 d, a, b, c` views {a0,a1,a2,a3,b0,b1,b2,b3} as source + // bytes 0..7 and uses nibble j of `c` (nibble 0 is the least significant) to pick source byte + // for result byte j. All the selectors here are written most-significant nibble first, i.e. + // 0x1404 means d3=src1, d2=src4, d1=src0, d0=src4. + __device__ __forceinline__ static void decode_quad(uint32_t mag_sel, uint32_t sgn_sel, + uint32_t& lo2, uint32_t& hi2) { + uint32_t sb; + // Sign bytes. Source bytes are {0x00,0x80,0x00,0x00, 0,0,0,0}, so a `sgn_sel` nibble of 0 + // picks 0x00 and 1 picks 0x80 -- bit 7 of the AType high byte, i.e. the float sign bit. + asm("prmt.b32 %0, %1, %2, %3;" : "=r"(sb) : "r"(0x00008000u), "r"(0u), "r"(sgn_sel)); + if constexpr (std::is_same_v) { + uint32_t hb; + // Same magnitude table as decode(): codes 0..3 -> {0x00,0x38,0x3C,0x3E}, + // codes 4..7 -> {0x40,0x42,0x44,0x46}. hb byte j is element j's half high byte. + asm("prmt.b32 %0, %1, %2, %3;" : "=r"(hb) : "r"(0x3E3C3800u), "r"(0x46444240u), "r"(mag_sel)); + hb |= sb; + // half low byte is always 0, so expand {b0,b1,b2,b3} to {0,b0,0,b1} and {0,b2,0,b3} by + // pulling the zero bytes from the second (all-zero) prmt operand. + // 0x1404: d = {hb1, 0, hb0, 0} (byte 3..0) = half2{elem0, elem1}. + asm("prmt.b32 %0, %1, %2, %3;" : "=r"(lo2) : "r"(hb), "r"(0u), "n"(0x1404)); + // 0x3424: d = {hb3, 0, hb2, 0} = half2{elem2, elem3}. + asm("prmt.b32 %0, %1, %2, %3;" : "=r"(hi2) : "r"(hb), "r"(0u), "n"(0x3424)); + } else { + uint32_t hb, lb; + // Same two bf16 tables as decode(): high byte {0x00,0x3F,0x3F,0x3F, 0x40,0x40,0x40,0x40} + // and low byte {0x00,0x00,0x80,0xC0, 0x00,0x40,0x80,0xC0}. Byte j is element j. + asm("prmt.b32 %0, %1, %2, %3;" : "=r"(hb) : "r"(0x3F3F3F00u), "r"(0x40404040u), "r"(mag_sel)); + asm("prmt.b32 %0, %1, %2, %3;" : "=r"(lb) : "r"(0xC0800000u), "r"(0xC0804000u), "r"(mag_sel)); + hb |= sb; + // bf16 needs both bytes, so source bytes are {lb0..lb3, hb0..hb3}. + // 0x5140: d = {hb1, lb1, hb0, lb0} = bfloat162{elem0, elem1}. + asm("prmt.b32 %0, %1, %2, %3;" : "=r"(lo2) : "r"(lb), "r"(hb), "n"(0x5140)); + // 0x7362: d = {hb3, lb3, hb2, lb2} = bfloat162{elem2, elem3}. + asm("prmt.b32 %0, %1, %2, %3;" : "=r"(hi2) : "r"(lb), "r"(hb), "n"(0x7362)); + } + } +#endif + template __device__ __forceinline__ static void convert(void* src, void* dst) { - uint8_t const* s = reinterpret_cast(src); - AType* d = reinterpret_cast(dst); if constexpr (!PairInterleaved) { static_assert(N % 2 == 0); +#if defined(__CUDA_ARCH__) + if constexpr (N % 8 == 0) { + // Packed path: decode a whole 32-bit weight word (eight codes) at a time. Bit-identical + // to the scalar path below, but ~3x fewer instructions, which matters because the QMoE + // FP4 GEMV is instruction-issue bound (ncu on Qwen3.6-35B-A3B NVFP4 decode: ~74% SM + // throughput at ~12% DRAM throughput, with the dequantize sequence about half of all + // issued instructions). Measured on H200/sm_90: FP4 GEMV kernel SASS shrinks ~30% and + // the two QMoE GEMVs drop 33.2 -> 26.2 us (fc1 swiglu) and 30.2 -> 22.2 us (fc2). + // Only valid for the plain (non pair-interleaved) nibble order: nibble j of the word + // is logical element j, which is exactly what the prmt selectors below assume. + // + // Both operands are re-typed to uint32_t, so callers must supply 4-byte-aligned + // buffers; the GEMV kernels declare their tiles `alignas(alignof(uint32_t))`. + uint32_t const* sw = reinterpret_cast(src); + uint32_t* dw = reinterpret_cast(dst); #pragma unroll - for (int i = 0; i < N; i += 2) { - uint8_t byte = s[i >> 1]; - d[i] = decode(static_cast(byte & 0x0F)); - d[i + 1] = decode(static_cast((byte >> 4) & 0x0F)); + for (int i = 0; i < N / 8; ++i) { + uint32_t const w = sw[i]; + uint32_t const mag = w & 0x77777777u; + uint32_t const sgn = (w >> 3) & 0x11111111u; + decode_quad(mag, sgn, dw[i * 4 + 0], dw[i * 4 + 1]); + decode_quad(mag >> 16, sgn >> 16, dw[i * 4 + 2], dw[i * 4 + 3]); + } + } else +#endif + { + uint8_t const* s = reinterpret_cast(src); + AType* d = reinterpret_cast(dst); +#pragma unroll + for (int i = 0; i < N; i += 2) { + uint8_t byte = s[i >> 1]; + d[i] = decode(static_cast(byte & 0x0F)); + d[i + 1] = decode(static_cast((byte >> 4) & 0x0F)); + } } } else { + uint8_t const* s = reinterpret_cast(src); + AType* d = reinterpret_cast(dst); // The pair-interleave permutes whole 32-bit words, so N must cover complete words. static_assert(N % 8 == 0, "Pair-interleaved FP4 decode needs a multiple of 8 elements"); // Packing writes element i to nibble slot (i even ? i/2 : (i - 1)/2 + 4), so logical diff --git a/onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemv/dispatcher.h b/onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemv/dispatcher.h index 79f9fe1065b3e..a77df8735b26e 100644 --- a/onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemv/dispatcher.h +++ b/onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemv/dispatcher.h @@ -309,6 +309,9 @@ __global__ void kernel(TypeA* act, TypeA* act_scale, uint8_t* weight, TypeA* sca (interleaved_offset_n * interleaved_k + tid * StepK) / Details::kElemsPerByteW, CtaK / Details::kElemsPerByteW, interleaved_k / Details::kElemsPerByteW); + // Kept as CtaN scalar loads rather than the vectorized ScalesAccess/load_scales form used by + // the FP4 MoE GEMV: that form needs kInterleave == 1 to make the CtaN scales contiguous, and + // every layout reaching this dense kernel is ColumnMajorInterleaved (kInterleave 2 or 4). GMemIterator scales_iterator( scales, (GroupSize != 0 ? real_offset_k / GroupSize * n : 0) + real_offset_n, diff --git a/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv_device.cuh b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv_device.cuh index 2d4b400381c36..c79e8c9d2998a 100644 --- a/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv_device.cuh +++ b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv_device.cuh @@ -8,6 +8,7 @@ #pragma once #include +#include #include #include "core/common/common.h" @@ -18,6 +19,74 @@ namespace onnxruntime::llm { namespace kernels { namespace fpA_intB_gemv { +// Accumulator element of the K-paired inner loop, see accumulate_column_tile below. 16-bit +// accumulation keeps the two k lanes of the vec2 apart until the epilogue; fp32 accumulation +// reduces each pair to one float right away. +template +using TileAccType = + std::conditional_t, float, + typename MathWrapper::Type2>; + +// Accumulates the K tile of one output column into `acc`. +// +// This replaces the dequantize/pack_to_vec2/mma sequence of the dense fpA_intB GEMV, which pairs +// the products across two *columns* so that one hfma2 serves both. That pairing needs the decoded +// weights shuffled into column-major register pairs, costing one prmt per weight pair -- about a +// ninth of the QMoE decode GEMV's issued instructions. Pairing along K instead needs no shuffle at +// all: the activation tile is already in k order and the converters emit k pairs contiguously +// (Mapper sends a logical pair to a physical pair, and to an even physical index, so the vec2 load +// stays aligned). +// +// Apply the group scale to each decoded weight before multiplying by the activation. Besides +// matching the original dequantize-then-mma order, this prevents an unscaled fp16 product from +// overflowing even when the final scaled product is representable. The fp32 policy converts each +// scaled pair before multiplying so BF16 products are not accumulated in BF16 first. +// +// Numerics: this is a reassociation, not a rewrite. The 16-bit policy keeps the even-k and odd-k +// partial sums in the two halves of one vec2 and only adds them together in collapse_tile_acc, +// where the previous code summed a column's K terms in one serial chain. Floating-point addition +// is not associative, so results are close but not bit-identical to that order; the per-thread +// chain is now half as long, which if anything reduces accumulated rounding error. +template +__device__ __forceinline__ void accumulate_column_tile(TileAccT& acc, void const* w, void const* act, TypeA scale) { + using Math = MathWrapper; + using Type = typename Math::Type; + using Type2 = typename Math::Type2; + static_assert(K % 2 == 0); + typename Details::LayoutDetails::Mapper mapper; + + Type2 const zero = Math::to_vec2(static_cast(0.f)); + Type2 const scale2 = Math::to_vec2(scale); +#pragma unroll + for (int j = 0; j < K / 2; ++j) { + Type2 const w2 = *reinterpret_cast(reinterpret_cast(w) + mapper(2 * j)); + Type2 const scaled_w2 = Math::fma2(w2, scale2, zero); + Type2 const a2 = reinterpret_cast(act)[j]; + if constexpr (std::is_same_v) { + float2 const scaled_w_f2 = Math::to_float2(scaled_w2); + float2 const a_f2 = Math::to_float2(a2); + acc += scaled_w_f2.x * a_f2.x + scaled_w_f2.y * a_f2.y; + } else { + acc = Math::fma2(scaled_w2, a2, acc); + } + } +} + +// Collapses the K-paired accumulators into the one-value-per-column form the epilogues expect. +template +__device__ __forceinline__ void collapse_tile_acc(AccT* tile_acc, TileAccT const* tile_k_acc) { + using Math = MathWrapper; +#pragma unroll + for (int i = 0; i < CtaN; ++i) { + if constexpr (std::is_same_v) { + tile_acc[i] = tile_k_acc[i]; + } else { + float2 const p = Math::to_float2(tile_k_acc[i]); + tile_acc[i] = static_cast(p.x + p.y); + } + } +} + template __global__ void moe_gemv_kernel(TypeA* act, uint8_t* weight, TypeA* scales, TypeA* bias, TypeA* out, @@ -73,48 +142,60 @@ __global__ void moe_gemv_kernel(TypeA* act, uint8_t* weight, TypeA* scales, Type GMemIterator weight_iterator( weight, (interleaved_offset_n * interleaved_k + tid * StepK) / Details::kElemsPerByteW, CtaK / Details::kElemsPerByteW, interleaved_k / Details::kElemsPerByteW); - GMemIterator scales_iterator( - scales, - (GroupSize != 0 ? real_offset_k / GroupSize * n : 0) + real_offset_n, - (GroupSize != 0 ? CtaK / Details::kInterleave / GroupSize * n : 0), Details::kInterleave); + using ScalesAccessT = ScalesAccess; + GMemIterator + scales_iterator( + scales, + (GroupSize != 0 ? real_offset_k / GroupSize * n : 0) + real_offset_n, + (GroupSize != 0 ? CtaK / Details::kInterleave / GroupSize * n : 0), Details::kInterleave); out += offset_m * n + tile_id_n * CtaN * Details::kInterleave; if constexpr (EnableBias) { bias += tile_id_n * CtaN * Details::kInterleave; } - AccT tile_acc[CtaM * CtaN]; - fill(tile_acc, static_cast(0.f)); + using Converter = typename ConverterWrapper
::Converter; + using TileAccT = TileAccType; + using Math = MathWrapper; + + TileAccT tile_k_acc[CtaN]; + if constexpr (std::is_same_v) { + fill(tile_k_acc, 0.f); + } else { + fill(tile_k_acc, Math::to_vec2(static_cast(0.f))); + } - TypeA vec_scale[CtaN]; + // load_scales() writes through a ScalesAccessT::TVec* (float4 when vectorized), and the + // iterators/converters below write tile_a through AccessTypeA* and tile_w through uint32_t*, + // so these arrays need the alignment of the widest access, not just of TypeA. + alignas(alignof(typename ScalesAccessT::TVec)) TypeA vec_scale[CtaN]; if constexpr (GroupSize == 0) { -#pragma unroll - for (int i = 0; i < CtaN; ++i) { - scales_iterator.load(vec_scale + i, 0, i); - } + load_scales(scales_iterator, vec_scale, 0); } for (int idx_k = tid * StepK, iter = 0; idx_k < interleaved_k; idx_k += CtaK, ++iter) { - TypeA tile_a[StepK], tile_w[StepK], tile_w_pack2[CtaN * StepK]; - uint8_t tile_w_quantized[StepK / Details::kElemsPerByteW]; + alignas(alignof(AccessTypeA)) TypeA tile_a[StepK]; + // Issue all CtaN weight loads before consuming any of them: interleaving a load with its own + // decode leaves a single load in flight and makes the kernel long-scoreboard bound. + AccessTypeW tile_w_quantized[CtaN * Details::kAccessNumW]; if constexpr (GroupSize != 0) { -#pragma unroll - for (int i = 0; i < CtaN; ++i) { - scales_iterator.load(vec_scale + i, iter, i); - } + load_scales(scales_iterator, vec_scale, iter); } + act_iterator.load(tile_a, iter, 0); #pragma unroll for (int i = 0; i < CtaN; ++i) { - weight_iterator.load(tile_w_quantized, iter, i); - dequantize(tile_w, tile_w_quantized, vec_scale + i, nullptr, 1.0f); - pack_to_vec2(tile_w_pack2, tile_w, i); + weight_iterator.load(tile_w_quantized + i * Details::kAccessNumW, iter, i); } #pragma unroll - for (int i = 0; i < CtaM; ++i) { - act_iterator.load(tile_a, iter, i); - mma(tile_acc + i * CtaN, tile_w_pack2, tile_a); + for (int i = 0; i < CtaN; ++i) { + alignas(alignof(uint32_t)) TypeA tile_w[StepK]; + Converter::template convert(tile_w_quantized + i * Details::kAccessNumW, tile_w); + accumulate_column_tile(tile_k_acc[i], tile_w, tile_a, vec_scale[i]); } } + + AccT tile_acc[CtaM * CtaN]; + collapse_tile_acc(tile_acc, tile_k_acc); epilogue(out, n, tile_acc, bias, 1.0f); #endif } @@ -234,48 +315,57 @@ __global__ void moe_gemv_interleaved_swiglu_kernel( GMemIterator weight_iterator( weight, (interleaved_offset_n * interleaved_k + tid * StepK) / Details::kElemsPerByteW, CtaK / Details::kElemsPerByteW, interleaved_k / Details::kElemsPerByteW); - GMemIterator scales_iterator( - scales, - (GroupSize != 0 ? real_offset_k / GroupSize * n : 0) + real_offset_n, - (GroupSize != 0 ? CtaK / Details::kInterleave / GroupSize * n : 0), Details::kInterleave); + using ScalesAccessT = ScalesAccess; + GMemIterator + scales_iterator( + scales, + (GroupSize != 0 ? real_offset_k / GroupSize * n : 0) + real_offset_n, + (GroupSize != 0 ? CtaK / Details::kInterleave / GroupSize * n : 0), Details::kInterleave); out += offset_m * inter_size + tile_id_n * CtaN * Details::kInterleave / 2; if constexpr (EnableBias) { bias += tile_id_n * CtaN * Details::kInterleave; } - AccT tile_acc[CtaM * CtaN]; - fill(tile_acc, static_cast(0.f)); + using Converter = typename ConverterWrapper
::Converter; + using TileAccT = TileAccType; + using Math = MathWrapper; + + TileAccT tile_k_acc[CtaN]; + if constexpr (std::is_same_v) { + fill(tile_k_acc, 0.f); + } else { + fill(tile_k_acc, Math::to_vec2(static_cast(0.f))); + } - TypeA vec_scale[CtaN]; + // See moe_gemv_kernel: these are written through wider pointer casts than TypeA. + alignas(alignof(typename ScalesAccessT::TVec)) TypeA vec_scale[CtaN]; if constexpr (GroupSize == 0) { -#pragma unroll - for (int i = 0; i < CtaN; ++i) { - scales_iterator.load(vec_scale + i, 0, i); - } + load_scales(scales_iterator, vec_scale, 0); } for (int idx_k = tid * StepK, iter = 0; idx_k < interleaved_k; idx_k += CtaK, ++iter) { - TypeA tile_a[StepK], tile_w[StepK], tile_w_pack2[CtaN * StepK]; - uint8_t tile_w_quantized[StepK / Details::kElemsPerByteW]; + alignas(alignof(AccessTypeA)) TypeA tile_a[StepK]; + // See moe_gemv_kernel: keep all CtaN weight loads in flight at once. + AccessTypeW tile_w_quantized[CtaN * Details::kAccessNumW]; if constexpr (GroupSize != 0) { -#pragma unroll - for (int i = 0; i < CtaN; ++i) { - scales_iterator.load(vec_scale + i, iter, i); - } + load_scales(scales_iterator, vec_scale, iter); } + act_iterator.load(tile_a, iter, 0); #pragma unroll for (int i = 0; i < CtaN; ++i) { - weight_iterator.load(tile_w_quantized, iter, i); - dequantize(tile_w, tile_w_quantized, vec_scale + i, nullptr, 1.0f); - pack_to_vec2(tile_w_pack2, tile_w, i); + weight_iterator.load(tile_w_quantized + i * Details::kAccessNumW, iter, i); } #pragma unroll - for (int i = 0; i < CtaM; ++i) { - act_iterator.load(tile_a, iter, i); - mma(tile_acc + i * CtaN, tile_w_pack2, tile_a); + for (int i = 0; i < CtaN; ++i) { + alignas(alignof(uint32_t)) TypeA tile_w[StepK]; + Converter::template convert(tile_w_quantized + i * Details::kAccessNumW, tile_w); + accumulate_column_tile(tile_k_acc[i], tile_w, tile_a, vec_scale[i]); } } + + AccT tile_acc[CtaM * CtaN]; + collapse_tile_acc(tile_acc, tile_k_acc); swiglu_epilogue(out, tile_acc, bias, activation_params); #endif } diff --git a/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv_fp4.cu b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv_fp4.cu index 1bbe9a8bbb461..5901df5f059d2 100644 --- a/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv_fp4.cu +++ b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv_fp4.cu @@ -10,6 +10,7 @@ #include "contrib_ops/cuda/llm/moe_gemm/moe_gemv.h" // Shared device-side kernels + launch/dispatch helpers (fpA_intB_gemv namespace). #include "contrib_ops/cuda/llm/moe_gemm/moe_gemv_device.cuh" +#include "contrib_ops/cuda/llm/common/cuda_runtime_utils.h" #include "core/platform/env_var_utils.h" namespace onnxruntime::llm { @@ -54,6 +55,50 @@ bool Fp4MoeGemvInterleavedHalfAccum() { static constexpr int kInterleavedCtaN = 4; static constexpr int kInterleavedThreads = 128; +// Opt-out for the shape-derived default tiling (env ORT_FP4_GEMV_DEFAULT_TILING=0), which +// restores the fixed kDefaultCtaN/kDefaultThreads tiling for every shape. +static bool Fp4MoeGemvUseDefaultTilingHeuristic() { + static bool const enabled = + onnxruntime::ParseEnvironmentVariableWithDefault("ORT_FP4_GEMV_DEFAULT_TILING", 1) == 1; + return enabled; +} + +// Blocks per SM required before a 64-thread block is worth it. The grid is fixed by +// (expanded_num_rows, n / CtaN) and does not depend on Threads, so halving the block size halves +// the threads each block contributes to residency. These kernels are register-limited (about 72 +// registers/thread on sm_90, so roughly 900 resident threads/SM), which a 128-thread block reaches +// with ~7 resident blocks and a 64-thread block only with ~14. Requiring 16 blocks/SM therefore +// keeps the SM just as full while halving the epilogue's reduction width. +static constexpr int64_t kMinBlocksPerSmForThreads64 = 16; + +MoeGemvConfig Fp4MoeGemvDefaultConfig(int64_t expanded_num_rows, int64_t n, int64_t k, + int multi_processor_count) { + // The interleaved path pins its own CtaN/Threads and ignores `config`. + if (Fp4MoeGemvUseInterleaved() || !Fp4MoeGemvUseDefaultTilingHeuristic()) { + return MoeGemvConfig::kDefault; + } + // Each block walks K in strides of CtaK = StepK * Threads, with StepK = 128 / activation_bits. + constexpr int64_t kStepK = 128 / 16; + constexpr int64_t kDefaultCtaK = kStepK * kDefaultThreads; + + // (a) Idle threads. When CtaK > k the tail of every block never enters the K loop at all, so a + // 128-thread block leaves half its threads doing nothing (NVFP4 fc2 has k = inter_size, + // e.g. 512 against CtaK = 1024). Narrowing the block is a pure win here. + if (kDefaultCtaK > k) { + return MoeGemvConfig::kThreads64; + } + + // (b) Epilogue cost. The MAC work per block is fixed by (CtaN, k) regardless of Threads, but the + // epilogue reduces partial sums across Threads/32 warps through shared memory. A narrower + // block does the same math with half the barriers and a shallower reduction tree, so prefer + // it whenever the grid is large enough that the SM still fills up (see the constant above). + const int64_t blocks = expanded_num_rows * (n / kDefaultCtaN); + if (blocks >= kMinBlocksPerSmForThreads64 * multi_processor_count) { + return MoeGemvConfig::kThreads64; + } + return MoeGemvConfig::kDefault; +} + // --- MXFP4 (e2m1) GEMV: non-interleaved ColumnMajor layout (kInterleave = 1) --- // Weights are the QMoERepackFP4ColToRow output ([experts, n, k/2] row-major, two e2m1 // codes per byte, even-K in the low nibble). Block scales are the @@ -131,7 +176,7 @@ bool is_moe_gemv_fp4_supported(int sm, int64_t expanded_num_rows, int64_t n, int if (k % group_size != 0) { return false; } - if (expanded_num_rows <= 0 || expanded_num_rows > kMaxProfiledExpandedRows) { + if (expanded_num_rows <= 0 || expanded_num_rows > kMaxProfiledExpandedRowsFp4) { return false; } if (n < kMinProfiledProblemDim || k < kMinProfiledProblemDim) { @@ -228,7 +273,8 @@ void launch_moe_gemv_fp4_symmetric(const T* act, const uint8_t* weight, const T* // AccT follows the Fp4GemvAccT policy (fp16->fp16 accum, bf16->fp32 accum): bf16 has only 7 // mantissa bits, so 16-bit accumulation over K loses too much precision and fails tolerance // (e.g. NVFP4 block-16 decode at k=512). CtaN/Threads remain pure parallelization/tiling knobs - // and the accumulation dtype is identical for every config, so this sweep stays bit-exact. + // and the accumulation dtype is identical for every config, so every config computes the same + // dot products; Threads additionally sets the K partition, so it perturbs the summation order. auto launch = [&](auto cta_n, auto threads) { fiv::dispatch_moe_gemv_group_size>( const_cast(act), const_cast(weight), const_cast(scales), const_cast(bias), out, @@ -284,7 +330,8 @@ void launch_moe_gemv_fp4_symmetric_interleaved_swiglu( } using Details = Fp4KernelDetails; // AccT follows the Fp4GemvAccT policy (fp16->fp16, bf16->fp32); see launch_moe_gemv_fp4_symmetric. - // The CtaN/Threads sweep stays bit-exact across configs since the accumulation dtype is fixed. + // The accumulation dtype is fixed across the CtaN/Threads sweep; Threads still changes the K + // partition, so it is a tiling knob, not a bit-exact one. auto launch = [&](auto cta_n, auto threads) { fiv::dispatch_moe_gemv_interleaved_swiglu_group_size>( const_cast(act), const_cast(weight), const_cast(scales), const_cast(bias), out, diff --git a/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv_fp4.h b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv_fp4.h index 7239bd94f94f5..ea6bc68903035 100644 --- a/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv_fp4.h +++ b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv_fp4.h @@ -16,14 +16,21 @@ namespace onnxruntime::llm { namespace kernels { namespace moe_gemv { -// Tiling/parallelization knob selected by the FP4 GEMV autotuner. CtaN/Threads are pure -// tiling knobs (numerically bit-exact), so the sweep only picks the fastest. +// Tiling/parallelization knob selected by the FP4 GEMV autotuner. Every config computes the same +// dot products with the same accumulation dtype, so the sweep only picks the fastest. CtaN is +// bit-exact (it only changes how many output columns a block owns). Threads is *not*: a block +// walks K in strides of StepK * Threads and the epilogue reduces across Threads/32 warps, so +// changing it changes the summation order and the low bits of the result can move. enum class MoeGemvConfig { kDefault, kCtaN16, kThreads64, }; +// Cover Qwen-style top_k=8 MTP decode. An (N+1)-token verification for +// num_speculative_tokens=N expands to (N+1)*8 rows, up to 64 for N=7. +inline constexpr int64_t kMaxProfiledExpandedRowsFp4 = 64; + // True when the opt-in interleaved MXFP4 GEMV path is enabled (env ORT_FP4_GEMV_INTERLEAVED=1). // It combines three changes over the default path: (a) the INT4-style ColumnMajorInterleaved FP4 // weight layout (kInterleave=4, kStepK=32) for 4x fewer K-trips, (b) dtype-conditional accumulation @@ -33,6 +40,16 @@ enum class MoeGemvConfig { // compute dispatch query this so the prepacked weights and the kernel always agree. bool Fp4MoeGemvUseInterleaved(); +// Shape-derived default tiling for the non-interleaved ColumnMajor FP4 GEMV. Used whenever the +// runtime does not have a profiled result for the shape, which is the shipping default because +// ORT_FP4_GEMV_AUTOTUNE is off (it synchronizes the inference stream) and is skipped entirely +// during CUDA-graph capture. Only Threads is derived; CtaN stays at the default, so the result +// never changes which shapes is_moe_gemv_fp4_supported accepts. Note that Threads sets the K +// partition, so the choice moves the last bits of the output (see MoeGemvConfig above). Set +// ORT_FP4_GEMV_DEFAULT_TILING=0 to fall back to the fixed default tiling. +MoeGemvConfig Fp4MoeGemvDefaultConfig(int64_t expanded_num_rows, int64_t n, int64_t k, + int multi_processor_count); + // FP4 GEMV shape support for the non-interleaved ColumnMajor layout (kInterleave = 1). Shared by // both MXFP4 (group_size == 32) and NVFP4 (group_size == 16). Requires sm >= 80, n divisible by // the kernel tile width (kCtaN) selected by `config`, and the profiled small-decode row/dim diff --git a/onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc b/onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc index 4fda47aafb6e7..d97caeed8e96f 100644 --- a/onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc +++ b/onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc @@ -1432,13 +1432,19 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const { using MoeGemvConfig = gemv::MoeGemvConfig; - // Choose the fc1 (SwiGLU) and fc2 GEMV tiling configs. CtaN/Threads are pure tiling - // knobs (numerically bit-exact), so the only goal is picking the fastest. Reuse a - // cached per-shape result when available; otherwise profile on a non-captured (warmup) - // call and freeze the choice for CUDA-graph replay. During capture (or when autotune is - // off) fall back to the default tiling. - MoeGemvConfig fc1_config = MoeGemvConfig::kDefault; - MoeGemvConfig fc2_config = MoeGemvConfig::kDefault; + // Choose the fc1 (SwiGLU) and fc2 GEMV tiling configs. Every config computes the same + // dot products with the same accumulation dtype, so the only goal is picking the + // fastest; Threads does set the K partition, so the low bits of the output can move + // between configs (see MoeGemvConfig). Reuse a cached per-shape result when available; + // otherwise start from the shape-derived analytic default and, when autotune is on, + // profile on a non-captured (warmup) call and freeze the choice for CUDA-graph replay. + // During capture (or when autotune is off, which is the shipping default) the analytic + // default is what actually runs. + const int multi_processor_count = GetDeviceProp().multiProcessorCount; + MoeGemvConfig fc1_config = + gemv::Fp4MoeGemvDefaultConfig(expanded, fc1_n, hidden, multi_processor_count); + MoeGemvConfig fc2_config = + gemv::Fp4MoeGemvDefaultConfig(expanded, hidden, inter, multi_processor_count); const int64_t row_bucket = onnxruntime::llm::kernels::cutlass_kernels::MoeGemmProfiler::bucketM(expanded); diff --git a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_wide_tile.wgsl.template b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_wide_tile.wgsl.template index c8ddc9f2590d1..3be1efa0f1268 100644 --- a/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_wide_tile.wgsl.template +++ b/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_wide_tile.wgsl.template @@ -192,9 +192,7 @@ $MAIN { let row = ((workgroup_idx / uniforms.num_N_tile) % uniforms.num_M_tile) * kTileM; let col = (workgroup_idx % uniforms.num_N_tile) * kTileN; - // Utilizing an f32 accumulator mitigated precision loss with minimal - // performance impact compared to an f16 accumulator. - var results : array; + var results : array; for (var block_idx = 0u; block_idx < uniforms.n_blocks_per_col; block_idx++) { // Load `a` elements into workgroup memory, TileM x KAVecSizeForBlock32 (block32) let a_row_idx = local_idx / KAVecSizeForBlock32; @@ -215,8 +213,8 @@ $MAIN { let a_data0 = a_data_tile[m_idx][b_idx * 2u]; let a_data1 = a_data_tile[m_idx][b_idx * 2u + 1u]; - results[m_idx] += f32(dot(a_data0, b_dequantized[0])) + - f32(dot(a_data1, b_dequantized[1])); + results[m_idx] += dot(a_data0, b_dequantized[0]) + + dot(a_data1, b_dequantized[1]); } } workgroupBarrier(); @@ -238,9 +236,9 @@ $MAIN { #endif for (var m_idx = 0u; m_idx < kTileM; m_idx++) { #if has_bias - write_output(batch, row + m_idx, col + local_idx, output_element_t(results[m_idx]) + bias_value); + write_output(batch, row + m_idx, col + local_idx, results[m_idx] + bias_value); #else - write_output(batch, row + m_idx, col + local_idx, output_element_t(results[m_idx])); + write_output(batch, row + m_idx, col + local_idx, results[m_idx]); #endif } } // MAIN diff --git a/onnxruntime/core/graph/contrib_ops/bert_defs.cc b/onnxruntime/core/graph/contrib_ops/bert_defs.cc index 2b654c4c21877..bbf8caf06faf1 100644 --- a/onnxruntime/core/graph/contrib_ops/bert_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/bert_defs.cc @@ -1423,13 +1423,27 @@ cumulative_sequence_length records cumulated length of each sequence length. // Input 'block_table': (batch_size, max_blocks_per_sequence) // Input 'cos_cache': (max_seq_len, head_size / 2) // Input 'sin_cache': (max_seq_len, head_size / 2) -// Output 'output': (token_count, hidden_size) +// Input 'slot_mapping': (token_count) +// Input 'head_sink': (num_heads) +// Input 'q_norm_weight': (head_size) +// Input 'k_norm_weight': (head_size) +// Input 'k_scale': (1) for PER_TENSOR, (kv_num_heads, 1, head_size) for PER_CHANNEL +// Input 'v_scale': (1) for PER_TENSOR, (kv_num_heads, 1, head_size) for PER_CHANNEL +// Input 'attention_metadata': (2), CPU memory: [max_query_len_bound, max_kv_len_bound] +// Output 'output': (token_count, num_heads * v_head_size) // Output 'key_cache_out': (num_blocks, block_size, kv_num_heads, head_size) -// Output 'value_cache_out': (num_blocks, block_size, kv_num_heads, head_size) +// Output 'value_cache_out': (num_blocks, block_size, kv_num_heads, head_size), absent for LATENT void PagedAttentionTypeAndShapeInference(ONNX_NAMESPACE::InferenceContext& ctx) { // Type inference ONNX_NAMESPACE::propagateElemTypeFromInputToOutput(ctx, 0, 0); + const std::string kv_cache_layout = getAttribute(ctx, "kv_cache_layout", "SEPARATE"); + if (kv_cache_layout != "SEPARATE" && kv_cache_layout != "LATENT") { + fail_shape_inference("kv_cache_layout must be 'SEPARATE' or 'LATENT'."); + } + const bool is_latent_kv = kv_cache_layout == "LATENT"; + const int64_t v_head_size_attr = getAttribute(ctx, "v_head_size", 0); + // Shape inference for output tensor if (hasInputShape(ctx, 0)) { auto& query_shape = getInputShape(ctx, 0); @@ -1439,7 +1453,37 @@ void PagedAttentionTypeAndShapeInference(ONNX_NAMESPACE::InferenceContext& ctx) fail_shape_inference("Input 0 (query) shall be 2 dimensions"); } - if (ctx.hasInput(2)) { + if (is_latent_kv) { + // Absorbed MLA: query is unpacked and (unlike SEPARATE mode) the output head width is + // v_head_size rather than head_size, so the output shape cannot simply be copied from query. + int64_t num_heads = getAttribute(ctx, "num_heads", 0); + int64_t q_hidden_size = query_dims[1].has_dim_value() ? query_dims[1].dim_value() : 0; + if (num_heads <= 0) { + fail_shape_inference("num_heads must be a positive integer."); + } + if (v_head_size_attr == 0) { + // V is as wide as the latent row, so the output width matches the query width. + propagateShapeFromInputToOutput(ctx, 0, 0); + } else if (q_hidden_size > 0) { + if (q_hidden_size % num_heads != 0) { + fail_shape_inference("Query hidden size must be divisible by num_heads."); + } + int64_t head_size = q_hidden_size / num_heads; + if (v_head_size_attr > head_size) { + fail_shape_inference("v_head_size must not exceed head_size."); + } + ONNX_NAMESPACE::TensorShapeProto output_shape; + *output_shape.add_dim() = query_dims[0]; + output_shape.add_dim()->set_dim_value(num_heads * v_head_size_attr); + updateOutputShape(ctx, 0, output_shape); + } else { + // Symbolic query hidden size: only the token dimension is known. + ONNX_NAMESPACE::TensorShapeProto output_shape; + *output_shape.add_dim() = query_dims[0]; + output_shape.add_dim(); + updateOutputShape(ctx, 0, output_shape); + } + } else if (ctx.hasInput(2)) { ONNX_NAMESPACE::TensorShapeProto output_shape; propagateShapeFromInputToOutput(ctx, 0, 0); } else { // packed QKV @@ -1461,12 +1505,23 @@ void PagedAttentionTypeAndShapeInference(ONNX_NAMESPACE::InferenceContext& ctx) // Shape inference for KV Cache output tensors if (ctx.getNumOutputs() > 1) { // has kv cache output - if (ctx.getNumOutputs() != 3) { + if (is_latent_kv) { + // A single physical cache: there is no value cache to alias out. + if (ctx.getNumOutputs() > 2) { + fail_shape_inference("value_cache_out must be absent when kv_cache_layout is 'LATENT'."); + } + } else if (ctx.getNumOutputs() != 3) { fail_shape_inference("Key cache and value cache output tensors must be both present or both absent."); + } else if (!ctx.hasInput(4)) { + // value_cache is schema-optional (it must be absent for LATENT), so a SEPARATE node could omit it + // while still declaring value_cache_out. Fail with a clear message instead of reading input 4. + fail_shape_inference("value_cache is required when value_cache_out is present."); } - // types - ONNX_NAMESPACE::propagateElemTypeFromInputToOutput(ctx, 0, 1); - ONNX_NAMESPACE::propagateElemTypeFromInputToOutput(ctx, 0, 2); + // Types: the cache outputs alias the cache inputs, so their element type comes from inputs 3/4 + // (T_CACHE) rather than from the query (T) — the two differ for a quantized cache. This has to + // run before the shape propagation below, which requires the output TypeProto to already be a + // tensor type. + ONNX_NAMESPACE::propagateElemTypeFromInputToOutput(ctx, 3, 1); // shapes auto& key_cache_shape = getInputShape(ctx, 3); auto& key_cache_dims = key_cache_shape.dim(); @@ -1475,9 +1530,11 @@ void PagedAttentionTypeAndShapeInference(ONNX_NAMESPACE::InferenceContext& ctx) } // KV cache in and out share the same buffer, thus they have the same shape ONNX_NAMESPACE::propagateShapeFromInputToOutput(ctx, 3, 1); - ONNX_NAMESPACE::propagateShapeFromInputToOutput(ctx, 4, 2); - ONNX_NAMESPACE::propagateElemTypeFromInputToOutput(ctx, 3, 1); - ONNX_NAMESPACE::propagateElemTypeFromInputToOutput(ctx, 4, 2); + + if (ctx.getNumOutputs() > 2) { + ONNX_NAMESPACE::propagateElemTypeFromInputToOutput(ctx, 4, 2); + ONNX_NAMESPACE::propagateShapeFromInputToOutput(ctx, 4, 2); + } } } @@ -1507,6 +1564,61 @@ ONNX_MS_OPERATOR_SET_SCHEMA( "Rotate using interleaved pattern. Default value is 0 (False).", AttributeProto::INT, OPTIONAL_VALUE) + .Attr("qk_norm_epsilon", + "Epsilon used by the Q/K RMSNorm when 'q_norm_weight' and 'k_norm_weight' are provided. " + "Default value is 1e-6.", + AttributeProto::FLOAT, + OPTIONAL_VALUE) + .Attr("k_quant_type", + "Quantization granularity of the key cache: 'NONE', 'PER_TENSOR' or 'PER_CHANNEL'. " + "Must be non-'NONE' exactly when 'key_cache' has a quantized element type, and then " + "'k_scale' is required. Default value is 'NONE'.", + AttributeProto::STRING, + std::string("NONE")) + .Attr("v_quant_type", + "Quantization granularity of the value cache: 'NONE', 'PER_TENSOR' or 'PER_CHANNEL'. " + "Must be non-'NONE' exactly when 'value_cache' has a quantized element type, and then " + "'v_scale' is required. Default value is 'NONE'.", + AttributeProto::STRING, + std::string("NONE")) + .Attr("k_cache_dtype", + "Logical element type stored in 'key_cache', named after the ONNX element type it denotes: '' " + "(the default) means the cache tensor's own element type is also the logical type. 'float16', " + "'bfloat16', 'int8' and 'float8e4m3fn' name that same type explicitly and must agree with the " + "tensor. 'int4' and 'float4e2m1' name sub-byte types packed two per byte into a uint8 cache, " + "where the last cache dimension holds (head_size + 1) / 2 bytes and logical element 2*i " + "occupies the low-order bits of byte i. Every value is a signed, zero-symmetric type: " + "quantization uses a scale with no zero point, so unsigned logical types are not expressible.", + AttributeProto::STRING, + std::string("")) + .Attr("v_cache_dtype", + "Logical element type stored in 'value_cache', with the same values and packing rule as " + "'k_cache_dtype'. Default value is '' (use the cache tensor's element type).", + AttributeProto::STRING, + std::string("")) + .Attr("kv_cache_layout", + "Physical layout of the KV cache: 'SEPARATE' or 'LATENT'. 'SEPARATE' (the default) uses " + "distinct 'key_cache' and 'value_cache' tensors. 'LATENT' selects absorbed Multi-head Latent " + "Attention: there is a single cache, 'value' and 'value_cache' must be absent, 'kv_num_heads' " + "must be 1, and V for every head is the leading 'v_head_size' channels of the same 'key_cache' " + "row that supplies K. Default value is 'SEPARATE'.", + AttributeProto::STRING, + std::string("SEPARATE")) + .Attr("v_head_size", + "Width of the value head, which may be narrower than head_size. Only valid when " + "'kv_cache_layout' is 'LATENT' (DeepSeek-V3 uses head_size=576 and v_head_size=512). When " + "v_head_size differs from head_size the 'scale' attribute is required, because the " + "1/sqrt(head_size) default no longer matches the pre-absorption head width. Default value is 0, " + "meaning the same as head_size.", + AttributeProto::INT, + OPTIONAL_VALUE) + .Attr("rotary_offset", + "First channel within head_size covered by rotary embedding, so RoPE is applied to " + "[rotary_offset, rotary_offset + rotary_dim) and channels outside that range are copied " + "through. Must be a multiple of 8. MLA sets this to kv_lora_rank so that RoPE only touches the " + "positional suffix of the latent row. Default value is 0.", + AttributeProto::INT, + OPTIONAL_VALUE) .Input(0, "query", "Query with shape (num_tokens, hidden_size), or packed QKV with shape (num_tokens, d) " @@ -1519,19 +1631,22 @@ ONNX_MS_OPERATOR_SET_SCHEMA( OpSchema::Optional) .Input(2, "value", - "Value with shape (num_tokens, kv_hidden_size)", + "Value with shape (num_tokens, kv_hidden_size). Must be absent when 'kv_cache_layout' is 'LATENT'.", "T", OpSchema::Optional) .Input(3, "key_cache", "Block-based key cache with shape (num_blocks, block_size, kv_num_heads, head_size). This is updated in " - "place within the op.", - "T") + "place within the op. When 'kv_cache_layout' is 'LATENT' this is the only cache, and V is read from its " + "leading v_head_size channels.", + "T_CACHE") .Input(4, "value_cache", "Block-based value cache with shape (num_blocks, block_size, kv_num_heads, head_size). This is updated " - "in place within the op. This should be the same shape as key_cache.", - "T") + "in place within the op. This should be the same shape as key_cache. Must be absent when " + "'kv_cache_layout' is 'LATENT'.", + "T_CACHE", + OpSchema::Optional) .Input(5, "cumulative_sequence_length", "A tensor with shape (batch_size + 1). It specifies the cumulative sequence lengths between the packed " @@ -1556,23 +1671,84 @@ ONNX_MS_OPERATOR_SET_SCHEMA( "2D tensor with shape (max total seqlen, head_size / 2).", "T", OpSchema::Optional) + .Input(10, + "slot_mapping", + "1D tensor with shape (num_tokens). For each query token, the flat slot index " + "(block_id * block_size + offset_in_block) at which its key/value is written into the KV cache. " + "A value of -1 skips the cache write for that token, which lets a scheduler suppress stores for " + "prefix-cache hits or rejected speculative tokens. When absent, slots are derived from " + "'past_seqlens', 'cumulative_sequence_length' and 'block_table' as before. 'block_table' is still " + "required, because it defines the read path.", + "S", + OpSchema::Optional) + .Input(11, + "head_sink", + "1D tensor with shape (num_heads). Each head has a learnable sink logit that participates in the " + "softmax denominator but contributes no value, so attention can 'do nothing'.", + "T", + OpSchema::Optional) + .Input(12, + "q_norm_weight", + "1D tensor with shape (head_size). RMSNorm gain applied to each query head before rotary " + "embedding. Must be provided together with 'k_norm_weight'.", + "T", + OpSchema::Optional) + .Input(13, + "k_norm_weight", + "1D tensor with shape (head_size). RMSNorm gain applied to each key head before rotary embedding " + "and before the key is written to the KV cache. Must be provided together with 'q_norm_weight'.", + "T", + OpSchema::Optional) + .Input(14, + "k_scale", + "Dequantization scale of the key cache. Shape is (1) when 'k_quant_type' is 'PER_TENSOR' and " + "(kv_num_heads, 1, head_size) when it is 'PER_CHANNEL'. Quantization is symmetric (no zero point).", + "T_KV_SCALE", + OpSchema::Optional) + .Input(15, + "v_scale", + "Dequantization scale of the value cache. Shape is (1) when 'v_quant_type' is 'PER_TENSOR' and " + "(kv_num_heads, 1, head_size) when it is 'PER_CHANNEL'. Quantization is symmetric (no zero point).", + "T_KV_SCALE", + OpSchema::Optional) + .Input(16, + "attention_metadata", + "1D tensor with shape (2) holding [max_query_len_bound, max_kv_len_bound] in CPU memory. " + "max_query_len_bound is an upper bound on the number of new tokens any one sequence " + "contributes; max_kv_len_bound is an upper bound on past_seqlens[i] + query_len[i]. Both are " + "replay-wide upper bounds, never exact per-step values: they must hold for every step this node " + "-- or a CUDA Graph capturing it -- will serve, and 0 means 'unknown'. They may only select the " + "backend and size launch dimensions and workspaces; they never enter a mask comparison, so " + "over-estimating only costs empty work. The op can otherwise obtain these only by copying " + "'cumulative_sequence_length' and 'past_seqlens' back from the device and synchronizing the " + "stream on every call, which stalls the pipeline once per node per step and makes the op " + "impossible to capture into a CUDA Graph. Schedulers already track these bounds on the host, so " + "supplying them is normally free. When absent, the op falls back to the device readback. " + "The values are trusted: an under-sized bound violates the contract and may omit attention work.", + "S", + OpSchema::Optional) .Output(0, "output", - "3D output tensor with shape (num_tokens, hidden_size)", + "2D output tensor with shape (num_tokens, num_heads * v_head_size), which is " + "(num_tokens, hidden_size) unless 'kv_cache_layout' is 'LATENT' with a narrower v_head_size.", "T") .Output(1, "key_cache_out", "Block-based key cache with shape (num_blocks, block_size, kv_num_heads, head_size). This is always " "the same tensor as key_cache.", - "T", + "T_CACHE", OpSchema::Optional) .Output(2, "value_cache_out", "Block-based value cache with shape (num_blocks, block_size, kv_num_heads, head_size). This is always " - "the same tensor as value_cache.", - "T", + "the same tensor as value_cache. Must be absent when 'kv_cache_layout' is 'LATENT'.", + "T_CACHE", OpSchema::Optional) .TypeConstraint("T", {"tensor(float16)", "tensor(bfloat16)"}, "Constrain input and output to float tensors.") + .TypeConstraint("T_CACHE", + {"tensor(float16)", "tensor(bfloat16)", "tensor(int8)", "tensor(float8e4m3fn)"}, + "Constrain the KV cache to float or quantized tensors.") + .TypeConstraint("T_KV_SCALE", {"tensor(float)"}, "Constrain KV cache scales to float tensors.") .TypeConstraint("S", {"tensor(int32)"}, "Constrain Positional inputs to int tensor.") .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { PagedAttentionTypeAndShapeInference(ctx); diff --git a/onnxruntime/core/providers/cuda/nn/layer_norm.cc b/onnxruntime/core/providers/cuda/nn/layer_norm.cc index d479261855e2d..315a1f6c48bd9 100644 --- a/onnxruntime/core/providers/cuda/nn/layer_norm.cc +++ b/onnxruntime/core/providers/cuda/nn/layer_norm.cc @@ -6,6 +6,7 @@ #include "core/providers/cuda/nn/layer_norm_impl.h" #include "core/providers/cpu/nn/layer_norm_helper.h" #include "core/providers/cuda/cuda_common.h" +#include namespace onnxruntime { namespace cuda { @@ -88,6 +89,18 @@ Status LayerNorm::ComputeInternal(OpKernelContext* ctx) con return Status::OK(); } + // Validate that norm_size won't cause integer overflow in host-side arithmetic. + // HostApplyLayerNorm computes: (n2 + 4 * warp_size - 1) where warp_size is typically 32. + // Maximum safe value: INT_MAX - (4 * max_warp_size) to prevent overflow. + // Using 256 as a conservative upper bound for 4 * warp_size. + constexpr int MAX_WARP_FACTOR = 256; + const int MAX_NORM_SIZE = std::numeric_limits::max() - MAX_WARP_FACTOR; + + ORT_RETURN_IF(params.num_rows > 0 && + (params.norm_size > MAX_NORM_SIZE || + params.norm_size > std::numeric_limits::max() / params.num_rows), + "LayerNormalization input is too large for CUDA kernel indexing: norm_size exceeds safe limits or num_rows * norm_size exceeds INT_MAX"); + HostApplyLayerNorm( GetDeviceProp(), Stream(ctx), Y_data, mean_data, inv_var_data, X_data, onnxruntime::narrow(params.num_rows), onnxruntime::narrow(params.norm_size), epsilon_, diff --git a/onnxruntime/core/providers/cuda/nn/rms_norm.cc b/onnxruntime/core/providers/cuda/nn/rms_norm.cc index 8db7cac1687bd..1f0eb1d60ed08 100644 --- a/onnxruntime/core/providers/cuda/nn/rms_norm.cc +++ b/onnxruntime/core/providers/cuda/nn/rms_norm.cc @@ -6,6 +6,7 @@ #include "core/providers/cuda/nn/layer_norm_impl.h" #include "core/providers/cuda/cuda_common.h" #include "core/providers/cpu/nn/layer_norm_helper.h" +#include #include namespace onnxruntime { @@ -76,6 +77,18 @@ Status RMSNorm::ComputeInternal(OpKernelContext* ctx) const { return Status::OK(); } + // Validate that norm_size won't cause integer overflow in host-side arithmetic. + // HostApplyLayerNorm computes: (n2 + 4 * warp_size - 1) where warp_size is typically 32. + // Maximum safe value: INT_MAX - (4 * max_warp_size) to prevent overflow. + // Using 256 as a conservative upper bound for 4 * warp_size. + constexpr int MAX_WARP_FACTOR = 256; + const int MAX_NORM_SIZE = std::numeric_limits::max() - MAX_WARP_FACTOR; + + ORT_RETURN_IF(params.num_rows > 0 && + (params.norm_size > MAX_NORM_SIZE || + params.norm_size > std::numeric_limits::max() / params.num_rows), + "RMSNormalization input is too large for CUDA kernel indexing: norm_size exceeds safe limits or num_rows * norm_size exceeds INT_MAX"); + // For RMSNorm, we don't need mean and inv_var data, so we can pass nullptr. CudaU* mean_data = nullptr; CudaU* inv_var_data = nullptr; diff --git a/onnxruntime/core/providers/js/README.md b/onnxruntime/core/providers/js/README.md new file mode 100644 index 0000000000000..dc4f132fdee6a --- /dev/null +++ b/onnxruntime/core/providers/js/README.md @@ -0,0 +1,9 @@ +# JS execution provider (JSEP) — deprecated + +This directory is the native half of **JSEP**, the JavaScript/TypeScript WebGPU path in `onnxruntime-web`. It is +**deprecated** and will be removed. The replacement is the native WebGPU execution provider in +[`onnxruntime/core/providers/webgpu/`](../webgpu). + +**Bug fixes and security fixes only.** New operators, new features and performance work belong in the WebGPU EP. + +See [docs/JSEP_Deprecation.md](../../../../docs/JSEP_Deprecation.md) for more details. diff --git a/onnxruntime/core/providers/webgpu/math/matmul.h b/onnxruntime/core/providers/webgpu/math/matmul.h index 85a4a45a1f79a..7628f60a770d5 100644 --- a/onnxruntime/core/providers/webgpu/math/matmul.h +++ b/onnxruntime/core/providers/webgpu/math/matmul.h @@ -50,10 +50,19 @@ class MatMul final : public WebGpuKernel { const MatMul& parent_; }; - MatMul(const OpKernelInfo& info) : WebGpuKernel{info} {} + MatMul(const OpKernelInfo& info) : WebGpuKernel{info} { + // Whether the B (weight) input is a constant initializer. The subgroup-matrix + // opt impl uses this to decide it can safely pad B once and cache the result + // (odd-N handling); a non-constant B changes per run and must not be cached. + const Tensor* b = nullptr; + b_is_constant_ = info.TryGetConstantInput(1, &b); + } Status ComputeInternal(ComputeContext& context) const override; + // True when input 1 (B) is a constant initializer. See b_is_constant_. + bool IsBConstant() const { return b_is_constant_; } + constexpr static uint32_t MATMUL_PACKED_WORKGROUP_SIZE_X = 8; constexpr static uint32_t MATMUL_PACKED_WORKGROUP_SIZE_Y = 8; constexpr static uint32_t MATMUL_PACKED_WORKGROUP_SIZE_Z = 1; @@ -64,6 +73,8 @@ class MatMul final : public WebGpuKernel { // impl_ after initialization means this device has no optimized path. mutable std::unique_ptr impl_; mutable std::once_flag impl_init_flag_; + + bool b_is_constant_ = false; }; class MatMulNaiveProgram final : public Program { diff --git a/onnxruntime/core/providers/webgpu/math/subgroup_matrix_matmul.cc b/onnxruntime/core/providers/webgpu/math/subgroup_matrix_matmul.cc index 7fbe6ea28248d..ca632c8968245 100644 --- a/onnxruntime/core/providers/webgpu/math/subgroup_matrix_matmul.cc +++ b/onnxruntime/core/providers/webgpu/math/subgroup_matrix_matmul.cc @@ -6,7 +6,9 @@ #include "core/providers/webgpu/math/subgroup_matrix_matmul.h" #include +#include #include +#include #include #include #include @@ -16,6 +18,7 @@ #include "core/providers/webgpu/math/subgroup_matrix_config.h" #include "core/providers/webgpu/shader_helper.h" #include "core/providers/webgpu/vendor/intel/math/subgroup_matrix_tiling_selector.h" +#include "core/providers/webgpu/webgpu_utils.h" namespace onnxruntime { namespace webgpu { @@ -26,6 +29,24 @@ namespace { // TODO: use subgroup-size-control to enforce the subgroup size is 32. constexpr uint32_t kSubgroupMatrixSubgroupSize = 32; +// Copies a row-major f16 weight B [K, N] into a column-padded [K, N_b] buffer +// (N_b >= N), zero-filling columns [N, N_b). Gives B an even row stride so the +// subgroup-matrix f16 load's 4-byte row-start alignment holds for odd N. +class SubgroupMatrixMatMulPadBProgram final : public Program { + public: + SubgroupMatrixMatMulPadBProgram() : Program{"SubgroupMatrixMatMulPadB"} {} + Status GenerateShaderCode(ShaderHelper& shader) const override { + const auto& input_b = shader.AddInput("input_b", ShaderUsage::UseValueTypeAlias); + const auto& output = shader.AddOutput("output", ShaderUsage::UseValueTypeAlias); + return WGSL_TEMPLATE_APPLY(shader, "math/subgroup_matrix_matmul_pad_b.wgsl.template", + WGSL_TEMPLATE_VARIABLE(input_b, input_b), + WGSL_TEMPLATE_VARIABLE(output, output)); + } + WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES({"output_size", ProgramUniformVariableDataType::Uint32}, + {"N", ProgramUniformVariableDataType::Uint32}, + {"N_b", ProgramUniformVariableDataType::Uint32}); +}; + // Subgroup-matrix MatMul implementation. Loads both A and B directly from global // memory and runs the subgroup-matrix kernel during Compute. The class is // intended to support all subgroup-matrix configs; for now only 8x16x16 is @@ -104,11 +125,16 @@ class SubgroupMatrixMatMulImpl final : public MatMul::MatMulOptImpl { } // The B right-operand is loaded with subgroupMatrixLoad using a row stride of - // N (uniforms.N). Intel's f16 subgroup-matrix load reads columns in 32-bit - // (2xf16) pairs and requires each K-row to start 4-byte aligned, i.e. an even - // element stride. An odd N offsets every other K-row by 2 bytes and corrupts - // the odd output columns. Fall back to the generic MatMul path for odd N. - if (N % 2 != 0) { + // N_b. Intel's f16 subgroup-matrix load reads columns in 32-bit (2xf16) pairs + // and requires each K-row to start 4-byte aligned, i.e. an even element stride. + // An odd N would offset every other K-row by 2 bytes and corrupt the odd output + // columns. For a constant weight (2D or batched) we can pad B to an even stride + // (N_b = N + 1); a non-constant odd-N B falls back here. This is only the cheap + // eligibility check - the actual padding is deferred until after tiling selection + // so we never allocate/dispatch a padded copy for a problem that will fall back + // anyway (e.g. K % 16 != 0, which the tiling selector declines). B has already + // been validated as 2D or a well-formed batched shape above. + if (N % 2 != 0 && !parent_.IsBConstant()) { return Status::OK(); } @@ -117,6 +143,15 @@ class SubgroupMatrixMatMulImpl final : public MatMul::MatMulOptImpl { return Status::OK(); } + // The optimized path will run: now materialize the even-strided B for odd N. + const Tensor* b_used = b; + uint32_t N_b = N; + if (N % 2 != 0) { + ORT_RETURN_IF_ERROR(EnsurePaddedB(context, *b, N)); + b_used = padded_b_.get(); + N_b = padded_b_stride_; + } + TensorShapeVector output_dims{a_shape.GetDims().begin(), a_shape.GetDims().end()}; output_dims.back() = static_cast(N); TensorShape output_shape{output_dims}; @@ -142,9 +177,9 @@ class SubgroupMatrixMatMulImpl final : public MatMul::MatMulOptImpl { program.SetDispatchGroupSize(dispatch_x, dispatch_y, batch); program.CacheHint(has_bias, config_index_, sg_mat_count_m, sg_mat_count_n, split_k) .AddInputs({{a, ProgramTensorMetadataDependency::TypeAndRank, 1}, - {b, ProgramTensorMetadataDependency::TypeAndRank, 1}}) + {b_used, ProgramTensorMetadataDependency::TypeAndRank, 1}}) .AddOutput({output, ProgramTensorMetadataDependency::Rank, output->Shape(), 1}) - .AddUniformVariables({{M}, {N}, {K}, {dispatch_x}}); + .AddUniformVariables({{M}, {N}, {K}, {dispatch_x}, {N_b}}); if (has_bias) { program.AddInput({bias, ProgramTensorMetadataDependency::None}); } @@ -155,8 +190,61 @@ class SubgroupMatrixMatMulImpl final : public MatMul::MatMulOptImpl { } private: + // Lazily builds an even-strided copy of a constant weight B [..., K, N] with odd N + // by widening its last dim to N_b = N + 1 (zero-filling the extra column) and + // caches it, so the per-run pad cost is paid once. Works for a 2D weight [K, N] + // and a batched weight [batch, K, N] alike: the pad pass treats B as a flat + // [rows, N] -> [rows, N_b] copy over rows = numel / N (= K, or batch*K), which is + // exactly the even-stride layout the kernel indexes via N_b. Runs the GPU pad pass + // on first use under call_once; the cached tensor is held for the kernel's + // lifetime. Only valid when B is a constant initializer (checked by the caller) - + // a runtime B changes per run and must not be cached. + Status EnsurePaddedB(ComputeContext& context, const Tensor& b, uint32_t N) const { + ORT_RETURN_IF_NOT(N < std::numeric_limits::max(), + "Cannot pad odd-N B because N+1 exceeds uint32_t range."); + const uint32_t n_b = N + 1; + TensorShapeVector padded_dims{b.Shape().GetDims().begin(), b.Shape().GetDims().end()}; + padded_dims.back() = static_cast(n_b); + const TensorShape padded_shape{padded_dims}; + const int64_t output_size_i64 = padded_shape.Size(); + ORT_RETURN_IF_NOT(output_size_i64 <= static_cast(std::numeric_limits::max()), + "Cannot pad odd-N B because the padded tensor has ", output_size_i64, + " elements, exceeding uint32_t shader indexing range."); + const uint32_t output_size = narrow(output_size_i64); + + std::call_once(pad_once_, [&]() { + auto padded = std::make_unique(context.CreateGPUTensor(b.DataType(), padded_shape)); + Status s = Status::OK(); + // A zero-element padded tensor (e.g. a zero-batch or empty constant B) needs no + // pad pass - dispatching 0 workgroups is pointless and some drivers reject it. + // Just cache the empty tensor; the main kernel dispatches nothing for it. + if (output_size != 0) { + SubgroupMatrixMatMulPadBProgram program; + program.SetWorkgroupSize(WORKGROUP_SIZE) + .SetDispatchGroupSize(CeilDiv(output_size, WORKGROUP_SIZE)) + .AddInput({&b, ProgramTensorMetadataDependency::TypeAndRank, 1}) + .AddOutput({padded.get(), ProgramTensorMetadataDependency::TypeAndRank, padded->Shape(), 1}) + .AddUniformVariables({{output_size}, {N}, {n_b}}); + s = context.RunProgram(program); + } + if (s.IsOK()) { + padded_b_ = std::move(padded); + padded_b_stride_ = n_b; + } + }); + // padded_b_ persists the outcome across calls: call_once runs the body only on + // the first call, so a failed pad stays failed (and null) on later calls. + return padded_b_ ? Status::OK() + : ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Failed to pad odd-N B for subgroup-matrix MatMul."); + } + const int32_t config_index_; SubgroupMatrixTilingSelector tiling_selector_; + + // Cached even-strided B for odd N; built once by EnsurePaddedB. + mutable std::once_flag pad_once_; + mutable std::unique_ptr padded_b_; + mutable uint32_t padded_b_stride_ = 0; }; Status GenerateShaderCode8x16x16(ShaderHelper& shader, const ShaderVariableHelper& output, diff --git a/onnxruntime/core/providers/webgpu/math/subgroup_matrix_matmul.h b/onnxruntime/core/providers/webgpu/math/subgroup_matrix_matmul.h index 6fe3e9972d0aa..e6f00d9aaa4cd 100644 --- a/onnxruntime/core/providers/webgpu/math/subgroup_matrix_matmul.h +++ b/onnxruntime/core/providers/webgpu/math/subgroup_matrix_matmul.h @@ -58,10 +58,13 @@ class SubgroupMatrixMatMulProgram final : public Program), row-major with stride K. // B: weight, loaded directly from global memory in plain row-major KxN layout as a -// right operand (subgroup_matrix_right) with stride N. +// right operand (subgroup_matrix_right) with row stride +// N_b. N_b is the padded row stride of the (possibly column-padded) B buffer and +// equals N unless the host padded B to an even stride to satisfy the load's +// 4-byte row-start alignment (see the odd-N handling in subgroup_matrix_matmul.cc). // // Workgroup: split_k subgroups x 32 lanes (split_k in {1,2,4,8}). // Tile: kTileM x kTileN, chosen adaptively by the host config provider. Each @@ -20,9 +23,11 @@ // distributed round-robin across the subgroups, each accumulates its own partial // tile in shared memory, and the partials are summed at write-out. // -// Preconditions enforced by the host: K % sg_mat_k == 0. M and N may be any -// size; partial tiles are handled by bounds-checked stores. Out-of-range A/B -// loads return zero (WebGPU bounds checking). +// Preconditions enforced by the host: K % sg_mat_k == 0, and B's row stride N_b +// is even (the load requires 4-byte-aligned row starts). M and N may be any size; +// partial tiles are handled by bounds-checked stores. Out-of-range A/B loads +// return zero (WebGPU bounds checking). Output width/stride uses N (not N_b), so +// any padded B columns in [N, N_b) are read but never written. // // Batching: the host dispatches one num_n_tile x num_m_tile tile grid per batch // slice; the batch slice is recovered from the flattened workgroup_idx (the @@ -75,7 +80,7 @@ $MAIN { // Flat-element offsets into A/B/output for this batch slice, derived from // M/N/K. For a shared 2D weight batch_id is 0, so B collapses to its base. let a_batch_offset = batch_id * uniforms.M * uniforms.K; - let b_batch_offset = batch_id * uniforms.K * uniforms.N; + let b_batch_offset = batch_id * uniforms.K * uniforms.N_b; let out_batch_offset = batch_id * uniforms.M * uniforms.N; let k_blocks = uniforms.K / kSgMatK; let sg_index = local_idx / kSubgroupSize; // which split-K subgroup (0..kSplitK-1) @@ -179,25 +184,25 @@ $MAIN { #endif for (var kb: u32 = sg_index; kb < k_blocks; kb = kb + kSplitK) { - // Load the B right tiles for this K block (KxN row-major, stride N). - let b_base = b_batch_offset + kb * kSgMatK * uniforms.N + global_base_n; + // Load the B right tiles for this K block (KxN_b row-major, stride N_b). + let b_base = b_batch_offset + kb * kSgMatK * uniforms.N_b + global_base_n; var sg_mat_b0: subgroup_matrix_right = subgroupMatrixLoad>( - &input_b, b_base + 0 * kSgMatN, false, uniforms.N); + &input_b, b_base + 0 * kSgMatN, false, uniforms.N_b); #if sg_mat_count_n >= 2 var sg_mat_b1: subgroup_matrix_right = subgroupMatrixLoad>( - &input_b, b_base + 1 * kSgMatN, false, uniforms.N); + &input_b, b_base + 1 * kSgMatN, false, uniforms.N_b); #endif #if sg_mat_count_n >= 3 var sg_mat_b2: subgroup_matrix_right = subgroupMatrixLoad>( - &input_b, b_base + 2 * kSgMatN, false, uniforms.N); + &input_b, b_base + 2 * kSgMatN, false, uniforms.N_b); #endif #if sg_mat_count_n >= 4 var sg_mat_b3: subgroup_matrix_right = subgroupMatrixLoad>( - &input_b, b_base + 3 * kSgMatN, false, uniforms.N); + &input_b, b_base + 3 * kSgMatN, false, uniforms.N_b); #endif // Load the A left tiles (one per M block), row-major with stride K. diff --git a/onnxruntime/core/providers/webgpu/math/subgroup_matrix_matmul_pad_b.wgsl.template b/onnxruntime/core/providers/webgpu/math/subgroup_matrix_matmul_pad_b.wgsl.template new file mode 100644 index 0000000000000..f5d0d6996bced --- /dev/null +++ b/onnxruntime/core/providers/webgpu/math/subgroup_matrix_matmul_pad_b.wgsl.template @@ -0,0 +1,21 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +// Copies a row-major f16 weight B [K, N] into a column-padded [K, N_b] buffer +// (N_b >= N), zero-filling columns [N, N_b). One thread per padded-output element. +// Gives B an even row stride so the subgroup-matrix f16 load's 4-byte row-start +// alignment holds for odd N. See EnsurePaddedB in subgroup_matrix_matmul.cc. + +#use guardAgainstOutOfBoundsWorkgroupSizes +#use .getByOffset .setByOffset + +$MAIN { + guardAgainstOutOfBoundsWorkgroupSizes(uniforms.output_size); + let r = global_idx / uniforms.N_b; // padded-output row + let c = global_idx % uniforms.N_b; // padded-output column + var v = output_value_t(0); + if (c < uniforms.N) { // real column -> copy; else zero pad + v = output_value_t(input_b.getByOffset(r * uniforms.N + c)); + } + output.setByOffset(global_idx, v); +} // MAIN diff --git a/onnxruntime/python/tools/symbolic_shape_infer.py b/onnxruntime/python/tools/symbolic_shape_infer.py index 9fc33c2d0d054..1f40c0e729f40 100755 --- a/onnxruntime/python/tools/symbolic_shape_infer.py +++ b/onnxruntime/python/tools/symbolic_shape_infer.py @@ -2548,7 +2548,52 @@ def _infer_GroupNorm(self, node): # noqa: N802 self._propagate_shape_and_type(node) def _infer_PagedAttention(self, node): # noqa: N802 - self._propagate_shape_and_type(node) + # Output 0 is (token_count, num_heads * v_head_size). That equals the query shape except in + # two cases: packed QKV (query is wider than the output) and kv_cache_layout="LATENT" with a + # v_head_size narrower than head_size. + kv_cache_layout = get_attribute(node, "kv_cache_layout", b"SEPARATE") + if isinstance(kv_cache_layout, bytes): + kv_cache_layout = kv_cache_layout.decode() + is_latent_kv = kv_cache_layout == "LATENT" + is_packed_qkv = not is_latent_kv and (len(node.input) < 2 or not node.input[1]) + + if is_latent_kv or is_packed_qkv: + num_heads = get_attribute(node, "num_heads") + kv_num_heads = get_attribute(node, "kv_num_heads") + query_shape = self._get_shape(node, 0) + output_dtype = self.known_vi_[node.input[0]].type.tensor_type.elem_type + # The head width is only recoverable when the query hidden size divides evenly. Otherwise the + # node is malformed (or the attributes are missing) and we fall back to generic propagation + # instead of silently emitting a truncated output width. + head_size = None + if query_shape is not None and len(query_shape) == 2 and is_literal(query_shape[1]) and num_heads: + divisor = num_heads if is_latent_kv else (num_heads + 2 * (kv_num_heads or 0)) + if divisor > 0 and query_shape[1] % divisor == 0: + head_size = query_shape[1] // divisor + if head_size is not None: + v_head_size = (get_attribute(node, "v_head_size", 0) or head_size) if is_latent_kv else head_size + vi = self.known_vi_[node.output[0]] + vi.CopyFrom( + helper.make_tensor_value_info( + node.output[0], output_dtype, [query_shape[0], num_heads * v_head_size] + ) + ) + else: + self._propagate_shape_and_type(node) + else: + self._propagate_shape_and_type(node) + + # The cache outputs alias the cache inputs, so they carry the cache element type and shape. + # value_cache (input 4) and value_cache_out (output 2) are absent in LATENT mode, so guard the + # aliased input as well: a node may declare the output while omitting the corresponding input. + for output_index, input_index in ((1, 3), (2, 4)): + if ( + len(node.output) > output_index + and node.output[output_index] + and len(node.input) > input_index + and node.input[input_index] + ): + self._propagate_shape_and_type(node, input_index, output_index) def _infer_GroupQueryAttention(self, node): # noqa: N802 output_dtype = self.known_vi_[node.input[0]].type.tensor_type.elem_type diff --git a/onnxruntime/test/contrib_ops/decoder_masked_multihead_attention_op_test.cc b/onnxruntime/test/contrib_ops/decoder_masked_multihead_attention_op_test.cc index 2451f7e03a281..571876a94c021 100644 --- a/onnxruntime/test/contrib_ops/decoder_masked_multihead_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/decoder_masked_multihead_attention_op_test.cc @@ -961,5 +961,36 @@ TEST(DecoderMaskedMultiHeadAttentionTest, cpu_cache_indirection_beam_index_out_o {}, nullptr, &execution_providers); } +TEST(DecoderMaskedMultiHeadAttentionTest, cpu_cache_indirection_batch_beam_not_divisible_by_num_beams) { + // num_beams = 2 does not evenly divide batch_beam_size = 3. + // cache_indirection dim 0 is a valid-looking 1 (= 3 / 2 with truncating division), + // but the shape is inconsistent with batch_beam_size and must be rejected up front. + OpTester tester("DecoderMaskedMultiHeadAttention", 1, onnxruntime::kMSDomain); + tester.AddAttribute("num_heads", 1); + tester.AddAttribute("past_present_share_buffer", 1); + + tester.AddInput("query", {3, 1, 4}, std::vector(12, 0.1f)); + tester.AddInput("key", {3, 1, 4}, std::vector(12, 0.2f)); + tester.AddInput("value", {3, 1, 4}, std::vector(12, 0.3f)); + tester.AddOptionalInputEdge(); + tester.AddOptionalInputEdge(); + tester.AddInput("past_key", {3, 1, 4, 4}, std::vector(48, 0.4f)); + tester.AddInput("past_value", {3, 1, 4, 4}, std::vector(48, 0.5f)); + tester.AddInput("past_sequence_length", {1}, {2}); + tester.AddInput("beam_width", {1}, {2}); + tester.AddInput("cache_indirection", {1, 2, 4}, std::vector(8, 0)); + tester.AddOptionalInputEdge(); + + tester.AddOutput("output", {3, 1, 4}, std::vector(12, 0.0f)); + tester.AddOutput("present_key", {3, 1, 4, 4}, std::vector(48, 0.0f)); + tester.AddOutput("present_value", {3, 1, 4, 4}, std::vector(48, 0.0f)); + tester.AddOptionalOutputEdge(); + + std::vector> execution_providers; + execution_providers.push_back(DefaultCpuExecutionProvider()); + tester.Run(OpTester::ExpectResult::kExpectFailure, "must equal batch_beam_size", + {}, nullptr, &execution_providers); +} + } // namespace test } // namespace onnxruntime diff --git a/onnxruntime/test/contrib_ops/multihead_attention_op_test.cc b/onnxruntime/test/contrib_ops/multihead_attention_op_test.cc index efe7d582475cc..a5fcbb25ca93f 100644 --- a/onnxruntime/test/contrib_ops/multihead_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/multihead_attention_op_test.cc @@ -588,6 +588,35 @@ TEST(MultiHeadAttentionTest, CacheIndirectionBeamIndexOutOfRange) { {}, nullptr, &execution_providers); } +TEST(MultiHeadAttentionTest, CacheIndirectionBatchBeamNotDivisibleByNumBeams) { + // num_beams = 2 does not evenly divide batch_beam_size = 3. + // cache_indirection dim 0 is a valid-looking 1 (= 3 / 2 with truncating division), + // but the shape is inconsistent with batch_beam_size and must be rejected up front. + OpTester tester("MultiHeadAttention", 1, onnxruntime::kMSDomain); + tester.AddAttribute("num_heads", 1); + + tester.AddInput("query", {3, 1, 4}, std::vector(12, 0.1f)); + tester.AddInput("key", {3, 1, 4}, std::vector(12, 0.2f)); + tester.AddInput("value", {3, 1, 4}, std::vector(12, 0.3f)); + tester.AddOptionalInputEdge(); + tester.AddOptionalInputEdge(); + tester.AddOptionalInputEdge(); + tester.AddInput("past_key", {3, 1, 4, 4}, std::vector(48, 0.4f)); + tester.AddInput("past_value", {3, 1, 4, 4}, std::vector(48, 0.5f)); + tester.AddInput("past_sequence_length", {1}, {2}); + tester.AddInput("cache_indirection", {1, 2, 4}, std::vector(8, 0)); + + tester.AddOutput("output", {3, 1, 4}, std::vector(12, 0.0f)); + tester.AddOutput("present_key", {3, 1, 4, 4}, std::vector(48, 0.0f)); + tester.AddOutput("present_value", {3, 1, 4, 4}, std::vector(48, 0.0f)); + tester.AddOptionalOutputEdge(); + + std::vector> execution_providers; + execution_providers.push_back(DefaultCpuExecutionProvider()); + tester.Run(OpTester::ExpectResult::kExpectFailure, "must equal batch_beam_size", + {}, nullptr, &execution_providers); +} + TEST(MultiHeadAttentionTest, OutputQKWithPaddingMaskAndAttentionBias) { OpTester tester("MultiHeadAttention", 1, onnxruntime::kMSDomain); tester.AddAttribute("num_heads", 1); diff --git a/onnxruntime/test/providers/cpu/math/matmul_test.cc b/onnxruntime/test/providers/cpu/math/matmul_test.cc index adef2a7cb17dc..a239e5755292e 100644 --- a/onnxruntime/test/providers/cpu/math/matmul_test.cc +++ b/onnxruntime/test/providers/cpu/math/matmul_test.cc @@ -861,6 +861,16 @@ TEST(MathOpTest, MatMulSubgroupMatrix) { {"SplitKScratchCap (64x64,K=256)", {64, 256}, 256, 64}, // Batched A folds to M=8; one tile + K=256 -> split_k=8. {"SplitKBatched (2*4 -> M=8)", {2, 4, 256}, 256, 16}, + // Odd N: the subgroup f16 load needs an even B row stride, so a constant odd-N + // weight is padded once to N+1 (even) and cached; output is still written at + // the real, odd N. Covers small/large odd N, min K, partial M, batched-A fold, + // and split-K, all with odd N. + {"OddN15 (N=15)", {32, 64}, 64, 15}, + {"OddN33 (N=33)", {16, 64}, 64, 33}, + {"OddN MinK (K=16,N=17)", {8, 16}, 16, 17}, + {"OddN PartialM (M=40,N=31)", {40, 64}, 64, 31}, + {"OddN BatchedA (2*32 -> M=64,N=63)", {2, 32, 64}, 64, 63}, + {"OddN SplitK (K=256,N=17)", {8, 256}, 256, 17}, }; for (const auto& c : cases) { diff --git a/onnxruntime/test/providers/webgpu/matmul_large_test.cc b/onnxruntime/test/providers/webgpu/matmul_large_test.cc index 070b57343ce6c..143034563be00 100644 --- a/onnxruntime/test/providers/webgpu/matmul_large_test.cc +++ b/onnxruntime/test/providers/webgpu/matmul_large_test.cc @@ -41,7 +41,8 @@ static void ComputeExpectedResult(const std::vector& a_vals, const std::v } template -void RunTestTyped(std::initializer_list a_dims, std::initializer_list b_dims) { +void RunTestTyped(std::initializer_list a_dims, std::initializer_list b_dims, + bool b_is_constant = false) { static_assert(std::is_same_v || std::is_same_v, "unexpected type for T"); auto webgpu_ep = DefaultWebGpuExecutionProvider(); @@ -68,11 +69,11 @@ void RunTestTyped(std::initializer_list a_dims, std::initializer_list) { test.AddInput("A", a_dims, a_vals); - test.AddInput("B", b_dims, b_vals); + test.AddInput("B", b_dims, b_vals, b_is_constant); test.AddOutput("Y", output_dims, expected_vals); } else { test.AddInput("A", a_dims, FloatsToMLFloat16s(a_vals)); - test.AddInput("B", b_dims, FloatsToMLFloat16s(b_vals)); + test.AddInput("B", b_dims, FloatsToMLFloat16s(b_vals), b_is_constant); test.AddOutput("Y", output_dims, FloatsToMLFloat16s(expected_vals)); test.SetOutputAbsErr("Y", 0.055f); test.SetOutputRelErr("Y", 0.02f); @@ -148,6 +149,23 @@ TEST(MatMul_Large, DISABLED_BatchedB_LargeBatchSmallTile) { RunBothTypes({128, 8, 256}, {128, 256, 16}); } +// Constant f16 weight with odd N. The Intel f16 subgroup-matrix load needs an +// even B row stride, so a non-constant odd-N B falls back to the generic path. When +// B is a constant initializer, the first Compute lazily pads it to an even stride +// (N+1) and the subgroup kernel consumes the cached copy via the N_b uniform (output +// is still written at the real, odd N). Marking B constant here exercises that +// padded path for both shared 2D and batched weights across several odd N (1023, +// 33, 65), with even K and both aligned and partial M. Results must match the +// reference. f16 only: the subgroup kernel is f16, so a float B would take the +// generic path. +TEST(MatMul_Large, DISABLED_ConstantWeightOddN) { + RunTestTyped({128, 64}, {64, 1023}, /*b_is_constant=*/true); + RunTestTyped({127, 64}, {64, 1023}, /*b_is_constant=*/true); + RunTestTyped({64, 96}, {96, 33}, /*b_is_constant=*/true); + RunTestTyped({130, 80}, {80, 65}, /*b_is_constant=*/true); + RunTestTyped({2, 127, 64}, {2, 64, 1023}, /*b_is_constant=*/true); +} + // Broadcasted batch dims that are NOT identical but share the same batch // *product* (A=[2,1,...], B=[1,2,...] -> [2,2,...]; A=[1,4,...], B=[4,1,...] -> // [4,4,...]). A product-only batch check would wrongly route these onto the diff --git a/onnxruntime/test/python/transformers/test_paged_attention_cuda.py b/onnxruntime/test/python/transformers/test_paged_attention_cuda.py index 75f34511874e8..fb7ca843d4d47 100644 --- a/onnxruntime/test/python/transformers/test_paged_attention_cuda.py +++ b/onnxruntime/test/python/transformers/test_paged_attention_cuda.py @@ -20,7 +20,6 @@ from onnx import TensorProto, helper from packaging import version from parameterized import parameterized -from test_gqa_cpu import smooth_softmax_ref from onnxruntime import InferenceSession, OrtValue, SessionOptions, get_available_providers @@ -28,6 +27,16 @@ pipeline_mode = True # Reduces number of tests so pipeline doesn't time out +# Element type of the paged KV cache, keyed by Config.kv_cache_type. +KV_CACHE_TENSOR_PROTO = { + "float16": TensorProto.FLOAT16, + "int8": TensorProto.INT8, + "fp8": TensorProto.FLOAT8E4M3FN, +} + +# Largest magnitude representable by each quantized cache type. +KV_CACHE_QMAX = {"int8": 127.0, "fp8": 448.0} + class Config: batch_size = 0 @@ -43,6 +52,22 @@ class Config: packed = False softcap = 0.0 ep = "CUDAExecutionProvider" + # Optional features layered on top of the original schema. They default to off so that every + # pre-existing parameterized test keeps its generated name and behavior. + use_slot_mapping = False + use_head_sink = False + use_qk_norm = False + qk_norm_epsilon = 1e-6 + # Quantized paged KV cache. "float16" keeps the cache unquantized; "int8" and "fp8" store the + # cache in the corresponding narrow type and require a matching non-"NONE" quant type. + kv_cache_type = "float16" + k_quant_type = "NONE" + v_quant_type = "NONE" + # "" means the cache tensor's own element type is the logical element type. Sub-byte names + # ("int4", "float4e2m1") describe a uint8 packed cache and are not supported yet. Every legal + # value is signed: quantization is symmetric with no zero point, so "uint4"/"uint8" are invalid. + k_cache_dtype = "" + v_cache_dtype = "" def __init__( self, @@ -84,6 +109,64 @@ def __repr__(self): ) +def kv_scale_shape(config, quant_type): + """Shape of the k_scale / v_scale input for a given quantization granularity.""" + if quant_type == "PER_TENSOR": + return [1] + if quant_type == "PER_CHANNEL": + return [config.kv_num_heads, 1, config.head_size] + raise ValueError(f"Unsupported quant_type: {quant_type}") + + +def compute_kv_scale(tensors, quant_type, kv_cache_type, kv_num_heads, head_size): + """Symmetric (zero-point-free) scale for the paged KV cache. + + 'tensors' are every tensor that will end up in the cache -- the pre-existing pages *and* the new + tokens the kernel is about to write -- so that the chosen scale never clips and the reference can + reproduce the kernel bit-for-bit. Each tensor's two trailing dimensions are (kv_num_heads, + head_size), which is exactly the PER_CHANNEL scale layout. + """ + qmax = KV_CACHE_QMAX[kv_cache_type] + if quant_type == "PER_TENSOR": + amax = max(float(t.abs().max().item()) for t in tensors) + return torch.tensor([max(amax, 1e-6) / qmax], dtype=torch.float32, device="cuda") + amax = None + for t in tensors: + per_channel = t.reshape(-1, kv_num_heads, head_size).abs().amax(dim=0) + amax = per_channel if amax is None else torch.maximum(amax, per_channel) + scale = torch.clamp(amax.to(torch.float32), min=1e-6) / qmax + return scale.reshape(kv_num_heads, 1, head_size) + + +def broadcast_kv_scale(scale, quant_type, kv_num_heads, head_size): + """View the scale so that it broadcasts against any tensor shaped (..., kv_num_heads, head_size).""" + if quant_type == "PER_TENSOR": + return scale + return scale.reshape(1, 1, kv_num_heads, head_size) + + +def quantize_kv(tensor_float, scale, kv_cache_type): + """Mirror of QuantizeToCache in paged_attention_impl.cu. + + The kernel multiplies by the reciprocal of the scale rather than dividing, so the reference does + the same: with round-to-nearest-even the two differ by one LSB often enough to make an exact + comparison of the updated cache flaky otherwise. + """ + scaled = tensor_float.to(torch.float32) * torch.reciprocal(scale) + if kv_cache_type == "fp8": + return torch.clamp(scaled, -KV_CACHE_QMAX["fp8"], KV_CACHE_QMAX["fp8"]).to(torch.float8_e4m3fn) + return torch.clamp(torch.round(scaled), -128.0, 127.0).to(torch.int8) + + +def dequantize_kv(quantized, scale): + """Mirror of DequantizeFromCache in paged_attention_impl.cu.""" + return (quantized.to(torch.float32) * scale).to(torch.float16) + + +def quantize_dequantize_kv(tensor_float, scale, kv_cache_type): + return dequantize_kv(quantize_kv(tensor_float, scale, kv_cache_type), scale) + + def create_paged_attention_graph( config, num_tokens, @@ -91,6 +174,25 @@ def create_paged_attention_graph( max_blocks_per_sequence, local_window_size=-1, ): + cache_proto_type = KV_CACHE_TENSOR_PROTO[config.kv_cache_type] + # The scale inputs and the quantization attributes are emitted independently of the cache dtype + # so that invalid combinations (quantized cache without a quant type, and vice versa) can be + # built and their rejection tested. + has_k_scale = config.k_quant_type != "NONE" + has_v_scale = config.v_quant_type != "NONE" + # Optional host-side [max_query_len_bound, max_kv_len_bound]. When present the kernel can skip + # the device readback of the cumulative length arrays, so results must be identical either way. + has_attention_metadata = getattr(config, "use_attention_metadata", False) + quant_attrs = ( + { + "k_quant_type": config.k_quant_type, + "v_quant_type": config.v_quant_type, + "k_cache_dtype": config.k_cache_dtype, + "v_cache_dtype": config.v_cache_dtype, + } + if (has_k_scale or has_v_scale or config.kv_cache_type != "float16") + else {} + ) nodes = [ helper.make_node( "PagedAttention", @@ -105,6 +207,13 @@ def create_paged_attention_graph( "block_table", "cos_cache" if config.rotary else "", "sin_cache" if config.rotary else "", + "slot_mapping" if config.use_slot_mapping else "", + "head_sink" if config.use_head_sink else "", + "q_norm_weight" if config.use_qk_norm else "", + "k_norm_weight" if config.use_qk_norm else "", + "k_scale" if has_k_scale else "", + "v_scale" if has_v_scale else "", + "attention_metadata" if has_attention_metadata else "", ], ["output", "key_cache_out", "value_cache_out"], "PagedAttention_0", @@ -114,7 +223,9 @@ def create_paged_attention_graph( do_rotary=config.rotary, rotary_interleaved=config.rotary_interleaved, softcap=config.softcap, + qk_norm_epsilon=config.qk_norm_epsilon, domain="com.microsoft", + **quant_attrs, ), ] @@ -131,7 +242,7 @@ def create_paged_attention_graph( ), helper.make_tensor_value_info( "key_cache", - TensorProto.FLOAT16, + cache_proto_type, [ num_blocks, config.paged_kv_block_size, @@ -141,7 +252,7 @@ def create_paged_attention_graph( ), helper.make_tensor_value_info( "value_cache", - TensorProto.FLOAT16, + cache_proto_type, [ num_blocks, config.paged_kv_block_size, @@ -203,6 +314,31 @@ def create_paged_attention_graph( ], ), ] + if config.use_slot_mapping: + graph_input += [ + helper.make_tensor_value_info("slot_mapping", TensorProto.INT32, [num_tokens]), + ] + if config.use_head_sink: + graph_input += [ + helper.make_tensor_value_info("head_sink", TensorProto.FLOAT16, [config.num_heads]), + ] + if config.use_qk_norm: + graph_input += [ + helper.make_tensor_value_info("q_norm_weight", TensorProto.FLOAT16, [config.head_size]), + helper.make_tensor_value_info("k_norm_weight", TensorProto.FLOAT16, [config.head_size]), + ] + if has_k_scale: + graph_input += [ + helper.make_tensor_value_info("k_scale", TensorProto.FLOAT, kv_scale_shape(config, config.k_quant_type)), + ] + if has_v_scale: + graph_input += [ + helper.make_tensor_value_info("v_scale", TensorProto.FLOAT, kv_scale_shape(config, config.v_quant_type)), + ] + if has_attention_metadata: + graph_input += [ + helper.make_tensor_value_info("attention_metadata", TensorProto.INT32, [2]), + ] graph_output = [ helper.make_tensor_value_info( @@ -212,7 +348,7 @@ def create_paged_attention_graph( ), helper.make_tensor_value_info( "key_cache_out", - TensorProto.FLOAT16, + cache_proto_type, [ num_blocks, config.paged_kv_block_size, @@ -222,7 +358,7 @@ def create_paged_attention_graph( ), helper.make_tensor_value_info( "value_cache_out", - TensorProto.FLOAT16, + cache_proto_type, [ num_blocks, config.paged_kv_block_size, @@ -263,10 +399,17 @@ def paged_attention_func( sin=None, window_size=-1, sdpa_kernel=0, + slot_mapping=None, + head_sink=None, + q_norm_weight=None, + k_norm_weight=None, + k_scale=None, + v_scale=None, ): num_tokens = cumulative_sequence_length[-1].item() num_blocks = key_cache.shape[0] max_blocks_per_sequence = block_table.shape[1] + quantized = config.kv_cache_type != "float16" onnx_model_str = create_paged_attention_graph( config, num_tokens, @@ -276,12 +419,24 @@ def paged_attention_func( ) ort_inputs = { "query": query.detach().cpu().numpy(), - "key_cache": OrtValue.ortvalue_from_numpy(key_cache.detach().cpu().numpy(), "cuda", 0), - "value_cache": OrtValue.ortvalue_from_numpy(value_cache.detach().cpu().numpy(), "cuda", 0), "cumulative_sequence_length": cumulative_sequence_length.detach().cpu().numpy(), "past_seqlens": past_seqlens.detach().cpu().numpy(), "block_table": block_table.detach().cpu().numpy(), } + if getattr(config, "use_attention_metadata", False): + override = getattr(config, "attention_metadata_override", None) + if override is not None: + ort_inputs["attention_metadata"] = override + else: + cum_q = cumulative_sequence_length.detach().cpu().numpy().astype(numpy.int64) + query_lens = cum_q[1:] - cum_q[:-1] + kv_lens = past_seqlens.detach().cpu().numpy().astype(numpy.int64) + query_lens + # The exact per-step maxima are valid upper bounds for a single Run, which is what these + # tests do. A real scheduler would pass looser, replay-wide bounds instead. + ort_inputs["attention_metadata"] = numpy.array([query_lens.max(), kv_lens.max()], dtype=numpy.int32) + if not quantized: + ort_inputs["key_cache"] = OrtValue.ortvalue_from_numpy(key_cache.detach().cpu().numpy(), "cuda", 0) + ort_inputs["value_cache"] = OrtValue.ortvalue_from_numpy(value_cache.detach().cpu().numpy(), "cuda", 0) sess_options = SessionOptions() if sdpa_kernel != 0 and config.ep == "CUDAExecutionProvider": providers = [(config.ep, {"sdpa_kernel": str(sdpa_kernel)})] @@ -299,20 +454,60 @@ def paged_attention_func( ort_inputs["sin_cache"] = sin.detach().cpu().numpy() io_binding.bind_cpu_input("cos_cache", ort_inputs["cos_cache"]) io_binding.bind_cpu_input("sin_cache", ort_inputs["sin_cache"]) + for name, tensor in ( + ("slot_mapping", slot_mapping), + ("head_sink", head_sink), + ("q_norm_weight", q_norm_weight), + ("k_norm_weight", k_norm_weight), + ("k_scale", k_scale), + ("v_scale", v_scale), + ): + if tensor is not None: + ort_inputs[name] = tensor.detach().cpu().numpy() + io_binding.bind_cpu_input(name, ort_inputs[name]) + if "attention_metadata" in ort_inputs: + io_binding.bind_cpu_input("attention_metadata", ort_inputs["attention_metadata"]) io_binding.bind_cpu_input("query", ort_inputs["query"]) - io_binding.bind_input( - "key_cache", "cuda", 0, numpy.float16, ort_inputs["key_cache"].shape(), ort_inputs["key_cache"].data_ptr() - ) - io_binding.bind_input( - "value_cache", "cuda", 0, numpy.float16, ort_inputs["value_cache"].shape(), ort_inputs["value_cache"].data_ptr() - ) + if quantized: + # A quantized cache has no numpy dtype, so bind the torch device buffers directly. + cache_proto_type = KV_CACHE_TENSOR_PROTO[config.kv_cache_type] + key_cache = key_cache.contiguous() + value_cache = value_cache.contiguous() + io_binding.bind_input("key_cache", "cuda", 0, cache_proto_type, tuple(key_cache.shape), key_cache.data_ptr()) + io_binding.bind_input( + "value_cache", "cuda", 0, cache_proto_type, tuple(value_cache.shape), value_cache.data_ptr() + ) + else: + io_binding.bind_input( + "key_cache", "cuda", 0, numpy.float16, ort_inputs["key_cache"].shape(), ort_inputs["key_cache"].data_ptr() + ) + io_binding.bind_input( + "value_cache", + "cuda", + 0, + numpy.float16, + ort_inputs["value_cache"].shape(), + ort_inputs["value_cache"].data_ptr(), + ) io_binding.bind_cpu_input("cumulative_sequence_length", ort_inputs["cumulative_sequence_length"]) io_binding.bind_cpu_input("past_seqlens", ort_inputs["past_seqlens"]) io_binding.bind_cpu_input("block_table", ort_inputs["block_table"]) io_binding.bind_output("output") - io_binding.bind_ortvalue_output("key_cache_out", ort_inputs["key_cache"]) - io_binding.bind_ortvalue_output("value_cache_out", ort_inputs["value_cache"]) + if quantized: + # Each cache output must alias its input, which is what the op requires anyway. + io_binding.bind_output( + "key_cache_out", "cuda", 0, cache_proto_type, tuple(key_cache.shape), key_cache.data_ptr() + ) + io_binding.bind_output( + "value_cache_out", "cuda", 0, cache_proto_type, tuple(value_cache.shape), value_cache.data_ptr() + ) + else: + io_binding.bind_ortvalue_output("key_cache_out", ort_inputs["key_cache"]) + io_binding.bind_ortvalue_output("value_cache_out", ort_inputs["value_cache"]) ort_session.run_with_iobinding(io_binding) + if quantized: + output = torch.tensor(numpy.array(io_binding.copy_outputs_to_cpu()[0])) + return output, key_cache, value_cache output, key_cache_out, value_cache_out = io_binding.copy_outputs_to_cpu() output = torch.tensor(numpy.array(output)) return output, key_cache_out, value_cache_out @@ -353,7 +548,7 @@ def attention_ref( softcap=0.0, upcast=True, reorder_ops=False, - use_smooth_softmax=False, + head_sink=None, ): """ Arguments: @@ -405,9 +600,12 @@ def attention_ref( ) scores.masked_fill_(local_mask, float("-inf")) - if use_smooth_softmax: - head_sink = None - attention = smooth_softmax_ref(scores, head_sink) + if head_sink is not None: + # Append one extra logit per (batch, head, query) to the softmax denominator that + # contributes no value. head_sink is the learned logit. + b, n, s, _ = scores.shape + sink = head_sink.to(scores.dtype).reshape(1, n, 1, 1).expand(b, -1, s, -1) + attention = torch.softmax(torch.cat([scores, sink], dim=-1), dim=-1)[..., :-1] else: attention = torch.softmax(scores, dim=-1) @@ -429,6 +627,19 @@ def attention_ref( return output.to(dtype=dtype_og), attention.to(dtype=dtype_og) +def rms_norm_ref(x, weight, epsilon): + """Per-head RMSNorm reference matching the fused CUDA prologue: reduce in fp32 over the last + dimension (head_size), scale, then cast back to the input dtype. + + Arguments: + x: (..., head_size) + weight: (head_size) + """ + x_f32 = x.float() + inv_rms = torch.rsqrt(x_f32.pow(2).mean(dim=-1, keepdim=True) + epsilon) + return (x_f32 * inv_rms * weight.float()).to(dtype=x.dtype) + + def rotary_embedding(*args, **kwargs): # Use local import since triton is not available in Windows. from rotary_flash import apply_rotary_emb # noqa: PLC0415 @@ -491,6 +702,32 @@ def generate_block_kvcache(config: Config, device, dtype): return k_cache, v_cache, block_table, k_cache_paged, v_cache_paged +def gather_paged_to_batch(config: Config, paged, block_table): + """Gather a paged [num_blocks, block_size, kv_num_heads, head_size] cache into the dense + [batch_size, total_sequence_length, kv_num_heads, head_size] view the reference works with.""" + return rearrange( + paged[block_table.to(dtype=torch.long).flatten()], + "(b nblocks) block_size ... -> b (nblocks block_size) ...", + b=config.batch_size, + )[:, : config.total_sequence_length] + + +def derive_slot_mapping(config: Config, past_seqlens, new_seqlens, cum_seqlens, block_table): + """Reproduce, on the host, the flat cache slot that the kernel derives for every query token + when 'slot_mapping' is absent: block_table[b, pos // block_size] * block_size + pos % block_size + where pos = past_seqlens[b] + index_of_token_within_its_sequence.""" + token_count = int(cum_seqlens[-1].item()) + slot_mapping = torch.empty(token_count, dtype=torch.int32, device="cuda") + block_table_cpu = block_table.cpu() + for b in range(config.batch_size): + start = int(cum_seqlens[b].item()) + for j in range(int(new_seqlens[b].item())): + pos = int(past_seqlens[b].item()) + j + block_id = int(block_table_cpu[b, pos // config.paged_kv_block_size].item()) + slot_mapping[start + j] = block_id * config.paged_kv_block_size + pos % config.paged_kv_block_size + return slot_mapping + + def parity_check_paged_attention( config: Config, rtol=1e-3, @@ -558,11 +795,35 @@ def parity_check_paged_attention( # Generate kv cache and associated block-based data structures k_cache, v_cache, block_table, k_cache_paged, v_cache_paged = generate_block_kvcache(config, "cuda", torch.float16) + # Optional per-head attention sink. + head_sink = None + if config.use_head_sink: + # Spread over [-2, 6]: exp(sink) then ranges from negligible to far larger than a typical + # softmax denominator, so a kernel that ignored the sink could not pass within tolerance. + head_sink = (torch.rand(config.num_heads, device="cuda") * 8.0 - 2.0).to(dtype=torch.float16) + + # Optional QK-Norm. The kernel applies RMSNorm to every Q and K head before rotary embedding, + # so the reference has to normalize before computing q_ro / k_ro below, and the normalized + + # rotated K is what must land in the KV cache. + q_norm_weight = None + k_norm_weight = None + if config.use_qk_norm: + q_norm_weight = torch.randn(config.head_size, device="cuda", dtype=torch.float16) + k_norm_weight = torch.randn(config.head_size, device="cuda", dtype=torch.float16) + q = rms_norm_ref(q, q_norm_weight, config.qk_norm_epsilon) + k_new = rms_norm_ref(k_new, k_norm_weight, config.qk_norm_epsilon) + + # Optional explicit slot mapping. Reproducing exactly what the kernel derives from + # past_seqlens / cumulative_sequence_length / block_table must give identical results. + slot_mapping = None + if config.use_slot_mapping: + slot_mapping = derive_slot_mapping(config, past_seqlens, new_seqlens, cum_seqlens, block_table) + # Set window size for local / causal window_size = (-1, -1) left_window_size = -1 if config.local: - left_window_size = random.randint(0, config.total_sequence_length - 1) # random.randint is inclusive + left_window_size = random.randint(1, config.total_sequence_length - 1) # random.randint is inclusive window_size = (left_window_size, 0) else: left_window_size = -1 @@ -581,6 +842,27 @@ def parity_check_paged_attention( cos, sin = None, None q_ro, k_ro = q, k_new + # Quantized paged KV cache. The pages the kernel reads back are the dequantized values, and the + # new tokens it writes go through the same quantize step, so the reference models both. The scale + # covers the existing pages *and* the incoming tokens so that nothing clips. + k_scale = v_scale = None + k_scale_b = v_scale_b = None + if config.kv_cache_type != "float16": + k_scale = compute_kv_scale( + [k_cache_paged, k_ro], config.k_quant_type, config.kv_cache_type, config.kv_num_heads, config.head_size + ) + v_scale = compute_kv_scale( + [v_cache_paged, v_new], config.v_quant_type, config.kv_cache_type, config.kv_num_heads, config.head_size + ) + k_scale_b = broadcast_kv_scale(k_scale, config.k_quant_type, config.kv_num_heads, config.head_size) + v_scale_b = broadcast_kv_scale(v_scale, config.v_quant_type, config.kv_num_heads, config.head_size) + k_cache_paged = quantize_kv(k_cache_paged, k_scale_b, config.kv_cache_type) + v_cache_paged = quantize_kv(v_cache_paged, v_scale_b, config.kv_cache_type) + k_cache = gather_paged_to_batch(config, dequantize_kv(k_cache_paged, k_scale_b), block_table) + v_cache = gather_paged_to_batch(config, dequantize_kv(v_cache_paged, v_scale_b), block_table) + k_ro = quantize_dequantize_kv(k_ro, k_scale_b, config.kv_cache_type) + v_new = quantize_dequantize_kv(v_new, v_scale_b, config.kv_cache_type) + # Update reference kv cache k_cache_ref = k_cache.clone() v_cache_ref = v_cache.clone() @@ -613,6 +895,7 @@ def parity_check_paged_attention( causal=True, window_size=window_size, softcap=config.softcap, + head_sink=head_sink, ) out_ref = out_ref.detach().cpu().numpy() @@ -634,12 +917,30 @@ def parity_check_paged_attention( sin, left_window_size, sdpa_kernel=sdpa_kernel, + slot_mapping=slot_mapping, + head_sink=head_sink, + q_norm_weight=q_norm_weight, + k_norm_weight=k_norm_weight, + k_scale=k_scale, + v_scale=v_scale, ) + if config.kv_cache_type != "float16": + updated_k_cache_paged = dequantize_kv(updated_k_cache_paged, k_scale_b).cpu().numpy() + updated_v_cache_paged = dequantize_kv(updated_v_cache_paged, v_scale_b).cpu().numpy() num_tokens = q_unpad.shape[0] out = torch.reshape(out, (num_tokens, config.num_heads, config.head_size)) out = out.detach().cpu().numpy() err_msg = f" with {config}" + # The updated cache is compared to the reference at one quantization step of slack: the host + # computes rotary / RMSNorm slightly differently from the kernel, and a 1-ULP fp16 difference in + # the pre-quantization value is enough to move the rounded result by a whole step. + cache_rtol, cache_atol = rtol, atol + if config.kv_cache_type == "int8": + cache_atol = atol + max(float(k_scale.max().item()), float(v_scale.max().item())) + elif config.kv_cache_type == "fp8": + cache_rtol = rtol + 2.0**-3 # float8e4m3fn has 3 mantissa bits + # Make sure past-present buffer updating correctly present_k = rearrange( updated_k_cache_paged[block_table.to(dtype=torch.long).flatten().cpu()], @@ -655,16 +956,16 @@ def parity_check_paged_attention( numpy.testing.assert_allclose( present_k[i, : total_seqlens[i]], k_cache_ref[i, : total_seqlens[i]].detach().cpu().numpy(), - rtol=rtol, - atol=atol, + rtol=cache_rtol, + atol=cache_atol, equal_nan=True, err_msg=err_msg, ) numpy.testing.assert_allclose( present_v[i, : total_seqlens[i]], v_cache_ref[i, : total_seqlens[i]].detach().cpu().numpy(), - rtol=rtol, - atol=atol, + rtol=cache_rtol, + atol=cache_atol, equal_nan=True, err_msg=err_msg, ) @@ -674,10 +975,17 @@ def parity_check_paged_attention( numpy.testing.assert_allclose(out_i, out_ref_i, rtol=rtol, atol=atol, equal_nan=True, err_msg=err_msg) +def has_cuda_device(): + """Every test in this file allocates torch tensors on "cuda" and runs the CUDA EP. + + Some pipelines install a CPU-only torch wheel, where any `device="cuda"` allocation raises + "AssertionError: Torch not compiled with CUDA enabled". Gate every test class on this so those + runs skip instead of erroring.""" + return torch.cuda.is_available() and "CUDAExecutionProvider" in get_available_providers() + + def has_flash_attention(): - if not torch.cuda.is_available(): - return False - if "CUDAExecutionProvider" not in get_available_providers(): + if not has_cuda_device(): return False major, _ = torch.cuda.get_device_capability() return major >= 8 and ( @@ -689,9 +997,7 @@ def has_flash_attention(): def has_memory_efficient_attention(): # CUTLASS fMHA (MemoryEfficientAttention) gate — these tests are fp16-only, # so sm>=53 is sufficient. bf16 MEA would require sm>=80 but is not covered here. - if not torch.cuda.is_available(): - return False - if "CUDAExecutionProvider" not in get_available_providers(): + if not has_cuda_device(): return False major, minor = torch.cuda.get_device_capability() return (major * 10 + minor) >= 53 @@ -704,6 +1010,13 @@ def has_memory_efficient_attention(): # where FlashAttention would otherwise be preferred. SDPA_KERNEL_EFFICIENT_ATTENTION = 2 +# Bit value matching AttentionBackend::DECODER_ATTENTION in +# onnxruntime/contrib_ops/cpu/bert/attention_common.h. Passing this as the CUDA provider option +# `sdpa_kernel` leaves the paged decode kernel as the only enabled backend, which is how the +# unquantized decode path is reached (it is otherwise only auto-selected for a quantized cache or +# when FlashAttention is unavailable). +SDPA_KERNEL_DECODER_ATTENTION = 512 + def paged_attention_test_cases(): batches = [4] if pipeline_mode else [1, 3, 5] @@ -785,6 +1098,212 @@ def test_paged_attention_mea(self, _, config): ) +@unittest.skipIf(not has_cuda_device(), reason="CUDA is not available, skipping tests.") +class TestPagedAttentionFeatures(unittest.TestCase): + """Coverage for the optional inputs and attributes added on top of the original schema: + slot_mapping, head_sink and QK-Norm, plus the block_size and batch_size + limits that were lifted at the same time.""" + + def _config(self, **overrides): + kwargs = { + "batch_size": 4, + "sequence_length": 33, + "total_sequence_length": 128, + "num_heads": 8, + "kv_num_heads": 2, + "head_size": 64, + "paged_kv_block_size": 256, + "local": False, + "rotary": False, + "rotary_interleaved": False, + "packed": False, + "softcap": 0.0, + } + feature_overrides = {k: overrides.pop(k) for k in list(overrides) if k not in kwargs} + kwargs.update(overrides) + config = Config(**kwargs) + for key, value in feature_overrides.items(): + setattr(config, key, value) + return config + + # ---- slot_mapping ------------------------------------------------------------------- + + @unittest.skipIf(not has_flash_attention(), reason="Flash Attention is not available") + def test_slot_mapping_matches_derived_mapping(self): + # An explicit slot_mapping that reproduces the derived mapping must be a no-op. + parity_check_paged_attention(self._config(use_slot_mapping=True), rtol=5e-3, atol=5e-3) + + @unittest.skipIf(not has_flash_attention(), reason="Flash Attention is not available") + def test_slot_mapping_with_rotary_and_packed(self): + config = self._config(use_slot_mapping=True, rotary=True, packed=True) + parity_check_paged_attention(config, rtol=5e-3, atol=5e-3) + + @unittest.skipIf( + not has_memory_efficient_attention(), + reason="MemoryEfficientAttention (fp16) requires sm>=53", + ) + def test_slot_mapping_mea(self): + parity_check_paged_attention( + self._config(use_slot_mapping=True), + rtol=5e-3, + atol=5e-3, + sdpa_kernel=SDPA_KERNEL_EFFICIENT_ATTENTION, + ) + + @unittest.skipIf(not has_flash_attention(), reason="Flash Attention is not available") + def test_slot_mapping_negative_one_skips_cache_write(self): + # -1 tells the kernel not to store this token's K/V, which is how a scheduler suppresses + # writes for prefix-cache hits or rejected speculative tokens. The cache must be + # bit-identical to its pre-run contents at those slots. + config = self._config(use_slot_mapping=True) + token_count = 2 + num_blocks = 4 + block_size = config.paged_kv_block_size + query = torch.randn(token_count, config.num_heads * config.head_size, device="cuda", dtype=torch.float16) + key = torch.randn(token_count, config.kv_num_heads * config.head_size, device="cuda", dtype=torch.float16) + value = torch.randn(token_count, config.kv_num_heads * config.head_size, device="cuda", dtype=torch.float16) + key_cache = torch.randn( + num_blocks, block_size, config.kv_num_heads, config.head_size, device="cuda", dtype=torch.float16 + ) + value_cache = torch.randn_like(key_cache) + key_cache_before = key_cache.clone().cpu().numpy() + value_cache_before = value_cache.clone().cpu().numpy() + + # One sequence of 2 new tokens on top of 1 cached token. + config.batch_size = 1 + cum_seqlens = torch.tensor([0, token_count], dtype=torch.int32, device="cuda") + past_seqlens = torch.tensor([1], dtype=torch.int32, device="cuda") + block_table = torch.tensor([[0, 1]], dtype=torch.int32, device="cuda") + # Store the first token at slot 1 of block 0; skip the second one entirely. + slot_mapping = torch.tensor([1, -1], dtype=torch.int32, device="cuda") + + _, key_cache_out, value_cache_out = paged_attention_func( + config, + query, + key, + value, + key_cache, + value_cache, + cum_seqlens, + past_seqlens, + block_table, + slot_mapping=slot_mapping, + ) + key_cache_out = numpy.array(key_cache_out) + value_cache_out = numpy.array(value_cache_out) + + # Slot 1 of block 0 holds the first token's K/V. + numpy.testing.assert_allclose( + key_cache_out[0, 1], key[0].reshape(config.kv_num_heads, config.head_size).cpu().numpy() + ) + numpy.testing.assert_allclose( + value_cache_out[0, 1], value[0].reshape(config.kv_num_heads, config.head_size).cpu().numpy() + ) + # Everything else, including the slot the second token would have used, is untouched. + key_cache_before[0, 1] = key_cache_out[0, 1] + value_cache_before[0, 1] = value_cache_out[0, 1] + numpy.testing.assert_array_equal(key_cache_out, key_cache_before) + numpy.testing.assert_array_equal(value_cache_out, value_cache_before) + + # ---- head_sink --------------------------------------------------------------------- + + @unittest.skipIf(not has_flash_attention(), reason="Flash Attention is not available") + def test_head_sink(self): + parity_check_paged_attention(self._config(use_head_sink=True), rtol=5e-3, atol=5e-3) + + @unittest.skipIf(not has_flash_attention(), reason="Flash Attention is not available") + def test_head_sink_local_and_softcap(self): + # The LSE epilogue must compose with sliding window and softcap, both of which are already + # baked into the log-sum-exp that FlashAttention returns. + config = self._config(use_head_sink=True, local=True, softcap=50.0) + parity_check_paged_attention(config, rtol=5e-3, atol=5e-3) + + @unittest.skipIf(not has_flash_attention(), reason="Flash Attention is not available") + def test_head_sink_with_rotary(self): + parity_check_paged_attention(self._config(use_head_sink=True, rotary=True), rtol=5e-3, atol=5e-3) + + @unittest.skipIf( + not has_memory_efficient_attention(), + reason="MemoryEfficientAttention (fp16) requires sm>=53", + ) + def test_head_sink_rejected_on_memory_efficient_path(self): + # The CUTLASS kernel does not expose a log-sum-exp, so the sink cannot be applied. The op + # must fail loudly rather than silently ignore the input. + with self.assertRaises(Exception) as ctx: + parity_check_paged_attention( + self._config(use_head_sink=True), + rtol=5e-3, + atol=5e-3, + sdpa_kernel=SDPA_KERNEL_EFFICIENT_ATTENTION, + ) + self.assertIn("head_sink", str(ctx.exception)) + + # ---- QK-Norm ----------------------------------------------------------------------- + + @unittest.skipIf(not has_flash_attention(), reason="Flash Attention is not available") + def test_qk_norm(self): + parity_check_paged_attention(self._config(use_qk_norm=True), rtol=5e-3, atol=5e-3) + + @unittest.skipIf(not has_flash_attention(), reason="Flash Attention is not available") + def test_qk_norm_with_rotary(self): + # QK-Norm is applied before rotary, and the normalized+rotated K is what lands in the cache. + # The cache parity assertions in parity_check_paged_attention cover that ordering. + parity_check_paged_attention(self._config(use_qk_norm=True, rotary=True), rtol=5e-3, atol=5e-3) + + @unittest.skipIf(not has_flash_attention(), reason="Flash Attention is not available") + def test_qk_norm_with_rotary_interleaved_and_packed(self): + config = self._config(use_qk_norm=True, rotary=True, rotary_interleaved=True, packed=True) + parity_check_paged_attention(config, rtol=5e-3, atol=5e-3) + + @unittest.skipIf(not has_flash_attention(), reason="Flash Attention is not available") + def test_qk_norm_non_power_of_two_head_size(self): + # head_size=80 rounds up to a 128-thread block in the fused prologue, so the lanes past + # head_size must contribute zero to the RMS reduction and skip all global accesses. + parity_check_paged_attention(self._config(use_qk_norm=True, head_size=80), rtol=5e-3, atol=5e-3) + + @unittest.skipIf( + not has_memory_efficient_attention(), + reason="MemoryEfficientAttention (fp16) requires sm>=53", + ) + def test_qk_norm_mea(self): + parity_check_paged_attention( + self._config(use_qk_norm=True), + rtol=5e-3, + atol=5e-3, + sdpa_kernel=SDPA_KERNEL_EFFICIENT_ATTENTION, + ) + + @unittest.skipIf(not has_flash_attention(), reason="Flash Attention is not available") + def test_qk_norm_and_head_sink_together(self): + parity_check_paged_attention( + self._config(use_qk_norm=True, use_head_sink=True, rotary=True, use_slot_mapping=True), + rtol=5e-3, + atol=5e-3, + ) + + # ---- lifted limits ----------------------------------------------------------------- + + @parameterized.expand([(16,), (32,), (64,), (128,)]) + @unittest.skipIf( + not has_memory_efficient_attention(), + reason="MemoryEfficientAttention (fp16) requires sm>=53", + ) + def test_small_block_size(self, block_size): + # block_size used to be required to be a multiple of 256. Smaller pages are now accepted; + # FlashAttention cannot address them (a kBlockN tile would straddle a page), so the kernel + # transparently falls back to the gather-based memory-efficient backend. + config = self._config(paged_kv_block_size=block_size, total_sequence_length=128) + parity_check_paged_attention(config, rtol=5e-3, atol=5e-3) + + @unittest.skipIf(not has_flash_attention(), reason="Flash Attention is not available") + def test_large_batch_size(self): + # cumulative_seqlens_kv used to be produced by independent 256-thread cub::BlockScan blocks, + # so any batch beyond 256 sequences got silently wrong KV offsets. + config = self._config(batch_size=300, sequence_length=4, total_sequence_length=64) + parity_check_paged_attention(config, rtol=5e-3, atol=5e-3) + + +@unittest.skipIf(not has_cuda_device(), reason="CUDA is not available, skipping tests.") class TestPagedAttentionRotaryZeroTokenRegression(unittest.TestCase): """Regression tests for the FA `max_query_len` heuristic when one or more batches have zero new tokens. @@ -884,5 +1403,1383 @@ def test_fa_kblockm_boundary_zero_token_no_rotary(self): parity_check_paged_attention(config, rtol=5e-3, atol=5e-3, new_seqlens_override=new_seqlens) +def has_fp8_kv_cache(): + """The float8e4m3fn PagedAttention kernels are only built when onnxruntime_USE_FP8_KV_CACHE is on.""" + if not hasattr(torch, "float8_e4m3fn") or not hasattr(TensorProto, "FLOAT8E4M3FN"): + return False + if not has_flash_attention(): + return False + config = Config(1, 1, 16, 1, 1, 64, 16, False, False, False, False, 0.0) + config.kv_cache_type = "fp8" + config.k_quant_type = "PER_TENSOR" + config.v_quant_type = "PER_TENSOR" + try: + InferenceSession( + create_paged_attention_graph(config, 1, 1, 1), + SessionOptions(), + providers=[config.ep], + ) + except Exception: + return False + return True + + +@unittest.skipIf(not has_flash_attention(), reason="Flash Attention is not available, skipping tests.") +class TestPagedAttentionQuantizedCache(unittest.TestCase): + """Coverage for the quantized paged KV cache (int8 / float8e4m3fn, PER_TENSOR / PER_CHANNEL). + + Both backends read a quantized cache through the dequantize-on-gather path, so these tests also + cover FlashAttention's non-paged varlen entry point, which is only reachable this way.""" + + def setUp(self): + # Quantization amplifies host/device rounding differences, so the inputs are re-seeded here + # to keep these tests independent of the order the rest of the module ran in. + torch.manual_seed(0) + + def _config(self, **overrides): + kwargs = { + "batch_size": 4, + "sequence_length": 33, + "total_sequence_length": 128, + "num_heads": 8, + "kv_num_heads": 2, + "head_size": 64, + "paged_kv_block_size": 256, + "local": False, + "rotary": False, + "rotary_interleaved": False, + "packed": False, + "softcap": 0.0, + } + feature_overrides = {k: overrides.pop(k) for k in list(overrides) if k not in kwargs} + kwargs.update(overrides) + config = Config(**kwargs) + for key, value in feature_overrides.items(): + setattr(config, key, value) + return config + + def _int8_config(self, quant_type="PER_TENSOR", **overrides): + return self._config(kv_cache_type="int8", k_quant_type=quant_type, v_quant_type=quant_type, **overrides) + + # ---- int8 --------------------------------------------------------------------------- + + @parameterized.expand([("per_tensor", "PER_TENSOR"), ("per_channel", "PER_CHANNEL")]) + def test_int8_cache(self, _, quant_type): + parity_check_paged_attention(self._int8_config(quant_type), rtol=5e-3, atol=5e-3) + + @parameterized.expand([("per_tensor", "PER_TENSOR"), ("per_channel", "PER_CHANNEL")]) + def test_int8_cache_with_rotary_and_packed(self, _, quant_type): + config = self._int8_config(quant_type, rotary=True, packed=True) + parity_check_paged_attention(config, rtol=5e-3, atol=5e-3) + + def test_int8_cache_with_qk_norm_and_slot_mapping(self): + # QK-Norm rescales K before it is written, and slot_mapping changes where it is written; + # both happen upstream of the quantization step in ReshapeAndCache. + config = self._int8_config("PER_CHANNEL", use_qk_norm=True, use_slot_mapping=True, rotary=True) + parity_check_paged_attention(config, rtol=5e-3, atol=5e-3) + + def test_int8_cache_local_and_softcap(self): + config = self._int8_config("PER_TENSOR", local=True, softcap=50.0) + parity_check_paged_attention(config, rtol=5e-3, atol=5e-3) + + def test_int8_cache_mixed_granularity(self): + # k_quant_type and v_quant_type are independent attributes. + config = self._config(kv_cache_type="int8", k_quant_type="PER_CHANNEL", v_quant_type="PER_TENSOR") + parity_check_paged_attention(config, rtol=5e-3, atol=5e-3) + + @parameterized.expand([("16", 16), ("32", 32), ("64", 64)]) + def test_int8_cache_small_block_size(self, _, block_size): + # A quantized cache never reaches FlashAttention's paged kernel, so the page-alignment + # constraint that forces small block sizes onto the MEA fallback does not apply here. + config = self._int8_config("PER_CHANNEL", paged_kv_block_size=block_size) + parity_check_paged_attention(config, rtol=5e-3, atol=5e-3) + + @unittest.skipIf( + not has_memory_efficient_attention(), + reason="MemoryEfficientAttention (fp16) requires sm>=53", + ) + def test_int8_cache_mea(self): + parity_check_paged_attention( + self._int8_config("PER_CHANNEL"), + rtol=5e-3, + atol=5e-3, + sdpa_kernel=SDPA_KERNEL_EFFICIENT_ATTENTION, + ) + + # ---- float8e4m3fn ------------------------------------------------------------------- + + @parameterized.expand([("per_tensor", "PER_TENSOR"), ("per_channel", "PER_CHANNEL")]) + @unittest.skipIf(not has_fp8_kv_cache(), reason="FP8 KV cache kernels are not built") + def test_fp8_cache(self, _, quant_type): + config = self._config(kv_cache_type="fp8", k_quant_type=quant_type, v_quant_type=quant_type) + parity_check_paged_attention(config, rtol=5e-3, atol=5e-3) + + @unittest.skipIf(not has_fp8_kv_cache(), reason="FP8 KV cache kernels are not built") + def test_fp8_cache_with_rotary_and_qk_norm(self): + config = self._config( + kv_cache_type="fp8", + k_quant_type="PER_CHANNEL", + v_quant_type="PER_CHANNEL", + rotary=True, + use_qk_norm=True, + ) + # Rotary and QK-Norm both diverge from the kernel by ~1 fp16 ULP, which is enough to move a + # value across an e4m3 rounding boundary; one such step is a ~12% change in a K element. + parity_check_paged_attention(config, rtol=5e-3, atol=2e-2) + + # ---- validation --------------------------------------------------------------------- + + def _run_minimal(self, config, k_scale=None, v_scale=None): + """Run PagedAttention on a trivially small input, bypassing the parity reference. Used to + check that invalid quantization configurations are rejected.""" + cache_torch_dtype = { + "float16": torch.float16, + "int8": torch.int8, + "fp8": getattr(torch, "float8_e4m3fn", None), + }[config.kv_cache_type] + cache_shape = (1, config.paged_kv_block_size, config.kv_num_heads, config.head_size) + key_cache = torch.zeros(cache_shape, dtype=torch.float16, device="cuda").to(cache_torch_dtype) + value_cache = torch.zeros(cache_shape, dtype=torch.float16, device="cuda").to(cache_torch_dtype) + num_tokens = config.sequence_length + query = torch.zeros((num_tokens, config.num_heads * config.head_size), dtype=torch.float16, device="cuda") + key = torch.zeros((num_tokens, config.kv_num_heads * config.head_size), dtype=torch.float16, device="cuda") + paged_attention_func( + config, + query, + key, + key.clone(), + key_cache, + value_cache, + torch.tensor([0, num_tokens], dtype=torch.int32, device="cuda"), + torch.zeros(config.batch_size, dtype=torch.int32, device="cuda"), + torch.zeros((config.batch_size, 1), dtype=torch.int32, device="cuda"), + k_scale=k_scale, + v_scale=v_scale, + ) + + def _minimal_config(self, **overrides): + return self._config( + batch_size=1, sequence_length=4, total_sequence_length=16, paged_kv_block_size=16, **overrides + ) + + def test_quantized_cache_without_quant_type_is_rejected(self): + config = self._minimal_config(kv_cache_type="int8") + with self.assertRaises(Exception) as ctx: + self._run_minimal(config) + self.assertIn("k_quant_type", str(ctx.exception)) + + def test_quant_type_without_quantized_cache_is_rejected(self): + config = self._minimal_config(k_quant_type="PER_TENSOR", v_quant_type="PER_TENSOR") + scale = torch.ones(1, dtype=torch.float32, device="cuda") + with self.assertRaises(Exception) as ctx: + self._run_minimal(config, k_scale=scale, v_scale=scale) + self.assertIn("not quantized", str(ctx.exception)) + + @parameterized.expand([("key", "k_cache_dtype"), ("value", "v_cache_dtype")]) + def test_unsupported_cache_dtype_is_rejected(self, _, attribute_name): + config = self._minimal_config(kv_cache_type="int8", k_quant_type="PER_TENSOR", v_quant_type="PER_TENSOR") + setattr(config, attribute_name, "int4") + scale = torch.ones(1, dtype=torch.float32, device="cuda") + with self.assertRaises(Exception) as ctx: + self._run_minimal(config, k_scale=scale, v_scale=scale) + self.assertIn(attribute_name, str(ctx.exception)) + + @parameterized.expand([("key", "k_cache_dtype"), ("value", "v_cache_dtype")]) + def test_cache_dtype_disagreeing_with_cache_tensor_is_rejected(self, _, attribute_name): + config = self._minimal_config(kv_cache_type="int8", k_quant_type="PER_TENSOR", v_quant_type="PER_TENSOR") + setattr(config, attribute_name, "float16") + scale = torch.ones(1, dtype=torch.float32, device="cuda") + with self.assertRaises(Exception) as ctx: + self._run_minimal(config, k_scale=scale, v_scale=scale) + self.assertIn(attribute_name, str(ctx.exception)) + + def test_cache_dtype_naming_the_cache_tensor_type_is_accepted(self): + # '' and an explicit spelling of the tensor's own element type mean the same thing. + config = self._minimal_config( + kv_cache_type="int8", + k_quant_type="PER_TENSOR", + v_quant_type="PER_TENSOR", + k_cache_dtype="int8", + v_cache_dtype="int8", + ) + scale = torch.ones(1, dtype=torch.float32, device="cuda") + self._run_minimal(config, k_scale=scale, v_scale=scale) + + @parameterized.expand([("key", "k_cache_dtype"), ("value", "v_cache_dtype")]) + def test_unknown_cache_dtype_is_rejected(self, _, attribute_name): + # "uint4" is deliberately not in the vocabulary: an unsigned logical type implies a zero + # point of 8, and this operator quantizes symmetrically with a scale only. + config = self._minimal_config(kv_cache_type="int8", k_quant_type="PER_TENSOR", v_quant_type="PER_TENSOR") + setattr(config, attribute_name, "uint4") + scale = torch.ones(1, dtype=torch.float32, device="cuda") + with self.assertRaises(Exception) as ctx: + self._run_minimal(config, k_scale=scale, v_scale=scale) + self.assertIn("Invalid KV cache data type", str(ctx.exception)) + + +@unittest.skipIf(not has_flash_attention(), reason="Flash Attention is not available, skipping tests.") +class TestPagedAttentionPagedDecode(unittest.TestCase): + """Coverage for the paged decode backend: a flash-decoding style kernel that scores the paged + KV cache in place (dequantizing in registers) instead of gathering it into a dense buffer. + + It is selected by the static shape test `token_count == batch_size`, which is a heuristic for + "one new token per sequence" rather than a proof, so the kernel has to stay correct for ragged + steps too (see the ragged cases below). The unquantized cases pin the backend with + `sdpa_kernel`; the quantized cases reach it through the normal auto-selection.""" + + def setUp(self): + torch.manual_seed(0) + + def _config(self, **overrides): + kwargs = { + "batch_size": 4, + "sequence_length": 1, + "total_sequence_length": 128, + "num_heads": 8, + "kv_num_heads": 2, + "head_size": 64, + "paged_kv_block_size": 256, + "local": False, + "rotary": False, + "rotary_interleaved": False, + "packed": False, + "softcap": 0.0, + } + feature_overrides = {k: overrides.pop(k) for k in list(overrides) if k not in kwargs} + kwargs.update(overrides) + config = Config(**kwargs) + for key, value in feature_overrides.items(): + setattr(config, key, value) + return config + + def _check_decode(self, config, rtol=5e-3, atol=5e-3, new_seqlens_override=None): + parity_check_paged_attention( + config, + rtol=rtol, + atol=atol, + sdpa_kernel=SDPA_KERNEL_DECODER_ATTENTION, + new_seqlens_override=new_seqlens_override, + ) + + # ---- shapes ------------------------------------------------------------------------- + + @parameterized.expand([("32", 32), ("64", 64), ("80", 80), ("96", 96), ("128", 128), ("256", 256)]) + def test_decode_head_size(self, _, head_size): + # head_size straddles the two PV thread mappings: one channel group when head_size >= 128 + # (the CTA width), several groups of head_size threads below it. + self._check_decode(self._config(head_size=head_size)) + + @parameterized.expand([("mha", 8, 8), ("gqa_4x", 8, 2), ("gqa_2x", 6, 3), ("mqa", 9, 1)]) + def test_decode_head_grouping(self, _, num_heads, kv_num_heads): + self._check_decode(self._config(num_heads=num_heads, kv_num_heads=kv_num_heads)) + + @parameterized.expand([("16", 16), ("32", 32), ("64", 64), ("256", 256), ("512", 512)]) + def test_decode_block_size(self, _, block_size): + # The decode kernel resolves a page per KV token, so unlike FlashAttention's paged kernel it + # has no page-alignment constraint and a KV tile may straddle any number of pages. + self._check_decode(self._config(paged_kv_block_size=block_size)) + + def test_decode_long_context_multi_tile(self): + # Far more than one 128-token tile per split, so the online-softmax rescaling across tiles + # is exercised rather than a single-shot tile. + self._check_decode(self._config(total_sequence_length=4000, paged_kv_block_size=256)) + + def test_decode_multi_split(self): + # One sequence and few heads leaves the GPU mostly idle, so the host picks num_splits > 1 + # and the cross-split reduction (rather than a single partial) produces the output. + self._check_decode(self._config(batch_size=1, num_heads=2, kv_num_heads=1, total_sequence_length=4000)) + + def test_decode_short_context(self): + # Fewer KV tokens than a tile, and short enough that some splits are empty. + self._check_decode(self._config(batch_size=1, num_heads=2, kv_num_heads=1, total_sequence_length=8)) + + def test_decode_zero_new_tokens(self): + # Sequences with no new token contribute no output row; the kernel must skip them without + # disturbing the rows of the sequences that do have one. + new_seqlens = torch.tensor([0, 1, 0, 1], dtype=torch.int32) + self._check_decode(self._config(), new_seqlens_override=new_seqlens) + + # ---- ragged steps ------------------------------------------------------------------- + # + # The host selects this backend from `token_count <= batch_size` alone, which does not prove one + # token per sequence. Every CTA therefore resolves its own sequence and in-sequence position + # from cumulative_sequence_length on device and masks against that token's own causal length. + # These cases are the ones a wrong implementation passes only by accident. + + def test_decode_ragged_shape_test_holds(self): + # token_count == batch_size == 4, but the tokens are distributed 3 / 0 / 0 / 1. + new_seqlens = torch.tensor([3, 0, 0, 1], dtype=torch.int32) + self._check_decode(self._config(sequence_length=3), new_seqlens_override=new_seqlens) + + def test_decode_ragged_all_tokens_in_one_sequence(self): + # The extreme case: one sequence owns the whole step, so per-token causal masking is the + # only thing that can produce the right answer. + new_seqlens = torch.tensor([4, 0, 0, 0], dtype=torch.int32) + self._check_decode(self._config(sequence_length=4), new_seqlens_override=new_seqlens) + + def test_decode_ragged_local_window(self): + new_seqlens = torch.tensor([3, 0, 0, 1], dtype=torch.int32) + self._check_decode(self._config(sequence_length=3, local=True), new_seqlens_override=new_seqlens) + + def test_decode_ragged_multi_split(self): + # Few (token, head) pairs and a long context force num_splits > 1, so the per-token split + # boundaries and the cross-split reduction are exercised on a ragged step. + new_seqlens = torch.tensor([2, 0], dtype=torch.int32) + self._check_decode( + self._config(batch_size=2, sequence_length=2, num_heads=2, kv_num_heads=1, total_sequence_length=4000), + new_seqlens_override=new_seqlens, + ) + + def test_decode_ragged_int8_cache(self): + # Auto-selected: a quantized decode-shaped step. XQA cannot serve it (its output layout is + # one row per batch index), so it must fall through to this kernel and still be correct. + new_seqlens = torch.tensor([3, 0, 0, 1], dtype=torch.int32) + config = self._config( + sequence_length=3, kv_cache_type="int8", k_quant_type="PER_CHANNEL", v_quant_type="PER_CHANNEL" + ) + parity_check_paged_attention(config, rtol=5e-3, atol=5e-3, new_seqlens_override=new_seqlens) + + # ---- masking and score transforms --------------------------------------------------- + + def test_decode_local_window(self): + self._check_decode(self._config(local=True)) + + def test_decode_softcap(self): + self._check_decode(self._config(softcap=50.0)) + + def test_decode_local_and_softcap(self): + self._check_decode(self._config(local=True, softcap=50.0)) + + def test_decode_head_sink(self): + self._check_decode(self._config(use_head_sink=True)) + + def test_decode_head_sink_local_and_softcap(self): + self._check_decode(self._config(use_head_sink=True, local=True, softcap=50.0)) + + # ---- prologue interaction ----------------------------------------------------------- + + def test_decode_rotary(self): + self._check_decode(self._config(rotary=True)) + + def test_decode_rotary_interleaved_and_packed(self): + self._check_decode(self._config(rotary=True, rotary_interleaved=True, packed=True)) + + def test_decode_qk_norm_and_slot_mapping(self): + self._check_decode(self._config(use_qk_norm=True, use_slot_mapping=True, rotary=True)) + + # ---- quantized cache (auto-selected, no sdpa_kernel override) ------------------------ + + @parameterized.expand([("per_tensor", "PER_TENSOR"), ("per_channel", "PER_CHANNEL")]) + def test_decode_int8_cache(self, _, quant_type): + # A quantized cache auto-selects the decode backend at sequence_length 1: the scales fold + # into Q and into the output, so the pages are read once at int8 width. + config = self._config(kv_cache_type="int8", k_quant_type=quant_type, v_quant_type=quant_type) + parity_check_paged_attention(config, rtol=5e-3, atol=5e-3) + + def test_decode_int8_cache_mixed_granularity(self): + config = self._config(kv_cache_type="int8", k_quant_type="PER_CHANNEL", v_quant_type="PER_TENSOR") + parity_check_paged_attention(config, rtol=5e-3, atol=5e-3) + + def test_decode_int8_cache_local_softcap_and_sink(self): + config = self._config( + kv_cache_type="int8", + k_quant_type="PER_CHANNEL", + v_quant_type="PER_CHANNEL", + local=True, + softcap=50.0, + use_head_sink=True, + ) + parity_check_paged_attention(config, rtol=5e-3, atol=5e-3) + + @parameterized.expand([("per_tensor", "PER_TENSOR"), ("per_channel", "PER_CHANNEL")]) + @unittest.skipIf(not has_fp8_kv_cache(), reason="FP8 KV cache kernels are not built") + def test_decode_fp8_cache(self, _, quant_type): + config = self._config(kv_cache_type="fp8", k_quant_type=quant_type, v_quant_type=quant_type) + parity_check_paged_attention(config, rtol=5e-3, atol=5e-3) + + @unittest.skipIf(not has_fp8_kv_cache(), reason="FP8 KV cache kernels are not built") + def test_decode_fp8_cache_with_rotary(self): + config = self._config(kv_cache_type="fp8", k_quant_type="PER_CHANNEL", v_quant_type="PER_CHANNEL", rotary=True) + parity_check_paged_attention(config, rtol=5e-3, atol=2e-2) + + +@unittest.skipIf(not has_cuda_device(), reason="CUDA is not available, skipping tests.") +class TestPagedAttentionXqaDecode(unittest.TestCase): + """Coverage for the XQA decode backend. + + XQA is the tensor-core decode kernel (shared with GroupQueryAttention) reading the paged cache + in place. It is auto-selected ahead of the portable decode kernel when the cache is quantized + and the step fits its constraints: exactly one new token per sequence, head_size in {64, 128}, + a query/KV group size in {4, 8, 16, 32}, no softcap, and block_size a multiple of 128 (a block + is presented to the kernel as several 128-token pages). Anything outside that falls back, which + is what the ORT_ENABLE_XQA=0 comparison below pins down. + + Every case here is also run through the fallback so a bug in XQA shows up as a parity failure + rather than being masked by both paths sharing the reference.""" + + def setUp(self): + torch.manual_seed(0) + if not has_fp8_kv_cache(): + self.skipTest("quantized KV cache kernels are not built") + + def _config(self, **overrides): + kwargs = { + "batch_size": 4, + "sequence_length": 1, + "total_sequence_length": 1024, + "num_heads": 8, + "kv_num_heads": 2, + "head_size": 64, + "paged_kv_block_size": 256, + "local": False, + "rotary": False, + "rotary_interleaved": False, + "packed": False, + "softcap": 0.0, + } + feature_overrides = {k: overrides.pop(k) for k in list(overrides) if k not in kwargs} + kwargs.update(overrides) + config = Config(**kwargs) + for key, value in feature_overrides.items(): + setattr(config, key, value) + return config + + def _check_xqa(self, quant_type="PER_TENSOR", kv_cache_type="int8", rtol=5e-3, atol=5e-3, **overrides): + config = self._config( + kv_cache_type=kv_cache_type, + k_quant_type=quant_type, + v_quant_type=quant_type, + **overrides, + ) + parity_check_paged_attention(config, rtol=rtol, atol=atol) + + # ---- shapes ------------------------------------------------------------------------- + + @parameterized.expand([("64", 64), ("128", 128)]) + def test_xqa_head_size(self, _, head_size): + self._check_xqa(head_size=head_size) + + @parameterized.expand([("grp4", 8, 2), ("grp8", 8, 1), ("grp16", 16, 1), ("grp32", 32, 1)]) + def test_xqa_head_grouping(self, _, num_heads, kv_num_heads): + self._check_xqa(num_heads=num_heads, kv_num_heads=kv_num_heads) + + @parameterized.expand([("128", 128), ("256", 256), ("512", 512)]) + def test_xqa_block_size(self, _, block_size): + # Each block is remapped to block_size / 128 consecutive XQA pages. + self._check_xqa(paged_kv_block_size=block_size) + + def test_xqa_batch_one(self): + self._check_xqa(batch_size=1) + + def test_xqa_long_context_multi_block(self): + # Long enough that XQA splits the sequence and reduces across CTAs through its scratch. + self._check_xqa(total_sequence_length=8192, batch_size=2) + + def test_xqa_short_context(self): + self._check_xqa(total_sequence_length=8, batch_size=1) + + def test_xqa_context_not_page_aligned(self): + # The live length is not a multiple of 128, so the last page is partially valid. + self._check_xqa(total_sequence_length=1000) + + # ---- quantization granularity ------------------------------------------------------- + + @parameterized.expand( + [ + ("int8_per_tensor", "int8", "PER_TENSOR"), + ("int8_per_channel", "int8", "PER_CHANNEL"), + ("fp8_per_tensor", "fp8", "PER_TENSOR"), + ("fp8_per_channel", "fp8", "PER_CHANNEL"), + ] + ) + def test_xqa_quant_type(self, _, kv_cache_type, quant_type): + self._check_xqa(kv_cache_type=kv_cache_type, quant_type=quant_type) + + def test_xqa_mixed_granularity(self): + # k PER_CHANNEL folds into Q, v PER_TENSOR stays a kernel argument: the two scales take + # different routes, so an asymmetric config catches a mix-up between them. + config = self._config(kv_cache_type="int8", k_quant_type="PER_CHANNEL", v_quant_type="PER_TENSOR") + parity_check_paged_attention(config, rtol=5e-3, atol=5e-3) + + # ---- masking and score transforms --------------------------------------------------- + + def test_xqa_local_window(self): + self._check_xqa(local=True) + + def test_xqa_head_sink(self): + self._check_xqa(use_head_sink=True) + + def test_xqa_head_sink_and_local(self): + self._check_xqa(use_head_sink=True, local=True) + + # ---- prologue interaction ----------------------------------------------------------- + + def test_xqa_rotary(self): + self._check_xqa(rotary=True, atol=2e-2) + + def test_xqa_rotary_interleaved_and_packed(self): + self._check_xqa(rotary=True, rotary_interleaved=True, packed=True, atol=2e-2) + + def test_xqa_qk_norm_and_slot_mapping(self): + self._check_xqa(use_qk_norm=True, use_slot_mapping=True, rotary=True, atol=2e-2) + + # ---- fallback ----------------------------------------------------------------------- + + def test_softcap_falls_back(self): + # XQA has no softcap, so this must land on the portable decode kernel and still be correct. + self._check_xqa(softcap=50.0) + + def test_unsupported_head_size_falls_back(self): + self._check_xqa(head_size=96) + + def test_unsupported_block_size_falls_back(self): + self._check_xqa(paged_kv_block_size=64) + + def test_multi_token_step_falls_back(self): + # More than one new token in a sequence: XQA emits one row per sequence, so this has to use + # a backend that handles a ragged step. + config = self._config( + sequence_length=2, kv_cache_type="int8", k_quant_type="PER_TENSOR", v_quant_type="PER_TENSOR" + ) + parity_check_paged_attention(config, rtol=5e-3, atol=5e-3) + + +@unittest.skipIf(not has_cuda_device(), reason="CUDA is not available, skipping tests.") +class TestPagedAttentionAttentionMetadata(unittest.TestCase): + """Coverage for the optional 'attention_metadata' input. + + 'cumulative_sequence_length' and 'past_seqlens' are device tensors, but the op needs an upper + bound on the query length (to size grids) and on the KV length (to size workspaces) on the host. + Without this input it falls back to the static capacities, and in two narrow prefill-side cases + -- a dense gather, or XQA, which needs proof of exactly one token per sequence -- it copies the + cumulative arrays back and blocks the stream instead. 'attention_metadata' supplies both bounds + in CPU memory so neither ever happens. + + It is purely an optimization, so the contract under test is that it changes nothing: every case + here mirrors a config that is also covered without the input and must produce the same results. + The cases span the backends because they consume the bounds differently.""" + + def setUp(self): + torch.manual_seed(0) + + def _config(self, **overrides): + kwargs = { + "batch_size": 4, + "sequence_length": 1, + "total_sequence_length": 1024, + "num_heads": 8, + "kv_num_heads": 2, + "head_size": 64, + "paged_kv_block_size": 256, + "local": False, + "rotary": False, + "rotary_interleaved": False, + "packed": False, + "softcap": 0.0, + } + feature_overrides = {k: overrides.pop(k) for k in list(overrides) if k not in kwargs} + kwargs.update(overrides) + config = Config(**kwargs) + config.use_attention_metadata = True + for key, value in feature_overrides.items(): + setattr(config, key, value) + return config + + def test_prefill_flash_attention(self): + # A prefill bound is > 1, so it only sizes the FlashAttention grid; the input must not + # perturb anything there either. + parity_check_paged_attention(self._config(sequence_length=256, total_sequence_length=256)) + + def test_decode_unquantized(self): + parity_check_paged_attention(self._config(), sdpa_kernel=SDPA_KERNEL_DECODER_ATTENTION) + + def test_decode_xqa_int8(self): + config = self._config(kv_cache_type="int8", k_quant_type="PER_CHANNEL", v_quant_type="PER_CHANNEL") + parity_check_paged_attention(config, rtol=5e-3, atol=5e-3) + + @unittest.skipIf(not has_fp8_kv_cache(), reason="FP8 KV cache kernels are not built") + def test_decode_xqa_fp8(self): + config = self._config(kv_cache_type="fp8", k_quant_type="PER_TENSOR", v_quant_type="PER_TENSOR") + parity_check_paged_attention(config, rtol=5e-2, atol=5e-2) + + def test_decode_softcap_portable_kernel(self): + # softcap rules out XQA, so this exercises the split-KV decode kernel, whose split count is + # derived from max_kv_len -- here from the bound rather than from a readback. + config = self._config(kv_cache_type="int8", k_quant_type="PER_TENSOR", v_quant_type="PER_TENSOR", softcap=50.0) + parity_check_paged_attention(config, rtol=5e-3, atol=5e-3) + + def test_decode_local_window(self): + config = self._config(local=True, kv_cache_type="int8", k_quant_type="PER_TENSOR", v_quant_type="PER_TENSOR") + parity_check_paged_attention(config, rtol=5e-3, atol=5e-3) + + def test_chunked_prefill_dense_gather(self): + # A quantized cache on a multi-token step gathers the live context into a dense buffer. With + # a bound available that buffer is sized by batch_size * max_kv_len_bound instead of the + # exact total_kv_tokens, so the gather kernel has to tolerate indices past the real end of + # the packed layout. + config = self._config( + sequence_length=8, + kv_cache_type="int8", + k_quant_type="PER_TENSOR", + v_quant_type="PER_TENSOR", + ) + parity_check_paged_attention(config, rtol=5e-3, atol=5e-3) + + def test_ragged_new_token_counts(self): + # token_count == batch_size even though one sequence contributes two tokens and another + # none. The shape test alone selects a decode-shaped backend, so every backend it can reach + # must handle the raggedness. + new_seqlens = torch.tensor([2, 0, 1, 1], dtype=torch.int32) + parity_check_paged_attention(self._config(sequence_length=2), new_seqlens_override=new_seqlens) + + def test_ragged_new_token_counts_decode_kernel(self): + # Same step, pinned to the paged decode kernel: with a bound of 2 there is no readback, so + # the kernel is the only thing that knows the tokens are unevenly distributed. + new_seqlens = torch.tensor([2, 0, 1, 1], dtype=torch.int32) + parity_check_paged_attention( + self._config(sequence_length=2), + new_seqlens_override=new_seqlens, + sdpa_kernel=SDPA_KERNEL_DECODER_ATTENTION, + rtol=5e-3, + atol=5e-3, + ) + + def test_zero_new_tokens_in_batch(self): + # token_count < batch_size: still decode-shaped, but XQA is excluded because it would emit a + # row for the sequence that contributed nothing. + new_seqlens = torch.tensor([1, 0, 1, 1], dtype=torch.int32) + parity_check_paged_attention(self._config(), new_seqlens_override=new_seqlens) + + def test_batch_one(self): + parity_check_paged_attention(self._config(batch_size=1), sdpa_kernel=SDPA_KERNEL_DECODER_ATTENTION) + + def test_negative_entry_is_rejected(self): + config = self._config() + config.attention_metadata_override = numpy.array([-1, 1024], dtype=numpy.int32) + with self.assertRaises(Exception) as ctx: + parity_check_paged_attention(config) + self.assertIn("must be non-negative", str(ctx.exception)) + + def test_unknown_bounds_fall_back_to_readback(self): + # All-zero means "no bound", which must behave exactly like supplying the static capacities. + config = self._config(kv_cache_type="int8", k_quant_type="PER_TENSOR", v_quant_type="PER_TENSOR") + config.attention_metadata_override = numpy.array([0, 0], dtype=numpy.int32) + parity_check_paged_attention(config, rtol=5e-3, atol=5e-3) + + def test_over_large_bounds_are_clamped(self): + # Bounds beyond the static limits are legal (they are still upper bounds) and must be + # clamped rather than used to size an allocation. + config = self._config(kv_cache_type="int8", k_quant_type="PER_TENSOR", v_quant_type="PER_TENSOR") + config.attention_metadata_override = numpy.array([1 << 20, 1 << 20], dtype=numpy.int32) + parity_check_paged_attention(config, rtol=5e-3, atol=5e-3) + + +#################################################################################################### +# Multi-head Latent Attention (kv_cache_layout="LATENT") +# +# See docs/contrib_ops/cuda/paged_attention.md §12. In the absorbed form there is a single physical +# cache holding the latent row [compressed_kv | k_pe] of width head_size = kv_lora_rank + +# qk_rope_head_dim. K is the whole row; V of every head is its leading v_head_size = kv_lora_rank +# channels. 'value' and 'value_cache' are therefore absent and kv_num_heads is 1. +#################################################################################################### + + +class MLAConfig: + """Shape bundle for a LATENT-layout PagedAttention node. Deliberately separate from Config, + whose fields (kv_num_heads, packed, ...) encode SEPARATE-mode assumptions.""" + + def __init__( + self, + batch_size, + num_heads, + kv_lora_rank, + qk_rope_head_dim, + block_size, + qk_nope_head_dim=None, + kv_cache_type="float16", + k_quant_type="NONE", + ): + self.batch_size = batch_size + self.num_heads = num_heads + self.kv_lora_rank = kv_lora_rank + self.qk_rope_head_dim = qk_rope_head_dim + # Absorbed geometry, i.e. what the op actually sees. + self.head_size = kv_lora_rank + qk_rope_head_dim + self.v_head_size = kv_lora_rank + self.rotary_offset = kv_lora_rank + self.kv_num_heads = 1 + self.block_size = block_size + # Un-absorbed geometry, used only by the equivalence test. + self.qk_nope_head_dim = qk_nope_head_dim + self.kv_cache_type = kv_cache_type + self.k_quant_type = k_quant_type + + @property + def softmax_scale(self): + # DeepSeek scales by the pre-absorption QK width, which is nope + rope, NOT the absorbed + # head_size. This is exactly why the op requires an explicit 'scale' for MLA. + width = (self.qk_nope_head_dim or self.kv_lora_rank) + self.qk_rope_head_dim + return width**-0.5 + + +def create_mla_graph( + mla_config, + num_tokens, + num_blocks, + max_blocks_per_sequence, + do_rotary=False, + rotary_interleaved=False, + rotary_offset=None, + scale=None, + local_window_size=-1, + softcap=0.0, + with_value=False, + with_value_cache=False, + with_value_cache_out=False, + with_head_sink=False, + with_qk_norm=False, + kv_cache_layout="LATENT", + v_head_size=None, +): + """Build a single-node LATENT PagedAttention model. Every deviation from a valid MLA graph is a + keyword here so that the negative tests can construct rejected models.""" + cache_proto_type = KV_CACHE_TENSOR_PROTO[mla_config.kv_cache_type] + head_size = mla_config.head_size + v_head_size = mla_config.v_head_size if v_head_size is None else v_head_size + rotary_offset = mla_config.rotary_offset if rotary_offset is None else rotary_offset + rotary_dim = mla_config.qk_rope_head_dim + + attrs = { + "num_heads": mla_config.num_heads, + "kv_num_heads": mla_config.kv_num_heads, + "kv_cache_layout": kv_cache_layout, + "v_head_size": v_head_size, + "local_window_size": local_window_size, + "softcap": softcap, + "domain": "com.microsoft", + } + if scale is not None: + attrs["scale"] = scale + if do_rotary: + attrs["do_rotary"] = 1 + attrs["rotary_interleaved"] = 1 if rotary_interleaved else 0 + attrs["rotary_offset"] = rotary_offset + if mla_config.k_quant_type != "NONE": + attrs["k_quant_type"] = mla_config.k_quant_type + + node_outputs = ["output", "key_cache_out"] + if with_value_cache_out: + node_outputs.append("value_cache_out") + + nodes = [ + helper.make_node( + "PagedAttention", + [ + "query", + "key", + "value" if with_value else "", + "key_cache", + "value_cache" if with_value_cache else "", + "cumulative_sequence_length", + "past_seqlens", + "block_table", + "cos_cache" if do_rotary else "", + "sin_cache" if do_rotary else "", + "", # slot_mapping + "head_sink" if with_head_sink else "", + "q_norm_weight" if with_qk_norm else "", + "k_norm_weight" if with_qk_norm else "", + "k_scale" if mla_config.k_quant_type != "NONE" else "", + ], + node_outputs, + "PagedAttention_MLA", + **attrs, + ), + ] + + cache_dims = [num_blocks, mla_config.block_size, mla_config.kv_num_heads, head_size] + graph_input = [ + helper.make_tensor_value_info("query", TensorProto.FLOAT16, [num_tokens, mla_config.num_heads * head_size]), + helper.make_tensor_value_info("key", TensorProto.FLOAT16, [num_tokens, mla_config.kv_num_heads * head_size]), + helper.make_tensor_value_info("key_cache", cache_proto_type, cache_dims), + helper.make_tensor_value_info("cumulative_sequence_length", TensorProto.INT32, [mla_config.batch_size + 1]), + helper.make_tensor_value_info("past_seqlens", TensorProto.INT32, [mla_config.batch_size]), + helper.make_tensor_value_info( + "block_table", TensorProto.INT32, [mla_config.batch_size, max_blocks_per_sequence] + ), + ] + if with_value: + # SEPARATE mode requires value to match key's hidden size; in LATENT its mere presence is + # the violation under test. + graph_input.append( + helper.make_tensor_value_info( + "value", TensorProto.FLOAT16, [num_tokens, mla_config.kv_num_heads * head_size] + ) + ) + if with_value_cache: + graph_input.append(helper.make_tensor_value_info("value_cache", cache_proto_type, cache_dims)) + if do_rotary: + # The rotary caches are indexed by rotary_dim // 2, and the op derives rotary_dim from their + # width. MLA rotates only the qk_rope_head_dim suffix, so these are much narrower than the + # head_size-derived caches a full-width RoPE would use. + cache_width = rotary_dim // 2 + graph_input += [ + helper.make_tensor_value_info("cos_cache", TensorProto.FLOAT16, [None, cache_width]), + helper.make_tensor_value_info("sin_cache", TensorProto.FLOAT16, [None, cache_width]), + ] + if with_head_sink: + graph_input.append(helper.make_tensor_value_info("head_sink", TensorProto.FLOAT16, [mla_config.num_heads])) + if with_qk_norm: + graph_input += [ + helper.make_tensor_value_info("q_norm_weight", TensorProto.FLOAT16, [head_size]), + helper.make_tensor_value_info("k_norm_weight", TensorProto.FLOAT16, [head_size]), + ] + if mla_config.k_quant_type != "NONE": + scale_shape = [1] if mla_config.k_quant_type == "PER_TENSOR" else [mla_config.kv_num_heads, 1, head_size] + graph_input.append(helper.make_tensor_value_info("k_scale", TensorProto.FLOAT, scale_shape)) + + graph_output = [ + helper.make_tensor_value_info("output", TensorProto.FLOAT16, [num_tokens, mla_config.num_heads * v_head_size]), + helper.make_tensor_value_info("key_cache_out", cache_proto_type, cache_dims), + ] + if with_value_cache_out: + graph_output.append(helper.make_tensor_value_info("value_cache_out", cache_proto_type, cache_dims)) + + graph = helper.make_graph(nodes, "PagedAttention_MLA_Graph", graph_input, graph_output) + return helper.make_model(graph).SerializeToString() + + +def run_mla( + mla_config, + query, + key, + key_cache, + cumulative_sequence_length, + past_seqlens, + block_table, + cos=None, + sin=None, + k_scale=None, + **graph_kwargs, +): + """Run a LATENT PagedAttention model and return (output, key_cache) with the cache updated in + place. key_cache is bound on device so the in-place scatter is observable.""" + num_tokens = int(cumulative_sequence_length[-1].item()) + onnx_model_str = create_mla_graph( + mla_config, + num_tokens, + key_cache.shape[0], + block_table.shape[1], + do_rotary=cos is not None, + **graph_kwargs, + ) + ort_session = InferenceSession(onnx_model_str, SessionOptions(), providers=["CUDAExecutionProvider"]) + io_binding = ort_session.io_binding() + + io_binding.bind_cpu_input("query", query.detach().cpu().numpy()) + io_binding.bind_cpu_input("key", key.detach().cpu().numpy()) + io_binding.bind_cpu_input("cumulative_sequence_length", cumulative_sequence_length.detach().cpu().numpy()) + io_binding.bind_cpu_input("past_seqlens", past_seqlens.detach().cpu().numpy()) + io_binding.bind_cpu_input("block_table", block_table.detach().cpu().numpy()) + if cos is not None: + io_binding.bind_cpu_input("cos_cache", cos.detach().cpu().numpy()) + io_binding.bind_cpu_input("sin_cache", sin.detach().cpu().numpy()) + if k_scale is not None: + io_binding.bind_cpu_input("k_scale", k_scale.detach().cpu().numpy()) + + cache_proto_type = KV_CACHE_TENSOR_PROTO[mla_config.kv_cache_type] + key_cache = key_cache.contiguous() + io_binding.bind_input("key_cache", "cuda", 0, cache_proto_type, tuple(key_cache.shape), key_cache.data_ptr()) + io_binding.bind_output("output") + io_binding.bind_output("key_cache_out", "cuda", 0, cache_proto_type, tuple(key_cache.shape), key_cache.data_ptr()) + ort_session.run_with_iobinding(io_binding) + output = torch.tensor(numpy.array(io_binding.copy_outputs_to_cpu()[0])) + return output, key_cache + + +def mla_reference( + mla_config, + query, # [token_count, num_heads, head_size] + latent_cache, # [batch, total_seqlen, head_size] dense view of the paged cache + past_seqlens, + new_seqlens, + cum_seqlens, + scale, + local_window_size=-1, + softcap=0.0, +): + """Straightforward fp32 MLA: K is the whole latent row, V its leading v_head_size channels.""" + v_head_size = mla_config.v_head_size + token_count = int(cum_seqlens[-1].item()) + out = torch.zeros(token_count, mla_config.num_heads, v_head_size, dtype=torch.float32, device="cuda") + q = query.to(torch.float32) + for b in range(mla_config.batch_size): + start = int(cum_seqlens[b].item()) + for j in range(int(new_seqlens[b].item())): + kv_end = int(past_seqlens[b].item()) + j + 1 + kv_begin = max(0, kv_end - local_window_size) if local_window_size > 0 else 0 + k = latent_cache[b, kv_begin:kv_end].to(torch.float32) # [L, head_size] + v = k[:, :v_head_size] + logits = torch.einsum("nh,lh->nl", q[start + j], k) * scale + if softcap > 0.0: + logits = softcap * torch.tanh(logits / softcap) + probs = torch.softmax(logits, dim=-1) + out[start + j] = torch.einsum("nl,lv->nv", probs, v) + return out + + +def make_mla_batch(mla_config, past_seqlens, new_seqlens, device="cuda"): + """Allocate a shuffled paged latent cache pre-filled with the 'past' tokens, plus the block + table and cumulative sequence lengths. Returns everything both paged and densified.""" + total_seqlens = past_seqlens + new_seqlens + max_total = int(total_seqlens.max().item()) + blocks_per_seq = math.ceil(max_total / mla_config.block_size) + num_blocks = blocks_per_seq * mla_config.batch_size + # A shuffled permutation makes block-table indirection load-bearing: a kernel that ignored it + # and read blocks sequentially would fail. + block_table = torch.randperm(num_blocks, dtype=torch.int32, device=device).reshape( + mla_config.batch_size, blocks_per_seq + ) + latent_paged = torch.randn( + num_blocks, + mla_config.block_size, + mla_config.kv_num_heads, + mla_config.head_size, + device=device, + dtype=torch.float16, + ) + cum_seqlens = torch.zeros(mla_config.batch_size + 1, dtype=torch.int32, device=device) + cum_seqlens[1:] = torch.cumsum(new_seqlens, dim=0) + return latent_paged, block_table, cum_seqlens, blocks_per_seq * mla_config.block_size + + +def densify_latent(mla_config, latent_paged, block_table, total_len): + return rearrange( + latent_paged[block_table.to(dtype=torch.long).flatten()], + "(b nblocks) block_size h d -> b (nblocks block_size) (h d)", + b=mla_config.batch_size, + )[:, :total_len] + + +def apply_offset_rope(x, cos, sin, positions, rotary_offset, rotary_dim, interleaved): + """Reference for the 'rotary_offset' attribute: rotate x[..., offset:offset+rotary_dim] using + the same half-rotated / interleaved conventions as the op, copying everything else through.""" + out = x.clone().to(torch.float32) + seg = out[..., rotary_offset : rotary_offset + rotary_dim] + c = cos[positions].to(torch.float32) # [tokens, rotary_dim // 2] + s = sin[positions].to(torch.float32) + c = c[..., : rotary_dim // 2] + s = s[..., : rotary_dim // 2] + while c.dim() < seg.dim(): + c = c.unsqueeze(-2) + s = s.unsqueeze(-2) + if interleaved: + even = seg[..., 0::2] + odd = seg[..., 1::2] + rotated = torch.stack([even * c - odd * s, even * s + odd * c], dim=-1).flatten(-2) + else: + half = rotary_dim // 2 + x1 = seg[..., :half] + x2 = seg[..., half:] + rotated = torch.cat([x1 * c - x2 * s, x1 * s + x2 * c], dim=-1) + out[..., rotary_offset : rotary_offset + rotary_dim] = rotated + return out.to(x.dtype) + + +@unittest.skipIf(not has_cuda_device(), reason="CUDA is not available, skipping tests.") +class TestPagedAttentionMLA(unittest.TestCase): + """Correctness of kv_cache_layout='LATENT' (design doc §12, phase P4).""" + + def setUp(self): + # These tests build random weights; a fixed seed keeps them order-independent. + torch.manual_seed(20240727) + + def _config(self, **kwargs): + # Small but structurally faithful: kv_lora_rank and qk_rope_head_dim keep DeepSeek's ratio + # while staying cheap enough for a unit test. + defaults = dict(batch_size=2, num_heads=4, kv_lora_rank=64, qk_rope_head_dim=32, block_size=16) + defaults.update(kwargs) + return MLAConfig(**defaults) + + def _run_case(self, mla_config, past_seqlens, new_seqlens, local_window_size=-1, softcap=0.0): + """Shared body: build a paged latent cache, run the op, compare against mla_reference.""" + device = "cuda" + past_seqlens = torch.tensor(past_seqlens, dtype=torch.int32, device=device) + new_seqlens = torch.tensor(new_seqlens, dtype=torch.int32, device=device) + latent_paged, block_table, cum_seqlens, total_len = make_mla_batch(mla_config, past_seqlens, new_seqlens) + token_count = int(cum_seqlens[-1].item()) + + query = torch.randn(token_count, mla_config.num_heads, mla_config.head_size, device=device, dtype=torch.float16) + new_key = torch.randn(token_count, mla_config.head_size, device=device, dtype=torch.float16) + + scale = mla_config.softmax_scale + out, latent_paged = run_mla( + mla_config, + query.reshape(token_count, -1), + new_key, + latent_paged, + cum_seqlens, + past_seqlens, + block_table, + scale=scale, + local_window_size=local_window_size, + softcap=softcap, + ) + + # The op scattered the new keys into the cache in place, so densifying afterwards gives the + # exact K/V the reference must see (including any quantization error). + dense = densify_latent(mla_config, latent_paged, block_table, total_len) + ref = mla_reference( + mla_config, + query, + dense, + past_seqlens, + new_seqlens, + cum_seqlens, + scale, + local_window_size=local_window_size, + softcap=softcap, + ) + out = out.reshape(token_count, mla_config.num_heads, mla_config.v_head_size).to(device).to(torch.float32) + torch.testing.assert_close(out, ref, rtol=2e-3, atol=2e-3) + return out + + def test_prefill(self): + config = self._config() + self._run_case(config, past_seqlens=[0, 0], new_seqlens=[13, 7]) + + def test_decode(self): + config = self._config() + self._run_case(config, past_seqlens=[31, 18], new_seqlens=[1, 1]) + + def test_chunked_prefill(self): + # Mixed batch: one sequence extends an existing cache, one is pure decode, one adds nothing. + config = self._config(batch_size=3) + self._run_case(config, past_seqlens=[20, 9, 5], new_seqlens=[6, 1, 0]) + + def test_local_window(self): + config = self._config() + self._run_case(config, past_seqlens=[24, 24], new_seqlens=[5, 5], local_window_size=8) + + def test_softcap(self): + config = self._config() + self._run_case(config, past_seqlens=[12, 12], new_seqlens=[4, 4], softcap=30.0) + + def test_deepseek_v3_geometry(self): + # The real absorbed shape: head_size 576, v_head_size 512, kv_num_heads 1. + config = self._config(num_heads=8, kv_lora_rank=512, qk_rope_head_dim=64, block_size=16) + self.assertEqual(config.head_size, 576) + self.assertEqual(config.v_head_size, 512) + self._run_case(config, past_seqlens=[17, 3], new_seqlens=[1, 4]) + + def test_absorbed_matches_non_absorbed(self): + """The point of MLA: attention over the latent row with absorbed projections equals + standard MHA over the up-projected K/V. Both sides run through PagedAttention, so this also + pins the LATENT path against the already-verified SEPARATE path.""" + device = "cuda" + batch_size, num_heads = 2, 4 + kv_lora_rank, qk_rope_head_dim = 64, 32 + qk_nope_head_dim, v_head_dim = 32, 32 + mla_config = MLAConfig( + batch_size, + num_heads, + kv_lora_rank, + qk_rope_head_dim, + block_size=16, + qk_nope_head_dim=qk_nope_head_dim, + ) + past_seqlens = torch.tensor([9, 5], dtype=torch.int32, device=device) + new_seqlens = torch.tensor([4, 6], dtype=torch.int32, device=device) + latent_paged, block_table, cum_seqlens, total_len = make_mla_batch(mla_config, past_seqlens, new_seqlens) + token_count = int(cum_seqlens[-1].item()) + scale = mla_config.softmax_scale + + # Up-projection weights, shared by both spellings. + w_uk = torch.randn(kv_lora_rank, num_heads, qk_nope_head_dim, device=device, dtype=torch.float32) * 0.1 + w_uv = torch.randn(kv_lora_rank, num_heads, v_head_dim, device=device, dtype=torch.float32) * 0.1 + + # Non-absorbed query: q_nope [N, qk_nope_head_dim] and q_pe [N, qk_rope_head_dim]. + q_nope = torch.randn(token_count, num_heads, qk_nope_head_dim, device=device, dtype=torch.float32) * 0.5 + q_pe = torch.randn(token_count, num_heads, qk_rope_head_dim, device=device, dtype=torch.float32) * 0.5 + new_key = torch.randn(token_count, mla_config.head_size, device=device, dtype=torch.float16) + + # --- absorbed: q_latent = [q_nope @ W_UK^T | q_pe], attention over the latent row --- + q_absorbed_nope = torch.einsum("tnp,cnp->tnc", q_nope, w_uk) # [T, N, kv_lora_rank] + q_absorbed = torch.cat([q_absorbed_nope, q_pe], dim=-1).to(torch.float16) + out_absorbed, latent_paged = run_mla( + mla_config, + q_absorbed.reshape(token_count, -1), + new_key, + latent_paged, + cum_seqlens, + past_seqlens, + block_table, + scale=scale, + ) + out_absorbed = out_absorbed.reshape(token_count, num_heads, kv_lora_rank).to(device).to(torch.float32) + # Absorbed output lives in latent space; project it out with W_UV to compare. + out_absorbed = torch.einsum("tnc,cnv->tnv", out_absorbed, w_uv) + + # --- non-absorbed: up-project the cache into per-head K/V and run plain MHA --- + dense = densify_latent(mla_config, latent_paged, block_table, total_len).to(torch.float32) + compressed_kv = dense[..., :kv_lora_rank] # [B, L, kv_lora_rank] + k_pe = dense[..., kv_lora_rank:] # [B, L, qk_rope_head_dim] + k_nope = torch.einsum("blc,cnp->blnp", compressed_kv, w_uk) + k_full = torch.cat([k_nope, k_pe.unsqueeze(2).expand(-1, -1, num_heads, -1)], dim=-1) + v_full = torch.einsum("blc,cnv->blnv", compressed_kv, w_uv) + q_full = torch.cat([q_nope, q_pe], dim=-1) + + out_ref = torch.zeros(token_count, num_heads, v_head_dim, dtype=torch.float32, device=device) + for b in range(batch_size): + start = int(cum_seqlens[b].item()) + for j in range(int(new_seqlens[b].item())): + kv_end = int(past_seqlens[b].item()) + j + 1 + logits = torch.einsum("nd,lnd->nl", q_full[start + j], k_full[b, :kv_end]) * scale + probs = torch.softmax(logits, dim=-1) + out_ref[start + j] = torch.einsum("nl,lnv->nv", probs, v_full[b, :kv_end]) + + torch.testing.assert_close(out_absorbed, out_ref, rtol=5e-3, atol=5e-3) + + def test_rotary_offset_matches_graph_applied_rope(self): + """do_rotary=1 with rotary_offset=kv_lora_rank must equal applying the same RoPE to the + k_pe suffix outside the op and running with do_rotary=0.""" + device = "cuda" + for interleaved in (False, True): + with self.subTest(interleaved=interleaved): + config = self._config() + past_seqlens = torch.tensor([6, 2], dtype=torch.int32, device=device) + new_seqlens = torch.tensor([3, 5], dtype=torch.int32, device=device) + latent_paged, block_table, cum_seqlens, total_len = make_mla_batch(config, past_seqlens, new_seqlens) + token_count = int(cum_seqlens[-1].item()) + rotary_dim = config.qk_rope_head_dim + cache_width = rotary_dim // 2 + max_pos = 128 + angle = torch.rand(max_pos, cache_width, device=device) * 2 - 1 + cos = torch.cos(angle).to(torch.float16) + sin = torch.sin(angle).to(torch.float16) + + query = torch.randn(token_count, config.num_heads, config.head_size, device=device, dtype=torch.float16) + new_key = torch.randn(token_count, config.head_size, device=device, dtype=torch.float16) + scale = config.softmax_scale + + # Positions the op uses: past_seqlens[b] + index within the sequence. + positions = torch.empty(token_count, dtype=torch.long, device=device) + for b in range(config.batch_size): + start = int(cum_seqlens[b].item()) + n = int(new_seqlens[b].item()) + positions[start : start + n] = int(past_seqlens[b].item()) + torch.arange(n, device=device) + + out_in_op, cache_in_op = run_mla( + config, + query.reshape(token_count, -1), + new_key, + latent_paged.clone(), + cum_seqlens, + past_seqlens, + block_table, + cos=cos, + sin=sin, + rotary_interleaved=interleaved, + scale=scale, + ) + + q_roped = apply_offset_rope(query, cos, sin, positions, config.rotary_offset, rotary_dim, interleaved) + k_roped = apply_offset_rope(new_key, cos, sin, positions, config.rotary_offset, rotary_dim, interleaved) + out_pre, cache_pre = run_mla( + config, + q_roped.reshape(token_count, -1), + k_roped, + latent_paged.clone(), + cum_seqlens, + past_seqlens, + block_table, + scale=scale, + ) + + torch.testing.assert_close(out_in_op.to(torch.float32), out_pre.to(torch.float32), rtol=3e-3, atol=3e-3) + # The scattered latent rows must match too: RoPE touches only the k_pe suffix. + torch.testing.assert_close( + cache_in_op.to(torch.float32), cache_pre.to(torch.float32), rtol=3e-3, atol=3e-3 + ) + + def test_v_aliases_k(self): + """With value_cache absent, V must be the leading v_head_size channels of key_cache. Zero + the tail of every latent row and confirm the output is unchanged: the tail is K-only.""" + device = "cuda" + config = self._config() + past_seqlens = torch.tensor([0, 0], dtype=torch.int32, device=device) + new_seqlens = torch.tensor([4, 4], dtype=torch.int32, device=device) + latent_paged, block_table, cum_seqlens, total_len = make_mla_batch(config, past_seqlens, new_seqlens) + token_count = int(cum_seqlens[-1].item()) + query = torch.randn(token_count, config.num_heads, config.head_size, device=device, dtype=torch.float16) + # Zero the query's rope slice so the k_pe channels cannot influence the logits either. + query[..., config.v_head_size :] = 0 + new_key = torch.randn(token_count, config.head_size, device=device, dtype=torch.float16) + scale = config.softmax_scale + + out_a, cache_a = run_mla( + config, + query.reshape(token_count, -1), + new_key, + latent_paged.clone(), + cum_seqlens, + past_seqlens, + block_table, + scale=scale, + ) + key_tail_zeroed = new_key.clone() + key_tail_zeroed[:, config.v_head_size :] = 0 + paged_tail_zeroed = latent_paged.clone() + paged_tail_zeroed[..., config.v_head_size :] = 0 + out_b, _ = run_mla( + config, + query.reshape(token_count, -1), + key_tail_zeroed, + paged_tail_zeroed, + cum_seqlens, + past_seqlens, + block_table, + scale=scale, + ) + torch.testing.assert_close(out_a.to(torch.float32), out_b.to(torch.float32), rtol=2e-3, atol=2e-3) + + # And the leading channels really are V: scaling them scales the output linearly. + self.assertGreater(out_a.abs().max().item(), 1e-3) + + def test_fp8_latent_cache(self): + if not has_fp8_kv_cache(): + self.skipTest("FP8 KV cache kernels are not built") + device = "cuda" + config = self._config(kv_cache_type="fp8", k_quant_type="PER_TENSOR") + past_seqlens = torch.tensor([8, 8], dtype=torch.int32, device=device) + new_seqlens = torch.tensor([2, 2], dtype=torch.int32, device=device) + total_seqlens = past_seqlens + new_seqlens + max_total = int(total_seqlens.max().item()) + blocks_per_seq = math.ceil(max_total / config.block_size) + num_blocks = blocks_per_seq * config.batch_size + block_table = torch.randperm(num_blocks, dtype=torch.int32, device=device).reshape( + config.batch_size, blocks_per_seq + ) + cum_seqlens = torch.zeros(config.batch_size + 1, dtype=torch.int32, device=device) + cum_seqlens[1:] = torch.cumsum(new_seqlens, dim=0) + token_count = int(cum_seqlens[-1].item()) + + k_scale = torch.tensor([0.01], dtype=torch.float32, device=device) + latent_float = torch.randn( + num_blocks, config.block_size, config.kv_num_heads, config.head_size, device=device, dtype=torch.float16 + ) + latent_paged = quantize_kv(latent_float, k_scale, "fp8") + + query = torch.randn(token_count, config.num_heads, config.head_size, device=device, dtype=torch.float16) + new_key = torch.randn(token_count, config.head_size, device=device, dtype=torch.float16) + scale = config.softmax_scale + out, latent_paged = run_mla( + config, + query.reshape(token_count, -1), + new_key, + latent_paged, + cum_seqlens, + past_seqlens, + block_table, + k_scale=k_scale, + scale=scale, + ) + # Dequantize the cache the op left behind and score against it, so only the attention math + # (not the quantization error) is under test. + dense_q = densify_latent(config, latent_paged, block_table, blocks_per_seq * config.block_size) + dense = dequantize_kv(dense_q, k_scale) + ref = mla_reference(config, query, dense, past_seqlens, new_seqlens, cum_seqlens, scale) + out = out.reshape(token_count, config.num_heads, config.v_head_size).to(device).to(torch.float32) + torch.testing.assert_close(out, ref, rtol=5e-3, atol=5e-3) + + # ---- rejected configurations (design doc §12.9, §12.10) ---- + + def _expect_rejected(self, message_fragment, **graph_kwargs): + """Build a LATENT model with one deliberate violation and assert the op rejects it. Schema + violations surface at session creation, input violations in ComputeInternal, so both are + wrapped.""" + config = graph_kwargs.pop("config", None) or self._config() + num_tokens, num_blocks, max_blocks = 4, 4, 2 + head_size = config.head_size + + def build_and_run(): + model = create_mla_graph(config, num_tokens, num_blocks, max_blocks, **graph_kwargs) + session = InferenceSession(model, SessionOptions(), providers=["CUDAExecutionProvider"]) + feeds = { + "query": torch.randn(num_tokens, config.num_heads * head_size).to(torch.float16).numpy(), + "key": torch.randn(num_tokens, config.kv_num_heads * head_size).to(torch.float16).numpy(), + "key_cache": torch.randn(num_blocks, config.block_size, config.kv_num_heads, head_size) + .to(torch.float16) + .numpy(), + "cumulative_sequence_length": numpy.array([0, 2, 4, 4, 4][: config.batch_size + 1], dtype=numpy.int32), + "past_seqlens": numpy.zeros(config.batch_size, dtype=numpy.int32), + "block_table": numpy.arange(config.batch_size * max_blocks, dtype=numpy.int32).reshape( + config.batch_size, max_blocks + ), + } + if graph_kwargs.get("with_value"): + feeds["value"] = torch.randn(num_tokens, config.kv_num_heads * head_size).to(torch.float16).numpy() + if graph_kwargs.get("with_value_cache"): + feeds["value_cache"] = feeds["key_cache"].copy() + if graph_kwargs.get("do_rotary"): + cache_width = config.qk_rope_head_dim // 2 + feeds["cos_cache"] = torch.ones(64, cache_width).to(torch.float16).numpy() + feeds["sin_cache"] = torch.zeros(64, cache_width).to(torch.float16).numpy() + if graph_kwargs.get("with_head_sink"): + feeds["head_sink"] = torch.zeros(config.num_heads).to(torch.float16).numpy() + if graph_kwargs.get("with_qk_norm"): + feeds["q_norm_weight"] = torch.ones(head_size).to(torch.float16).numpy() + feeds["k_norm_weight"] = torch.ones(head_size).to(torch.float16).numpy() + session.run(None, feeds) + + with self.assertRaises(Exception) as ctx: + build_and_run() + self.assertIn(message_fragment, str(ctx.exception)) + + def test_reject_missing_scale(self): + # v_head_size != head_size with no explicit scale would silently use 1/sqrt(head_size). + self._expect_rejected("explicit 'scale'") + + def test_reject_value_input(self): + self._expect_rejected("'value'", scale=0.1, with_value=True) + + def test_reject_value_cache(self): + self._expect_rejected("'value_cache' must be absent", scale=0.1, with_value_cache=True) + + def test_reject_value_cache_output(self): + self._expect_rejected("value_cache_out must be absent", scale=0.1, with_value_cache_out=True) + + def test_reject_head_sink(self): + self._expect_rejected("'head_sink'", scale=0.1, with_head_sink=True) + + def test_reject_qk_norm(self): + self._expect_rejected("q_norm_weight", scale=0.1, with_qk_norm=True) + + def test_reject_multi_kv_head(self): + config = self._config() + config.kv_num_heads = 2 + self._expect_rejected("'kv_num_heads' must be 1", config=config, scale=0.1) + + def test_reject_v_head_size_in_separate_layout(self): + # v_head_size != head_size is only meaningful for LATENT. + self._expect_rejected( + "may only differ from head_size", + scale=0.1, + kv_cache_layout="SEPARATE", + with_value=True, + with_value_cache=True, + with_value_cache_out=True, + ) + + def test_reject_unaligned_rotary_offset(self): + self._expect_rejected("multiple of 8", scale=0.1, do_rotary=True, rotary_offset=60) + + def test_reject_rotary_offset_overflow(self): + config = self._config() + self._expect_rejected( + "must not exceed head_size", config=config, scale=0.1, do_rotary=True, rotary_offset=config.head_size + ) + + if __name__ == "__main__": unittest.main(verbosity=2) diff --git a/onnxruntime/test/python/transformers/test_qmoe_nvfp4_cuda.py b/onnxruntime/test/python/transformers/test_qmoe_nvfp4_cuda.py index c79e8cc4368bc..6a0e7c69d3ae2 100644 --- a/onnxruntime/test/python/transformers/test_qmoe_nvfp4_cuda.py +++ b/onnxruntime/test/python/transformers/test_qmoe_nvfp4_cuda.py @@ -19,6 +19,8 @@ # -------------------------------------------------------------------------- import os +import subprocess +import sys import unittest import numpy @@ -308,6 +310,8 @@ def _run_nvfp4_moe_test( use_swiglu=False, block_size=NVFP4_BLOCK_SIZE, gemv_mode=None, + input_scale=1.0, + atol_override=None, ): self._skip_if_no_fp4() @@ -389,7 +393,7 @@ def _run_nvfp4_moe_test( else: os.environ["ORT_ENABLE_FP4_GEMV"] = prev_gemv_env - input_tensor = torch.randn(num_tokens, hidden_size, device=device, dtype=torch_dtype) + input_tensor = torch.randn(num_tokens, hidden_size, device=device, dtype=torch_dtype) * input_scale router_logits = torch.randn(num_tokens, num_experts, device=device, dtype=torch_dtype) output_tensor = torch.zeros(num_tokens, hidden_size, device=device, dtype=torch_dtype) @@ -409,6 +413,7 @@ def _run_nvfp4_moe_test( iobinding.synchronize_outputs() ort_output = output_tensor.clone() + self.assertTrue(torch.isfinite(ort_output).all().item(), "NVFP4 MoE output contains NaN or infinity") ref_output = self._compute_reference( input_tensor, @@ -432,6 +437,8 @@ def _run_nvfp4_moe_test( ) atol = 0.15 if torch_dtype == torch.bfloat16 else 0.12 + if atol_override is not None: + atol = atol_override # The native block-scaled FP4xFP4 CUTLASS prefill kernel (Blackwell / SM120+, taken # only when the per-run token count reaches the prefill threshold) additionally # quantizes the *activations* to 4-bit NVFP4 (block-16 with E4M3 block scales). The @@ -631,10 +638,13 @@ def test_nvfp4_fp16_larger_dims(self): # ================================================================ # Fused FP4 GEMV decode fast path (block size 16). The GEMV support window requires - # n, k >= 512 and expanded rows (num_tokens * top_k) <= 8, plus SwiGLU fusion, so these + # n, k >= 512 and expanded rows (num_tokens * top_k) <= 64, plus SwiGLU fusion, so these # decode-shaped SwiGLU cases route through the NVFP4 GEMV kernel (gemv_mode="1"). The # gemv_mode="0" companion forces the dequant fallback on the identical shape; both must # match the exact dequantized reference. + # + # "MTP" below is multi-token prediction (speculative decode): verifying N speculative + # tokens runs N+1 tokens at once, so a top_k=8 model expands to (N+1)*8 rows. # ================================================================ def test_nvfp4_fp16_gemv_decode_swiglu(self): @@ -661,6 +671,27 @@ def test_nvfp4_bf16_gemv_decode_swiglu(self): gemv_mode="1", ) + def test_nvfp4_fp16_gemv_scales_weights_before_multiply(self): + # Overflow guard for accumulate_column_tile(): it must apply the group scale to the + # decoded weight *before* multiplying by the activation. FP4 codes reach 6.0 and the + # group scales are well below 1, so multiplying first can overflow FP16 (max 65504) + # even when the scaled product is representable. Driving the activations to ~1e4 puts + # the unscaled products past that limit, which the isfinite() check in the helper + # catches. atol is raised because the outputs themselves are ~1e4 times larger; 3.0 + # there is ~3e-4 relative, i.e. far tighter than the default 0.12 at unit scale. + self._run_nvfp4_moe_test( + hidden_size=512, + inter_size=512, + num_experts=4, + top_k=2, + num_tokens=1, + onnx_dtype=TensorProto.FLOAT16, + use_swiglu=True, + gemv_mode="1", + input_scale=10000.0, + atol_override=3.0, + ) + def test_nvfp4_fp16_gemv_disabled_swiglu(self): self._run_nvfp4_moe_test( hidden_size=512, @@ -673,6 +704,121 @@ def test_nvfp4_fp16_gemv_disabled_swiglu(self): gemv_mode="0", ) + @parameterized.expand( + [ + (TensorProto.FLOAT16, 2), + (TensorProto.BFLOAT16, 2), + (TensorProto.FLOAT16, 3), + ] + ) + def test_nvfp4_gemv_mtp_swiglu(self, onnx_dtype, num_tokens): + self._run_nvfp4_moe_test( + hidden_size=512, + inter_size=512, + num_experts=8, + top_k=8, + num_tokens=num_tokens, + onnx_dtype=onnx_dtype, + use_swiglu=True, + gemv_mode="1", + ) + + def test_nvfp4_fp16_gemv_mtp_fallback_swiglu(self): + self._run_nvfp4_moe_test( + hidden_size=512, + inter_size=512, + num_experts=8, + top_k=8, + num_tokens=3, + onnx_dtype=TensorProto.FLOAT16, + use_swiglu=True, + gemv_mode="0", + ) + + def test_nvfp4_fp16_gemv_expanded_rows_at_window_limit(self): + # 8 tokens x top_k 8 = 64 expanded rows, exactly kMaxProfiledExpandedRowsFp4, so this + # is the largest shape is_moe_gemv_fp4_supported still accepts onto the GEMV path. + self._run_nvfp4_moe_test( + hidden_size=512, + inter_size=512, + num_experts=8, + top_k=8, + num_tokens=8, + onnx_dtype=TensorProto.FLOAT16, + use_swiglu=True, + gemv_mode="1", + ) + + def test_nvfp4_fp16_gemv_expanded_rows_above_window_limit(self): + # 9 tokens x top_k 8 = 72 > kMaxProfiledExpandedRowsFp4, so the GEMV path is rejected + # even with gemv_mode="1" and the run must still be correct on the dequant fallback. + self._run_nvfp4_moe_test( + hidden_size=512, + inter_size=512, + num_experts=8, + top_k=8, + num_tokens=9, + onnx_dtype=TensorProto.FLOAT16, + use_swiglu=True, + gemv_mode="1", + ) + + def test_nvfp4_fp16_gemv_long_k_tiling(self): + # Two gaps in one case. + # (a) Tiling: every other GEMV test uses k = 512 < kDefaultCtaK (StepK 8 * 128 threads), + # which only exercises the "idle threads" clause of Fp4MoeGemvDefaultConfig. k = 1024 + # reaches the second clause, where the choice is driven by blocks-per-SM instead; 64 + # expanded rows keep the grid large enough to actually take the 64-thread branch. + # (b) Accumulation order: the GEMV keeps even-k and odd-k partial sums apart until the + # epilogue, so it does not sum a column in the same order as the dequant fallback. + # Comparing the two ORT outputs directly at twice the K of the other parity test + # bounds that reordering drift where it has the most terms to cancel against. + shape = dict( + hidden_size=1024, + inter_size=1024, + num_experts=8, + top_k=8, + num_tokens=8, + onnx_dtype=TensorProto.FLOAT16, + use_swiglu=True, + ) + gemv_out = self._run_nvfp4_moe_test(**shape, gemv_mode="1") + fallback_out = self._run_nvfp4_moe_test(**shape, gemv_mode="0") + max_diff = (gemv_out.float() - fallback_out.float()).abs().max().item() + print(f"NVFP4 GEMV-vs-fallback parity (k=1024): FP16 SwiGLU max_diff={max_diff:.6f}") + self.assertLess( + max_diff, + 0.12, + f"NVFP4 fused GEMV diverged from dequant fallback at k=1024: max_diff={max_diff:.6f}", + ) + + def test_nvfp4_gemv_default_tiling_optout(self): + # ORT_FP4_GEMV_DEFAULT_TILING=0 restores the fixed kDefault tiling. The kernel latches + # it in a function-local static on first use, so it can only be exercised in a fresh + # process; re-run one GEMV test there with the opt-out set. + self._skip_if_no_fp4() + env = dict(os.environ) + env["ORT_FP4_GEMV_DEFAULT_TILING"] = "0" + proc = subprocess.run( + [ + sys.executable, + "-m", + "unittest", + "-v", + f"{os.path.splitext(os.path.basename(__file__))[0]}.TestQMoENVFP4.test_nvfp4_fp16_gemv_long_k_tiling", + ], + cwd=os.path.dirname(os.path.abspath(__file__)), + env=env, + capture_output=True, + text=True, + check=False, + ) + self.assertEqual( + proc.returncode, + 0, + f"ORT_FP4_GEMV_DEFAULT_TILING=0 run failed:\n{proc.stdout}\n{proc.stderr}", + ) + def test_nvfp4_fp16_gemv_vs_fallback_parity(self): # Dispatch-equivalence regression: run the identical decode-shaped NVFP4 case twice — once # forcing the fused FP4 GEMV path (ORT_ENABLE_FP4_GEMV=1) and once forcing the dequant diff --git a/plugin-ep-cuda/README.md b/plugin-ep-cuda/README.md index a5bff0274170d..fd3e730d437b7 100644 --- a/plugin-ep-cuda/README.md +++ b/plugin-ep-cuda/README.md @@ -12,6 +12,8 @@ For more information about plugin EPs, see the documentation of truth shared by all packages built from this directory. The packages do not declare a hard dependency on a specific ONNX Runtime package; instead, this version string is injected into each package's README at build/pack time, and the native plugin EP code validates compatibility at registration time. +- [`paths.txt`](paths.txt) - Specifies directories and paths related to the CUDA EP. These paths are used to filter the + commits considered when identifying changes between releases, e.g., for generating release notes. - [`python/`](python/) - Sources and build script for the `onnxruntime-ep-cuda12`/`onnxruntime-ep-cuda13` Python wheels. - [`csharp/`](csharp/) - Sources and packaging script for the `Microsoft.ML.OnnxRuntime.EP.Cuda` NuGet package. diff --git a/plugin-ep-cuda/paths.txt b/plugin-ep-cuda/paths.txt new file mode 100644 index 0000000000000..b39da54676e30 --- /dev/null +++ b/plugin-ep-cuda/paths.txt @@ -0,0 +1,51 @@ +:(top)onnxruntime/core/providers/cuda +:(top)onnxruntime/contrib_ops/cuda +:(top)plugin-ep-cuda +:(top)docs/cuda_plugin_ep +:(top)onnxruntime/docs/contrib_ops/cuda +:(top)onnxruntime/test/providers/cuda +:(top)onnxruntime/test/contrib_ops/cuda_kernels +:(top)cmake/external/cuDNN.cmake +:(top)cmake/external/cuda_configuration.cmake +:(top)cmake/external/cudnn_frontend.cmake +:(top)cmake/external/cutlass.cmake +:(top)cmake/patches/cudnn_frontend +:(top)cmake/patches/cutlass +:(top)cmake/onnxruntime_cuda_source_filters.cmake +:(top)cmake/onnxruntime_providers_cuda.cmake +:(top)cmake/onnxruntime_providers_cuda_plugin.cmake +:(top)tools/ci_build/github/azure-pipelines/c-api-noopenmp-packaging-pipelines-cuda13.yml +:(top)tools/ci_build/github/azure-pipelines/plugin-cuda-pipeline.yml +:(top)tools/ci_build/github/azure-pipelines/plugin-cuda-test-pipeline.yml +:(top)tools/ci_build/github/azure-pipelines/py-cuda-package-test-pipeline.yml +:(top)tools/ci_build/github/azure-pipelines/py-cuda-packaging-pipeline.yml +:(top)tools/ci_build/github/azure-pipelines/py-cuda13-packaging-pipeline.yml +:(top)tools/ci_build/github/azure-pipelines/stages/nuget-combine-cuda-stage.yml +:(top)tools/ci_build/github/azure-pipelines/stages/nuget-cuda-packaging-stage.yml +:(top)tools/ci_build/github/azure-pipelines/stages/nuget-cuda-publishing-stage.yml +:(top)tools/ci_build/github/azure-pipelines/stages/nuget-linux-cuda-packaging-stage.yml +:(top)tools/ci_build/github/azure-pipelines/stages/nuget-win-cuda-packaging-stage.yml +:(top)tools/ci_build/github/azure-pipelines/stages/plugin-cuda-nuget-packaging-stage.yml +:(top)tools/ci_build/github/azure-pipelines/stages/plugin-cuda-packaging-stage.yml +:(top)tools/ci_build/github/azure-pipelines/stages/plugin-linux-cuda-stage.yml +:(top)tools/ci_build/github/azure-pipelines/stages/plugin-linux-cuda-test-stage.yml +:(top)tools/ci_build/github/azure-pipelines/stages/plugin-win-cuda-stage.yml +:(top)tools/ci_build/github/azure-pipelines/stages/plugin-win-cuda-test-stage.yml +:(top)tools/ci_build/github/azure-pipelines/stages/py-gpu-packaging-stage.yml +:(top)tools/ci_build/github/azure-pipelines/stages/py-linux-gpu-stage.yml +:(top)tools/ci_build/github/azure-pipelines/stages/py-win-gpu-stage.yml +:(top).github/workflows/linux_cuda_ci.yml +:(top).github/workflows/linux_cuda_no_cudnn.yml +:(top).github/workflows/linux_cuda_plugin_ci.yml +:(top).github/workflows/windows_cuda.yml +:(top).github/workflows/windows_cuda_no_cudnn.yml +:(top).github/workflows/windows_cuda_plugin.yml +:(top)onnxruntime/test/python/transformers/cuda_plugin_ep_helper.py +:(top)onnxruntime/test/python/transformers/test_cuda_plugin_ep.py +:(top)onnxruntime/test/python/transformers/test_moe_cuda.py +:(top)onnxruntime/test/python/transformers/test_paged_attention_cuda.py +:(top)onnxruntime/test/python/transformers/test_qmoe_cuda.py +:(top)onnxruntime/test/python/transformers/test_qmoe_fp4_cuda.py +:(top)onnxruntime/test/python/transformers/test_qmoe_fp8_cuda.py +:(top)onnxruntime/test/python/transformers/test_qmoe_nvfp4_cuda.py +:(top)onnxruntime/test/python/transformers/test_qmoe_wfp4afp8_cuda.py \ No newline at end of file diff --git a/plugin-ep-webgpu/README.md b/plugin-ep-webgpu/README.md index b4b0b38fffe21..d7459b9f2a187 100644 --- a/plugin-ep-webgpu/README.md +++ b/plugin-ep-webgpu/README.md @@ -15,6 +15,8 @@ For more information about plugin EPs, see the of truth shared by all packages built from this directory. The packages do not declare a hard dependency on a specific ONNX Runtime package; instead, this version string is injected into each package's README at build/pack time, and the native plugin EP code validates compatibility at registration time. +- [`paths.txt`](paths.txt) — Specifies directories and paths that are related to the WebGPU EP. These paths are used to + filter the commits considered when identifying changes between releases, e.g., for generating release notes. - [`python/`](python/) — Sources and build script for the `onnxruntime-ep-webgpu` Python wheel. See [`python/README.md`](python/README.md) for build and test instructions. - [`csharp/`](csharp/) — Sources and packaging script for the `Microsoft.ML.OnnxRuntime.EP.WebGpu` NuGet package. See diff --git a/plugin-ep-webgpu/paths.txt b/plugin-ep-webgpu/paths.txt new file mode 100644 index 0000000000000..e3c471a6e7dcc --- /dev/null +++ b/plugin-ep-webgpu/paths.txt @@ -0,0 +1,11 @@ +:(top)cgmanifests/webgpu +:(top)cmake/onnxruntime_providers_webgpu.cmake +:(top)cmake/patches/dawn +:(top)include/onnxruntime/core/providers/webgpu +:(top)onnxruntime/contrib_ops/webgpu +:(top)onnxruntime/core/providers/webgpu +:(top)onnxruntime/core/session/plugin_ep/ep_factory_webgpu.cc +:(top)onnxruntime/core/session/plugin_ep/ep_factory_webgpu.h +:(top)onnxruntime/test/providers/webgpu +:(top)onnxruntime/test/webgpu +:(top)plugin-ep-webgpu diff --git a/tools/ci_build/github/azure-pipelines/plugin-cuda-pipeline.yml b/tools/ci_build/github/azure-pipelines/plugin-cuda-pipeline.yml index 4b713fbc1a3d7..fade2cc416e49 100644 --- a/tools/ci_build/github/azure-pipelines/plugin-cuda-pipeline.yml +++ b/tools/ci_build/github/azure-pipelines/plugin-cuda-pipeline.yml @@ -55,6 +55,22 @@ parameters: - dev default: dev +# CUDA 13.x always compresses fatbins with -compress-mode=size. Enabling this makes the +# CUDA 12.8 build match it, which shrinks the binary a lot but raises the minimum driver +# to the CUDA 12.4 level (Linux >= 550.54.14, Windows >= 551.61). +- name: enable_cuda_fatbin_size_compression + displayName: 'Compress CUDA fatbins for size (CUDA 12.8 only; needs r550+ driver)' + type: boolean + default: false + +# The fpA_intB GEMM/GEMV kernels are the largest single component of the binary after +# flash attention. They are opt-in at run time (ORT_FPA_INTB_GEMM / ep.cuda.fpa_intb_gemm), +# but MatMulNBits still uses them automatically when the weights are prepacked. +- name: enable_fpa_intb_gemm + displayName: 'Build fpA_intB GEMM CUDA kernels' + type: boolean + default: true + - name: cmake_build_type type: string default: 'Release' @@ -72,9 +88,9 @@ variables: # Non-dev package versions (release, RC) must use Release build type - name: invalidBuildTypeConfig value: ${{ and(ne(parameters.package_type, 'dev'), ne(parameters.cmake_build_type, 'Release')) }} - # aarch64 is only available for CUDA 13.x - - name: invalidAArch64Config - value: ${{ and(eq(parameters.build_linux_aarch64, true), ne(parameters.cuda_version, '13.x')) }} + # ARM64 is only available for CUDA 13.x + - name: invalidArm64Config + value: ${{ and(or(eq(parameters.build_windows_arm64, true), eq(parameters.build_linux_aarch64, true)), ne(parameters.cuda_version, '13.x')) }} extends: template: v1/1ES.Official.PipelineTemplate.yml@1esPipelines @@ -106,7 +122,7 @@ extends: stages: # Validate parameter combinations - - ${{ if or(eq(variables['invalidBuildTypeConfig'], 'True'), eq(variables['invalidAArch64Config'], 'True')) }}: + - ${{ if or(eq(variables['invalidBuildTypeConfig'], 'True'), eq(variables['invalidArm64Config'], 'True')) }}: - stage: Validate_Parameters displayName: 'Validate Parameters' dependsOn: [] @@ -123,11 +139,11 @@ extends: echo "##vso[task.logissue type=error]Non-dev package version requires Release build type." exit 1 displayName: 'ERROR: Non-dev package version requires Release build type' - - ${{ if eq(variables['invalidAArch64Config'], 'True') }}: + - ${{ if eq(variables['invalidArm64Config'], 'True') }}: - script: | - echo "##vso[task.logissue type=error]Linux aarch64 build is only available for CUDA 13.x." + echo "##vso[task.logissue type=error]Windows ARM64 and Linux aarch64 builds are only available for CUDA 13.x." exit 1 - displayName: 'ERROR: aarch64 requires CUDA 13.x' + displayName: 'ERROR: ARM64 requires CUDA 13.x' - ${{ else }}: - template: stages/plugin-cuda-packaging-stage.yml parameters: @@ -140,14 +156,18 @@ extends: package_type: ${{ parameters.package_type }} version_file: ${{ variables.epVersionFile }} cmake_build_type: ${{ parameters.cmake_build_type }} + enable_cuda_fatbin_size_compression: ${{ parameters.enable_cuda_fatbin_size_compression }} + enable_fpa_intb_gemm: ${{ parameters.enable_fpa_intb_gemm }} ${{ if eq(parameters.cuda_version, '12.8') }}: python_package_name: 'onnxruntime-ep-cuda12' docker_base_image: 'onnxruntimebuildcache.azurecr.io/internal/azureml/onnxruntime/build/cuda12_x64_almalinux8_gcc14:20251017.1' - cmake_x64_cuda_archs: '61-real;75-real;86-real;89-real;120-real;120-virtual' - cmake_arm64_cuda_archs: '61-real;75-real;86-real;89-real;120-real;120-virtual' + cmake_windows_x64_cuda_archs: '61-real;75-real;86-real;89-real;120-real' + cmake_linux_x64_cuda_archs: '75-real;80-real;86-real;89-real;90-real;120-real' ${{ if eq(parameters.cuda_version, '13.x') }}: python_package_name: 'onnxruntime-ep-cuda13' docker_base_image: 'onnxruntimebuildcache.azurecr.io/internal/azureml/onnxruntime/build/cuda13_x64_almalinux8_gcc14:20251107.1' docker_base_image_aarch64: 'onnxruntimebuildcache.azurecr.io/public/azureml/onnxruntime_build_cuda13_aarch64_almalinux9_gcc14:20260323.1' - cmake_x64_cuda_archs: '75-real;80-real;86-real;89-real;90-real;120-real;120-virtual' - cmake_arm64_cuda_archs: '110-real;120-real;121-real;120-virtual' + cmake_windows_x64_cuda_archs: '75-real;80-real;86-real;89-real;120-real' + cmake_windows_arm64_cuda_archs: '120-real;121-real' + cmake_linux_x64_cuda_archs: '75-real;80-real;86-real;89-real;90-real;120-real' + cmake_linux_aarch64_cuda_archs: '89-real;90-real;100-real;103-real;120-real;121-real' diff --git a/tools/ci_build/github/azure-pipelines/stages/plugin-cuda-packaging-stage.yml b/tools/ci_build/github/azure-pipelines/stages/plugin-cuda-packaging-stage.yml index 5f9837fddb69e..3522b92dd27be 100644 --- a/tools/ci_build/github/azure-pipelines/stages/plugin-cuda-packaging-stage.yml +++ b/tools/ci_build/github/azure-pipelines/stages/plugin-cuda-packaging-stage.yml @@ -51,6 +51,16 @@ parameters: - RelWithDebInfo - MinSizeRel +- name: enable_cuda_fatbin_size_compression + type: boolean + displayName: 'Compress CUDA fatbins with -compress-mode=size' + default: false + +- name: enable_fpa_intb_gemm + type: boolean + displayName: 'Build fpA_intB GEMM CUDA kernels' + default: true + - name: docker_base_image type: string displayName: 'Linux x86_64 docker base image' @@ -61,15 +71,25 @@ parameters: displayName: 'Linux aarch64 docker base image' default: '' -- name: cmake_x64_cuda_archs +- name: cmake_windows_x64_cuda_archs + type: string + displayName: 'CMAKE_CUDA_ARCHITECTURES for Windows x64' + default: '61-real;75-real;86-real;89-real;120-real' + +- name: cmake_windows_arm64_cuda_archs + type: string + displayName: 'CMAKE_CUDA_ARCHITECTURES for Windows ARM64' + default: '120-real;121-real' + +- name: cmake_linux_x64_cuda_archs type: string - displayName: 'CMAKE_CUDA_ARCHITECTURES for x64' - default: '61-real;75-real;86-real;89-real;120-real;120-virtual' + displayName: 'CMAKE_CUDA_ARCHITECTURES for Linux x64' + default: '75-real;80-real;86-real;89-real;90-real;120-real' -- name: cmake_arm64_cuda_archs +- name: cmake_linux_aarch64_cuda_archs type: string - displayName: 'CMAKE_CUDA_ARCHITECTURES for ARM64' - default: '110-real;120-real;121-real;120-virtual' + displayName: 'CMAKE_CUDA_ARCHITECTURES for Linux aarch64' + default: '89-real;90-real;100-real;103-real;120-real;121-real' - name: python_version type: string @@ -147,11 +167,13 @@ stages: parameters: arch: 'x64' cuda_version: ${{ parameters.cuda_version }} - cmake_cuda_archs: ${{ parameters.cmake_x64_cuda_archs }} + cmake_cuda_archs: ${{ parameters.cmake_windows_x64_cuda_archs }} package_version: ${{ parameters.package_type }} version_file: ${{ parameters.version_file }} python_package_name: ${{ parameters.python_package_name }} cmake_build_type: ${{ parameters.cmake_build_type }} + enable_cuda_fatbin_size_compression: ${{ parameters.enable_cuda_fatbin_size_compression }} + enable_fpa_intb_gemm: ${{ parameters.enable_fpa_intb_gemm }} # Windows ARM64 - ${{ if eq(parameters.build_windows_arm64, true) }}: @@ -160,11 +182,13 @@ stages: arch: 'arm64' cuda_version: ${{ parameters.cuda_version }} arm64_cuda_version: ${{ parameters.arm64_cuda_version }} - cmake_cuda_archs: ${{ parameters.cmake_arm64_cuda_archs }} + cmake_cuda_archs: ${{ parameters.cmake_windows_arm64_cuda_archs }} package_version: ${{ parameters.package_type }} version_file: ${{ parameters.version_file }} python_package_name: ${{ parameters.python_package_name }} cmake_build_type: ${{ parameters.cmake_build_type }} + enable_cuda_fatbin_size_compression: ${{ parameters.enable_cuda_fatbin_size_compression }} + enable_fpa_intb_gemm: ${{ parameters.enable_fpa_intb_gemm }} # Linux x64 - ${{ if eq(parameters.build_linux_x64, true) }}: @@ -174,7 +198,7 @@ stages: arch: 'x64' machine_pool: 'onnxruntime-Ubuntu2404-AMD-CPU' cuda_version: ${{ parameters.cuda_version }} - cmake_cuda_archs: ${{ parameters.cmake_x64_cuda_archs }} + cmake_cuda_archs: ${{ parameters.cmake_linux_x64_cuda_archs }} package_version: ${{ parameters.package_type }} version_file: ${{ parameters.version_file }} python_package_name: ${{ parameters.python_package_name }} @@ -183,6 +207,8 @@ stages: python_version: ${{ parameters.python_version }} docker_python_exe_path: ${{ parameters.docker_python_exe_path }} artifact_name: cuda_plugin_linux_x64 + enable_cuda_fatbin_size_compression: ${{ parameters.enable_cuda_fatbin_size_compression }} + enable_fpa_intb_gemm: ${{ parameters.enable_fpa_intb_gemm }} # Linux aarch64 (CUDA 13.x only) - ${{ if eq(parameters.build_linux_aarch64, true) }}: @@ -192,7 +218,7 @@ stages: arch: 'aarch64' machine_pool: 'onnxruntime-linux-ARM64-CPU-2019' cuda_version: ${{ parameters.cuda_version }} - cmake_cuda_archs: ${{ parameters.cmake_arm64_cuda_archs }} + cmake_cuda_archs: ${{ parameters.cmake_linux_aarch64_cuda_archs }} package_version: ${{ parameters.package_type }} version_file: ${{ parameters.version_file }} python_package_name: ${{ parameters.python_package_name }} @@ -201,6 +227,8 @@ stages: python_version: ${{ parameters.python_version }} docker_python_exe_path: ${{ parameters.docker_python_exe_path }} artifact_name: cuda_plugin_linux_aarch64 + enable_cuda_fatbin_size_compression: ${{ parameters.enable_cuda_fatbin_size_compression }} + enable_fpa_intb_gemm: ${{ parameters.enable_fpa_intb_gemm }} # NuGet packaging (runs after all platform builds) - template: plugin-cuda-nuget-packaging-stage.yml diff --git a/tools/ci_build/github/azure-pipelines/stages/plugin-linux-cuda-stage.yml b/tools/ci_build/github/azure-pipelines/stages/plugin-linux-cuda-stage.yml index e665ac35dd8ed..a1452b89ebb9a 100644 --- a/tools/ci_build/github/azure-pipelines/stages/plugin-linux-cuda-stage.yml +++ b/tools/ci_build/github/azure-pipelines/stages/plugin-linux-cuda-stage.yml @@ -57,6 +57,14 @@ parameters: type: string default: 'cuda_plugin_linux_x64' +- name: enable_cuda_fatbin_size_compression + type: boolean + default: false + +- name: enable_fpa_intb_gemm + type: boolean + default: true + stages: - stage: ${{ parameters.stage_name }} dependsOn: [] @@ -77,6 +85,16 @@ stages: artifactName: ${{ parameters.artifact_name }} variables: - template: ../templates/common-variables.yml + - name: FatbinCompressDefine + ${{ if eq(parameters.enable_cuda_fatbin_size_compression, true) }}: + value: 'onnxruntime_CUDA_FATBIN_COMPRESS_SIZE=ON' + ${{ else }}: + value: '' + - name: FpaIntBGemmDefine + ${{ if eq(parameters.enable_fpa_intb_gemm, true) }}: + value: 'onnxruntime_USE_FPA_INTB_GEMM=ON' + ${{ else }}: + value: 'onnxruntime_USE_FPA_INTB_GEMM=OFF' steps: - checkout: self clean: true @@ -115,7 +133,7 @@ stages: workingDirectory: $(Build.SourcesDirectory) displayName: 'Build CUDA Plugin (Python ${{ parameters.python_version }}, ${{ parameters.arch }}, CUDA ${{ parameters.cuda_version }})' env: - EXTRA_CMAKE_DEFINES: $(PluginEpVersionDefine) + EXTRA_CMAKE_DEFINES: $(PluginEpVersionDefine) $(FatbinCompressDefine) $(FpaIntBGemmDefine) - script: | set -e -x @@ -125,8 +143,19 @@ stages: echo "Error: Expected plugin binary not found at '$plugin_path'. Failing build to avoid publishing an invalid package." exit 1 fi + + dynamic_search_path="$(readelf -d "$plugin_path" | sed -nE 's/.*\((RPATH|RUNPATH)\).*\[([^]]*)\].*/\2/p')" + if [[ -n "$dynamic_search_path" && ":$dynamic_search_path:" == *::* ]]; then + echo "Error: Plugin binary contains an empty RPATH/RUNPATH component: '$dynamic_search_path'" + exit 1 + fi + if [[ "$dynamic_search_path" == *cuda* ]]; then + echo "Error: Plugin binary contains a hard-coded CUDA RPATH/RUNPATH: '$dynamic_search_path'" + exit 1 + fi + cp "$plugin_path" "$(Build.ArtifactStagingDirectory)/bin/" - displayName: 'Copy plugin binaries' + displayName: 'Verify and copy plugin binaries' - script: | set -e -x diff --git a/tools/ci_build/github/azure-pipelines/stages/plugin-win-cuda-stage.yml b/tools/ci_build/github/azure-pipelines/stages/plugin-win-cuda-stage.yml index 2ec3e21391037..c16f40365f471 100644 --- a/tools/ci_build/github/azure-pipelines/stages/plugin-win-cuda-stage.yml +++ b/tools/ci_build/github/azure-pipelines/stages/plugin-win-cuda-stage.yml @@ -47,6 +47,14 @@ parameters: type: string default: '61-real;75-real;86-real;89-real;120-real;120-virtual' +- name: enable_cuda_fatbin_size_compression + type: boolean + default: false + +- name: enable_fpa_intb_gemm + type: boolean + default: true + stages: - stage: Win_plugin_cuda_${{ parameters.arch }}_Build dependsOn: [] @@ -85,6 +93,16 @@ stages: value: '-Dorg.gradle.daemon=false' - name: VSGenerator value: 'Visual Studio 17 2022' + - name: FatbinCompressOption + ${{ if eq(parameters.enable_cuda_fatbin_size_compression, true) }}: + value: '--cmake_extra_defines onnxruntime_CUDA_FATBIN_COMPRESS_SIZE=ON' + ${{ else }}: + value: '' + - name: FpaIntBGemmOption + ${{ if eq(parameters.enable_fpa_intb_gemm, true) }}: + value: '--cmake_extra_defines onnxruntime_USE_FPA_INTB_GEMM=ON' + ${{ else }}: + value: '--cmake_extra_defines onnxruntime_USE_FPA_INTB_GEMM=OFF' steps: - checkout: self clean: true @@ -192,6 +210,8 @@ stages: --cmake_extra_defines CMAKE_CUDA_ARCHITECTURES="${{ parameters.cmake_cuda_archs }}" --cmake_extra_defines onnxruntime_BUILD_CUDA_EP_AS_PLUGIN=ON --cmake_extra_defines $(PluginEpVersionDefine) + $(FatbinCompressOption) + $(FpaIntBGemmOption) $(TelemetryOption) workingDirectory: '$(Build.BinariesDirectory)' @@ -221,6 +241,8 @@ stages: --cmake_extra_defines CMAKE_CUDA_ARCHITECTURES="${{ parameters.cmake_cuda_archs }}" --cmake_extra_defines onnxruntime_BUILD_CUDA_EP_AS_PLUGIN=ON --cmake_extra_defines $(PluginEpVersionDefine) + $(FatbinCompressOption) + $(FpaIntBGemmOption) $(TelemetryOption) workingDirectory: '$(Build.BinariesDirectory)' @@ -252,6 +274,8 @@ stages: --cmake_extra_defines onnxruntime_BUILD_CUDA_EP_AS_PLUGIN=ON --cmake_extra_defines onnxruntime_BUILD_UNIT_TESTS=OFF --cmake_extra_defines $(PluginEpVersionDefine) + $(FatbinCompressOption) + $(FpaIntBGemmOption) $(TelemetryOption) workingDirectory: '$(Build.BinariesDirectory)' @@ -269,8 +293,29 @@ stages: Write-Error "Expected plugin binary not found at '$pluginPath'. Failing build to avoid publishing an invalid package." exit 1 } - Write-Host "Verified plugin binary exists at: $pluginPath" - displayName: 'Verify plugin binary exists' + + $versionInfo = (Get-Item $pluginPath).VersionInfo + $requiredProperties = @( + 'FileDescription', + 'FileVersion', + 'OriginalFilename', + 'ProductName', + 'ProductVersion', + 'LegalCopyright', + 'Language' + ) + foreach ($property in $requiredProperties) { + if ([string]::IsNullOrWhiteSpace($versionInfo.$property)) { + throw "Plugin DLL version property '$property' is missing." + } + } + if ($versionInfo.OriginalFilename -ne 'onnxruntime_providers_cuda.dll') { + throw "Unexpected OriginalFilename '$($versionInfo.OriginalFilename)'." + } + + Write-Host "Verified plugin DLL version information:" + $versionInfo | Format-List $requiredProperties + displayName: 'Verify plugin binary and version information' - task: CopyFiles@2 displayName: 'Copy plugin binaries to staging directory' diff --git a/tools/python/compile_contributors.py b/tools/python/compile_contributors.py index 92ba59747493e..df62d72fa0a4b 100644 --- a/tools/python/compile_contributors.py +++ b/tools/python/compile_contributors.py @@ -11,7 +11,7 @@ Usage: python compile_contributors.py [--base ] [--target ] [--dir ] - [--paths [ ...]] + [--paths [ ...]] [--paths-file ] Example: python compile_contributors.py --base origin/rel-1.23.2 --target origin/rel-1.24.1 --dir rel-1.24.1_report @@ -21,6 +21,10 @@ python compile_contributors.py --base origin/main~500 --target origin/main \ --paths ":(top)path/to/component_a" ":(top)path/to/component_b" + # Or load pathspecs from a file (one pathspec per line; blank lines and # comments are ignored): + python compile_contributors.py --base origin/main~500 --target origin/main \ + --paths-file plugin-ep-webgpu/paths.txt + Outputs: - detail.csv: Detailed breakdown of PRs, authors, and commit links. - logs.txt: Processing logs and summary (professional humans-only contributor list for release notes). @@ -314,6 +318,18 @@ def get_prs_from_log(log_output, prs_base=None, log_file=None, scan_depth=100): return all_prs +def read_pathspecs_file(pathspecs_file): + """Read git pathspecs from file, trimming whitespace and skipping blanks/comments.""" + pathspecs = [] + with open(pathspecs_file, encoding="utf-8") as f: + for line in f: + entry = line.strip() + if not entry or entry.startswith("#"): + continue + pathspecs.append(entry) + return pathspecs + + def main(): parser = argparse.ArgumentParser(description="Compile contributor list from Git log comparison.") parser.add_argument("--base", default="origin/rel-1.23.2", help="Base branch/commit to compare from") @@ -333,8 +349,31 @@ def main(): "sub-PRs are still expanded regardless of paths." ), ) + parser.add_argument( + "--paths-file", + default=None, + metavar="FILE", + help=( + "Optional file containing paths (git pathspec, one per line) to limit history to. " + "Blank lines and lines starting with '#' are ignored. " + "Can be combined with --paths." + ), + ) args = parser.parse_args() + selected_paths = [] + if args.paths: + selected_paths.extend(args.paths) + if args.paths_file: + try: + selected_paths.extend(read_pathspecs_file(args.paths_file)) + except OSError as e: + parser.error(f"Could not read --paths-file '{args.paths_file}': {e}") + + # Preserve order while removing duplicates. + if selected_paths: + selected_paths = list(dict.fromkeys(selected_paths)) + # Early validation if not check_preflight(): return @@ -345,7 +384,7 @@ def main(): scan_depth = args.scan_depth # Build a pathspec suffix (e.g. ["--", "onnxruntime/core/providers/webgpu", ...]) once, # so it can be appended to each `git log` invocation below. - paths_args = (["--", *args.paths]) if args.paths else [] + paths_args = (["--", *selected_paths]) if selected_paths else [] if not os.path.exists(output_dir): os.makedirs(output_dir) @@ -353,8 +392,8 @@ def main(): logs_path = os.path.join(output_dir, "logs.txt") with open(logs_path, "w", encoding="utf-8") as log_file: log_event(f"Starting comparison: {branch_base} -> {branch_target}", log_file) - if args.paths: - log_event(f"Limiting history to paths: {args.paths}", log_file) + if selected_paths: + log_event(f"Limiting history to paths: {selected_paths}", log_file) # 1. Fetch base branch PRs (scan depth controlled by scan_depth) log_event(f"Fetching base branch history for {branch_base} (last {scan_depth})...", log_file)