Add kernels for DeepSeek v4.1

README.md

@@ -18,6 +18,7 @@ FlashMLA is DeepSeek's library of optimized attention kernels, powering the [Dee
 
 ## News
 
+- **2026.09.10 Release of DeepSeek v4.1's Attention Kernels**: We've released attention kernels for [DeepSeek-V4.1](https://huggingface.co/deepseek-ai/DeepSeek-V4.1-Flash), including both prefill and decoding (with FP8 or FP4 KV cache). We've also released a [fused-norm-rope-attn-rope-cast kernel](#fused-norm--rope--attn--rope--cast-kernel) which fuses Q-norm (only used in V4, not V4.1), Q-RoPE, core attention, O-RoPE (conjugate), and cast-to-fp8, while retaining the same performance.
 - **2025.09.29 Release of Sparse Attention Kernels**: With the launch of [DeepSeek-V3.2](https://github.com/deepseek-ai/DeepSeek-V3.2-Exp), we are releasing the corresponding token-level sparse attention kernels. These kernels power the model's DeepSeek Sparse Attention (DSA) and achieve up to 640 TFlops during prefilling and 410 TFlops during decoding. We also release a deep-dive blog for our new FP8 sparse decoding kernel. Check it out [here](docs/20250929-hopper-fp8-sparse-deep-dive.md).
 - **2025.08.01 Kernels for MHA on SM100**: Thanks to [NVIDIA's PR](https://github.com/deepseek-ai/FlashMLA/pull/76) for MHA forward / backward kernels on SM100!
 - **2025.04.22 Deep-Dive Blog**: We'd love to share the technical details behind the new FlashMLA kernel! Check out our deep-dive write-up [here](docs/20250422-new-kernel-deep-dive.md).
@@ -32,7 +33,7 @@ python tests/test_flash_mla_dense_decoding.py
 python tests/test_flash_mla_sparse_decoding.py
 ```
 
-The dense MLA decoding kernel achieves up to 3000 GB/s in memory-bound configuration and 660 TFLOPS in computation-bound configuration on H800 SXM5 with CUDA 12.8. The token-level sparse MLA decoding kernel (which uses an FP8 KV cache while performing the matrix multiplication in bfloat16) achieves 410 TFLOPS in compute-bound configuration on H800 SXM5 with CUDA 12.8, and achieves up to 350 TFlops on B200 (which is not really optimized yet).
+The dense MLA decoding kernel achieves up to 3000 GB/s in memory-bound configuration and 660 TFLOPS in computation-bound configuration on H800 SXM5 with CUDA 12.8. The token-level sparse MLA decoding kernel (which uses an FP8 KV cache while performing the matrix multiplication in bfloat16) achieves 410 TFLOPS in compute-bound configuration on H800 SXM5 with CUDA 12.8, and achieves up to 700 TFlops on B200.
 
 #### Test & benchmark MHA prefill (Dense):
 
@@ -50,6 +51,17 @@ python tests/test_flash_mla_sparse_prefill.py
 
 It achieves up to 640 TFlops in forward computation on H800 SXM5 with CUDA 12.8, and achieves up to 1450 TFlops on B200, CUDA 12.9.
 
+#### Test & benchmark the fused norm RoPE attn RoPE cast kernel (Sparse):
+
+```bash
+python tests/test_fused_norm_rope_attn_rope_cast.py
+```
+
+[TileLang](https://github.com/tile-ai/tilelang), [Tile-Kernels](https://github.com/deepseek-ai/TileKernels/), and [DeepGEMM](https://github.com/deepseek-ai/DeepGEMM) are required for running this test script.
+
+This kernel fuses Q-norm (only used in V4, not V4.1), Q-RoPE, core attention, O-RoPE (conjugate) and cast-to-FP8 into a single kernel, saving times for those small kernels. Although it fuses many small operations, this kernel still keeps the same or even slightly higher TFlops (at the cost of having to permute the Q_b and Wv weights in advance). It achieves up to 1430 TFlops during prefill and 670 TFlops during decoding on B200.
+
+
 ## Requirements
 
 - SM90 / SM100 (See the support matrix below)
@@ -58,16 +70,17 @@ It achieves up to 640 TFlops in forward computation on H800 SXM5 with CUDA 12.8,
 
 Support matrix:
 
-| Kernel | GPU Architecture | MLA Mode [2] | KVCache Format |
+| Kernel | GPU Architecture | MLA Mode [1] | Supported Models |
 | :---: | :---: | :---: | :---: |
-| Dense Decoding | SM90 | MQA | BF16 |
-| Sparse Decoding | SM90 & SM100 | MQA | FP8 [1] |
-| Dense Prefill | SM100 | MHA |  |
-| Sparse Prefill | SM90 & SM100 | MQA |  |
+| Dense Decoding | SM90 | MQA | DeepSeek V3 / V3.1 |
+| Sparse Decoding | SM90 & SM100 | MQA | DeepSeek V3.2 / V4 / V4.1 [2] |
+| Dense Prefill | SM100 | MHA | DeepSeek V3 / V3.1 / V3.2 |
+| Sparse Prefill | SM90 & SM100 | MQA | DeepSeek V3.2 / V4 / V4.1 |
+| Fused Norm RoPE Attn RoPE Cast | SM100 | MQA | DeepSeek V4 / V4.1 |
 
-[1]: For more details on using FP8 KV cache, see documents below.
+[1]: Here "MLA Mode" refers to the mode used for MLA calculation. MQA stands for Multi-Query Attention mode (i.e. `head_dim_k` = 576 (for DeepSeek V3/V3.1/V3.2) or 512 (for DeepSeek V4/V4.1) with `head_dim_v` = 512), while MHA stands for Multi-Head Attention mode (i.e. `head_dim_k` = 192 / 128 with `head_dim_v` = 128). For a detailed explanation of these modes, please refer to the appendix of [DeepSeek V3.2's Paper](https://github.com/deepseek-ai/DeepSeek-V3.2-Exp).
 
-[2]: Here "MLA Mode" refers to the mode used for MLA calculation. MQA stands for Multi-Query Attention mode (i.e. `head_dim_k` =  576 with `head_dim_v` = 512), while MHA stands for Multi-Head Attention mode (i.e. `head_dim_k` = 192 / 128 with `head_dim_v` = 128). For a detailed explanation of these modes, please refer to the appendix of [DeepSeek V3.2's Paper](https://github.com/deepseek-ai/DeepSeek-V3.2-Exp).
+[2] Sparse Decoding for DeepSeek V4.1 is only available on SM100
 
 ## Installation
 
@@ -113,13 +126,18 @@ Where
 - `h_q` is the number of query heads.
 
 **FP8 KV Cache:**
-If `is_fp8_kvcache` is set to `True`, the kernel reads the KV cache in the "FP8 with scale" format (described below). It dequantizes the cache to bfloat16 and performs attention computation in bfloat16. The output is also in bfloat16.
+If `is_fp8_kvcache` is set to `True`, the kernel reads the KV cache in the "FP8 with scale" format (described below). It dequantizes the cache to bfloat16 and performs attention computation in bfloat16. The output is also in bfloat16. In this repository, `is_fp8_kvcache=True` is only supported together with `indices` (i.e. sparse attention); the dense decoding kernel reads bf16 / fp16 KV caches. 
 
-In the "FP8 with scale" format, each token's KV cache is 656 Bytes, structured as:
+In the "FP8 with scale" format for DeepSeek V3.2 (`head_dim` = 576, sparse attention), a page block is `page_block_size` token-major rows of 656 Bytes each:
 -   **First 512 bytes:** The "quantized NoPE" part, containing 512 `float8_e4m3` values.
 -   **Next 16 bytes:** Scale factors, containing 4 `float32` values. The first `float32` is the scale for the first 128 `float8_e4m3` values, the second for the next 128, and so on.
 -   **Last 128 bytes:** The "RoPE" part, containing 64 `bfloat16` values. This part is not quantized for accuracy.
 
+For DeepSeek V4 / V4.1 (`head_dim` = 512), the format is detected from the last dimension of `k_cache` (i.e. the bytes per token): 584 (V4), 528 (V4.1) or 288 (V4.1 fp4). In all three, a page block stores `page_block_size` data rows first and `page_block_size` scale rows afterwards:
+-   **V4**: 584 Bytes per token. The data row is 448 Bytes of quantized NoPE (`float8_e4m3`) followed by 128 Bytes, i.e. the 64 `bfloat16` RoPE values (not quantized). The scale row is 8 Bytes, of which the first 7 are `float8_e8m0` scales and the 8th byte is padding; each scale covers 64 consecutive `float8_e4m3` values of the NoPE part.
+-   **V4.1**: 528 Bytes per token. The data row is 512 Bytes of `float8_e4m3`, i.e. the 64 RoPE dimensions are quantized as well and there is no `bfloat16` part. The scale row is 16 Bytes of `float8_e8m0`, each scale covering 32 consecutive `float8_e4m3` values.
+-   **V4.1 fp4**: 288 Bytes per token. The data row is 256 Bytes containing 512 `e2m1` values, 2 values per byte (the even-indexed one in the low nibble). The scale row is 32 Bytes of `float8_e4m3`, each scale covering 16 consecutive `e2m1` values. This format is only valid for `extra_k_cache`, and only when `k_cache` is in the V4.1 format; otherwise `extra_k_cache` must have the same format as `k_cache`. In pratice we expect the sliding window (SWA) kv cache to be in FP8 and the compress attention (CA) kv cache to be in FP4.
+
 See `tests/quant.py` for quantization and dequantization details.
 
 **Sparse Attention (`indices` tensor):**
@@ -134,7 +152,7 @@ The kernel returns `(out, lse)`, where:
 -   `out` is the attention result.
 -   `lse` is the log-sum-exp value of the attention scores for each query head.
 
-See `tests/test_flash_mla_decoding.py` for a complete example.
+See `tests/test_flash_mla_dense_decoding.py` and `tests/test_flash_mla_sparse_decoding.py` for complete examples.
 
 ### Sparse MLA Prefill
 
@@ -169,7 +187,7 @@ out = S @ focused_kv  # [s_q, h_q, d_qk]
 return (out, max_logits, lse)
 ```
 
-See `tests/test_flash_mla_prefill.py` for a complete example.
+See `tests/test_flash_mla_sparse_prefill.py` for a complete example.
 
 ### Dense MHA Prefill
 
@@ -180,6 +198,114 @@ This kernel implements the standard dense Multi-Head Attention (MHA) forward and
 
 The usage is similar to the `flash_attn` package. See `tests/test_fmha_sm100.py` for a complete example.
 
+### Fused norm + RoPE + attn + RoPE + cast kernel
+
+In the DeepSeek-V4.1 release, we also provide a fused kernel that combines Q-norm (only used in V4, not in V4.1), Q-RoPE, core attention, O-RoPE (conjugate) and the cast to FP8 into a single kernel. It removes the extra time spent on these small kernels while keeping the same or even slightly higher TFlops, at the cost of having to permute the Q_b and Wv weights in advance.
+
+In DeepSeek-V4.1 attention, Q (`[hidden_size]`) is first projected to `[q_lora_rank]` (the Q_a projection) and then to `[num_attention_heads, head_dim]` (the Q_b projection). After core attention, the output (`[num_attention_heads, head_dim]`) is reshaped to `[o_groups, num_attention_heads // o_groups * head_dim]`, and each of its rows is projected to `[o_lora_rank]` (the Wv projection), giving an `[o_groups, o_lora_rank]` matrix. That matrix is reshaped to `[o_groups * o_lora_rank]` and finally projected to `[hidden_size]` (the Wo projection). This kernel requires the Q_b and Wv weights to be permuted.
+
+To permute the Q_b weight:
+
+```python
+import torch
+import tile_kernels
+from flash_mla import fused_norm_rope_attn_rope_cast
+
+h_q, d_q = 64, 512          # Q heads and Q head dimension
+q_lora_rank = 1536
+scale_gran = 128
+
+# q_b_proj: [h_q * d_q, q_lora_rank], bfloat16
+q_b_proj = torch.randn((h_q * d_q, q_lora_rank), dtype=torch.bfloat16, device='cuda')
+
+# Quantize the weight to FP8 with per-token scale factors, in DeepGEMM's layout
+q_b_proj_fp8, q_b_sf = tile_kernels.quant.per_token_cast(
+    q_b_proj, 'e4m3', scale_gran,
+    use_tma_aligned_col_major_sf=True, round_sf=True, use_packed_ue8m0=True,
+)
+
+# Permute the weight and its scale factors into the layout required by the fused kernel
+q_b_proj_fp8, q_b_sf = fused_norm_rope_attn_rope_cast.permute_q_b_proj(
+    (q_b_proj_fp8, q_b_sf), h_q, d_q,
+)
+```
+
+To permute the Wv weight:
+
+```python
+import deep_gemm
+import torch
+import tile_kernels
+from flash_mla import fused_norm_rope_attn_rope_cast
+
+n_wv_group, wv_group_size, d_o = 8, 8, 512      # n_wv_group * wv_group_size == h_q
+wv_proj_out_dim = 512                           # o_lora_rank
+scale_gran = 32
+
+# wv_proj: [n_wv_group, wv_proj_out_dim, wv_group_size * d_o], bfloat16
+wv_proj = torch.randn((n_wv_group * wv_proj_out_dim, wv_group_size * d_o),
+                      dtype=torch.bfloat16, device='cuda')
+
+# Quantize the weight to FP8, and put its scale factors into the layout that DeepGEMM's einsum expects
+wv_proj_fp8, wv_sf = tile_kernels.quant.per_token_cast(
+    wv_proj, 'e4m3', scale_gran,
+    use_tma_aligned_col_major_sf=False, round_sf=True, use_packed_ue8m0=False,
+)
+wv_sf = deep_gemm.transform_sf_into_required_layout(
+    wv_sf.view(n_wv_group, wv_proj_out_dim, wv_group_size * d_o // scale_gran),
+    wv_proj_out_dim, wv_group_size * d_o,
+    num_groups=n_wv_group, recipe=(1, 1, scale_gran), is_sfa=False,
+)
+wv_proj_fp8 = wv_proj_fp8.view(n_wv_group, wv_proj_out_dim, wv_group_size * d_o)
+
+# Permute the weight and its scale factors into the layout required by the fused kernel
+wv_proj_fp8, wv_sf = fused_norm_rope_attn_rope_cast.permute_wv_proj(
+    (wv_proj_fp8, wv_sf), wv_group_size, d_o,
+)
+```
+
+And finally, to use the fused kernel:
+
+```python
+# q: [s_q, h_q, d_qk], bfloat16, i.e. the Q_b projection computed with the permuted weight above
+out_fp8, out_sf, max_logits, lse = fused_norm_rope_attn_rope_cast.prefill(
+    enable_q_norm,              # False for DeepSeek-V4.1
+    rms_norm_eps,               # e.g. 1e-4
+    token_positions,            # [s_q], int32
+    False, 64, cos_sin_cache,   # non-neox RoPE with rope_dim = 64
+    n_wv_group,                 # h_q // wv_group_size
+    32,                         # num_per_channels
+    True, True, True,           # use_tma_aligned_col_major_sf, round_sf, use_packed_ue8m0
+    q, kv, indices,             # bf16 Q, bf16 KV [s_kv, h_kv, d_qk], int32 indices [s_q, h_kv, topk]
+    sm_scale=sm_scale,
+    attn_sink=attn_sink,        # optional, [h_q], float32
+    topk_length=topk_length,    # optional, [s_q], int32
+)
+
+# For decoding, call `decode` instead, passing the paged quantized KV cache:
+#   q: [s_q, h_q, d_qk], bf16
+#   k_cache: [num_blocks, page_block_size, h_kv, bytes_per_token], fp8_e4m3
+#   indices_in_kvcache: [s_q, topk], int32
+out_fp8, out_sf, lse = fused_norm_rope_attn_rope_cast.decode(
+    enable_q_norm, rms_norm_eps,
+    token_positions, False, 64, cos_sin_cache,
+    n_wv_group, 32, True, True, True,
+    q, k_cache, indices_in_kvcache,
+    sm_scale=sm_scale,
+    attn_sink=attn_sink,
+    topk_length=topk_length,
+    extra_k_cache=extra_k_cache,                        # optional, same layout as k_cache
+    extra_indices_in_kvcache=extra_indices_in_kvcache,  # optional, [s_q, extra_topk], int32
+    extra_topk_length=extra_topk_length,                # optional, [s_q], int32
+)
+
+# The FP8 output is consumed directly by the Wv projection, using the permuted Wv weight
+wv_out = torch.empty((s_q, n_wv_group, wv_proj_out_dim), dtype=torch.bfloat16, device='cuda')
+deep_gemm.fp8_einsum("bhr,hdr->bhd", (out_fp8, out_sf), (wv_proj_fp8, wv_sf), wv_out, recipe=(1, 1, 32))
+```
+
+You may refer to the fused kernel's test script ([tests/test_fused_norm_rope_attn_rope_cast.py](tests/test_fused_norm_rope_attn_rope_cast.py)) for a complete example.
+
 ## Acknowledgement
 
 FlashMLA is inspired by [FlashAttention 2&3](https://github.com/dao-AILab/flash-attention/) and [cutlass](https://github.com/nvidia/cutlass) projects.

csrc/api/api.cpp

@@ -1,15 +1,18 @@
 #include <pybind11/pybind11.h>
 
-#include "sparse_fwd.h"
-#include "sparse_decode.h"
-#include "dense_decode.h"
-#include "dense_fwd.h"
+void register_sparse_prefill(pybind11::module_& m);
+void register_sparse_decode(pybind11::module_& m);
+void register_dense_fwd(pybind11::module_& m);
+void register_dense_bwd(pybind11::module_& m);
+void register_dense_decode(pybind11::module_& m);
+void register_fused_norm_rope_attn_rope_cast_fwd(pybind11::module_& m);
 
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
     m.doc() = "FlashMLA";
-    m.def("sparse_decode_fwd", &sparse_attn_decode_interface);
-    m.def("dense_decode_fwd", &dense_attn_decode_interface);
-    m.def("sparse_prefill_fwd", &sparse_attn_prefill_interface);
-    m.def("dense_prefill_fwd", &FMHACutlassSM100FwdRun);
-    m.def("dense_prefill_bwd", &FMHACutlassSM100BwdRun);
+    register_sparse_prefill(m);
+    register_sparse_decode(m);
+    register_dense_fwd(m);
+    register_dense_bwd(m);
+    register_dense_decode(m);
+    register_fused_norm_rope_attn_rope_cast_fwd(m);
 }

csrc/api/common.h

@@ -9,6 +9,8 @@
 
 #include <cutlass/bfloat16.h>
 
+#include "kernels/kv_cache_format.h"
+
 static constexpr float LOG_2_E = 1.44269504f;
 
 // Instantiation for tensor.data_ptr<cutlass::bfloat16_t>()
@@ -90,8 +92,8 @@ 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 { \
         TORCH_CHECK(false, "Unsupported model type: ", (int)MODEL_TYPE); \
@@ -140,6 +142,29 @@ 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)
+// =============================================
+
+// 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;
+        }
+    }
+    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) || ...);
+    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_bwd.cpp

@@ -0,0 +1,9 @@
+#include "common.h"
+
+#include "kernels/sm100/prefill/dense/interface.h"
+
+void register_dense_bwd(pybind11::module_& m) {
+    m.def("dense_prefill_bwd",
+        &FMHACutlassSM100BwdRun,
+        "Run Dense Attention Prefill Backward (cutlass FMHA)");
+}

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

@@ -1,14 +1,12 @@
-#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<at::Tensor, at::Tensor, std::optional<at::Tensor>, std::optional<at::Tensor>>
 dense_attn_decode_interface(
@@ -173,12 +171,12 @@ dense_attn_decode_interface(
     params.stream = at::cuda::getCurrentCUDAStream().stream();
 
     if (q_dtype == torch::kBFloat16) {
-        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 == torch::kHalf) {
 #ifdef FLASH_MLA_DISABLE_FP16
         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 {
         TORCH_CHECK(false, "Unsupported dtype for dense MLA on SM90");
@@ -223,3 +221,9 @@ dense_attn_decode_interface(
 
     return {out, lse, tile_scheduler_metadata, num_splits};
 }
+
+void register_dense_decode(pybind11::module_& m) {
+    m.def("dense_decode_fwd",
+        &dense_attn_decode_interface,
+        "Run Dense Attention Decode Forward");
+}