Sync DeepSeek upstream through V4.1 kernels (#221)

.gitignore

@@ -10,3 +10,4 @@ compile_commands.json
 .cache
 /dev
 /.clangd
+/.venv/

csrc/api/api.cpp

@@ -2,16 +2,17 @@
 
 #include <torch/csrc/stable/library.h>
 
-#include "sparse_fwd.h"
-#include "sparse_decode.h"
-#include "dense_decode.h"
-#include "dense_fwd.h"
+#include "interfaces.h"
 
 STABLE_TORCH_LIBRARY(_flashmla_C, m) {
     m.def("sparse_decode_fwd(Tensor q, Tensor kv, Tensor indices, Tensor? topk_length, Tensor? attn_sink, Tensor(a)? tile_scheduler_metadata, Tensor(b)? num_splits, Tensor? extra_kv, Tensor? extra_indices, Tensor? extra_topk_length, int d_v, float sm_scale, Tensor(c!)? out_) -> (Tensor(c!), Tensor, Tensor(a)?, Tensor(b)?)");
     m.def("dense_decode_fwd(Tensor q, Tensor kcache, int head_size_v, Tensor seqlens_k, Tensor block_table, float softmax_scale, bool is_causal, Tensor(a)? tile_scheduler_metadata, Tensor(b)? num_splits, Tensor(c!)? out_) -> (Tensor(c!), Tensor, Tensor(a)?, Tensor(b)?)");
     m.def("sparse_prefill_fwd(Tensor q, Tensor kv, Tensor indices, float sm_scale, int d_v, Tensor? attn_sink, Tensor? topk_length, Tensor(a!)? out_) -> Tensor[]");
     m.def("dense_prefill_fwd(Tensor workspace_buffer, Tensor q, Tensor k, Tensor v, Tensor cumulative_seqlen_q, Tensor cumulative_seqlen_kv, Tensor(a!) o, Tensor(b!) lse, int mask_mode_code, float softmax_scale, int max_seqlen_q, int max_seqlen_kv, bool is_varlen) -> ()");
+    m.def("fused_norm_rope_attn_rope_cast_fwd(Tensor q, Tensor kv, Tensor indices, float sm_scale, int d_v, Tensor? attn_sink, Tensor? topk_length, bool enable_q_norm, float rms_norm_eps, Tensor token_positions, bool is_rope_neox_style, int rope_dim, Tensor cos_sin_cache, int n_wv_group, int num_per_channels, bool use_tma_aligned_col_major_sf, bool round_sf, bool use_packed_ue8m0) -> Tensor[]");
+    m.def("fused_norm_rope_attn_rope_cast_decode(Tensor q, Tensor kv, Tensor indices, float sm_scale, int d_v, Tensor? attn_sink, Tensor? topk_length, Tensor? extra_kv, Tensor? extra_indices, Tensor? extra_topk_length, bool enable_q_norm, float rms_norm_eps, Tensor token_positions, bool is_rope_neox_style, int rope_dim, Tensor cos_sin_cache, int n_wv_group, int num_per_channels, bool use_tma_aligned_col_major_sf, bool round_sf, bool use_packed_ue8m0) -> Tensor[]");
+    m.def("permute_q_b_proj(Tensor q_b_proj, Tensor scale_factors, int h_q, int d_q) -> Tensor[]");
+    m.def("permute_wv_proj(Tensor wv_proj, Tensor scale_factors, int wv_group_size, int d_o) -> Tensor[]");
 #ifdef FLASH_MLA_ENABLE_DENSE_BWD
     // Dense prefill backward is only registered when its kernel is compiled
     // (standalone setup.py). vLLM's integrated build is inference-only and does
@@ -25,6 +26,10 @@ STABLE_TORCH_LIBRARY_IMPL(_flashmla_C, CUDA, m) {
     m.impl("dense_decode_fwd", TORCH_BOX(&dense_attn_decode_interface));
     m.impl("sparse_prefill_fwd", TORCH_BOX(&sparse_attn_prefill_interface));
     m.impl("dense_prefill_fwd", TORCH_BOX(&FMHACutlassSM100FwdRun));
+    m.impl("fused_norm_rope_attn_rope_cast_fwd", TORCH_BOX(&fused_norm_rope_attn_rope_cast_fwd));
+    m.impl("fused_norm_rope_attn_rope_cast_decode", TORCH_BOX(&fused_norm_rope_attn_rope_cast_decode));
+    m.impl("permute_q_b_proj", TORCH_BOX(&permute_q_b_proj));
+    m.impl("permute_wv_proj", TORCH_BOX(&permute_wv_proj));
 #ifdef FLASH_MLA_ENABLE_DENSE_BWD
     m.impl("dense_prefill_bwd", TORCH_BOX(&FMHACutlassSM100BwdRun));
 #endif

csrc/api/common.h

@@ -21,6 +21,8 @@
 
 #include <cutlass/bfloat16.h>
 
+#include "kernels/kv_cache_format.h"
+
 using torch::stable::Tensor;
 using torch::headeronly::ScalarType;
 
@@ -103,26 +105,14 @@ inline int int64_stride_to_int(int64_t orig_stride) {
     if (MODEL_TYPE == ModelType::V32) { \
         static constexpr ModelType CONSTEXPR_NAME = ModelType::V32; \
         return __VA_ARGS__(); \
-    } else if (MODEL_TYPE == ModelType::MODEL1) { \
-        static constexpr ModelType CONSTEXPR_NAME = ModelType::MODEL1; \
+    } else if (MODEL_TYPE == ModelType::V4) { \
+        static constexpr ModelType CONSTEXPR_NAME = ModelType::V4; \
         return __VA_ARGS__(); \
     } else { \
         STD_TORCH_CHECK(false, "Unsupported model type: ", (int)MODEL_TYPE); \
     } \
 } ();
 
-// Like DISPATCH_MODEL_TYPE, but also covers the NVFP4 format (SM100-only kernel).
-// Kept separate so that SM90 kernel templates are never instantiated for NVFP4.
-#define DISPATCH_MODEL_TYPE_SM100(MODEL_TYPE, CONSTEXPR_NAME, ...) \
-[&] () { \
-    if (MODEL_TYPE == ModelType::V32_NVFP4_FP8ROPE) { \
-        static constexpr ModelType CONSTEXPR_NAME = ModelType::V32_NVFP4_FP8ROPE; \
-        return __VA_ARGS__(); \
-    } else { \
-        return DISPATCH_MODEL_TYPE(MODEL_TYPE, CONSTEXPR_NAME, __VA_ARGS__); \
-    } \
-} ();
-
 // The following code is adapted from https://ykiko.me/en/articles/680412313/, which converts enum values to string names.
 template<auto value>
 constexpr auto get_static_enum_name(){
@@ -165,6 +155,42 @@ static constexpr std::string get_dynamic_enum_name(T value){
     return (std::string)names[static_cast<std::size_t>(value)];
 }
 
+// =============================================
+// Paged quantized KV cache formats (decoding)
+// =============================================
+
+// V3.2 geometry has either the original 656-byte fp8/bf16 record or the
+// SM100-only 352-byte NVFP4-NoPE/fp8-RoPE record.
+inline ModelType detect_kv_cache_format_for_headdim_576(int bytes_per_token) {
+    for (ModelType mt : {ModelType::V32, ModelType::V32_NVFP4_FP8ROPE}) {
+        if (bytes_per_token == kv_cache_bytes_per_token(mt)) {
+            return mt;
+        }
+    }
+    STD_TORCH_CHECK(false, "Unsupported bytes_per_token for d_qk=576: ", bytes_per_token, ". Expected ",
+        kv_cache_bytes_per_token(ModelType::V32), " (V3.2 fp8) or ",
+        kv_cache_bytes_per_token(ModelType::V32_NVFP4_FP8ROPE), " (V3.2 NVFP4 NoPE + fp8 RoPE)");
+}
+
+// The format of a paged quantized KV cache with d_qk = 512 (V4 / V4.1 / V4.1 fp4), detected by bytes_per_token (kv.size(3))
+inline ModelType detect_kv_cache_format_for_headdim_512(int bytes_per_token) {
+    for (ModelType mt : {ModelType::V4, ModelType::V41, ModelType::V41_FP4}) {
+        if (bytes_per_token == kv_cache_bytes_per_token(mt)) {
+            return mt;
+        }
+    }
+    STD_TORCH_CHECK(false, "Unsupported bytes_per_token for d_qk=512: ", bytes_per_token, ". Expected ",
+        kv_cache_bytes_per_token(ModelType::V4), " (V4), ", kv_cache_bytes_per_token(ModelType::V41), " (V4.1) or ",
+        kv_cache_bytes_per_token(ModelType::V41_FP4), " (V4.1 fp4)");
+}
+
+// Dispatches the runtime (kv, extra_kv) format pair
+template<typename... Pairs, typename Fn>
+inline void dispatch_kv_formats(KVFormatPairs<Pairs...>, ModelType kv, ModelType extra_kv, Fn &&fn) {
+    bool matched = ((kv == Pairs::kv && extra_kv == Pairs::extra_kv ? (fn.template operator()<Pairs::kv, Pairs::extra_kv>(), true) : false) || ...);
+    STD_TORCH_CHECK(matched, "Unsupported KV cache formats for this implementation: kv ", get_dynamic_enum_name(kv), ", extra_kv ", get_dynamic_enum_name(extra_kv));
+}
+
 // A shortcut macro to declare supported features in an implementation class.
 #define DECLARE_SUPPORTED_FEATURES(...) \
 protected: \

csrc/api/dense_decode.h → csrc/api/dense_decode.cpp

@@ -1,16 +1,14 @@
-#pragma once
-
 #include <cutlass/half.h>
 #include <cutlass/fast_math.h>
 
 #include "common.h"
-#include "params.h"
+#include "kernels/params.h"
 
-#include "sm90/decode/dense/splitkv_mla.h"
-#include "smxx/decode/get_decoding_sched_meta/get_decoding_sched_meta.h"
-#include "smxx/decode/combine/combine.h"
+#include "kernels/sm90/decode/dense/splitkv_mla.h"
+#include "kernels/smxx/decode/get_decoding_sched_meta/get_decoding_sched_meta.h"
+#include "kernels/smxx/decode/combine/combine.h"
 
-static std::tuple<Tensor, Tensor, std::optional<Tensor>, std::optional<Tensor>>
+std::tuple<Tensor, Tensor, std::optional<Tensor>, std::optional<Tensor>>
 dense_attn_decode_interface(
     Tensor q,                                    // batch_size x seqlen_q x num_heads x head_size
     const Tensor &kcache,                        // num_blocks x page_block_size x num_heads_k x head_size (when is_fp8 is False) or num_blocks x num_heads_k x (page_block_size*656) (when is_fp8 is True)
@@ -58,7 +56,7 @@ dense_attn_decode_interface(
     const int num_heads_q = q.size(2);
     const int head_size_k = q.size(3);
     STD_TORCH_CHECK(head_size_k == 576 || head_size_k == 512, "Only head_size_k == 576 or 512 is supported");
-    STD_TORCH_CHECK(head_size_v == 512, "Only head_size_v == 512 is supported");
+    STD_TORCH_CHECK(head_size_v == 512, "Only head_size_v == 576 is supported");
 
     const int max_num_blocks_per_seq = block_table.size(1);
     const int num_blocks = kcache.size(0);
@@ -184,12 +182,12 @@ dense_attn_decode_interface(
     params.stream = get_current_cuda_stream(q);
 
     if (q_dtype == ScalarType::BFloat16) {
-        sm90::run_flash_splitkv_mla_kernel<cutlass::bfloat16_t>(params);
+        sm90::decode::dense::run_flash_splitkv_mla_kernel<cutlass::bfloat16_t>(params);
     } else if (q_dtype == ScalarType::Half) {
 #ifdef FLASH_MLA_DISABLE_FP16
         STD_TORCH_CHECK(false, "FlashMLA is compiled with -DFLASH_MLA_DISABLE_FP16. Please remove this flag from your environment and re-compile FlashMLA.");
 #else
-        sm90::run_flash_splitkv_mla_kernel<cutlass::half_t>(params);
+        sm90::decode::dense::run_flash_splitkv_mla_kernel<cutlass::half_t>(params);
 #endif
     } else {
         STD_TORCH_CHECK(false, "Unsupported dtype for dense MLA on SM90");