[Tianlei Wu (tianleiwu)]

Description

Add native FP16 (MLFloat16) KV cache and FP16 query/key support to DynamicSparseAttention<BFloat16> and FP32 rotary cache / FP16 output support to RMSNormRotary for long-context Qwen3.8-Flash-Next models:

1. com.microsoft::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_cache/sin_cache can be stored in exact FP32.
   • Add output_fp16 attribute (0 default BFloat16, 1 MLFloat16) and extend TOut type constraint to tensor(bfloat16) and tensor(float16).
   • Add RMSNormRotaryExactKernel<kFp32Cos, kOutputFp16> which computes RMSNorm + RoPE in exact FP32 without intermediate BFloat16 rounding before converting to the output type, reducing Q/K relative RMSE by 9.4x (0.1954% -> 0.0207%) and @32K attention logit RMSE by 9.5x (0.2789% -> 0.0293%).

2. com.microsoft::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>, while fixing output 0 (output) type inference to propagate from input 2 (value) when present.
   • Instantiate Tensor Core MmaFp16 prefill kernel (DynamicSparseAttentionPrefillKernel<false, true>, mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32) and grouped decode kernel (GroupedDecodeDynamicSparseAttentionKernel<..., kFp16Cache=true>) for BFloat16 activations with FP16 KV cache.
   • Skip PrepareQueryKernel when query already matches the target attention dtype (need_prepare_query == false), making the FP16 KV cache path slightly faster than BF16 (143 us vs 152 us per layer at S=8, L=32K..192K).

3. Unit Tests:

   • Added RMSNormRotaryFp32CosAndFp16OutputParity in onnxruntime/test/contrib_ops/rms_norm_rotary_test.cc.
   • Added SelectedOnlyFp16CacheBFloat16_CUDA and MLFloat16 cache/query coverage in SelectedOnlyPrefillBFloat16_CUDA, FloatingCacheGroupedReferenceParity_CUDA, and GroupedDecodeRatio12ReferenceParity_CUDA in onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc.

fp32_fp16kv_v2 (--qsa-qkv-dtype v_rc5_fp16 + MLFloat16 value in DynamicSparseAttention<BFloat16>)

In v1 (*_fp32_fp16kv), V (Input(2, "value")) entering DynamicSparseAttention<BFloat16> originated from merged_proj (SplitQsa) in BFloat16 (7-bit mantissa), leaving the lower 3 mantissa bits of present_value (FP16) zero-padded (0.1658% truncation error). Extending DynamicSparseAttention<BFloat16> to accept MLFloat16 value (output_bf16=1), extending RMSNormRotary to accept MLFloat16 or BFloat16 input X, and computing QSA v_proj from the bit-exact CutlassFpAIntB dequantized weight w_bf16_deq in FLOAT16 (--qsa-qkv-dtype v_rc5_fp16, +30.0 MiB across 12 layers, 0.000000% weight shift vs. rc5 calibration, 5.7x lower V rounding error 0.1658% -> 0.0291%) with native INT8 PLE (--precision-profile fp32_fp16kv_v2) yields:

Model Variant	Peak VRAM at 256K (GiB)	MMLU-Pro 1200 at 16K (Strict)	LCB v6 at 64K (N=175)	LCB v6 at 128K (N=175)	LCB v6 at 256K (N=175)	MMLU / LCB Tok-Wtd TPS
rc5_top16_fp32_fp16kv_v2	69.20	84.00% (1008/1200)	73.14% (128/175, +5.71pp vs Base)	87.43% (153/175)	88.00% (154/175l, +1.14pp vs _fp32, +2.29pp vs v1)	280.3 / 244.2
rc5_all_int4_fp32_fp16kv_v2	80.92	86.00% (1032/1200) (strict 85.33% = 1024/1200)	69.71% (122/175)	81.71% (143/175)	84.57% (148/175, +0.57pp vs v1)	272.6 / 248.4
rc5_2b375_fp32_fp16kv_v2	55.12	83.50% (1002/1200, +0.58pp vs Base, +1.08pp vs _fp32)	68.57% (120/175, +4.00pp vs Base, +4.57pp vs _fp32)	82.29% (144/175, +2.29pp vs Base, +6.29pp vs _fp32)	84.57% (148/175, +1.71pp vs Base, +5.71pp vs _fp32)	311.7 / 255.6

[github-actions]

You can commit the suggested changes from lintrunner.

[Copilot]

🟡 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

• FP16 cache path overflows BF16 auxiliary values · New
• WebGPU accepts mixed types that produce invalid WGSL · New
• FP16-cache normalization and RoPE branches lack parity tests · New
• Regenerate operator docs for updated schemas and registrations · New

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}

[Tianlei Wu (tianleiwu)]

Feedback fixes pushed in d09014c095.

Targeted validation

Suite	Passed	Failed	Skipped
DynamicSparseAttention CUDA	77	0	0
RMSNormRotary CUDA	3	1 (pre-existing; see below)	0
Native WebGPU (lavapipe)	14	0	0

94 targeted tests passed. The new auxiliary-range, FP16-cache norm/RoPE, WebGPU mixed-type rejection, and FP16-input/default-BF16-output regressions pass. The existing FP32 rotary-cache/FP16-output precision grid also passes. Nsight confirms mixed auxiliary inputs use generic fused/split kernels while main-cache-only cases still exercise grouped kernels.

The remaining RMSNormRotaryExactUnfusedParity failure is unchanged baseline behavior: it reproduces on the original PR head b99616f2f8 with default NVCC flags. The discrepancy comes from BF16 FMA contraction versus separately rounded reference arithmetic. An isolated control recompilation of the original, unchanged RMSNormRotary CUDA translation unit with --fmad=false passes; the main validation build retained its normal flags, and this feedback commit does not change legacy kernel arithmetic or test tolerances.

Contrib schema docs match fresh generation under CI comparison rules. The two CUDA rows and WebGPU DSA row exactly match documentation generated from real rebuilt registries; unrelated provider sections are unchanged. Targeted lintrunner and whitespace checks pass.