From 9e8a15a574dea1727d8787d85aa368cf9bc68985 Mon Sep 17 00:00:00 2001
From: Akshay Sonawane <asonawane@microsoft.com>
Date: Mon, 19 Oct 2026 17:56:51 +0000
Subject: [PATCH 1/2] Optimize quantized MoE and block-scaled FP8 CUDA
 projections

Part of the grouped generic CUDA operator stack. Preserve the qualified source snapshot.
---
 cmake/onnxruntime_unittests.cmake             |   6 +
 .../contrib_ops/cuda/cuda_contrib_kernels.cc  |   6 +
 .../cuda/llm/moe_gemm/moe_gemv_fp4.cu         | 124 ++++++---
 .../cuda/llm/moe_gemm/moe_gemv_fp4.h          |   9 +
 .../cuda/llm/moe_gemm/moe_kernels.cu          |  13 +-
 .../cuda/math/matmul_block_scaled_fp8.cc      |  94 +++++--
 .../cuda/math/matmul_block_scaled_fp8.cu      | 139 +++++++---
 .../cuda/math/matmul_block_scaled_fp8.h       |  17 +-
 .../contrib_ops/cuda/moe/moe_quantization.cc  |  15 +-
 .../contrib_ops/cuda/moe/qmoe_kernels.cu      |  36 +++
 .../contrib_ops/cuda/moe/qmoe_kernels.h       |   4 +
 .../quantize_dequantize_linear.cc             |   3 +-
 .../core/graph/contrib_ops/contrib_defs.cc    |  39 ++-
 .../graph/contrib_ops/quantization_defs.cc    |  61 +++--
 .../cuda_kernels/qmoe_fp8_compaction_test.cc  | 243 ++++++++++++++++++
 .../test/contrib_ops/fused_tensor_ops_test.cc |  15 ++
 .../matmul_block_scaled_fp8_test.cc           | 119 +++++++++
 onnxruntime/test/contrib_ops/moe_test.cc      |  21 +-
 18 files changed, 844 insertions(+), 120 deletions(-)

diff --git a/cmake/onnxruntime_unittests.cmake b/cmake/onnxruntime_unittests.cmake
index a50658db46486..17526724dede9 100644
--- a/cmake/onnxruntime_unittests.cmake
+++ b/cmake/onnxruntime_unittests.cmake
@@ -1158,6 +1158,12 @@ if (onnxruntime_ENABLE_CUDA_EP_INTERNAL_TESTS AND onnxruntime_BUILD_CUDA_EP_AS_P
         "${ONNXRUNTIME_ROOT}/contrib_ops/cuda/llm/moe_gemm/moe_gemm_kernels_fp16_uint2.cu"
         "${ONNXRUNTIME_ROOT}/contrib_ops/cuda/moe/qmoe_kernels.cu"
       )
+      if (TARGET onnxruntime_providers_cuda_plugin_llm_fp4)
+        list(APPEND onnxruntime_providers_cuda_plugin_ut_impl_src
+          "${ONNXRUNTIME_ROOT}/contrib_ops/cuda/llm/moe_gemm/moe_gemv_fp4.cu")
+        set_property(SOURCE "${TEST_SRC_DIR}/contrib_ops/cuda_kernels/qmoe_fp8_compaction_test.cc"
+          APPEND PROPERTY COMPILE_DEFINITIONS ORT_TEST_NVFP4_GEMV)
+      endif()
     endif()
 
     onnxruntime_add_object_library(onnxruntime_providers_cuda_plugin_ut_impl ${onnxruntime_providers_cuda_plugin_ut_impl_src})
diff --git a/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc b/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc
index 11ef53dd60d45..c7fe3de628179 100644
--- a/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc
+++ b/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc
@@ -179,6 +179,9 @@ class CUDA_MS_OP_TYPED_CLASS_NAME(1, BFloat16, ArgMaxLastAxis);
 class CUDA_MS_OP_CLASS_NAME(1, EmbeddingRMSNorm);
 class CUDA_MS_OP_CLASS_NAME(1, RowSelect);
 class CUDA_MS_OP_CLASS_NAME(1, SigmoidGateMul);
+#if !defined(DISABLE_FLOAT8_TYPES)
+class CUDA_MS_OP_TYPED_CLASS_NAME(1, Float8E4M3FN_BFloat16, DequantizeLinear);
+#endif
 class CUDA_MS_OP_CLASS_NAME(1, UntileHyperConnectionWeights);
 class CUDA_MS_OP_TYPED_CLASS_NAME(1, float, LinearAttentionGate);
 class CUDA_MS_OP_TYPED_CLASS_NAME(1, MLFloat16, LinearAttentionGate);
@@ -560,6 +563,9 @@ Status RegisterCudaContribKernels(KernelRegistry& kernel_registry) {
       BuildKernelCreateInfo<CUDA_MS_OP_CLASS_NAME(1, EmbeddingRMSNorm)>,
       BuildKernelCreateInfo<CUDA_MS_OP_CLASS_NAME(1, RowSelect)>,
       BuildKernelCreateInfo<CUDA_MS_OP_CLASS_NAME(1, SigmoidGateMul)>,
+#if !defined(DISABLE_FLOAT8_TYPES)
+      BuildKernelCreateInfo<CUDA_MS_OP_TYPED_CLASS_NAME(1, Float8E4M3FN_BFloat16, DequantizeLinear)>,
+#endif
       BuildKernelCreateInfo<CUDA_MS_OP_CLASS_NAME(1, UntileHyperConnectionWeights)>,
       BuildKernelCreateInfo<CUDA_MS_OP_TYPED_CLASS_NAME(1, float, LinearAttentionGate)>,
       BuildKernelCreateInfo<CUDA_MS_OP_TYPED_CLASS_NAME(1, MLFloat16, LinearAttentionGate)>,
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 5d4a72f8d7e62..9d5006911b3b3 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
@@ -4,6 +4,7 @@
 #include "contrib_ops/cuda/llm/moe_gemm/moe_gemv_fp4.h"
 
 #include <cuda_fp16.h>
+#include <cuda_fp8.h>
 #include <type_traits>
 
 // Shared INT/FP4 profiled-shape thresholds.
@@ -162,19 +163,8 @@ template <typename T>
 using Fp4GemvAccT = std::conditional_t<std::is_same<T, half>::value, half, float>;
 
 __device__ __forceinline__ float DecodeE4M3Fn(uint8_t code) {
-  const int sign = code & 0x80;
-  const int exponent = (code >> 3) & 0x0f;
-  const int mantissa = code & 0x07;
-  if ((code & 0x7f) == 0) {
-    return sign ? -0.0f : 0.0f;
-  }
-  if (exponent == 0x0f && mantissa == 0x07) {
-    return __int_as_float(0x7fffffff);
-  }
-  const float value = exponent == 0
-                          ? static_cast<float>(mantissa) * 0.001953125f
-                          : __int_as_float(((exponent + 120) << 23) | (mantissa << 20));
-  return sign ? -value : value;
+  const __half_raw decoded = __nv_cvt_fp8_to_halfraw(code, __NV_E4M3);
+  return __half2float(__ushort_as_half(decoded.x));
 }
 
 template <typename T, bool FusedSwiGlu, bool EnableBias>
@@ -353,23 +343,31 @@ __global__ void MoeGemvFp4RawNPackedKernel(
 #endif
 }
 
-template <typename T, bool FusedSwiGlu, bool EnableBias>
+template <typename T, bool FusedSwiGlu, bool EnableBias, bool FusedFinalize = false>
 __global__ void MoeGemvFp4RawKPackedKernel(
     const T* act, const uint8_t* weight, const uint8_t* block_scales, const float* global_scales,
     const T* bias, T* out,
     const int64_t* expert_first_token_offset, const int* permuted_row_to_expert, int num_experts,
     int64_t weight_expert_stride, int64_t scale_expert_stride, int output_columns, int reduction_size,
     cutlass_kernels::ActivationParams activation_params,
-    const int* permuted_row_to_source_row, int num_rows) {
+    const int* permuted_row_to_source_row, int num_rows,
+    const int* unpermuted_row_to_permuted_row = nullptr, const float* routing_scales = nullptr,
+    int experts_per_token = 0) {
 #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)
-  constexpr int kKLanes = 8;
+  constexpr int kKLanes = FusedSwiGlu ? 16 : 8;
   constexpr int kBlockSize = 16;
-  constexpr int kColsPerBlock = 16;
+  constexpr int kColsPerBlock = FusedFinalize ? 8 : 128 / kKLanes;
+  constexpr int kReductionUnroll = FusedSwiGlu ? 2 : 1;
   constexpr unsigned kFullMask = 0xffffffffu;
   using PackedT = std::conditional_t<std::is_same_v<T, half>, half2, __nv_bfloat162>;
 
-  const int row = static_cast<int>(blockIdx.x);
-  int expert = permuted_row_to_expert != nullptr ? permuted_row_to_expert[row] : 0;
+  const int row = FusedFinalize
+                      ? unpermuted_row_to_permuted_row[threadIdx.y * num_rows + blockIdx.x]
+                      : static_cast<int>(blockIdx.x);
+  int expert = row >= 0 && permuted_row_to_expert != nullptr ? permuted_row_to_expert[row] : -1;
+  if constexpr (!FusedFinalize) {
+    if (permuted_row_to_expert == nullptr) expert = 0;
+  }
 #pragma unroll 1
   for (int expert_index = 0; expert_index < num_experts && permuted_row_to_expert == nullptr; ++expert_index) {
     if (row >= static_cast<int>(expert_first_token_offset[expert_index + 1])) {
@@ -378,13 +376,16 @@ __global__ void MoeGemvFp4RawKPackedKernel(
     }
     break;
   }
-  if (expert < 0 || expert >= num_experts) {
-    return;
+  const bool valid_expert = expert >= 0 && expert < num_experts;
+  if (!valid_expert) {
+    if constexpr (!FusedFinalize) return;
+    expert = 0;
   }
 
   const int column = static_cast<int>(blockIdx.y) * kColsPerBlock + threadIdx.x / kKLanes;
   const int k_lane = threadIdx.x % kKLanes;
-  const int source_row = permuted_row_to_source_row ? permuted_row_to_source_row[row] % num_rows : row;
+  const int source_row = permuted_row_to_source_row ? permuted_row_to_source_row[row] % num_rows
+                                                    : (FusedFinalize ? max(0, row) : row);
   const T* row_act = act + static_cast<int64_t>(source_row) * reduction_size;
   const uint8_t* expert_weight = weight + static_cast<int64_t>(expert) * weight_expert_stride;
   const uint8_t* expert_scales = block_scales + static_cast<int64_t>(expert) * scale_expert_stride;
@@ -392,8 +393,9 @@ __global__ void MoeGemvFp4RawKPackedKernel(
   const T* expert_bias = EnableBias ? bias + static_cast<int64_t>(expert) * output_columns : nullptr;
 
   float accumulators[2] = {};
+#pragma unroll(kReductionUnroll)
   for (int reduction_column = k_lane * kBlockSize;
-       column < output_columns && reduction_column < reduction_size;
+       valid_expert && column < output_columns && reduction_column < reduction_size;
        reduction_column += kKLanes * kBlockSize) {
     uint2 packed = *reinterpret_cast<const uint2*>(
         expert_weight + static_cast<int64_t>(column) * (reduction_size / 2) + reduction_column / 2);
@@ -421,14 +423,77 @@ __global__ void MoeGemvFp4RawKPackedKernel(
   for (int offset = kKLanes / 2; offset > 0; offset /= 2) {
     accumulator += __shfl_xor_sync(kFullMask, accumulator, offset, kKLanes);
   }
-  const float paired_accumulator = __shfl_xor_sync(kFullMask, accumulator, kKLanes);
-  if (k_lane == 0 && column < output_columns && column % 2 == 0) {
-    StoreMoeGemvFp4Pair<T, FusedSwiGlu, EnableBias>(
-        accumulator, paired_accumulator, expert_bias, out, row, column, output_columns, expert, activation_params);
+  if constexpr (FusedFinalize) {
+    __shared__ float expert_results[16][kColsPerBlock];
+    __shared__ bool expert_valid[16];
+    if (threadIdx.x == 0) expert_valid[threadIdx.y] = valid_expert;
+    if (k_lane == 0) {
+      if constexpr (EnableBias) {
+        if (valid_expert && column < output_columns) accumulator += static_cast<float>(expert_bias[column]);
+      }
+      const T rounded = static_cast<T>(accumulator);
+      expert_results[threadIdx.y][threadIdx.x / kKLanes] = static_cast<float>(rounded);
+    }
+    __syncthreads();
+    if (threadIdx.y == 0 && threadIdx.x < kColsPerBlock) {
+      const int output_column = static_cast<int>(blockIdx.y) * kColsPerBlock + threadIdx.x;
+      if (output_column < output_columns) {
+        float sum = 0.0f;
+        for (int route = 0; route < experts_per_token; ++route) {
+          if (expert_valid[route]) {
+            sum += routing_scales[blockIdx.x * experts_per_token + route] * expert_results[route][threadIdx.x];
+          }
+        }
+        out[static_cast<int64_t>(blockIdx.x) * output_columns + output_column] = static_cast<T>(sum);
+      }
+    }
+  } else {
+    const float paired_accumulator = __shfl_xor_sync(kFullMask, accumulator, kKLanes);
+    if (k_lane == 0 && column < output_columns && column % 2 == 0) {
+      StoreMoeGemvFp4Pair<T, FusedSwiGlu, EnableBias>(
+          accumulator, paired_accumulator, expert_bias, out, row, column, output_columns, expert, activation_params);
+    }
   }
 #endif
 }
 
+template <typename T>
+bool try_launch_moe_gemv_nvfp4_fused_finalize(
+    const T* act, const uint8_t* weight, const uint8_t* block_scales, const float* global_scales,
+    const T* bias, T* out, const int* unpermuted_row_to_permuted_row,
+    const int* permuted_row_to_expert, const float* routing_scales, int num_experts,
+    int num_rows, int experts_per_token, int n, int k, cudaStream_t stream) {
+  if (num_rows < 1 || num_rows > 2 || experts_per_token < 1 || experts_per_token > 16 ||
+      n <= 0 || n % 2 != 0 || k <= 0 || k % 16 != 0) {
+    return false;
+  }
+  const dim3 grid(static_cast<unsigned>(num_rows), static_cast<unsigned>((n + 7) / 8));
+  const dim3 block(64, static_cast<unsigned>(experts_per_token));
+  const int64_t weight_stride = static_cast<int64_t>(n) * k / 2;
+  const int64_t scale_stride = static_cast<int64_t>(n) * (k / 16);
+  const cutlass_kernels::ActivationParams activation;
+  if (bias != nullptr) {
+    MoeGemvFp4RawKPackedKernel<T, false, true, true><<<grid, block, 0, stream>>>(
+        act, weight, block_scales, global_scales, bias, out, nullptr, permuted_row_to_expert,
+        num_experts, weight_stride, scale_stride, n, k, activation, nullptr, num_rows,
+        unpermuted_row_to_permuted_row, routing_scales, experts_per_token);
+  } else {
+    MoeGemvFp4RawKPackedKernel<T, false, false, true><<<grid, block, 0, stream>>>(
+        act, weight, block_scales, global_scales, bias, out, nullptr, permuted_row_to_expert,
+        num_experts, weight_stride, scale_stride, n, k, activation, nullptr, num_rows,
+        unpermuted_row_to_permuted_row, routing_scales, experts_per_token);
+  }
+  CUDA_CALL_THROW(cudaGetLastError());
+  return true;
+}
+
+template bool try_launch_moe_gemv_nvfp4_fused_finalize<half>(
+    const half*, const uint8_t*, const uint8_t*, const float*, const half*, half*,
+    const int*, const int*, const float*, int, int, int, int, int, cudaStream_t);
+template bool try_launch_moe_gemv_nvfp4_fused_finalize<__nv_bfloat16>(
+    const __nv_bfloat16*, const uint8_t*, const uint8_t*, const float*, const __nv_bfloat16*, __nv_bfloat16*,
+    const int*, const int*, const float*, int, int, int, int, int, cudaStream_t);
+
 template <typename T, bool FusedSwiGlu>
 void LaunchMoeGemvFp4RawNPacked(
     const T* act, const uint8_t* weight, const uint8_t* block_scales, const float* global_scales,
@@ -437,10 +502,11 @@ void LaunchMoeGemvFp4RawNPacked(
     int64_t expanded_num_rows, int64_t n, int64_t k, cutlass_kernels::ActivationParams activation_params,
     const int* permuted_row_to_source_row, int64_t num_rows, cudaStream_t stream, bool weights_row_major) {
   constexpr int kThreads = 128;
-  constexpr int kColsPerBlock = 16;
+  constexpr int kColsPerBlock = FusedSwiGlu ? 8 : 16;
   const int64_t weight_expert_stride = n * k / 2;
   const int64_t scale_expert_stride = n * (k / 16);
-  const dim3 grid(static_cast<unsigned>(expanded_num_rows), static_cast<unsigned>((n + kColsPerBlock - 1) / kColsPerBlock));
+  const int columns_per_block = weights_row_major ? kColsPerBlock : 16;
+  const dim3 grid(static_cast<unsigned>(expanded_num_rows), static_cast<unsigned>((n + columns_per_block - 1) / columns_per_block));
   if (weights_row_major) {
     if (bias != nullptr) {
       MoeGemvFp4RawKPackedKernel<T, FusedSwiGlu, true><<<grid, kThreads, 0, stream>>>(
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 ed6d6d72f3a24..695c99f5c287b 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
@@ -96,6 +96,15 @@ void launch_moe_gemv_fp4_symmetric(
     int64_t n, int64_t k, int group_size, int sm, MoeGemvConfig config, bool sm80_pair_interleaved,
     bool raw_n_packed, cudaStream_t stream, bool weights_row_major = false);
 
+// Native K-packed NVFP4 FC2 with top-k reduction; false leaves unsupported shapes to the caller.
+// The caller must use single-EP/single-TP routing. Each expert result is rounded to T before reduction.
+template <typename T>
+bool try_launch_moe_gemv_nvfp4_fused_finalize(
+    const T* act, const uint8_t* weight, const uint8_t* block_scales, const float* global_scales,
+    const T* bias, T* out, const int* unpermuted_row_to_permuted_row,
+    const int* permuted_row_to_expert, const float* routing_scales, int num_experts,
+    int num_rows, int experts_per_token, int n, int k, cudaStream_t stream);
+
 // Launches the MXFP4 MoE GEMV and fuses interleaved SwiGLU activation.
 //   weight/scales/bias use raw FC1 output width n = 2 * inter_size
 //   out is post-activation [expanded_num_rows, inter_size]
diff --git a/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_kernels.cu b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_kernels.cu
index 98b0f6b56cdcc..f86fd0ad33572 100644
--- a/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_kernels.cu
+++ b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_kernels.cu
@@ -1608,8 +1608,8 @@ __global__ void finalizeMoeRoutingKernel(const GemmOutputType* expanded_permuted
   constexpr int64_t FINALIZE_ELEM_PER_THREAD = 128 / std::min(sizeof_bits<OutputType>::value, sizeof_bits<GemmOutputType>::value);
   assert(orig_cols % FINALIZE_ELEM_PER_THREAD == 0);
 
-  const int64_t start_offset = threadIdx.x;
-  const int64_t stride = FINALIZE_THREADS_PER_BLOCK;
+  const int64_t start_offset = static_cast<int64_t>(blockIdx.y) * blockDim.x + threadIdx.x;
+  const int64_t stride = static_cast<int64_t>(blockDim.x) * gridDim.y;
   const int64_t num_elems_in_col = orig_cols / FINALIZE_ELEM_PER_THREAD;
 
   using BiasElem = cutlass::Array<ScaleBiasType, FINALIZE_ELEM_PER_THREAD>;
@@ -1628,7 +1628,7 @@ __global__ void finalizeMoeRoutingKernel(const GemmOutputType* expanded_permuted
   __shared__ int s_expert_id[STAGE_ROUTING ? FINALIZE_MAX_STAGED_TOPK : 1];
   __shared__ float s_row_scale[STAGE_ROUTING ? FINALIZE_MAX_STAGED_TOPK : 1];
   if constexpr (STAGE_ROUTING) {
-    for (int64_t k_idx = threadIdx.x; k_idx < experts_per_token; k_idx += FINALIZE_THREADS_PER_BLOCK) {
+    for (int64_t k_idx = threadIdx.x; k_idx < experts_per_token; k_idx += blockDim.x) {
       int64_t const k_offset = original_row * experts_per_token + k_idx;
       int const expert_id = token_selected_experts[k_offset] - start_expert_id;
       bool const is_local_expert = expert_id >= 0 && expert_id < num_experts_per_node;
@@ -1915,6 +1915,13 @@ void finalizeMoeRoutingKernelLauncher(const GemmOutputType* expanded_permuted_ro
       }
 #undef LAUNCH_FINALIZE_ONE_ROW
     } else if (experts_per_token <= FINALIZE_MAX_STAGED_TOPK) {
+      if (num_rows <= 2) {
+        constexpr int kDecodeThreads = 32;
+        config.blockDim = dim3(kDecodeThreads);
+        config.gridDim = dim3(static_cast<unsigned>(num_rows),
+                              static_cast<unsigned>((cols / kFinalizeElemPerThread + kDecodeThreads - 1) /
+                                                    kDecodeThreads));
+      }
       auto func = final_scales
                       ? &finalizeMoeRoutingKernel<OutputType, GemmOutputType, ScaleBiasType, ScaleMode::DEFAULT, true>
                       : &finalizeMoeRoutingKernel<OutputType, GemmOutputType, ScaleBiasType, ScaleMode::NO_SCALE, true>;
diff --git a/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8.cc b/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8.cc
index 28ef396fc8cbf..f98f3a7232f6c 100644
--- a/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8.cc
+++ b/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8.cc
@@ -48,6 +48,16 @@ MatMulBlockQuantizedFp8Weight::MatMulBlockQuantizedFp8Weight(const OpKernelInfo&
       max_dequant_scratch_bytes_(DequantScratchLimitBytes()),
       enable_deep_gemm_(ParseEnvironmentVariableWithDefault<bool>("ORT_FP8_MATMUL_DEEPGEMM", false)) {
   ORT_ENFORCE(block_size_ > 0, "block_size must be positive.");
+  if (info.GetAttrs<int64_t>("split_sizes", split_sizes_).IsOK()) {
+    ORT_ENFORCE(!split_sizes_.empty() && split_sizes_.size() <= Fp8GemvOutputSplit::kMaxOutputs,
+                "split_sizes requires one to eight widths.");
+    int64_t total = 0;
+    for (const auto width : split_sizes_) {
+      ORT_ENFORCE(width > 0 && width <= INT64_MAX - total,
+                  "split_sizes must contain positive widths with no sum overflow.");
+      total += width;
+    }
+  }
 }
 
 template <typename T>
@@ -94,27 +104,75 @@ Status MatMulBlockQuantizedFp8Weight::ComputeImpl(OpKernelContext* context) cons
   TensorShape b_logical_shape({n, k});
   ORT_RETURN_IF_ERROR(helper.Compute(a_shape, b_logical_shape, transa, transb));
 
-  Tensor* Y = context->Output(0, helper.OutputShape());
-  if (Y->Shape().Size() == 0) {
-    return Status::OK();
-  }
-
   const int m_i = SafeInt<int>(helper.M());
   const int n_i = SafeInt<int>(helper.N());
   const int k_i = SafeInt<int>(helper.K());
 
+  const bool use_gemv = m_i > 0 &&
+                        m_i <= MatMulBlockScaledFp8GemvMaxM(k_i, SafeInt<int>(block_size_), GetDeviceProp()) &&
+                        (k_i % 16 == 0) && (block_size_ % 16 == 0);
+  const bool direct_split = !split_sizes_.empty() && use_gemv &&
+                            MatMulBlockScaledFp8GemvSupportsSplitOutput(
+                                m_i, k_i, SafeInt<int>(block_size_), GetDeviceProp());
+  Fp8GemvOutputSplit split;
+  void* y_data;
+  IAllocatorUniquePtr<CudaT> y_scratch;
+  const auto output_shape = helper.OutputShape();
+  if (split_sizes_.empty()) {
+    ORT_ENFORCE(context->OutputCount() == 1, "Without split_sizes exactly one output is required.");
+    y_data = context->Output(0, output_shape)->MutableDataRaw();
+  } else {
+    ORT_ENFORCE(context->OutputCount() == static_cast<int>(split_sizes_.size()),
+                "Output count must match split_sizes.");
+    int64_t total = 0;
+    for (const auto width : split_sizes_) {
+      total += width;
+    }
+    ORT_ENFORCE(total == n, "split_sizes must sum to the weight's N dimension.");
+    split.count = static_cast<int>(split_sizes_.size());
+    for (int i = 0; i < split.count; ++i) {
+      TensorShapeVector shape(output_shape.GetDims().begin(), output_shape.GetDims().end());
+      shape.back() = split_sizes_[i];
+      split.outputs[i] = context->Output(i, TensorShape(shape))->MutableDataRaw();
+      split.widths[i] = SafeInt<int>(split_sizes_[i]);
+    }
+    if (direct_split) {
+      y_data = split.outputs[0];
+    } else {
+      y_scratch = GetScratchBuffer<CudaT>(SafeInt<size_t>(output_shape.Size()), GetComputeStream(context));
+      y_data = y_scratch.get();
+    }
+  }
+  if (output_shape.Size() == 0) {
+    return Status::OK();
+  }
+  const auto finish_output = [&]() -> Status {
+    if (split.count != 0 && !direct_split) {
+      size_t offset = 0;
+      for (int i = 0; i < split.count; ++i) {
+        const size_t width_bytes = static_cast<size_t>(split.widths[i]) * sizeof(CudaT);
+        CUDA_RETURN_IF_ERROR(cudaMemcpy2DAsync(
+            split.outputs[i], width_bytes, static_cast<const uint8_t*>(y_data) + offset,
+            static_cast<size_t>(n_i) * sizeof(CudaT), width_bytes, m_i,
+            cudaMemcpyDeviceToDevice, Stream(context)));
+        offset += width_bytes;
+      }
+    }
+    return Status::OK();
+  };
+
   if (k_i == 0) {
-    CUDA_RETURN_IF_ERROR(cudaMemsetAsync(Y->MutableDataRaw(), 0, Y->SizeInBytes(), Stream(context)));
+    CUDA_RETURN_IF_ERROR(cudaMemsetAsync(y_data, 0, SafeInt<size_t>(output_shape.Size()) * sizeof(CudaT), Stream(context)));
     if (bias != nullptr) {
       ORT_RETURN_IF_ERROR(LaunchAddBiasBlockScaledFp8(
-          Y->MutableDataRaw(),
+          y_data,
           bias->DataRaw(),
           m_i,
           n_i,
           std::is_same<T, BFloat16>::value,
           Stream(context)));
     }
-    return Status::OK();
+    return finish_output();
   }
 
   // Optional W8A8 activation path: statically quantize A to FP8 E4M3 and dequantize back so the
@@ -123,9 +181,6 @@ Status MatMulBlockQuantizedFp8Weight::ComputeImpl(OpKernelContext* context) cons
   //
   // The decode GEMV absorbs this round trip in-register, so the standalone kernel (and its [M, K]
   // scratch buffer) is only materialized for the cuBLAS path below.
-  const bool use_gemv = m_i > 0 &&
-                        m_i <= MatMulBlockScaledFp8GemvMaxM(k_i, SafeInt<int>(block_size_), GetDeviceProp()) &&
-                        (k_i % 16 == 0) && (block_size_ % 16 == 0);
   const bool fuse_act_qdq = use_gemv && !FusedFp8ActivationQdqDisabled();
 
 #if defined(USE_DEEP_GEMM)
@@ -154,14 +209,14 @@ Status MatMulBlockQuantizedFp8Weight::ComputeImpl(OpKernelContext* context) cons
         const size_t weight_offset = SafeInt<size_t>(offset) * k_i;
         const size_t scale_offset = SafeInt<size_t>(offset) * (k_i / 128);
         ORT_RETURN_IF_ERROR(LaunchMatMulFp8DeepGemm(
-            Y->MutableData<T>() + offset, a_quant.get(), a_scales.get(),
+            static_cast<T*>(y_data) + offset, a_quant.get(), a_scales.get(),
             static_cast<const uint8_t*>(b->DataRaw()) + weight_offset,
             b_scale->Data<float>() + scale_offset,
             bias != nullptr ? bias->Data<T>() + offset : nullptr,
             packed_b_scales.get(), accum.get(), m_i, rows, k_i, aligned_m, n_i,
             std::is_same<T, BFloat16>::value, GetDeviceProp().multiProcessorCount, Stream(context)));
       }
-      return Status::OK();
+      return finish_output();
     }
   }
 #endif
@@ -186,8 +241,8 @@ Status MatMulBlockQuantizedFp8Weight::ComputeImpl(OpKernelContext* context) cons
   // A fused warp-per-column kernel reads the FP8 weight directly, avoiding both the [N, K]
   // dequant scratch buffer and the cuBLAS GEMM (which is underutilized at M == 1).
   if (use_gemv) {
-    return LaunchMatMulBlockScaledFp8Gemv(
-        Y->MutableDataRaw(),
+    ORT_RETURN_IF_ERROR(LaunchMatMulBlockScaledFp8Gemv(
+        y_data,
         a_ptr,
         b->DataRaw(),
         b_scale->Data<float>(),
@@ -199,7 +254,8 @@ Status MatMulBlockQuantizedFp8Weight::ComputeImpl(OpKernelContext* context) cons
         SafeInt<int>(block_size_),
         std::is_same<T, BFloat16>::value,
         GetDeviceProp(),
-        Stream(context));
+        Stream(context), direct_split ? split : Fp8GemvOutputSplit{}));
+    return finish_output();
   }
 
   // Dequantize the FP8 weight into a scratch buffer of the activation type, then GEMM. The scratch
@@ -243,7 +299,7 @@ Status MatMulBlockQuantizedFp8Weight::ComputeImpl(OpKernelContext* context) cons
         reinterpret_cast<const CudaT*>(a_ptr),
         helper.Lda(transa),
         &zero,
-        reinterpret_cast<CudaT*>(Y->MutableDataRaw()) + n_offset,
+        static_cast<CudaT*>(y_data) + n_offset,
         helper.Ldc(),
         GetDeviceProp(),
         UseTF32()));
@@ -251,7 +307,7 @@ Status MatMulBlockQuantizedFp8Weight::ComputeImpl(OpKernelContext* context) cons
 
   if (bias != nullptr) {
     ORT_RETURN_IF_ERROR(LaunchAddBiasBlockScaledFp8(
-        Y->MutableDataRaw(),
+        y_data,
         bias->DataRaw(),
         m_i,
         n_i,
@@ -259,7 +315,7 @@ Status MatMulBlockQuantizedFp8Weight::ComputeImpl(OpKernelContext* context) cons
         Stream(context)));
   }
 
-  return Status::OK();
+  return finish_output();
 #endif
 }
 
diff --git a/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8.cu b/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8.cu
index d045519626b27..599a8ec957852 100644
--- a/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8.cu
+++ b/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8.cu
@@ -514,7 +514,20 @@ struct Fp8GemvMma<__nv_bfloat16> {
   }
 };
 
-template <int KSplit, int MTiles, typename AType>
+template <typename T>
+__device__ __forceinline__ T* Fp8OutputColumn(T* output, int col, int& stride,
+                                             const Fp8GemvOutputSplit& split) {
+  for (int i = 0; i < split.count; ++i) {
+    if (col < split.widths[i]) {
+      stride = split.widths[i];
+      return static_cast<T*>(split.outputs[i]) + col;
+    }
+    col -= split.widths[i];
+  }
+  return output + col;
+}
+
+template <int KSplit, int MTiles, bool SplitOutput, typename AType>
 __device__ __forceinline__ void Fp8MmaGemvBody(AType* __restrict__ output,
                                                const AType* __restrict__ input_a,
                                                const __nv_fp8_e4m3* __restrict__ input_b,
@@ -525,7 +538,8 @@ __device__ __forceinline__ void Fp8MmaGemvBody(AType* __restrict__ output,
                                                int n,
                                                int k,
                                                int block_size,
-                                               int k_blocks) {
+                                               int k_blocks,
+                                               const Fp8GemvOutputSplit& split) {
   using Mma = Fp8GemvMma<AType>;
 
   const bool act_qdq = act_scale != nullptr;
@@ -560,6 +574,7 @@ __device__ __forceinline__ void Fp8MmaGemvBody(AType* __restrict__ output,
   float accb[MTiles][4] = {};  // unscaled, within the current K block
   int cur_kb = -1;
 
+#pragma unroll 2
   for (int wi = warp; wi < windows; wi += KSplit) {
     const int k0 = wi << 6;
     const int kb = k0 / block_size;
@@ -670,23 +685,31 @@ __device__ __forceinline__ void Fp8MmaGemvBody(AType* __restrict__ output,
   // M tile.
   const float bias_lo = (bias != nullptr && lo_ok) ? to_float<AType>(bias[col_lo]) : 0.f;
   const float bias_hi = (bias != nullptr && hi_ok) ? to_float<AType>(bias[col_hi]) : 0.f;
+  int stride_lo = n;
+  int stride_hi = n;
+  AType* output_lo = output + col_lo;
+  AType* output_hi = output + col_hi;
+  if constexpr (SplitOutput) {
+    output_lo = lo_ok ? Fp8OutputColumn(output, col_lo, stride_lo, split) : nullptr;
+    output_hi = hi_ok ? Fp8OutputColumn(output, col_hi, stride_hi, split) : nullptr;
+  }
 #pragma unroll
   for (int mt = 0; mt < MTiles; ++mt) {
     const int row = (t << 1) + (mt << 3);
     if (row < m) {
       if (lo_ok) {
-        output[static_cast<size_t>(row) * n + col_lo] = from_float<AType>(acc[mt][0] + bias_lo);
+        output_lo[static_cast<size_t>(row) * stride_lo] = from_float<AType>(acc[mt][0] + bias_lo);
       }
       if (hi_ok) {
-        output[static_cast<size_t>(row) * n + col_hi] = from_float<AType>(acc[mt][2] + bias_hi);
+        output_hi[static_cast<size_t>(row) * stride_hi] = from_float<AType>(acc[mt][2] + bias_hi);
       }
     }
     if (row + 1 < m) {
       if (lo_ok) {
-        output[static_cast<size_t>(row + 1) * n + col_lo] = from_float<AType>(acc[mt][1] + bias_lo);
+        output_lo[static_cast<size_t>(row + 1) * stride_lo] = from_float<AType>(acc[mt][1] + bias_lo);
       }
       if (hi_ok) {
-        output[static_cast<size_t>(row + 1) * n + col_hi] = from_float<AType>(acc[mt][3] + bias_hi);
+        output_hi[static_cast<size_t>(row + 1) * stride_hi] = from_float<AType>(acc[mt][3] + bias_hi);
       }
     }
   }
@@ -695,6 +718,12 @@ __device__ __forceinline__ void Fp8MmaGemvBody(AType* __restrict__ output,
 // Two entry points over one body. The pinned one carries a residency hint; see
 // `Fp8MmaGemvPinsResidency` for when the launcher picks it and why the plain one has to stay.
 // clang-format off
+#if CUDA_VERSION >= 12000
+// Avoid a per-thread local copy when the inlined store helper references this descriptor.
+#define ORT_FP8_SPLIT_PARAM const __grid_constant__ Fp8GemvOutputSplit split
+#else
+#define ORT_FP8_SPLIT_PARAM const Fp8GemvOutputSplit split
+#endif
 #define ORT_FP8_MMA_GEMV_PARAMS                  \
   AType* __restrict__ output,                    \
       const AType* __restrict__ input_a,         \
@@ -702,24 +731,25 @@ __device__ __forceinline__ void Fp8MmaGemvBody(AType* __restrict__ output,
       const float* __restrict__ weight_scale,    \
       const AType* __restrict__ bias,            \
       const float* __restrict__ act_scale,       \
-      int m, int n, int k, int block_size, int k_blocks
+      int m, int n, int k, int block_size, int k_blocks, ORT_FP8_SPLIT_PARAM
 
 #define ORT_FP8_MMA_GEMV_ARGS \
-  output, input_a, input_b, weight_scale, bias, act_scale, m, n, k, block_size, k_blocks
+  output, input_a, input_b, weight_scale, bias, act_scale, m, n, k, block_size, k_blocks, split
 // clang-format on
 
-template <int KSplit, int MTiles, typename AType>
+template <int KSplit, int MTiles, bool SplitOutput, typename AType>
 __global__ void MatMulBlockScaledFp8MmaGemvKernel(ORT_FP8_MMA_GEMV_PARAMS) {
-  Fp8MmaGemvBody<KSplit, MTiles, AType>(ORT_FP8_MMA_GEMV_ARGS);
+  Fp8MmaGemvBody<KSplit, MTiles, SplitOutput, AType>(ORT_FP8_MMA_GEMV_ARGS);
 }
 
-template <int KSplit, int MTiles, typename AType>
+template <int KSplit, int MTiles, bool SplitOutput, typename AType>
 __global__ __launch_bounds__(32 * KSplit, 3) void MatMulBlockScaledFp8MmaGemvKernelPinned(ORT_FP8_MMA_GEMV_PARAMS) {
-  Fp8MmaGemvBody<KSplit, MTiles, AType>(ORT_FP8_MMA_GEMV_ARGS);
+  Fp8MmaGemvBody<KSplit, MTiles, SplitOutput, AType>(ORT_FP8_MMA_GEMV_ARGS);
 }
 
 #undef ORT_FP8_MMA_GEMV_ARGS
 #undef ORT_FP8_MMA_GEMV_PARAMS
+#undef ORT_FP8_SPLIT_PARAM
 
 // Kill switch for A/B testing the tensor-core path against the FMA path in the same binary.
 bool Fp8GemvMmaEnabled() {
@@ -915,6 +945,20 @@ bool Fp8MmaGb10TuningEnabled() {
   return enabled;
 }
 
+bool MatMulBlockScaledFp8GemvSupportsSplitOutput(int m, int k, int block_size,
+                                                const cudaDeviceProp& device_prop) {
+#if !defined(DISABLE_FLOAT8_TYPES) && defined(CUDA_VERSION) && CUDA_VERSION >= 11080
+  return device_prop.major >= 8 && m > 0 && m <= kFp8MmaGemvTileM &&
+         k >= 256 && k % 64 == 0 && block_size % 64 == 0 && Fp8GemvMmaEnabled();
+#else
+  ORT_UNUSED_PARAMETER(m);
+  ORT_UNUSED_PARAMETER(k);
+  ORT_UNUSED_PARAMETER(block_size);
+  ORT_UNUSED_PARAMETER(device_prop);
+  return false;
+#endif
+}
+
 static Status LaunchMatMulBlockScaledFp8GemvImpl(void* y,
                                                  const void* a,
                                                  const void* b_fp8,
@@ -928,7 +972,8 @@ static Status LaunchMatMulBlockScaledFp8GemvImpl(void* y,
                                                  bool is_bf16,
                                                  const cudaDeviceProp& device_prop,
                                                  cudaStream_t stream,
-                                                 bool enable_gb10_ksplit32) {
+                                                 bool enable_gb10_ksplit32,
+                                                 const Fp8GemvOutputSplit& split) {
 #if !defined(DISABLE_FLOAT8_TYPES) && defined(CUDA_VERSION) && CUDA_VERSION >= 11080
   if (m <= 0 || n <= 0 || k <= 0) {
     return Status::OK();
@@ -941,6 +986,19 @@ static Status LaunchMatMulBlockScaledFp8GemvImpl(void* y,
                     "MatMulBlockQuantizedFp8Weight GEMV requires block_size divisible by 16, got ", block_size, ".");
   const int k_blocks = (k + block_size - 1) / block_size;
   const auto* b = reinterpret_cast<const __nv_fp8_e4m3*>(b_fp8);
+  ORT_RETURN_IF_NOT(split.count >= 0 && split.count <= Fp8GemvOutputSplit::kMaxOutputs,
+                    "FP8 GEMV split output count must be between zero and eight.");
+  if (split.count != 0) {
+    int64_t total = 0;
+    for (int i = 0; i < split.count; ++i) {
+      ORT_RETURN_IF_NOT(split.widths[i] > 0 && split.outputs[i] != nullptr,
+                        "FP8 GEMV split outputs require positive widths and non-null buffers.");
+      total += split.widths[i];
+    }
+    ORT_RETURN_IF_NOT(total == n, "FP8 GEMV split widths must sum to N.");
+  }
+  ORT_RETURN_IF_NOT(split.count == 0 || MatMulBlockScaledFp8GemvSupportsSplitOutput(m, k, block_size, device_prop),
+                    "Direct split output requires the tensor-core GEMV path with M <= 32.");
 
   if (m > kFp8MmaGemvTileM) {
     ORT_RETURN_IF_NOT(device_prop.major >= 8 && k % 64 == 0 && k >= 256 &&
@@ -950,12 +1008,12 @@ static Status LaunchMatMulBlockScaledFp8GemvImpl(void* y,
     const size_t element_size = is_bf16 ? sizeof(__nv_bfloat16) : sizeof(half);
     ORT_RETURN_IF_ERROR(LaunchMatMulBlockScaledFp8GemvImpl(
         y, a, b_fp8, weight_scale, bias, act_scale, kFp8MmaGemvTileM, n, k, block_size,
-        is_bf16, device_prop, stream, false));
+        is_bf16, device_prop, stream, false, {}));
     return LaunchMatMulBlockScaledFp8GemvImpl(
         static_cast<uint8_t*>(y) + static_cast<size_t>(kFp8MmaGemvTileM) * n * element_size,
         static_cast<const uint8_t*>(a) + static_cast<size_t>(kFp8MmaGemvTileM) * k * element_size,
         b_fp8, weight_scale, bias, act_scale, m - kFp8MmaGemvTileM, n, k, block_size,
-        is_bf16, device_prop, stream, false);
+        is_bf16, device_prop, stream, false, {});
   }
 
   // Tensor-core path (SM80+). Beats the FMA kernel at every M on H200: 1.06-1.23x at M == 1 and
@@ -982,28 +1040,35 @@ static Status LaunchMatMulBlockScaledFp8GemvImpl(void* y,
         n, k_split, mtiles, device_prop.multiProcessorCount, device_prop.major, device_prop.minor);
     const auto launch_mma = [&]<int KSplit, int MTiles>() {
       const dim3 mma_threads{32, KSplit};
-#define ORT_FP8_LAUNCH_MMA(kernel_name)                                                      \
-  do {                                                                                       \
-    if (is_bf16) {                                                                           \
-      kernel_name<KSplit, MTiles><<<mma_blocks, mma_threads, 0, stream>>>(                   \
-          reinterpret_cast<__nv_bfloat16*>(y), reinterpret_cast<const __nv_bfloat16*>(a), b, \
-          weight_scale, reinterpret_cast<const __nv_bfloat16*>(bias), act_scale, m, n, k,    \
-          block_size, k_blocks);                                                             \
-    } else {                                                                                 \
-      kernel_name<KSplit, MTiles><<<mma_blocks, mma_threads, 0, stream>>>(                   \
-          reinterpret_cast<half*>(y), reinterpret_cast<const half*>(a), b,                   \
-          weight_scale, reinterpret_cast<const half*>(bias), act_scale, m, n, k,             \
-          block_size, k_blocks);                                                             \
-    }                                                                                        \
+      const auto launch_variant = [&]<bool SplitOutput>() {
+#define ORT_FP8_LAUNCH_MMA(kernel_name)                                                       \
+  do {                                                                                      \
+    if (is_bf16) {                                                                          \
+      kernel_name<KSplit, MTiles, SplitOutput><<<mma_blocks, mma_threads, 0, stream>>>(       \
+          reinterpret_cast<__nv_bfloat16*>(y), reinterpret_cast<const __nv_bfloat16*>(a), b,  \
+          weight_scale, reinterpret_cast<const __nv_bfloat16*>(bias), act_scale, m, n, k,     \
+          block_size, k_blocks, split);                                                     \
+    } else {                                                                               \
+      kernel_name<KSplit, MTiles, SplitOutput><<<mma_blocks, mma_threads, 0, stream>>>(       \
+          reinterpret_cast<half*>(y), reinterpret_cast<const half*>(a), b,                    \
+          weight_scale, reinterpret_cast<const half*>(bias), act_scale, m, n, k,              \
+          block_size, k_blocks, split);                                                     \
+    }                                                                                      \
   } while (0)
-      if constexpr (KSplit == 16 && MTiles == 1) {
-        if (pin_residency) {
-          ORT_FP8_LAUNCH_MMA(MatMulBlockScaledFp8MmaGemvKernelPinned);
-          return;
+        if constexpr (KSplit == 16 && MTiles == 1) {
+          if (pin_residency) {
+            ORT_FP8_LAUNCH_MMA(MatMulBlockScaledFp8MmaGemvKernelPinned);
+            return;
+          }
         }
-      }
-      ORT_FP8_LAUNCH_MMA(MatMulBlockScaledFp8MmaGemvKernel);
+        ORT_FP8_LAUNCH_MMA(MatMulBlockScaledFp8MmaGemvKernel);
 #undef ORT_FP8_LAUNCH_MMA
+      };
+      if (split.count) {
+        launch_variant.template operator()<true>();
+      } else {
+        launch_variant.template operator()<false>();
+      }
     };
     // Only 1, 2 and 4 row tiles are instantiated; an M of 17..24 rounds up to 4 and masks the
     // remainder, which costs nothing next to the weight traffic it shares.
@@ -1109,6 +1174,7 @@ static Status LaunchMatMulBlockScaledFp8GemvImpl(void* y,
   ORT_UNUSED_PARAMETER(device_prop);
   ORT_UNUSED_PARAMETER(stream);
   ORT_UNUSED_PARAMETER(enable_gb10_ksplit32);
+  ORT_UNUSED_PARAMETER(split);
   return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "MatMulBlockQuantizedFp8Weight requires CUDA 11.8 or later.");
 #endif
 }
@@ -1125,10 +1191,11 @@ Status LaunchMatMulBlockScaledFp8Gemv(void* y,
                                       int block_size,
                                       bool is_bf16,
                                       const cudaDeviceProp& device_prop,
-                                      cudaStream_t stream) {
+                                      cudaStream_t stream,
+                                      const Fp8GemvOutputSplit& split) {
   return LaunchMatMulBlockScaledFp8GemvImpl(
       y, a, b_fp8, weight_scale, bias, act_scale, m, n, k, block_size,
-      is_bf16, device_prop, stream, true);
+      is_bf16, device_prop, stream, true, split);
 }
 
 }  // namespace onnxruntime::contrib::cuda
diff --git a/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8.h b/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8.h
index 7128b86971dcc..35e2245543313 100644
--- a/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8.h
+++ b/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8.h
@@ -4,6 +4,7 @@
 #pragma once
 
 #include <cstddef>
+#include <vector>
 
 #include "core/providers/cuda/cuda_kernel.h"
 
@@ -31,10 +32,21 @@ class MatMulBlockQuantizedFp8Weight final : public onnxruntime::cuda::CudaKernel
   Status ComputeImpl(OpKernelContext* context) const;
 
   int64_t block_size_;
+  std::vector<int64_t> split_sizes_;
   size_t max_dequant_scratch_bytes_;
   bool enable_deep_gemm_;
 };
 
+struct Fp8GemvOutputSplit {
+  static constexpr int kMaxOutputs = 8;
+  void* outputs[kMaxOutputs]{};
+  int widths[kMaxOutputs]{};
+  int count = 0;
+};
+
+bool MatMulBlockScaledFp8GemvSupportsSplitOutput(int m, int k, int block_size,
+                                                const cudaDeviceProp& device_prop);
+
 #if defined(USE_DEEP_GEMM)
 Status LaunchPrepareMatMulFp8DeepGemm(void* a_quant, float* a_scales, const void* a,
                                       const float* a_scale, int m, int k, int aligned_m,
@@ -83,6 +95,8 @@ Status LaunchQuantizeDequantizeActivationFp8(void* a_out,
 // directly (no [N, K] dequant buffer). a is [M, K] activation (FP16/BF16), b_fp8 is [N, K]
 // FP8 E4M3, weight_scale is [N, ceil(K/block_size)] fp32, bias is an optional [N] vector (may be
 // null). Output y is [M, N] in the activation type. Requires k % 16 == 0 and block_size % 16 == 0.
+// A nonempty split descriptor routes tensor-core stores into consecutive column partitions;
+// callers must check MatMulBlockScaledFp8GemvSupportsSplitOutput before using that mode.
 // device_prop selects the tensor-core (mma.m16n8k16) variant, which needs SM80+.
 // act_scale is an optional device fp32 scalar; when non-null the kernel applies the W8A8
 // activation quantize/dequantize inline, bit-identically to LaunchQuantizeDequantizeActivationFp8,
@@ -100,7 +114,8 @@ Status LaunchMatMulBlockScaledFp8Gemv(void* y,
                                       int block_size,
                                       bool is_bf16,
                                       const cudaDeviceProp& device_prop,
-                                      cudaStream_t stream);
+                                      cudaStream_t stream,
+                                      const Fp8GemvOutputSplit& split = {});
 
 // Largest M selected for GEMV dispatch. The tensor-core default is conservative because the
 // crossover with dequantize + cuBLAS depends on the matrix shape; ORT_FP8_GEMV_MAX_M can raise the
diff --git a/onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc b/onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc
index a6b9f0cd36fa4..ac7c00d848ff6 100644
--- a/onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc
+++ b/onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc
@@ -2460,7 +2460,10 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const {
     IAllocatorUniquePtr<int> blocked_expert_counts;
     IAllocatorUniquePtr<int> blocked_expert_counts_cumsum;
     IAllocatorUniquePtr<int> blocked_row_to_unpermuted_row;
-    if (!ck::fusedBuildExpertMapsSortFirstToken(
+    if (!TryLaunchSmallDecodeExpertMaps(
+            expert_indices, p_r2u, unpermuted_row_to_permuted_row, p_exp, p_efto,
+            static_cast<int>(num_rows), num_experts, static_cast<int>(k_), stream) &&
+        !ck::fusedBuildExpertMapsSortFirstToken(
             expert_indices, p_r2u, unpermuted_row_to_permuted_row, p_exp, p_efto,
             num_rows, num_experts, static_cast<int>(k_), 0, num_experts, stream)) {
       const int64_t num_tokens_per_block = ck::computeNumTokensPerBlock(num_rows, num_experts);
@@ -2629,6 +2632,16 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const {
       }
 
       launch_fc1(fc1_config);
+      if (use_raw_nvfp4_gemv && nvfp4_weights_row_major_ &&
+          parallelism_config.tp_size == 1 && parallelism_config.ep_size == 1 &&
+          gemv::try_launch_moe_gemv_nvfp4_fused_finalize<T>(
+              static_cast<const T*>(p_fc1_buf.get()), gemv_fc2_weight, gemv_fc2_raw_block_scales,
+              gemv_fc2_raw_global_scales, static_cast<const T*>(fc2_bias),
+              static_cast<T*>(output->MutableDataRaw()), unpermuted_row_to_permuted_row,
+              p_exp, expert_scales, num_experts, static_cast<int>(num_rows), static_cast<int>(k_),
+              static_cast<int>(hidden), static_cast<int>(inter), stream)) {
+        return;
+      }
       launch_fc2(fc2_config);
       ck::finalizeMoeRoutingKernelLauncher<T, T, T>(
           static_cast<const T*>(p_fc2_buf.get()), static_cast<T*>(output->MutableDataRaw()),
diff --git a/onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.cu b/onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.cu
index 5d4c2dd3a4a19..c6342f99fb6e6 100644
--- a/onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.cu
+++ b/onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.cu
@@ -220,6 +220,42 @@ __global__ void SoftmaxTopKMergeKernel(const T* logits, float* topk_scales, int*
   }
 }
 
+__global__ void SmallDecodeExpertMapsKernel(
+    const int* indices, int* row_to_source, int* source_to_row, int* sorted_experts,
+    int64_t* expert_offsets, int rows, int experts, int k) {
+  __shared__ int selected[32];
+  const int expanded = rows * k;
+  if (threadIdx.x < expanded) selected[threadIdx.x] = indices[threadIdx.x];
+  __syncthreads();
+  if (threadIdx.x < expanded) {
+    const int original = threadIdx.x;
+    const int expert = selected[original];
+    int rank = 0;
+    for (int i = 0; i < expanded; ++i) {
+      rank += selected[i] < expert || (selected[i] == expert && i < original);
+    }
+    const int source = (original % k) * rows + original / k;
+    row_to_source[rank] = source;
+    source_to_row[source] = rank;
+    sorted_experts[rank] = expert;
+  }
+  for (int expert = threadIdx.x; expert <= experts; expert += blockDim.x) {
+    int count = 0;
+    for (int i = 0; i < expanded; ++i) count += selected[i] < expert;
+    expert_offsets[expert] = count;
+  }
+}
+
+bool TryLaunchSmallDecodeExpertMaps(
+    const int* indices, int* row_to_source, int* source_to_row, int* sorted_experts,
+    int64_t* expert_offsets, int rows, int experts, int k, cudaStream_t stream) {
+  if (rows < 1 || rows > 2 || experts < 1 || experts > 512 || k < 1 || k > 16 || k > experts) return false;
+  SmallDecodeExpertMapsKernel<<<1, 128, 0, stream>>>(
+      indices, row_to_source, source_to_row, sorted_experts, expert_offsets, rows, experts, k);
+  CUDA_CALL_THROW(cudaGetLastError());
+  return true;
+}
+
 // Warp-bitonic softmax + top-k for num_experts <= 32. Each warp handles one
 // row, with lane `l` owning expert `l`. The whole softmax reduction and the
 // sort are done with warp shuffles (no shared memory). This is the fastest path
diff --git a/onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.h b/onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.h
index 7bb9b771a92d2..bb767699ce29d 100644
--- a/onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.h
+++ b/onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.h
@@ -12,6 +12,10 @@ namespace onnxruntime {
 namespace contrib {
 namespace cuda {
 
+bool TryLaunchSmallDecodeExpertMaps(
+    const int* indices, int* row_to_source, int* source_to_row, int* sorted_experts,
+    int64_t* expert_offsets, int rows, int experts, int k, cudaStream_t stream);
+
 void LaunchSoftmaxTopK(
     const float* logits,
     float* topk_scales,
diff --git a/onnxruntime/contrib_ops/cuda/quantization/quantize_dequantize_linear.cc b/onnxruntime/contrib_ops/cuda/quantization/quantize_dequantize_linear.cc
index f612a9b6a6f54..40f54aba301cf 100644
--- a/onnxruntime/contrib_ops/cuda/quantization/quantize_dequantize_linear.cc
+++ b/onnxruntime/contrib_ops/cuda/quantization/quantize_dequantize_linear.cc
@@ -34,7 +34,8 @@ REGISTER_Q_KERNEL_TYPED(uint8_t, MLFloat16)
       kCudaExecutionProvider,                                      \
       (*KernelDefBuilder::Create())                                \
           .TypeConstraint("T1", DataTypeImpl::GetTensorType<T>())  \
-          .TypeConstraint("T2", DataTypeImpl::GetTensorType<U>()), \
+          .TypeConstraint("T2", DataTypeImpl::GetTensorType<U>())  \
+          .TypeConstraint("TS", DataTypeImpl::GetTensorType<U>()), \
       DequantizeLinear<T, U>);
 
 REGISTER_DQ_KERNEL_TYPED(int8_t, MLFloat16)
diff --git a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc
index abc7e84188962..2fe6a490ee90a 100644
--- a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc
+++ b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc
@@ -3189,9 +3189,13 @@ are `A_scaled = fp8_e4m3(A / a_scale) * a_scale` (W8A8). Otherwise, A retains it
 precision (weight-only W8A16).
 
 The operator multiplies the activation by the transpose of B_scaled and adds the optional bias.
-The output has shape [..., N] and the same element type as A.)DOC")
+The output has shape [..., N] and the same element type as A. When split_sizes is provided,
+return up to eight consecutive last-dimension partitions instead, equivalent to Split(Y, split_sizes).
+The partition widths must be positive and sum to N.)DOC")
         .Attr("block_size", "Number of consecutive K values that share one weight scale. Default 128.",
               AttributeProto::INT, static_cast<int64_t>(128))
+        .Attr("split_sizes", "Optional last-dimension output widths (one to eight positive sizes).",
+              AttributeProto::INTS, OPTIONAL_VALUE)
         .Input(0, "A", "Row-major FP16/BF16 activation of shape [..., K].", "T")
         .Input(1, "B", "Row-major FP8 E4M3 weight of shape [N, K].", "T1")
         .Input(2, "b_scale", "Per-block FP32 weight scales of shape [N, ceil(K / block_size)].", "T2")
@@ -3201,13 +3205,31 @@ The output has shape [..., N] and the same element type as A.)DOC")
                "its FP16/BF16 precision.",
                "T2", OpSchema::Optional)
         .Input(4, "bias", "Optional bias of shape [N].", "T", OpSchema::Optional)
-        .Output(0, "Y", "Output of shape [..., N] in the activation type.", "T")
+        .Output(0, "Y", "Output(s) of shape [..., N] or [..., split_sizes[i]] in the activation type.",
+                "T", OpSchema::Variadic)
         .TypeConstraint("T", {"tensor(float16)", "tensor(bfloat16)"},
                         "Constrain activation, bias and output to FP16 or BF16.")
         .TypeConstraint("T1", {"tensor(float8e4m3fn)"}, "Constrain weight to FP8 E4M3.")
         .TypeConstraint("T2", {"tensor(float)"}, "Constrain scales to FP32.")
         .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) {
-          propagateElemTypeFromInputToOutput(ctx, 0, 0);
+          const auto* split = ctx.getAttribute("split_sizes");
+          if ((split == nullptr && ctx.getNumOutputs() != 1) ||
+              (split != nullptr && (split->ints_size() < 1 || split->ints_size() > 8 ||
+                                   static_cast<size_t>(split->ints_size()) != ctx.getNumOutputs()))) {
+            fail_shape_inference("Output count must match split_sizes, or be one when split_sizes is absent.");
+          }
+          int64_t split_total = 0;
+          if (split != nullptr) {
+            for (const auto width : split->ints()) {
+              if (width <= 0 || width > INT64_MAX - split_total) {
+                fail_shape_inference("split_sizes must contain positive widths with no sum overflow.");
+              }
+              split_total += width;
+            }
+          }
+          for (size_t i = 0; i < ctx.getNumOutputs(); ++i) {
+            propagateElemTypeFromInputToOutput(ctx, 0, i);
+          }
           if (!hasNInputShapes(ctx, 2)) {
             return;
           }
@@ -3225,7 +3247,16 @@ The output has shape [..., N] and the same element type as A.)DOC")
             *output_shape.add_dim() = a_shape.dim(i);
           }
           *output_shape.add_dim() = b_shape.dim(0);
-          updateOutputShape(ctx, 0, output_shape);
+          if (split != nullptr && b_shape.dim(0).has_dim_value() &&
+              split_total != b_shape.dim(0).dim_value()) {
+            fail_shape_inference("split_sizes must sum to the weight's N dimension.");
+          }
+          for (size_t i = 0; i < ctx.getNumOutputs(); ++i) {
+            if (split != nullptr) {
+              output_shape.mutable_dim(output_shape.dim_size() - 1)->set_dim_value(split->ints(static_cast<int>(i)));
+            }
+            updateOutputShape(ctx, i, output_shape);
+          }
         }));
 
 static void MatmulWithQuantWeightShapeInference(ONNX_NAMESPACE::InferenceContext& ctx,
diff --git a/onnxruntime/core/graph/contrib_ops/quantization_defs.cc b/onnxruntime/core/graph/contrib_ops/quantization_defs.cc
index 70cd688832b02..46bfebf9ec693 100644
--- a/onnxruntime/core/graph/contrib_ops/quantization_defs.cc
+++ b/onnxruntime/core/graph/contrib_ops/quantization_defs.cc
@@ -182,12 +182,16 @@ ONNX_MS_OPERATOR_SET_SCHEMA(
         }));
 
 static const char* DequantizeLinear_ver1_doc = R"DOC(
-The linear dequantization operator. It consumes a quantized data, a scale, a zero point and computes the full precision data.
-The dequantization formula is y = (x - x_zero_point) * x_scale.
-Scale and zero point must have same shape. They must be either scalar (per tensor) or 1-D tensor (per 'axis').)DOC";
+The linear dequantization operator. It consumes quantized data, a scale, and an optional zero point.
+The dequantization formula is y = (x - x_zero_point) * x_scale. Scale and zero point are scalar
+or per-axis vectors. FP8 E4M3FN inputs can be dequantized directly to BF16 with FP32 scales.)DOC";
 
 ONNX_MS_OPERATOR_SET_SCHEMA(DequantizeLinear, 1,
                             OpSchema()
+                                .Attr("output_dtype",
+                                      "Optional output element type. When omitted, the scale "
+                                      "element type determines the output type.",
+                                      AttributeProto::INT, false)
                                 .Attr("axis",
                                       "The axis along which same quantization parameters are applied. It's optional."
                                       "If it's not specified, it means per-tensor quantization and input 'x_scale' and "
@@ -199,25 +203,36 @@ ONNX_MS_OPERATOR_SET_SCHEMA(DequantizeLinear, 1,
                                 .Input(1, "x_scale",
                                        "Scale for input 'x'. It can be a scalar, which means a per-tensor/layer "
                                        "dequantization, or a 1-D tensor for per-axis dequantization.",
-                                       "T2")
+                                       "TS")
                                 .Input(2, "x_zero_point",
                                        "Zero point for input 'x'. Shape must match x_scale. It's optional. "
                                        "Zero point is 0 when it's not specified.",
                                        "T1", OpSchema::Optional)
                                 .Output(0, "y", "N-D full precision output tensor. It has same shape as input 'x'.",
                                         "T2")
-                                .TypeConstraint("T1", {"tensor(int8)", "tensor(uint8)", "tensor(int16)",
-                                                       "tensor(uint16)", "tensor(int32)", "tensor(int4)",
-                                                       "tensor(uint4)"},
+                                .TypeConstraint("T1", {"tensor(int8)", "tensor(uint8)", "tensor(int16)", "tensor(uint16)", "tensor(int32)", "tensor(int4)", "tensor(uint4)", "tensor(float8e4m3fn)"},
                                                 "Constrain 'x' and 'x_zero_point' to 8-bit integer tensors, "
-                                                "16-bit integer tensors, or 32-bit signed integer tensors.")
-                                .TypeConstraint("T2", {"tensor(float16)", "tensor(float)"},
-                                                "Constrain 'y', 'x_scale' to float tensors.")
+                                                "16-bit integer tensors, 32-bit signed integer tensors, or "
+                                                "FP8 E4M3FN tensors.")
+                                .TypeConstraint("T2", {"tensor(float16)", "tensor(float)", "tensor(bfloat16)"},
+                                                "Constrain the output element type.")
+                                .TypeConstraint("TS", {"tensor(float16)", "tensor(float)", "tensor(bfloat16)"},
+                                                "Constrain the scale element type.")
                                 .SetDoc(DequantizeLinear_ver1_doc)
                                 .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) {
-                                  auto y_type = ctx.getOutputType(0);
-                                  // only float is supported
-                                  y_type->mutable_tensor_type()->set_elem_type(ONNX_NAMESPACE::TensorProto::FLOAT);
+                                  const int64_t output_dtype =
+                                      getAttribute(ctx, "output_dtype", static_cast<int64_t>(0));
+                                  if (output_dtype != 0) {
+                                    if (output_dtype != ONNX_NAMESPACE::TensorProto::FLOAT &&
+                                        output_dtype != ONNX_NAMESPACE::TensorProto::FLOAT16 &&
+                                        output_dtype != ONNX_NAMESPACE::TensorProto::BFLOAT16) {
+                                      fail_shape_inference(
+                                          "DequantizeLinear: output_dtype must be float, float16, or bfloat16");
+                                    }
+                                    updateOutputElemType(ctx, 0, static_cast<int32_t>(output_dtype));
+                                  } else if (ctx.getInputType(1) != nullptr) {
+                                    propagateElemTypeFromInputToOutput(ctx, 1, 0);
+                                  }
 
                                   if (!hasInputShape(ctx, 0)) return;
 
@@ -852,16 +867,16 @@ ONNX_MS_OPERATOR_SET_SCHEMA(
         {"tensor(uint8)", "tensor(int8)"},
         "Constrain input and output types to 8 bit signed and unsigned tensors.")
       .TypeAndShapeInferenceFunction([](InferenceContext& ctx) {
-        propagateElemTypeFromInputToOutput(ctx, 1, 0);
-        if (hasNInputShapes(ctx, 9)) {
-          std::vector<const onnx::TensorShapeProto*> shapes;
-          shapes.push_back(&ctx.getInputType(0)->tensor_type().shape());
-          shapes.push_back(&ctx.getInputType(1)->tensor_type().shape());
-          shapes.push_back(&ctx.getInputType(4)->tensor_type().shape());
-          multidirectionalBroadcastShapeInference(
-              shapes, *ctx.getOutputType(0)->mutable_tensor_type()->mutable_shape());
-        }
-      }));
+  propagateElemTypeFromInputToOutput(ctx, 1, 0);
+  if (hasNInputShapes(ctx, 9)) {
+    std::vector<const onnx::TensorShapeProto*> shapes;
+    shapes.push_back(&ctx.getInputType(0)->tensor_type().shape());
+    shapes.push_back(&ctx.getInputType(1)->tensor_type().shape());
+    shapes.push_back(&ctx.getInputType(4)->tensor_type().shape());
+    multidirectionalBroadcastShapeInference(
+        shapes, *ctx.getOutputType(0)->mutable_tensor_type()->mutable_shape());
+  }
+}));
 
 ONNX_MS_OPERATOR_SET_SCHEMA(
     QGemm, 1,
diff --git a/onnxruntime/test/contrib_ops/cuda_kernels/qmoe_fp8_compaction_test.cc b/onnxruntime/test/contrib_ops/cuda_kernels/qmoe_fp8_compaction_test.cc
index 00cbba0e19f61..e414fe31f873e 100644
--- a/onnxruntime/test/contrib_ops/cuda_kernels/qmoe_fp8_compaction_test.cc
+++ b/onnxruntime/test/contrib_ops/cuda_kernels/qmoe_fp8_compaction_test.cc
@@ -1,6 +1,7 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
+#include <algorithm>
 #include <array>
 #include <cmath>
 #include <cstdio>
@@ -11,6 +12,9 @@
 #include <gtest/gtest.h>
 
 #include "contrib_ops/cuda/moe/qmoe_kernels.h"
+#ifdef ORT_TEST_NVFP4_GEMV
+#include "contrib_ops/cuda/llm/moe_gemm/moe_gemv_fp4.h"
+#endif
 #include "core/providers/cuda/cuda_common.h"
 
 namespace onnxruntime {
@@ -18,6 +22,9 @@ namespace test {
 namespace {
 
 namespace qmoe = contrib::cuda;
+#ifdef ORT_TEST_NVFP4_GEMV
+namespace gemv = llm::kernels::moe_gemv;
+#endif
 
 struct CudaBuffer {
   ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(CudaBuffer);
@@ -48,6 +55,242 @@ struct CudaBuffer {
   size_t bytes;
 };
 
+#ifdef ORT_TEST_NVFP4_GEMV
+template <typename T>
+void TestRoutingMaps(int rows, int experts, int k, bool normalize, bool tied) {
+  std::vector<T> logits(rows * experts);
+  for (int i = 0; i < rows * experts; ++i) {
+    logits[i] = T(tied ? 0.0f : float((i * 73 + i / experts * 11) % 101 - 50) / 16.0f);
+  }
+  const int expanded = rows * k;
+  CudaBuffer input(logits.size() * sizeof(T));
+  CudaBuffer scales(expanded * sizeof(float)), indices(expanded * sizeof(int));
+  CudaBuffer r2u(expanded * sizeof(int)), u2r(expanded * sizeof(int)), sorted(expanded * sizeof(int));
+  CudaBuffer offsets((experts + 1) * sizeof(int64_t));
+  input.Upload(logits.data());
+  qmoe::LaunchSoftmaxTopK(input.As<T>(), scales.As<float>(), indices.As<int>(),
+                         rows, experts, k, normalize, nullptr);
+  ASSERT_TRUE(qmoe::TryLaunchSmallDecodeExpertMaps(
+      indices.As<int>(), r2u.As<int>(), u2r.As<int>(), sorted.As<int>(),
+      offsets.As<int64_t>(), rows, experts, k, nullptr));
+  std::vector<int> expected_indices(expanded), actual_r2u(expanded),
+      actual_u2r(expanded), actual_sorted(expanded);
+  std::vector<int64_t> actual_offsets(experts + 1);
+  indices.Download(expected_indices.data());
+  r2u.Download(actual_r2u.data());
+  u2r.Download(actual_u2r.data());
+  sorted.Download(actual_sorted.data());
+  offsets.Download(actual_offsets.data());
+  std::vector<int> order(expanded);
+  for (int i = 0; i < expanded; ++i) order[i] = i;
+  std::stable_sort(order.begin(), order.end(), [&](int a, int b) {
+    return expected_indices[a] < expected_indices[b];
+  });
+  for (int rank = 0; rank < expanded; ++rank) {
+    const int original = order[rank];
+    const int source = (original % k) * rows + original / k;
+    EXPECT_EQ(actual_r2u[rank], source);
+    EXPECT_EQ(actual_u2r[source], rank);
+    EXPECT_EQ(actual_sorted[rank], expected_indices[original]);
+  }
+  for (int expert = 0; expert <= experts; ++expert) {
+    EXPECT_EQ(actual_offsets[expert],
+              std::count_if(expected_indices.begin(), expected_indices.end(),
+                            [&](int selected) { return selected < expert; }));
+  }
+}
+
+TEST(CUDA_EP_Unittest, QMoENvfp4SmallDecodeMapsPreserveStableOrder) {
+  for (int rows : {1, 2}) {
+    for (int experts : {16, 257, 512}) {
+      for (int k : {1, 10, 16}) {
+        for (bool normalize : {false, true}) {
+          for (bool tied : {false, true}) {
+            TestRoutingMaps<half>(rows, experts, k, normalize, tied);
+            TestRoutingMaps<__nv_bfloat16>(rows, experts, k, normalize, tied);
+          }
+        }
+      }
+    }
+  }
+  EXPECT_FALSE(qmoe::TryLaunchSmallDecodeExpertMaps(
+      nullptr, nullptr, nullptr, nullptr, nullptr, 3, 512, 10, nullptr));
+  EXPECT_FALSE(qmoe::TryLaunchSmallDecodeExpertMaps(
+      nullptr, nullptr, nullptr, nullptr, nullptr, 2, 513, 10, nullptr));
+  EXPECT_FALSE(qmoe::TryLaunchSmallDecodeExpertMaps(
+      nullptr, nullptr, nullptr, nullptr, nullptr, 0, 512, 10, nullptr));
+  EXPECT_FALSE(qmoe::TryLaunchSmallDecodeExpertMaps(
+      nullptr, nullptr, nullptr, nullptr, nullptr, 2, 512, 17, nullptr));
+  EXPECT_FALSE(qmoe::TryLaunchSmallDecodeExpertMaps(
+      nullptr, nullptr, nullptr, nullptr, nullptr, 2, 8, 10, nullptr));
+}
+
+template <typename T>
+void TestNvfp4KPackedGateUp(int k, bool use_bias, bool varied_scales = false, int n = 48) {
+  constexpr int experts = 3;
+  constexpr int rows = 4;
+  constexpr std::array<int, rows> routes{2, 0, 2, 1};
+  constexpr std::array<int, rows> source_rows{1, 0, 0, 1};
+  constexpr std::array<float, experts> globals{0.25f, 0.125f, 0.5f};
+  constexpr std::array<float, 8> magnitudes{0, 0.5f, 1, 1.5f, 2, 3, 4, 6};
+  std::vector<T> activations(2 * k), bias(experts * n);
+  std::vector<uint8_t> weights(experts * n * k / 2), scales(experts * n * k / 16);
+  for (size_t i = 0; i < activations.size(); ++i) {
+    activations[i] = T(float(int(i % 11) - 5) / 64.0f);
+  }
+  for (size_t i = 0; i < weights.size(); ++i) weights[i] = static_cast<uint8_t>(i * 37 + 19);
+  for (size_t i = 0; i < scales.size(); ++i) {
+    scales[i] = varied_scales ? static_cast<uint8_t>(i % 127) : (i % 2 ? 0x38 : 0x30);
+  }
+  for (size_t i = 0; i < bias.size(); ++i) bias[i] = T(float(int(i % 7) - 3) / 16.0f);
+  CudaBuffer a(activations.size() * sizeof(T)), w(weights.size()), s(scales.size());
+  CudaBuffer g(sizeof(globals)), b(bias.size() * sizeof(T)), r(sizeof(routes)), sources(sizeof(source_rows));
+  CudaBuffer output(rows * n / 2 * sizeof(T));
+  a.Upload(activations.data());
+  w.Upload(weights.data());
+  s.Upload(scales.data());
+  g.Upload(globals.data());
+  b.Upload(bias.data());
+  r.Upload(routes.data());
+  sources.Upload(source_rows.data());
+  llm::kernels::cutlass_kernels::ActivationParams activation;
+  activation.alpha = 1.0f;
+  activation.beta = 0.25f;
+  activation.limit = 7.0f;
+  gemv::launch_moe_gemv_fp4_symmetric_interleaved_swiglu<T>(
+      a.As<T>(), w.As<uint8_t>(), nullptr, s.As<uint8_t>(), g.As<float>(),
+      use_bias ? b.As<T>() : nullptr, output.As<T>(), nullptr, r.As<int>(),
+      experts, rows, n / 2, k, 16, 90, activation, gemv::MoeGemvConfig::kDefault,
+      false, true, sources.As<int>(), 2, nullptr, true);
+  CUDA_CALL_THROW(cudaGetLastError());
+  std::vector<T> actual(rows * n / 2);
+  output.Download(actual.data());
+  for (int row = 0; row < rows; ++row) {
+    for (int col = 0; col < n; col += 2) {
+      float sums[2]{};
+      for (int pair = 0; pair < 2; ++pair) {
+        const int column = col + pair;
+        for (int d = 0; d < k; ++d) {
+          const uint8_t packed = weights[(routes[row] * n + column) * (k / 2) + d / 2];
+          const int code = (d % 2 ? packed >> 4 : packed) & 15;
+          const uint8_t scale_code = scales[(routes[row] * n + column) * (k / 16) + d / 16];
+          const int exponent = scale_code >> 3;
+          const int mantissa = scale_code & 7;
+          const float block_scale = exponent == 0 ? std::ldexp(float(mantissa), -9)
+                                                  : std::ldexp(1.0f + mantissa * 0.125f, exponent - 7);
+          const float rounded_scale = static_cast<float>(T(block_scale * globals[routes[row]]));
+          const float weight = static_cast<float>(T(magnitudes[code & 7] * (code & 8 ? -1 : 1) * rounded_scale));
+          sums[pair] += weight * static_cast<float>(activations[source_rows[row] * k + d]);
+        }
+        if (use_bias) sums[pair] += static_cast<float>(bias[routes[row] * n + column]);
+      }
+      const float gate = std::min(sums[0], activation.limit);
+      const float up = std::clamp(sums[1], -activation.limit, activation.limit) + activation.beta;
+      const float expected = static_cast<float>(T(gate / (1.0f + std::exp(-gate)) * up));
+      EXPECT_NEAR(static_cast<float>(actual[row * n / 2 + col / 2]), expected,
+                  0.001f + std::abs(expected) * 0.01f)
+          << "row=" << row << ", column=" << col << ", K=" << k << ", bias=" << use_bias;
+    }
+  }
+}
+
+TEST(CUDA_EP_Unittest, QMoENvfp4KPackedGateUpTiling) {
+  for (int k : {144, 2560}) {
+    for (bool use_bias : {false, true}) {
+      TestNvfp4KPackedGateUp<half>(k, use_bias);
+      TestNvfp4KPackedGateUp<__nv_bfloat16>(k, use_bias);
+    }
+  }
+}
+
+TEST(CUDA_EP_Unittest, QMoENvfp4KPackedGateUpScaleRange) {
+  TestNvfp4KPackedGateUp<half>(2560, true, true);
+  TestNvfp4KPackedGateUp<__nv_bfloat16>(2560, true, true);
+}
+
+TEST(CUDA_EP_Unittest, QMoENvfp4KPackedGateUpPartialTiles) {
+  for (int n : {2, 10}) {
+    for (int k : {16, 272}) {
+      for (bool use_bias : {false, true}) {
+        TestNvfp4KPackedGateUp<half>(k, use_bias, false, n);
+        TestNvfp4KPackedGateUp<__nv_bfloat16>(k, use_bias, false, n);
+      }
+    }
+  }
+}
+
+template <typename T>
+void TestNvfp4FusedFinalize(int num_rows, int topk, int k, bool use_bias) {
+  constexpr int n = 50, experts = 3;
+  const int expanded = num_rows * topk;
+  std::vector<T> activations(expanded * k), biases(experts * n);
+  std::vector<uint8_t> weights(experts * n * k / 2), scales(experts * n * k / 16, 0x30);
+  std::vector<float> globals{0.3f, 0.17f, 0.51f}, routing(expanded);
+  std::vector<int> mapping(expanded), expert_ids(expanded);
+  for (size_t i = 0; i < activations.size(); ++i) activations[i] = T(float(int(i % 11) - 5) / 64.0f);
+  for (size_t i = 0; i < weights.size(); ++i) weights[i] = static_cast<uint8_t>(i * 37 + 19);
+  for (size_t i = 0; i < biases.size(); ++i) biases[i] = T(float(int(i % 7) - 3) / 16.0f);
+  for (int i = 0; i < expanded; ++i) {
+    mapping[i] = (i * 7) % expanded;
+    expert_ids[i] = i % experts;
+    routing[i] = float(i % topk + 1) / float(topk * (topk + 1) / 2);
+  }
+  if (topk > 1) mapping[expanded - 1] = -1;
+  if (expanded > 1) expert_ids[expanded / 2] = experts;
+  CudaBuffer a(activations.size() * sizeof(T)), w(weights.size()), s(scales.size());
+  CudaBuffer g(globals.size() * sizeof(float)), b(biases.size() * sizeof(T));
+  CudaBuffer m(mapping.size() * sizeof(int)), e(expert_ids.size() * sizeof(int)), r(routing.size() * sizeof(float));
+  CudaBuffer reference(expanded * n * sizeof(T)), output(num_rows * n * sizeof(T));
+  a.Upload(activations.data());
+  w.Upload(weights.data());
+  s.Upload(scales.data());
+  g.Upload(globals.data());
+  b.Upload(biases.data());
+  m.Upload(mapping.data());
+  e.Upload(expert_ids.data());
+  r.Upload(routing.data());
+  gemv::launch_moe_gemv_fp4_symmetric<T>(
+      a.As<T>(), w.As<uint8_t>(), nullptr, s.As<uint8_t>(), g.As<float>(),
+      use_bias ? b.As<T>() : nullptr, reference.As<T>(), nullptr, e.As<int>(),
+      experts, expanded, n, k, 16, 90, gemv::MoeGemvConfig::kDefault, false, true, nullptr, true);
+  ASSERT_TRUE(gemv::try_launch_moe_gemv_nvfp4_fused_finalize<T>(
+      a.As<T>(), w.As<uint8_t>(), s.As<uint8_t>(), g.As<float>(),
+      use_bias ? b.As<T>() : nullptr, output.As<T>(), m.As<int>(), e.As<int>(), r.As<float>(),
+      experts, num_rows, topk, n, k, nullptr));
+  CUDA_CALL_THROW(cudaGetLastError());
+  std::vector<T> reference_values(expanded * n), actual(num_rows * n);
+  reference.Download(reference_values.data());
+  output.Download(actual.data());
+  for (int row = 0; row < num_rows; ++row) {
+    for (int col = 0; col < n; ++col) {
+      float sum = 0.0f;
+      for (int rank = 0; rank < topk; ++rank) {
+        const int permuted = mapping[rank * num_rows + row];
+        if (permuted >= 0 && expert_ids[permuted] < experts) {
+          sum += routing[row * topk + rank] * static_cast<float>(reference_values[permuted * n + col]);
+        }
+      }
+      const float expected = static_cast<float>(T(sum));
+      EXPECT_NEAR(static_cast<float>(actual[row * n + col]), expected, 1e-5f + std::abs(expected) * 0.01f)
+          << "row=" << row << ", col=" << col << ", topk=" << topk << ", bias=" << use_bias;
+    }
+  }
+}
+
+TEST(CUDA_EP_Unittest, QMoENvfp4FusedFinalizePreservesRoutingAndRounding) {
+  for (int rows : {1, 2}) {
+    for (int topk : {1, 4, 10, 16}) {
+      for (int k : {144, 640}) {
+        for (bool use_bias : {false, true}) {
+          TestNvfp4FusedFinalize<half>(rows, topk, k, use_bias);
+          TestNvfp4FusedFinalize<__nv_bfloat16>(rows, topk, k, use_bias);
+        }
+      }
+    }
+  }
+}
+#endif
+
 TEST(CUDA_EP_Unittest, QMoEFp8CompactExpertsPreservesRoutesAndClearsUnusedSlots) {
   constexpr int num_experts = 512;
   constexpr int capacity = 6;
diff --git a/onnxruntime/test/contrib_ops/fused_tensor_ops_test.cc b/onnxruntime/test/contrib_ops/fused_tensor_ops_test.cc
index 441a7bf62ccce..4aaeb82ed47ad 100644
--- a/onnxruntime/test/contrib_ops/fused_tensor_ops_test.cc
+++ b/onnxruntime/test/contrib_ops/fused_tensor_ops_test.cc
@@ -234,6 +234,21 @@ TEST(FusedTensorOpsTest, SigmoidGateMulSiluEmptyAndInvalidAttributes) {
   }
 }
 
+#if !defined(DISABLE_FLOAT8_TYPES)
+TEST(FusedTensorOpsTest, DequantizeLinearFp8ToBf16) {
+  auto providers = CudaProvider();
+  if (providers.empty()) GTEST_SKIP();
+
+  OpTester test("DequantizeLinear", 1, kMSDomain);
+  test.AddAttribute<int64_t>("axis", 1);
+  test.AddAttribute<int64_t>("output_dtype", ONNX_NAMESPACE::TensorProto_DataType_BFLOAT16);
+  test.AddInput<Float8E4M3FN>("weight", {1, 8}, std::vector<Float8E4M3FN>(8, Float8E4M3FN(1.0f)));
+  test.AddInput<float>("scale", {8}, std::vector<float>(8, 2.0f));
+  test.AddOptionalInputEdge<Float8E4M3FN>();
+  test.AddOutput<BFloat16>("weight_bf16", {1, 8}, ToTensorType<BFloat16>(std::vector<float>(8, 2.0f)));
+  test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers);
+}
+#endif
 
 TEST(FusedTensorOpsTest, UntileHyperConnectionWeights) {
   constexpr int64_t kInputElements = 160ll * 20480;
diff --git a/onnxruntime/test/contrib_ops/matmul_block_scaled_fp8_test.cc b/onnxruntime/test/contrib_ops/matmul_block_scaled_fp8_test.cc
index 9dd9b0e722a29..94dbfa82a4636 100644
--- a/onnxruntime/test/contrib_ops/matmul_block_scaled_fp8_test.cc
+++ b/onnxruntime/test/contrib_ops/matmul_block_scaled_fp8_test.cc
@@ -2,6 +2,7 @@
 // Licensed under the MIT License.
 
 #include <cstdlib>
+#include <utility>
 
 #include "gtest/gtest.h"
 #include "test/common/cuda_op_test_utils.h"
@@ -69,6 +70,124 @@ std::vector<Float8E4M3FN> MakeConstRowWeight(const std::vector<float>& row_value
 }
 }  // namespace
 
+template <typename T>
+void TestFp8SplitOutputs() {
+  if (!HasCudaEnvironment(800)) {
+    GTEST_SKIP() << "CUDA device is required for split-output FP8 matmul.";
+  }
+  for (const auto& [m, k] : {std::pair<int64_t, int64_t>{0, 256}, {1, 256}, {2, 256}, {2, 320},
+                            {2, 1088}, {9, 256},
+                            {32, 256}, {33, 256}, {65, 256}, {2, 16}, {2, 30}, {2, 0}}) {
+    for (const bool activation_scale : {false, true}) {
+      SCOPED_TRACE("M=" + std::to_string(m) + " K=" + std::to_string(k) +
+                   " W8A8=" + std::to_string(activation_scale));
+      constexpr int64_t n = 17;
+      constexpr int64_t block_size = 64;
+      const int64_t k_blocks = (k + block_size - 1) / block_size;
+      const std::vector<int64_t> widths = m == 2 && k == 256
+                                             ? std::vector<int64_t>{1, 1, 1, 1, 1, 1, 1, 10}
+                                             : (m == 33 ? std::vector<int64_t>{17}
+                                                        : std::vector<int64_t>{1, 7, 9});
+      std::vector<T> a;
+      std::vector<Float8E4M3FN> b;
+      std::vector<float> scales;
+      std::vector<T> bias;
+      for (int64_t row = 0; row < m; ++row) {
+        for (int64_t i = 0; i < k; ++i) {
+          a.emplace_back(static_cast<float>((row + i) % 5 - 2) / 4.0f);
+        }
+      }
+      for (int64_t col = 0; col < n; ++col) {
+        bias.emplace_back(static_cast<float>(col - 8) / 8.0f);
+        for (int64_t i = 0; i < k; ++i) {
+          b.emplace_back(static_cast<float>((col + i) % 3 - 1));
+        }
+        for (int64_t kb = 0; kb < k_blocks; ++kb) {
+          scales.push_back(static_cast<float>(1 + (col + kb) % 2) / 4.0f);
+        }
+      }
+      TensorShapeVector a_shape = m == 1 ? TensorShapeVector{k} : TensorShapeVector{m, k};
+      if (m == 9) {
+        a_shape = {3, 3, k};
+      }
+      OpTester test("MatMulBlockQuantizedFp8Weight", 1, kMSDomain);
+      test.AddAttribute<int64_t>("block_size", block_size);
+      test.AddAttribute<std::vector<int64_t>>("split_sizes", widths);
+      test.AddInput<T>("A", a_shape, a);
+      test.AddInput<Float8E4M3FN>("B", {n, k}, b);
+      test.AddInput<float>("b_scale", {n, k_blocks}, scales);
+      if (activation_scale) {
+        test.AddInput<float>("a_scale", {}, {0.5f});
+      } else {
+        test.AddOptionalInputEdge<float>();
+      }
+      if (activation_scale) {
+        test.AddInput<T>("bias", {n}, bias);
+      }
+      int64_t start = 0;
+      for (size_t output = 0; output < widths.size(); ++output) {
+        std::vector<T> expected;
+        for (int64_t row = 0; row < m; ++row) {
+          for (int64_t col = start; col < start + widths[output]; ++col) {
+            float sum = activation_scale ? static_cast<float>(bias[col]) : 0.0f;
+            for (int64_t i = 0; i < k; ++i) {
+              sum += static_cast<float>(a[row * k + i]) * b[col * k + i].ToFloat() *
+                     scales[col * k_blocks + i / block_size];
+            }
+            expected.emplace_back(sum);
+          }
+        }
+        auto shape = a_shape;
+        shape.back() = widths[output];
+        test.AddOutput<T>(("Y" + std::to_string(output)).c_str(), shape, expected);
+        start += widths[output];
+      }
+      test.SetOutputTolerance(0.0f, 0.0f);
+      std::vector<std::unique_ptr<IExecutionProvider>> providers;
+      providers.push_back(DefaultCudaExecutionProvider());
+      test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers);
+    }
+  }
+}
+
+TEST(MatMulBlockQuantizedFp8WeightOpTest, SplitOutputsFp16) {
+  TestFp8SplitOutputs<MLFloat16>();
+}
+
+TEST(MatMulBlockQuantizedFp8WeightOpTest, SplitOutputsBFloat16) {
+  TestFp8SplitOutputs<BFloat16>();
+}
+
+TEST(MatMulBlockQuantizedFp8WeightOpTest, RejectsInvalidSplitSizes) {
+  if (!HasCudaEnvironment(800)) {
+    GTEST_SKIP() << "CUDA device is required for split-output FP8 matmul.";
+  }
+  for (const std::vector<int64_t>& widths : {std::vector<int64_t>{}, {0, 17}, {-1, 18}, {8, 8},
+                                            {1, 1, 1, 1, 1, 1, 1, 1, 9}}) {
+    OpTester test("MatMulBlockQuantizedFp8Weight", 1, kMSDomain);
+    test.AddAttribute<std::vector<int64_t>>("split_sizes", widths);
+    test.AddInput<BFloat16>("A", {1, 16}, std::vector<BFloat16>(16));
+    test.AddInput<Float8E4M3FN>("B", {17, 16}, std::vector<Float8E4M3FN>(17 * 16));
+    test.AddInput<float>("b_scale", {17, 1}, std::vector<float>(17, 1.0f));
+    for (size_t i = 0; i < std::max<size_t>(1, widths.size()); ++i) {
+      const int64_t width = i < widths.size() ? std::max<int64_t>(0, widths[i]) : 17;
+      test.AddOutput<BFloat16>(("Y" + std::to_string(i)).c_str(), {1, width}, std::vector<BFloat16>(width));
+    }
+    std::vector<std::unique_ptr<IExecutionProvider>> providers;
+    providers.push_back(DefaultCudaExecutionProvider());
+    test.Run(OpTester::ExpectResult::kExpectFailure, "split_sizes", {}, nullptr, &providers);
+  }
+  OpTester count_test("MatMulBlockQuantizedFp8Weight", 1, kMSDomain);
+  count_test.AddAttribute<std::vector<int64_t>>("split_sizes", {8, 9});
+  count_test.AddInput<BFloat16>("A", {1, 16}, std::vector<BFloat16>(16));
+  count_test.AddInput<Float8E4M3FN>("B", {17, 16}, std::vector<Float8E4M3FN>(17 * 16));
+  count_test.AddInput<float>("b_scale", {17, 1}, std::vector<float>(17, 1.0f));
+  count_test.AddOutput<BFloat16>("Y", {1, 17}, std::vector<BFloat16>(17));
+  std::vector<std::unique_ptr<IExecutionProvider>> providers;
+  providers.push_back(DefaultCudaExecutionProvider());
+  count_test.Run(OpTester::ExpectResult::kExpectFailure, "split_sizes", {}, nullptr, &providers);
+}
+
 TEST(MatMulBlockQuantizedFp8WeightOpTest, GemvTensorCoreKSplitSelection) {
   struct Case {
     int n;
diff --git a/onnxruntime/test/contrib_ops/moe_test.cc b/onnxruntime/test/contrib_ops/moe_test.cc
index 12148390f63aa..0d968eaf07a82 100644
--- a/onnxruntime/test/contrib_ops/moe_test.cc
+++ b/onnxruntime/test/contrib_ops/moe_test.cc
@@ -2542,8 +2542,8 @@ template <typename ElementType>
 static void RunQMoEPackedPrefillRoutingTest(int64_t num_rows, int64_t hidden_size, int64_t inter_size,
                                             int64_t top_k, bool skewed, bool scale_initializers = true,
                                             int64_t row_tile_size = 0, int64_t block_size = 64,
-                                            int64_t fc1_bits = 2, int64_t fc2_bits = 4) {
-  constexpr int64_t num_experts = 8;
+                                            int64_t fc1_bits = 2, int64_t fc2_bits = 4,
+                                            int64_t num_experts = 8) {
   constexpr float alpha = 1.5f;
   constexpr float beta = 0.125f;
   constexpr float limit = 0.5f;
@@ -2556,7 +2556,7 @@ static void RunQMoEPackedPrefillRoutingTest(int64_t num_rows, int64_t hidden_siz
     return 0.125f * static_cast<float>(expert + 1) +
            0.015625f * static_cast<float>(column % 5) + 0.03125f * static_cast<float>(block);
   };
-  const auto make_weights = [](int64_t rows, int64_t columns, int64_t bits, bool interleaved) {
+  const auto make_weights = [num_experts](int64_t rows, int64_t columns, int64_t bits, bool interleaved) {
     const int64_t pack_size = 8 / bits;
     const uint8_t zero = static_cast<uint8_t>(1u << (bits - 1));
     std::vector<uint8_t> weights(static_cast<size_t>(num_experts * rows * columns / pack_size), 0);
@@ -2684,6 +2684,7 @@ TEST(MoETest, QMoETest_MixedWidthCudaPackedPrefillRouting) {
   if (!HasCudaEnvironment(800)) {
     GTEST_SKIP() << "SM80 or later is required for packed INT2 prefill.";
   }
+
   ScopedEnvironmentVariables scoped_env_vars{
       {{"ORT_DISABLE_MOE_GEMV", "1"}, {"ORT_ENABLE_QMOE_INT2_PREFILL", "1"}}};
   for (int64_t top_k : {1, 2, 4}) {
@@ -2697,6 +2698,20 @@ TEST(MoETest, QMoETest_MixedWidthCudaPackedPrefillRouting) {
   }
 }
 
+TEST(MoETest, QMoETest_CudaColumnTiledFinalize) {
+  if (!HasCudaEnvironment(800)) {
+    GTEST_SKIP() << "SM80 or later is required for packed INT2 decode.";
+  }
+  for (int64_t rows : {1, 2}) {
+    for (int64_t hidden : {2560, 2624}) {
+      RunQMoEPackedPrefillRoutingTest<MLFloat16>(rows, hidden, 128, 10, false, true, 0, 64, 2, 4, 16);
+#if defined(ENABLE_BF16)
+      RunQMoEPackedPrefillRoutingTest<BFloat16>(rows, hidden, 128, 10, false, true, 0, 64, 2, 4, 16);
+#endif
+    }
+  }
+}
+
 TEST(MoETest, QMoETest_MixedWidthCudaPackedPrefillRuntimeScales) {
   if (!HasCudaEnvironment(800)) {
     GTEST_SKIP() << "SM80 or later is required for packed INT2 prefill.";

From 7347338e7a350b83b78cfc2cc02277f2f434cabe Mon Sep 17 00:00:00 2001
From: Akshay Sonawane <asonawane@microsoft.com>
Date: Mon, 19 Oct 2026 18:12:49 +0000
Subject: [PATCH 2/2] Fix lint

---
 .../packed_sparse_attention_indexer_common.h  |   6 +-
 .../cuda/bert/branchwise_rms_norm_impl.cu     |   4 +-
 .../cuda/bert/causal_conv_with_state_impl.cu  |   4 +-
 .../cuda/bert/dynamic_sparse_attention.cc     |  30 +--
 .../bert/dynamic_sparse_attention_impl.cu     |   2 +-
 .../contrib_ops/cuda/bert/fused_tensor_ops.cc |  46 ++--
 .../contrib_ops/cuda/bert/fused_tensor_ops.cu |  18 +-
 .../contrib_ops/cuda/bert/gated_delta_net.cc  |  28 +-
 .../cuda/bert/gated_delta_net_impl.cu         |   8 +-
 .../cuda/bert/gated_delta_net_impl.h          |   2 +-
 .../cuda/bert/rotary_embedding_impl.cu        |   9 +-
 .../cuda/math/matmul_block_scaled_fp8.cu      |  30 +--
 .../cuda/math/matmul_block_scaled_fp8.h       |   2 +-
 .../packed_sparse_attention_indexer_impl.cu   |  14 +-
 .../core/graph/contrib_ops/bert_defs.cc       | 242 +++++++++---------
 .../core/graph/contrib_ops/contrib_defs.cc    |   2 +-
 .../cuda/reduction/reduction_functions.cu     |   2 +-
 .../cuda/reduction/reduction_functions.h      |   2 +-
 .../causal_conv_with_state_op_test.cc         |  14 +-
 .../cuda_kernels/qmoe_fp8_compaction_test.cc  |   2 +-
 .../dynamic_sparse_attention_op_test.cc       |  40 +--
 .../test/contrib_ops/fused_tensor_ops_test.cc |  14 +-
 .../contrib_ops/hyper_connection_op_test.cc   |  10 +-
 .../matmul_block_scaled_fp8_test.cc           |  13 +-
 .../contrib_ops/rotary_embedding_op_test.cc   |  10 +-
 ...untime_test_python_symbolic_shape_infer.py |   3 -
 26 files changed, 289 insertions(+), 268 deletions(-)

diff --git a/onnxruntime/contrib_ops/cpu/sparse/packed_sparse_attention_indexer_common.h b/onnxruntime/contrib_ops/cpu/sparse/packed_sparse_attention_indexer_common.h
index fdc636bc9626d..72fbac76026f6 100644
--- a/onnxruntime/contrib_ops/cpu/sparse/packed_sparse_attention_indexer_common.h
+++ b/onnxruntime/contrib_ops/cpu/sparse/packed_sparse_attention_indexer_common.h
@@ -71,9 +71,9 @@ enum OutputIndex : int {
   kPresentStateLengths = 5,  // [batch_size, 2], int32
   kStateUpdate = 6,          // qsa only, optional: [batch_size, state_update_capacity, head_size]
   kIndexShareStatus = 7,
-  kMidKvBuffer = 8,          // qsa only, optional: buffer snapshot after the first request token
-  kMidStateLengths = 9,      // qsa only, optional: state lengths after the first request token
-  kValidationStatus = 10,    // qsa only, optional: per-request overflow/metadata status
+  kMidKvBuffer = 8,        // qsa only, optional: buffer snapshot after the first request token
+  kMidStateLengths = 9,    // qsa only, optional: state lengths after the first request token
+  kValidationStatus = 10,  // qsa only, optional: per-request overflow/metadata status
   kOutputCount = 8,
   kMaxOutputCount = 11,
 };
diff --git a/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm_impl.cu b/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm_impl.cu
index c07197fa3bf14..240b057b827e0 100644
--- a/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm_impl.cu
+++ b/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm_impl.cu
@@ -168,9 +168,9 @@ INSTANTIATE(__nv_bfloat16)
 
 #undef INSTANTIATE
 
-#define INSTANTIATE_STANDARD(T) \
+#define INSTANTIATE_STANDARD(T)                                                                    \
   template Status LaunchStandardBranchwiseRMSNorm<T>(cudaStream_t, const T*, const void*, int, T*, \
-                                                      int, int, int, bool, float, float*, float*, float*, void*, size_t, int);
+                                                     int, int, int, bool, float, float*, float*, float*, void*, size_t, int);
 INSTANTIATE_STANDARD(float)
 INSTANTIATE_STANDARD(half)
 INSTANTIATE_STANDARD(__nv_bfloat16)
diff --git a/onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state_impl.cu b/onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state_impl.cu
index c5284935a5e04..eea53e8bbe272 100644
--- a/onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state_impl.cu
+++ b/onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state_impl.cu
@@ -480,8 +480,8 @@ Status LaunchCausalConvProjectionKernel(
   return CUDA_CALL(cudaGetLastError());
 }
 
-#define INSTANTIATE_PROJECTION(T) \
-  template Status LaunchCausalConvProjectionKernel<T>( \
+#define INSTANTIATE_PROJECTION(T)                                       \
+  template Status LaunchCausalConvProjectionKernel<T>(                  \
       cudaStream_t, const T*, const T*, const T*, const T*, T*, T*, T*, \
       int, int, int, int, int, int, int, int, bool, bool);
 INSTANTIATE_PROJECTION(float)
diff --git a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention.cc b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention.cc
index b11283cff7599..f7a65fbad3663 100644
--- a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention.cc
+++ b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention.cc
@@ -19,21 +19,21 @@ namespace cuda {
 
 using namespace onnxruntime::cuda;
 
-#define REGISTER_KERNEL_TYPED(T)                                       \
-  ONNX_OPERATOR_TYPED_KERNEL_EX(                                       \
-      DynamicSparseAttention,                                          \
-      kMSDomain,                                                       \
-      1,                                                               \
-      T,                                                               \
-      kCudaExecutionProvider,                                          \
-      (*KernelDefBuilder::Create())                                    \
-          .TypeConstraint("T", DataTypeImpl::GetTensorType<T>())       \
+#define REGISTER_KERNEL_TYPED(T)                                            \
+  ONNX_OPERATOR_TYPED_KERNEL_EX(                                            \
+      DynamicSparseAttention,                                               \
+      kMSDomain,                                                            \
+      1,                                                                    \
+      T,                                                                    \
+      kCudaExecutionProvider,                                               \
+      (*KernelDefBuilder::Create())                                         \
+          .TypeConstraint("T", DataTypeImpl::GetTensorType<T>())            \
           .TypeConstraint("TCache", BuildKernelDefConstraints<T, int8_t>()) \
-          .TypeConstraint("S", DataTypeImpl::GetTensorType<float>())   \
-          .TypeConstraint("M", DataTypeImpl::GetTensorType<int32_t>()) \
-          .MayInplace(3, 1)                                            \
-          .MayInplace(4, 2)                                            \
-          .InputMemoryType(OrtMemTypeCPUInput, 10),                    \
+          .TypeConstraint("S", DataTypeImpl::GetTensorType<float>())        \
+          .TypeConstraint("M", DataTypeImpl::GetTensorType<int32_t>())      \
+          .MayInplace(3, 1)                                                 \
+          .MayInplace(4, 2)                                                 \
+          .InputMemoryType(OrtMemTypeCPUInput, 10),                         \
       DynamicSparseAttention<T>);
 
 REGISTER_KERNEL_TYPED(float)
@@ -142,7 +142,7 @@ Status DynamicSparseAttention<T>::ComputeInternal(OpKernelContext* context) cons
                     "DynamicSparseAttention: INT8 caches require per-channel key_scale and value_scale.");
   ORT_RETURN_IF_NOT((past_key == nullptr || past_key->IsDataType<T>() || quantized_cache) &&
                         (past_value == nullptr || (quantized_cache ? past_value->IsDataType<int8_t>()
-                                                                 : past_value->IsDataType<T>())),
+                                                                   : past_value->IsDataType<T>())),
                     "DynamicSparseAttention: cache types must match activations or both be INT8.");
 
   DynamicSparseAttentionParameters parameters;
diff --git a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu
index 0e3e8746de44d..ccdfdf76bb7ee 100644
--- a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu
+++ b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu
@@ -956,7 +956,7 @@ __global__ void GroupedDecodeDynamicSparseAttentionKernel(
   const int selected_count = max(0, min(selected_counts[row], max_selected));
   const int candidate_count = local_count + selected_count;
   constexpr int kCandidatesPerSplit = kQuantizedCache ? kGroupedQuantizedDecodeCandidatesPerSplit
-                                                     : kGroupedDecodeCandidatesPerSplit;
+                                                      : kGroupedDecodeCandidatesPerSplit;
   const int candidate_begin = split * kCandidatesPerSplit;
   const int32_t* row_indices =
       max_selected == 0 ? selected_indices : selected_indices + static_cast<int64_t>(row) * max_selected;
diff --git a/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cc b/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cc
index 24d2a153bf450..555c98f1fb16d 100644
--- a/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cc
+++ b/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cc
@@ -13,12 +13,12 @@
 
 namespace onnxruntime::contrib::cuda {
 
-#define REGISTER_ARGMAX_LAST_AXIS(T)                                       \
-  ONNX_OPERATOR_TYPED_KERNEL_EX(                                           \
-      ArgMaxLastAxis, kMSDomain, 1, T, kCudaExecutionProvider,             \
-      (*KernelDefBuilder::Create())                                        \
-            .TypeConstraint("T", DataTypeImpl::GetTensorType<T>())         \
-            .TypeConstraint("I", DataTypeImpl::GetTensorType<int64_t>()),  \
+#define REGISTER_ARGMAX_LAST_AXIS(T)                                    \
+  ONNX_OPERATOR_TYPED_KERNEL_EX(                                        \
+      ArgMaxLastAxis, kMSDomain, 1, T, kCudaExecutionProvider,          \
+      (*KernelDefBuilder::Create())                                     \
+          .TypeConstraint("T", DataTypeImpl::GetTensorType<T>())        \
+          .TypeConstraint("I", DataTypeImpl::GetTensorType<int64_t>()), \
       ArgMaxLastAxis<T>);
 
 REGISTER_ARGMAX_LAST_AXIS(float)
@@ -30,38 +30,38 @@ REGISTER_ARGMAX_LAST_AXIS(BFloat16)
 ONNX_OPERATOR_KERNEL_EX(
     EmbeddingRMSNorm, kMSDomain, 1, kCudaExecutionProvider,
     (*KernelDefBuilder::Create())
-          .TypeConstraint("I", DataTypeImpl::GetTensorType<int64_t>())
-          .TypeConstraint("T", DataTypeImpl::GetTensorType<BFloat16>())
-          .TypeConstraint("F", DataTypeImpl::GetTensorType<float>()),
+        .TypeConstraint("I", DataTypeImpl::GetTensorType<int64_t>())
+        .TypeConstraint("T", DataTypeImpl::GetTensorType<BFloat16>())
+        .TypeConstraint("F", DataTypeImpl::GetTensorType<float>()),
     EmbeddingRMSNorm);
 
 ONNX_OPERATOR_KERNEL_EX(
     RowSelect, kMSDomain, 1, kCudaExecutionProvider,
     (*KernelDefBuilder::Create())
-          .TypeConstraint("I", DataTypeImpl::GetTensorType<int64_t>())
-          .TypeConstraint("T", DataTypeImpl::GetTensorType<BFloat16>()),
+        .TypeConstraint("I", DataTypeImpl::GetTensorType<int64_t>())
+        .TypeConstraint("T", DataTypeImpl::GetTensorType<BFloat16>()),
     RowSelect);
 
 ONNX_OPERATOR_KERNEL_EX(
     SigmoidGateMul, kMSDomain, 1, kCudaExecutionProvider,
     (*KernelDefBuilder::Create())
-          .TypeConstraint("T", DataTypeImpl::GetTensorType<BFloat16>()),
+        .TypeConstraint("T", DataTypeImpl::GetTensorType<BFloat16>()),
     SigmoidGateMul);
 
 #if !defined(DISABLE_FLOAT8_TYPES)
 ONNX_OPERATOR_TYPED_KERNEL_EX(
     DequantizeLinear, kMSDomain, 1, Float8E4M3FN_BFloat16, kCudaExecutionProvider,
     (*KernelDefBuilder::Create())
-          .TypeConstraint("T1", DataTypeImpl::GetTensorType<Float8E4M3FN>())
-          .TypeConstraint("T2", DataTypeImpl::GetTensorType<BFloat16>())
-          .TypeConstraint("TS", DataTypeImpl::GetTensorType<float>()),
+        .TypeConstraint("T1", DataTypeImpl::GetTensorType<Float8E4M3FN>())
+        .TypeConstraint("T2", DataTypeImpl::GetTensorType<BFloat16>())
+        .TypeConstraint("TS", DataTypeImpl::GetTensorType<float>()),
     DequantizeLinearFp8Bf16);
 #endif
 
 ONNX_OPERATOR_KERNEL_EX(
     UntileHyperConnectionWeights, kMSDomain, 1, kCudaExecutionProvider,
     (*KernelDefBuilder::Create())
-          .TypeConstraint("T", DataTypeImpl::GetTensorType<BFloat16>()),
+        .TypeConstraint("T", DataTypeImpl::GetTensorType<BFloat16>()),
     UntileHyperConnectionWeights);
 
 template <typename T>
@@ -153,12 +153,10 @@ SigmoidGateMul::SigmoidGateMul(const OpKernelInfo& info)
   const int64_t silu = info.GetAttrOrDefault<int64_t>("silu_gate", 0);
   ORT_ENFORCE(silu == 0 || silu == 1, "SigmoidGateMul silu_gate must be zero or one");
   ORT_ENFORCE(!(broadcast_gate_ && silu_gate_), "SigmoidGateMul broadcast_gate and silu_gate are exclusive");
-  ORT_ENFORCE(broadcast_gate_ || silu_gate_ || (head_dim_ > 0 && heads_ > 0 &&
-                  q_stride_ >= 2 * head_dim_ &&
-                  tok_stride_ > 0 && q_offset_ >= 0),
+  ORT_ENFORCE(broadcast_gate_ || silu_gate_ || (head_dim_ > 0 && heads_ > 0 && q_stride_ >= 2 * head_dim_ && tok_stride_ > 0 && q_offset_ >= 0),
               "SigmoidGateMul received invalid layout attributes");
   ORT_ENFORCE(!(broadcast_gate_ || silu_gate_) || (head_dim_ == 0 && heads_ == 0 && q_stride_ == 0 &&
-                                      tok_stride_ == 0 && q_offset_ == 0),
+                                                   tok_stride_ == 0 && q_offset_ == 0),
               "SigmoidGateMul broadcast_gate/silu_gate cannot be combined with packed layout attributes");
 }
 
@@ -184,12 +182,12 @@ Status SigmoidGateMul::ComputeInternal(OpKernelContext* context) const {
   } else {
     const int64_t row_width = static_cast<int64_t>(heads_) * head_dim_;
     ORT_RETURN_IF_NOT(count % row_width == 0,
-                    "SigmoidGateMul attention size must be divisible by heads * head_dim");
+                      "SigmoidGateMul attention size must be divisible by heads * head_dim");
     const int64_t tokens = count / row_width;
     ORT_RETURN_IF_NOT(projection->Shape().Size() >=
-                        (tokens - 1) * tok_stride_ + q_offset_ +
-                            static_cast<int64_t>(heads_) * q_stride_,
-                    "SigmoidGateMul projection is too small for the configured layout");
+                          (tokens - 1) * tok_stride_ + q_offset_ +
+                              static_cast<int64_t>(heads_) * q_stride_,
+                      "SigmoidGateMul projection is too small for the configured layout");
   }
   Tensor* output = context->Output(0, attention->Shape());
   return LaunchSigmoidGateMul(
diff --git a/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cu b/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cu
index 6db5544942e1f..789edc4a26615 100644
--- a/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cu
+++ b/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cu
@@ -169,8 +169,8 @@ __global__ void SigmoidGateMulKernel(const __nv_bfloat16* attention,
     const BFloat16 a(gate);
     // Match CUDA Sigmoid's input-type intermediates, including its negative-input subtraction.
     const BFloat16 value = a > BFloat16(0.0f)
-                              ? BFloat16(1.0f) / (BFloat16(1.0f) + _Exp(-a))
-                              : BFloat16(1.0f) - BFloat16(1.0f) / (BFloat16(1.0f) + _Exp(a));
+                               ? BFloat16(1.0f) / (BFloat16(1.0f) + _Exp(-a))
+                               : BFloat16(1.0f) - BFloat16(1.0f) / (BFloat16(1.0f) + _Exp(a));
     sigmoid_bf16 = __float2bfloat16_rn(static_cast<float>(value));
   }
   if constexpr (SiluGate) {
@@ -254,13 +254,13 @@ Status LaunchArgMaxLastAxis(cudaStream_t stream, const void* input, int input_ty
   switch (input_type) {
     case ONNX_NAMESPACE::TensorProto_DataType_FLOAT:
       return LaunchTypedArgmax(stream, static_cast<const float*>(input), rows, width, output,
-                                partial_values, partial_indices);
+                               partial_values, partial_indices);
     case ONNX_NAMESPACE::TensorProto_DataType_FLOAT16:
       return LaunchTypedArgmax(stream, static_cast<const half*>(input), rows, width, output,
-                                partial_values, partial_indices);
+                               partial_values, partial_indices);
     case ONNX_NAMESPACE::TensorProto_DataType_BFLOAT16:
       return LaunchTypedArgmax(stream, static_cast<const __nv_bfloat16*>(input), rows, width, output,
-                                partial_values, partial_indices);
+                               partial_values, partial_indices);
     default:
       return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "ArgMaxLastAxis input type is not supported");
   }
@@ -299,11 +299,11 @@ Status LaunchSigmoidGateMul(cudaStream_t stream, const uint16_t* attention,
   const int64_t block_count = (count - 1) / kThreads + 1;
   ORT_RETURN_IF_NOT(block_count <= std::numeric_limits<int>::max(),
                     "SigmoidGateMul requires too many CUDA blocks");
-#define LAUNCH_GATE(BROADCAST, RESIDUAL, SILU) \
+#define LAUNCH_GATE(BROADCAST, RESIDUAL, SILU)                                                             \
   SigmoidGateMulKernel<BROADCAST, RESIDUAL, SILU><<<static_cast<int>(block_count), kThreads, 0, stream>>>( \
-      reinterpret_cast<const __nv_bfloat16*>(attention), \
-      reinterpret_cast<const __nv_bfloat16*>(projection), \
-      reinterpret_cast<__nv_bfloat16*>(output), count, head_dim, heads, q_stride, \
+      reinterpret_cast<const __nv_bfloat16*>(attention),                                                   \
+      reinterpret_cast<const __nv_bfloat16*>(projection),                                                  \
+      reinterpret_cast<__nv_bfloat16*>(output), count, head_dim, heads, q_stride,                          \
       token_stride, q_offset, gate_row_width, reinterpret_cast<const __nv_bfloat16*>(residual), sigmoid_in_fp32)
   if (silu_gate) {
     if (residual != nullptr) {
diff --git a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net.cc b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net.cc
index f346bd3f6027d..380e99746af31 100644
--- a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net.cc
+++ b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net.cc
@@ -22,22 +22,22 @@ using onnxruntime::cuda::CudaKernel;
 using onnxruntime::cuda::OrtToCudaType;
 namespace gdn = gated_delta_net;
 
-#define REGISTER_KERNEL_TYPED(T)                                        \
-  ONNX_OPERATOR_TYPED_KERNEL_EX(                                        \
-      GatedDeltaNet,                                                    \
-      kMSDomain,                                                        \
-      1,                                                                \
-      T,                                                                \
-      kCudaExecutionProvider,                                           \
-      (*KernelDefBuilder::Create())                                     \
-          .TypeConstraint("T", DataTypeImpl::GetTensorType<T>())        \
-          .TypeConstraint("TS", DataTypeImpl::GetTensorType<float>())   \
-          .TypeConstraint("TG", {DataTypeImpl::GetTensorType<float>(), \
+#define REGISTER_KERNEL_TYPED(T)                                           \
+  ONNX_OPERATOR_TYPED_KERNEL_EX(                                           \
+      GatedDeltaNet,                                                       \
+      kMSDomain,                                                           \
+      1,                                                                   \
+      T,                                                                   \
+      kCudaExecutionProvider,                                              \
+      (*KernelDefBuilder::Create())                                        \
+          .TypeConstraint("T", DataTypeImpl::GetTensorType<T>())           \
+          .TypeConstraint("TS", DataTypeImpl::GetTensorType<float>())      \
+          .TypeConstraint("TG", {DataTypeImpl::GetTensorType<float>(),     \
                                  DataTypeImpl::GetTensorType<MLFloat16>(), \
                                  DataTypeImpl::GetTensorType<BFloat16>()}) \
-          .TypeConstraint("TI", DataTypeImpl::GetTensorType<int32_t>()) \
-          .InputMemoryType(OrtMemTypeCPUInput, 10)                      \
-          .MayInplace(6, 1),                                            \
+          .TypeConstraint("TI", DataTypeImpl::GetTensorType<int32_t>())    \
+          .InputMemoryType(OrtMemTypeCPUInput, 10)                         \
+          .MayInplace(6, 1),                                               \
       GatedDeltaNet<T>);
 
 REGISTER_KERNEL_TYPED(float)
diff --git a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.cu b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.cu
index f308442329c5d..048faa9c2f3f5 100644
--- a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.cu
+++ b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.cu
@@ -745,8 +745,8 @@ Status LaunchRecurrent(const Descriptor& desc, bool warp_specialized, const Vari
 
 template <typename T, bool InputTypeGates>
 static Status LaunchGatedDeltaNetTypedGates(const Descriptor& desc, const Plan& plan, const VariantPack<T>& pack,
-                           float scale, int max_threads_per_block,
-                           size_t max_shared_memory_per_block, cudaStream_t stream) {
+                                            float scale, int max_threads_per_block,
+                                            size_t max_shared_memory_per_block, cudaStream_t stream) {
   KernelParams p{};
   p.total_tokens = desc.total_tokens;
   p.uniform_len = desc.batch > 0 ? desc.total_tokens / desc.batch : 0;
@@ -800,13 +800,13 @@ static Status LaunchGatedDeltaNetTypedGates(const Descriptor& desc, const Plan&
       // The chunked engine only accepts head_size_qk == head_size_v == 128, so the warp
       // kernel is always the right shape here.
       return LaunchRecurrent<T, InputTypeGates>(desc, /*warp_specialized=*/true, pack, p,
-                                max_shared_memory_per_block, stream);
+                                                max_shared_memory_per_block, stream);
     }
   }
 
   (void)max_threads_per_block;
   return LaunchRecurrent<T, InputTypeGates>(desc, plan.warp_specialized, pack, p, max_shared_memory_per_block,
-                            stream);
+                                            stream);
 }
 
 template <typename T>
diff --git a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.h b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.h
index 2bd82be0b53bf..3035ca02fa083 100644
--- a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.h
+++ b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.h
@@ -40,7 +40,7 @@ struct VariantPack {
 template <typename T>
 __device__ __forceinline__ float LoadGate(const void* data, int64_t index, bool gates_in_input_type) {
   return gates_in_input_type ? static_cast<float>(static_cast<const T*>(data)[index])
-                            : static_cast<const float*>(data)[index];
+                             : static_cast<const float*>(data)[index];
 }
 #endif
 
diff --git a/onnxruntime/contrib_ops/cuda/bert/rotary_embedding_impl.cu b/onnxruntime/contrib_ops/cuda/bert/rotary_embedding_impl.cu
index 9f162783b3a30..55648cc944249 100644
--- a/onnxruntime/contrib_ops/cuda/bert/rotary_embedding_impl.cu
+++ b/onnxruntime/contrib_ops/cuda/bert/rotary_embedding_impl.cu
@@ -189,8 +189,7 @@ Status LaunchRotaryEmbeddingKernel(cudaStream_t stream, T* output, const T* inpu
                                    const int position_ids_format, const bool interleaved,
                                    const int max_threads_per_block,
                                    int4 in_strides, int4 out_strides,  // strides in bnsh coord
-                                   const bool use_fp32_arithmetic
-) {
+                                   const bool use_fp32_arithmetic) {
   // Note: Current implementation assumes head_size <= max_threads_per_block
   // because head_size is currently large for LLaMA-2. For smaller head_size
   // and num_heads values, we can create a block as `block(num_heads, head_size, 1)`
@@ -221,8 +220,10 @@ Status LaunchRotaryEmbeddingKernel(cudaStream_t stream, T* output, const T* inpu
           interleaved, in_strides, out_strides);
     }
   };
-  if (use_fp32_arithmetic) launch(std::true_type{});
-  else launch(std::false_type{});
+  if (use_fp32_arithmetic)
+    launch(std::true_type{});
+  else
+    launch(std::false_type{});
 
   return CUDA_CALL(cudaGetLastError());
 }
diff --git a/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8.cu b/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8.cu
index 599a8ec957852..16ae3c0d5c949 100644
--- a/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8.cu
+++ b/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8.cu
@@ -516,7 +516,7 @@ struct Fp8GemvMma<__nv_bfloat16> {
 
 template <typename T>
 __device__ __forceinline__ T* Fp8OutputColumn(T* output, int col, int& stride,
-                                             const Fp8GemvOutputSplit& split) {
+                                              const Fp8GemvOutputSplit& split) {
   for (int i = 0; i < split.count; ++i) {
     if (col < split.widths[i]) {
       stride = split.widths[i];
@@ -946,7 +946,7 @@ bool Fp8MmaGb10TuningEnabled() {
 }
 
 bool MatMulBlockScaledFp8GemvSupportsSplitOutput(int m, int k, int block_size,
-                                                const cudaDeviceProp& device_prop) {
+                                                 const cudaDeviceProp& device_prop) {
 #if !defined(DISABLE_FLOAT8_TYPES) && defined(CUDA_VERSION) && CUDA_VERSION >= 11080
   return device_prop.major >= 8 && m > 0 && m <= kFp8MmaGemvTileM &&
          k >= 256 && k % 64 == 0 && block_size % 64 == 0 && Fp8GemvMmaEnabled();
@@ -1041,19 +1041,19 @@ static Status LaunchMatMulBlockScaledFp8GemvImpl(void* y,
     const auto launch_mma = [&]<int KSplit, int MTiles>() {
       const dim3 mma_threads{32, KSplit};
       const auto launch_variant = [&]<bool SplitOutput>() {
-#define ORT_FP8_LAUNCH_MMA(kernel_name)                                                       \
-  do {                                                                                      \
-    if (is_bf16) {                                                                          \
-      kernel_name<KSplit, MTiles, SplitOutput><<<mma_blocks, mma_threads, 0, stream>>>(       \
-          reinterpret_cast<__nv_bfloat16*>(y), reinterpret_cast<const __nv_bfloat16*>(a), b,  \
-          weight_scale, reinterpret_cast<const __nv_bfloat16*>(bias), act_scale, m, n, k,     \
-          block_size, k_blocks, split);                                                     \
-    } else {                                                                               \
-      kernel_name<KSplit, MTiles, SplitOutput><<<mma_blocks, mma_threads, 0, stream>>>(       \
-          reinterpret_cast<half*>(y), reinterpret_cast<const half*>(a), b,                    \
-          weight_scale, reinterpret_cast<const half*>(bias), act_scale, m, n, k,              \
-          block_size, k_blocks, split);                                                     \
-    }                                                                                      \
+#define ORT_FP8_LAUNCH_MMA(kernel_name)                                                      \
+  do {                                                                                       \
+    if (is_bf16) {                                                                           \
+      kernel_name<KSplit, MTiles, SplitOutput><<<mma_blocks, mma_threads, 0, stream>>>(      \
+          reinterpret_cast<__nv_bfloat16*>(y), reinterpret_cast<const __nv_bfloat16*>(a), b, \
+          weight_scale, reinterpret_cast<const __nv_bfloat16*>(bias), act_scale, m, n, k,    \
+          block_size, k_blocks, split);                                                      \
+    } else {                                                                                 \
+      kernel_name<KSplit, MTiles, SplitOutput><<<mma_blocks, mma_threads, 0, stream>>>(      \
+          reinterpret_cast<half*>(y), reinterpret_cast<const half*>(a), b,                   \
+          weight_scale, reinterpret_cast<const half*>(bias), act_scale, m, n, k,             \
+          block_size, k_blocks, split);                                                      \
+    }                                                                                        \
   } while (0)
         if constexpr (KSplit == 16 && MTiles == 1) {
           if (pin_residency) {
diff --git a/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8.h b/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8.h
index 35e2245543313..b54a37da08a87 100644
--- a/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8.h
+++ b/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8.h
@@ -45,7 +45,7 @@ struct Fp8GemvOutputSplit {
 };
 
 bool MatMulBlockScaledFp8GemvSupportsSplitOutput(int m, int k, int block_size,
-                                                const cudaDeviceProp& device_prop);
+                                                 const cudaDeviceProp& device_prop);
 
 #if defined(USE_DEEP_GEMM)
 Status LaunchPrepareMatMulFp8DeepGemm(void* a_quant, float* a_scales, const void* a,
diff --git a/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.cu b/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.cu
index 8f82da6379ab5..4410df08e47cd 100644
--- a/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.cu
+++ b/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.cu
@@ -667,7 +667,7 @@ __global__ void QsaMergeTileTopKKernel(const uint64_t* input_keys, uint64_t* out
 }
 
 __global__ void QsaMergeTwoLevelsTopKKernel(const uint64_t* input_keys, uint64_t* output_keys,
-                                          int key_stride, int list_width, int rows) {
+                                            int key_stride, int list_width, int rows) {
   __shared__ uint64_t inputs[4 * kBoundedTopKMax];
   __shared__ uint64_t intermediate[2 * kBoundedTopKMax];
   const int groups_per_row = (key_stride + 4 * list_width - 1) / (4 * list_width);
@@ -1515,7 +1515,7 @@ Status LaunchQsaPackedSparseAttentionIndexer(
     const int64_t tile_work =
         static_cast<int64_t>(params.total_tokens) * std::min(tile_count, kHierarchicalScoreMaxBlocksPerRow);
     QsaScoreTileTopKKernel<T, TState><<<static_cast<int>(std::min<int64_t>(tile_work, kSaiMaxGridDimX)),
-                                kHierarchicalTileScoreThreads, 0, stream>>>(
+                                        kHierarchicalTileScoreThreads, 0, stream>>>(
         present_key_state, query_rotated, cumulative_sequence_lengths, past_sequence_lengths, position_ids,
         present_state_lengths, overflow_flags, merge_input, tile_count, params);
     static const bool use_radix_topk =
@@ -1530,7 +1530,7 @@ Status LaunchQsaPackedSparseAttentionIndexer(
           const int groups_per_row = (key_stride + 4 * list_width - 1) / (4 * list_width);
           const int64_t merge_work = static_cast<int64_t>(params.total_tokens) * groups_per_row;
           QsaMergeTwoLevelsTopKKernel<<<static_cast<int>(std::min<int64_t>(merge_work, kSaiMaxGridDimX)),
-                                       kHierarchicalMergeThreads, 0, stream>>>(
+                                        kHierarchicalMergeThreads, 0, stream>>>(
               merge_input, merge_output, key_stride, list_width, params.total_tokens);
           list_width *= 4;
         } else {
@@ -1663,10 +1663,10 @@ Status LaunchCsaPackedSparseAttentionIndexer(
   return CUDA_CALL(cudaGetLastError());
 }
 
-#define INSTANTIATE_QSA(T, TState)                                                                                \
-  template Status LaunchQsaPackedSparseAttentionIndexer<T, TState>(                                               \
-      cudaStream_t, const PackedSparseAttentionIndexerParams&, const T*, const T*, const T*, const T*,            \
-      const T*, const T*, const int32_t*, const int32_t*, const int64_t*, const TState*, const T*,                \
+#define INSTANTIATE_QSA(T, TState)                                                                     \
+  template Status LaunchQsaPackedSparseAttentionIndexer<T, TState>(                                    \
+      cudaStream_t, const PackedSparseAttentionIndexerParams&, const T*, const T*, const T*, const T*, \
+      const T*, const T*, const int32_t*, const int32_t*, const int64_t*, const TState*, const T*,     \
       const int32_t*, const int32_t*, bool, int32_t*, int32_t*, TState*, T*, int32_t*, T*, float*, int32_t*);
 
 #define INSTANTIATE_CSA(T)                                                                                        \
diff --git a/onnxruntime/core/graph/contrib_ops/bert_defs.cc b/onnxruntime/core/graph/contrib_ops/bert_defs.cc
index cf92e6f11b939..7f490a9a1c881 100644
--- a/onnxruntime/core/graph/contrib_ops/bert_defs.cc
+++ b/onnxruntime/core/graph/contrib_ops/bert_defs.cc
@@ -3477,55 +3477,55 @@ that are multiplied to query and key before the inner product of query and key i
 )DOC";
 static OpSchema RotaryEmbeddingSchema() {
   return OpSchema()
-        .SetDoc(RotaryEmbedding_ver1_doc)
-        .Attr("scale",
-              "Custom scale will be used if specified. Default value is 1.0",
-              AttributeProto::FLOAT,
-              OPTIONAL_VALUE)
-        .Attr("interleaved",
-              "Indicates whether the input has real and imaginary parts interleaved. "
-              "Default value is 0 (False), meaning the first half of the input consists of real values "
-              "and the second half consists of imaginary values.",
-              AttributeProto::INT,
-              OPTIONAL_VALUE)
-        .Attr("rotary_embedding_dim",
-              "Rotary embedding dimension. Default value is 0.",
-              AttributeProto::INT,
-              OPTIONAL_VALUE)
-        .Attr("num_heads",
-              "Number of attention heads. Default value is 0. Must use with rotary_embedding_dim",
-              AttributeProto::INT,
-              OPTIONAL_VALUE)
-        .Attr("is_packed_batching",
-              "ragged batch inputs or not. Default value is 0",
-              AttributeProto::INT,
-              OPTIONAL_VALUE)
-        .Input(0,
-               "input",
-               "3D tensor with shape (batch_size, sequence_length, hidden_size) or 4D with shape (batch_size, num_heads, sequence_length, head_size)",
-               "T")
-        .Input(1,
-               "position_ids",
-               "1D tensor with shape (1) or 2D tensor with shape (batch_size, sequence_length)",
-               "M")
-        .Input(2,
-               "cos_cache",
-               "2D tensor with shape (max_sequence_length, head_size / 2) or (max_sequence_length, rotary_embedding_dim / 2)",
-               "T")
-        .Input(3,
-               "sin_cache",
-               "2D tensor with shape (max_sequence_length, head_size / 2) or (max_sequence_length, rotary_embedding_dim / 2)",
-               "T")
-        .Output(0,
-                "output",
-                "tensor with same shape as input.",
-                "T")
-        .TypeConstraint("T", {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, "Constrain input and output types to float tensors.")
-        .TypeConstraint("M", {"tensor(int64)"}, "Constrain input and output types to integer tensors")
-        .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) {
-          propagateElemTypeFromInputToOutput(ctx, 0, 0);
-          propagateShapeFromInputToOutput(ctx, 0, 0);
-        });
+      .SetDoc(RotaryEmbedding_ver1_doc)
+      .Attr("scale",
+            "Custom scale will be used if specified. Default value is 1.0",
+            AttributeProto::FLOAT,
+            OPTIONAL_VALUE)
+      .Attr("interleaved",
+            "Indicates whether the input has real and imaginary parts interleaved. "
+            "Default value is 0 (False), meaning the first half of the input consists of real values "
+            "and the second half consists of imaginary values.",
+            AttributeProto::INT,
+            OPTIONAL_VALUE)
+      .Attr("rotary_embedding_dim",
+            "Rotary embedding dimension. Default value is 0.",
+            AttributeProto::INT,
+            OPTIONAL_VALUE)
+      .Attr("num_heads",
+            "Number of attention heads. Default value is 0. Must use with rotary_embedding_dim",
+            AttributeProto::INT,
+            OPTIONAL_VALUE)
+      .Attr("is_packed_batching",
+            "ragged batch inputs or not. Default value is 0",
+            AttributeProto::INT,
+            OPTIONAL_VALUE)
+      .Input(0,
+             "input",
+             "3D tensor with shape (batch_size, sequence_length, hidden_size) or 4D with shape (batch_size, num_heads, sequence_length, head_size)",
+             "T")
+      .Input(1,
+             "position_ids",
+             "1D tensor with shape (1) or 2D tensor with shape (batch_size, sequence_length)",
+             "M")
+      .Input(2,
+             "cos_cache",
+             "2D tensor with shape (max_sequence_length, head_size / 2) or (max_sequence_length, rotary_embedding_dim / 2)",
+             "T")
+      .Input(3,
+             "sin_cache",
+             "2D tensor with shape (max_sequence_length, head_size / 2) or (max_sequence_length, rotary_embedding_dim / 2)",
+             "T")
+      .Output(0,
+              "output",
+              "tensor with same shape as input.",
+              "T")
+      .TypeConstraint("T", {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, "Constrain input and output types to float tensors.")
+      .TypeConstraint("M", {"tensor(int64)"}, "Constrain input and output types to integer tensors")
+      .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) {
+        propagateElemTypeFromInputToOutput(ctx, 0, 0);
+        propagateShapeFromInputToOutput(ctx, 0, 0);
+      });
 }
 
 ONNX_MS_OPERATOR_SET_SCHEMA(RotaryEmbedding, 1, RotaryEmbeddingSchema());
@@ -4689,15 +4689,18 @@ ONNX_MS_OPERATOR_SET_SCHEMA(
     CausalConvWithState, 1,
     OpSchema()
         .SetDoc(CausalConvWithState_ver1_doc)
-        .Attr("input_channels", "CUDA channels-last input channel slice width; 0 selects all channels. "
+        .Attr("input_channels",
+              "CUDA channels-last input channel slice width; 0 selects all channels. "
               "Nondefault extended attributes require rank-3 input and state_window=0.",
               AttributeProto::INT, static_cast<int64_t>(0))
         .Attr("input_channel_offset", "Starting channel of the channels-last input slice (CUDA only).",
               AttributeProto::INT, static_cast<int64_t>(0))
-        .Attr("state_length", "Carry history length (CUDA only); 0 uses (kernel_size-1)*dilation. "
+        .Attr("state_length",
+              "Carry history length (CUDA only); 0 uses (kernel_size-1)*dilation. "
               "An explicit length must cover the receptive field; convolution reads its trailing history.",
               AttributeProto::INT, static_cast<int64_t>(0))
-        .Attr("state_channels_last", "State layout override (CUDA only): -1 follows channels_last, "
+        .Attr("state_channels_last",
+              "State layout override (CUDA only): -1 follows channels_last, "
               "0 uses [B,C,state_length], 1 uses [B,state_length,C].",
               AttributeProto::INT, static_cast<int64_t>(-1))
         .Attr("activation",
@@ -4781,9 +4784,11 @@ ONNX_MS_OPERATOR_SET_SCHEMA(
                 "(seq_len - W + j). Extended attributes override the state length and layout "
                 "as for past_state.",
                 "T")
-        .Output(2, "first_state", "Optional carry state after the first input position, with the same "
+        .Output(2, "first_state",
+                "Optional carry state after the first input position, with the same "
                 "layout and length as present_state. CUDA only; requires rank-3 channels-last input "
-                "and state_window=0.", "T", OpSchema::Optional)
+                "and state_window=0.",
+                "T", OpSchema::Optional)
         .TypeConstraint("T",
                         {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"},
                         "Constrain input and output types to float tensors.")
@@ -5763,20 +5768,26 @@ ONNX_MS_OPERATOR_SET_SCHEMA(
         .Attr("q_stride", "Per-head stride in the packed projection.", AttributeProto::INT, static_cast<int64_t>(0))
         .Attr("tok_stride", "Per-token stride in the packed projection.", AttributeProto::INT, static_cast<int64_t>(0))
         .Attr("q_offset", "Offset of the query projection within each token.", AttributeProto::INT, static_cast<int64_t>(0))
-        .Attr("broadcast_gate", "When one, projection has the attention shape with final dimension one. "
-              "Packed layout attributes must be zero. CUDA only.", AttributeProto::INT, static_cast<int64_t>(0))
-        .Attr("sigmoid_in_fp32", "One computes sigmoid in FP32 before BF16 conversion. Zero preserves "
+        .Attr("broadcast_gate",
+              "When one, projection has the attention shape with final dimension one. "
+              "Packed layout attributes must be zero. CUDA only.",
+              AttributeProto::INT, static_cast<int64_t>(0))
+        .Attr("sigmoid_in_fp32",
+              "One computes sigmoid in FP32 before BF16 conversion. Zero preserves "
               "the input-type intermediate arithmetic of CUDA Sigmoid, including subtraction on negative inputs.",
               AttributeProto::INT, static_cast<int64_t>(1))
-        .Attr("silu_gate", "When one, projection and attention have identical shapes. Multiply the "
+        .Attr("silu_gate",
+              "When one, projection and attention have identical shapes. Multiply the "
               "BF16-rounded sigmoid by projection and round to BF16 before multiplying attention. "
               "Cannot be combined with broadcast_gate or packed layout attributes.",
               AttributeProto::INT, static_cast<int64_t>(0))
         .Input(0, "attention", "BF16 attention values.", "T")
         .Input(1, "projection", "BF16 packed projection containing gate values.", "T")
         .Input(2, "residual", "Optional BF16 residual with the attention shape.", "T", OpSchema::Optional)
-        .Output(0, "output", "Attention multiplied by BF16-rounded sigmoid gates, rounded to BF16 "
-                "before optionally adding residual and rounding again.", "T")
+        .Output(0, "output",
+                "Attention multiplied by BF16-rounded sigmoid gates, rounded to BF16 "
+                "before optionally adding residual and rounding again.",
+                "T")
         .TypeConstraint("T", {"tensor(bfloat16)"}, "Inputs and output use BF16.")
         .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) {
           propagateElemTypeFromInputToOutput(ctx, 0, 0);
@@ -5947,64 +5958,66 @@ WebGPU currently supports only the default mode without Y_fp32.
 
 static OpSchema BranchwiseRMSNormSchema() {
   return OpSchema()
-        .SetDoc(BranchwiseRMSNorm_ver1_doc)
-        .Attr("epsilon", "Epsilon added before reciprocal square root.",
-              AttributeProto::FLOAT, 1e-5f)
-        .Attr("num_branches",
-              "Number of branches for flattened input. Omit or set to zero for grouped input.",
-              AttributeProto::INT, static_cast<int64_t>(0))
-        .Attr("use_onnx_arithmetic",
-              "Use explicit FP32 mean/sqrt/divide/multiply arithmetic instead of reciprocal RMS. "
-              "Supported values are 0 (default) and 1.",
-              AttributeProto::INT, static_cast<int64_t>(0))
-        .Input(0, "X", "Grouped (..., C, H) or flattened (..., C * H) input.", "T")
-        .Input(1, "scale", "Optional scale with shape (C * H), (C, H), or (H).",
-               "M", OpSchema::Optional)
-        .Output(0, "Y", "RMS-normalized output with the same shape as X.", "T")
-        .Output(1, "Y_fp32", "Optional float32 widening of the already T-rounded Y. "
-                "Requires use_onnx_arithmetic=1.", "F", OpSchema::Optional)
-        .TypeConstraint("T",
-                        {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"},
-                        "Constrain input and output to floating-point tensors.")
-        .TypeConstraint("M",
-                        {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"},
-                        "Constrain scale to floating-point tensors.")
-        .TypeConstraint("F", {"tensor(float)"}, "Widened output is float32.")
-        .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) {
-          propagateElemTypeFromInputToOutput(ctx, 0, 0);
-          const auto start = getAttribute(ctx, "input_start", 0);
-          const auto width = getAttribute(ctx, "input_width", 0);
-          if (start < 0 || width < 0 || (width == 0 && start != 0)) {
-            fail_shape_inference("input_start and input_width must describe a nonnegative slice");
+      .SetDoc(BranchwiseRMSNorm_ver1_doc)
+      .Attr("epsilon", "Epsilon added before reciprocal square root.",
+            AttributeProto::FLOAT, 1e-5f)
+      .Attr("num_branches",
+            "Number of branches for flattened input. Omit or set to zero for grouped input.",
+            AttributeProto::INT, static_cast<int64_t>(0))
+      .Attr("use_onnx_arithmetic",
+            "Use explicit FP32 mean/sqrt/divide/multiply arithmetic instead of reciprocal RMS. "
+            "Supported values are 0 (default) and 1.",
+            AttributeProto::INT, static_cast<int64_t>(0))
+      .Input(0, "X", "Grouped (..., C, H) or flattened (..., C * H) input.", "T")
+      .Input(1, "scale", "Optional scale with shape (C * H), (C, H), or (H).",
+             "M", OpSchema::Optional)
+      .Output(0, "Y", "RMS-normalized output with the same shape as X.", "T")
+      .Output(1, "Y_fp32",
+              "Optional float32 widening of the already T-rounded Y. "
+              "Requires use_onnx_arithmetic=1.",
+              "F", OpSchema::Optional)
+      .TypeConstraint("T",
+                      {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"},
+                      "Constrain input and output to floating-point tensors.")
+      .TypeConstraint("M",
+                      {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"},
+                      "Constrain scale to floating-point tensors.")
+      .TypeConstraint("F", {"tensor(float)"}, "Widened output is float32.")
+      .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) {
+        propagateElemTypeFromInputToOutput(ctx, 0, 0);
+        const auto start = getAttribute(ctx, "input_start", 0);
+        const auto width = getAttribute(ctx, "input_width", 0);
+        if (start < 0 || width < 0 || (width == 0 && start != 0)) {
+          fail_shape_inference("input_start and input_width must describe a nonnegative slice");
+        }
+        if (width != 0 && (getAttribute(ctx, "use_onnx_arithmetic", 0) != 1 ||
+                           getAttribute(ctx, "num_branches", 0) != 0)) {
+          fail_shape_inference("input_width requires use_onnx_arithmetic=1 and grouped input");
+        }
+        if (hasInputShape(ctx, 0)) {
+          auto shape = getInputShape(ctx, 0);
+          if (width != 0) {
+            if (shape.dim_size() < 2) fail_shape_inference("grouped input requires rank at least two");
+            auto* last = shape.mutable_dim(shape.dim_size() - 1);
+            if (last->has_dim_value() && (start > last->dim_value() || width > last->dim_value() - start)) {
+              fail_shape_inference("input slice exceeds the last dimension");
+            }
+            last->set_dim_value(width);
           }
-          if (width != 0 && (getAttribute(ctx, "use_onnx_arithmetic", 0) != 1 ||
-                             getAttribute(ctx, "num_branches", 0) != 0)) {
-            fail_shape_inference("input_width requires use_onnx_arithmetic=1 and grouped input");
+          updateOutputShape(ctx, 0, shape);
+        }
+        if (ctx.getNumOutputs() > 1) {
+          if (getAttribute(ctx, "use_onnx_arithmetic", 0) != 1) {
+            fail_shape_inference("Y_fp32 requires use_onnx_arithmetic=1");
           }
+          updateOutputElemType(ctx, 1, ONNX_NAMESPACE::TensorProto_DataType_FLOAT);
           if (hasInputShape(ctx, 0)) {
             auto shape = getInputShape(ctx, 0);
-            if (width != 0) {
-              if (shape.dim_size() < 2) fail_shape_inference("grouped input requires rank at least two");
-              auto* last = shape.mutable_dim(shape.dim_size() - 1);
-              if (last->has_dim_value() && (start > last->dim_value() || width > last->dim_value() - start)) {
-                fail_shape_inference("input slice exceeds the last dimension");
-              }
-              last->set_dim_value(width);
-            }
-            updateOutputShape(ctx, 0, shape);
+            if (width != 0) shape.mutable_dim(shape.dim_size() - 1)->set_dim_value(width);
+            updateOutputShape(ctx, 1, shape);
           }
-          if (ctx.getNumOutputs() > 1) {
-            if (getAttribute(ctx, "use_onnx_arithmetic", 0) != 1) {
-              fail_shape_inference("Y_fp32 requires use_onnx_arithmetic=1");
-            }
-            updateOutputElemType(ctx, 1, ONNX_NAMESPACE::TensorProto_DataType_FLOAT);
-            if (hasInputShape(ctx, 0)) {
-              auto shape = getInputShape(ctx, 0);
-              if (width != 0) shape.mutable_dim(shape.dim_size() - 1)->set_dim_value(width);
-              updateOutputShape(ctx, 1, shape);
-            }
-          }
-        });
+        }
+      });
 }
 
 ONNX_MS_OPERATOR_SET_SCHEMA(BranchwiseRMSNorm, 1, BranchwiseRMSNormSchema());
@@ -6013,7 +6026,8 @@ ONNX_MS_OPERATOR_SET_SCHEMA(
     BranchwiseRMSNormSchema()
         .Attr("input_start", "CUDA-only starting offset in the last input dimension.",
               AttributeProto::INT, static_cast<int64_t>(0))
-        .Attr("input_width", "CUDA-only width selected from the last input dimension before normalization. "
+        .Attr("input_width",
+              "CUDA-only width selected from the last input dimension before normalization. "
               "Zero retains the full input. Requires grouped input and use_onnx_arithmetic=1. "
               "Outputs retain the leading dimensions and use this last dimension.",
               AttributeProto::INT, static_cast<int64_t>(0)));
diff --git a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc
index 2fe6a490ee90a..3ead362cd46b6 100644
--- a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc
+++ b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc
@@ -3215,7 +3215,7 @@ The partition widths must be positive and sum to N.)DOC")
           const auto* split = ctx.getAttribute("split_sizes");
           if ((split == nullptr && ctx.getNumOutputs() != 1) ||
               (split != nullptr && (split->ints_size() < 1 || split->ints_size() > 8 ||
-                                   static_cast<size_t>(split->ints_size()) != ctx.getNumOutputs()))) {
+                                    static_cast<size_t>(split->ints_size()) != ctx.getNumOutputs()))) {
             fail_shape_inference("Output count must match split_sizes, or be one when split_sizes is absent.");
           }
           int64_t split_total = 0;
diff --git a/onnxruntime/core/providers/cuda/reduction/reduction_functions.cu b/onnxruntime/core/providers/cuda/reduction/reduction_functions.cu
index bf960242503d4..75024a420d3a9 100644
--- a/onnxruntime/core/providers/cuda/reduction/reduction_functions.cu
+++ b/onnxruntime/core/providers/cuda/reduction/reduction_functions.cu
@@ -539,7 +539,7 @@ struct RoundedSquare {
 
 template <typename TIn>
 Status reduce_square_matrix_columns(cudaStream_t stream, const TIn* input, float* output,
-                                   int m, int n, void* buffer, size_t buffer_size, int input_row_stride) {
+                                    int m, int n, void* buffer, size_t buffer_size, int input_row_stride) {
   return detail::call_reduce_matrix_columns<TIn, float, RoundedSquare, Identity, false>(
       stream, input, output, m, n, buffer, buffer_size, input_row_stride);
 }
diff --git a/onnxruntime/core/providers/cuda/reduction/reduction_functions.h b/onnxruntime/core/providers/cuda/reduction/reduction_functions.h
index e5bb1b03820f0..9c877c0743b31 100644
--- a/onnxruntime/core/providers/cuda/reduction/reduction_functions.h
+++ b/onnxruntime/core/providers/cuda/reduction/reduction_functions.h
@@ -109,7 +109,7 @@ Status reduce_matrix_columns(cudaStream_t stream, const TIn* input, TOut* output
 /** Sum FP32-rounded squares across columns, preserving reduce_matrix_columns' summation order. */
 template <typename TIn>
 Status reduce_square_matrix_columns(cudaStream_t stream, const TIn* input, float* output,
-                                   int m, int n, void* buffer, size_t buffer_size, int input_row_stride = 0);
+                                    int m, int n, void* buffer, size_t buffer_size, int input_row_stride = 0);
 
 /**
  * Computes ReduceSum for an arbitrary set of axes. This is the general fallback
diff --git a/onnxruntime/test/contrib_ops/causal_conv_with_state_op_test.cc b/onnxruntime/test/contrib_ops/causal_conv_with_state_op_test.cc
index 8e015e80d5b32..7cac5d1d53c16 100644
--- a/onnxruntime/test/contrib_ops/causal_conv_with_state_op_test.cc
+++ b/onnxruntime/test/contrib_ops/causal_conv_with_state_op_test.cc
@@ -365,12 +365,16 @@ void TestPackedChannelHistory(int length, bool state_last, bool use_past, bool u
   test.AddAttribute<std::string>("activation", silu ? "silu" : "none");
   test.AddInput<T>("input", {batch, length, width}, input);
   test.AddInput<T>("weight", {channels, 1, kernel}, weight);
-  if (use_bias) test.AddInput<T>("bias", {channels}, bias);
-  else test.AddOptionalInputEdge<T>();
+  if (use_bias)
+    test.AddInput<T>("bias", {channels}, bias);
+  else
+    test.AddOptionalInputEdge<T>();
   const std::vector<int64_t> state_shape = state_last ? std::vector<int64_t>{batch, history, channels}
-                                                    : std::vector<int64_t>{batch, channels, history};
-  if (use_past) test.AddInput<T>("past", state_shape, past);
-  else test.AddOptionalInputEdge<T>();
+                                                      : std::vector<int64_t>{batch, channels, history};
+  if (use_past)
+    test.AddInput<T>("past", state_shape, past);
+  else
+    test.AddOptionalInputEdge<T>();
   test.AddOutput<T>("output", {batch, length, channels}, output);
   test.AddOutput<T>("present", state_shape, present);
   if (first_output) test.AddOutput<T>("first", state_shape, first);
diff --git a/onnxruntime/test/contrib_ops/cuda_kernels/qmoe_fp8_compaction_test.cc b/onnxruntime/test/contrib_ops/cuda_kernels/qmoe_fp8_compaction_test.cc
index e414fe31f873e..f54836ae48a94 100644
--- a/onnxruntime/test/contrib_ops/cuda_kernels/qmoe_fp8_compaction_test.cc
+++ b/onnxruntime/test/contrib_ops/cuda_kernels/qmoe_fp8_compaction_test.cc
@@ -69,7 +69,7 @@ void TestRoutingMaps(int rows, int experts, int k, bool normalize, bool tied) {
   CudaBuffer offsets((experts + 1) * sizeof(int64_t));
   input.Upload(logits.data());
   qmoe::LaunchSoftmaxTopK(input.As<T>(), scales.As<float>(), indices.As<int>(),
-                         rows, experts, k, normalize, nullptr);
+                          rows, experts, k, normalize, nullptr);
   ASSERT_TRUE(qmoe::TryLaunchSmallDecodeExpertMaps(
       indices.As<int>(), r2u.As<int>(), u2r.As<int>(), sorted.As<int>(),
       offsets.As<int64_t>(), rows, experts, k, nullptr));
diff --git a/onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc b/onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc
index 659bc157e9982..bc648f4f5ba37 100644
--- a/onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc
+++ b/onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc
@@ -168,11 +168,11 @@ void RunDynamicSparseAttentionCase(
   }
   if (c.has_past) {
     tester.AddInput<CacheT>("past_key",
-                       {c.batch_size, c.kv_num_heads, c.cache_sequence_length, c.head_size},
-                       ToTensorData<CacheT>(c.past_key));
+                            {c.batch_size, c.kv_num_heads, c.cache_sequence_length, c.head_size},
+                            ToTensorData<CacheT>(c.past_key));
     tester.AddInput<CacheT>("past_value",
-                       {c.batch_size, c.kv_num_heads, c.cache_sequence_length, c.head_size},
-                       ToTensorData<CacheT>(c.past_value));
+                            {c.batch_size, c.kv_num_heads, c.cache_sequence_length, c.head_size},
+                            ToTensorData<CacheT>(c.past_value));
   } else {
     tester.AddOptionalInputEdge<T>();
     tester.AddOptionalInputEdge<T>();
@@ -239,20 +239,24 @@ void RunDynamicSparseAttentionCase(
     tester.AddOptionalInputEdge<int32_t>();
   }
   if (!c.key_scale.empty() || !c.value_scale.empty()) {
-    if (c.key_scale.empty()) tester.AddOptionalInputEdge<float>();
-    else tester.AddInput<float>("key_scale", {int64_t(c.key_scale.size())}, c.key_scale);
-    if (c.value_scale.empty()) tester.AddOptionalInputEdge<float>();
-    else tester.AddInput<float>("value_scale", {int64_t(c.value_scale.size())}, c.value_scale);
+    if (c.key_scale.empty())
+      tester.AddOptionalInputEdge<float>();
+    else
+      tester.AddInput<float>("key_scale", {int64_t(c.key_scale.size())}, c.key_scale);
+    if (c.value_scale.empty())
+      tester.AddOptionalInputEdge<float>();
+    else
+      tester.AddInput<float>("value_scale", {int64_t(c.value_scale.size())}, c.value_scale);
   }
 
   tester.AddOutput<T>("output", {c.batch_size, c.sequence_length, hidden_size},
                       ToTensorData<T>(c.expected_output));
   tester.AddOutput<CacheT>("present_key",
-                      {c.batch_size, c.kv_num_heads, c.cache_sequence_length, c.head_size},
-                      ToTensorData<CacheT>(c.expected_present_key));
+                           {c.batch_size, c.kv_num_heads, c.cache_sequence_length, c.head_size},
+                           ToTensorData<CacheT>(c.expected_present_key));
   tester.AddOutput<CacheT>("present_value",
-                      {c.batch_size, c.kv_num_heads, c.cache_sequence_length, c.head_size},
-                      ToTensorData<CacheT>(c.expected_present_value));
+                           {c.batch_size, c.kv_num_heads, c.cache_sequence_length, c.head_size},
+                           ToTensorData<CacheT>(c.expected_present_value));
   if (c.expected_validation_status.has_value()) {
     tester.AddOutput<int32_t>("validation_status", {1}, {*c.expected_validation_status});
   }
@@ -420,11 +424,11 @@ TEST(DynamicSparseAttentionTest, Int8CacheRejectsMissingOrMalformedScales_CUDA)
     } else {
       c.key_scale.resize(8);
     }
-    const char* error = variant == 0 ? "INT8 caches require per-channel"
-                       : variant == 1 ? "must be provided together"
-                                      : "INT8 cache scales must have shape";
+    const char* error = variant == 0   ? "INT8 caches require per-channel"
+                        : variant == 1 ? "must be provided together"
+                                       : "INT8 cache scales must have shape";
     RunDynamicSparseAttentionCase<float, int8_t>(c, std::move(cuda_ep),
-                                                OpTester::ExpectResult::kExpectFailure, error);
+                                                 OpTester::ExpectResult::kExpectFailure, error);
   }
 }
 
@@ -517,14 +521,14 @@ DynamicSparseAttentionCase MakeGroupedInt8CacheCase(int sequence_length, bool au
         double dot = 0.0;
         for (int d = 0; d < c.head_size; ++d) {
           const float key = from_auxiliary ? c.auxiliary_key[offset + d]
-                                          : c.expected_present_key[offset + d] * c.key_scale[kv * c.head_size + d];
+                                           : c.expected_present_key[offset + d] * c.key_scale[kv * c.head_size + d];
           dot += c.query[out + d] * key;
         }
         const double weight = std::exp(dot * c.scale);
         denominator += weight;
         for (int d = 0; d < c.head_size; ++d) {
           const float value = from_auxiliary ? c.auxiliary_value[offset + d]
-                                            : c.expected_present_value[offset + d] * c.value_scale[kv * c.head_size + d];
+                                             : c.expected_present_value[offset + d] * c.value_scale[kv * c.head_size + d];
           numerator[d] += static_cast<float>(weight * value);
         }
       }
diff --git a/onnxruntime/test/contrib_ops/fused_tensor_ops_test.cc b/onnxruntime/test/contrib_ops/fused_tensor_ops_test.cc
index 4aaeb82ed47ad..9d6797761e3a0 100644
--- a/onnxruntime/test/contrib_ops/fused_tensor_ops_test.cc
+++ b/onnxruntime/test/contrib_ops/fused_tensor_ops_test.cc
@@ -146,8 +146,8 @@ void RunSigmoidGateMulBroadcastResidualRounding(bool sigmoid_in_fp32, bool silu_
         const BFloat16 denominator(1.0f + BFloat16(exp_value).ToFloat());
         const BFloat16 reciprocal(1.0f / denominator.ToFloat());
         const BFloat16 rounded_sigmoid = sigmoid_in_fp32
-                                            ? BFloat16(sigmoid)
-                                            : (gate > 0 ? reciprocal : BFloat16(1.0f - reciprocal.ToFloat()));
+                                             ? BFloat16(sigmoid)
+                                             : (gate > 0 ? reciprocal : BFloat16(1.0f - reciprocal.ToFloat()));
         const BFloat16 activation = silu_gate ? BFloat16(gate * rounded_sigmoid.ToFloat()) : rounded_sigmoid;
         const BFloat16 product(values[i].ToFloat() * activation.ToFloat());
         expected[i] = residual ? BFloat16(product.ToFloat() + residuals[i].ToFloat()) : product;
@@ -186,7 +186,7 @@ TEST(FusedTensorOpsTest, SigmoidGateMulRejectsBroadcastAndResidualShapes) {
     test.AddAttribute<int64_t>("broadcast_gate", 1);
     test.AddInput<BFloat16>("attention", {2, 3}, ToBFloat16(std::vector<float>(6, 1.0f)));
     test.AddInput<BFloat16>("projection", bad_residual ? std::vector<int64_t>{2, 1} : std::vector<int64_t>{1, 2},
-                           ToBFloat16({1.0f, 1.0f}));
+                            ToBFloat16({1.0f, 1.0f}));
     if (bad_residual) test.AddInput<BFloat16>("residual", {1}, ToBFloat16({1.0f}));
     test.AddOutput<BFloat16>("output", {2, 3}, ToBFloat16(std::vector<float>(6, 0.0f)));
     test.Run(OpTester::ExpectResult::kExpectFailure,
@@ -208,8 +208,8 @@ TEST(FusedTensorOpsTest, SigmoidGateMulBroadcastEmptyAndInvalidAttributes) {
     test.AddInput<BFloat16>("residual", {0, 3}, {});
     test.AddOutput<BFloat16>("output", {0, 3}, {});
     const std::vector<std::string> errors = {"", "broadcast_gate must be zero or one",
-                                           "sigmoid_in_fp32 must be zero or one",
-                                           "cannot be combined with packed layout attributes"};
+                                             "sigmoid_in_fp32 must be zero or one",
+                                             "cannot be combined with packed layout attributes"};
     test.Run(mode == 0 ? OpTester::ExpectResult::kExpectSuccess : OpTester::ExpectResult::kExpectFailure,
              errors[mode], {}, nullptr, &providers);
   }
@@ -227,8 +227,8 @@ TEST(FusedTensorOpsTest, SigmoidGateMulSiluEmptyAndInvalidAttributes) {
     test.AddInput<BFloat16>("projection", {0, mode == 4 ? 1 : 3}, {});
     test.AddOutput<BFloat16>("output", {0, 3}, {});
     const std::vector<std::string> errors = {"", "silu_gate must be zero or one", "are exclusive",
-                                           "cannot be combined with packed layout attributes",
-                                           "silu gate must have the same shape"};
+                                             "cannot be combined with packed layout attributes",
+                                             "silu gate must have the same shape"};
     test.Run(mode == 0 ? OpTester::ExpectResult::kExpectSuccess : OpTester::ExpectResult::kExpectFailure,
              errors[mode], {}, nullptr, &providers);
   }
diff --git a/onnxruntime/test/contrib_ops/hyper_connection_op_test.cc b/onnxruntime/test/contrib_ops/hyper_connection_op_test.cc
index 7942e8fc22425..d7e1bc7acdd72 100644
--- a/onnxruntime/test/contrib_ops/hyper_connection_op_test.cc
+++ b/onnxruntime/test/contrib_ops/hyper_connection_op_test.cc
@@ -113,7 +113,10 @@ class StandardRMSChainTester : public OpTester {
  public:
   StandardRMSChainTester(int output_type, float epsilon, int64_t start = 0, int64_t width = 0)
       : OpTester("StandardRMSChain", 17, kOnnxDomain, false),
-        output_type_(output_type), epsilon_(epsilon), start_(start), width_(width) {}
+        output_type_(output_type),
+        epsilon_(epsilon),
+        start_(start),
+        width_(width) {}
 
   void AddNodes(Graph& graph, std::vector<NodeArg*>& inputs, std::vector<NodeArg*>& outputs,
                 std::vector<std::function<void(Node&)>>&) override {
@@ -132,7 +135,8 @@ class StandardRMSChainTester : public OpTester {
       source = value("selected");
       graph.AddNode("slice", "Slice", "",
                     {inputs[0], constant("starts", start_), constant("ends", start_ + width_),
-                     constant("axes", -1)}, {source});
+                     constant("axes", -1)},
+                    {source});
     }
     auto* x = value("float_x");
     graph.AddNode("cast", "Cast", "", {source}, {x})
@@ -200,7 +204,7 @@ void TestStandardRMSArithmetic(int type, bool packed = false) {
         std::vector<OrtValue> expected;
         for (OpTester* tester : {static_cast<OpTester*>(&reference), &fused}) {
           const std::vector<int64_t> input_shape = rows == 48 ? std::vector<int64_t>{1, 2, 24, stride}
-                                                            : std::vector<int64_t>{1, rows, stride};
+                                                              : std::vector<int64_t>{1, rows, stride};
           auto output_shape = input_shape;
           output_shape.back() = width;
           tester->AddInput<T>("X", input_shape, input);
diff --git a/onnxruntime/test/contrib_ops/matmul_block_scaled_fp8_test.cc b/onnxruntime/test/contrib_ops/matmul_block_scaled_fp8_test.cc
index 94dbfa82a4636..0695e717ea6eb 100644
--- a/onnxruntime/test/contrib_ops/matmul_block_scaled_fp8_test.cc
+++ b/onnxruntime/test/contrib_ops/matmul_block_scaled_fp8_test.cc
@@ -75,9 +75,7 @@ void TestFp8SplitOutputs() {
   if (!HasCudaEnvironment(800)) {
     GTEST_SKIP() << "CUDA device is required for split-output FP8 matmul.";
   }
-  for (const auto& [m, k] : {std::pair<int64_t, int64_t>{0, 256}, {1, 256}, {2, 256}, {2, 320},
-                            {2, 1088}, {9, 256},
-                            {32, 256}, {33, 256}, {65, 256}, {2, 16}, {2, 30}, {2, 0}}) {
+  for (const auto& [m, k] : {std::pair<int64_t, int64_t>{0, 256}, {1, 256}, {2, 256}, {2, 320}, {2, 1088}, {9, 256}, {32, 256}, {33, 256}, {65, 256}, {2, 16}, {2, 30}, {2, 0}}) {
     for (const bool activation_scale : {false, true}) {
       SCOPED_TRACE("M=" + std::to_string(m) + " K=" + std::to_string(k) +
                    " W8A8=" + std::to_string(activation_scale));
@@ -85,9 +83,9 @@ void TestFp8SplitOutputs() {
       constexpr int64_t block_size = 64;
       const int64_t k_blocks = (k + block_size - 1) / block_size;
       const std::vector<int64_t> widths = m == 2 && k == 256
-                                             ? std::vector<int64_t>{1, 1, 1, 1, 1, 1, 1, 10}
-                                             : (m == 33 ? std::vector<int64_t>{17}
-                                                        : std::vector<int64_t>{1, 7, 9});
+                                              ? std::vector<int64_t>{1, 1, 1, 1, 1, 1, 1, 10}
+                                              : (m == 33 ? std::vector<int64_t>{17}
+                                                         : std::vector<int64_t>{1, 7, 9});
       std::vector<T> a;
       std::vector<Float8E4M3FN> b;
       std::vector<float> scales;
@@ -162,8 +160,7 @@ TEST(MatMulBlockQuantizedFp8WeightOpTest, RejectsInvalidSplitSizes) {
   if (!HasCudaEnvironment(800)) {
     GTEST_SKIP() << "CUDA device is required for split-output FP8 matmul.";
   }
-  for (const std::vector<int64_t>& widths : {std::vector<int64_t>{}, {0, 17}, {-1, 18}, {8, 8},
-                                            {1, 1, 1, 1, 1, 1, 1, 1, 9}}) {
+  for (const std::vector<int64_t>& widths : {std::vector<int64_t>{}, {0, 17}, {-1, 18}, {8, 8}, {1, 1, 1, 1, 1, 1, 1, 1, 9}}) {
     OpTester test("MatMulBlockQuantizedFp8Weight", 1, kMSDomain);
     test.AddAttribute<std::vector<int64_t>>("split_sizes", widths);
     test.AddInput<BFloat16>("A", {1, 16}, std::vector<BFloat16>(16));
diff --git a/onnxruntime/test/contrib_ops/rotary_embedding_op_test.cc b/onnxruntime/test/contrib_ops/rotary_embedding_op_test.cc
index 9ebbc10c97bf2..a6957e18bc194 100644
--- a/onnxruntime/test/contrib_ops/rotary_embedding_op_test.cc
+++ b/onnxruntime/test/contrib_ops/rotary_embedding_op_test.cc
@@ -27,7 +27,10 @@ class FloatRotaryChainTester : public OpTester {
  public:
   FloatRotaryChainTester(int type, int heads, int rotary_dim, bool interleaved)
       : OpTester("RotaryEmbedding", 1, kMSDomain, false),
-        type_(type), heads_(heads), rotary_dim_(rotary_dim), interleaved_(interleaved) {}
+        type_(type),
+        heads_(heads),
+        rotary_dim_(rotary_dim),
+        interleaved_(interleaved) {}
 
   void AddNodes(Graph& graph, std::vector<NodeArg*>& inputs, std::vector<NodeArg*>& outputs,
                 std::vector<std::function<void(Node&)>>&) override {
@@ -74,7 +77,7 @@ void TestFloatRotaryArithmetic(int type) {
           const int rotary_dim = 64;
           const int cache_rows = batch * sequence + 3;
           std::vector<int64_t> shape = transposed ? std::vector<int64_t>{batch, heads, sequence, width}
-                                                 : std::vector<int64_t>{batch, sequence, heads * width};
+                                                  : std::vector<int64_t>{batch, sequence, heads * width};
           std::vector<T> input(batch * sequence * heads * width), cosine(cache_rows * rotary_dim / 2),
               sine(cosine.size());
           const bool scalar_position = !transposed;
@@ -94,8 +97,7 @@ void TestFloatRotaryArithmetic(int type) {
           std::vector<OrtValue> expected;
           for (OpTester* tester : {static_cast<OpTester*>(&reference), &fused}) {
             tester->AddInput<T>("input", shape, input);
-            tester->AddInput<int64_t>("positions", scalar_position ? std::vector<int64_t>{1}
-                                                                  : std::vector<int64_t>{batch, sequence}, positions);
+            tester->AddInput<int64_t>("positions", scalar_position ? std::vector<int64_t>{1} : std::vector<int64_t>{batch, sequence}, positions);
             tester->AddInput<T>("cosine", {cache_rows, rotary_dim / 2}, cosine);
             tester->AddInput<T>("sine", {cache_rows, rotary_dim / 2}, sine);
             tester->AddOutput<T>("output", shape, std::vector<T>(input.size()));
diff --git a/onnxruntime/test/python/onnxruntime_test_python_symbolic_shape_infer.py b/onnxruntime/test/python/onnxruntime_test_python_symbolic_shape_infer.py
index 96bb35672a6cf..6ac581d72cab8 100644
--- a/onnxruntime/test/python/onnxruntime_test_python_symbolic_shape_infer.py
+++ b/onnxruntime/test/python/onnxruntime_test_python_symbolic_shape_infer.py
@@ -882,9 +882,6 @@ def infer(query_width, packed):
         self.assertEqual(separate_graph.output[0].type.tensor_type.shape.dim[2].dim_value, 33)
         self.assertFalse(separate_graph.output[1].type.tensor_type.HasField("shape"))
 
-
-
-
     def test_unsqueeze_opset_11(self):
         graph = helper.make_graph(
             [