feat: trtllm-gen FMHA features for sm107 (spcompress, fp16softmax)

csrc/trtllm_fmha_kernel_launcher.cu

@@ -121,7 +121,8 @@ void trtllm_paged_attention_launcher(
     bool uses_shared_paged_kv_idx, bool enable_block_sparse_attention, int64_t sm_count,
     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, cudaStream_t stream) {
+    int64_t lse_stride_heads, int64_t bf16q_fp8kv_transform_mode, bool use_fp16_softmax,
+    bool uses_spcompress, 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
@@ -307,6 +308,10 @@ void trtllm_paged_attention_launcher(
   runner_params.mSkipsSoftmaxWhenPossible = skips_softmax;
   runner_params.mSkipSoftmaxThresholdScaleFactor = skip_softmax_threshold_scale_factor;
 
+  // Cubin-variant selectors (FP16 softmax accumulator, sparse compression).
+  runner_params.mUseFp16Softmax = use_fp16_softmax;
+  runner_params.mUsesSpcompress = uses_spcompress;
+
   auto [foundKernels, kinfo] = fmha_runner->isSupportedWithInfo(runner_params);
   if (!foundKernels) {
     std::ostringstream err_msg;
@@ -352,7 +357,8 @@ void trtllm_paged_attention_decode(
     Optional<TensorView> value_block_scales, Optional<float> skip_softmax_threshold_scale_factor,
     Optional<bool> uses_shared_paged_kv_idx, Optional<TensorView> lse, int64_t lse_stride_tokens,
     int64_t lse_stride_heads, bool enable_block_sparse_attention,
-    Optional<TensorView> sparse_mla_top_k_lens, int64_t bf16q_fp8kv_transform_mode) {
+    Optional<TensorView> sparse_mla_top_k_lens, int64_t bf16q_fp8kv_transform_mode,
+    Optional<bool> use_fp16_softmax) {
   auto q_data_type = dl_dtype_to_tllm_data_type(query.dtype());
   auto kv_data_type = dl_dtype_to_tllm_data_type(key_cache.dtype());
   TVM_FFI_ICHECK_EQ(key_cache.ndim(), value_cache.ndim());
@@ -483,6 +489,9 @@ void trtllm_paged_attention_decode(
   bool const is_single_pool_dynamic_sparse_mla = sparse_mla_top_k_lens_ptr != nullptr &&
                                                  sparse_mla_top_k > 0 && head_dim_q == 512 &&
                                                  head_dim_o == 512 && is_shared_kv;
+  bool const use_fp16_softmax_value = use_fp16_softmax.value_or(false);
+  // Spcompress is a context-phase cubin variant; decode never selects it.
+  bool const uses_spcompress_value = false;
 
   if (enable_block_sparse_attention) {
     // Block-sparse attention uses per-KV-head page tables and sequence lengths. The kernel
@@ -523,7 +532,8 @@ void trtllm_paged_attention_decode(
       skip_softmax_threshold_scale_factor_value, skips_softmax, uses_shared_paged_kv_idx_value,
       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, stream);
+      lse_stride_tokens, lse_stride_heads, bf16q_fp8kv_transform_mode, use_fp16_softmax_value,
+      uses_spcompress_value, stream);
 }
 
 void trtllm_paged_attention_context(
@@ -536,7 +546,8 @@ void trtllm_paged_attention_context(
     bool enable_pdl, int64_t workspace_size, Optional<TensorView> attention_sinks,
     Optional<TensorView> key_block_scales, Optional<TensorView> value_block_scales,
     Optional<float> skip_softmax_threshold_scale_factor, Optional<bool> uses_shared_paged_kv_idx,
-    bool is_causal, Optional<TensorView> lse, int64_t lse_stride_tokens, int64_t lse_stride_heads) {
+    Optional<bool> use_fp16_softmax, Optional<bool> uses_spcompress, bool is_causal,
+    Optional<TensorView> lse, int64_t lse_stride_tokens, int64_t lse_stride_heads) {
   auto q_data_type = dl_dtype_to_tllm_data_type(query.dtype());
   auto kv_data_type = dl_dtype_to_tllm_data_type(key_cache.dtype());
   auto o_data_type = dl_dtype_to_tllm_data_type(out.dtype());
@@ -639,6 +650,8 @@ void trtllm_paged_attention_context(
   float const skip_softmax_threshold_scale_factor_value =
       skip_softmax_threshold_scale_factor.value_or(0.0f);
   bool const skips_softmax = skip_softmax_threshold_scale_factor_value != 0.0f;
+  bool const use_fp16_softmax_value = use_fp16_softmax.value_or(false);
+  bool const uses_spcompress_value = uses_spcompress.value_or(false);
 
   TVM_FFI_CHECK(
       is_causal || window_left == -1,
@@ -663,7 +676,8 @@ void trtllm_paged_attention_context(
       skip_softmax_threshold_scale_factor_value, skips_softmax, uses_shared_paged_kv_idx_value,
       /*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, stream);
+      lse_stride_tokens, lse_stride_heads, /*bf16q_fp8kv_transform_mode=*/0, use_fp16_softmax_value,
+      uses_spcompress_value, stream);
 }
 
 void trtllm_ragged_attention_launcher(
@@ -677,10 +691,10 @@ void trtllm_ragged_attention_launcher(
     int64_t sm_count, bool enable_pdl, bool is_causal, int64_t k_stride_keys_values,
     int64_t k_stride_heads, int64_t k_stride_batch, int64_t v_stride_keys_values,
     int64_t v_stride_heads, int64_t v_stride_batch, float skip_softmax_threshold_scale_factor,
-    bool skips_softmax, int64_t workspace_size, const float* sage_attn_sfs_q,
-    const float* sage_attn_sfs_k, const float* sage_attn_sfs_p, const float* sage_attn_sfs_v,
-    int num_elts_sage_q, int num_elts_sage_k, int num_elts_sage_p, int num_elts_sage_v,
-    int64_t lse_stride_tokens, int64_t lse_stride_heads, cudaStream_t stream) {
+    bool skips_softmax, bool use_fp16_softmax, bool uses_spcompress, int64_t workspace_size,
+    const float* sage_attn_sfs_q, const float* sage_attn_sfs_k, const float* sage_attn_sfs_p,
+    const float* sage_attn_sfs_v, int num_elts_sage_q, int num_elts_sage_k, int num_elts_sage_p,
+    int num_elts_sage_v, int64_t lse_stride_tokens, int64_t lse_stride_heads, 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
@@ -761,6 +775,8 @@ void trtllm_ragged_attention_launcher(
   runner_params.mSkipsSoftmaxWhenPossible = skips_softmax;
   runner_params.mSkipSoftmaxThresholdScaleFactor = skip_softmax_threshold_scale_factor;
 
+  runner_params.mUseFp16Softmax = use_fp16_softmax;
+  runner_params.mUsesSpcompress = uses_spcompress;
   // SageAttention scaling factors.
   runner_params.ptrSageAttnSfsQ = sage_attn_sfs_q;
   runner_params.ptrSageAttnSfsK = sage_attn_sfs_k;
@@ -785,6 +801,7 @@ void trtllm_ragged_attention(
     TensorView cum_seq_lens_kv, int64_t sm_count, bool enable_pdl, bool is_causal,
     int64_t workspace_size, Optional<TensorView> attention_sinks,
     Optional<float> skip_softmax_threshold_scale_factor, Optional<TensorView> lse,
+    Optional<bool> use_fp16_softmax, Optional<bool> uses_spcompress,
     Optional<TensorView> sage_attn_sfs_q, Optional<TensorView> sage_attn_sfs_k,
     Optional<TensorView> sage_attn_sfs_p, Optional<TensorView> sage_attn_sfs_v,
     int64_t num_elts_per_sage_attn_blk_q, int64_t num_elts_per_sage_attn_blk_k,
@@ -869,6 +886,8 @@ void trtllm_ragged_attention(
   float const skip_softmax_threshold_scale_factor_value =
       skip_softmax_threshold_scale_factor.value_or(0.0f);
   bool const skips_softmax = skip_softmax_threshold_scale_factor_value != 0.0f;
+  bool const use_fp16_softmax_value = use_fp16_softmax.value_or(false);
+  bool const uses_spcompress_value = uses_spcompress.value_or(false);
 
   trtllm_ragged_attention_launcher(
       out.data_ptr(), query.data_ptr(), key.data_ptr(), value.data_ptr(),
@@ -879,9 +898,9 @@ void trtllm_ragged_attention(
       bmm1_scale_value, bmm2_scale_value, bmm1_scale_log2_ptr, bmm2_scale_ptr, o_sf_scale,
       batch_size, window_left, sm_count, enable_pdl, is_causal, k_stride_keys_values,
       k_stride_heads, k_stride_batch, v_stride_keys_values, v_stride_heads, v_stride_batch,
-      skip_softmax_threshold_scale_factor_value, skips_softmax, workspace_size, sage_attn_sfs_q_ptr,
-      sage_attn_sfs_k_ptr, sage_attn_sfs_p_ptr, sage_attn_sfs_v_ptr,
-      static_cast<int>(num_elts_per_sage_attn_blk_q),
+      skip_softmax_threshold_scale_factor_value, skips_softmax, use_fp16_softmax_value,
+      uses_spcompress_value, workspace_size, sage_attn_sfs_q_ptr, sage_attn_sfs_k_ptr,
+      sage_attn_sfs_p_ptr, sage_attn_sfs_v_ptr, static_cast<int>(num_elts_per_sage_attn_blk_q),
       static_cast<int>(num_elts_per_sage_attn_blk_k),
       static_cast<int>(num_elts_per_sage_attn_blk_p),
       static_cast<int>(num_elts_per_sage_attn_blk_v), lse_stride_tokens, lse_stride_heads, stream);
@@ -1015,7 +1034,8 @@ void trtllm_paged_attention_decode_sparse_mla_dsv4(
       enable_pdl, workspace_size,
       /*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, stream);
+      /*lse_stride_heads=*/0, /*bf16q_fp8kv_transform_mode=*/0, /*use_fp16_softmax=*/false,
+      /*uses_spcompress=*/false, stream);
 }
 
 namespace trtllm_cubin_loader {

flashinfer/decode.py

@@ -2906,6 +2906,7 @@ def _paged_run(
             False,  # enable_block_sparse_attention
             None,  # sparse_mla_top_k_lens
             bf16q_fp8kv_transform_mode,
+            None,  # use_fp16_softmax
         )
         return out
 
@@ -3600,6 +3601,7 @@ def trtllm_batch_decode_with_kv_cache(
             enable_block_sparse_attention,
             None,  # sparse_mla_top_k_lens
             bf16q_fp8kv_transform_mode_value,
+            None,  # use_fp16_softmax
         )
 
         result_out = (

flashinfer/mla/_core.py

@@ -42,6 +42,7 @@
     MaskMode,
     _check_block_tables_shape,
     check_shape_dtype_device,
+    check_trtllm_gen_sm107_only_feature,
     determine_mla_backend,
     device_support_pdl,
     get_compute_capability,
@@ -2837,6 +2838,7 @@ def _cute_dsl_incompatibility_reason(
     kv_lora_rank: int,
     page_size: int,
     is_var_seq: bool,
+    use_fp16_softmax: Optional[bool] = None,
     cute_dsl_impl: str = "auto",
     cum_seq_lens_q: Optional[torch.Tensor] = None,
     max_q_len: Optional[int] = None,
@@ -2891,6 +2893,8 @@ def _cute_dsl_incompatibility_reason(
             "cute-dsl backend (MLA decode kernel) does not support separate KV "
             "page indices (uses_shared_paged_kv_idx=False)"
         )
+    if use_fp16_softmax:
+        return "cute-dsl backend (MLA decode kernel) does not support use_fp16_softmax"
     # LSE is supported on the monolithic path; the modular path raises a
     # clear NotImplementedError in wrappers/batch_mla.py if it gets picked
     # for an LSE request (e.g. when ``sinks`` forces the modular dispatch).
@@ -3134,6 +3138,7 @@ def __init__(
         uses_shared_paged_kv_idx: bool,
         return_lse: bool,
         lse: Optional[torch.Tensor],
+        use_fp16_softmax: Optional[bool] = None,
     ):
         self._run = get_trtllm_gen_fmha_module().trtllm_paged_attention_decode
         self.kv_cache = kv_cache
@@ -3161,6 +3166,7 @@ def __init__(
         self.uses_shared_paged_kv_idx = uses_shared_paged_kv_idx
         self.return_lse = return_lse
         self.lse = lse
+        self.use_fp16_softmax = use_fp16_softmax
 
     def __hash__(self):
         # The default `TunableRunner.__hash__` walks `self.__dict__` and falls
@@ -3292,6 +3298,7 @@ def forward(
             False,  # enable_block_sparse_attention
             sparse_mla_top_k_lens,
             0,  # bf16q_fp8kv_transform_mode
+            self.use_fp16_softmax,
         )
         return out
 
@@ -3524,6 +3531,7 @@ def trtllm_batch_decode_with_kv_cache_mla(
     cp_world: int = 1,
     cp_rank: int = 0,
     causal_seqlens_kv_global: Optional[torch.Tensor] = None,
+    use_fp16_softmax: Optional[bool] = None,
 ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
     r"""Decode MLA with TRTLLM-GEN, CuteDSL, XQA, or SM120/SM121 sparse kernels.
 
@@ -3634,7 +3642,14 @@ def trtllm_batch_decode_with_kv_cache_mla(
         Whether K and V page indices are shared as a unified index.
         True (default) uses vLLM/FlashInfer layout with a 2D page table.
         False uses TRT-LLM layout with a 3D page table ``[batch_size, 2, max_num_pages_per_seq]``.
-        False is only supported by TRTLLM-GEN.
+        False is only supported for trtllm-gen backend.
+    use_fp16_softmax : Optional[bool]
+        Select the trtllm-gen ``Fp16Softmax`` cubin variant. MLA decode is the
+        primary consumer of this flag — `Fp16Softmax` generation cubins are
+        only shipped for MLA head dims (``head_dim_qk/v ∈ {576/512, 320/256}``).
+        When ``None`` (default) or ``False`` the standard FP32-accumulator
+        softmax cubin is used. Only supported by ``backend="trtllm-gen"``;
+        passing ``True`` to other backends raises ``ValueError``.
     lse : Optional[torch.Tensor] = None
         Optional pre-allocated buffer for Log-Sum-Exp values. Supported by
         ``trtllm-gen``, ``cute-dsl``, and ``sparse`` backends. Must have
@@ -3811,6 +3826,10 @@ def trtllm_batch_decode_with_kv_cache_mla(
         causal_seqlens_kv_global=causal_seqlens_kv_global,
     )
 
+    check_trtllm_gen_sm107_only_feature(
+        use_fp16_softmax, "use_fp16_softmax", query.device
+    )
+
     if backend == "auto":
         cc = get_compute_capability(query.device)
         if cc[0] == 12 and sparse_mla_top_k > 0:
@@ -3853,6 +3872,10 @@ def trtllm_batch_decode_with_kv_cache_mla(
             raise ValueError(
                 "XQA MLA does not support separate KV page indices (uses_shared_paged_kv_idx=False)"
             )
+        if use_fp16_softmax:
+            raise ValueError(
+                "use_fp16_softmax is only supported by backend='trtllm-gen'"
+            )
         if return_lse or lse is not None:
             raise NotImplementedError(
                 "XQA MLA backend does not support return_lse/lse output"
@@ -4151,6 +4174,7 @@ def trtllm_batch_decode_with_kv_cache_mla(
             False,  # enable_block_sparse_attention
             sparse_mla_top_k_lens,
             0,  # bf16q_fp8kv_transform_mode
+            use_fp16_softmax,
         )
         return out
 
@@ -4224,6 +4248,7 @@ def trtllm_batch_decode_with_kv_cache_mla(
         kv_lora_rank,
         page_size,
         is_var_seq,
+        use_fp16_softmax=use_fp16_softmax,
         cute_dsl_impl=cute_dsl_impl,
         enable_dcp=enable_dcp,
         cp_world=cp_world,
@@ -4283,6 +4308,7 @@ def trtllm_batch_decode_with_kv_cache_mla(
                 uses_shared_paged_kv_idx=uses_shared_paged_kv_idx,
                 return_lse=return_lse,
                 lse=lse,
+                use_fp16_softmax=use_fp16_softmax,
             )
         )
     if "cute-dsl" in runner_names: