Repository navigation
[CUDA] Support FP16 KV cache and FP16 query/key in DynamicSparseAttention<BFloat16> and FP32 rotary cache / FP16 output in RMSNormRotary - #33997
Open
Tianlei Wu (tianleiwu) wants to merge 2 commits into
Conversation
…tion<BFloat16> and FP32 rotary cache / FP16 output in RMSNormRotary - RMSNormRotary (bert_defs.cc, rms_norm_rotary.h/.cc/.cu): - Extend TCos type constraint to allow tensor(float) in addition to tensor(bfloat16) so 256K RoPE cos/sin caches can be kept in FP32. - Add output_fp16 attribute and TOut type constraint (tensor(bfloat16), tensor(float16)) and instantiate RMSNormRotaryExactKernel<kFp32Cos, kOutputFp16>, computing RMSNorm + RoPE in exact FP32 without intermediate BFloat16 rounding and outputting full 10-bit mantissa FP16 Q/K tensors. - DynamicSparseAttention (bert_defs.cc, dynamic_sparse_attention.cc, dynamic_sparse_attention_impl.h/.cu, dynamic_sparse_attention_prefill.cuh): - Extend TCache to support tensor(float16) and TQuery to support tensor(float16) in DynamicSparseAttention<BFloat16>. - Instantiate Tensor Core MmaFp16 prefill kernel (DynamicSparseAttentionPrefillKernel<false, true>) and grouped decode kernel (GroupedDecodeDynamicSparseAttentionKernel<..., kFp16Cache=true>) for BFloat16 activations with FP16 KV cache. - Skip PrepareQueryKernel when query is already in the target attention dtype (need_prepare_query == false). - Add unit tests in rms_norm_rotary_test.cc and dynamic_sparse_attention_op_test.cc. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
Comment on lines
+22
to
+37
| #define REGISTER_KERNEL_TYPED(T, ...) \ | ||
| ONNX_OPERATOR_TYPED_KERNEL_EX( \ | ||
| DynamicSparseAttention, \ | ||
| kMSDomain, \ | ||
| 1, \ | ||
| T, \ | ||
| kCudaExecutionProvider, \ | ||
| (*KernelDefBuilder::Create()) \ | ||
| .TypeConstraint("T", DataTypeImpl::GetTensorType<T>()) \ | ||
| .TypeConstraint("TQuery", BuildKernelDefConstraints<__VA_ARGS__>())\ | ||
| .TypeConstraint("TCache", BuildKernelDefConstraints<__VA_ARGS__>())\ | ||
| .TypeConstraint("S", DataTypeImpl::GetTensorType<float>()) \ | ||
| .TypeConstraint("M", DataTypeImpl::GetTensorType<int32_t>()) \ | ||
| .MayInplace(3, 1) \ | ||
| .MayInplace(4, 2) \ | ||
| .InputMemoryType(OrtMemTypeCPUInput, 10), \ |
Contributor
There was a problem hiding this comment.
Suggested change
| #define REGISTER_KERNEL_TYPED(T, ...) \ | |
| ONNX_OPERATOR_TYPED_KERNEL_EX( \ | |
| DynamicSparseAttention, \ | |
| kMSDomain, \ | |
| 1, \ | |
| T, \ | |
| kCudaExecutionProvider, \ | |
| (*KernelDefBuilder::Create()) \ | |
| .TypeConstraint("T", DataTypeImpl::GetTensorType<T>()) \ | |
| .TypeConstraint("TQuery", BuildKernelDefConstraints<__VA_ARGS__>())\ | |
| .TypeConstraint("TCache", BuildKernelDefConstraints<__VA_ARGS__>())\ | |
| .TypeConstraint("S", DataTypeImpl::GetTensorType<float>()) \ | |
| .TypeConstraint("M", DataTypeImpl::GetTensorType<int32_t>()) \ | |
| .MayInplace(3, 1) \ | |
| .MayInplace(4, 2) \ | |
| .InputMemoryType(OrtMemTypeCPUInput, 10), \ | |
| #define REGISTER_KERNEL_TYPED(T, ...) \ | |
| ONNX_OPERATOR_TYPED_KERNEL_EX( \ | |
| DynamicSparseAttention, \ | |
| kMSDomain, \ | |
| 1, \ | |
| T, \ | |
| kCudaExecutionProvider, \ | |
| (*KernelDefBuilder::Create()) \ | |
| .TypeConstraint("T", DataTypeImpl::GetTensorType<T>()) \ | |
| .TypeConstraint("TQuery", BuildKernelDefConstraints<__VA_ARGS__>()) \ | |
| .TypeConstraint("TCache", BuildKernelDefConstraints<__VA_ARGS__>()) \ | |
| .TypeConstraint("S", DataTypeImpl::GetTensorType<float>()) \ | |
| .TypeConstraint("M", DataTypeImpl::GetTensorType<int32_t>()) \ | |
| .MayInplace(3, 1) \ | |
| .MayInplace(4, 2) \ | |
| .InputMemoryType(OrtMemTypeCPUInput, 10), \ |
Contributor
There was a problem hiding this comment.
🟡 Changes recommended
Grouped attention can overflow BF16 auxiliary values when converting to FP16, while the schema change permits unsupported mixed types that generate invalid WebGPU shaders.
4 open issues
What changed in this PR
This PR extends ONNX Runtime’s CUDA contrib operators for long-context models, adding mixed-precision sparse attention and higher-precision RMS normalization with rotary embeddings.
Changes:
- Add FP16 query/key and KV-cache support to DynamicSparseAttention, including prefill/decode kernels and a query-preparation bypass.
- Add FP32 rotary caches and optional FP16 output to RMSNormRotary using FP32 intermediate arithmetic.
- Update operator schemas, output type inference, and CUDA parity tests.
| File | Description |
|---|---|
| onnxruntime/test/contrib_ops/rms_norm_rotary_test.cc | Adds FP32-cache and FP16-output parity coverage. |
| onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc | Adds mixed query/cache dtype coverage. |
| onnxruntime/core/graph/contrib_ops/bert_defs.cc | Extends schemas and updates output type inference. |
| onnxruntime/contrib_ops/cuda/bert/rms_norm_rotary.h | Extends launch parameters and output configuration. |
| onnxruntime/contrib_ops/cuda/bert/rms_norm_rotary.cu | Implements the FP32-intermediate normalization/rotary kernel. |
| onnxruntime/contrib_ops/cuda/bert/rms_norm_rotary.cc | Registers new types and dispatches cache/output modes. |
| onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention.cc | Registers and validates mixed-precision inputs. |
| onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_prefill.cuh | Adds FP16 Tensor Core prefill support. |
| onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.h | Adds mixed-precision execution flags. |
| onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu | Updates preparation, cache writes, attention kernels, and dispatch. |
Review effort: Balanced
Comment on lines
+947
to
+948
| shared_key[shared_offset + d] = FromFloat<half>(ToFloat(auxiliary_key[offset + d])); | ||
| shared_value[shared_offset + d] = FromFloat<half>(ToFloat(auxiliary_value[offset + d])); |
Comment on lines
+704
to
+705
| "TQuery") | ||
| .Input(1, "key", "Current main key [batch, sequence, kv_num_heads * head_size].", "TQuery", |
Comment on lines
+400
to
+401
| result_f32 = ApplyRotaryFloat(reduction, h, static_cast<int>(position), cos_cache, sin_cache, | ||
| rotary_dim, rotary_offset, rotary_interleaved); |
Comment on lines
+6427
to
+6428
| .TypeConstraint("TCos", {"tensor(bfloat16)", "tensor(float)"}, "BF16 or FP32 rotary cache.") | ||
| .TypeConstraint("TOut", {"tensor(bfloat16)", "tensor(float16)"}, "BF16 or FP16 output.") |
…nd FP16 X input in RMSNormRotary Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Add native
FP16(MLFloat16) KV cache andFP16query/key support toDynamicSparseAttention<BFloat16>andFP32rotary cache /FP16output support toRMSNormRotaryfor long-context Qwen3.8-Flash-Next models:com.microsoft::RMSNormRotary(bert_defs.cc,rms_norm_rotary.h/.cc/.cu):TCostype constraint to allowtensor(float)in addition totensor(bfloat16)so 256K RoPEcos_cache/sin_cachecan be stored in exactFP32.output_fp16attribute (0defaultBFloat16,1MLFloat16) and extendTOuttype constraint totensor(bfloat16)andtensor(float16).RMSNormRotaryExactKernel<kFp32Cos, kOutputFp16>which computes RMSNorm + RoPE in exactFP32without intermediateBFloat16rounding before converting to the output type, reducing Q/K relative RMSE by 9.4x (0.1954% -> 0.0207%) and@32Kattention logit RMSE by 9.5x (0.2789% -> 0.0293%).com.microsoft::DynamicSparseAttention(bert_defs.cc,dynamic_sparse_attention.cc,dynamic_sparse_attention_impl.h/.cu,dynamic_sparse_attention_prefill.cuh):TCacheto supporttensor(float16)andTQueryto supporttensor(float16)inDynamicSparseAttention<BFloat16>, while fixing output 0 (output) type inference to propagate from input 2 (value) when present.MmaFp16prefill kernel (DynamicSparseAttentionPrefillKernel<false, true>,mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32) and grouped decode kernel (GroupedDecodeDynamicSparseAttentionKernel<..., kFp16Cache=true>) forBFloat16activations withFP16KV cache.PrepareQueryKernelwhenqueryalready matches the target attention dtype (need_prepare_query == false), making theFP16KV cache path slightly faster thanBF16(143 usvs152 usper layer atS=8, L=32K..192K).Unit Tests:
RMSNormRotaryFp32CosAndFp16OutputParityinonnxruntime/test/contrib_ops/rms_norm_rotary_test.cc.SelectedOnlyFp16CacheBFloat16_CUDAandMLFloat16cache/query coverage inSelectedOnlyPrefillBFloat16_CUDA,FloatingCacheGroupedReferenceParity_CUDA, andGroupedDecodeRatio12ReferenceParity_CUDAinonnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc.