[Perf] Ship the SM70 grouped long-context route on by default

cmake/external_projects/vllm_flash_attn.cmake

@@ -113,6 +113,23 @@ if(VLLM_FLASH_ATTN_SM70 AND TARGET _vllm_fa2_C)
     "${SM70_V37_DIR}/register.cpp")
 endif()
 
+# The grouped E4M3 FP32 long-context route ships inside the same extension, so
+# the accelerated path is available without an externally built DSO.
+if(VLLM_FLASH_ATTN_SM70 AND TARGET _vllm_fa2_C)
+  set(SM70_GROUPED_LONG_DIR
+      "${CMAKE_CURRENT_LIST_DIR}/../../csrc/attention/sm70_grouped_long")
+  set(SM70_GROUPED_LONG_SRC
+      "${SM70_GROUPED_LONG_DIR}/kernel/grouped-attention.cu")
+  # Flags mirror the manifest the operator was qualified with. As with v37, the
+  # properties must be set in the target scope or SM70 silently loses them.
+  set_source_files_properties(${SM70_GROUPED_LONG_SRC}
+    TARGET_DIRECTORY _vllm_fa2_C
+    PROPERTIES COMPILE_OPTIONS
+      "-gencode=arch=compute_70,code=sm_70;-O3;-std=c++17;--use_fast_math;--expt-relaxed-constexpr;--expt-extended-lambda;-U__CUDA_NO_HALF_OPERATORS__;-U__CUDA_NO_HALF_CONVERSIONS__;-U__CUDA_NO_HALF2_OPERATORS__")
+  target_include_directories(_vllm_fa2_C PRIVATE "${SM70_GROUPED_LONG_DIR}/include")
+  target_sources(_vllm_fa2_C PRIVATE ${SM70_GROUPED_LONG_SRC})
+endif()
+
 # Restore the install prefix after FA's install rules
 install(CODE "set(CMAKE_INSTALL_PREFIX \"\${OLD_CMAKE_INSTALL_PREFIX}\")" ALL_COMPONENTS)
 install(CODE "set(CMAKE_INSTALL_LOCAL_ONLY TRUE)" ALL_COMPONENTS)

csrc/attention/sm70_grouped_long/LICENSE

@@ -0,0 +1,28 @@
+BSD 3-Clause License
+
+Copyright (c) 2025, D.Skryabin
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+1. Redistributions of source code must retain the above copyright notice, this
+   list of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright notice,
+   this list of conditions and the following disclaimer in the documentation
+   and/or other materials provided with the distribution.
+
+3. Neither the name of the copyright holder nor the names of its contributors
+   may be used to endorse or promote products derived from this software without
+   specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
+ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
+WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR
+ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
+INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS
+OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION HOWEVER CAUSED AND ON ANY
+THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT INCLUDING
+NEGLIGENCE OR OTHERWISE ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN
+IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

csrc/attention/sm70_grouped_long/include/fused_mha.h

@@ -0,0 +1,173 @@
+#ifndef FUSED_MHA_H
+#define FUSED_MHA_H
+
+#include <cuda_runtime.h>
+#include <stdexcept>
+#include <string>
+#include <torch/extension.h>
+#include <ATen/ATen.h>
+
+std::vector<at::Tensor> flash_attention_forward(
+    at::Tensor& q, const at::Tensor& k, const at::Tensor& v,
+    std::optional<at::Tensor>& out_, std::optional<at::Tensor>& alibi_slopes_,
+    const float p_dropout, const float softmax_scale, bool is_causal,
+    int window_size_left, int window_size_right, const float softcap,
+    const bool return_softmax, std::optional<at::Generator> gen_);
+
+at::Tensor flash_attention_qk_scores(const at::Tensor& q, const at::Tensor& k,
+                                     const float softmax_scale,
+                                     const bool is_causal);
+
+at::Tensor flash_attention_decode_paged(
+    const at::Tensor& q, const at::Tensor& k_cache, const at::Tensor& v_cache,
+    std::optional<at::Tensor>& out_, const at::Tensor& block_table,
+    const at::Tensor& seq_lens, at::Tensor& tmp_out, at::Tensor& max_logits,
+    at::Tensor& exp_sums, const at::Tensor& active_num_partitions,
+    const float softmax_scale, const int partition_size,
+    const int launch_num_partitions, const std::string& kv_cache_dtype,
+    const float k_scale, const float v_scale, const int window_size_left,
+    const int window_size_right, const std::optional<at::Tensor>& anchor_lens,
+    const int64_t anchored_window);
+
+at::Tensor flash_attention_decode_paged_xqa(
+    const at::Tensor& q, const at::Tensor& k_cache, const at::Tensor& v_cache,
+    std::optional<at::Tensor>& out_, const at::Tensor& block_table,
+    const at::Tensor& seq_lens, at::Tensor& tmp_out, at::Tensor& max_logits,
+    at::Tensor& exp_sums, const at::Tensor& active_num_partitions,
+    const float softmax_scale, const int partition_size,
+    const int launch_num_partitions, const std::string& kv_cache_dtype,
+    const float k_scale, const float v_scale, const int window_size_left,
+    const int window_size_right, const int batch_context_max_seq_len);
+
+at::Tensor flash_attention_decode_paged_xqa_staged(
+    const at::Tensor& q, const at::Tensor& k_cache, const at::Tensor& v_cache,
+    std::optional<at::Tensor>& out_, const at::Tensor& block_table,
+    const at::Tensor& seq_lens, at::Tensor& tmp_out, at::Tensor& max_logits,
+    at::Tensor& exp_sums, at::Tensor& online_rescales,
+    const at::Tensor& active_num_partitions, const float softmax_scale,
+    const int partition_size, const int launch_num_partitions,
+    const std::string& kv_cache_dtype, const float k_scale, const float v_scale,
+    const int window_size_left, const int window_size_right);
+
+at::Tensor flash_attention_grouped_verify_paged(
+    const at::Tensor& q, const at::Tensor& k_cache, const at::Tensor& v_cache,
+    std::optional<at::Tensor>& out_, const at::Tensor& block_table,
+    const at::Tensor& seq_lens, at::Tensor& partial_out,
+    at::Tensor& partial_lse, const float softmax_scale,
+    const std::string& kv_cache_dtype, const float k_scale, const float v_scale,
+    const bool one_pass);
+
+int64_t flash_attention_grouped_verify_max_query_tokens();
+
+int64_t flash_attention_grouped_verify_request_major_abi_version();
+
+at::Tensor flash_attention_grouped_e4m3_fp32_paged(
+    const at::Tensor& q, const at::Tensor& k, const at::Tensor& v,
+    at::Tensor& out, const at::Tensor& block_table,
+    const at::Tensor& row_lengths, at::Tensor& partial, at::Tensor& lse,
+    float scale, float k_scale, float v_scale);
+
+int64_t flash_attention_grouped_e4m3_fp32_precision_version();
+
+int64_t flash_attention_grouped_sparse_page4_abi_version();
+
+at::Tensor flash_attention_grouped_sparse_page4(
+    const at::Tensor& q, const at::Tensor& k_cache, const at::Tensor& v_cache,
+    std::optional<at::Tensor>& out_, const at::Tensor& block_table,
+    const at::Tensor& token_masks, const at::Tensor& seq_lens, at::Tensor& lse,
+    const float softmax_scale, const std::string& kv_cache_dtype,
+    const float k_scale, const float v_scale);
+
+at::Tensor flash_attention_grouped_sparse_page4_plan(
+    const at::Tensor& logical_indices, const at::Tensor& block_table,
+    const at::Tensor& token_to_req, const at::Tensor& query_positions,
+    const at::Tensor& sequence_lengths, at::Tensor& output_blocks,
+    at::Tensor& output_masks, at::Tensor& output_seq_lens, const int page_size,
+    const int physical_page_stride, const int num_cache_blocks);
+
+at::Tensor flash_attention_decode_paged_wmma(
+    const at::Tensor& q, const at::Tensor& k_cache, const at::Tensor& v_cache,
+    std::optional<at::Tensor>& out_, const at::Tensor& block_table,
+    const at::Tensor& seq_lens, const float softmax_scale,
+    const std::string& kv_cache_dtype, const float k_scale,
+    const float v_scale);
+
+at::Tensor flash_attention_decode_qk_scores(
+    const at::Tensor& q, const at::Tensor& k_cache,
+    const at::Tensor& block_table, const at::Tensor& seq_lens,
+    const float softmax_scale, const int partition_size,
+    const std::string& kv_cache_dtype, const float k_scale);
+
+at::Tensor flash_attention_turboquant_decode_paged(
+    const at::Tensor& q_rot, const at::Tensor& kv_cache,
+    std::optional<at::Tensor>& out_, const at::Tensor& block_table,
+    const at::Tensor& seq_lens, at::Tensor& tmp_out, at::Tensor& max_logits,
+    at::Tensor& exp_sums, const at::Tensor& centroids,
+    const float softmax_scale, const int partition_size, const int mse_bits,
+    const int value_quant_bits, const bool norm_correction);
+
+at::Tensor flash_attention_prefill_paged(
+    const at::Tensor& q, const at::Tensor& k_cache, const at::Tensor& v_cache,
+    std::optional<at::Tensor>& out_, const at::Tensor& block_table,
+    const at::Tensor& seq_lens, const float softmax_scale,
+    const std::string& kv_cache_dtype, const float k_scale, const float v_scale,
+    const bool is_causal, const int window_size_left,
+    const int window_size_right, const std::optional<at::Tensor>& anchor_lens,
+    const int64_t anchored_window);
+
+std::vector<at::Tensor>
+flash_attention_prefill_paged_d256_bm32_allp_pair_scratch(
+    const at::Tensor& q, const at::Tensor& k_cache, const at::Tensor& v_cache,
+    std::optional<at::Tensor>& out_, std::optional<at::Tensor>& softmax_lse_,
+    const at::Tensor& block_table, const at::Tensor& seq_lens,
+    const float softmax_scale);
+
+std::vector<at::Tensor>
+flash_attention_prefill_paged_d256_bm32_allp_pair_scratch_splitkv3(
+    const at::Tensor& q, const at::Tensor& k_cache, const at::Tensor& v_cache,
+    std::optional<at::Tensor>& out_, std::optional<at::Tensor>& softmax_lse_,
+    at::Tensor& split_tmp_out, at::Tensor& split_tmp_row_max,
+    at::Tensor& split_tmp_row_sum, const at::Tensor& block_table,
+    const int64_t actual_n, const float softmax_scale);
+
+at::Tensor flash_attention_prefill_paged_bfla(
+    const at::Tensor& q, const at::Tensor& k_cache, const at::Tensor& v_cache,
+    std::optional<at::Tensor>& out_, const at::Tensor& block_table,
+    const at::Tensor& seq_lens, const at::Tensor& bfla_block_mask,
+    const int bfla_mask_block_n, const float softmax_scale,
+    const std::string& kv_cache_dtype, const float k_scale, const float v_scale,
+    const bool is_causal, const int window_size_left,
+    const int window_size_right);
+
+at::Tensor flash_attention_prefill_paged_splitkv(
+    const at::Tensor& q, const at::Tensor& k_cache, const at::Tensor& v_cache,
+    std::optional<at::Tensor>& out_, const at::Tensor& block_table,
+    const at::Tensor& seq_lens, const float softmax_scale,
+    const std::string& kv_cache_dtype, const float k_scale, const float v_scale,
+    const bool is_causal, const int window_size_left,
+    const int window_size_right, const int split_kv_tokens,
+    const int max_seq_len_hint);
+
+void flash_attention_fp8_e5m2_paged_kv_to_fp16(
+    const at::Tensor& key_cache, const at::Tensor& value_cache,
+    const at::Tensor& block_table, const at::Tensor& seq_lens,
+    at::Tensor& key_out, at::Tensor& value_out, const float key_scale,
+    const float value_scale);
+
+void flash_attention_fp8_e4m3_paged_kv_to_fp16(
+    const at::Tensor& key_cache, const at::Tensor& value_cache,
+    const at::Tensor& block_table, const at::Tensor& seq_lens,
+    at::Tensor& key_out, at::Tensor& value_out, const float key_scale,
+    const float value_scale);
+
+std::vector<at::Tensor> flash_attention_backward(
+    const at::Tensor& dout, const at::Tensor& q, const at::Tensor& k,
+    const at::Tensor& v, const at::Tensor& out, const at::Tensor& softmax_lse,
+    std::optional<at::Tensor>& dq_, std::optional<at::Tensor>& dk_,
+    std::optional<at::Tensor>& dv_, std::optional<at::Tensor>& alibi_slopes_,
+    const float p_dropout, const float softmax_scale, const bool is_causal,
+    int window_size_left, int window_size_right, const float softcap,
+    const bool deterministic, std::optional<at::Generator> gen_,
+    std::optional<at::Tensor>& rng_state);
+
+#endif