feat(mla): expose DSv4 TRTLLM-GEN RopeQuant

csrc/fmhaReduction.cu

@@ -14,6 +14,7 @@
  * limitations under the License.
  */
 
+#include <cuda_fp8.h>
 #include <cuda_runtime_api.h>
 #include <float.h>
 
@@ -31,8 +32,53 @@ namespace kernels {
 
 #define NumThreadsPerCta 512
 
+// DSv4 output epilogue: inverse RoPE, UE8M0 block scales, and E4M3 values. This matches the
+// fused-cubin output contract when a BF16 split-KV kernel needs a separate reduction.
+namespace dsv4 {
+
+constexpr int kHeadDim = 512;
+constexpr int kRopeDim = 64;
+constexpr int kQuantBlock = 128;
+constexpr int kHeadsPerGroup = 8;
+constexpr float kE4m3Max = 448.f;
+constexpr float kAmaxFloor = 1e-10f;
+
+__device__ __forceinline__ void inverseRopePair(float& even, float& odd, float cosVal,
+                                                float sinVal) {
+  float const e = even;
+  float const o = odd;
+  even = e * cosVal + o * sinVal;
+  odd = o * cosVal - e * sinVal;
+}
+
+__device__ __forceinline__ int32_t ue8m0ExponentFromAmax(float amax) {
+  amax = fmaxf(amax, kAmaxFloor);
+  int32_t exponent = ilogbf(amax) - 8;
+  if (ldexpf(kE4m3Max, exponent) < amax) {
+    ++exponent;
+  }
+  return min(max(exponent + 127, 0), 255);
+}
+
+__device__ __forceinline__ int64_t valueOffset(int32_t group, int32_t token, int32_t headInGroup,
+                                               int32_t dim, int32_t numTokens) {
+  return (((static_cast<int64_t>(group) * numTokens + token) * kHeadsPerGroup) + headInGroup) *
+             kHeadDim +
+         dim;
+}
+
+__device__ __forceinline__ int64_t scaleByteOffset(int32_t group, int32_t headInGroup,
+                                                   int32_t token, int32_t block,
+                                                   int64_t scaleBufM) {
+  return ((static_cast<int64_t>(group) * kHeadsPerGroup + headInGroup) * scaleBufM + token) *
+             sizeof(int32_t) +
+         block;
+}
+
+}  // namespace dsv4
+
 template <int32_t TileSizePerCtaQ, int32_t HeadDimPerCta, bool IsE4m3Bmm, typename DtypeO,
-          typename DtypePartialO>
+          typename DtypePartialO, bool Dsv4OutputEpilogue = false>
 __global__ void __launch_bounds__(NumThreadsPerCta, 2)
     fmhaReductionKernel(KernelParams const params, bool isTokenSparse, bool groupsTokensHeadsQ,
                         bool supportsVarSparseMlaTopKLens, int32_t numCtasForReduction,
@@ -284,6 +330,65 @@ __global__ void __launch_bounds__(NumThreadsPerCta, 2)
       mul(f2, f2, normalizedScale2);
     }
 
+    if constexpr (Dsv4OutputEpilogue) {
+      static_assert(
+          HeadDimPerCta % dsv4::kQuantBlock == 0 && dsv4::kQuantBlock % NumEltsPer16BVec == 0,
+          "A quant block must be owned by whole lanes of one CTA");
+      static_assert(NumEltsPer16BVec * sizeof(__nv_fp8_e4m3) == sizeof(uint2),
+                    "The E4M3 vector store assumes 8 elements");
+
+      int64_t const absRowIdx{softmaxStatsOffset + softmaxStatsRowIdx};
+      int32_t const tokenIdx{static_cast<int32_t>(absRowIdx / params.mNumHeadsQ)};
+      int32_t const headIdx{static_cast<int32_t>(absRowIdx % params.mNumHeadsQ)};
+      int32_t const dimIdx{headDimCtaIdxV * HeadDimPerCta + headDimIdx};
+
+      int32_t constexpr RopeStart{dsv4::kHeadDim - dsv4::kRopeDim};
+      if (dimIdx >= RopeStart) {
+        int32_t const position{params.ptrSeqLensKv[batchIdx] - seqLenQ + tokenIdx - seqOffsetQ};
+        float const* cosSin{params.ptrDsv4InvRopeCosSinCache +
+                            static_cast<int64_t>(position) * dsv4::kRopeDim};
+#pragma unroll
+        for (int32_t ii = 0; ii < NumEltsPer16BVec; ii += 2) {
+          int32_t const pairIdx{(dimIdx + ii - RopeStart) >> 1};
+          dsv4::inverseRopePair(outputVals[ii], outputVals[ii + 1], cosSin[pairIdx],
+                                cosSin[dsv4::kRopeDim / 2 + pairIdx]);
+        }
+      }
+
+      float amax{0.f};
+#pragma unroll
+      for (int32_t ii = 0; ii < NumEltsPer16BVec; ++ii) {
+        amax = fmaxf(amax, fabsf(outputVals[ii]));
+      }
+#pragma unroll
+      for (int32_t offset = 1; offset < dsv4::kQuantBlock / NumEltsPer16BVec; offset <<= 1) {
+        amax = fmaxf(amax, __shfl_xor_sync(0xffffffffu, amax, offset));
+      }
+      int32_t const biasedExp{dsv4::ue8m0ExponentFromAmax(amax)};
+
+      if (isValidRow) {
+        int32_t const group{headIdx / dsv4::kHeadsPerGroup};
+        int32_t const headInGroup{headIdx % dsv4::kHeadsPerGroup};
+        float const invScale{__frcp_rn(__int_as_float(biasedExp << 23))};
+        __nv_fp8_e4m3 quantized[NumEltsPer16BVec];
+#pragma unroll
+        for (int32_t ii = 0; ii < NumEltsPer16BVec; ++ii) {
+          quantized[ii] = static_cast<__nv_fp8_e4m3>(
+              fminf(fmaxf(outputVals[ii] * invScale, -dsv4::kE4m3Max), dsv4::kE4m3Max));
+        }
+        *reinterpret_cast<uint2*>(
+            reinterpret_cast<__nv_fp8_e4m3*>(params.ptrO) +
+            dsv4::valueOffset(group, tokenIdx, headInGroup, dimIdx, params.mSumOfSeqLensQ)) =
+            *reinterpret_cast<uint2 const*>(quantized);
+        if (dimIdx % dsv4::kQuantBlock == 0) {
+          reinterpret_cast<uint8_t*>(params.ptrDsv4OScale)[dsv4::scaleByteOffset(
+              group, headInGroup, tokenIdx, dimIdx / dsv4::kQuantBlock, params.mDsv4ScaleBufM)] =
+              static_cast<uint8_t>(biasedExp);
+        }
+      }
+      continue;
+    }
+
     // Convert the float values to DtypeO, and Store it to global memory.
     if (isValidRow) {
       convertAndStoreToGmem<DtypeO>(reinterpret_cast<char*>(oPtr + gmemStoreOffset), outputVals);
@@ -341,6 +446,8 @@ void runFmhaReduction(TllmGenFmhaKernelMetaInfo const& kernelMeta, KernelParams
     return;
   }
 
+  bool const dsv4OutputEpilogue = params.ptrDsv4OScale != nullptr;
+
   // This should only be enabled when using keepsMmaAbForGeneration kernel.
   FLASHINFER_CHECK(
       isKeepsMmaAbForGenerationKernel(static_cast<FmhaKernelType>(kernelMeta.mKernelType)),
@@ -404,7 +511,19 @@ void runFmhaReduction(TllmGenFmhaKernelMetaInfo const& kernelMeta, KernelParams
   // Select the kernel function pointer.
   void (*kernel)(KernelParams const, bool, bool, bool, int32_t, int32_t, int32_t, int32_t) =
       nullptr;
-  if (headDimPerCtaV == 64) {
+  if (dsv4OutputEpilogue) {
+    FLASHINFER_CHECK(kernelMeta.mGroupsHeadsQ && !kernelMeta.mGroupsTokensHeadsQ &&
+                         kernelMeta.mTileSizeQ == 64 && headDimPerCtaV >= 128 &&
+                         params.mNumHeadsQ % numHeadsPerCta == 0,
+                     "Not implemented");
+    if (headDimPerCtaV == 128) {
+      kernel = &fmhaReductionKernel<64, 128, true, __nv_bfloat16, __nv_bfloat16, true>;
+    } else if (headDimPerCtaV == 256) {
+      kernel = &fmhaReductionKernel<64, 256, true, __nv_bfloat16, __nv_bfloat16, true>;
+    } else {
+      kernel = &fmhaReductionKernel<64, 512, true, __nv_bfloat16, __nv_bfloat16, true>;
+    }
+  } else if (headDimPerCtaV == 64) {
     SELECT_FMHA_REDUCTION_KERNEL_WITH_HEAD_DIM_PER_CTA(64);
   } else if (headDimPerCtaV == 128) {
     SELECT_FMHA_REDUCTION_KERNEL_WITH_HEAD_DIM_PER_CTA(128);

csrc/trtllm_fmha_kernel_launcher.cu

@@ -181,7 +181,8 @@ void trtllm_paged_attention_launcher(
     bool enable_pdl, int64_t workspace_size, int64_t k_sf_stride_heads, int64_t k_sf_stride_batch,
     int64_t v_sf_stride_heads, int64_t v_sf_stride_batch, bool is_causal, int64_t lse_stride_tokens,
     int64_t lse_stride_heads, int64_t bf16q_fp8kv_transform_mode, bool use_fp16_softmax,
-    bool uses_spcompress, cudaStream_t stream) {
+    bool uses_spcompress, float const* dsv4_inv_rope_cos_sin_cache, float* dsv4_output_scale,
+    int64_t dsv4_scale_buf_m, cudaStream_t stream) {
   if (num_qo_heads % num_kv_heads != 0) {
     std::ostringstream err_msg;
     err_msg << "num_qo_heads must be a multiple of num_kv_heads, got num_kv_heads: " << num_kv_heads
@@ -249,6 +250,10 @@ void trtllm_paged_attention_launcher(
   runner_params.scaleSoftmaxLog2 = bmm1_scale * M_LOG2E;
   runner_params.scaleSoftmaxLog2Ptr = bmm1_scale_log2_ptr;
   runner_params.oSfPtr = out_scale_factor;
+  runner_params.dsv4InvRopeCosSinCachePtr = dsv4_inv_rope_cos_sin_cache;
+  runner_params.dsv4OScalePtr = dsv4_output_scale;
+  runner_params.mDsv4ScaleBufM = dsv4_scale_buf_m;
+  runner_params.mFusesDsv4InvRopeFp8Quant = dsv4_inv_rope_cos_sin_cache != nullptr;
   runner_params.mSfStartTokenIdx = o_sf_start_index;
   runner_params.mScaleSfO = o_sf_scale;
   TVM_FFI_ICHECK(o_sf_vec_size == 16 || o_sf_vec_size == -1)
@@ -592,7 +597,8 @@ void trtllm_paged_attention_decode(
       enable_block_sparse_attention, sm_count, enable_pdl, workspace_size, k_sf_stride_heads,
       k_sf_stride_batch, v_sf_stride_heads, v_sf_stride_batch, /*is_causal=*/true,
       lse_stride_tokens, lse_stride_heads, bf16q_fp8kv_transform_mode, use_fp16_softmax_value,
-      uses_spcompress_value, stream);
+      uses_spcompress_value, /*dsv4_inv_rope_cos_sin_cache=*/nullptr,
+      /*dsv4_output_scale=*/nullptr, /*dsv4_scale_buf_m=*/0, stream);
 }
 
 void trtllm_paged_attention_context(
@@ -736,7 +742,8 @@ void trtllm_paged_attention_context(
       /*enable_block_sparse_attention=*/false, sm_count, enable_pdl, workspace_size,
       k_sf_stride_heads, k_sf_stride_batch, v_sf_stride_heads, v_sf_stride_batch, is_causal,
       lse_stride_tokens, lse_stride_heads, /*bf16q_fp8kv_transform_mode=*/0, use_fp16_softmax_value,
-      uses_spcompress_value, stream);
+      uses_spcompress_value, /*dsv4_inv_rope_cos_sin_cache=*/nullptr,
+      /*dsv4_output_scale=*/nullptr, /*dsv4_scale_buf_m=*/0, stream);
 }
 
 void trtllm_ragged_attention_launcher(
@@ -973,10 +980,14 @@ void trtllm_paged_attention_decode_sparse_mla_dsv4(
     TensorView seq_lens, TensorView sparse_mla_top_k_lens, Variant<double, ffi::Tensor> bmm1_scale,
     Variant<double, ffi::Tensor> bmm2_scale, int64_t batch_size, int64_t max_q_len,
     int64_t sm_count, bool enable_pdl, int64_t workspace_size, Optional<TensorView> attention_sinks,
-    Optional<TensorView> cum_seq_lens_q) {
+    Optional<TensorView> cum_seq_lens_q, Optional<TensorView> dsv4_inv_rope_cos_sin_cache,
+    Optional<TensorView> dsv4_output_scale) {
   auto q_data_type = dl_dtype_to_tllm_data_type(query.dtype());
   auto kv_data_type = dl_dtype_to_tllm_data_type(primary_kv_cache.dtype());
   auto o_data_type = dl_dtype_to_tllm_data_type(out.dtype());
+  bool const fuses_dsv4_inv_rope_fp8_quant = dsv4_inv_rope_cos_sin_cache.has_value();
+  TVM_FFI_ICHECK_EQ(fuses_dsv4_inv_rope_fp8_quant, dsv4_output_scale.has_value())
+      << "dsv4_inv_rope_cos_sin_cache and dsv4_output_scale must be provided together";
 
   TVM_FFI_ICHECK(query.ndim() == 3) << "query must have shape [B*Q, H, D]";
   TVM_FFI_ICHECK(primary_kv_cache.ndim() == 4)
@@ -996,8 +1007,15 @@ void trtllm_paged_attention_decode_sparse_mla_dsv4(
   TVM_FFI_ICHECK_EQ(kv_data_type, q_data_type) << "primary_kv_cache dtype must match query dtype";
   TVM_FFI_ICHECK_EQ(dl_dtype_to_tllm_data_type(sliding_window_kv_cache.dtype()), q_data_type)
       << "sliding_window_kv_cache dtype must match query dtype";
-  TVM_FFI_ICHECK_EQ(o_data_type, Data_type::DATA_TYPE_BF16)
-      << "DeepSeek V4 sparse MLA output must be BF16";
+  if (fuses_dsv4_inv_rope_fp8_quant) {
+    TVM_FFI_ICHECK_EQ(q_data_type, Data_type::DATA_TYPE_E4M3)
+        << "DeepSeek V4 RopeQuant requires FP8 E4M3 query and KV inputs";
+    TVM_FFI_ICHECK_EQ(o_data_type, Data_type::DATA_TYPE_E4M3)
+        << "DeepSeek V4 RopeQuant output must be FP8 E4M3";
+  } else {
+    TVM_FFI_ICHECK_EQ(o_data_type, Data_type::DATA_TYPE_BF16)
+        << "DeepSeek V4 sparse MLA output must be BF16";
+  }
 
   int const sum_seq_q = query.size(0);
   int const num_qo_heads = query.size(1);
@@ -1006,6 +1024,9 @@ void trtllm_paged_attention_decode_sparse_mla_dsv4(
   int* cum_seq_lens_q_ptr =
       is_varlen_q ? static_cast<int*>(cum_seq_lens_q.value().data_ptr()) : nullptr;
   TVM_FFI_ICHECK_EQ(num_kv_heads, 1) << "DeepSeek V4 sparse MLA expects one KV head";
+  if (fuses_dsv4_inv_rope_fp8_quant) {
+    TVM_FFI_ICHECK_EQ(num_qo_heads, 128) << "DeepSeek V4 RopeQuant requires 128 query heads";
+  }
   TVM_FFI_ICHECK_EQ(sliding_window_kv_cache.size(-3), 1)
       << "sliding_window_kv_cache must have one KV head";
   TVM_FFI_ICHECK_EQ(seq_lens.size(0), batch_size);
@@ -1030,12 +1051,63 @@ void trtllm_paged_attention_decode_sparse_mla_dsv4(
       is_4bit(kv_data_type) ? primary_kv_cache.size(-1) * 2 : primary_kv_cache.size(-1);
   int const head_dim_sw = is_4bit(kv_data_type) ? sliding_window_kv_cache.size(-1) * 2
                                                 : sliding_window_kv_cache.size(-1);
-  int const head_dim_o = is_4bit(o_data_type) ? out.size(-1) * 2 : out.size(-1);
+  int const head_dim_o = fuses_dsv4_inv_rope_fp8_quant
+                             ? 512
+                             : (is_4bit(o_data_type) ? out.size(-1) * 2 : out.size(-1));
   TVM_FFI_ICHECK_EQ(head_dim_q, 512);
   TVM_FFI_ICHECK_EQ(head_dim_k, 512);
   TVM_FFI_ICHECK_EQ(head_dim_sw, 512);
   TVM_FFI_ICHECK_EQ(head_dim_o, 512);
 
+  float const* dsv4_inv_rope_cos_sin_cache_ptr = nullptr;
+  float* dsv4_output_scale_ptr = nullptr;
+  int64_t dsv4_scale_buf_m = 0;
+  if (fuses_dsv4_inv_rope_fp8_quant) {
+    auto const& cos_sin_cache = dsv4_inv_rope_cos_sin_cache.value();
+    auto const& output_scale = dsv4_output_scale.value();
+    int64_t constexpr heads_per_group = 8;
+    int64_t const num_groups = num_qo_heads / heads_per_group;
+    int64_t const group_width = heads_per_group * head_dim_o;
+
+    TVM_FFI_ICHECK_EQ(out.ndim(), 3)
+        << "RopeQuant out must have shape [sum_q, num_head_groups, group_width]";
+    TVM_FFI_ICHECK_EQ(out.size(0), sum_seq_q);
+    TVM_FFI_ICHECK_EQ(out.size(1), num_groups);
+    TVM_FFI_ICHECK_EQ(out.size(2), group_width);
+    TVM_FFI_ICHECK_EQ(out.stride(0), group_width);
+    TVM_FFI_ICHECK_EQ(out.stride(1), sum_seq_q * group_width);
+    TVM_FFI_ICHECK_EQ(out.stride(2), 1);
+
+    TVM_FFI_ICHECK_EQ(cos_sin_cache.dtype(), dl_float32)
+        << "dsv4_inv_rope_cos_sin_cache must be float32";
+    TVM_FFI_ICHECK(cos_sin_cache.ndim() == 2 && cos_sin_cache.size(1) == 64 &&
+                   cos_sin_cache.IsContiguous())
+        << "dsv4_inv_rope_cos_sin_cache must be contiguous [max_position, 64]";
+
+    TVM_FFI_ICHECK_EQ(output_scale.dtype(), dl_int32)
+        << "dsv4_output_scale must contain packed UE8M0 values in int32 storage";
+    TVM_FFI_ICHECK_EQ(output_scale.ndim(), 3)
+        << "dsv4_output_scale must have shape [sum_q, num_head_groups, 8]";
+    TVM_FFI_ICHECK_EQ(output_scale.size(0), sum_seq_q);
+    TVM_FFI_ICHECK_EQ(output_scale.size(1), num_groups);
+    TVM_FFI_ICHECK_EQ(output_scale.size(2), heads_per_group);
+    TVM_FFI_ICHECK_EQ(output_scale.stride(0), 1);
+    dsv4_scale_buf_m = output_scale.stride(2);
+    TVM_FFI_ICHECK(dsv4_scale_buf_m >= sum_seq_q && dsv4_scale_buf_m <= INT_MAX &&
+                   dsv4_scale_buf_m % 4 == 0)
+        << "dsv4_output_scale token stride must be a multiple of 4 and cover sum_q";
+    TVM_FFI_ICHECK_EQ(output_scale.stride(1), heads_per_group * dsv4_scale_buf_m);
+
+    for (auto const& tensor : {out, output_scale, cos_sin_cache}) {
+      TVM_FFI_ICHECK(tensor.device().device_type == query.device().device_type &&
+                     tensor.device().device_id == query.device().device_id)
+          << "RopeQuant outputs and cos/sin cache must be on the same device as query";
+    }
+
+    dsv4_inv_rope_cos_sin_cache_ptr = static_cast<float const*>(cos_sin_cache.data_ptr());
+    dsv4_output_scale_ptr = static_cast<float*>(output_scale.data_ptr());
+  }
+
   int const sparse_mla_top_k = sparse_indices.size(-1);
   TVM_FFI_ICHECK(sparse_mla_top_k >= kDsv4SparseMlaSlidingWindowTopK)
       << "sparse topK must include 128 sliding-window entries";
@@ -1155,7 +1227,8 @@ void trtllm_paged_attention_decode_sparse_mla_dsv4(
       /*k_sf_stride_heads=*/0, /*k_sf_stride_batch=*/0, /*v_sf_stride_heads=*/0,
       /*v_sf_stride_batch=*/0, /*is_causal=*/true, /*lse_stride_tokens=*/0,
       /*lse_stride_heads=*/0, /*bf16q_fp8kv_transform_mode=*/0, /*use_fp16_softmax=*/false,
-      /*uses_spcompress=*/false, stream);
+      /*uses_spcompress=*/false, dsv4_inv_rope_cos_sin_cache_ptr, dsv4_output_scale_ptr,
+      dsv4_scale_buf_m, stream);
 }
 
 namespace trtllm_cubin_loader {