diff --git a/.gitignore b/.gitignore
index 036d4367..bdd83f48 100644
--- a/.gitignore
+++ b/.gitignore
@@ -10,3 +10,4 @@ compile_commands.json
.cache
/dev
/.clangd
+/.venv/
diff --git a/README.md b/README.md
index 1945725f..ba420f6d 100644
--- a/README.md
+++ b/README.md
@@ -16,8 +16,94 @@ FlashMLA is DeepSeek's library of optimized attention kernels, powering the [Dee
- Dense attention for the prefill stage
- Dense attention for the decoding stage
+## Changes in the vLLM fork
+
+This repository is [vLLM](https://github.com/vllm-project/vllm)'s fork of [deepseek-ai/FlashMLA](https://github.com/deepseek-ai/FlashMLA). It tracks upstream (currently synced through [deepseek-ai/FlashMLA@07a1089](https://github.com/deepseek-ai/FlashMLA/commit/07a1089857b63e74e3133630c02b083b75e8d4b2), which includes the DeepSeek V4.1 kernels) and adds the following on top of it. vLLM compiles these sources directly through its `cmake/external_projects/flashmla.cmake`; set `FLASH_MLA_SRC_DIR` to build vLLM against a local checkout.
+
+- **PyTorch stable ABI.** The API layer in `csrc/api/` uses the libtorch stable ABI (`torch/csrc/stable`, `STABLE_TORCH_LIBRARY`) instead of `torch/extension.h` and pybind11. Operators are registered under `torch.ops._flashmla_C`, and the module defines `PyInit__flashmla_C` so that vLLM can import it as `vllm._flashmla_C`. The extension is built against CPython's limited API, so a single `abi3` wheel works for every CPython >= 3.10; it requires PyTorch >= 2.10 at runtime. `flash_mla/flash_mla_interface.py` calls the operators through `torch.ops._flashmla_C` and is vendored into vLLM as `vllm/third_party/flashmla/flash_mla_interface.py`.
+- **Registered operators.** `sparse_decode_fwd`, `dense_decode_fwd`, `sparse_prefill_fwd`, `dense_prefill_fwd`, `fused_norm_rope_attn_rope_cast_fwd`, `fused_norm_rope_attn_rope_cast_decode`, `permute_q_b_proj` and `permute_wv_proj`. `dense_prefill_bwd` is registered only when the extension is compiled with `FLASH_MLA_ENABLE_DENSE_BWD` (set by the standalone `setup.py`; vLLM's inference-only build leaves it out).
+- **Optional output buffers.** `flash_mla_with_kvcache(..., out=...)` (sparse and dense decoding) and `flash_mla_sparse_fwd(..., out=...)` write into a caller-provided tensor instead of allocating a new one.
+- **Dense FP8 KV-cache decoding on SM90.** `csrc/extension/sm90/dense_fp8/` contains vLLM's Hopper MLA decoding kernel for FP8 KV caches. vLLM builds it as a separate `_flashmla_extension_C` module (`fwd_kvcache_mla_fp8`, `get_mla_decoding_metadata_dense_fp8`); it is not part of the standalone `setup.py` build.
+- **NVFP4 KV cache for sparse decoding on SM100.** Besides the 656-byte FP8 layout, sparse decoding with `head_dim == 576` accepts a 352-byte-per-token layout: 256 bytes of e2m1 NoPE values, 64 bytes of unscaled e4m3 RoPE values and 32 e4m3 scale factors (one per 16 NoPE values). The layout is detected from `k_cache.shape[-1]`. See `KVCacheLayout.V32_NVFP4_FP8ROPE` in `tests/quant.py` for the exact wire format and `csrc/kernels/sm100/decode/sparse/nvfp4_head64/` for the kernel; 64 and 128 query heads are supported.
+- **Robustness fixes.** Thread-safe cached device properties and stable-ABI stream and tensor helpers in `csrc/kerutils/include/kerutils/supplemental/`, plus fixes to the dense FP8 decoding metadata (for example `num_sm_parts` is clamped to at least 1).
+- **Tests.** `tests/test_api_registration.py`, `tests/test_output_buffer_api.py` and `tests/test_nvfp4_quant.py`, plus NVFP4 cases in `tests/test_flash_mla_sparse_decoding.py`.
+
+### Detailed diff against upstream
+
+Everything below is the complete set of files that differ from upstream; regenerate it with `git diff --stat 07a1089857b63e74e3133630c02b083b75e8d4b2 HEAD -- . ':!README.md'`. Files that are not listed are identical to upstream.
+
+| Area | Files | Difference from upstream |
+| :--- | :--- | :--- |
+| Operator registration | `csrc/api/api.cpp`, `csrc/api/interfaces.h` (new), `csrc/api/dense_fwd.cpp` and `csrc/api/dense_bwd.cpp` (removed) | The pybind11 `PYBIND11_MODULE` and per-file `register_*` shims are replaced by `STABLE_TORCH_LIBRARY` / `STABLE_TORCH_LIBRARY_IMPL` with explicit operator schemas. All interface functions are declared once in `interfaces.h`. `dense_prefill_bwd` is only registered under `FLASH_MLA_ENABLE_DENSE_BWD`, and `PyInit__flashmla_C` lets vLLM import the library as a Python module. |
+| API implementations | `csrc/api/sparse_decode.cpp`, `csrc/api/dense_decode.cpp`, `csrc/api/sparse_prefill.cpp`, `csrc/api/fused_norm_rope_attn_rope_cast_fwd.cpp`, `csrc/api/common.h` | `at::Tensor`, `TORCH_CHECK`, `torch::empty`, `at::cuda::CUDAGuard` and `at::cuda::getCurrentCUDAStream` are replaced by `torch::stable::Tensor`, `STD_TORCH_CHECK`, `torch::stable::new_empty`, `torch::stable::accelerator::DeviceGuard` and the stable stream helper. Scalar arguments are widened to `int64_t` / `double` as stable schemas require. Sparse decode, dense decode and sparse prefill take an optional `out_` buffer. `Arch` reads the cached device properties. Sparse decode detects the NVFP4 layout from bytes per token (`detect_kv_cache_format_for_headdim_576`), advertises `NVFP4_FP8ROPE_KVCACHE_FORMAT` on the SM100 head-64 and head-64x2 implementations, and always uses split-KV scheduling for it. |
+| Dense MHA prefill entry points | `csrc/kernels/sm100/prefill/dense/interface.h`, `fmha_cutlass_fwd_sm100.cu` / `.cuh`, `fmha_cutlass_bwd_sm100.cu` / `.cuh`, `common/utils.hpp` | `FMHACutlassSM100FwdRun` and `FMHACutlassSM100BwdRun` take `torch::stable::Tensor` and `int64_t` / `double` scalars so they can be registered directly as stable operators. |
+| Stable-ABI helpers | `csrc/kerutils/include/kerutils/supplemental/cuda_stream.h`, `device_prop.h`, `torch_tensors.h` | `get_current_cuda_stream(tensor)` through the AOTI shim, a thread-safe per-device `cudaDeviceProp` cache (`std::once_flag`) replacing `at::cuda::getCurrentDeviceProperties()`, and the `KU_CHECK_*` / `get_optional_tensor_ptr` helpers rewritten over `torch::stable::Tensor`. |
+| NVFP4 KV cache (SM100 sparse decode) | `csrc/kernels/params.h`, `csrc/kernels/kv_cache_format.h`, `csrc/kernels/sm100/decode/sparse/nvfp4_head64/config.h`, `kernel.cuh`, `kernel.h`, `instantiations/v32_nvfp4_fp8rope.cu`, `csrc/kernels/sm100/helpers.h` | New `ModelType::V32_NVFP4_FP8ROPE`; `KVCacheFormat` describes the 352-byte record (256 B e2m1 NoPE, 64 B e4m3 RoPE, 32 B e4m3 scales) and `kv_cache_bytes_per_token` returns it. The kernel is a dedicated head-64 decode kernel, derived from the pre-V4.1 SM100 head-64 kernel and extended with e2m1 dequantization; `helpers.h` gains a bf16-scale overload of `fp8x2_to_bf16x2_with_scale` that it uses. |
+| SM90 dense FP8 decoding extension | `csrc/extension/torch_api.cpp`, `csrc/extension/sm90/dense_fp8/*` | vLLM-only sources for the `_flashmla_extension_C` module (`fwd_kvcache_mla_fp8`, `get_mla_decoding_metadata_dense_fp8`), already ported to the stable ABI. Not compiled by `setup.py`; vLLM's CMake builds them. |
+| Python package | `flash_mla/__init__.py`, `flash_mla/flash_mla_interface.py`, `flash_mla/fused_norm_rope_attn_rope_cast.py` | `__init__` loads `_flashmla_C*.so` with `torch.ops.load_library`; call sites use `torch.ops._flashmla_C` instead of the `flash_mla.cuda` pybind module; `flash_mla_with_kvcache` and `flash_mla_sparse_fwd` accept `out=` and document the NVFP4 layout. |
+| Build | `setup.py`, `.gitignore` | The extension is named `flash_mla._flashmla_C`; it is compiled with `-DTORCH_TARGET_VERSION=0x020a000000000000 -DUSE_CUDA -DFLASH_MLA_ENABLE_DENSE_BWD`, `py_limited_api=True` and `bdist_wheel.py_limited_api = cp310`; the NVFP4 instantiation is added and the removed `dense_fwd.cpp` / `dense_bwd.cpp` are dropped from the source list; `.venv/` is ignored. |
+| Tests | `tests/lib.py`, `tests/quant.py`, `tests/test_flash_mla_sparse_decoding.py`, `tests/test_api_registration.py`, `tests/test_output_buffer_api.py`, `tests/test_nvfp4_quant.py` | `KVCacheLayout.V32_NVFP4_FP8ROPE` quantization and dequantization (including the scale-byte permutation) and per-layout byte accounting for the bandwidth numbers; NVFP4 correctness, corner and performance cases; new tests for operator registration, `out=` forwarding and the NVFP4 wire format. |
+
+
+File-level diffstat against deepseek-ai/FlashMLA@07a1089
+
+```text
+ .gitignore | 1 +
+ csrc/api/api.cpp | 55 ++++--
+ csrc/api/common.h | 69 ++++---
+ csrc/api/dense_bwd.cpp | 9 -
+ csrc/api/dense_decode.cpp | 151 ++++++++-------
+ csrc/api/dense_fwd.cpp | 9 -
+ csrc/api/fused_norm_rope_attn_rope_cast_fwd.cpp | 281 ++++++++++++++--------------
+ csrc/api/interfaces.h | 102 ++++++++++
+ csrc/api/sparse_decode.cpp | 183 ++++++++++--------
+ csrc/api/sparse_prefill.cpp | 71 +++----
+ csrc/extension/sm90/dense_fp8/flash_fwd_mla_fp8_sm90.cu | 10 +
+ csrc/extension/sm90/dense_fp8/flash_fwd_mla_kernel.h | 709 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
+ csrc/extension/sm90/dense_fp8/flash_fwd_mla_metadata.cu | 77 ++++++++
+ csrc/extension/sm90/dense_fp8/flash_mla.h | 85 +++++++++
+ csrc/extension/sm90/dense_fp8/fp8_transpose_v.h | 89 +++++++++
+ csrc/extension/sm90/dense_fp8/named_barrier.h | 21 +++
+ csrc/extension/sm90/dense_fp8/pybind.cpp | 246 +++++++++++++++++++++++++
+ csrc/extension/sm90/dense_fp8/softmax.h | 202 ++++++++++++++++++++
+ csrc/extension/sm90/dense_fp8/static_switch.h | 70 +++++++
+ csrc/extension/sm90/dense_fp8/utils.h | 279 ++++++++++++++++++++++++++++
+ csrc/extension/torch_api.cpp | 47 +++++
+ csrc/kernels/kv_cache_format.h | 19 +-
+ csrc/kernels/params.h | 5 +-
+ csrc/kernels/sm100/decode/sparse/head64/config.h | 2 +-
+ csrc/kernels/sm100/decode/sparse/nvfp4_head64/config.h | 270 +++++++++++++++++++++++++++
+ csrc/kernels/sm100/decode/sparse/nvfp4_head64/instantiations/v32_nvfp4_fp8rope.cu | 8 +
+ csrc/kernels/sm100/decode/sparse/nvfp4_head64/kernel.cuh | 1103 +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
+ csrc/kernels/sm100/decode/sparse/nvfp4_head64/kernel.h | 10 +
+ csrc/kernels/sm100/helpers.h | 8 +
+ csrc/kernels/sm100/prefill/dense/common/utils.hpp | 1 -
+ csrc/kernels/sm100/prefill/dense/fmha_cutlass_bwd_sm100.cu | 27 +--
+ csrc/kernels/sm100/prefill/dense/fmha_cutlass_bwd_sm100.cuh | 50 ++---
+ csrc/kernels/sm100/prefill/dense/fmha_cutlass_fwd_sm100.cu | 29 +--
+ csrc/kernels/sm100/prefill/dense/fmha_cutlass_fwd_sm100.cuh | 39 ++--
+ csrc/kernels/sm100/prefill/dense/interface.h | 20 +-
+ csrc/kerutils/include/kerutils/supplemental/cuda_stream.h | 19 ++
+ csrc/kerutils/include/kerutils/supplemental/device_prop.h | 56 ++++++
+ csrc/kerutils/include/kerutils/supplemental/torch_tensors.h | 29 +--
+ flash_mla/__init__.py | 10 +
+ flash_mla/flash_mla_interface.py | 25 ++-
+ flash_mla/fused_norm_rope_attn_rope_cast.py | 4 +-
+ setup.py | 21 ++-
+ tests/lib.py | 10 +-
+ tests/quant.py | 78 +++++++-
+ tests/test_api_registration.py | 24 +++
+ tests/test_flash_mla_sparse_decoding.py | 49 +++++
+ tests/test_nvfp4_quant.py | 28 +++
+ tests/test_output_buffer_api.py | 87 +++++++++
+ 48 files changed, 4306 insertions(+), 491 deletions(-)
+```
+
+
+
## 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 +118,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 +136,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 +155,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 +211,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 +237,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 +272,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 +283,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.
diff --git a/csrc/api/api.cpp b/csrc/api/api.cpp
index fde82b3f..c6109cf1 100644
--- a/csrc/api/api.cpp
+++ b/csrc/api/api.cpp
@@ -2,16 +2,17 @@
#include
-#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
diff --git a/csrc/api/common.h b/csrc/api/common.h
index f0546a27..d369a77b 100644
--- a/csrc/api/common.h
+++ b/csrc/api/common.h
@@ -21,6 +21,8 @@
#include
+#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
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(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
+inline void dispatch_kv_formats(KVFormatPairs, ModelType kv, ModelType extra_kv, Fn &&fn) {
+ bool matched = ((kv == Pairs::kv && extra_kv == Pairs::extra_kv ? (fn.template operator()(), 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: \
diff --git a/csrc/api/dense_decode.h b/csrc/api/dense_decode.cpp
similarity index 95%
rename from csrc/api/dense_decode.h
rename to csrc/api/dense_decode.cpp
index d52ea8cb..34499600 100644
--- a/csrc/api/dense_decode.h
+++ b/csrc/api/dense_decode.cpp
@@ -1,16 +1,14 @@
-#pragma once
-
#include
#include
#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, std::optional>
+std::tuple, std::optional>
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(params);
+ sm90::decode::dense::run_flash_splitkv_mla_kernel(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(params);
+ sm90::decode::dense::run_flash_splitkv_mla_kernel(params);
#endif
} else {
STD_TORCH_CHECK(false, "Unsupported dtype for dense MLA on SM90");
diff --git a/csrc/api/dense_fwd.h b/csrc/api/dense_fwd.h
deleted file mode 100644
index c5a4acfb..00000000
--- a/csrc/api/dense_fwd.h
+++ /dev/null
@@ -1,5 +0,0 @@
-#pragma once
-
-#include "common.h"
-
-#include "sm100/prefill/dense/interface.h"
diff --git a/csrc/api/fused_norm_rope_attn_rope_cast_fwd.cpp b/csrc/api/fused_norm_rope_attn_rope_cast_fwd.cpp
new file mode 100644
index 00000000..81e2093e
--- /dev/null
+++ b/csrc/api/fused_norm_rope_attn_rope_cast_fwd.cpp
@@ -0,0 +1,542 @@
+#include "common.h"
+
+#include "kernels/params.h"
+
+#include "kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/kernel.h"
+#include "kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/permute_q_b_proj/kernel.h"
+#include "kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/permute_wv_proj/kernel.h"
+
+// Local aliases: `kernels/defines.h` declares this type as `fp8`, and the fused kernel headers declare
+// it inside their own namespace, so this translation unit needs the explicit name at file scope.
+using bf16 = cutlass::bfloat16_t;
+using fp8_e4m3 = cutlass::float_e4m3_t;
+
+using Params = sm100::prefill::fused_norm_rope_attn_rope_cast_fwd::core_attn::ParamT;
+using DecodeParams = sm100::prefill::fused_norm_rope_attn_rope_cast_fwd::core_attn::ParamT;
+using Config = sm100::prefill::fused_norm_rope_attn_rope_cast_fwd::core_attn::Config;
+
+static Tensor allocate_scale_factor(
+ uint32_t batch_size,
+ uint32_t hidden_dim,
+ uint32_t scale_gran,
+ const Tensor &like,
+ std::optional extra_dim = std::nullopt) {
+ // Allocate a scale factor tensor, which should have shape ([extra_dim], batch_size, hidden_dim / (scale_gran*4))
+ // Meet DeepGEMM's SF requirement under use_tma_aligned_col_major_sf == True, round_sf == True, and use_packed_ue8m0 == True
+ STD_TORCH_CHECK(hidden_dim % (scale_gran*4) == 0);
+ uint32_t sf_align_requirement = 16u / sizeof(int32_t);
+ uint32_t aligned_batch_size_for_sf = (batch_size + sf_align_requirement - 1) / sf_align_requirement * sf_align_requirement;
+ uint32_t leading_dim = extra_dim.value_or(1);
+ Tensor sf = torch::stable::new_empty(
+ like,
+ {leading_dim, hidden_dim / (scale_gran * 4), aligned_batch_size_for_sf},
+ ScalarType::Int);
+ KU_CHECK_CONTIGUOUS(sf);
+ sf = torch::stable::transpose(sf, 1, 2);
+ sf = torch::stable::narrow(sf, 1, 0, batch_size); // [leading_dim, batch_size, hidden_dim / (scale_gran*4)], int32
+ if (!extra_dim.has_value())
+ sf = torch::stable::squeeze(sf, 0);
+ return sf;
+}
+
+std::vector fused_norm_rope_attn_rope_cast_fwd(
+ const Tensor &q,
+ const Tensor &kv,
+ const Tensor &indices,
+ double sm_scale,
+ int64_t d_v,
+ const std::optional &attn_sink,
+ const std::optional &topk_length,
+ bool enable_q_norm,
+ double rms_norm_eps,
+ const Tensor &token_positions,
+ bool is_rope_neox_style,
+ int64_t rope_dim,
+ const Tensor &cos_sin_cache,
+
+ int64_t n_wv_group,
+ int64_t num_per_channels,
+ bool use_tma_aligned_col_major_sf,
+ bool round_sf,
+ bool use_packed_ue8m0
+) {
+ Arch arch = Arch();
+ bool is_sm100f = arch.is_sm100f();
+ STD_TORCH_CHECK(is_sm100f, "Fused Norm + RoPE + Core Attn + RoPE + Cast (fused_norm_rope_attn_rope_cast_fwd) is only supported on SM100f architectures.");
+
+ KU_CHECK_NDIM(q, 3);
+ KU_CHECK_NDIM(kv, 3);
+ KU_CHECK_NDIM(indices, 3);
+ KU_CHECK_NDIM(attn_sink, 1);
+ KU_CHECK_NDIM(topk_length, 1);
+ KU_CHECK_NDIM(token_positions, 1);
+ KU_CHECK_NDIM(cos_sin_cache, 2);
+
+ int s_q = q.size(0);
+ int s_kv = kv.size(0);
+ int h_q = q.size(1);
+ int h_kv = kv.size(1);
+ int d_qk = q.size(2);
+ int topk = indices.size(2);
+ uint32_t wv_group_size = h_q / n_wv_group;
+
+ STD_TORCH_CHECK(h_q % n_wv_group == 0, "h_q %% n_wv_group != 0");
+ STD_TORCH_CHECK(is_rope_neox_style == false, "Only `is_rope_neox_style == False` is supported");
+ STD_TORCH_CHECK(use_tma_aligned_col_major_sf == true, "`use_tma_aligned_col_major_sf` must be True");
+ STD_TORCH_CHECK(round_sf == true, "`round_sf` must be True");
+ STD_TORCH_CHECK(use_packed_ue8m0 == true, "`use_packed_ue8m0` must be True");
+
+ KU_CHECK_DEVICE(q);
+ KU_CHECK_DEVICE(kv);
+ KU_CHECK_DEVICE(indices);
+ KU_CHECK_DEVICE(attn_sink);
+ KU_CHECK_DEVICE(topk_length);
+ KU_CHECK_DEVICE(token_positions);
+ KU_CHECK_DEVICE(cos_sin_cache);
+
+ KU_CHECK_DTYPE(q, ScalarType::BFloat16);
+ KU_CHECK_DTYPE(kv, ScalarType::BFloat16);
+ KU_CHECK_DTYPE(indices, ScalarType::Int);
+ KU_CHECK_DTYPE(attn_sink, ScalarType::Float);
+ KU_CHECK_DTYPE(topk_length, ScalarType::Int);
+ KU_CHECK_DTYPE(token_positions, ScalarType::Int);
+ KU_CHECK_DTYPE(cos_sin_cache, ScalarType::Float);
+
+ KU_CHECK_SHAPE(q, s_q, h_q, d_qk);
+ KU_CHECK_SHAPE(kv, s_kv, h_kv, d_qk);
+ KU_CHECK_SHAPE(indices, s_q, h_kv, topk);
+ KU_CHECK_SHAPE(attn_sink, h_q);
+ KU_CHECK_SHAPE(topk_length, s_q);
+ KU_CHECK_SHAPE(token_positions, s_q);
+ KU_CHECK_SHAPE(cos_sin_cache, cos_sin_cache.size(0), rope_dim);
+
+ KU_CHECK_LAST_DIM_CONTIGUOUS(q);
+ // q is in the permuted layout (see permute_q_b_proj), so the kernel assumes that the h_q*d_qk elements of one token are contiguous (only q.stride(0) is used by the kernel)
+ STD_TORCH_CHECK(q.stride(1) == d_qk, "q must be contiguous within each token (i.e. q.stride(1) == d_qk), since q is in the permuted layout, got q.stride(1) = ", q.stride(1));
+ KU_CHECK_LAST_DIM_CONTIGUOUS(kv);
+ KU_CHECK_LAST_DIM_CONTIGUOUS(indices);
+ KU_CHECK_LAST_DIM_CONTIGUOUS(attn_sink);
+ KU_CHECK_CONTIGUOUS(topk_length);
+ KU_CHECK_CONTIGUOUS(token_positions);
+ KU_CHECK_CONTIGUOUS(cos_sin_cache);
+
+ STD_TORCH_CHECK(num_per_channels == 32, "num_per_channels must be 32, got ", num_per_channels);
+
+ torch::stable::accelerator::DeviceGuard device_guard(q.get_device_index());
+
+ STD_TORCH_CHECK(d_v % (num_per_channels * 4) == 0); // 4 is the number of uint8 in uint32, since `use_packed_ue8m0` is `True`
+ Tensor out_fp8 = torch::stable::new_empty(q, {s_q, n_wv_group, wv_group_size * d_v}, ScalarType::Float8_e4m3fn);
+ uint32_t out_sf_scale_gran = 32; // Since the weight is per-32 scaled and deep_gemm.einsum requires A and B to have the same scale granularity, the output sf is always stored in a per-32 scaled format, although it will be actually per-128 scaled when num_per_channels is 128
+ Tensor out_sf = allocate_scale_factor(s_q, wv_group_size * d_v, out_sf_scale_gran, q, n_wv_group);
+ out_sf = torch::stable::transpose(out_sf, 0, 1); // [s_q, n_wv_group, wv_group_size * d_v / (out_sf_scale_gran*4)]
+ Tensor max_logits = torch::stable::new_empty(q, {s_q, h_q}, ScalarType::Float);
+ Tensor lse = torch::stable::new_empty(q, {s_q, h_q}, ScalarType::Float);
+ KU_CHECK_CONTIGUOUS(out_fp8);
+ STD_TORCH_CHECK(out_sf.stride(0) == 1);
+ KU_CHECK_CONTIGUOUS(max_logits);
+ KU_CHECK_CONTIGUOUS(lse);
+
+ Params params = {
+ s_q, s_kv, h_q, h_kv, d_qk, d_v, topk,
+ sm_scale, sm_scale * LOG_2_E,
+
+ (bf16*)q.data_ptr(),
+ (bf16*)kv.data_ptr(),
+ (int*)indices.data_ptr(),
+ ku::get_optional_tensor_ptr(attn_sink),
+ ku::get_optional_tensor_ptr(topk_length),
+
+ int64_stride_to_int(q.stride(0)), int64_stride_to_int(q.stride(1)),
+ int64_stride_to_int(kv.stride(0)), int64_stride_to_int(kv.stride(1)),
+ int64_stride_to_int(indices.stride(0)), int64_stride_to_int(indices.stride(1)),
+
+ nullptr,
+ (float*)max_logits.data_ptr(),
+ (float*)lse.data_ptr(),
+
+ arch.num_sms,
+ get_current_cuda_stream(q),
+
+ enable_q_norm,
+ rms_norm_eps,
+ (uint32_t*)token_positions.data_ptr(),
+ is_rope_neox_style,
+ rope_dim,
+ (float*)cos_sin_cache.data_ptr(),
+
+ n_wv_group,
+ wv_group_size,
+ num_per_channels,
+ use_tma_aligned_col_major_sf,
+ round_sf,
+ use_packed_ue8m0,
+
+ (fp8_e4m3*)out_fp8.data_ptr(),
+ (uint32_t*)out_sf.data_ptr(),
+ (uint32_t)int64_stride_to_int(out_sf.stride(1)),
+ (uint32_t)int64_stride_to_int(out_sf.stride(2))
+ };
+
+ STD_TORCH_CHECK(h_q == 64 || h_q == 128, "Only h_q == 64 or 128 is supported for fused_norm_rope_attn_rope_cast_fwd, got ", h_q);
+ DISPATCH_NUM_HEADS(h_q, H_Q, ([&]() {
+ DISPATCH_BOOLEAN_FLAG(enable_q_norm, ENABLE_Q_NORM, ([&]() {
+ sm100::prefill::fused_norm_rope_attn_rope_cast_fwd::core_attn::run_fused_norm_rope_attn_rope_cast_fwd_kernel(params);
+ }));
+ }));
+
+ return {out_fp8, out_sf, max_logits, lse};
+}
+
+
+std::vector fused_norm_rope_attn_rope_cast_decode(
+ const Tensor &q, // [s_q, h_q, d_qk]
+ const Tensor &kv, // [num_blocks, page_block_size, h_kv, bytes_per_token], paged quantized KV cache
+ const Tensor &indices, // [s_q, topk]
+ double sm_scale,
+ int64_t d_v,
+ const std::optional &attn_sink, // [h_q]
+ const std::optional &topk_length, // [s_q]
+ const std::optional &extra_kv, // [extra_num_blocks, extra_page_block_size, h_kv, bytes_per_token]
+ const std::optional &extra_indices, // [s_q, extra_topk]
+ const std::optional &extra_topk_length, // [s_q]
+ bool enable_q_norm,
+ double rms_norm_eps,
+ const Tensor &token_positions, // [s_q]
+ bool is_rope_neox_style,
+ int64_t rope_dim,
+ const Tensor &cos_sin_cache, // [*, rope_dim]
+
+ int64_t n_wv_group,
+ int64_t num_per_channels,
+ bool use_tma_aligned_col_major_sf,
+ bool round_sf,
+ bool use_packed_ue8m0
+) {
+ Arch arch = Arch();
+ STD_TORCH_CHECK(arch.is_sm100f(), "Fused Norm + RoPE + Core Attn + RoPE + Cast (fused_norm_rope_attn_rope_cast_decode) is only supported on SM100f architectures.");
+
+ KU_CHECK_NDIM(q, 3);
+ KU_CHECK_NDIM(kv, 4);
+ KU_CHECK_NDIM(indices, 2);
+ KU_CHECK_NDIM(attn_sink, 1);
+ KU_CHECK_NDIM(topk_length, 1);
+ KU_CHECK_NDIM(extra_kv, 4);
+ KU_CHECK_NDIM(extra_indices, 2);
+ KU_CHECK_NDIM(extra_topk_length, 1);
+ KU_CHECK_NDIM(token_positions, 1);
+ KU_CHECK_NDIM(cos_sin_cache, 2);
+
+ int s_q = q.size(0);
+ int h_q = q.size(1);
+ int d_qk = q.size(2);
+ int num_blocks = kv.size(0);
+ int page_block_size = kv.size(1);
+ int h_kv = kv.size(2);
+ int topk = indices.size(1);
+
+ bool have_extra_kvcache = extra_kv.has_value();
+
+ int extra_num_blocks = 0, extra_page_block_size = 0, extra_topk = 0;
+ if (have_extra_kvcache) {
+ extra_num_blocks = extra_kv->size(0);
+ extra_page_block_size = extra_kv->size(1);
+ }
+ if (extra_indices.has_value()) {
+ extra_topk = extra_indices->size(-1);
+ }
+
+ // Metadata sanity check
+ STD_TORCH_CHECK(s_q > 0);
+ STD_TORCH_CHECK(h_q == 64 || h_q == 128, "Only h_q == 64 or 128 is supported for fused_norm_rope_attn_rope_cast_decode, got ", h_q);
+ STD_TORCH_CHECK(h_kv == 1, "Currently only MQA (i.e. h_kv == 1) is supported");
+ STD_TORCH_CHECK(d_qk == 512, "Only head_size_k == 512 (V4 / V4.1) is supported");
+ STD_TORCH_CHECK(d_v == 512, "Only head_size_v == 512 is supported");
+ STD_TORCH_CHECK(topk > 0);
+ STD_TORCH_CHECK(h_q % n_wv_group == 0, "h_q %% n_wv_group != 0");
+ uint32_t wv_group_size = h_q / n_wv_group;
+
+ if (have_extra_kvcache) {
+ STD_TORCH_CHECK(extra_indices.has_value(), "extra_indices must be provided when extra_kv is provided");
+ } else {
+ STD_TORCH_CHECK(!extra_indices.has_value(), "extra_indices must not be provided when extra_kv is not provided");
+ STD_TORCH_CHECK(!extra_topk_length.has_value(), "extra_topk_length must not be provided when extra_kv is not provided");
+ }
+
+ STD_TORCH_CHECK(is_rope_neox_style == false, "Only `is_rope_neox_style == False` is supported");
+ STD_TORCH_CHECK(use_tma_aligned_col_major_sf == true, "`use_tma_aligned_col_major_sf` must be True");
+ STD_TORCH_CHECK(round_sf == true, "`round_sf` must be True");
+ STD_TORCH_CHECK(use_packed_ue8m0 == true, "`use_packed_ue8m0` must be True");
+
+ // Check device
+ KU_CHECK_DEVICE(q);
+ KU_CHECK_DEVICE(kv);
+ KU_CHECK_DEVICE(indices);
+ KU_CHECK_DEVICE(attn_sink);
+ KU_CHECK_DEVICE(topk_length);
+ KU_CHECK_DEVICE(extra_kv);
+ KU_CHECK_DEVICE(extra_indices);
+ KU_CHECK_DEVICE(extra_topk_length);
+ KU_CHECK_DEVICE(token_positions);
+ KU_CHECK_DEVICE(cos_sin_cache);
+
+ // Check data type
+ KU_CHECK_DTYPE(q, ScalarType::BFloat16);
+ STD_TORCH_CHECK(kv.scalar_type() == ScalarType::Float8_e4m3fn || kv.scalar_type() == ScalarType::Char || kv.scalar_type() == ScalarType::Byte, "kv must have dtype fp8_e4m3fn, int8 or uint8");
+ if (have_extra_kvcache) {
+ STD_TORCH_CHECK(extra_kv->scalar_type() == ScalarType::Float8_e4m3fn || extra_kv->scalar_type() == ScalarType::Char || extra_kv->scalar_type() == ScalarType::Byte, "extra_kv must have dtype fp8_e4m3fn, int8 or uint8");
+ }
+ KU_CHECK_DTYPE(indices, ScalarType::Int);
+ KU_CHECK_DTYPE(attn_sink, ScalarType::Float);
+ KU_CHECK_DTYPE(topk_length, ScalarType::Int);
+ KU_CHECK_DTYPE(extra_indices, ScalarType::Int);
+ KU_CHECK_DTYPE(extra_topk_length, ScalarType::Int);
+ KU_CHECK_DTYPE(token_positions, ScalarType::Int);
+ KU_CHECK_DTYPE(cos_sin_cache, ScalarType::Float);
+
+ // Check layout
+ KU_CHECK_LAST_DIM_CONTIGUOUS(q);
+ // q is in the permuted layout (see permute_q_b_proj), so the kernel assumes that the h_q*d_qk elements of one token are contiguous (only q.stride(0) is used by the kernel)
+ STD_TORCH_CHECK(q.stride(1) == d_qk, "q must be contiguous within each token (i.e. q.stride(1) == d_qk), since q is in the permuted layout, got q.stride(1) = ", q.stride(1));
+ KU_CHECK_LAST_DIM_CONTIGUOUS(kv);
+ KU_CHECK_LAST_DIM_CONTIGUOUS(indices);
+ KU_CHECK_CONTIGUOUS(attn_sink);
+ KU_CHECK_CONTIGUOUS(topk_length);
+ KU_CHECK_LAST_DIM_CONTIGUOUS(extra_kv);
+ KU_CHECK_LAST_DIM_CONTIGUOUS(extra_indices);
+ KU_CHECK_CONTIGUOUS(extra_topk_length);
+ KU_CHECK_CONTIGUOUS(token_positions);
+ KU_CHECK_CONTIGUOUS(cos_sin_cache);
+
+ // The formats of `kv` and `extra_kv` (V4 / V4.1 / V4.1 fp4, see KVCacheFormat), detected by bytes_per_token
+ ModelType model_type = detect_kv_cache_format_for_headdim_512(kv.size(3));
+ ModelType extra_model_type = have_extra_kvcache ? detect_kv_cache_format_for_headdim_512(extra_kv->size(3)) : model_type;
+ STD_TORCH_CHECK(model_type != ModelType::V41_FP4, "The fp4 KV cache is only supported as extra_kv");
+ STD_TORCH_CHECK(is_valid_kv_format_pair(model_type, extra_model_type), "extra_kv must have the format of kv, or the V4.1 fp4 format when kv has the V4.1 format, got ",
+ get_dynamic_enum_name(model_type), " and ", get_dynamic_enum_name(extra_model_type));
+ KU_CHECK_SHAPE(kv, num_blocks, page_block_size, h_kv, kv_cache_bytes_per_token(model_type));
+ KU_CHECK_SHAPE(extra_kv, extra_num_blocks, extra_page_block_size, h_kv, kv_cache_bytes_per_token(extra_model_type));
+ STD_TORCH_CHECK(kv.stride(1) == kv_cache_bytes_per_token(model_type), "The whole block must be contiguous for the paged KV cache");
+ if (have_extra_kvcache) {
+ STD_TORCH_CHECK(extra_kv->stride(1) == kv_cache_bytes_per_token(extra_model_type), "The whole block must be contiguous for the paged extra KV cache");
+ }
+ STD_TORCH_CHECK(num_per_channels == 32, "num_per_channels must be 32, got ", num_per_channels);
+
+ // Check shape
+ KU_CHECK_SHAPE(q, s_q, h_q, d_qk);
+ KU_CHECK_SHAPE(indices, s_q, topk);
+ KU_CHECK_SHAPE(attn_sink, h_q);
+ KU_CHECK_SHAPE(topk_length, s_q);
+ KU_CHECK_SHAPE(extra_indices, s_q, extra_topk);
+ KU_CHECK_SHAPE(extra_topk_length, s_q);
+ KU_CHECK_SHAPE(token_positions, s_q);
+ KU_CHECK_SHAPE(cos_sin_cache, cos_sin_cache.size(0), rope_dim);
+
+ torch::stable::accelerator::DeviceGuard device_guard(q.get_device_index());
+
+ STD_TORCH_CHECK(d_v % (num_per_channels * 4) == 0); // 4 is the number of uint8 in uint32, since `use_packed_ue8m0` is `True`
+ Tensor out_fp8 = torch::stable::new_empty(q, {s_q, n_wv_group, wv_group_size * d_v}, ScalarType::Float8_e4m3fn);
+ uint32_t out_sf_scale_gran = 32; // Since the weight is per-32 scaled and deep_gemm.einsum requires A and B to have the same scale granularity, the output sf is always stored in a per-32 scaled format, although it will be actually per-128 scaled when num_per_channels is 128
+ Tensor out_sf = allocate_scale_factor(s_q, wv_group_size * d_v, out_sf_scale_gran, q, n_wv_group);
+ out_sf = torch::stable::transpose(out_sf, 0, 1); // [s_q, n_wv_group, wv_group_size * d_v / (out_sf_scale_gran*4)]
+ Tensor lse = torch::stable::new_empty(q, {s_q, h_q}, ScalarType::Float);
+ KU_CHECK_CONTIGUOUS(out_fp8);
+ STD_TORCH_CHECK(out_sf.stride(0) == 1);
+ KU_CHECK_CONTIGUOUS(lse);
+
+ SparseAttnDecodeParams base_params = {
+ 1 /* b */, s_q, h_q, h_kv, d_qk, d_v,
+ sm_scale, sm_scale * LOG_2_E,
+ num_blocks, page_block_size, topk,
+ model_type, extra_model_type,
+
+ (bf16*)q.data_ptr(),
+ (bf16*)kv.data_ptr(),
+ (int*)indices.data_ptr(),
+ ku::get_optional_tensor_ptr(topk_length),
+ ku::get_optional_tensor_ptr(attn_sink),
+ (float*)lse.data_ptr(),
+ nullptr, // `out` (bf16) is unused; the FP8 output goes to `out_fp8` + `out_sf`
+
+ extra_num_blocks, extra_page_block_size, extra_topk,
+ ku::get_optional_tensor_ptr(extra_kv),
+ ku::get_optional_tensor_ptr(extra_indices),
+ ku::get_optional_tensor_ptr(extra_topk_length),
+
+ 0, int64_stride_to_int(q.stride(0)), int64_stride_to_int(q.stride(1)), // stride_q_b is unused since b == 1
+ int64_stride_to_int(kv.stride(0)), int64_stride_to_int(kv.stride(1)),
+ 0, int64_stride_to_int(indices.stride(0)), // stride_indices_b is unused since b == 1
+ 0, int64_stride_to_int(lse.stride(0)), // stride_lse_b is unused since b == 1
+ 0, 0, 0, // stride_o_b, stride_o_s_q, stride_o_h_q: unused since `out` is unused
+
+ have_extra_kvcache ? int64_stride_to_int(extra_kv->stride(0)) : 0,
+ have_extra_kvcache ? int64_stride_to_int(extra_kv->stride(1)) : 0,
+ 0, // stride_extra_indices_b is unused since b == 1
+ have_extra_kvcache ? int64_stride_to_int(extra_indices->stride(0)) : 0,
+ get_current_cuda_stream(q),
+
+ false, // enable_split_kv: split-KV is not supported by this kernel
+ // The remaining split-KV related fields are zero-initialized
+ };
+
+ DecodeParams params = {
+ base_params,
+
+ enable_q_norm,
+ rms_norm_eps,
+ (uint32_t*)token_positions.data_ptr(),
+ is_rope_neox_style,
+ rope_dim,
+ (float*)cos_sin_cache.data_ptr(),
+
+ n_wv_group,
+ wv_group_size,
+ num_per_channels,
+ use_tma_aligned_col_major_sf,
+ round_sf,
+ use_packed_ue8m0,
+
+ (fp8_e4m3*)out_fp8.data_ptr(),
+ (uint32_t*)out_sf.data_ptr(),
+ (uint32_t)int64_stride_to_int(out_sf.stride(1)),
+ (uint32_t)int64_stride_to_int(out_sf.stride(2))
+ };
+
+ DISPATCH_NUM_HEADS(h_q, H_Q, ([&]() {
+ DISPATCH_BOOLEAN_FLAG(enable_q_norm, ENABLE_Q_NORM, ([&]() {
+ if (extra_model_type == ModelType::V41_FP4) {
+ sm100::prefill::fused_norm_rope_attn_rope_cast_fwd::core_attn::run_fused_norm_rope_attn_rope_cast_fwd_kernel(params);
+ } else if (model_type == ModelType::V4) {
+ sm100::prefill::fused_norm_rope_attn_rope_cast_fwd::core_attn::run_fused_norm_rope_attn_rope_cast_fwd_kernel(params);
+ } else if (model_type == ModelType::V41) {
+ sm100::prefill::fused_norm_rope_attn_rope_cast_fwd::core_attn::run_fused_norm_rope_attn_rope_cast_fwd_kernel(params);
+ } else {
+ STD_TORCH_CHECK(false, "Unsupported model_type: ", get_dynamic_enum_name(model_type));
+ }
+ }));
+ }));
+
+ return {out_fp8, out_sf, lse};
+}
+
+
+std::vector permute_q_b_proj(
+ const Tensor &q_b_proj,
+ const Tensor &scale_factors,
+ int64_t h_q,
+ int64_t d_q
+) {
+ KU_CHECK_NDIM(q_b_proj, 2);
+ KU_CHECK_NDIM(scale_factors, 2);
+
+ int h_q_d_q = h_q * d_q;
+ int q_lora_rank = q_b_proj.size(1);
+
+ int gran = q_lora_rank / (4 * scale_factors.size(1));
+ STD_TORCH_CHECK(gran == 32 || gran == 128, "gran must be 32 or 128, got ", gran);
+ STD_TORCH_CHECK(q_lora_rank % (gran * 4) == 0, "q_lora_rank must be divisible by gran * 4");
+
+ KU_CHECK_DEVICE(q_b_proj);
+ KU_CHECK_DEVICE(scale_factors);
+
+ KU_CHECK_DTYPE(q_b_proj, ScalarType::Float8_e4m3fn);
+ KU_CHECK_DTYPE(scale_factors, ScalarType::Int);
+
+ KU_CHECK_SHAPE(q_b_proj, h_q_d_q, q_lora_rank);
+ KU_CHECK_SHAPE(scale_factors, h_q_d_q, q_lora_rank / gran / 4);
+
+ KU_CHECK_LAST_DIM_CONTIGUOUS(q_b_proj);
+ STD_TORCH_CHECK(scale_factors.stride(0) == 1, "scale_factors must be contiguous on the first dimension");
+
+ torch::stable::accelerator::DeviceGuard device_guard(q_b_proj.get_device_index());
+
+ Tensor q_b_proj_permuted = torch::stable::empty_like(q_b_proj);
+ Tensor scale_factors_permuted = allocate_scale_factor(h_q_d_q, q_lora_rank, gran, q_b_proj);
+ KU_CHECK_CONTIGUOUS(q_b_proj_permuted);
+
+ sm100::prefill::fused_norm_rope_attn_rope_cast_fwd::permute_q_b_proj::Params params = {
+ (uint32_t)h_q,
+ (uint32_t)d_q,
+ (uint32_t)q_lora_rank,
+ (uint32_t)gran,
+
+ (fp8_e4m3*)q_b_proj.data_ptr(),
+ (uint64_t)q_b_proj.stride(0),
+ (int32_t*)scale_factors.data_ptr(),
+ (uint64_t)scale_factors.stride(1),
+
+ (fp8_e4m3*)q_b_proj_permuted.data_ptr(),
+ (uint64_t)q_b_proj_permuted.stride(0),
+ (int32_t*)scale_factors_permuted.data_ptr(),
+ (uint64_t)scale_factors_permuted.stride(1),
+
+ get_current_cuda_stream(q_b_proj),
+ };
+
+ sm100::prefill::fused_norm_rope_attn_rope_cast_fwd::permute_q_b_proj::run_permute_q_b_proj_kernel(params);
+
+ return {q_b_proj_permuted, scale_factors_permuted};
+}
+
+
+std::vector permute_wv_proj(
+ const Tensor &wv_proj,
+ const Tensor &scale_factors,
+ int64_t wv_group_size,
+ int64_t d_o
+) {
+ KU_CHECK_NDIM(wv_proj, 3);
+ KU_CHECK_NDIM(scale_factors, 3);
+
+ KU_CHECK_DEVICE(wv_proj);
+ KU_CHECK_DEVICE(scale_factors);
+
+ KU_CHECK_DTYPE(wv_proj, ScalarType::Float8_e4m3fn);
+ KU_CHECK_DTYPE(scale_factors, ScalarType::Int);
+
+ int n_wv_group = wv_proj.size(0);
+ int d_proj_out = wv_proj.size(1);
+ KU_CHECK_SHAPE(wv_proj, n_wv_group, d_proj_out, wv_group_size * d_o);
+
+ int input_gran = wv_group_size * d_o / (4 * scale_factors.size(2));
+ int output_gran = 32; // Fixed to 32, otherwise permution between chunk (which has 32 elements) will be impossible
+ STD_TORCH_CHECK(input_gran == 32, "input scale granularity must be 32, got ", input_gran);
+ STD_TORCH_CHECK((wv_group_size * d_o) % (input_gran * 4) == 0, "q_lora_rank must be divisible by gran * 4");
+ KU_CHECK_SHAPE(scale_factors, n_wv_group, d_proj_out, (wv_group_size * d_o) / input_gran / 4);
+
+ KU_CHECK_LAST_DIM_CONTIGUOUS(wv_proj);
+ STD_TORCH_CHECK(scale_factors.stride(1) == 1, "scale_factors must be contiguous on the second dimension");
+
+ torch::stable::accelerator::DeviceGuard device_guard(wv_proj.get_device_index());
+
+ Tensor wv_proj_permuted = torch::stable::empty_like(wv_proj);
+ Tensor scale_factors_permuted = allocate_scale_factor(d_proj_out, wv_group_size * d_o, output_gran, wv_proj, n_wv_group);
+ KU_CHECK_CONTIGUOUS(wv_proj_permuted);
+
+ sm100::prefill::fused_norm_rope_attn_rope_cast_fwd::permute_wv_proj::Params params = {
+ (uint32_t)d_o,
+ (uint32_t)input_gran,
+ (uint32_t)wv_group_size,
+ (uint32_t)n_wv_group,
+ (uint32_t)d_proj_out,
+
+ (fp8_e4m3*)wv_proj.data_ptr(),
+ (uint64_t)wv_proj.stride(0),
+ (uint64_t)wv_proj.stride(1),
+ (int32_t*)scale_factors.data_ptr(),
+ (uint64_t)scale_factors.stride(0),
+ (uint64_t)scale_factors.stride(2),
+
+ (fp8_e4m3*)wv_proj_permuted.data_ptr(),
+ (uint64_t)wv_proj_permuted.stride(0),
+ (uint64_t)wv_proj_permuted.stride(1),
+ (int32_t*)scale_factors_permuted.data_ptr(),
+ (uint64_t)scale_factors_permuted.stride(0),
+ (uint64_t)scale_factors_permuted.stride(2),
+
+ get_current_cuda_stream(wv_proj),
+ };
+
+ sm100::prefill::fused_norm_rope_attn_rope_cast_fwd::permute_wv_proj::run_permute_wv_proj_kernel(params);
+
+ return {wv_proj_permuted, scale_factors_permuted};
+}
diff --git a/csrc/api/interfaces.h b/csrc/api/interfaces.h
new file mode 100644
index 00000000..e712c500
--- /dev/null
+++ b/csrc/api/interfaces.h
@@ -0,0 +1,102 @@
+#pragma once
+
+#include
+#include
+#include
+
+#include "common.h"
+#include "kernels/sm100/prefill/dense/interface.h"
+
+std::vector sparse_attn_prefill_interface(
+ const Tensor &q,
+ const Tensor &kv,
+ const Tensor &indices,
+ double sm_scale,
+ int64_t d_v,
+ const std::optional &attn_sink,
+ const std::optional &topk_length,
+ const std::optional &out_);
+
+std::tuple, std::optional>
+sparse_attn_decode_interface(
+ const Tensor &q,
+ const Tensor &kv,
+ const Tensor &indices,
+ const std::optional &topk_length,
+ const std::optional &attn_sink,
+ std::optional tile_scheduler_metadata,
+ std::optional num_splits,
+ const std::optional &extra_kv,
+ const std::optional &extra_indices,
+ const std::optional &extra_topk_length,
+ int64_t d_v,
+ double sm_scale,
+ const std::optional &out_);
+
+std::tuple, std::optional>
+dense_attn_decode_interface(
+ Tensor q,
+ const Tensor &kcache,
+ int64_t head_size_v,
+ const Tensor &seqlens_k,
+ const Tensor &block_table,
+ double softmax_scale,
+ bool is_causal,
+ std::optional tile_scheduler_metadata,
+ std::optional num_splits,
+ const std::optional &out_);
+
+std::vector fused_norm_rope_attn_rope_cast_fwd(
+ const Tensor &q,
+ const Tensor &kv,
+ const Tensor &indices,
+ double sm_scale,
+ int64_t d_v,
+ const std::optional &attn_sink,
+ const std::optional &topk_length,
+ bool enable_q_norm,
+ double rms_norm_eps,
+ const Tensor &token_positions,
+ bool is_rope_neox_style,
+ int64_t rope_dim,
+ const Tensor &cos_sin_cache,
+ int64_t n_wv_group,
+ int64_t num_per_channels,
+ bool use_tma_aligned_col_major_sf,
+ bool round_sf,
+ bool use_packed_ue8m0);
+
+std::vector fused_norm_rope_attn_rope_cast_decode(
+ const Tensor &q,
+ const Tensor &kv,
+ const Tensor &indices,
+ double sm_scale,
+ int64_t d_v,
+ const std::optional &attn_sink,
+ const std::optional &topk_length,
+ const std::optional &extra_kv,
+ const std::optional &extra_indices,
+ const std::optional &extra_topk_length,
+ bool enable_q_norm,
+ double rms_norm_eps,
+ const Tensor &token_positions,
+ bool is_rope_neox_style,
+ int64_t rope_dim,
+ const Tensor &cos_sin_cache,
+ int64_t n_wv_group,
+ int64_t num_per_channels,
+ bool use_tma_aligned_col_major_sf,
+ bool round_sf,
+ bool use_packed_ue8m0);
+
+std::vector permute_q_b_proj(
+ const Tensor &q_b_proj,
+ const Tensor &scale_factors,
+ int64_t h_q,
+ int64_t d_q);
+
+std::vector permute_wv_proj(
+ const Tensor &wv_proj,
+ const Tensor &scale_factors,
+ int64_t wv_group_size,
+ int64_t d_o);
diff --git a/csrc/api/sparse_decode.h b/csrc/api/sparse_decode.cpp
similarity index 56%
rename from csrc/api/sparse_decode.h
rename to csrc/api/sparse_decode.cpp
index cb2c150e..ad3d75a8 100644
--- a/csrc/api/sparse_decode.h
+++ b/csrc/api/sparse_decode.cpp
@@ -1,14 +1,22 @@
-#pragma once
-
#include "common.h"
-#include "params.h"
+#include "kernels/params.h"
-#include "sm90/decode/sparse_fp8/splitkv_mla.h"
-#include "sm100/decode/head64/kernel.h"
-#include "sm100/prefill/sparse/fwd_for_small_topk/head128/phase1.h"
-#include "smxx/decode/get_decoding_sched_meta/get_decoding_sched_meta.h"
-#include "smxx/decode/combine/combine.h"
+#include "kernels/sm90/decode/sparse/splitkv_mla.h"
+#include "kernels/sm100/decode/sparse/head64/kernel.h"
+#include "kernels/sm100/decode/sparse/nvfp4_head64/kernel.h"
+#include "kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/phase1.h"
+#include "kernels/smxx/decode/get_decoding_sched_meta/get_decoding_sched_meta.h"
+#include "kernels/smxx/decode/combine/combine.h"
+
+template
+static constexpr SparseAttnFwdMode get_decode_fwd_mode() {
+ if constexpr (ENABLE_SPLIT_KV) {
+ return SparseAttnFwdMode::DecodeWithSplitKV;
+ } else {
+ return SparseAttnFwdMode::Decode;
+ }
+}
// Feature set of sparse decoding kernels
enum class DecodeFeatures : int {
@@ -19,15 +27,15 @@ enum class DecodeFeatures : int {
HEAD_DIM_512,
V32_KVCACHE_FORMAT,
- MODEL1_KVCACHE_FORMAT,
+ V4_KVCACHE_FORMAT,
+ V41_KVCACHE_FORMAT,
+ V41_FP4_KVCACHE_FORMAT,
+ NVFP4_FP8ROPE_KVCACHE_FORMAT,
ATTN_SINK,
TOPK_LENGTH,
EXTRA_KVCACHE,
- EXTRA_TOPK_LENGTH,
-
- // NVFP4 (e2m1 + per-16 e4m3 scales) NoPE with e4m3 RoPE. SM100-only.
- NVFP4_FP8ROPE_KVCACHE_FORMAT
+ EXTRA_TOPK_LENGTH
};
struct DecodeImplMeta {
@@ -51,7 +59,7 @@ class Decode_Sm90_Impl : public DecodeImplBase {
DecodeFeatures::HEAD_DIM_512,
DecodeFeatures::HEAD_DIM_576,
DecodeFeatures::V32_KVCACHE_FORMAT,
- DecodeFeatures::MODEL1_KVCACHE_FORMAT,
+ DecodeFeatures::V4_KVCACHE_FORMAT,
DecodeFeatures::ATTN_SINK,
DecodeFeatures::TOPK_LENGTH,
DecodeFeatures::EXTRA_KVCACHE,
@@ -72,7 +80,7 @@ class Decode_Sm90_Impl : public DecodeImplBase {
void run_(const SparseAttnDecodeParams ¶ms, const std::vector &required_features) override {
DISPATCH_MODEL_TYPE(params.model_type, MODEL_TYPE, [&]() {
DISPATCH_NUM_HEADS(params.h_q, NUM_HEADS, [&]() {
- sm90::decode::sparse_fp8::run_flash_splitkv_mla_fp8_sparse_kernel(params);
+ sm90::decode::sparse::run_flash_splitkv_mla_fp8_sparse_kernel(params);
});
});
}
@@ -84,13 +92,17 @@ class Decode_Sm100_Head64_Impl : public DecodeImplBase {
DecodeFeatures::HEAD_DIM_512,
DecodeFeatures::HEAD_DIM_576,
DecodeFeatures::V32_KVCACHE_FORMAT,
- DecodeFeatures::MODEL1_KVCACHE_FORMAT,
+ DecodeFeatures::V4_KVCACHE_FORMAT,
+ DecodeFeatures::V41_KVCACHE_FORMAT,
+ DecodeFeatures::V41_FP4_KVCACHE_FORMAT,
DecodeFeatures::NVFP4_FP8ROPE_KVCACHE_FORMAT,
DecodeFeatures::ATTN_SINK,
DecodeFeatures::TOPK_LENGTH,
DecodeFeatures::EXTRA_KVCACHE,
DecodeFeatures::EXTRA_TOPK_LENGTH
)
+ using SupportedKVFormats = KVFormatPairs, KVFormatPair, KVFormatPair,
+ KVFormatPair>;
public:
DecodeImplMeta get_meta(int h_q, int s_q) override {
@@ -104,8 +116,19 @@ class Decode_Sm100_Head64_Impl : public DecodeImplBase {
protected:
void run_(const SparseAttnDecodeParams ¶ms, const std::vector &required_features) override {
- DISPATCH_MODEL_TYPE_SM100(params.model_type, MODEL_TYPE, [&]() {
- sm100::decode::head64::run_flash_splitkv_mla_fp8_sparse_kernel(params);
+ if (params.model_type == ModelType::V32_NVFP4_FP8ROPE) {
+ STD_TORCH_CHECK(params.extra_model_type == ModelType::V32_NVFP4_FP8ROPE,
+ "NVFP4 does not support a mixed extra KV-cache format");
+ STD_TORCH_CHECK(params.enable_split_kv, "NVFP4 requires split-KV scheduling");
+ sm100::decode::head64::run_flash_splitkv_mla_fp8_sparse_kernel(params);
+ return;
+ }
+ dispatch_kv_formats(SupportedKVFormats{}, params.model_type, params.extra_model_type, [&]() {
+ DISPATCH_BOOLEAN_FLAG(params.enable_split_kv, ENABLE_SPLIT_KV, ([&]() {
+ STD_TORCH_CHECK(params.h_q == 64, "Unsupported h_q: ", params.h_q);
+ using sm100::decode::sparse::head64::Config;
+ sm100::decode::sparse::head64::run_flash_splitkv_mla_fp8_sparse_kernel(params);
+ }));
});
}
};
@@ -119,13 +142,14 @@ class Decode_Sm100_Head64x2_Impl : public DecodeImplBase {
DecodeFeatures::HEAD_DIM_512,
DecodeFeatures::HEAD_DIM_576,
DecodeFeatures::V32_KVCACHE_FORMAT,
- DecodeFeatures::MODEL1_KVCACHE_FORMAT,
+ DecodeFeatures::V4_KVCACHE_FORMAT,
DecodeFeatures::NVFP4_FP8ROPE_KVCACHE_FORMAT,
DecodeFeatures::ATTN_SINK,
DecodeFeatures::TOPK_LENGTH,
DecodeFeatures::EXTRA_KVCACHE,
DecodeFeatures::EXTRA_TOPK_LENGTH
)
+ using SupportedKVFormats = KVFormatPairs, KVFormatPair>;
public:
DecodeImplMeta get_meta(int h_q, int s_q) override {
@@ -139,7 +163,10 @@ class Decode_Sm100_Head64x2_Impl : public DecodeImplBase {
protected:
void run_(const SparseAttnDecodeParams ¶ms, const std::vector &required_features) override {
- DISPATCH_MODEL_TYPE_SM100(params.model_type, MODEL_TYPE, [&]() {
+ if (params.model_type == ModelType::V32_NVFP4_FP8ROPE) {
+ STD_TORCH_CHECK(params.extra_model_type == ModelType::V32_NVFP4_FP8ROPE,
+ "NVFP4 does not support a mixed extra KV-cache format");
+ STD_TORCH_CHECK(params.enable_split_kv, "NVFP4 requires split-KV scheduling");
for (int start_head_idx = 0; start_head_idx < 128; start_head_idx += 64) {
SparseAttnDecodeParams cur_params = params;
cur_params.q += start_head_idx * params.stride_q_h_q;
@@ -148,15 +175,32 @@ class Decode_Sm100_Head64x2_Impl : public DecodeImplBase {
}
cur_params.lse += start_head_idx;
cur_params.out += start_head_idx * params.stride_o_h_q;
- if (cur_params.lse_accum) {
- cur_params.lse_accum += start_head_idx;
- }
- if (cur_params.o_accum) {
- cur_params.o_accum += start_head_idx * params.stride_o_accum_h_q;
- }
+ cur_params.lse_accum += start_head_idx;
+ cur_params.o_accum += start_head_idx * params.stride_o_accum_h_q;
cur_params.h_q = 64;
- sm100::decode::head64::run_flash_splitkv_mla_fp8_sparse_kernel(cur_params);
+ sm100::decode::head64::run_flash_splitkv_mla_fp8_sparse_kernel(cur_params);
}
+ return;
+ }
+ dispatch_kv_formats(SupportedKVFormats{}, params.model_type, params.extra_model_type, [&]() {
+ DISPATCH_BOOLEAN_FLAG(params.enable_split_kv, ENABLE_SPLIT_KV, ([&]() {
+ for (int start_head_idx = 0; start_head_idx < 128; start_head_idx += 64) {
+ SparseAttnDecodeParams cur_params = params;
+ cur_params.q += start_head_idx * params.stride_q_h_q;
+ if (cur_params.attn_sink) {
+ cur_params.attn_sink += start_head_idx;
+ }
+ cur_params.lse += start_head_idx;
+ cur_params.out += start_head_idx * params.stride_o_h_q;
+ if (cur_params.enable_split_kv) {
+ cur_params.lse_accum += start_head_idx;
+ cur_params.o_accum += start_head_idx * params.stride_o_accum_h_q;
+ }
+ cur_params.h_q = 64;
+ using sm100::decode::sparse::head64::Config;
+ sm100::decode::sparse::head64::run_flash_splitkv_mla_fp8_sparse_kernel(cur_params);
+ }
+ }));
});
}
};
@@ -166,12 +210,16 @@ class Decode_Sm100_Head128_Impl : public DecodeImplBase {
DECLARE_SUPPORTED_FEATURES(
DecodeFeatures::HEAD_128,
DecodeFeatures::HEAD_DIM_512,
- DecodeFeatures::MODEL1_KVCACHE_FORMAT,
+ DecodeFeatures::V4_KVCACHE_FORMAT,
+ DecodeFeatures::V41_KVCACHE_FORMAT,
+ DecodeFeatures::V41_FP4_KVCACHE_FORMAT,
DecodeFeatures::ATTN_SINK,
DecodeFeatures::TOPK_LENGTH,
DecodeFeatures::EXTRA_KVCACHE,
DecodeFeatures::EXTRA_TOPK_LENGTH
)
+ using SupportedKVFormats = KVFormatPairs, KVFormatPair,
+ KVFormatPair>;
public:
DecodeImplMeta get_meta(int h_q, int s_q) override {
@@ -185,11 +233,20 @@ class Decode_Sm100_Head128_Impl : public DecodeImplBase {
protected:
void run_(const SparseAttnDecodeParams ¶ms, const std::vector &required_features) override {
- sm100::fwd_for_small_topk::head128::run_fwd_for_small_topk_phase1_kernel(params);
+ SparseAttnDecodeParams hotfixed_params = params;
+ if (params.s_q == 1 && params.b > 1) {
+ // For this kernel, we require `params.stride_q_b % params.stride_q_s_q == 0`, since we "squeeze" the batch size dimention and the sequence length q dimension together during tensormap creation
+ hotfixed_params.stride_q_s_q = hotfixed_params.stride_q_b;
+ }
+ dispatch_kv_formats(SupportedKVFormats{}, params.model_type, params.extra_model_type, [&]() {
+ DISPATCH_BOOLEAN_FLAG(params.enable_split_kv, ENABLE_SPLIT_KV, ([&]() {
+ sm100::prefill::sparse_fwd_for_small_topk::head128::run_sparse_fwd_for_small_topk_phase1_kernel(), 512, MODEL_TYPE, EXTRA_MODEL_TYPE>(hotfixed_params);
+ }));
+ });
}
};
-static std::tuple, std::optional>
+std::tuple, std::optional>
sparse_attn_decode_interface(
const Tensor &q, // [b, s_q, h_q, d_qk]
const Tensor &kv, // [num_blocks, page_block_size, h_k, d_qk]
@@ -237,6 +294,9 @@ sparse_attn_decode_interface(
extra_topk = extra_indices->size(-1);
}
+ // Split-KV only pays off when a request has enough work. The sm90 kernel always splits.
+ bool enable_split_kv = arch.is_sm90a() || !(topk + extra_topk <= 640);
+
// metadata sanity check
STD_TORCH_CHECK(b > 0);
STD_TORCH_CHECK(s_q > 0);
@@ -295,39 +355,27 @@ sparse_attn_decode_interface(
// Check shape
KU_CHECK_SHAPE(q, b, s_q, h_q, d_qk);
- ModelType model_type;
- {
- // Infer the quantized KV cache format from the KV cache's bytes-per-token
- // (i.e. its last dim)
- const int bytes_per_token = static_cast(kv.size(3));
- if (d_qk == 576 && d_v == 512) {
- if (bytes_per_token == 512 + 64*2 + (512/128)*4) {
- // V3.2 style, 656B/token: 512B e4m3 NoPE | 128B bf16 RoPE | 16B fp32 NoPE SF
- model_type = ModelType::V32;
- } else if (bytes_per_token == 352) {
- // NVFP4 NoPE + FP8 RoPE, 352B/token: 256B e2m1 NoPE | 64B e4m3 RoPE (unscaled) | 32B e4m3 NoPE SF
- // Must match KernelTemplate::BYTES_PER_TOKEN
- model_type = ModelType::V32_NVFP4_FP8ROPE;
- } else {
- STD_TORCH_CHECK(false, "Cannot infer the sparse KV cache format: with d_qk == ", d_qk, " and d_v == ", d_v,
- ", kv.size(-1) (bytes per token) must be 656 (fp8) or 352 (nvfp4 nope + fp8 rope), but got ", bytes_per_token);
- }
- } else if (d_qk == 512 && d_v == 512) {
- if (bytes_per_token == 448 + 64*2 + (448/64)*1 + 1) {
- // MODEL1 style, 584B/token
- model_type = ModelType::MODEL1;
- } else {
- STD_TORCH_CHECK(false, "Cannot infer the sparse KV cache format: with d_qk == ", d_qk, " and d_v == ", d_v,
- ", kv.size(-1) (bytes per token) must be 584 (fp8), but got ", bytes_per_token);
- }
- } else {
- STD_TORCH_CHECK(false, "Unsupported head sizes for sparse decoding: d_qk == ", d_qk, ", d_v == ", d_v);
- }
- KU_CHECK_SHAPE(extra_kv, extra_num_blocks, extra_page_block_size, h_kv, bytes_per_token);
- STD_TORCH_CHECK(kv.stride(1) == bytes_per_token, "The whole block must be contiguous when is_fp8_cache is True for kv cache");
- if (extra_kv.has_value()) {
- STD_TORCH_CHECK(extra_kv->stride(1) == bytes_per_token, "The whole block must be contiguous when is_fp8_cache is True for extra kv cache");
- }
+ // The formats of `kv` and `extra_kv`
+ ModelType model_type, extra_model_type;
+ if (d_qk == 576 && d_v == 512) {
+ model_type = detect_kv_cache_format_for_headdim_576(kv.size(3));
+ extra_model_type = have_extra_kcache ? detect_kv_cache_format_for_headdim_576(extra_kv->size(3)) : model_type;
+ } else if (d_qk == 512 && d_v == 512) {
+ model_type = detect_kv_cache_format_for_headdim_512(kv.size(3));
+ extra_model_type = have_extra_kcache ? detect_kv_cache_format_for_headdim_512(extra_kv->size(3)) : model_type;
+ } else {
+ STD_TORCH_CHECK(false, "Unsupported head sizes for is_fp8_kvcache == True");
+ }
+ STD_TORCH_CHECK(model_type != ModelType::V41_FP4, "The fp4 KV cache is only supported as extra_kv");
+ STD_TORCH_CHECK(is_valid_kv_format_pair(model_type, extra_model_type), "invalid kv format pair, ", get_dynamic_enum_name(model_type), " and ", get_dynamic_enum_name(extra_model_type));
+ // The preserved NVFP4 kernel predates the no-split path and consumes the
+ // split scheduler metadata even for small top-k values.
+ enable_split_kv = enable_split_kv || model_type == ModelType::V32_NVFP4_FP8ROPE;
+ KU_CHECK_SHAPE(kv, num_blocks, page_block_size, h_kv, kv_cache_bytes_per_token(model_type));
+ KU_CHECK_SHAPE(extra_kv, extra_num_blocks, extra_page_block_size, h_kv, kv_cache_bytes_per_token(extra_model_type));
+ STD_TORCH_CHECK(kv.stride(1) == kv_cache_bytes_per_token(model_type), "The whole block must be contiguous when is_fp8_cache is True for kv cache");
+ if (have_extra_kcache) {
+ STD_TORCH_CHECK(extra_kv->stride(1) == kv_cache_bytes_per_token(extra_model_type), "The whole block must be contiguous when is_fp8_cache is True for extra kv cache");
}
KU_CHECK_SHAPE(indices, b, s_q, topk);
KU_CHECK_SHAPE(topk_length, b);
@@ -348,7 +396,6 @@ sparse_attn_decode_interface(
out = torch::stable::new_empty(q, {b, s_q, h_q, d_v});
}
Tensor lse = torch::stable::new_empty(q, {b, s_q, h_q}, ScalarType::Float);
-
std::vector features;
if (h_q == 64) {
features.push_back(DecodeFeatures::HEAD_64);
@@ -364,15 +411,6 @@ sparse_attn_decode_interface(
} else {
STD_TORCH_CHECK(false, "Unsupported d_qk: ", d_qk);
}
- if (model_type == ModelType::V32) {
- features.push_back(DecodeFeatures::V32_KVCACHE_FORMAT);
- } else if (model_type == ModelType::MODEL1) {
- features.push_back(DecodeFeatures::MODEL1_KVCACHE_FORMAT);
- } else if (model_type == ModelType::V32_NVFP4_FP8ROPE) {
- features.push_back(DecodeFeatures::NVFP4_FP8ROPE_KVCACHE_FORMAT);
- } else {
- STD_TORCH_CHECK(false, "Unsupported model type: ", (int)model_type);
- }
if (have_attn_sink) {
features.push_back(DecodeFeatures::ATTN_SINK);
}
@@ -385,6 +423,21 @@ sparse_attn_decode_interface(
if (have_extra_topk_length) {
features.push_back(DecodeFeatures::EXTRA_TOPK_LENGTH);
}
+ for (ModelType mt : {model_type, extra_model_type}) {
+ if (mt == ModelType::V32) {
+ features.push_back(DecodeFeatures::V32_KVCACHE_FORMAT);
+ } else if (mt == ModelType::V4) {
+ features.push_back(DecodeFeatures::V4_KVCACHE_FORMAT);
+ } else if (mt == ModelType::V41) {
+ features.push_back(DecodeFeatures::V41_KVCACHE_FORMAT);
+ } else if (mt == ModelType::V41_FP4) {
+ features.push_back(DecodeFeatures::V41_FP4_KVCACHE_FORMAT);
+ } else if (mt == ModelType::V32_NVFP4_FP8ROPE) {
+ features.push_back(DecodeFeatures::NVFP4_FP8ROPE_KVCACHE_FORMAT);
+ } else {
+ STD_TORCH_CHECK(false, "Unsupported model type: ", (int)mt);
+ }
+ }
DecodeImplBase* impl;
if (arch.is_sm100f()) {
@@ -413,7 +466,7 @@ sparse_attn_decode_interface(
b, s_q, h_q, h_kv, d_qk, d_v,
sm_scale, sm_scale * LOG_2_E,
num_blocks, page_block_size, topk,
- model_type,
+ model_type, extra_model_type,
(bf16*)q.data_ptr(),
(bf16*)kv.data_ptr(),
@@ -438,49 +491,50 @@ sparse_attn_decode_interface(
have_extra_kcache ? int64_stride_to_int(extra_kv->stride(1)) : 0,
have_extra_kcache ? int64_stride_to_int(extra_indices->stride(0)) : 0,
have_extra_kcache ? int64_stride_to_int(extra_indices->stride(1)) : 0,
- get_current_cuda_stream(q)
+ get_current_cuda_stream(q),
+
+ enable_split_kv,
};
- // Get MLA metadata if necessary
Tensor o_accum, lse_accum;
- if (!tile_scheduler_metadata.has_value()) {
- tile_scheduler_metadata = torch::stable::new_empty(q, {impl_meta.num_sm_parts, sizeof(DecodingSchedMeta)/4}, ScalarType::Int);
- num_splits = torch::stable::new_empty(q, {b+1}, ScalarType::Int);
+ if (enable_split_kv) {
+ // Get MLA metadata if necessary
+ if (!tile_scheduler_metadata.has_value()) {
+ tile_scheduler_metadata = torch::stable::new_empty(q, {impl_meta.num_sm_parts, sizeof(DecodingSchedMeta)/4}, ScalarType::Int);
+ num_splits = torch::stable::new_empty(q, {b+1}, ScalarType::Int);
+ KU_CHECK_CONTIGUOUS(tile_scheduler_metadata);
+ KU_CHECK_CONTIGUOUS(num_splits);
+
+ GetDecodeSchedMetaParams get_sched_meta_params = {
+ b, s_q,
+ impl_meta.block_size_topk,
+ impl_meta.fixed_overhead_num_blocks,
+ topk,
+ extra_topk,
+ ku::get_optional_tensor_ptr(topk_length),
+ ku::get_optional_tensor_ptr(extra_topk_length),
+ nullptr,
+ (DecodingSchedMeta*)tile_scheduler_metadata->data_ptr(),
+ num_splits->mutable_data_ptr(),
+ impl_meta.num_sm_parts,
+ get_current_cuda_stream(q)
+ };
+ smxx::decode::run_get_decoding_sched_meta_kernel(get_sched_meta_params);
+ }
+ KU_CHECK_DEVICE(tile_scheduler_metadata);
+ KU_CHECK_DEVICE(num_splits);
+ KU_CHECK_DTYPE(tile_scheduler_metadata, ScalarType::Int);
+ KU_CHECK_DTYPE(num_splits, ScalarType::Int);
KU_CHECK_CONTIGUOUS(tile_scheduler_metadata);
KU_CHECK_CONTIGUOUS(num_splits);
-
- GetDecodeSchedMetaParams get_sched_meta_params = {
- b, s_q,
- impl_meta.block_size_topk,
- impl_meta.fixed_overhead_num_blocks,
- topk,
- extra_topk,
- ku::get_optional_tensor_ptr(topk_length),
- ku::get_optional_tensor_ptr(extra_topk_length),
- nullptr,
- (DecodingSchedMeta*)tile_scheduler_metadata->data_ptr(),
- num_splits->mutable_data_ptr(),
- impl_meta.num_sm_parts,
- get_current_cuda_stream(q)
- };
- smxx::decode::run_get_decoding_sched_meta_kernel(get_sched_meta_params);
- }
- // Stick the metadata pointers to `params`
- KU_CHECK_DEVICE(tile_scheduler_metadata);
- KU_CHECK_DEVICE(num_splits);
- KU_CHECK_DTYPE(tile_scheduler_metadata, ScalarType::Int);
- KU_CHECK_DTYPE(num_splits, ScalarType::Int);
- KU_CHECK_CONTIGUOUS(tile_scheduler_metadata);
- KU_CHECK_CONTIGUOUS(num_splits);
- KU_CHECK_SHAPE(tile_scheduler_metadata, impl_meta.num_sm_parts, sizeof(DecodingSchedMeta)/sizeof(int));
- KU_CHECK_SHAPE(num_splits, b+1);
- params.tile_scheduler_metadata_ptr = (DecodingSchedMeta*)tile_scheduler_metadata->data_ptr();
- params.num_splits_ptr = num_splits->mutable_data_ptr();
- params.num_sm_parts = impl_meta.num_sm_parts;
-
- const bool needs_split_workspace = impl_meta.num_sm_parts > 1;
- if (needs_split_workspace) {
- const int total_num_splits = b + impl_meta.num_sm_parts;
+ KU_CHECK_SHAPE(tile_scheduler_metadata, impl_meta.num_sm_parts, sizeof(DecodingSchedMeta)/4);
+ KU_CHECK_SHAPE(num_splits, b+1);
+ // Stick the metadata pointers to `params`
+ params.tile_scheduler_metadata_ptr = (DecodingSchedMeta*)tile_scheduler_metadata->data_ptr();
+ params.num_splits_ptr = num_splits->mutable_data_ptr();
+ params.num_sm_parts = impl_meta.num_sm_parts;
+ // Allocate intermediate buffers for split-KV
+ const int total_num_splits = b + params.num_sm_parts;
lse_accum = torch::stable::new_empty(q, {total_num_splits, s_q, h_q}, ScalarType::Float);
o_accum = torch::stable::new_empty(q, {total_num_splits, s_q, h_q, d_v}, ScalarType::Float);
KU_CHECK_CONTIGUOUS(lse_accum);
@@ -495,33 +549,29 @@ sparse_attn_decode_interface(
}
impl->run(params, features);
+ if (enable_split_kv) {
+ CombineParams combine_params = {
+ b, s_q, h_q, d_v,
- if (!needs_split_workspace) {
- delete impl;
- return {out, torch::stable::transpose(lse, 1, 2), tile_scheduler_metadata, num_splits};
- }
-
- CombineParams combine_params = {
- b, s_q, h_q, d_v,
-
- params.lse,
- params.out,
- params.stride_lse_b, params.stride_lse_s_q,
- params.stride_o_b, params.stride_o_s_q, params.stride_o_h_q,
+ params.lse,
+ params.out,
+ params.stride_lse_b, params.stride_lse_s_q,
+ params.stride_o_b, params.stride_o_s_q, params.stride_o_h_q,
- params.lse_accum,
- params.o_accum,
- params.stride_lse_accum_split, params.stride_lse_accum_s_q,
- params.stride_o_accum_split, params.stride_o_accum_s_q, params.stride_o_accum_h_q,
+ params.lse_accum,
+ params.o_accum,
+ params.stride_lse_accum_split, params.stride_lse_accum_s_q,
+ params.stride_o_accum_split, params.stride_o_accum_s_q, params.stride_o_accum_h_q,
- params.tile_scheduler_metadata_ptr,
- params.num_splits_ptr,
- params.num_sm_parts,
+ params.tile_scheduler_metadata_ptr,
+ params.num_splits_ptr,
+ params.num_sm_parts,
- ku::get_optional_tensor_ptr(attn_sink),
- get_current_cuda_stream(q)
- };
- smxx::decode::run_flash_mla_combine_kernel(combine_params);
+ ku::get_optional_tensor_ptr(attn_sink),
+ get_current_cuda_stream(q)
+ };
+ smxx::decode::run_flash_mla_combine_kernel(combine_params);
+ }
delete impl;
diff --git a/csrc/api/sparse_fwd.h b/csrc/api/sparse_prefill.cpp
similarity index 89%
rename from csrc/api/sparse_fwd.h
rename to csrc/api/sparse_prefill.cpp
index 4e6030c5..1eb72311 100644
--- a/csrc/api/sparse_fwd.h
+++ b/csrc/api/sparse_prefill.cpp
@@ -1,13 +1,11 @@
-#pragma once
-
#include "common.h"
-#include "params.h"
+#include "kernels/params.h"
-#include "sm90/prefill/sparse/phase1.h"
-#include "sm100/prefill/sparse/fwd/head128/phase1.h"
-#include "sm100/prefill/sparse/fwd/head64/phase1.h"
-#include "sm100/prefill/sparse/fwd_for_small_topk/head128/phase1.h"
+#include "kernels/sm90/prefill/sparse/phase1.h"
+#include "kernels/sm100/prefill/sparse/fwd/head128/phase1.h"
+#include "kernels/sm100/prefill/sparse/fwd/head64/phase1.h"
+#include "kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/phase1.h"
enum class FwdFeatures : int {
HEAD_64,
@@ -17,7 +15,6 @@ enum class FwdFeatures : int {
HEAD_DIM_512,
ATTN_SINK,
- SINK_LSE,
TOPK_LENGTH
};
@@ -33,7 +30,6 @@ class Fwd_Sm90_Impl : public FwdImplBase {
FwdFeatures::HEAD_DIM_512,
FwdFeatures::HEAD_DIM_576,
FwdFeatures::ATTN_SINK,
- FwdFeatures::SINK_LSE,
FwdFeatures::TOPK_LENGTH
)
@@ -41,7 +37,7 @@ class Fwd_Sm90_Impl : public FwdImplBase {
void run_(const SparseAttnFwdParams ¶ms, const std::vector &required_features) override {
DISPATCH_HEAD_DIM(params.d_qk, HEAD_DIM_QK, [&]() {
DISPATCH_BOOLEAN_FLAG(params.topk_length != nullptr, HAVE_TOPK_LENGTH, [&]() {
- sm90::fwd::run_fwd_phase1_kernel(params);
+ sm90::prefill::sparse_fwd::run_fwd_phase1_kernel(params);
});
});
}
@@ -53,14 +49,14 @@ class Fwd_Sm100_Head64_Impl : public FwdImplBase {
FwdFeatures::HEAD_DIM_512,
FwdFeatures::HEAD_DIM_576,
FwdFeatures::ATTN_SINK,
- FwdFeatures::SINK_LSE,
FwdFeatures::TOPK_LENGTH
)
protected:
void run_(const SparseAttnFwdParams ¶ms, const std::vector &required_features) override {
DISPATCH_HEAD_DIM(params.d_qk, HEAD_DIM_QK, [&]() {
- sm100::fwd::head64::run_fwd_phase1_kernel(params);
+ STD_TORCH_CHECK(params.h_q == 64, "Unsupported h_q: ", params.h_q);
+ sm100::prefill::sparse_fwd::head64::run_sparse_fwd_phase1_kernel(params);
});
}
};
@@ -71,14 +67,13 @@ class Fwd_Sm100_Head128_Impl : public FwdImplBase {
FwdFeatures::HEAD_DIM_512,
FwdFeatures::HEAD_DIM_576,
FwdFeatures::ATTN_SINK,
- FwdFeatures::SINK_LSE,
FwdFeatures::TOPK_LENGTH
)
protected:
void run_(const SparseAttnFwdParams ¶ms, const std::vector &required_features) override {
DISPATCH_HEAD_DIM(params.d_qk, HEAD_DIM_QK, [&]() {
- sm100::fwd::head128::run_fwd_phase1_kernel(params);
+ sm100::prefill::sparse_fwd::head128::run_sparse_fwd_phase1_kernel(params);
});
}
};
@@ -88,17 +83,16 @@ class Fwd_Sm100_Head128_Small_TopK_Impl : public FwdImplBase {
FwdFeatures::HEAD_128,
FwdFeatures::HEAD_DIM_512,
FwdFeatures::ATTN_SINK,
- FwdFeatures::SINK_LSE,
FwdFeatures::TOPK_LENGTH
)
protected:
void run_(const SparseAttnFwdParams ¶ms, const std::vector &required_features) override {
- sm100::fwd_for_small_topk::head128::run_fwd_for_small_topk_phase1_kernel(params);
+ sm100::prefill::sparse_fwd_for_small_topk::head128::run_sparse_fwd_for_small_topk_phase1_kernel(params);
}
};
-static std::vector sparse_attn_prefill_interface(
+std::vector sparse_attn_prefill_interface(
const Tensor &q,
const Tensor &kv,
const Tensor &indices,
diff --git a/csrc/defines.h b/csrc/kernels/defines.h
similarity index 100%
rename from csrc/defines.h
rename to csrc/kernels/defines.h
diff --git a/csrc/kernels/kv_cache_format.h b/csrc/kernels/kv_cache_format.h
new file mode 100644
index 00000000..8c81552d
--- /dev/null
+++ b/csrc/kernels/kv_cache_format.h
@@ -0,0 +1,61 @@
+#pragma once
+
+#include "params.h"
+
+// Layout of paged quantized KV cache. A page block holds page_block_size tokens as two byte arrays, one row per token:
+// [page_block_size, TMA_K_STRIDE] token data: D_FP4 / 2 + D_FP8 bytes of quantized values,
+// then the D_BF16 bf16 (RoPE) values of V3.2 / V4
+// [page_block_size, NUM_SCALES_EACH_TOKEN * SCALE_BYTES] scales
+// Exception: V3.2 has no scale array; its 4 fp32 scales sit inside the token data, between the NoPE and the RoPE part.
+// Reference quantizer / dequantizer: tests/quant.py
+template
+struct KVCacheFormat {
+ static constexpr ModelType MODEL_TYPE = MT;
+ static constexpr bool IS_V32 = MT == ModelType::V32;
+ static constexpr bool IS_NVFP4 = MT == ModelType::V32_NVFP4_FP8ROPE;
+ static constexpr bool IS_FP4 = MT == ModelType::V41_FP4;
+ static constexpr int D_QK = IS_V32 || IS_NVFP4 ? 576 : 512;
+ static constexpr int D_ROPE = 64;
+ static constexpr int D_NOPE = D_QK - D_ROPE;
+ static constexpr int D_FP4 = IS_NVFP4 ? D_NOPE : (IS_FP4 ? D_QK : 0); // Dimensions stored as fp4 e2m1
+ static constexpr int D_FP8 = IS_NVFP4 ? D_ROPE : (IS_FP4 ? 0 : (MT == ModelType::V41 ? D_QK : D_NOPE)); // Dimensions stored as fp8 e4m3
+ static constexpr int D_BF16 = D_QK - D_FP4 - D_FP8; // Dimensions stored as bf16 (the RoPE part of V3.2 / V4), not needing dequant
+ static constexpr int QUANT_TILE_SIZE = IS_NVFP4 || IS_FP4 ? 16 : (MT == ModelType::V41 ? 32 : (MT == ModelType::V4 ? 64 : 128)); // Dimensions sharing one scale
+ static constexpr int NUM_SCALES_EACH_TOKEN = (IS_V32 || IS_NVFP4 ? D_NOPE : D_QK) / QUANT_TILE_SIZE; // 4 / 8 (7 + 1 byte padding) / 16 / 32
+ static constexpr int SCALE_BYTES = IS_V32 ? 4 : 1; // fp32 (V3.2), ue8m0 (V4 / V4.1) or e4m3 (fp4)
+ static constexpr int QUANT_BYTES = D_FP4 / 2 + D_FP8; // The quantized (fp4 / fp8) data of a token
+ // Bytes between two tokens in the data region of a page block: 656 / 576 / 512 / 256. The stride of the tensor maps of the
+ // quantized part, so it must be >= 256 for the int32 TMA coordinates to cover a whole KV cache
+ static constexpr int TMA_K_STRIDE = IS_NVFP4 ? QUANT_BYTES + NUM_SCALES_EACH_TOKEN :
+ QUANT_BYTES + (IS_V32 ? NUM_SCALES_EACH_TOKEN * SCALE_BYTES : 0) + 2 * D_BF16;
+ static constexpr int BYTES_PER_TOKEN = IS_NVFP4 ? TMA_K_STRIDE :
+ TMA_K_STRIDE + (IS_V32 ? 0 : NUM_SCALES_EACH_TOKEN * SCALE_BYTES); // 656 / 584 / 528 / 288 / 352
+ static_assert(!IS_NVFP4 || BYTES_PER_TOKEN == 352);
+};
+
+// Runtime counterpart of KVCacheFormat::BYTES_PER_TOKEN
+constexpr int kv_cache_bytes_per_token(ModelType mt) {
+ switch (mt) {
+ case ModelType::V32: return KVCacheFormat::BYTES_PER_TOKEN;
+ case ModelType::V4: return KVCacheFormat::BYTES_PER_TOKEN;
+ case ModelType::V41: return KVCacheFormat::BYTES_PER_TOKEN;
+ case ModelType::V41_FP4: return KVCacheFormat::BYTES_PER_TOKEN;
+ case ModelType::V32_NVFP4_FP8ROPE: return KVCacheFormat::BYTES_PER_TOKEN;
+ }
+ return 0;
+}
+
+// The (kv, extra_kv) format pairs that exist: extra_kv has the format of kv, or is the V4.1 fp4 cache next to a V4.1 (fp8) kv.
+constexpr bool is_valid_kv_format_pair(ModelType kv, ModelType extra_kv) {
+ return extra_kv == kv || (kv == ModelType::V41 && extra_kv == ModelType::V41_FP4);
+}
+
+template
+struct KVFormatPair {
+ static_assert(is_valid_kv_format_pair(KV, EXTRA_KV));
+ static constexpr ModelType kv = KV, extra_kv = EXTRA_KV;
+};
+
+// A list of KVFormatPair, see dispatch_kv_formats in csrc/api/common.h
+template
+struct KVFormatPairs {};
diff --git a/csrc/params.h b/csrc/kernels/params.h
similarity index 87%
rename from csrc/params.h
rename to csrc/kernels/params.h
index b0785139..c42487c6 100644
--- a/csrc/params.h
+++ b/csrc/kernels/params.h
@@ -3,10 +3,12 @@
#include "cutlass/bfloat16.h"
enum class ModelType {
- V32,
- MODEL1,
- // V3.2 geometry (d_qk=576) with NVFP4 (e2m1, per-16 e4m3 scales) NoPE and
- // e4m3 RoPE. SM100-only. See csrc/sm100/decode/head64/config.h for the layout.
+ V32, // DeepSeek V3.2 (d_qk=576)
+ V4, // DeepSeek V4 (d_qk=512)
+ V41, // DeepSeek V4.1 (d_qk=512, RoPE fp8, quant tile size 32)
+ V41_FP4, // DeepSeek V4.1 (d_qk=512, fp4 e2m1, quant tile size 16, e4m3 scales)
+ // V3.2 geometry with fp4 e2m1 NoPE, per-16 e4m3 scales, and
+ // unscaled fp8 e4m3 RoPE. Supported by the SM100 head64 kernel.
V32_NVFP4_FP8ROPE
};
@@ -69,7 +71,8 @@ struct SparseAttnDecodeParams {
int d_qk, d_v;
float sm_scale, sm_scale_div_log2;
int num_blocks, page_block_size, topk;
- ModelType model_type;
+ ModelType model_type; // Format of `kv`, see KVCacheFormat in kv_cache_format.h
+ ModelType extra_model_type; // Format of `extra_kv`: model_type, or V41_FP4 next to a V41 `kv` (see is_valid_kv_format_pair)
cutlass::bfloat16_t* __restrict__ q; // [b, s_q, h_q, d_qk]
cutlass::bfloat16_t* __restrict__ kv; // [num_blocks, page_block_size, d_qk]
@@ -94,7 +97,9 @@ struct SparseAttnDecodeParams {
int stride_extra_indices_b, stride_extra_indices_s_q;
cudaStream_t stream;
-
+
+ bool enable_split_kv;
+
// SplitKV-related parameters
float* __restrict__ lse_accum; // [num_splits, s_q, h_q]
float* __restrict__ o_accum; // [num_splits, s_q, h_q, d_v]
@@ -173,11 +178,12 @@ struct SparseAttnFwdParams {
// We have some kernels that implement both prefill and decode modes in a single kernel (with different template instantiations). The following enum helps to distinguish the modes.
enum class SparseAttnFwdMode {
Prefill, // Normal prefill mode
+ Decode, // To trigger decoding mode (without split KV) for kernels that support both prefill and decode
DecodeWithSplitKV, // To trigger decoding mode for kernels that support both prefill and decode
};
template
-inline constexpr bool is_decode_v = std::bool_constant::value;
+inline constexpr bool is_decode_v = std::bool_constant::value;
template
using SparseFwdArgT = std::conditional_t, SparseAttnDecodeParams, SparseAttnFwdParams>;
diff --git a/csrc/sm100/prefill/sparse/common_subroutine.h b/csrc/kernels/sm100/common_subroutine.h
similarity index 66%
rename from csrc/sm100/prefill/sparse/common_subroutine.h
rename to csrc/kernels/sm100/common_subroutine.h
index 36ddab4c..7ce378aa 100644
--- a/csrc/sm100/prefill/sparse/common_subroutine.h
+++ b/csrc/kernels/sm100/common_subroutine.h
@@ -22,8 +22,8 @@ char load_indices_and_generate_mask(
KU_LDG_256(
gIndices + lane_idx*8,
indices,
- ".nc",
- "no_allocate",
+ ".nc",
+ "evict_first",
"evict_normal",
"256B"
);
@@ -47,6 +47,8 @@ char load_indices_and_generate_mask(
/*
Get P from Tensor Memory, reduce P within shared memory, perform masking, and store back if necessary
+For head=64:
+
Initially, since dual gemm is used, we have two P pieces in Tensor Memory, one occupying rows 0 ~ 63 while the other occupying rows 64 ~ 127. We'd like to have them reduced into one single P piece, stored in registers with layout:
N N --- (topk)
@@ -63,17 +65,18 @@ Initially, since dual gemm is used, we have two P pieces in Tensor Memory, one o
(head)
where N = NUM_ELEMS_PER_THREAD
+
*/
template<
int NUM_ELEMS_PER_THREAD,
- int TMEM_COL_START,
int BARRIER_WARP02_SYNC_ID,
int BARRIER_WARP13_SYNC_ID,
bool STORE_BACK_P
>
CUTE_DEVICE
void retrieve_mask_and_reduce_p(
- char* k_validness_base,
+ uint32_t p_tmem_col_start,
+ char* k_validness_base, // The stride of row should be aligned to power of 2, for vectorized load
int local_warp_idx,
int lane_idx,
auto slot_bar_P_empty_arrival,
@@ -85,12 +88,15 @@ void retrieve_mask_and_reduce_p(
static_assert(BARRIER_WARP13_SYNC_ID == BARRIER_WARP02_SYNC_ID+1);
float p_peer[NUM_ELEMS_PER_THREAD];
+ auto load_op = [](uint32_t ts, void* d) {
+ ku::tmem_ld_32dp32bNx(ts, d);
+ };
if (local_warp_idx < 2) {
- ku::tmem_ld_32dp32bNx(TMEM_COL_START, p);
- ku::tmem_ld_32dp32bNx(TMEM_COL_START + NUM_ELEMS_PER_THREAD, p_peer);
+ ku::tmem_ld_st_decomposed(p_tmem_col_start, p, load_op);
+ ku::tmem_ld_st_decomposed(p_tmem_col_start + NUM_ELEMS_PER_THREAD, p_peer, load_op);
} else {
- ku::tmem_ld_32dp32bNx(TMEM_COL_START, p_peer);
- ku::tmem_ld_32dp32bNx(TMEM_COL_START + NUM_ELEMS_PER_THREAD, p);
+ ku::tmem_ld_st_decomposed(p_tmem_col_start, p_peer, load_op);
+ ku::tmem_ld_st_decomposed(p_tmem_col_start + NUM_ELEMS_PER_THREAD, p, load_op);
}
cutlass::arch::fence_view_async_tmem_load();
ku::tcgen05_before_thread_sync();
@@ -98,11 +104,27 @@ void retrieve_mask_and_reduce_p(
// Mask invalid tokens
// We put masking before reduction, since (-inf) + anything (except nan and +inf) is (-inf), which guarantees correctness, and this can overlap with smem load
- static_assert(NUM_ELEMS_PER_THREAD == 32);
- uint32_t is_k_valid = *(uint32_t*)(k_validness_base + (local_warp_idx>=2?NUM_ELEMS_PER_THREAD/8:0));
+ static_assert(NUM_ELEMS_PER_THREAD == 48 || NUM_ELEMS_PER_THREAD == 32 || NUM_ELEMS_PER_THREAD == 16);
+ using is_k_valid_mask_t = \
+ cute::conditional_t>>;
+ int k_validness_offset = local_warp_idx>=2 ? NUM_ELEMS_PER_THREAD/8 : 0; // |0|2| \n |1|3|
+ uint8_t is_k_valid_masks[NUM_ELEMS_PER_THREAD / 8];
+ if constexpr (NUM_ELEMS_PER_THREAD == 16) {
+ *(uint16_t*)is_k_valid_masks = *(uint16_t*)(k_validness_base + k_validness_offset);
+ } else if constexpr (NUM_ELEMS_PER_THREAD == 32) {
+ *(uint32_t*)is_k_valid_masks = *(uint32_t*)(k_validness_base + k_validness_offset);
+ } else if constexpr (NUM_ELEMS_PER_THREAD == 48) {
+ *(uint16_t*)(is_k_valid_masks + 0) = *(uint16_t*)(k_validness_base + k_validness_offset + 0);
+ *(uint16_t*)(is_k_valid_masks + 2) = *(uint16_t*)(k_validness_base + k_validness_offset + 2);
+ *(uint16_t*)(is_k_valid_masks + 4) = *(uint16_t*)(k_validness_base + k_validness_offset + 4);
+ }
CUTE_UNROLL
for (int i = 0; i < NUM_ELEMS_PER_THREAD; i += 1) {
- if (!(is_k_valid >> i & 1))
+ if (!(is_k_valid_masks[i/8] >> (i%8) & 1))
p[i] = -CUDART_INF_F;
}
@@ -136,10 +158,10 @@ void retrieve_mask_and_reduce_p(
/*
Rescale O in Tensor Memory.
-O should occupy 128 rows x (D_V/2) columns in Tensor Memory.
+O should occupy 128 rows x NUM_COLS columns in Tensor Memory.
*/
template<
- int D_V,
+ int NUM_COLS,
int CHUNK_SIZE,
int TMEM_COL_START
>
@@ -151,7 +173,7 @@ void rescale_O(
float2 o[CHUNK_SIZE/2];
CUTE_UNROLL
- for (int chunk_idx = 0; chunk_idx < (D_V/2)/CHUNK_SIZE; ++chunk_idx) {
+ for (int chunk_idx = 0; chunk_idx < NUM_COLS/CHUNK_SIZE; ++chunk_idx) {
// Load O
ku::tmem_ld_32dp32bNx(TMEM_COL_START + chunk_idx*CHUNK_SIZE, o);
cutlass::arch::fence_view_async_tmem_load();
@@ -167,17 +189,29 @@ void rescale_O(
}
}
-template
+template
CUTE_DEVICE
float get_max(
float p[NUM_ELEMS_PER_THREAD]
) {
- float local_max = -CUDART_INF_F;
+ // Use 2-way max to optimize SASS ordering
+ float local_max0 = -CUDART_INF_F;
+ float local_max1 = -CUDART_INF_F;
+ CUTE_UNROLL
+ for (int i = 0; i < NUM_ELEMS_PER_THREAD/2; ++i) {
+ float t = p[i];
+ if constexpr (DO_ABS)
+ t = fabsf(t);
+ local_max0 = max(local_max0, t);
+ }
CUTE_UNROLL
- for (int i = 0; i < NUM_ELEMS_PER_THREAD; ++i) {
- local_max = max(local_max, p[i]);
+ for (int i = 0; i < NUM_ELEMS_PER_THREAD-NUM_ELEMS_PER_THREAD/2; ++i) {
+ float t = p[i+NUM_ELEMS_PER_THREAD/2];
+ if constexpr (DO_ABS)
+ t = fabsf(t);
+ local_max1 = max(local_max1, t);
}
- return local_max;
+ return max(local_max0, local_max1);
}
/*
diff --git a/csrc/kernels/sm100/decode/sparse/head128/README.md b/csrc/kernels/sm100/decode/sparse/head128/README.md
new file mode 100644
index 00000000..ec8b4d62
--- /dev/null
+++ b/csrc/kernels/sm100/decode/sparse/head128/README.md
@@ -0,0 +1 @@
+Head128 decoding kernels are located at `csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512.cu` (for k_dim = 512) or simulated using 2x head64 kernel (for k_dim = 576)
\ No newline at end of file
diff --git a/csrc/kernels/sm100/decode/sparse/head64/config.h b/csrc/kernels/sm100/decode/sparse/head64/config.h
new file mode 100644
index 00000000..23f480e3
--- /dev/null
+++ b/csrc/kernels/sm100/decode/sparse/head64/config.h
@@ -0,0 +1,237 @@
+#pragma once
+
+#include "kernel.h"
+
+#include
+#include
+#include
+
+#include
+
+#include "kernels/defines.h"
+#include "kernels/kv_cache_format.h"
+#include "kernels/sm100/dequant_utils.cuh"
+
+
+namespace sm100::decode::sparse::head64 {
+
+using cutlass::arch::fence_view_async_shared;
+using cutlass::arch::NamedBarrier;
+using e8m0 = __nv_fp8_e8m0;
+using e4m3 = cutlass::float_e4m3_t;
+using namespace cute;
+
+enum NamedBarriers : uint32_t {
+ main_loop_sync = 0,
+ wg0_sync = 1,
+ wg0_warp02_sync = 2,
+ wg0_warp13_sync = 3,
+ everyone_sync = 4
+};
+
+template
+struct KernelTemplate {
+
+static constexpr uint32_t B_H = 64; // Head block size. This kernel only supports h_q == B_H
+static constexpr bool ENABLE_SPLITKV = CONFIG.ENABLE_SPLITKV;
+
+using OrigKVFormat = KVCacheFormat; // Format of `kv`
+using ExtraKVFormat = KVCacheFormat; // Format of `extra_kv`
+static_assert(is_valid_kv_format_pair(CONFIG.MODEL_TYPE, CONFIG.EXTRA_MODEL_TYPE));
+
+static constexpr int D_Q = OrigKVFormat::D_QK;
+static constexpr int D_K = D_Q;
+static constexpr int D_V = 512;
+static constexpr int D_NOPE = OrigKVFormat::D_NOPE;
+static constexpr int D_ROPE = OrigKVFormat::D_ROPE;
+static constexpr int D_FP8 = OrigKVFormat::D_FP8; // K dimensions stored as fp8 and needing dequant
+static constexpr int D_BF16 = OrigKVFormat::D_BF16; // K dimensions stored as bf16 and not needing dequant
+static constexpr int QUANT_TILE_SIZE = OrigKVFormat::QUANT_TILE_SIZE;
+static constexpr bool V_HAVE_ROPE = OrigKVFormat::MODEL_TYPE == ModelType::V32 ? false : true;
+static constexpr int NUM_SCALES_EACH_TOKEN = OrigKVFormat::NUM_SCALES_EACH_TOKEN; // Padding is included
+static constexpr int TMA_K_STRIDE = OrigKVFormat::TMA_K_STRIDE; // Stride of K's tensormap. This stride must 1) be a factor of the actual stride between tokens 2) large enough to cover the entire KV cache. Since TMA copy's coordinate can only be 32bit signed integers, this number must >= 128, perferrably >= 256. So we set this to 656 for V32, 576 for V4 and 512 for V41. Extra padding may be necessary for KV blocks.
+static_assert(D_NOPE + D_ROPE == D_Q);
+static_assert(D_FP8 + D_BF16 == D_Q);
+static_assert(V_HAVE_ROPE ? (D_NOPE + D_ROPE == D_V) : (D_NOPE == D_V));
+
+static constexpr int B_TOPK = 64;
+static constexpr int NUM_BUFS = 2;
+static constexpr int NUM_INDEX_BUFS = 4; // Number of buffers for indices (tma_coords) & is_token_valid & scales
+
+// Both caches are dequantized into the same bf16 tile (B_TOPK x D_FP8, plus the bf16 part) by the same warpgroup, see KVBlockDequantizer
+static_assert(ExtraKVFormat::D_FP8 + ExtraKVFormat::D_FP4 == D_FP8 && ExtraKVFormat::D_BF16 == D_BF16);
+// Bytes between two raw (quantized) rows in shared memory, i.e. the box of their tensor map. fp4 rows are padded by 32 B so that the 4 rows of a
+// gather4 group start in 4 different quarters of the 32 banks (288 / 4 = 72 = 8 mod 32) and the LDS.32 of the dequantizer has no
+// bank conflict; fp8 rows have no padding
+template static constexpr int RAW_TOKEN_SMEM_STRIDE = F::IS_FP4 ? F::QUANT_BYTES + 32 : F::D_FP8;
+// One tma_gather4 writes its 4 rows RAW_TOKEN_SMEM_STRIDE apart, and its destination must be 128 B aligned
+static_assert(4 * RAW_TOKEN_SMEM_STRIDE % 128 == 0 && 4 * RAW_TOKEN_SMEM_STRIDE % 128 == 0);
+// One row of 1 B scales per token. In a kernel with an fp4 extra_kv the rows are 32 B and the fp8 tokens use their first 16 B
+static constexpr int SCALE_SMEM_STRIDE = std::max(OrigKVFormat::NUM_SCALES_EACH_TOKEN, ExtraKVFormat::NUM_SCALES_EACH_TOKEN);
+template using Dequantizer = KVBlockDequantizer, SCALE_SMEM_STRIDE>;
+static constexpr int NUM_THREADS = 128*3; // 128 exp + 1/32 utcmma + 1/32 raw KV producer + 1/32 rope producer + 32 index+scale+valid_mask producer + 128 dequant
+static constexpr float MAX_INIT_VAL = -1e30f; // To avoid (-inf) - (-inf) = NaN
+
+static constexpr int D_Q_SW128 = 512;
+static constexpr int D_Q_SW64 = OrigKVFormat::MODEL_TYPE == ModelType::V32 ? 64 : 0;
+static_assert(D_Q_SW128 + D_Q_SW64 == D_Q);
+static constexpr int K_ROPE_SW = OrigKVFormat::MODEL_TYPE == ModelType::V41 ? 0 : (OrigKVFormat::MODEL_TYPE == ModelType::V32 ? 64 : 128); // RoPE part stored in SW64 (for V32) or SW128 (for V4), in bytes. 0 for V41 (no separate bf16 RoPE, loaded as fp8 in D_FP8)
+
+template<
+ typename Shape_Q_SW128, typename TMA_Q_SW128,
+ typename Shape_O, typename TMA_O
+>
+struct TmaParams {
+ Shape_Q_SW128 shape_Q_SW128; TMA_Q_SW128 tma_Q_SW128;
+ Shape_O shape_O; TMA_O tma_O;
+ CUtensorMap tensor_map_q_sw64; // Invalid if D_Q_SW64 == 0
+ CUtensorMap tensor_map_kv_quant_part; // The quantized (fp8) part of `kv`, one raw row per token
+ CUtensorMap tensor_map_kv_bf16_part; // The bf16 (RoPE) part of `kv`. Invalid if D_BF16 == 0
+ CUtensorMap tensor_map_extra_kv_quant_part; // Same for `extra_kv` (fp8 or fp4). Invalid if extra_topk == 0
+ CUtensorMap tensor_map_extra_kv_bf16_part;
+};
+
+// Tensor memory columns
+struct tmem_cols {
+ // 0 ~ 256: output
+ // 256 ~ 256 + B_H*D_Q/256: Q
+ // 400 ~ 464: P
+ static constexpr int O = 0;
+ static constexpr int Q = 256;
+ static constexpr int Q_Tail = 256 + B_H*D_NOPE/2/128;
+ static constexpr int P = 400;
+};
+
+template
+using SmemLayoutQTiles = decltype(coalesce(tile_to_shape(
+ UMMA::Layout_K_SW128_Atom{},
+ Shape, Int>{},
+ Step<_1, _2>{}
+), Shape<_1, _1>{}));
+
+using SmemLayoutQ_SW128 = SmemLayoutQTiles;
+
+using SmemLayoutOBuf = decltype(tile_to_shape(
+ UMMA::Layout_K_SW128_Atom{},
+ Shape, Int>{}
+));
+
+using SmemLayoutOBuf_TMA = decltype(tile_to_shape(
+ UMMA::Layout_K_SW128_Atom{},
+ Shape, Int<64>>{}
+)); // A TMA tile
+
+static_assert(D_V == 512);
+using SmemLayoutOAccumBuf = Layout<
+ Shape, Int>,
+ Stride, _1> // We use stride = 520 here to avoid bank conflict
+>;
+
+using SmemLayoutS = decltype(tile_to_shape(
+ UMMA::Layout_K_INTER_Atom{},
+ Shape, Int>{},
+ Step<_1, _2>{}
+));
+
+template
+using SmemLayoutKTiles_SW128 = decltype(coalesce(tile_to_shape(
+ UMMA::Layout_K_SW128_Atom{},
+ Shape, Int<64*NUM_TILES>>{},
+ Step<_1, _2>{}
+), Shape<_1, _1>{}));
+
+template
+using SmemLayoutKTiles_DualGemm_SW128 = decltype(coalesce(tile_to_shape(
+ UMMA::Layout_K_SW128_Atom{},
+ Shape, Int<64*NUM_TILES>>{},
+ Step<_1, _2>{}
+), Shape<_1, _1>{}));
+
+template
+using SmemLayoutKTilesTransposed_SW128 = decltype(composition(
+ SmemLayoutKTiles_SW128{},
+ Layout<
+ Shape, Int>,
+ Stride, _1>
+ >{}
+));
+
+template
+using SmemLayoutKTiles_SW64 = decltype(coalesce(tile_to_shape(
+ UMMA::Layout_K_SW64_Atom{},
+ Shape, Int<32*NUM_TILES>>{},
+ Step<_1, _2>{}
+), Shape<_1, _1>{}));
+
+template
+using SmemLayoutKTiles_DualGemm_SW64 = decltype(coalesce(tile_to_shape(
+ UMMA::Layout_K_SW64_Atom{},
+ Shape, Int<32*NUM_TILES>>{},
+ Step<_1, _2>{}
+), Shape<_1, _1>{}));
+
+template
+using SmemLayoutKTilesTransposed_SW64 = decltype(composition(
+ SmemLayoutKTiles_SW64{},
+ Layout<
+ Shape, Int>,
+ Stride, _1>
+ >{}
+));
+
+struct SharedMemoryPlan {
+ union {
+ struct {
+ array_aligned> q;
+ bf16 q_sw64[B_H*D_Q_SW64]; // NOTE D_Q_SW64 may be 0 but array_aligned will have a size of 16, so we use array here. The former tensor (`q`) promises its alignment.
+ union {
+ array_aligned> o_buf;
+ array_aligned> o_accum_buf;
+ } o;
+ } qo;
+ struct {
+ struct {
+ alignas(1024) bf16 quant_part[B_TOPK*D_FP8]; // Quantized-origin (fp8 / fp4) part, dequantized to bf16
+ alignas(1024) bf16 bf16_part[B_TOPK*D_BF16]; // bf16-origin part, swizzled as K_ROPE_SW
+ } dequant[NUM_BUFS];
+ static_assert(sizeof(dequant) >= sizeof(bf16) * (B_H*D_Q)); // So that Q does not cover raw_quant
+ array_aligned raw_quant[NUM_BUFS]; // Raw (quantized) rows of a KV block, RAW_TOKEN_SMEM_STRIDE apart. For V41, includes both NoPE and RoPE. 128 B aligned for gather4
+ static_assert(B_TOPK * RAW_TOKEN_SMEM_STRIDE <= B_TOPK * D_FP8);
+ } kv;
+ } u;
+ union {
+ float p_exchange_buf[4][32 * (B_TOPK/(128/B_H))];
+ array_aligned> s;
+ } s_p;
+ CUTE_ALIGNAS(16) float rowwise_max_buf[128];
+ char is_token_valid[NUM_INDEX_BUFS][B_TOPK/8];
+ int tma_coord[NUM_INDEX_BUFS][B_TOPK];
+ CUTE_ALIGNAS(16) uint8_t scales[NUM_INDEX_BUFS][B_TOPK][SCALE_SMEM_STRIDE]; // ue8m0 (fp8) or e4m3 (fp4), see KVBlockDequantizer
+ array_aligned tmem_start_addr;
+ transac_bar_t bar_last_store_done;
+ transac_bar_t bar_q_tma, bar_q_utccp;
+ transac_bar_t bar_bf16_part_load_ready[NUM_BUFS];
+ transac_bar_t bar_quant_part_dequant_ready[NUM_BUFS];
+ transac_bar_t bar_raw_ready[NUM_BUFS], bar_raw_free[NUM_BUFS];
+ transac_bar_t bar_valid_coord_scale_ready[NUM_INDEX_BUFS], bar_valid_coord_scale_free[NUM_INDEX_BUFS];
+ transac_bar_t bar_qk_done[NUM_BUFS], bar_so_ready[NUM_BUFS], bar_sv_done[NUM_BUFS];
+};
+
+using TiledMMA_P = decltype(make_tiled_mma(
+ SM100_MMA_F16BF16_WS_TS_NOELECT{}
+)); // *2 for dual gemm
+
+static constexpr int PV_GEMM_N = 256;
+using TiledMMA_O = decltype(make_tiled_mma(
+ SM100_MMA_F16BF16_WS_SS_NOELECT{}
+));
+
+template
+static __device__ void
+flash_fwd_splitkv_mla_fp8_sparse_kernel_devfunc(const SparseAttnDecodeParams ¶ms, const TmaParam &tma_params );
+
+static void run(const SparseAttnDecodeParams ¶ms);
+
+};
+
+}
diff --git a/csrc/kernels/sm100/decode/sparse/head64/instantiations/v32_h64.cu b/csrc/kernels/sm100/decode/sparse/head64/instantiations/v32_h64.cu
new file mode 100644
index 00000000..408b5817
--- /dev/null
+++ b/csrc/kernels/sm100/decode/sparse/head64/instantiations/v32_h64.cu
@@ -0,0 +1,8 @@
+#include "../kernel.cuh"
+
+namespace sm100::decode::sparse::head64 {
+
+template
+void run_flash_splitkv_mla_fp8_sparse_kernel(const SparseAttnDecodeParams ¶ms);
+
+}
diff --git a/csrc/kernels/sm100/decode/sparse/head64/instantiations/v32_h64_no_split.cu b/csrc/kernels/sm100/decode/sparse/head64/instantiations/v32_h64_no_split.cu
new file mode 100644
index 00000000..44a5fcb8
--- /dev/null
+++ b/csrc/kernels/sm100/decode/sparse/head64/instantiations/v32_h64_no_split.cu
@@ -0,0 +1,8 @@
+#include "../kernel.cuh"
+
+namespace sm100::decode::sparse::head64 {
+
+template
+void run_flash_splitkv_mla_fp8_sparse_kernel(const SparseAttnDecodeParams ¶ms);
+
+}
diff --git a/csrc/kernels/sm100/decode/sparse/head64/instantiations/v41_h64.cu b/csrc/kernels/sm100/decode/sparse/head64/instantiations/v41_h64.cu
new file mode 100644
index 00000000..81a6b788
--- /dev/null
+++ b/csrc/kernels/sm100/decode/sparse/head64/instantiations/v41_h64.cu
@@ -0,0 +1,8 @@
+#include "../kernel.cuh"
+
+namespace sm100::decode::sparse::head64 {
+
+template
+void run_flash_splitkv_mla_fp8_sparse_kernel(const SparseAttnDecodeParams ¶ms);
+
+}
diff --git a/csrc/kernels/sm100/decode/sparse/head64/instantiations/v41_h64_no_split.cu b/csrc/kernels/sm100/decode/sparse/head64/instantiations/v41_h64_no_split.cu
new file mode 100644
index 00000000..3c46019b
--- /dev/null
+++ b/csrc/kernels/sm100/decode/sparse/head64/instantiations/v41_h64_no_split.cu
@@ -0,0 +1,8 @@
+#include "../kernel.cuh"
+
+namespace sm100::decode::sparse::head64 {
+
+template
+void run_flash_splitkv_mla_fp8_sparse_kernel(const SparseAttnDecodeParams ¶ms);
+
+}
diff --git a/csrc/kernels/sm100/decode/sparse/head64/instantiations/v41fp4_h64.cu b/csrc/kernels/sm100/decode/sparse/head64/instantiations/v41fp4_h64.cu
new file mode 100644
index 00000000..9646d6cd
--- /dev/null
+++ b/csrc/kernels/sm100/decode/sparse/head64/instantiations/v41fp4_h64.cu
@@ -0,0 +1,8 @@
+#include "../kernel.cuh"
+
+namespace sm100::decode::sparse::head64 {
+
+template
+void run_flash_splitkv_mla_fp8_sparse_kernel(const SparseAttnDecodeParams ¶ms);
+
+}
diff --git a/csrc/kernels/sm100/decode/sparse/head64/instantiations/v41fp4_h64_no_split.cu b/csrc/kernels/sm100/decode/sparse/head64/instantiations/v41fp4_h64_no_split.cu
new file mode 100644
index 00000000..61d2b683
--- /dev/null
+++ b/csrc/kernels/sm100/decode/sparse/head64/instantiations/v41fp4_h64_no_split.cu
@@ -0,0 +1,8 @@
+#include "../kernel.cuh"
+
+namespace sm100::decode::sparse::head64 {
+
+template
+void run_flash_splitkv_mla_fp8_sparse_kernel(const SparseAttnDecodeParams ¶ms);
+
+}
diff --git a/csrc/kernels/sm100/decode/sparse/head64/instantiations/v4_h64.cu b/csrc/kernels/sm100/decode/sparse/head64/instantiations/v4_h64.cu
new file mode 100644
index 00000000..de5943a2
--- /dev/null
+++ b/csrc/kernels/sm100/decode/sparse/head64/instantiations/v4_h64.cu
@@ -0,0 +1,8 @@
+#include "../kernel.cuh"
+
+namespace sm100::decode::sparse::head64 {
+
+template
+void run_flash_splitkv_mla_fp8_sparse_kernel(const SparseAttnDecodeParams ¶ms);
+
+}
diff --git a/csrc/kernels/sm100/decode/sparse/head64/instantiations/v4_h64_no_split.cu b/csrc/kernels/sm100/decode/sparse/head64/instantiations/v4_h64_no_split.cu
new file mode 100644
index 00000000..1e0876c2
--- /dev/null
+++ b/csrc/kernels/sm100/decode/sparse/head64/instantiations/v4_h64_no_split.cu
@@ -0,0 +1,8 @@
+#include "../kernel.cuh"
+
+namespace sm100::decode::sparse::head64 {
+
+template
+void run_flash_splitkv_mla_fp8_sparse_kernel(const SparseAttnDecodeParams ¶ms);
+
+}
diff --git a/csrc/kernels/sm100/decode/sparse/head64/kernel.cuh b/csrc/kernels/sm100/decode/sparse/head64/kernel.cuh
new file mode 100644
index 00000000..0e5b68d9
--- /dev/null
+++ b/csrc/kernels/sm100/decode/sparse/head64/kernel.cuh
@@ -0,0 +1,908 @@
+/*
+Sparse MLA Decoding — SM100, h_q == 64
+
+Decoding kernel for h_q == 64 heads. Uses FP8 KV cache with dequantization,
+UTCMMA/TMEM architecture, and TMA for data transfers. Supports split-KV partitioning
+for large topk values. Config defines per-model-type dimensions (V3.2: d_qk=576,
+V4: d_qk=512, RoPE bf16, tile_size=64; V4.1: d_qk=512, RoPE fp8, tile_size=32).
+
+Template parameter: Config (see kernel.h) — the formats of `kv` and `extra_kv` and whether split-KV is enabled
+
+NUM_THREADS=384 (3 warpgroups)
+
+I/O: See SparseAttnDecodeParams in kernels/params.h
+*/
+#include "kernel.h"
+
+#include
+#include
+#include
+#include
+#include
+#include
+
+#include "kerutils/kerutils.cuh"
+
+#include "kernels/utils.h"
+#include "kernels/defines.h"
+#include "kernels/sm100/helpers.h"
+#include "kernels/sm100/common_subroutine.h"
+
+#include "config.h"
+
+namespace sm100::decode::sparse::head64 {
+
+template
+template
+__device__ void
+KernelTemplate
+::flash_fwd_splitkv_mla_fp8_sparse_kernel_devfunc(const SparseAttnDecodeParams ¶ms, const TmaParam &tma_params ) {
+#if defined(KERUTILS_ENABLE_SM100A)
+ const int s_q_idx = blockIdx.x;
+ const int partition_idx = blockIdx.y;
+ const int warpgroup_idx = cutlass::canonical_warp_group_idx();
+ const int idx_in_warpgroup = threadIdx.x % 128;
+ const int warp_idx = cutlass::canonical_warp_idx_sync();
+ const int lane_idx = threadIdx.x % 32;
+
+ extern __shared__ char wksp_buf[];
+ SharedMemoryPlan &plan = *reinterpret_cast(wksp_buf);
+
+ if (warp_idx == 0 && elect_one_sync()) {
+ cute::prefetch_tma_descriptor(tma_params.tma_Q_SW128.get_tma_descriptor());
+ cute::prefetch_tma_descriptor(tma_params.tma_O.get_tma_descriptor());
+ cute::prefetch_tma_descriptor(&tma_params.tensor_map_q_sw64);
+ cute::prefetch_tma_descriptor(&tma_params.tensor_map_kv_quant_part);
+ if constexpr (D_BF16 > 0) {
+ cute::prefetch_tma_descriptor(&tma_params.tensor_map_kv_bf16_part);
+ }
+ }
+
+ if (warp_idx == 0) {
+ if (elect_one_sync()) {
+ plan.bar_last_store_done.init(128);
+ plan.bar_q_tma.init(1);
+ plan.bar_q_utccp.init(1);
+ for (int i = 0; i < NUM_BUFS; ++i) {
+ if constexpr (D_BF16 > 0)
+ plan.bar_bf16_part_load_ready[i].init(1);
+ plan.bar_quant_part_dequant_ready[i].init(128);
+ plan.bar_raw_ready[i].init(1);
+ plan.bar_raw_free[i].init(128);
+ plan.bar_qk_done[i].init(1);
+ plan.bar_so_ready[i].init(128);
+ plan.bar_sv_done[i].init(1);
+ }
+ for (int i = 0; i < NUM_INDEX_BUFS; ++i) {
+ plan.bar_valid_coord_scale_ready[i].init(32);
+ plan.bar_valid_coord_scale_free[i].init(128+128+1+(D_BF16 > 0));
+ }
+ cutlass::arch::fence_barrier_init();
+ }
+ cute::TMEM::Allocator1Sm().allocate(512, plan.tmem_start_addr.data());
+ KU_TRAP_ONLY_DEVICE_ASSERT(plan.tmem_start_addr.data()[0] == 0);
+ cute::TMEM::Allocator1Sm().release_allocation_lock();
+ }
+ __syncthreads();
+
+ struct MainLoopArgs {
+ int batch_idx, start_block_idx, end_block_idx;
+ bool is_no_split; int n_split_idx;
+ bool bar_phase_batch_rel; // Bar phase of barriers that are used once per batch
+ int topk_length, extra_topk_length, num_orig_kv_blocks;
+ bool is_last_batch;
+ };
+
+ auto run_main_loop = [&](auto f) {
+ // NOTE Putting the following code outside the warpgroup specialization switch results in register spilling.
+ // [[maybe_unused]] int begin_req_idx, end_req_idx, sched_begin_block_idx, sched_end_block_idx, begin_n_split_idx, is_first_req_splitted, is_last_req_splitted;
+ DecodingSchedMeta sched_meta;
+ if constexpr (ENABLE_SPLITKV) {
+ KU_LDG_256(
+ params.tile_scheduler_metadata_ptr + partition_idx,
+ &sched_meta,
+ ".nc",
+ "no_allocate",
+ "evict_normal",
+ "256B"
+ );
+ } else {
+ sched_meta.begin_req_idx = sched_meta.end_req_idx = partition_idx;
+ }
+
+ if (sched_meta.begin_req_idx >= params.b) {
+ return;
+ }
+
+ bool bar_phase_batch_rel = 0;
+ #pragma unroll 1
+ for (int batch_idx = sched_meta.begin_req_idx; batch_idx <= sched_meta.end_req_idx; ++batch_idx, bar_phase_batch_rel ^= 1) {
+ int start_block_idx, end_block_idx;
+ bool is_split; int n_split_idx;
+ int topk_length = params.topk_length ? __ldg(params.topk_length + batch_idx) : params.topk;
+ int orig_topk_padded = max(ku::ceil(topk_length, (int)B_TOPK), (int)B_TOPK);
+ int extra_topk_length = params.extra_topk_length ? __ldg(params.extra_topk_length + batch_idx) : params.extra_topk;
+ int total_topk_padded = orig_topk_padded + ku::ceil(extra_topk_length, (int)B_TOPK); // % B_TOPK == 0
+ if constexpr (ENABLE_SPLITKV) {
+ start_block_idx = batch_idx == sched_meta.begin_req_idx ? sched_meta.begin_block_idx : 0;
+ end_block_idx = batch_idx == sched_meta.end_req_idx ? sched_meta.end_block_idx : total_topk_padded / B_TOPK;
+ is_split = batch_idx == sched_meta.begin_req_idx ? sched_meta.is_first_req_splitted : (batch_idx == sched_meta.end_req_idx ? sched_meta.is_last_req_splitted : false);
+ n_split_idx = batch_idx == sched_meta.begin_req_idx ? (__ldg(params.num_splits_ptr+batch_idx) + sched_meta.begin_split_idx) : __ldg(params.num_splits_ptr+batch_idx);
+ } else {
+ start_block_idx = 0;
+ end_block_idx = total_topk_padded / B_TOPK;
+ is_split = false;
+ n_split_idx = 0;
+ }
+
+ MainLoopArgs args = {
+ batch_idx, start_block_idx, end_block_idx,
+ !is_split, n_split_idx,
+ bar_phase_batch_rel,
+ topk_length, extra_topk_length,
+ orig_topk_padded / B_TOPK,
+ batch_idx == sched_meta.end_req_idx
+ };
+
+ f(args);
+ NamedBarrier(NUM_THREADS, NamedBarriers::everyone_sync).arrive_and_wait_unaligned();
+ }
+ };
+
+ struct RingState {
+ int buf_idx = 0;
+ bool bar_phase = 0;
+ int index_buf_idx = 0;
+ bool index_bar_phase = 0;
+ CUTE_DEVICE void update() {
+ bar_phase ^= (buf_idx == NUM_BUFS-1);
+ buf_idx = (buf_idx+1) % NUM_BUFS;
+ index_bar_phase ^= (index_buf_idx == NUM_INDEX_BUFS-1);
+ index_buf_idx = (index_buf_idx+1) % NUM_INDEX_BUFS;
+ }
+ };
+ RingState rs;
+
+ if (warpgroup_idx == 0) {
+ // Scale & Exp warpgroup
+ // The same technique (and highly similar code) as the sm100 sparse prefill head64 kernel
+ cutlass::arch::warpgroup_reg_alloc<224>();
+
+ constexpr int B_EPI = 64; // Must be equal to the size of the swizzle atom
+ Tensor sO = make_tensor(make_smem_ptr(plan.u.qo.o.o_buf.data()), SmemLayoutOBuf{});
+ bf16* sO_bases[B_EPI/8]; // 64 is the size of the swizzle atom (in number of elements) while 8 is the width of each write
+ CUTE_UNROLL
+ for (int i = 0; i < B_EPI/8; ++i)
+ sO_bases[i] = &sO(idx_in_warpgroup%B_H, (idx_in_warpgroup/B_H)*(PV_GEMM_N/(128/B_H)) + i*8);
+
+ const float2 scale = float2 {params.sm_scale_div_log2, params.sm_scale_div_log2};
+ bf16* sS_base = plan.s_p.s.data() + lane_idx*8 + (warp_idx&1)*(B_H/2)*8 + (warp_idx/2)*B_H*(B_TOPK/2);
+
+ // NOTE Padding rows (idx%B_H >= h_q) must not read
+ // attn_sink (OOB); their output is discarded anyway.
+ float attn_sink = (params.attn_sink == nullptr || idx_in_warpgroup%B_H >= params.h_q)
+ ? -CUDART_INF_F : __ldg((float*)params.attn_sink + (idx_in_warpgroup%B_H)) * CUDART_L2E_F;
+ constexpr int NUM_ELEMS_PER_THREAD = B_TOPK / 2;
+
+ run_main_loop([&](const MainLoopArgs &args) {
+ cute::tma_store_wait<0>();
+ plan.bar_last_store_done.arrive();
+
+ float mi = MAX_INIT_VAL;
+ float li = 0.0f;
+ float real_mi = -CUDART_INF_F;
+
+ CUTE_NO_UNROLL
+ for (int block_idx = args.start_block_idx; block_idx < args.end_block_idx; ++block_idx) {
+ NamedBarrier::arrive_and_wait(128, NamedBarriers::wg0_sync); // Make sure all intermediate buffers (including p_exchange_buf, rowwise max_buf) are free
+ plan.bar_valid_coord_scale_ready[rs.index_buf_idx].wait(rs.index_bar_phase); // Put the barrier wait here for more code reordering space
+ plan.bar_qk_done[rs.buf_idx].wait(rs.bar_phase);
+ ku::tcgen05_after_thread_sync();
+
+
+ float p[NUM_ELEMS_PER_THREAD];
+ retrieve_mask_and_reduce_p<
+ NUM_ELEMS_PER_THREAD,
+ NamedBarriers::wg0_warp02_sync,
+ NamedBarriers::wg0_warp13_sync,
+ false
+ >(
+ tmem_cols::P,
+ plan.is_token_valid[rs.index_buf_idx],
+ warp_idx, lane_idx,
+ [&]() {},
+ plan.s_p.p_exchange_buf,
+ p
+ );
+
+ float cur_pi_max = get_max(p);
+ cur_pi_max *= params.sm_scale_div_log2;
+
+ plan.rowwise_max_buf[idx_in_warpgroup] = cur_pi_max;
+ NamedBarrier::arrive_and_wait(128, NamedBarriers::wg0_sync); // This also separates "reading p_exchange_buf" and "writing S"
+ plan.bar_valid_coord_scale_free[rs.index_buf_idx].arrive();
+ cur_pi_max = max(cur_pi_max, plan.rowwise_max_buf[idx_in_warpgroup^64]);
+ real_mi = max(real_mi, cur_pi_max);
+ bool should_scale_o = __any_sync(0xffffffff, cur_pi_max - mi > 6.0f);
+ // By this point:
+ // - cur_pi_max, real_mi, and mi is identical within each row (i.e. thread 0+64, 1+65, ... for HEAD64)
+ // - should_scale_o is identical among every warp, and is identical among threads that controls the same row
+
+ // Calc scale factor, and scale li
+ float new_max, scale_for_old;
+ if (!should_scale_o) {
+ // Don't scale O
+ scale_for_old = 1.0f;
+ new_max = mi;
+ } else {
+ new_max = max(cur_pi_max, mi);
+ scale_for_old = exp2f(mi - new_max);
+ }
+ mi = new_max; // mi is still identical within each row
+
+ // Calculate S
+ __nv_bfloat162 s[NUM_ELEMS_PER_THREAD/2];
+ float2 neg_new_max = float2 {-new_max, -new_max};
+ float2 cur_sum = float2 {0.0f, 0.0f};
+ CUTE_UNROLL
+ for (int i = 0; i < NUM_ELEMS_PER_THREAD/2; i += 1) {
+ float2 d = ku::float2_fma(float2{p[i*2], p[i*2+1]}, scale, neg_new_max);
+ d.x = exp2f(d.x);
+ d.y = exp2f(d.y);
+ cur_sum = ku::float2_add(cur_sum, d);
+ s[i] = __float22bfloat162_rn(d);
+ }
+ li = fma(li, scale_for_old, (cur_sum.x + cur_sum.y));
+
+ // Write S
+ CUTE_UNROLL
+ for (int i = 0; i < NUM_ELEMS_PER_THREAD/8; i += 1) {
+ *(uint128_t*)(sS_base + B_H*8*i) = *(uint128_t*)(s + i*4);
+ }
+
+ // Scale O
+ if (block_idx != args.start_block_idx && should_scale_o) {
+ ku::tcgen05_after_thread_sync();
+
+ rescale_O<
+ D_V / (128/B_H),
+ 64,
+ tmem_cols::O
+ >(scale_for_old);
+ ku::tcgen05_before_thread_sync();
+ }
+
+ fence_view_async_shared();
+ plan.bar_so_ready[rs.buf_idx].arrive();
+
+ if (block_idx != args.end_block_idx-1) {
+ rs.update(); // Don't update rs for the last round since we want to wait for the last SV gemm
+ }
+
+ }
+
+ if (real_mi == -CUDART_INF_F) {
+ // real_mi == -CUDART_INF_F <=> No valid TopK indices
+ // We set li to 0 to fit the definition that li := exp(x[i] - mi)
+ li = 0.0f;
+ mi = -CUDART_INF_F;
+ }
+
+ // Exchange li
+ plan.rowwise_max_buf[idx_in_warpgroup] = li;
+ NamedBarrier::arrive_and_wait(128, NamedBarriers::wg0_sync);
+ li += plan.rowwise_max_buf[idx_in_warpgroup^64];
+
+ // Store li
+ if (idx_in_warpgroup < params.h_q) {
+ if (args.is_no_split) {
+ float cur_lse = fma(mi, CUDART_LN2_F, logf(li));
+ cur_lse = cur_lse == -CUDART_INF_F ? +CUDART_INF_F : cur_lse;
+ float* gSoftmaxLse = (float*)params.lse + args.batch_idx*params.stride_lse_b + s_q_idx*params.stride_lse_s_q + idx_in_warpgroup;
+ *gSoftmaxLse = cur_lse;
+ } else {
+ float cur_lse = log2f(li) + mi;
+ float* gSoftmaxLseAccum = (float*)params.lse_accum + args.n_split_idx*params.stride_lse_accum_split + s_q_idx*params.stride_lse_accum_s_q + idx_in_warpgroup;
+ *gSoftmaxLseAccum = cur_lse;
+ }
+ }
+
+ plan.bar_sv_done[rs.buf_idx].wait(rs.bar_phase);
+ rs.update();
+ ku::tcgen05_after_thread_sync();
+
+ if (args.is_last_batch) {
+ cudaTriggerProgrammaticLaunchCompletion();
+ }
+
+ static constexpr int O_FOLD = 2;
+ if (args.is_no_split) {
+ Tensor tma_gO = flat_divide(
+ tma_params.tma_O.get_tma_tensor(tma_params.shape_O)(_, _, s_q_idx, args.batch_idx),
+ Shape, Int<64>>{}
+ )(_, _, _0{}, _);
+ auto thr_tma = tma_params.tma_O.get_slice(_0{});
+ Tensor tma_sO = flat_divide(
+ sO,
+ Shape, Int<64>>{}
+ )(_, _, _0{}, _);
+
+ float o_scale = li == 0.0f ? 0.0f : __fdividef(1.0f, li + exp2f(attn_sink - mi));
+ float2 o_scale_float2 = {o_scale, o_scale};
+ float2 o[B_EPI/2];
+ __nv_bfloat162 o_bf16[B_EPI/2];
+ CUTE_UNROLL
+ for (int i = 0; i < (D_V/O_FOLD) / B_EPI; ++i) {
+ // Load
+ ku::tmem_ld_32dp32bNx(tmem_cols::O + i*B_EPI, o);
+ cutlass::arch::fence_view_async_tmem_load();
+ // Scale & Convert
+ CUTE_UNROLL
+ for (int j = 0; j < B_EPI/2; ++j) {
+ o[j] = ku::float2_mul(o[j], o_scale_float2);
+ o_bf16[j] = __float22bfloat162_rn(o[j]);
+ }
+ // Store
+ int col_base = (i >= PV_GEMM_N/O_FOLD/B_EPI ? D_V/2 : 0) + (i*B_EPI%(PV_GEMM_N/O_FOLD));
+ CUTE_UNROLL
+ for (int j = 0; j < B_EPI / 8; ++j)
+ *(__int128_t*)(sO_bases[j] + col_base*B_H) = *(__int128_t*)(&o_bf16[j*4]);
+ // Sync
+ fence_view_async_shared();
+ NamedBarrier::arrive_and_wait(128, NamedBarriers::wg0_sync);
+ // S -> G
+ if (warp_idx < O_FOLD && elect_one_sync()) {
+ int copy_block_idx = warp_idx*(PV_GEMM_N/2/B_EPI) + col_base/B_EPI;
+ cute::copy(
+ tma_params.tma_O,
+ thr_tma.partition_S(tma_sO(_, _, copy_block_idx)),
+ thr_tma.partition_D(tma_gO(_, _, copy_block_idx))
+ );
+ }
+ }
+ cute::tma_store_arrive();
+ } else {
+ float o_scale = li == 0.0f ? 0.0f : __fdividef(1.0f, li); // Here we leave attn_sink to the combine kernel, otherwise attn_sink will take effect for multiple times
+ float2 o_scale_float2 = {o_scale, o_scale};
+ constexpr int B_EPI = 64;
+ float2 o[B_EPI/2];
+ Tensor sO = make_tensor(make_smem_ptr(plan.u.qo.o.o_accum_buf.data()), SmemLayoutOAccumBuf{});
+ CUTE_UNROLL
+ for (int i = 0; i < (D_V/O_FOLD) / B_EPI; ++i) {
+ // Load
+ ku::tmem_ld_32dp32bNx(tmem_cols::O + i*B_EPI, o);
+ cutlass::arch::fence_view_async_tmem_load();
+ // Scale & Convert
+ CUTE_UNROLL
+ for (int j = 0; j < B_EPI/2; ++j)
+ o[j] = ku::float2_mul(o[j], o_scale_float2);
+ // Store
+ int col_base = (idx_in_warpgroup/B_H)*(PV_GEMM_N/O_FOLD) + (i >= PV_GEMM_N/O_FOLD/B_EPI ? D_V/2 : 0) + (i*B_EPI%(PV_GEMM_N/O_FOLD));
+ CUTE_UNROLL
+ for (int j = 0; j < B_EPI / 4; ++j)
+ *(__int128_t*)&sO(idx_in_warpgroup%B_H, col_base + j*4) = *(__int128_t*)(&o[j*2]);
+ }
+ fence_view_async_shared();
+ NamedBarrier::arrive_and_wait(128, NamedBarriers::wg0_sync);
+ if (elect_one_sync()) {
+ CUTE_UNROLL
+ for (int local_row = 0; local_row < B_H/4; ++local_row) {
+ int smem_row = local_row*4 + warp_idx;
+ if (smem_row < params.h_q) { // Don't write padding rows (o_accum only has h_q rows)
+ SM90_BULK_COPY_S2G::copy(
+ &sO(smem_row, _0{}),
+ (float*)params.o_accum + args.n_split_idx*params.stride_o_accum_split + s_q_idx*params.stride_o_accum_s_q + smem_row*params.stride_o_accum_h_q,
+ D_V*sizeof(float)
+ );
+ }
+ }
+ cute::tma_store_arrive();
+ }
+ }
+ });
+
+ if (warp_idx == 0) {
+ cute::TMEM::Allocator1Sm().free(0, 512);
+ }
+ } else if (warpgroup_idx == 1) {
+ cutlass::arch::warpgroup_reg_dealloc<72>();
+ const int warp_idx = cutlass::canonical_warp_idx_sync(); // Missing this leads to reg spilling
+
+ if (warp_idx == 4 && elect_one_sync()) {
+
+ // MMA Warp
+ run_main_loop([&](const MainLoopArgs &args) {
+ if (args.start_block_idx >= args.end_block_idx) {
+ ku::trap();
+ }
+ // Issue Q (SW128) G->S
+ {
+ Tensor gQ = tma_params.tma_Q_SW128.get_tma_tensor(tma_params.shape_Q_SW128)(_, _, s_q_idx, args.batch_idx);
+ Tensor sQ = make_tensor(make_smem_ptr(plan.u.qo.q.data()), SmemLayoutQ_SW128{});
+ ku::launch_tma_copy(
+ tma_params.tma_Q_SW128,
+ gQ,
+ sQ,
+ plan.bar_q_tma,
+ TMA::CacheHintSm90::EVICT_FIRST
+ );
+ }
+ // Issue Q (SW64) G -> S
+ if constexpr (D_Q_SW64 > 0) {
+ cute::SM90_TMA_LOAD_5D::copy(
+ &tma_params.tensor_map_q_sw64,
+ (uint64_t*)&plan.bar_q_tma,
+ (uint64_t)TMA::CacheHintSm90::EVICT_FIRST,
+ plan.u.qo.q_sw64,
+ 0, 0, 0,
+ s_q_idx, args.batch_idx
+ );
+ }
+ plan.bar_q_tma.arrive_and_expect_tx(B_H*D_Q*sizeof(bf16));
+ plan.bar_q_tma.wait(args.bar_phase_batch_rel);
+ ku::tcgen05_after_thread_sync();
+ // Issue Q (SW128) UTCCP
+ {
+ UMMA::SmemDescriptor sQ_desc = UMMA::make_umma_desc(
+ make_tensor(
+ make_smem_ptr(plan.u.qo.q.data()),
+ tile_to_shape(
+ UMMA::Layout_K_SW128_Atom{},
+ Shape, Int<64>>{} // *2 to leverage dual GEMM
+ )
+ )
+ );
+ static_assert(D_Q_SW128%128 == 0);
+ #pragma unroll D_Q_SW128/128
+ for (int tile_idx = 0; tile_idx < D_Q_SW128/128; ++tile_idx) {
+ // Each tile: 64 x (64*2) logically, 128 x 64 bf16 on TMEM
+ CUTE_UNROLL
+ for (int subtile_idx = 0; subtile_idx < B_H/16; ++subtile_idx) {
+ // Each subtile: 64 x (16*2) logically, 128 x 16 bf16 (128dp256b) on TMEM
+ SM100_UTCCP_128dp256bit_1cta::copy(
+ sQ_desc + (tile_idx*(B_H*128) + subtile_idx*16) * 2 / 16,
+ tmem_cols::Q + tile_idx*32 + subtile_idx*8
+ );
+ }
+ }
+ }
+ // Issue Q (SW64) UTCCP
+ if constexpr (D_Q_SW64 > 0) {
+ UMMA::SmemDescriptor sQ_SW64_desc = UMMA::make_umma_desc(
+ make_tensor(
+ make_smem_ptr(plan.u.qo.q_sw64),
+ tile_to_shape(
+ UMMA::Layout_K_SW64_Atom{},
+ Shape