diff --git a/README.md b/README.md index 1945725f..de65f762 100644 --- a/README.md +++ b/README.md @@ -18,6 +18,7 @@ FlashMLA is DeepSeek's library of optimized attention kernels, powering the [Dee ## News +- **2026.09.10 Release of DeepSeek v4.1's Attention Kernels**: We've released attention kernels for [DeepSeek-V4.1](https://huggingface.co/deepseek-ai/DeepSeek-V4.1-Flash), including both prefill and decoding (with FP8 or FP4 KV cache). We've also released a [fused-norm-rope-attn-rope-cast kernel](#fused-norm--rope--attn--rope--cast-kernel) which fuses Q-norm (only used in V4, not V4.1), Q-RoPE, core attention, O-RoPE (conjugate), and cast-to-fp8, while retaining the same performance. - **2025.09.29 Release of Sparse Attention Kernels**: With the launch of [DeepSeek-V3.2](https://github.com/deepseek-ai/DeepSeek-V3.2-Exp), we are releasing the corresponding token-level sparse attention kernels. These kernels power the model's DeepSeek Sparse Attention (DSA) and achieve up to 640 TFlops during prefilling and 410 TFlops during decoding. We also release a deep-dive blog for our new FP8 sparse decoding kernel. Check it out [here](docs/20250929-hopper-fp8-sparse-deep-dive.md). - **2025.08.01 Kernels for MHA on SM100**: Thanks to [NVIDIA's PR](https://github.com/deepseek-ai/FlashMLA/pull/76) for MHA forward / backward kernels on SM100! - **2025.04.22 Deep-Dive Blog**: We'd love to share the technical details behind the new FlashMLA kernel! Check out our deep-dive write-up [here](docs/20250422-new-kernel-deep-dive.md). @@ -32,7 +33,7 @@ python tests/test_flash_mla_dense_decoding.py python tests/test_flash_mla_sparse_decoding.py ``` -The dense MLA decoding kernel achieves up to 3000 GB/s in memory-bound configuration and 660 TFLOPS in computation-bound configuration on H800 SXM5 with CUDA 12.8. The token-level sparse MLA decoding kernel (which uses an FP8 KV cache while performing the matrix multiplication in bfloat16) achieves 410 TFLOPS in compute-bound configuration on H800 SXM5 with CUDA 12.8, and achieves up to 350 TFlops on B200 (which is not really optimized yet). +The dense MLA decoding kernel achieves up to 3000 GB/s in memory-bound configuration and 660 TFLOPS in computation-bound configuration on H800 SXM5 with CUDA 12.8. The token-level sparse MLA decoding kernel (which uses an FP8 KV cache while performing the matrix multiplication in bfloat16) achieves 410 TFLOPS in compute-bound configuration on H800 SXM5 with CUDA 12.8, and achieves up to 700 TFlops on B200. #### Test & benchmark MHA prefill (Dense): @@ -50,6 +51,17 @@ python tests/test_flash_mla_sparse_prefill.py It achieves up to 640 TFlops in forward computation on H800 SXM5 with CUDA 12.8, and achieves up to 1450 TFlops on B200, CUDA 12.9. +#### Test & benchmark the fused norm RoPE attn RoPE cast kernel (Sparse): + +```bash +python tests/test_fused_norm_rope_attn_rope_cast.py +``` + +[TileLang](https://github.com/tile-ai/tilelang), [Tile-Kernels](https://github.com/deepseek-ai/TileKernels/), and [DeepGEMM](https://github.com/deepseek-ai/DeepGEMM) are required for running this test script. + +This kernel fuses Q-norm (only used in V4, not V4.1), Q-RoPE, core attention, O-RoPE (conjugate) and cast-to-FP8 into a single kernel, saving times for those small kernels. Although it fuses many small operations, this kernel still keeps the same or even slightly higher TFlops (at the cost of having to permute the Q_b and Wv weights in advance). It achieves up to 1430 TFlops during prefill and 670 TFlops during decoding on B200. + + ## Requirements - SM90 / SM100 (See the support matrix below) @@ -58,16 +70,17 @@ It achieves up to 640 TFlops in forward computation on H800 SXM5 with CUDA 12.8, Support matrix: -| Kernel | GPU Architecture | MLA Mode [2] | KVCache Format | +| Kernel | GPU Architecture | MLA Mode [1] | Supported Models | | :---: | :---: | :---: | :---: | -| Dense Decoding | SM90 | MQA | BF16 | -| Sparse Decoding | SM90 & SM100 | MQA | FP8 [1] | -| Dense Prefill | SM100 | MHA | | -| Sparse Prefill | SM90 & SM100 | MQA | | +| Dense Decoding | SM90 | MQA | DeepSeek V3 / V3.1 | +| Sparse Decoding | SM90 & SM100 | MQA | DeepSeek V3.2 / V4 / V4.1 [2] | +| Dense Prefill | SM100 | MHA | DeepSeek V3 / V3.1 / V3.2 | +| Sparse Prefill | SM90 & SM100 | MQA | DeepSeek V3.2 / V4 / V4.1 | +| Fused Norm RoPE Attn RoPE Cast | SM100 | MQA | DeepSeek V4 / V4.1 | -[1]: For more details on using FP8 KV cache, see documents below. +[1]: Here "MLA Mode" refers to the mode used for MLA calculation. MQA stands for Multi-Query Attention mode (i.e. `head_dim_k` = 576 (for DeepSeek V3/V3.1/V3.2) or 512 (for DeepSeek V4/V4.1) with `head_dim_v` = 512), while MHA stands for Multi-Head Attention mode (i.e. `head_dim_k` = 192 / 128 with `head_dim_v` = 128). For a detailed explanation of these modes, please refer to the appendix of [DeepSeek V3.2's Paper](https://github.com/deepseek-ai/DeepSeek-V3.2-Exp). -[2]: Here "MLA Mode" refers to the mode used for MLA calculation. MQA stands for Multi-Query Attention mode (i.e. `head_dim_k` = 576 with `head_dim_v` = 512), while MHA stands for Multi-Head Attention mode (i.e. `head_dim_k` = 192 / 128 with `head_dim_v` = 128). For a detailed explanation of these modes, please refer to the appendix of [DeepSeek V3.2's Paper](https://github.com/deepseek-ai/DeepSeek-V3.2-Exp). +[2] Sparse Decoding for DeepSeek V4.1 is only available on SM100 ## Installation @@ -113,13 +126,18 @@ Where - `h_q` is the number of query heads. **FP8 KV Cache:** -If `is_fp8_kvcache` is set to `True`, the kernel reads the KV cache in the "FP8 with scale" format (described below). It dequantizes the cache to bfloat16 and performs attention computation in bfloat16. The output is also in bfloat16. +If `is_fp8_kvcache` is set to `True`, the kernel reads the KV cache in the "FP8 with scale" format (described below). It dequantizes the cache to bfloat16 and performs attention computation in bfloat16. The output is also in bfloat16. In this repository, `is_fp8_kvcache=True` is only supported together with `indices` (i.e. sparse attention); the dense decoding kernel reads bf16 / fp16 KV caches. -In the "FP8 with scale" format, each token's KV cache is 656 Bytes, structured as: +In the "FP8 with scale" format for DeepSeek V3.2 (`head_dim` = 576, sparse attention), a page block is `page_block_size` token-major rows of 656 Bytes each: - **First 512 bytes:** The "quantized NoPE" part, containing 512 `float8_e4m3` values. - **Next 16 bytes:** Scale factors, containing 4 `float32` values. The first `float32` is the scale for the first 128 `float8_e4m3` values, the second for the next 128, and so on. - **Last 128 bytes:** The "RoPE" part, containing 64 `bfloat16` values. This part is not quantized for accuracy. +For DeepSeek V4 / V4.1 (`head_dim` = 512), the format is detected from the last dimension of `k_cache` (i.e. the bytes per token): 584 (V4), 528 (V4.1) or 288 (V4.1 fp4). In all three, a page block stores `page_block_size` data rows first and `page_block_size` scale rows afterwards: +- **V4**: 584 Bytes per token. The data row is 448 Bytes of quantized NoPE (`float8_e4m3`) followed by 128 Bytes, i.e. the 64 `bfloat16` RoPE values (not quantized). The scale row is 8 Bytes, of which the first 7 are `float8_e8m0` scales and the 8th byte is padding; each scale covers 64 consecutive `float8_e4m3` values of the NoPE part. +- **V4.1**: 528 Bytes per token. The data row is 512 Bytes of `float8_e4m3`, i.e. the 64 RoPE dimensions are quantized as well and there is no `bfloat16` part. The scale row is 16 Bytes of `float8_e8m0`, each scale covering 32 consecutive `float8_e4m3` values. +- **V4.1 fp4**: 288 Bytes per token. The data row is 256 Bytes containing 512 `e2m1` values, 2 values per byte (the even-indexed one in the low nibble). The scale row is 32 Bytes of `float8_e4m3`, each scale covering 16 consecutive `e2m1` values. This format is only valid for `extra_k_cache`, and only when `k_cache` is in the V4.1 format; otherwise `extra_k_cache` must have the same format as `k_cache`. In pratice we expect the sliding window (SWA) kv cache to be in FP8 and the compress attention (CA) kv cache to be in FP4. + See `tests/quant.py` for quantization and dequantization details. **Sparse Attention (`indices` tensor):** @@ -134,7 +152,7 @@ The kernel returns `(out, lse)`, where: - `out` is the attention result. - `lse` is the log-sum-exp value of the attention scores for each query head. -See `tests/test_flash_mla_decoding.py` for a complete example. +See `tests/test_flash_mla_dense_decoding.py` and `tests/test_flash_mla_sparse_decoding.py` for complete examples. ### Sparse MLA Prefill @@ -169,7 +187,7 @@ out = S @ focused_kv # [s_q, h_q, d_qk] return (out, max_logits, lse) ``` -See `tests/test_flash_mla_prefill.py` for a complete example. +See `tests/test_flash_mla_sparse_prefill.py` for a complete example. ### Dense MHA Prefill @@ -180,6 +198,114 @@ This kernel implements the standard dense Multi-Head Attention (MHA) forward and The usage is similar to the `flash_attn` package. See `tests/test_fmha_sm100.py` for a complete example. +### Fused norm + RoPE + attn + RoPE + cast kernel + +In the DeepSeek-V4.1 release, we also provide a fused kernel that combines Q-norm (only used in V4, not in V4.1), Q-RoPE, core attention, O-RoPE (conjugate) and the cast to FP8 into a single kernel. It removes the extra time spent on these small kernels while keeping the same or even slightly higher TFlops, at the cost of having to permute the Q_b and Wv weights in advance. + +In DeepSeek-V4.1 attention, Q (`[hidden_size]`) is first projected to `[q_lora_rank]` (the Q_a projection) and then to `[num_attention_heads, head_dim]` (the Q_b projection). After core attention, the output (`[num_attention_heads, head_dim]`) is reshaped to `[o_groups, num_attention_heads // o_groups * head_dim]`, and each of its rows is projected to `[o_lora_rank]` (the Wv projection), giving an `[o_groups, o_lora_rank]` matrix. That matrix is reshaped to `[o_groups * o_lora_rank]` and finally projected to `[hidden_size]` (the Wo projection). This kernel requires the Q_b and Wv weights to be permuted. + +To permute the Q_b weight: + +```python +import torch +import tile_kernels +from flash_mla import fused_norm_rope_attn_rope_cast + +h_q, d_q = 64, 512 # Q heads and Q head dimension +q_lora_rank = 1536 +scale_gran = 128 + +# q_b_proj: [h_q * d_q, q_lora_rank], bfloat16 +q_b_proj = torch.randn((h_q * d_q, q_lora_rank), dtype=torch.bfloat16, device='cuda') + +# Quantize the weight to FP8 with per-token scale factors, in DeepGEMM's layout +q_b_proj_fp8, q_b_sf = tile_kernels.quant.per_token_cast( + q_b_proj, 'e4m3', scale_gran, + use_tma_aligned_col_major_sf=True, round_sf=True, use_packed_ue8m0=True, +) + +# Permute the weight and its scale factors into the layout required by the fused kernel +q_b_proj_fp8, q_b_sf = fused_norm_rope_attn_rope_cast.permute_q_b_proj( + (q_b_proj_fp8, q_b_sf), h_q, d_q, +) +``` + +To permute the Wv weight: + +```python +import deep_gemm +import torch +import tile_kernels +from flash_mla import fused_norm_rope_attn_rope_cast + +n_wv_group, wv_group_size, d_o = 8, 8, 512 # n_wv_group * wv_group_size == h_q +wv_proj_out_dim = 512 # o_lora_rank +scale_gran = 32 + +# wv_proj: [n_wv_group, wv_proj_out_dim, wv_group_size * d_o], bfloat16 +wv_proj = torch.randn((n_wv_group * wv_proj_out_dim, wv_group_size * d_o), + dtype=torch.bfloat16, device='cuda') + +# Quantize the weight to FP8, and put its scale factors into the layout that DeepGEMM's einsum expects +wv_proj_fp8, wv_sf = tile_kernels.quant.per_token_cast( + wv_proj, 'e4m3', scale_gran, + use_tma_aligned_col_major_sf=False, round_sf=True, use_packed_ue8m0=False, +) +wv_sf = deep_gemm.transform_sf_into_required_layout( + wv_sf.view(n_wv_group, wv_proj_out_dim, wv_group_size * d_o // scale_gran), + wv_proj_out_dim, wv_group_size * d_o, + num_groups=n_wv_group, recipe=(1, 1, scale_gran), is_sfa=False, +) +wv_proj_fp8 = wv_proj_fp8.view(n_wv_group, wv_proj_out_dim, wv_group_size * d_o) + +# Permute the weight and its scale factors into the layout required by the fused kernel +wv_proj_fp8, wv_sf = fused_norm_rope_attn_rope_cast.permute_wv_proj( + (wv_proj_fp8, wv_sf), wv_group_size, d_o, +) +``` + +And finally, to use the fused kernel: + +```python +# q: [s_q, h_q, d_qk], bfloat16, i.e. the Q_b projection computed with the permuted weight above +out_fp8, out_sf, max_logits, lse = fused_norm_rope_attn_rope_cast.prefill( + enable_q_norm, # False for DeepSeek-V4.1 + rms_norm_eps, # e.g. 1e-4 + token_positions, # [s_q], int32 + False, 64, cos_sin_cache, # non-neox RoPE with rope_dim = 64 + n_wv_group, # h_q // wv_group_size + 32, # num_per_channels + True, True, True, # use_tma_aligned_col_major_sf, round_sf, use_packed_ue8m0 + q, kv, indices, # bf16 Q, bf16 KV [s_kv, h_kv, d_qk], int32 indices [s_q, h_kv, topk] + sm_scale=sm_scale, + attn_sink=attn_sink, # optional, [h_q], float32 + topk_length=topk_length, # optional, [s_q], int32 +) + +# For decoding, call `decode` instead, passing the paged quantized KV cache: +# q: [s_q, h_q, d_qk], bf16 +# k_cache: [num_blocks, page_block_size, h_kv, bytes_per_token], fp8_e4m3 +# indices_in_kvcache: [s_q, topk], int32 +out_fp8, out_sf, lse = fused_norm_rope_attn_rope_cast.decode( + enable_q_norm, rms_norm_eps, + token_positions, False, 64, cos_sin_cache, + n_wv_group, 32, True, True, True, + q, k_cache, indices_in_kvcache, + sm_scale=sm_scale, + attn_sink=attn_sink, + topk_length=topk_length, + extra_k_cache=extra_k_cache, # optional, same layout as k_cache + extra_indices_in_kvcache=extra_indices_in_kvcache, # optional, [s_q, extra_topk], int32 + extra_topk_length=extra_topk_length, # optional, [s_q], int32 +) + +# The FP8 output is consumed directly by the Wv projection, using the permuted Wv weight +wv_out = torch.empty((s_q, n_wv_group, wv_proj_out_dim), dtype=torch.bfloat16, device='cuda') +deep_gemm.fp8_einsum("bhr,hdr->bhd", (out_fp8, out_sf), (wv_proj_fp8, wv_sf), wv_out, recipe=(1, 1, 32)) +``` + +You may refer to the fused kernel's test script ([tests/test_fused_norm_rope_attn_rope_cast.py](tests/test_fused_norm_rope_attn_rope_cast.py)) for a complete example. + ## Acknowledgement FlashMLA is inspired by [FlashAttention 2&3](https://github.com/dao-AILab/flash-attention/) and [cutlass](https://github.com/nvidia/cutlass) projects. diff --git a/csrc/api/api.cpp b/csrc/api/api.cpp index f43f2a09..8b15e911 100644 --- a/csrc/api/api.cpp +++ b/csrc/api/api.cpp @@ -1,15 +1,18 @@ #include -#include "sparse_fwd.h" -#include "sparse_decode.h" -#include "dense_decode.h" -#include "dense_fwd.h" +void register_sparse_prefill(pybind11::module_& m); +void register_sparse_decode(pybind11::module_& m); +void register_dense_fwd(pybind11::module_& m); +void register_dense_bwd(pybind11::module_& m); +void register_dense_decode(pybind11::module_& m); +void register_fused_norm_rope_attn_rope_cast_fwd(pybind11::module_& m); PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.doc() = "FlashMLA"; - m.def("sparse_decode_fwd", &sparse_attn_decode_interface); - m.def("dense_decode_fwd", &dense_attn_decode_interface); - m.def("sparse_prefill_fwd", &sparse_attn_prefill_interface); - m.def("dense_prefill_fwd", &FMHACutlassSM100FwdRun); - m.def("dense_prefill_bwd", &FMHACutlassSM100BwdRun); + register_sparse_prefill(m); + register_sparse_decode(m); + register_dense_fwd(m); + register_dense_bwd(m); + register_dense_decode(m); + register_fused_norm_rope_attn_rope_cast_fwd(m); } diff --git a/csrc/api/common.h b/csrc/api/common.h index 2c930ed9..5d760545 100644 --- a/csrc/api/common.h +++ b/csrc/api/common.h @@ -9,6 +9,8 @@ #include +#include "kernels/kv_cache_format.h" + static constexpr float LOG_2_E = 1.44269504f; // Instantiation for tensor.data_ptr() @@ -90,8 +92,8 @@ inline int int64_stride_to_int(int64_t orig_stride) { if (MODEL_TYPE == ModelType::V32) { \ static constexpr ModelType CONSTEXPR_NAME = ModelType::V32; \ return __VA_ARGS__(); \ - } else if (MODEL_TYPE == ModelType::MODEL1) { \ - static constexpr ModelType CONSTEXPR_NAME = ModelType::MODEL1; \ + } else if (MODEL_TYPE == ModelType::V4) { \ + static constexpr ModelType CONSTEXPR_NAME = ModelType::V4; \ return __VA_ARGS__(); \ } else { \ TORCH_CHECK(false, "Unsupported model type: ", (int)MODEL_TYPE); \ @@ -140,6 +142,29 @@ static constexpr std::string get_dynamic_enum_name(T value){ return (std::string)names[static_cast(value)]; } +// ============================================= +// Paged quantized KV cache formats (decoding) +// ============================================= + +// The format of a paged quantized KV cache with d_qk = 512 (V4 / V4.1 / V4.1 fp4), detected by bytes_per_token (kv.size(3)) +inline ModelType detect_kv_cache_format_for_headdim_512(int bytes_per_token) { + for (ModelType mt : {ModelType::V4, ModelType::V41, ModelType::V41_FP4}) { + if (bytes_per_token == kv_cache_bytes_per_token(mt)) { + return mt; + } + } + TORCH_CHECK(false, "Unsupported bytes_per_token for d_qk=512: ", bytes_per_token, ". Expected ", + kv_cache_bytes_per_token(ModelType::V4), " (V4), ", kv_cache_bytes_per_token(ModelType::V41), " (V4.1) or ", + kv_cache_bytes_per_token(ModelType::V41_FP4), " (V4.1 fp4)"); +} + +// Dispatches the runtime (kv, extra_kv) format pair +template +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) || ...); + 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_bwd.cpp b/csrc/api/dense_bwd.cpp new file mode 100644 index 00000000..52bda621 --- /dev/null +++ b/csrc/api/dense_bwd.cpp @@ -0,0 +1,9 @@ +#include "common.h" + +#include "kernels/sm100/prefill/dense/interface.h" + +void register_dense_bwd(pybind11::module_& m) { + m.def("dense_prefill_bwd", + &FMHACutlassSM100BwdRun, + "Run Dense Attention Prefill Backward (cutlass FMHA)"); +} diff --git a/csrc/api/dense_decode.h b/csrc/api/dense_decode.cpp similarity index 94% rename from csrc/api/dense_decode.h rename to csrc/api/dense_decode.cpp index 7df178a6..fca4087d 100644 --- a/csrc/api/dense_decode.h +++ b/csrc/api/dense_decode.cpp @@ -1,14 +1,12 @@ -#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> dense_attn_decode_interface( @@ -173,12 +171,12 @@ dense_attn_decode_interface( params.stream = at::cuda::getCurrentCUDAStream().stream(); if (q_dtype == torch::kBFloat16) { - sm90::run_flash_splitkv_mla_kernel(params); + sm90::decode::dense::run_flash_splitkv_mla_kernel(params); } else if (q_dtype == torch::kHalf) { #ifdef FLASH_MLA_DISABLE_FP16 TORCH_CHECK(false, "FlashMLA is compiled with -DFLASH_MLA_DISABLE_FP16. Please remove this flag from your environment and re-compile FlashMLA."); #else - sm90::run_flash_splitkv_mla_kernel(params); + sm90::decode::dense::run_flash_splitkv_mla_kernel(params); #endif } else { TORCH_CHECK(false, "Unsupported dtype for dense MLA on SM90"); @@ -223,3 +221,9 @@ dense_attn_decode_interface( return {out, lse, tile_scheduler_metadata, num_splits}; } + +void register_dense_decode(pybind11::module_& m) { + m.def("dense_decode_fwd", + &dense_attn_decode_interface, + "Run Dense Attention Decode Forward"); +} diff --git a/csrc/api/dense_fwd.cpp b/csrc/api/dense_fwd.cpp new file mode 100644 index 00000000..ba1fc9a9 --- /dev/null +++ b/csrc/api/dense_fwd.cpp @@ -0,0 +1,9 @@ +#include "common.h" + +#include "kernels/sm100/prefill/dense/interface.h" + +void register_dense_fwd(pybind11::module_& m) { + m.def("dense_prefill_fwd", + &FMHACutlassSM100FwdRun, + "Run Dense Attention Prefill Forward (cutlass FMHA)"); +} 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..d188498f --- /dev/null +++ b/csrc/api/fused_norm_rope_attn_rope_cast_fwd.cpp @@ -0,0 +1,547 @@ +#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 at::Tensor allocate_scale_factor(uint32_t batch_size, uint32_t hidden_dim, uint32_t scale_gran, const at::TensorOptions &opts, 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 + 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); + at::Tensor sf = torch::empty({leading_dim, hidden_dim / (scale_gran * 4), aligned_batch_size_for_sf}, opts.dtype(torch::kInt32)); + KU_CHECK_CONTIGUOUS(sf); + sf = sf.transpose(1, 2).slice(1, 0, batch_size); // [leading_dim, batch_size, hidden_dim / (scale_gran*4)], int32 + if (!extra_dim.has_value()) + sf = sf.squeeze(0); + return sf; +} + +static std::vector fused_norm_rope_attn_rope_cast_fwd( + const at::Tensor &q, + const at::Tensor &kv, + const at::Tensor &indices, + float sm_scale, + int d_v, // TODO change int to uint32_t + const std::optional &attn_sink, + const std::optional &topk_length, + bool enable_q_norm, + float rms_norm_eps, + const at::Tensor &token_positions, + bool is_rope_neox_style, + uint32_t rope_dim, + const at::Tensor &cos_sin_cache, + + uint32_t n_wv_group, + uint32_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(); + 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; + + TORCH_CHECK(h_q % n_wv_group == 0, "h_q %% n_wv_group != 0"); + TORCH_CHECK(is_rope_neox_style == false, "Only `is_rope_neox_style == False` is supported"); + TORCH_CHECK(use_tma_aligned_col_major_sf == true, "`use_tma_aligned_col_major_sf` must be True"); + TORCH_CHECK(round_sf == true, "`round_sf` must be True"); + 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, torch::kBFloat16); + KU_CHECK_DTYPE(kv, torch::kBFloat16); + KU_CHECK_DTYPE(indices, torch::kInt32); + KU_CHECK_DTYPE(attn_sink, torch::kFloat32); + KU_CHECK_DTYPE(topk_length, torch::kInt32); + KU_CHECK_DTYPE(token_positions, torch::kInt32); + KU_CHECK_DTYPE(cos_sin_cache, torch::kFloat32); + + 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) + 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); + + TORCH_CHECK(num_per_channels == 32, "num_per_channels must be 32, got ", num_per_channels); + + at::cuda::CUDAGuard device_guard{(char)q.get_device()}; + auto opts = q.options(); + + TORCH_CHECK(d_v % (num_per_channels * 4) == 0); // 4 is the number of uint8 in uint32, since `use_packed_ue8m0` is `True` + at::Tensor out_fp8 = torch::empty({s_q, n_wv_group, wv_group_size * d_v}, opts.dtype(torch::kFloat8_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 + at::Tensor out_sf = allocate_scale_factor(s_q, wv_group_size * d_v, out_sf_scale_gran, opts, n_wv_group).transpose(0, 1); // [s_q, n_wv_group, wv_group_size * d_v / (out_sf_scale_gran*4)] + at::Tensor max_logits = torch::empty({s_q, h_q}, opts.dtype(torch::kFloat)); + at::Tensor lse = torch::empty({s_q, h_q}, opts.dtype(torch::kFloat)); + KU_CHECK_CONTIGUOUS(out_fp8); + 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, + at::cuda::getCurrentCUDAStream().stream(), + + 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)) + }; + + 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}; +} + + +static std::vector fused_norm_rope_attn_rope_cast_decode( + const at::Tensor &q, // [s_q, h_q, d_qk] + const at::Tensor &kv, // [num_blocks, page_block_size, h_kv, bytes_per_token], paged quantized KV cache + const at::Tensor &indices, // [s_q, topk] + float sm_scale, + int 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, + float rms_norm_eps, + const at::Tensor &token_positions, // [s_q] + bool is_rope_neox_style, + uint32_t rope_dim, + const at::Tensor &cos_sin_cache, // [*, rope_dim] + + uint32_t n_wv_group, + uint32_t num_per_channels, + bool use_tma_aligned_col_major_sf, + bool round_sf, + bool use_packed_ue8m0 +) { + Arch arch = Arch(); + 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 + TORCH_CHECK(s_q > 0); + 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); + TORCH_CHECK(h_kv == 1, "Currently only MQA (i.e. h_kv == 1) is supported"); + TORCH_CHECK(d_qk == 512, "Only head_size_k == 512 (V4 / V4.1) is supported"); + TORCH_CHECK(d_v == 512, "Only head_size_v == 512 is supported"); + TORCH_CHECK(topk > 0); + 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) { + TORCH_CHECK(extra_indices.has_value(), "extra_indices must be provided when extra_kv is provided"); + } else { + TORCH_CHECK(!extra_indices.has_value(), "extra_indices must not be provided when extra_kv is not provided"); + TORCH_CHECK(!extra_topk_length.has_value(), "extra_topk_length must not be provided when extra_kv is not provided"); + } + + TORCH_CHECK(is_rope_neox_style == false, "Only `is_rope_neox_style == False` is supported"); + TORCH_CHECK(use_tma_aligned_col_major_sf == true, "`use_tma_aligned_col_major_sf` must be True"); + TORCH_CHECK(round_sf == true, "`round_sf` must be True"); + 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, torch::kBFloat16); + TORCH_CHECK(kv.dtype() == torch::kFloat8_e4m3fn || kv.dtype() == torch::kInt8 || kv.dtype() == torch::kUInt8, "kv must have dtype fp8_e4m3fn, int8 or uint8"); + if (have_extra_kvcache) { + TORCH_CHECK(extra_kv->dtype() == torch::kFloat8_e4m3fn || extra_kv->dtype() == torch::kInt8 || extra_kv->dtype() == torch::kUInt8, "extra_kv must have dtype fp8_e4m3fn, int8 or uint8"); + } + KU_CHECK_DTYPE(indices, torch::kInt32); + KU_CHECK_DTYPE(attn_sink, torch::kFloat32); + KU_CHECK_DTYPE(topk_length, torch::kInt32); + KU_CHECK_DTYPE(extra_indices, torch::kInt32); + KU_CHECK_DTYPE(extra_topk_length, torch::kInt32); + KU_CHECK_DTYPE(token_positions, torch::kInt32); + KU_CHECK_DTYPE(cos_sin_cache, torch::kFloat32); + + // 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) + 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; + TORCH_CHECK(model_type != ModelType::V41_FP4, "The fp4 KV cache is only supported as extra_kv"); + 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)); + 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) { + 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"); + } + 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); + + at::cuda::CUDAGuard device_guard{(char)q.get_device()}; + auto opts = q.options(); + + TORCH_CHECK(d_v % (num_per_channels * 4) == 0); // 4 is the number of uint8 in uint32, since `use_packed_ue8m0` is `True` + at::Tensor out_fp8 = torch::empty({s_q, n_wv_group, wv_group_size * d_v}, opts.dtype(torch::kFloat8_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 + at::Tensor out_sf = allocate_scale_factor(s_q, wv_group_size * d_v, out_sf_scale_gran, opts, n_wv_group).transpose(0, 1); // [s_q, n_wv_group, wv_group_size * d_v / (out_sf_scale_gran*4)] + at::Tensor lse = torch::empty({s_q, h_q}, opts.dtype(torch::kFloat)); + KU_CHECK_CONTIGUOUS(out_fp8); + 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, + at::cuda::getCurrentCUDAStream().stream(), + + 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 { + TORCH_CHECK(false, "Unsupported model_type: ", get_dynamic_enum_name(model_type)); + } + })); + })); + + return {out_fp8, out_sf, lse}; +} + + +static std::vector permute_q_b_proj( + const at::Tensor &q_b_proj, + const at::Tensor &scale_factors, + int h_q, + int 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)); + TORCH_CHECK(gran == 32 || gran == 128, "gran must be 32 or 128, got ", gran); + 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, torch::kFloat8_e4m3fn); + KU_CHECK_DTYPE(scale_factors, torch::kInt32); + + 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); + TORCH_CHECK(scale_factors.stride(0) == 1, "scale_factors must be contiguous on the first dimension"); + + at::cuda::CUDAGuard device_guard{(char)q_b_proj.get_device()}; + + at::Tensor q_b_proj_permuted = torch::empty_like(q_b_proj); + at::Tensor scale_factors_permuted = allocate_scale_factor(h_q_d_q, q_lora_rank, gran, q_b_proj.options()); + 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), + + at::cuda::getCurrentCUDAStream().stream(), + }; + + 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}; +} + + +static std::vector permute_wv_proj( + const at::Tensor &wv_proj, + const at::Tensor &scale_factors, + int wv_group_size, + int 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, torch::kFloat8_e4m3fn); + KU_CHECK_DTYPE(scale_factors, torch::kInt32); + + 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 + TORCH_CHECK(input_gran == 32, "input scale granularity must be 32, got ", input_gran); + 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); + TORCH_CHECK(scale_factors.stride(1) == 1, "scale_factors must be contiguous on the second dimension"); + + at::cuda::CUDAGuard device_guard{(char)wv_proj.get_device()}; + + at::Tensor wv_proj_permuted = torch::empty_like(wv_proj); + at::Tensor scale_factors_permuted = allocate_scale_factor(d_proj_out, wv_group_size * d_o, output_gran, wv_proj.options(), 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), + + at::cuda::getCurrentCUDAStream().stream(), + }; + + 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}; +} + + +void register_fused_norm_rope_attn_rope_cast_fwd(pybind11::module_& m) { + m.def("fused_norm_rope_attn_rope_cast_fwd", + &fused_norm_rope_attn_rope_cast_fwd, + "Run Fused Norm + RoPE + Core Attn + RoPE + Cast"); + m.def("fused_norm_rope_attn_rope_cast_decode", + &fused_norm_rope_attn_rope_cast_decode, + "Run Fused Norm + RoPE + Core Attn + RoPE + Cast (Decoding, with paged quantized KV cache)"); + m.def("permute_q_b_proj", &permute_q_b_proj, + "Permute q_b_proj weight layout for fused norm rope attn rope cast fwd"); + m.def("permute_wv_proj", &permute_wv_proj, + "Permute wv_proj weight layout for fused norm rope attn rope cast fwd"); +} diff --git a/csrc/api/sparse_decode.h b/csrc/api/sparse_decode.cpp similarity index 54% rename from csrc/api/sparse_decode.h rename to csrc/api/sparse_decode.cpp index 6df5c841..7244905c 100644 --- a/csrc/api/sparse_decode.h +++ b/csrc/api/sparse_decode.cpp @@ -1,14 +1,21 @@ -#pragma once - #include "common.h" -#include "params.h" +#include "kernels/params.h" + +#include "kernels/sm90/decode/sparse/splitkv_mla.h" +#include "kernels/sm100/decode/sparse/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" -#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" +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,7 +26,9 @@ enum class DecodeFeatures : int { HEAD_DIM_512, V32_KVCACHE_FORMAT, - MODEL1_KVCACHE_FORMAT, + V4_KVCACHE_FORMAT, + V41_KVCACHE_FORMAT, + V41_FP4_KVCACHE_FORMAT, ATTN_SINK, TOPK_LENGTH, @@ -48,7 +57,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, @@ -69,7 +78,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); }); }); } @@ -81,12 +90,16 @@ 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::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 { @@ -100,8 +113,12 @@ class Decode_Sm100_Head64_Impl : public DecodeImplBase { protected: void run_(const SparseAttnDecodeParams ¶ms, const std::vector &required_features) override { - DISPATCH_MODEL_TYPE(params.model_type, MODEL_TYPE, [&]() { - sm100::decode::head64::run_flash_splitkv_mla_fp8_sparse_kernel(params); + dispatch_kv_formats(SupportedKVFormats{}, params.model_type, params.extra_model_type, [&]() { + DISPATCH_BOOLEAN_FLAG(params.enable_split_kv, ENABLE_SPLIT_KV, ([&]() { + 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); + })); }); } }; @@ -115,12 +132,13 @@ 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::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 { @@ -134,20 +152,25 @@ class Decode_Sm100_Head64x2_Impl : public DecodeImplBase { protected: void run_(const SparseAttnDecodeParams ¶ms, const std::vector &required_features) override { - DISPATCH_MODEL_TYPE(params.model_type, MODEL_TYPE, [&]() { - 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; + 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); } - cur_params.lse += start_head_idx; - cur_params.out += start_head_idx * params.stride_o_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); - } + })); }); } }; @@ -157,12 +180,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 { @@ -176,10 +203,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> sparse_attn_decode_interface( const at::Tensor &q, // [b, s_q, h_q, d_qk] @@ -227,6 +264,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 TORCH_CHECK(b > 0); TORCH_CHECK(s_q > 0); @@ -285,23 +325,23 @@ sparse_attn_decode_interface( // Check shape KU_CHECK_SHAPE(q, b, s_q, h_q, d_qk); - { - int bytes_per_token; - if (d_qk == 576 && d_v == 512) { - // V3.2 style - bytes_per_token = 512 + 64*2 + (512/128)*4; - } else if (d_qk == 512 && d_v == 512) { - // MODEL1 style - bytes_per_token = 448 + 64*2 + (448/64)*1 + 1; - } else { - TORCH_CHECK(false, "Unsupported head sizes for is_fp8_kvcache == True"); - } - KU_CHECK_SHAPE(kv, num_blocks, page_block_size, h_kv, bytes_per_token); - KU_CHECK_SHAPE(extra_kv, extra_num_blocks, extra_page_block_size, h_kv, bytes_per_token); - 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()) { - 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 = extra_model_type = ModelType::V32; + } 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 { + TORCH_CHECK(false, "Unsupported head sizes for is_fp8_kvcache == True"); + } + TORCH_CHECK(model_type != ModelType::V41_FP4, "The fp4 KV cache is only supported as extra_kv"); + 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)); + 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)); + 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) { + 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); @@ -315,15 +355,6 @@ sparse_attn_decode_interface( at::Tensor out = torch::empty({b, s_q, h_q, d_v}, opts); at::Tensor lse = torch::empty({b, s_q, h_q}, opts.dtype(at::kFloat)); - ModelType model_type; - if (d_qk == 576) { - model_type = ModelType::V32; - } else if (d_qk == 512) { - model_type = ModelType::MODEL1; - } else { - TORCH_CHECK(false, "Unsupported d_qk: ", d_qk); - } - std::vector features; if (h_q == 64) { features.push_back(DecodeFeatures::HEAD_64); @@ -339,13 +370,6 @@ sparse_attn_decode_interface( } else { 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 { - TORCH_CHECK(false, "Unsupported model type: ", (int)model_type); - } if (have_attn_sink) { features.push_back(DecodeFeatures::ATTN_SINK); } @@ -358,6 +382,19 @@ 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 { + TORCH_CHECK(false, "Unsupported model type: ", (int)mt); + } + } DecodeImplBase* impl; if (arch.is_sm100f()) { @@ -386,7 +423,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(), @@ -411,85 +448,95 @@ 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, - at::cuda::getCurrentCUDAStream().stream() + at::cuda::getCurrentCUDAStream().stream(), + + enable_split_kv, }; - // Get MLA metadata if necessary at::Tensor o_accum, lse_accum; - if (!tile_scheduler_metadata.has_value()) { - tile_scheduler_metadata = torch::empty({impl_meta.num_sm_parts, sizeof(DecodingSchedMeta)/4}, opts.dtype(torch::kInt32)); - num_splits = torch::empty({b+1}, opts.dtype(torch::kInt32)); + if (enable_split_kv) { + // Get MLA metadata if necessary + if (!tile_scheduler_metadata.has_value()) { + tile_scheduler_metadata = torch::empty({impl_meta.num_sm_parts, sizeof(DecodingSchedMeta)/4}, opts.dtype(torch::kInt32)); + num_splits = torch::empty({b+1}, opts.dtype(torch::kInt32)); + 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->data_ptr(), + impl_meta.num_sm_parts, + at::cuda::getCurrentCUDAStream().stream() + }; + 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, torch::kInt32); + KU_CHECK_DTYPE(num_splits, torch::kInt32); 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->data_ptr(), - impl_meta.num_sm_parts, - at::cuda::getCurrentCUDAStream().stream() - }; - smxx::decode::run_get_decoding_sched_meta_kernel(get_sched_meta_params); + 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->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::empty({total_num_splits, s_q, h_q}, opts.dtype(at::kFloat)); + o_accum = torch::empty({total_num_splits, s_q, h_q, d_v}, opts.dtype(at::kFloat)); + KU_CHECK_CONTIGUOUS(lse_accum); + KU_CHECK_CONTIGUOUS(o_accum); + params.lse_accum = lse_accum.data_ptr(); + params.o_accum = o_accum.data_ptr(); + params.stride_lse_accum_split = int64_stride_to_int(lse_accum.stride(0)); + params.stride_lse_accum_s_q = int64_stride_to_int(lse_accum.stride(1)); + params.stride_o_accum_split = int64_stride_to_int(o_accum.stride(0)); + params.stride_o_accum_s_q = int64_stride_to_int(o_accum.stride(1)); + params.stride_o_accum_h_q = int64_stride_to_int(o_accum.stride(2)); } - // Stick the metadata pointers to `params` - KU_CHECK_DEVICE(tile_scheduler_metadata); - KU_CHECK_DEVICE(num_splits); - KU_CHECK_DTYPE(tile_scheduler_metadata, torch::kInt32); - KU_CHECK_DTYPE(num_splits, torch::kInt32); - 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->data_ptr(); - params.num_sm_parts = impl_meta.num_sm_parts; - - // Allocate intermediate buffers for split-KV - const int total_num_splits = b + impl_meta.num_sm_parts; - lse_accum = torch::empty({total_num_splits, s_q, h_q}, opts.dtype(at::kFloat)); - o_accum = torch::empty({total_num_splits, s_q, h_q, d_v}, opts.dtype(at::kFloat)); - KU_CHECK_CONTIGUOUS(lse_accum); - KU_CHECK_CONTIGUOUS(o_accum); - params.lse_accum = lse_accum.data_ptr(); - params.o_accum = o_accum.data_ptr(); - params.stride_lse_accum_split = int64_stride_to_int(lse_accum.stride(0)); - params.stride_lse_accum_s_q = int64_stride_to_int(lse_accum.stride(1)); - params.stride_o_accum_split = int64_stride_to_int(o_accum.stride(0)); - params.stride_o_accum_s_q = int64_stride_to_int(o_accum.stride(1)); - params.stride_o_accum_h_q = int64_stride_to_int(o_accum.stride(2)); impl->run(params, features); - - CombineParams combine_params = { - b, s_q, h_q, d_v, + if (enable_split_kv) { + 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), - at::cuda::getCurrentCUDAStream().stream() - }; - smxx::decode::run_flash_mla_combine_kernel(combine_params); + ku::get_optional_tensor_ptr(attn_sink), + at::cuda::getCurrentCUDAStream().stream() + }; + smxx::decode::run_flash_mla_combine_kernel(combine_params); + } delete impl; return {out, lse.transpose(1, 2), tile_scheduler_metadata, num_splits}; } + +void register_sparse_decode(pybind11::module_& m) { + m.def("sparse_decode_fwd", + &sparse_attn_decode_interface, + "Run Sparse Attention Decode Forward"); +} diff --git a/csrc/api/sparse_fwd.h b/csrc/api/sparse_prefill.cpp similarity index 87% rename from csrc/api/sparse_fwd.h rename to csrc/api/sparse_prefill.cpp index 66d7111e..b9dd72cf 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); + 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,13 +83,12 @@ 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); } }; @@ -241,3 +235,9 @@ static std::vector sparse_attn_prefill_interface( return {out, max_logits, lse}; } + +void register_sparse_prefill(pybind11::module_& m) { + m.def("sparse_prefill_fwd", + &sparse_attn_prefill_interface, + "Run Sparse Attention Prefill Forward"); +} 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..e0127c1b --- /dev/null +++ b/csrc/kernels/kv_cache_format.h @@ -0,0 +1,56 @@ +#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_FP4 = MT == ModelType::V41_FP4; + static constexpr int D_QK = IS_V32 ? 576 : 512; + static constexpr int D_ROPE = 64; + static constexpr int D_NOPE = D_QK - D_ROPE; + static constexpr int D_FP4 = IS_FP4 ? D_QK : 0; // Dimensions stored as fp4 e2m1 + static constexpr int D_FP8 = IS_FP4 ? 0 : (MT == ModelType::V41 ? D_QK : D_NOPE); // Dimensions stored as fp8 e4m3, needing dequant + 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_FP4 ? 16 : (MT == ModelType::V41 ? 32 : (MT == ModelType::V4 ? 64 : 128)); // Dimensions sharing one scale + static constexpr int NUM_SCALES_EACH_TOKEN = (IS_V32 ? D_NOPE : D_QK) / QUANT_TILE_SIZE; // 4 / 8 (7 + 1 byte of 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 = QUANT_BYTES + (IS_V32 ? NUM_SCALES_EACH_TOKEN * SCALE_BYTES : 0) + 2 * D_BF16; + static constexpr int BYTES_PER_TOKEN = TMA_K_STRIDE + (IS_V32 ? 0 : NUM_SCALES_EACH_TOKEN * SCALE_BYTES); // 656 / 584 / 528 / 288 +}; + +// 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; + } + 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 89% rename from csrc/params.h rename to csrc/kernels/params.h index 4433e8d4..74e619ff 100644 --- a/csrc/params.h +++ b/csrc/kernels/params.h @@ -3,8 +3,10 @@ #include "cutlass/bfloat16.h" enum class ModelType { - V32, - MODEL1 + 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) }; struct __align__(4*8) DecodingSchedMeta { @@ -66,7 +68,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] @@ -91,7 +94,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] @@ -170,11 +175,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/sm100/decode/head64/config.h b/csrc/kernels/sm100/decode/sparse/head64/config.h similarity index 51% rename from csrc/sm100/decode/head64/config.h rename to csrc/kernels/sm100/decode/sparse/head64/config.h index 401f3acf..8a633300 100644 --- a/csrc/sm100/decode/head64/config.h +++ b/csrc/kernels/sm100/decode/sparse/head64/config.h @@ -8,10 +8,12 @@ #include -#include "defines.h" -#include "params.h" +#include "kernels/defines.h" +#include "kernels/kv_cache_format.h" +#include "kernels/sm100/dequant_utils.cuh" -namespace sm100::decode::head64 { + +namespace sm100::decode::sparse::head64 { using cutlass::arch::fence_view_async_shared; using cutlass::arch::NamedBarrier; @@ -27,32 +29,53 @@ enum NamedBarriers : uint32_t { everyone_sync = 4 }; -template +template struct KernelTemplate { -static constexpr int D_Q = MODEL_TYPE == ModelType::V32 ? 576 : 512; +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 = MODEL_TYPE == ModelType::V32 ? 512 : 448; -static constexpr int D_ROPE = 64; -static constexpr int QUANT_TILE_SIZE = MODEL_TYPE == ModelType::V32 ? 128 : 64; -static constexpr bool V_HAVE_ROPE = MODEL_TYPE == ModelType::V32 ? false : true; -static constexpr int NUM_SCALES_EACH_TOKEN = MODEL_TYPE == ModelType::V32 ? 4 : 8; // Padding is included -static constexpr int TMA_K_STRIDE = MODEL_TYPE == ModelType::V32 ? D_NOPE+2*D_ROPE+4*(D_NOPE/QUANT_TILE_SIZE) : D_NOPE+2*D_ROPE; // 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 and 576 for MODEL1. Extra padding may be necessary for KV blocks. +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_H = 64; 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 = MODEL_TYPE == ModelType::V32 ? 64 : 0; +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 = MODEL_TYPE == ModelType::V32 ? 64 : 128; // RoPE part stored in SW64 (for V32) or SW128 (for MODEL1), in bytes +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, @@ -62,16 +85,16 @@ 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_nope; - CUtensorMap tensor_map_kv_rope; - CUtensorMap tensor_map_extra_kv_nope; - CUtensorMap tensor_map_extra_kv_rope; + 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 + 64*D_Q/256: Q + // 256 ~ 256 + B_H*D_Q/256: Q // 400 ~ 464: P static constexpr int O = 0; static constexpr int Q = 256; @@ -113,14 +136,14 @@ using SmemLayoutS = decltype(tile_to_shape( template using SmemLayoutKTiles_SW128 = decltype(coalesce(tile_to_shape( UMMA::Layout_K_SW128_Atom{}, - Shape, Int<64*NUM_TILES>>{}, + 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>>{}, + Shape, Int<64*NUM_TILES>>{}, Step<_1, _2>{} ), Shape<_1, _1>{})); @@ -136,14 +159,14 @@ using SmemLayoutKTilesTransposed_SW128 = decltype(composition( template using SmemLayoutKTiles_SW64 = decltype(coalesce(tile_to_shape( UMMA::Layout_K_SW64_Atom{}, - Shape, Int<32*NUM_TILES>>{}, + 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>>{}, + Shape, Int<32*NUM_TILES>>{}, Step<_1, _2>{} ), Shape<_1, _1>{})); @@ -168,26 +191,27 @@ struct SharedMemoryPlan { } qo; struct { struct { - array_aligned nope; // NoPE part, dequantized - array_aligned rope; // RoPE part, dequantized. SW64 in v32 mode, SW128 in MODEL1 mode + 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 covers raw_nope - array_aligned raw_nope[NUM_BUFS]; // Raw (quantized) NoPE part + 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 { - float4 p_exchange_buf[4][16 * B_TOPK / 4]; + 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]; - e8m0 scales[NUM_INDEX_BUFS][B_TOPK][NUM_SCALES_EACH_TOKEN]; + 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_rope_ready[NUM_BUFS]; - transac_bar_t bar_nope_ready[NUM_BUFS]; + 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]; @@ -197,13 +221,14 @@ 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{} + 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); +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/sm100/decode/head64/kernel.cuh b/csrc/kernels/sm100/decode/sparse/head64/kernel.cuh similarity index 58% rename from csrc/sm100/decode/head64/kernel.cuh rename to csrc/kernels/sm100/decode/sparse/head64/kernel.cuh index 7c46921c..0e5b68d9 100644 --- a/csrc/sm100/decode/head64/kernel.cuh +++ b/csrc/kernels/sm100/decode/sparse/head64/kernel.cuh @@ -1,3 +1,17 @@ +/* +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 @@ -9,18 +23,20 @@ #include "kerutils/kerutils.cuh" -#include "utils.h" -#include "sm100/helpers.h" +#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::head64 { +namespace sm100::decode::sparse::head64 { -template +template template __device__ void -KernelTemplate -::flash_fwd_splitkv_mla_fp8_sparse_kernel_devfunc(const SparseAttnDecodeParams ¶ms, const TmaParam &tma_params) { +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; @@ -36,8 +52,10 @@ KernelTemplate 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_nope); - cute::prefetch_tma_descriptor(&tma_params.tensor_map_kv_rope); + 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) { @@ -46,8 +64,9 @@ KernelTemplate plan.bar_q_tma.init(1); plan.bar_q_utccp.init(1); for (int i = 0; i < NUM_BUFS; ++i) { - plan.bar_rope_ready[i].init(1); - plan.bar_nope_ready[i].init(128); + 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); @@ -56,7 +75,7 @@ KernelTemplate } 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+1); + plan.bar_valid_coord_scale_free[i].init(128+128+1+(D_BF16 > 0)); } cutlass::arch::fence_barrier_init(); } @@ -78,14 +97,18 @@ KernelTemplate // 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; - KU_LDG_256( - params.tile_scheduler_metadata_ptr + partition_idx, - &sched_meta, - ".nc", - "no_allocate", - "evict_normal", - "256B" - ); + 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; @@ -94,14 +117,23 @@ KernelTemplate 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 - int start_block_idx = batch_idx == sched_meta.begin_req_idx ? sched_meta.begin_block_idx : 0; - int end_block_idx = batch_idx == sched_meta.end_req_idx ? sched_meta.end_block_idx : total_topk_padded / B_TOPK; - bool 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); - int 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); + 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, @@ -141,12 +173,16 @@ KernelTemplate 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%64, (idx_in_warpgroup/64)*128 + i*8); + 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); - float attn_sink = params.attn_sink == nullptr ? -CUDART_INF_F : __ldg((float*)params.attn_sink + (idx_in_warpgroup%64)) * CUDART_L2E_F; + // 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>(); @@ -163,64 +199,23 @@ KernelTemplate plan.bar_qk_done[rs.buf_idx].wait(rs.bar_phase); ku::tcgen05_after_thread_sync(); - // Load P - float p[B_TOPK/2], p_peer[B_TOPK/2]; - if (warp_idx < 2) { - ku::tmem_ld_32dp32bNx(tmem_cols::P, p); - ku::tmem_ld_32dp32bNx(tmem_cols::P+32, p_peer); - } else { - ku::tmem_ld_32dp32bNx(tmem_cols::P, p_peer); - ku::tmem_ld_32dp32bNx(tmem_cols::P+32, p); - } - cutlass::arch::fence_view_async_tmem_load(); - ku::tcgen05_before_thread_sync(); - // Reduce within shared mem - { - // Store - // Warp 0, 1 store their right (col 32 ~ 63) part, while warp 2, 3 store their left (row 0 ~ 31) part - CUTE_UNROLL - for (int i = 0; i < (B_TOPK/2)/4; ++i) - plan.s_p.p_exchange_buf[warp_idx^2][i*32 + lane_idx] = *(float4*)(p_peer + i*4); - NamedBarrier::arrive_and_wait(64, NamedBarriers::wg0_warp02_sync+(warp_idx&1)); // Synchronize between warp 0 and warp 2, as well as warp 1 - warp 3 - // Load - CUTE_UNROLL - for (int i = 0; i < (B_TOPK/2)/4; ++i) { - float2 t[2]; - *(float4*)t = plan.s_p.p_exchange_buf[warp_idx][i*32 + lane_idx]; - float2* cur_p = (float2*)(p + i*4); - cur_p[0] = ku::float2_add(cur_p[0], t[0]); - cur_p[1] = ku::float2_add(cur_p[1], t[1]); - } - } - - // Since dual gemm is utilized, the layout of P in register now look like: - // - // 32 32 - // +-------+-------+ - // | | | - // 32 | Warp0 | Warp2 | - // | | | - // +-------+-------+ - // | | | - // 32 | Warp1 | Warp3 | - // | | | - // +-------+-------+ - - // Mask - uint32_t valid_mask = *((uint32_t*)plan.is_token_valid[rs.index_buf_idx] + (idx_in_warpgroup>=64?1:0)); - CUTE_UNROLL - for (int i = 0; i < B_TOPK/2; i += 1) { - if (!(valid_mask>>i&1)) - p[i] = -CUDART_INF_F; - } + 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 + ); - // Get rowwise max of Pi - float cur_pi_max = -CUDART_INF_F; - CUTE_UNROLL - for (int i = 0; i < (B_TOPK/2); i += 1) { - cur_pi_max = max(cur_pi_max, p[i]); - } + 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; @@ -230,8 +225,8 @@ KernelTemplate 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, ...) - // - should_scale_o is identical among every warp, and is identical among threads that controls the same row (i.e. among threads 0~31+64~95; and is identical among threads 32~63+96~127) + // - 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; @@ -246,11 +241,11 @@ KernelTemplate mi = new_max; // mi is still identical within each row // Calculate S - __nv_bfloat162 s[(B_TOPK/2)/2]; + __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 < (B_TOPK/2)/2; i += 1) { + 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); @@ -261,32 +256,19 @@ KernelTemplate // Write S CUTE_UNROLL - for (int i = 0; i < B_TOPK/2/8; i += 1) { + 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) { - float2 scale_for_old_float2 = float2 {scale_for_old, scale_for_old}; ku::tcgen05_after_thread_sync(); - static constexpr int CHUNK_SIZE = 64; - float2 o[CHUNK_SIZE/2]; - CUTE_UNROLL - for (int chunk_idx = 0; chunk_idx < (D_V/2)/CHUNK_SIZE; ++chunk_idx) { - // Load O - ku::tmem_ld_32dp32bNx(tmem_cols::O + chunk_idx*CHUNK_SIZE, o); - cutlass::arch::fence_view_async_tmem_load(); - - // Mult - for (int i = 0; i < CHUNK_SIZE/2; ++i) { - o[i] = ku::float2_mul(o[i], scale_for_old_float2); - } - - // Store O - ku::tmem_st_32dp32bNx(tmem_cols::O + chunk_idx*CHUNK_SIZE, o); - cutlass::arch::fence_view_async_tmem_store(); - } + rescale_O< + D_V / (128/B_H), + 64, + tmem_cols::O + >(scale_for_old); ku::tcgen05_before_thread_sync(); } @@ -296,6 +278,7 @@ KernelTemplate 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) { @@ -311,9 +294,9 @@ KernelTemplate li += plan.rowwise_max_buf[idx_in_warpgroup^64]; // Store li - if (idx_in_warpgroup < B_H) { + if (idx_in_warpgroup < params.h_q) { if (args.is_no_split) { - float cur_lse = fmaf(mi, CUDART_LN2_F, logf(li)); + 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; @@ -332,6 +315,7 @@ KernelTemplate 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), @@ -348,7 +332,7 @@ KernelTemplate float2 o[B_EPI/2]; __nv_bfloat162 o_bf16[B_EPI/2]; CUTE_UNROLL - for (int i = 0; i < (D_V/2) / B_EPI; ++i) { + 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(); @@ -359,7 +343,7 @@ KernelTemplate o_bf16[j] = __float22bfloat162_rn(o[j]); } // Store - int col_base = (i*B_EPI>=D_V/4 ? D_V/2 : 0) + (i*B_EPI%(D_V/4)); + 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]); @@ -367,18 +351,12 @@ KernelTemplate fence_view_async_shared(); NamedBarrier::arrive_and_wait(128, NamedBarriers::wg0_sync); // S -> G - if (warp_idx == 0 && elect_one_sync()) { + 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(_, _, col_base/64)), - thr_tma.partition_D(tma_gO(_, _, col_base/64)) - ); - } - if (warp_idx == 1 && elect_one_sync()) { - cute::copy( - tma_params.tma_O, - thr_tma.partition_S(tma_sO(_, _, col_base/64 + (D_V/4)/64)), - thr_tma.partition_D(tma_gO(_, _, col_base/64 + (D_V/4)/64)) + thr_tma.partition_S(tma_sO(_, _, copy_block_idx)), + thr_tma.partition_D(tma_gO(_, _, copy_block_idx)) ); } } @@ -390,7 +368,7 @@ KernelTemplate 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/2) / B_EPI; ++i) { + 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(); @@ -399,10 +377,10 @@ KernelTemplate 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/64)*128 + (i*B_EPI >= D_V/4 ? D_V/2 : 0) + (i*B_EPI%(D_V/4)); + 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%64, col_base + j*4) = *(__int128_t*)(&o[j*2]); + *(__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); @@ -410,11 +388,13 @@ KernelTemplate CUTE_UNROLL for (int local_row = 0; local_row < B_H/4; ++local_row) { int smem_row = local_row*4 + warp_idx; - 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) - ); + 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(); } @@ -473,11 +453,11 @@ KernelTemplate ) ); static_assert(D_Q_SW128%128 == 0); - CUTE_UNROLL + #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 < 64/16; ++subtile_idx) { + 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, @@ -529,29 +509,29 @@ KernelTemplate // Mainloop CUTE_NO_UNROLL for (int block_idx = args.start_block_idx; block_idx < args.end_block_idx; ++block_idx) { - if constexpr (MODEL_TYPE == ModelType::V32) { + if constexpr (OrigKVFormat::MODEL_TYPE == ModelType::V32) { // V3.2: RoPE behaves like an extra block with size 64, so we can do RoPE first // QK RoPE - plan.bar_rope_ready[rs.buf_idx].wait(rs.bar_phase); + plan.bar_bf16_part_load_ready[rs.buf_idx].wait(rs.bar_phase); ku::tcgen05_after_thread_sync(); Tensor tQ_rope = tiled_mma_P.get_slice(_0{}).make_fragment_A( partition_shape_A(tiled_mma_P, Shape, Int>{}) ); tQ_rope.data().get() = tmem_cols::Q_Tail; - Tensor sK_rope = make_tensor(make_smem_ptr(plan.u.kv.dequant[rs.buf_idx].rope.data()), SmemLayoutKTiles_DualGemm_SW64<2/2>{}); + Tensor sK_rope = make_tensor(make_smem_ptr(plan.u.kv.dequant[rs.buf_idx].bf16_part), SmemLayoutKTiles_DualGemm_SW64<2/2>{}); ku::utcmma_ts(tiled_mma_P, tQ_rope, sK_rope, tP, true); // QK NoPE - plan.bar_nope_ready[rs.buf_idx].wait(rs.bar_phase); + plan.bar_quant_part_dequant_ready[rs.buf_idx].wait(rs.bar_phase); ku::tcgen05_after_thread_sync(); Tensor tQ_nope = tiled_mma_P.get_slice(_0{}).make_fragment_A( partition_shape_A(tiled_mma_P, Shape, Int>{}) ); tQ_nope.data().get() = tmem_cols::Q; - Tensor sK_nope = make_tensor(make_smem_ptr(plan.u.kv.dequant[rs.buf_idx].nope.data()), SmemLayoutKTiles_DualGemm_SW128<512/64/2>{}); + Tensor sK_nope = make_tensor(make_smem_ptr(plan.u.kv.dequant[rs.buf_idx].quant_part), SmemLayoutKTiles_DualGemm_SW128<512/64/2>{}); ku::utcmma_ts(tiled_mma_P, tQ_nope, sK_nope, tP, false); } else { - // MODEL1: RoPE is the last 64 dims within the full 512 dim, which couples with the last 64 dim from the NoPE part when performing dual GEMM. i.e. + // V4: RoPE is the last 64 dims within the full 512 dim, which couples with the last 64 dim from the NoPE part when performing dual GEMM. i.e. // // logical view: |0|1|2|3|4|5|6|7| (where 7 is the RoPE part) // dual gemm's view: @@ -559,15 +539,16 @@ KernelTemplate // |1|3|5|7| // // So we must wait for both the NoPE and the RoPE part, and then perform dual GEMM - plan.bar_rope_ready[rs.buf_idx].wait(rs.bar_phase); - plan.bar_nope_ready[rs.buf_idx].wait(rs.bar_phase); + if constexpr (D_BF16 > 0) + plan.bar_bf16_part_load_ready[rs.buf_idx].wait(rs.bar_phase); + plan.bar_quant_part_dequant_ready[rs.buf_idx].wait(rs.bar_phase); ku::tcgen05_after_thread_sync(); Tensor tQ = tiled_mma_P.get_slice(_0{}).make_fragment_A( partition_shape_A(tiled_mma_P, Shape, Int>{}) ); tQ.data().get() = tmem_cols::Q; - Tensor sK = make_tensor(make_smem_ptr(plan.u.kv.dequant[rs.buf_idx].nope.data()), SmemLayoutKTiles_DualGemm_SW128<512/64/2>{}); + Tensor sK = make_tensor(make_smem_ptr(plan.u.kv.dequant[rs.buf_idx].quant_part), SmemLayoutKTiles_DualGemm_SW128<512/64/2>{}); ku::utcmma_ts(tiled_mma_P, tQ, sK, tP, true); } ku::umma_arrive_noelect(plan.bar_qk_done[rs.buf_idx]); @@ -576,22 +557,23 @@ KernelTemplate plan.bar_so_ready[rs.buf_idx].wait(rs.bar_phase); ku::tcgen05_after_thread_sync(); Tensor sS = make_tensor(make_smem_ptr(plan.s_p.s.data()), SmemLayoutS{}); - Tensor sV = make_tensor(make_smem_ptr(plan.u.kv.dequant[rs.buf_idx].nope.data()), SmemLayoutKTilesTransposed_SW128{}); // NOTE: For MODEL1, it "expands" to the RoPE part. + Tensor sV = make_tensor(make_smem_ptr(plan.u.kv.dequant[rs.buf_idx].quant_part), SmemLayoutKTilesTransposed_SW128{}); // NOTE: For V4, it "expands" to the RoPE part. ku::utcmma_ss(tiled_mma_O, sS, sV, tO, block_idx == args.start_block_idx); ku::umma_arrive_noelect(plan.bar_sv_done[rs.buf_idx]); rs.update(); } }); + } else if (warp_idx == 5 && elect_one_sync()) { - // Raw KV NoPE retrieval warp + // Raw (quantized) KV retrieval warp run_main_loop([&](const MainLoopArgs &args) { plan.bar_q_utccp.wait(args.bar_phase_batch_rel); plan.bar_last_store_done.wait(args.bar_phase_batch_rel); - CUTE_NO_UNROLL - for (int block_idx = args.start_block_idx; block_idx < args.end_block_idx; ++block_idx) { + for_each_kv_block(args.start_block_idx, args.end_block_idx, args.num_orig_kv_blocks, [&](int block_idx, bool is_extra_block) { plan.bar_valid_coord_scale_ready[rs.index_buf_idx].wait(rs.index_bar_phase); plan.bar_raw_free[rs.buf_idx].wait(rs.bar_phase^1); + const CUtensorMap *tensor_map = is_extra_block ? &tma_params.tensor_map_extra_kv_quant_part : &tma_params.tensor_map_kv_quant_part; int4 cur_indices = *(int4*)(plan.tma_coord[rs.index_buf_idx] + 0); int4 nxt_cur_indices; CUTE_UNROLL @@ -599,72 +581,74 @@ KernelTemplate if (row+4 < B_TOPK) nxt_cur_indices = *(int4*)(plan.tma_coord[rs.index_buf_idx] + row + 4); ku::tma_gather4( - block_idx >= args.num_orig_kv_blocks ? &tma_params.tensor_map_extra_kv_nope : &tma_params.tensor_map_kv_nope, + tensor_map, plan.bar_raw_ready[rs.buf_idx], - plan.u.kv.raw_nope[rs.buf_idx].data() + D_NOPE*row, + plan.u.kv.raw_quant[rs.buf_idx].data() + RAW_TOKEN_SMEM_STRIDE*row, 0, cur_indices, - (int64_t)TMA::CacheHintSm90::EVICT_LAST + (int64_t)TMA::CacheHintSm90::EVICT_FIRST ); cur_indices = nxt_cur_indices; } - plan.bar_raw_ready[rs.buf_idx].arrive_and_expect_tx(B_TOPK*D_NOPE*sizeof(e4m3)); + plan.bar_raw_ready[rs.buf_idx].arrive_and_expect_tx(B_TOPK*RAW_TOKEN_SMEM_STRIDE); plan.bar_valid_coord_scale_free[rs.index_buf_idx].arrive(); rs.update(); - } + }); }); } else if (warp_idx == 6 && elect_one_sync()) { // KV RoPE retrieval warp - run_main_loop([&](const MainLoopArgs &args) { - plan.bar_q_utccp.wait(args.bar_phase_batch_rel); - plan.bar_last_store_done.wait(args.bar_phase_batch_rel); - CUTE_NO_UNROLL - for (int block_idx = args.start_block_idx; block_idx < args.end_block_idx; ++block_idx) { - plan.bar_valid_coord_scale_ready[rs.index_buf_idx].wait(rs.index_bar_phase); - if constexpr (MODEL_TYPE == ModelType::V32) { - plan.bar_qk_done[rs.buf_idx].wait(rs.bar_phase^1); - } else { - plan.bar_sv_done[rs.buf_idx].wait(rs.bar_phase^1); - } - int4 cur_indices = *(int4*)(plan.tma_coord[rs.index_buf_idx] + 0); - int4 nxt_cur_indices; - CUTE_UNROLL - for (int row = 0; row < B_TOPK; row += 4) { - if (row+4 < B_TOPK) - nxt_cur_indices = *(int4*)(plan.tma_coord[rs.index_buf_idx] + row + 4); + if constexpr (D_BF16 > 0) { + run_main_loop([&](const MainLoopArgs &args) { + plan.bar_q_utccp.wait(args.bar_phase_batch_rel); + plan.bar_last_store_done.wait(args.bar_phase_batch_rel); + CUTE_NO_UNROLL + for (int block_idx = args.start_block_idx; block_idx < args.end_block_idx; ++block_idx) { + plan.bar_valid_coord_scale_ready[rs.index_buf_idx].wait(rs.index_bar_phase); + if constexpr (OrigKVFormat::MODEL_TYPE == ModelType::V32) { + plan.bar_qk_done[rs.buf_idx].wait(rs.bar_phase^1); + } else { + plan.bar_sv_done[rs.buf_idx].wait(rs.bar_phase^1); + } + int4 cur_indices = *(int4*)(plan.tma_coord[rs.index_buf_idx] + 0); + int4 nxt_cur_indices; CUTE_UNROLL - for (int t = 0; t < D_ROPE/(K_ROPE_SW/2); ++t) { - ku::tma_gather4( - block_idx >= args.num_orig_kv_blocks ? &tma_params.tensor_map_extra_kv_rope : &tma_params.tensor_map_kv_rope, - plan.bar_rope_ready[rs.buf_idx], - plan.u.kv.dequant[rs.buf_idx].rope.data() + (K_ROPE_SW/2)*row + t*B_TOPK*(K_ROPE_SW/2), - t*(K_ROPE_SW/2), - cur_indices, - (int64_t)TMA::CacheHintSm90::EVICT_LAST - ); + for (int row = 0; row < B_TOPK; row += 4) { + if (row+4 < B_TOPK) + nxt_cur_indices = *(int4*)(plan.tma_coord[rs.index_buf_idx] + row + 4); + CUTE_UNROLL + for (int t = 0; t < D_BF16/(K_ROPE_SW/2); ++t) { + ku::tma_gather4( + block_idx >= args.num_orig_kv_blocks ? &tma_params.tensor_map_extra_kv_bf16_part : &tma_params.tensor_map_kv_bf16_part, + plan.bar_bf16_part_load_ready[rs.buf_idx], + plan.u.kv.dequant[rs.buf_idx].bf16_part + (K_ROPE_SW/2)*row + t*B_TOPK*(K_ROPE_SW/2), + t*(K_ROPE_SW/2), + cur_indices, + (int64_t)TMA::CacheHintSm90::EVICT_FIRST + ); + } + cur_indices = nxt_cur_indices; } - cur_indices = nxt_cur_indices; + plan.bar_bf16_part_load_ready[rs.buf_idx].arrive_and_expect_tx(B_TOPK*D_BF16*sizeof(bf16)); + plan.bar_valid_coord_scale_free[rs.index_buf_idx].arrive(); + rs.update(); } - plan.bar_rope_ready[rs.buf_idx].arrive_and_expect_tx(B_TOPK*D_ROPE*sizeof(bf16)); - plan.bar_valid_coord_scale_free[rs.index_buf_idx].arrive(); - rs.update(); - } - }); + }); + } } else if (warp_idx == 7) { // Indices transformation warp // Responsible for generating: TMA coordinates, scale factors, and valid masks static_assert(B_TOPK == 64); - static constexpr int tma_coords_step_per_token = MODEL_TYPE == ModelType::V32 ? 656/TMA_K_STRIDE : 576/TMA_K_STRIDE; - int tma_coords_step_per_block = params.stride_kv_block / TMA_K_STRIDE; // must < 2G since k_batch_stride < 1T and TMA_K_STRIDE > 512 - int tma_coords_step_per_extra_block = params.stride_extra_kv_block / TMA_K_STRIDE; + static constexpr int tma_coords_step_per_token = 1; // A token's data is exactly one TMA_K_STRIDE row, for all formats + int tma_coords_step_per_block = params.stride_kv_block / TMA_K_STRIDE; // must < 2G since k_batch_stride < 1T and TMA_K_STRIDE >= 512 + int tma_coords_step_per_extra_block = params.stride_extra_kv_block / ExtraKVFormat::TMA_K_STRIDE; uint8_t* k_scales_ptr = - MODEL_TYPE == ModelType::V32 ? + OrigKVFormat::MODEL_TYPE == ModelType::V32 ? (uint8_t*)params.kv + D_NOPE : - (uint8_t*)params.kv + params.page_block_size*(D_NOPE+2*D_ROPE); + (uint8_t*)params.kv + params.page_block_size*TMA_K_STRIDE; uint8_t* extra_k_scales_ptr = - MODEL_TYPE == ModelType::V32 ? + OrigKVFormat::MODEL_TYPE == ModelType::V32 ? (uint8_t*)params.extra_kv + D_NOPE : - (uint8_t*)params.extra_kv + params.extra_page_block_size*(D_NOPE+2*D_ROPE); + (uint8_t*)params.extra_kv + params.extra_page_block_size*ExtraKVFormat::TMA_K_STRIDE; run_main_loop([&](const MainLoopArgs &args) { int* indices = (int*)params.indices + params.stride_indices_b*args.batch_idx + params.stride_indices_s_q*s_q_idx; @@ -672,56 +656,75 @@ KernelTemplate struct IsOrigBlock {}; struct IsExtraBlock {}; - auto process_one_block = [&](int block_idx, auto is_extra_block_t) { + // Prefetch the next block's indices while processing the current one, so that the + // index LDG latency overlaps with the scale LDG and the computation of the current block. + auto load_block_indices = [&](int block_idx, auto is_extra_block_t) -> int2 { static constexpr bool IS_EXTRA_BLOCK = std::is_same_v; + if constexpr (!IS_EXTRA_BLOCK) { + return __ldg((int2*)(indices + block_idx*B_TOPK + lane_idx*2)); + } else { + return __ldg((int2*)(extra_indices + (block_idx-args.num_orig_kv_blocks)*B_TOPK + lane_idx*2)); + } + }; + auto process_one_block = [&](int block_idx, int2 my_indices, auto is_extra_block_t) { + static constexpr bool IS_EXTRA_BLOCK = std::is_same_v; + using F = std::conditional_t; int cur_block_size = IS_EXTRA_BLOCK ? params.extra_page_block_size : params.page_block_size; int64_t cur_k_block_stride = IS_EXTRA_BLOCK ? params.stride_extra_kv_block : params.stride_kv_block; [[maybe_unused]] int cur_k_row_stride = IS_EXTRA_BLOCK ? params.stride_extra_kv_row : params.stride_kv_row; uint8_t* cur_k_scales_ptr = IS_EXTRA_BLOCK ? extra_k_scales_ptr : k_scales_ptr; int cur_tma_coords_step_per_block = IS_EXTRA_BLOCK ? tma_coords_step_per_extra_block : tma_coords_step_per_block; - int abs_pos, my_indices[2]; - if (!IS_EXTRA_BLOCK) { - abs_pos = block_idx*B_TOPK + lane_idx*2; - *(int2*)my_indices = __ldg((int2*)(indices + abs_pos)); - } else { - abs_pos = (block_idx-args.num_orig_kv_blocks)*B_TOPK + lane_idx*2; - *(int2*)my_indices = __ldg((int2*)(extra_indices + abs_pos)); - } - plan.bar_valid_coord_scale_free[rs.index_buf_idx].wait(rs.index_bar_phase^1); + int abs_pos = IS_EXTRA_BLOCK ? (block_idx-args.num_orig_kv_blocks)*B_TOPK + lane_idx*2 : block_idx*B_TOPK + lane_idx*2; - int tma_coords[2]; - e8m0 scales[2*NUM_SCALES_EACH_TOKEN]; - char valid_mask = 0; + // Issue the div/mod and the scale LDGs before waiting for the index buffer, so their + // latency overlaps with the (usually long) empty-barrier wait. + int block_idx_arr[2], idx_in_block_arr[2]; + alignas(16) uint8_t scales[2][SCALE_SMEM_STRIDE]; CUTE_UNROLL for (int i = 0; i < 2; ++i) { - int block_idx, idx_in_block; - block_idx = (unsigned int)my_indices[i] / cur_block_size; - idx_in_block = (unsigned int)my_indices[i] % cur_block_size; - bool is_token_valid = my_indices[i] != -1 && (abs_pos+i < (IS_EXTRA_BLOCK?args.extra_topk_length:args.topk_length)); - valid_mask |= is_token_valid << i; - tma_coords[i] = is_token_valid ? block_idx*cur_tma_coords_step_per_block + idx_in_block*tma_coords_step_per_token : -1; // If the token is invalid because it topk position exceeds topk_length, we must manually fill tma_coords with -1 to avoid copying-in NaN. - if constexpr (MODEL_TYPE == ModelType::V32) { - int64_t offset = is_token_valid ? block_idx*cur_k_block_stride + idx_in_block*cur_k_row_stride : 0; + int cur_idx = i == 0 ? my_indices.x : my_indices.y; + int kv_block_idx = (unsigned int)cur_idx / cur_block_size; + int idx_in_block = (unsigned int)cur_idx % cur_block_size; + block_idx_arr[i] = kv_block_idx; + idx_in_block_arr[i] = idx_in_block; + bool is_token_valid = cur_idx != -1 && (abs_pos+i < (IS_EXTRA_BLOCK?args.extra_topk_length:args.topk_length)); + if constexpr (F::IS_V32) { + int64_t offset = is_token_valid ? kv_block_idx*cur_k_block_stride + idx_in_block*cur_k_row_stride : 0; float4 cur_scale_fp32 = __ldg((float4*)(cur_k_scales_ptr + offset)); e8m0 res[4]; *(__nv_fp8x2_storage_t*)(res+0) = __nv_cvt_float2_to_e8m0x2(float2{cur_scale_fp32.x, cur_scale_fp32.y}, __NV_NOSAT, cudaRoundZero); *(__nv_fp8x2_storage_t*)(res+2) = __nv_cvt_float2_to_e8m0x2(float2{cur_scale_fp32.z, cur_scale_fp32.w}, __NV_NOSAT, cudaRoundZero); if (!is_token_valid) *(uint32_t*)res = (uint32_t)0; - *(uint32_t*)(scales+i*NUM_SCALES_EACH_TOKEN) = *(uint32_t*)(res); + *(uint32_t*)scales[i] = *(uint32_t*)(res); } else { - int64_t offset = block_idx*cur_k_block_stride + idx_in_block*8; // Each token has 7 scale factors with an extra 1B padding - uint64_t scalesx8 = is_token_valid ? __ldg((uint64_t*)(cur_k_scales_ptr + offset)) : 0; - *(uint64_t*)(scales+i*NUM_SCALES_EACH_TOKEN) = scalesx8; + // Each token's row of 1 B scales (padded to whole words, e.g. 7 + 1 B for V4) + int64_t offset = kv_block_idx*cur_k_block_stride + idx_in_block*F::NUM_SCALES_EACH_TOKEN; + ldg_or_zero(scales[i], cur_k_scales_ptr + offset, is_token_valid); } } + plan.bar_valid_coord_scale_free[rs.index_buf_idx].wait(rs.index_bar_phase^1); + + int tma_coords[2]; + char valid_mask = 0; + CUTE_UNROLL + for (int i = 0; i < 2; ++i) { + int cur_idx = i == 0 ? my_indices.x : my_indices.y; + bool is_token_valid = cur_idx != -1 && (abs_pos+i < (IS_EXTRA_BLOCK?args.extra_topk_length:args.topk_length)); + valid_mask |= is_token_valid << i; + tma_coords[i] = is_token_valid ? block_idx_arr[i]*cur_tma_coords_step_per_block + idx_in_block_arr[i]*tma_coords_step_per_token : -1; // If the token is invalid because it topk position exceeds topk_length, we must manually fill tma_coords with -1 to avoid copying-in NaN. + } valid_mask <<= lane_idx%4*2; valid_mask |= __shfl_xor_sync(0xFFFFFFFF, valid_mask, 0x1); valid_mask |= __shfl_xor_sync(0xFFFFFFFF, valid_mask, 0x2); - if constexpr (MODEL_TYPE == ModelType::V32) { - *(uint64_t*)(plan.scales[rs.index_buf_idx] + lane_idx*2) = *(uint64_t*)scales; + if constexpr (SCALE_SMEM_STRIDE == F::NUM_SCALES_EACH_TOKEN) { + // The lane's 2 tokens are contiguous in smem: one store + copy_bytes<2*SCALE_SMEM_STRIDE>(plan.scales[rs.index_buf_idx][lane_idx*2], scales[0]); } else { - *(__int128_t*)(plan.scales[rs.index_buf_idx] + lane_idx*2) = *(__int128_t*)scales; + // Mixed-format kernel, fp8 tokens: each token uses the first bytes of its 32 B row + CUTE_UNROLL + for (int i = 0; i < 2; ++i) + copy_bytes(plan.scales[rs.index_buf_idx][lane_idx*2 + i], scales[i]); } *(int2*)(plan.tma_coord[rs.index_buf_idx] + lane_idx*2) = *(int2*)tma_coords; if (lane_idx%4 == 0) @@ -731,14 +734,24 @@ KernelTemplate rs.update(); }; + const int orig_end = min(args.num_orig_kv_blocks, args.end_block_idx); + int2 my_indices = args.start_block_idx < orig_end ? load_block_indices(args.start_block_idx, IsOrigBlock{}) : int2{}; CUTE_NO_UNROLL - for (int block_idx = args.start_block_idx; block_idx < min(args.num_orig_kv_blocks, args.end_block_idx); ++block_idx) { - process_one_block(block_idx, IsOrigBlock{}); + for (int block_idx = args.start_block_idx; block_idx < orig_end; ++block_idx) { + bool has_next = block_idx+1 < orig_end; + int2 next_indices = has_next ? load_block_indices(block_idx+1, IsOrigBlock{}) : int2{}; + process_one_block(block_idx, my_indices, IsOrigBlock{}); + if (has_next) my_indices = next_indices; } + const int extra_start = max(args.start_block_idx, args.num_orig_kv_blocks); + my_indices = extra_start < args.end_block_idx ? load_block_indices(extra_start, IsExtraBlock{}) : int2{}; CUTE_NO_UNROLL - for (int block_idx = max(args.start_block_idx, args.num_orig_kv_blocks); block_idx < args.end_block_idx; ++block_idx) { - process_one_block(block_idx, IsExtraBlock{}); + for (int block_idx = extra_start; block_idx < args.end_block_idx; ++block_idx) { + bool has_next = block_idx+1 < args.end_block_idx; + int2 next_indices = has_next ? load_block_indices(block_idx+1, IsExtraBlock{}) : int2{}; + process_one_block(block_idx, my_indices, IsExtraBlock{}); + if (has_next) my_indices = next_indices; } }); } else { @@ -748,97 +761,38 @@ KernelTemplate // Dequant warpgroup cutlass::arch::warpgroup_reg_alloc<208>(); - // 8 threads per token - constexpr int GROUP_SIZE = 8, NUM_GROUPS = 128/8, ROWS_PER_GROUP = B_TOPK / NUM_GROUPS, COLS_PER_GROUP = D_NOPE/(GROUP_SIZE*8); - int group_idx = idx_in_warpgroup/GROUP_SIZE, idx_in_group = idx_in_warpgroup%GROUP_SIZE; - Tensor nope0 = make_tensor(make_smem_ptr(plan.u.kv.dequant[0].nope.data()), SmemLayoutKTiles_SW128{}); - bf16* nope0_base = &nope0(group_idx, idx_in_group*8); - bf16* nope1_base = nope0_base + (plan.u.kv.dequant[1].nope.data() - plan.u.kv.dequant[0].nope.data()); - e4m3* raw_nope0_base = plan.u.kv.raw_nope[rs.buf_idx].data() + group_idx*D_NOPE + idx_in_group*8; - e4m3* raw_nope1_base = raw_nope0_base + B_H*D_NOPE; + // Turns the raw (fp8 / fp4) rows of each KV block into the bf16 K tile, see KVBlockDequantizer. The format of a block is + // that of the cache it comes from, resolved at compile time by for_each_kv_block + static_assert(NUM_BUFS == 2); + const Dequantizer dequant_orig(idx_in_warpgroup); + const Dequantizer dequant_extra(idx_in_warpgroup); // Identical to dequant_orig unless the extra cache is fp4 + const uint8_t *raw_0 = plan.u.kv.raw_quant[0].data(), *raw_1 = plan.u.kv.raw_quant[1].data(); + const uint32_t dst_0 = cute::cast_smem_ptr_to_uint(plan.u.kv.dequant[0].quant_part); + const uint32_t dst_1 = cute::cast_smem_ptr_to_uint(plan.u.kv.dequant[1].quant_part); run_main_loop([&](const MainLoopArgs &args) { - // plan.bar_last_store_done.wait(args.bar_phase_batch_rel); // No need to wait since the raw nope producer must wait + // plan.bar_last_store_done.wait(args.bar_phase_batch_rel); // No need to wait since the raw KV producer must wait plan.bar_q_utccp.wait(args.bar_phase_batch_rel); - CUTE_NO_UNROLL - for (int block_idx = args.start_block_idx; block_idx < args.end_block_idx; ++block_idx) { + for_each_kv_block(args.start_block_idx, args.end_block_idx, args.num_orig_kv_blocks, [&](int, bool) { plan.bar_valid_coord_scale_ready[rs.index_buf_idx].wait(rs.index_bar_phase); plan.bar_raw_ready[rs.buf_idx].wait(rs.bar_phase); plan.bar_sv_done[rs.buf_idx].wait(rs.bar_phase^1); - uint32_t cur_nope_base_uint_addr = cute::cast_smem_ptr_to_uint(rs.buf_idx == 0 ? nope0_base : nope1_base); - e4m3* raw_nope_base = rs.buf_idx == 0 ? raw_nope0_base : raw_nope1_base; - auto st_128b = [&](int local_row_idx, int local_col_idx, __int128_t &data) { - asm volatile ("st.weak.shared::cta.b128 [%0], %1;\n" - : - : "r"(cur_nope_base_uint_addr + 2*(local_row_idx*NUM_GROUPS*64 + local_col_idx*B_TOPK*64)), "q"(data) // 2 for sizeof(bf16) - ); // We have this `asm volatile` here, otherwise the compiler generates ST.E instead of STS - }; - auto get_raw_fp8 = [&](int local_row_idx, int local_col_idx) -> uint64_t { - return *(uint64_t*)(raw_nope_base + local_row_idx*NUM_GROUPS*D_NOPE + local_col_idx*(GROUP_SIZE*8)); - }; - // The following code suffers from a 2-way bank conflict when reading from SMEM. - if constexpr (MODEL_TYPE == ModelType::V32) { - CUTE_UNROLL - for (int local_row_idx = 0; local_row_idx < ROWS_PER_GROUP; ++local_row_idx) { - int row_idx = local_row_idx*NUM_GROUPS + group_idx; - bf16 scales[4]; - e8m0 scales_e8m0[4]; - *(uint32_t*)scales_e8m0 = *(uint32_t*)plan.scales[rs.index_buf_idx][row_idx]; - *(__nv_bfloat162_raw*)(scales+0) = __nv_cvt_e8m0x2_to_bf162raw(*(unsigned short*)(scales_e8m0+0)); - *(__nv_bfloat162_raw*)(scales+2) = __nv_cvt_e8m0x2_to_bf162raw(*(unsigned short*)(scales_e8m0+2)); - - uint64_t cur_data_fp8x8 = get_raw_fp8(local_row_idx, 0); - CUTE_UNROLL - for (int local_col_idx = 0; local_col_idx < COLS_PER_GROUP; ++local_col_idx) { - ku::nve4m3x2 data_fp8[4]; - ku::nvbf16x2 data_bf16[4]; - *(uint64_t*)data_fp8 = cur_data_fp8x8; - if (local_col_idx+1 < COLS_PER_GROUP) - cur_data_fp8x8 = get_raw_fp8(local_row_idx, local_col_idx+1); - bf16 scale = scales[local_col_idx / (D_NOPE/(GROUP_SIZE*8)/4)]; - CUTE_UNROLL - for (int i = 0; i < 4; ++i) { - data_bf16[i] = fp8x2_to_bf16x2_with_scale(data_fp8[i], *(ku::nvbf16*)(&scale)); - } - st_128b(local_row_idx, local_col_idx, *(__int128_t*)data_bf16); - } - } - } else { - CUTE_UNROLL - for (int local_row_idx = 0; local_row_idx < ROWS_PER_GROUP; ++local_row_idx) { - int row_idx = local_row_idx*NUM_GROUPS + group_idx; - bf16 scales[8]; - e8m0 scales_e8m0[8]; - *(uint64_t*)scales_e8m0 = *(uint64_t*)plan.scales[rs.index_buf_idx][row_idx]; - *(__nv_bfloat162_raw*)(scales+0) = __nv_cvt_e8m0x2_to_bf162raw(*(unsigned short*)(scales_e8m0+0)); - *(__nv_bfloat162_raw*)(scales+2) = __nv_cvt_e8m0x2_to_bf162raw(*(unsigned short*)(scales_e8m0+2)); - *(__nv_bfloat162_raw*)(scales+4) = __nv_cvt_e8m0x2_to_bf162raw(*(unsigned short*)(scales_e8m0+4)); - *(__nv_bfloat162_raw*)(scales+6) = __nv_cvt_e8m0x2_to_bf162raw(*(unsigned short*)(scales_e8m0+6)); - - uint64_t cur_data_fp8x8 = get_raw_fp8(local_row_idx, 0); - CUTE_UNROLL - for (int local_col_idx = 0; local_col_idx < COLS_PER_GROUP; ++local_col_idx) { - ku::nve4m3x2 data_fp8[4]; - ku::nvbf16x2 data_bf16[4]; - *(uint64_t*)data_fp8 = cur_data_fp8x8; - if (local_col_idx+1 < COLS_PER_GROUP) - cur_data_fp8x8 = get_raw_fp8(local_row_idx, local_col_idx+1); - bf16 scale = scales[local_col_idx]; - CUTE_UNROLL - for (int i = 0; i < 4; ++i) { - data_bf16[i] = fp8x2_to_bf16x2_with_scale(data_fp8[i], *(ku::nvbf16*)(&scale)); - } - st_128b(local_row_idx, local_col_idx, *(__int128_t*)data_bf16); - } - } - } + const auto *dequant = [&] { + if constexpr (std::is_same_v) return &dequant_orig; else return &dequant_extra; + }(); + dequant->run( + rs.buf_idx == 0 ? raw_0 : raw_1, + plan.scales[rs.index_buf_idx][0], + rs.buf_idx == 0 ? dst_0 : dst_1, + []{} + ); cutlass::arch::fence_view_async_shared(); - plan.bar_nope_ready[rs.buf_idx].arrive(); + plan.bar_quant_part_dequant_ready[rs.buf_idx].arrive(); plan.bar_raw_free[rs.buf_idx].arrive(); plan.bar_valid_coord_scale_free[rs.index_buf_idx].arrive(); rs.update(); - } + }); }); } #else @@ -850,24 +804,24 @@ KernelTemplate template __global__ void __launch_bounds__(Kernel::NUM_THREADS, 1, 1) -flash_fwd_splitkv_mla_fp8_sparse_kernel(__grid_constant__ const SparseAttnDecodeParams params, __grid_constant__ const TmaParams tma_params) { - Kernel::flash_fwd_splitkv_mla_fp8_sparse_kernel_devfunc(params, tma_params); +flash_fwd_splitkv_mla_fp8_sparse_kernel(__grid_constant__ const SparseAttnDecodeParams params, __grid_constant__ const TmaParams tma_params ) { + Kernel::flash_fwd_splitkv_mla_fp8_sparse_kernel_devfunc(params, tma_params ); } -template -void KernelTemplate::run(const SparseAttnDecodeParams ¶ms) { +template +void KernelTemplate::run(const SparseAttnDecodeParams ¶ms) { KU_ASSERT(params.topk % B_TOPK == 0, "topk (%d) mod B_TOPK (%d) must be 0", params.topk, B_TOPK); KU_ASSERT(params.extra_topk % B_TOPK == 0, "extra_topk (%d) mod B_TOPK (%d) must be 0", params.extra_topk, B_TOPK); - KU_ASSERT(params.h_q == B_H); + KU_ASSERT(params.h_q == 64); KU_ASSERT(params.h_kv == 1); KU_ASSERT(params.d_qk == D_Q); KU_ASSERT(params.d_v == D_V); - if constexpr (MODEL_TYPE == ModelType::MODEL1) { - constexpr int BYTES_PER_TOKEN = D_NOPE + 2*D_ROPE + 8; - KU_ASSERT(params.stride_kv_row == BYTES_PER_TOKEN, "Each page block in KV cache must be contiguous for head64 sparse fp8 decoding attention in MODEL1"); // Each block must be contiguous - } + KU_ASSERT(params.model_type == OrigKVFormat::MODEL_TYPE && params.extra_model_type == ExtraKVFormat::MODEL_TYPE); - auto shape_Q_SW128 = make_shape(B_H, D_Q, params.s_q, params.b); + // NOTE The head dim of the gmem shape is h_q while the smem box covers + // B_H rows, so TMA zero-fills rows [h_q, B_H) on load (Q), and drops them on + // store (O). Same for tensor_map_q_sw64 below. + auto shape_Q_SW128 = make_shape(params.h_q, D_Q, params.s_q, params.b); auto tma_Q_SW128 = cute::make_tma_copy( SM90_TMA_LOAD{}, make_tensor( @@ -880,7 +834,7 @@ void KernelTemplate::run(const SparseAttnDecodeParams ¶ms) { SmemLayoutQ_SW128{} ); - auto shape_O = make_shape(B_H, D_V, params.s_q, params.b); + auto shape_O = make_shape(params.h_q, D_V, params.s_q, params.b); auto tma_O = cute::make_tma_copy( SM90_TMA_STORE{}, make_tensor( @@ -906,35 +860,19 @@ void KernelTemplate::run(const SparseAttnDecodeParams ¶ms) { ); } - auto get_nope_rope_tensormap = [&](bool is_extra, void* k_ptr, int num_blocks, int64_t k_batch_stride) -> std::pair { - static_assert(D_NOPE%8 == 0); - KU_ASSERT((int64_t)k_ptr % 16 == 0, "The base address of %sk_ptr (%p) must be 16B aligned for sparse fp8 attention on sm100f", is_extra?"extra_":"", k_ptr); - KU_ASSERT(k_batch_stride % TMA_K_STRIDE == 0, "%sk_cache.stride(0) (%ld) must be a multiple of %d. Padding might be necessary", is_extra?"extra_":"", k_batch_stride, TMA_K_STRIDE); - CUtensorMap tensor_map_kv_nope = ku::make_tensor_map( - {D_NOPE/8, (uint64_t)num_blocks * (k_batch_stride/TMA_K_STRIDE)}, - {TMA_K_STRIDE}, - {D_NOPE/8, 1}, - k_ptr, - CUtensorMapDataType::CU_TENSOR_MAP_DATA_TYPE_INT64, - CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_NONE, - CUtensorMapL2promotion::CU_TENSOR_MAP_L2_PROMOTION_L2_128B - ); // NOTE We combine 8 float8 into 1 int64 since boxdim cannot > 256 - CUtensorMap tensor_map_kv_rope = ku::make_tensor_map( - {D_ROPE, (uint64_t)num_blocks * (k_batch_stride/TMA_K_STRIDE)}, - {TMA_K_STRIDE}, - {K_ROPE_SW/2, 1}, - (uint8_t*)k_ptr + (MODEL_TYPE == ModelType::V32 ? (D_NOPE+16) : D_NOPE), - CUtensorMapDataType::CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, - K_ROPE_SW == 64 ? CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_64B : CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_128B, - CUtensorMapL2promotion::CU_TENSOR_MAP_L2_PROMOTION_L2_128B - ); - return {tensor_map_kv_nope, tensor_map_kv_rope}; - }; - - auto [tensor_map_kv_nope, tensor_map_kv_rope] = get_nope_rope_tensormap(false, params.kv, params.num_blocks, params.stride_kv_block); - CUtensorMap tensor_map_extra_kv_nope{}, tensor_map_extra_kv_rope{}; + CUtensorMap tensor_map_kv_quant_part = make_kv_quant_part_tensor_map>( + "k_cache", params.kv, params.num_blocks, params.stride_kv_block, params.stride_kv_row, 0); + CUtensorMap tensor_map_kv_bf16_part{}; + if constexpr (D_BF16 > 0) { + tensor_map_kv_bf16_part = make_kv_bf16_part_tensor_map(params.kv, params.num_blocks, params.stride_kv_block); + } + CUtensorMap tensor_map_extra_kv_quant_part{}, tensor_map_extra_kv_bf16_part{}; if (params.extra_topk > 0) { - std::tie(tensor_map_extra_kv_nope, tensor_map_extra_kv_rope) = get_nope_rope_tensormap(true, params.extra_kv, params.extra_num_blocks, params.stride_extra_kv_block); + tensor_map_extra_kv_quant_part = make_kv_quant_part_tensor_map>( + "extra_k_cache", params.extra_kv, params.extra_num_blocks, params.stride_extra_kv_block, params.stride_extra_kv_row, 0); + if constexpr (ExtraKVFormat::D_BF16 > 0) { + tensor_map_extra_kv_bf16_part = make_kv_bf16_part_tensor_map(params.extra_kv, params.extra_num_blocks, params.stride_extra_kv_block); + } } TmaParams< @@ -944,25 +882,27 @@ void KernelTemplate::run(const SparseAttnDecodeParams ¶ms) { shape_Q_SW128, tma_Q_SW128, shape_O, tma_O, tensor_map_q_sw64, - tensor_map_kv_nope, - tensor_map_kv_rope, - tensor_map_extra_kv_nope, - tensor_map_extra_kv_rope + tensor_map_kv_quant_part, + tensor_map_kv_bf16_part, + tensor_map_extra_kv_quant_part, + tensor_map_extra_kv_bf16_part }; - auto mla_kernel = &flash_fwd_splitkv_mla_fp8_sparse_kernel, decltype(tma_params)>; + auto mla_kernel = &flash_fwd_splitkv_mla_fp8_sparse_kernel, decltype(tma_params)>; constexpr size_t smem_size = sizeof(SharedMemoryPlan); - static_assert(smem_size < 227*1024); + static_assert(smem_size <= 227*1024); KU_CUDA_CHECK(cudaFuncSetAttribute(mla_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); + // NOTE Don't use PDL because of potential compiler bugs! - mla_kernel<<>>(params, tma_params); + mla_kernel<<>>(params, tma_params ); KU_CHECK_KERNEL_LAUNCH(); + } -template +template void run_flash_splitkv_mla_fp8_sparse_kernel(const SparseAttnDecodeParams ¶ms) { - KernelTemplate::run(params); + KernelTemplate::run(params); } } diff --git a/csrc/kernels/sm100/decode/sparse/head64/kernel.h b/csrc/kernels/sm100/decode/sparse/head64/kernel.h new file mode 100644 index 00000000..a72c6773 --- /dev/null +++ b/csrc/kernels/sm100/decode/sparse/head64/kernel.h @@ -0,0 +1,16 @@ +#pragma once + +#include "kernels/params.h" + +namespace sm100::decode::sparse::head64 { + +struct Config { + ModelType MODEL_TYPE; // Format of `kv`: V32, V4 or V41 + ModelType EXTRA_MODEL_TYPE; // Format of `extra_kv`: MODEL_TYPE, or V41_FP4 with MODEL_TYPE == V41 (see is_valid_kv_format_pair) + bool ENABLE_SPLITKV; +}; + +template +void run_flash_splitkv_mla_fp8_sparse_kernel(const SparseAttnDecodeParams ¶ms); + +} diff --git a/csrc/kernels/sm100/dequant_utils.cuh b/csrc/kernels/sm100/dequant_utils.cuh new file mode 100644 index 00000000..709989ac --- /dev/null +++ b/csrc/kernels/sm100/dequant_utils.cuh @@ -0,0 +1,201 @@ +#pragma once + +#include +#include + +#include "kernels/kv_cache_format.h" +#include "kernels/sm100/helpers.h" + +// Shared pieces of the sm100 decoding kernels for reading paged quantized KV caches (`kv` of format Orig plus an optional +// `extra_kv` of format Extra, see KVCacheFormat): iterating the KV blocks of a request, dequantizing a block, and the TMA tensor +// maps of the caches +namespace sm100 { + +// Runs f.template operator()(block_idx, is_extra_block) for the KV blocks [begin, end) of a request, F being KVCacheFormat of the +// cache the block comes from: blocks [0, num_orig_blocks) are from `kv`, the others from `extra_kv`. A block never mixes the two -- +// the caller keeps the orig/extra boundary on a block boundary (head64 asserts topk % B_TOPK == 0; phase1 rounds each cache's +// block count up and masks the tail slots). When the formats differ, the blocks of each cache get their own loop, so that the code +// of the two formats is never interleaved (ptxas would otherwise merge the live ranges of both) +template +CUTE_DEVICE void for_each_kv_block(int begin, int end, int num_orig_blocks, Fn &&f) { + if constexpr (std::is_same_v) { + CUTE_NO_UNROLL + for (int block_idx = begin; block_idx < end; ++block_idx) { + f.template operator()(block_idx, block_idx >= num_orig_blocks); + } + } else { + CUTE_NO_UNROLL + for (int block_idx = begin; block_idx < min(num_orig_blocks, end); ++block_idx) { + f.template operator()(block_idx, false); + } + CUTE_NO_UNROLL + for (int block_idx = max(begin, num_orig_blocks); block_idx < end; ++block_idx) { + f.template operator()(block_idx, true); + } + } +} + +// NUM_BYTES (4 / 8 / 16 / 32) from global memory with one load (two for 32), or zeros if !valid +template +CUTE_DEVICE void ldg_or_zero(uint8_t *dst, const uint8_t *src, bool valid) { + static_assert(NUM_BYTES == 4 || NUM_BYTES == 8 || NUM_BYTES == 16 || NUM_BYTES == 32); + if constexpr (NUM_BYTES == 4) { + *(uint32_t*)dst = valid ? __ldg((const uint32_t*)src) : 0u; + } else if constexpr (NUM_BYTES == 8) { + *(uint64_t*)dst = valid ? __ldg((const uint64_t*)src) : 0ull; + } else { + CUTE_UNROLL + for (int i = 0; i < NUM_BYTES / 16; ++i) { + *((int4*)dst + i) = valid ? __ldg((const int4*)src + i) : int4{0, 0, 0, 0}; + } + } +} + +// NUM_BYTES (4 / 8 / 16 / 32 / 64) between registers and shared memory with the widest accesses +template +CUTE_DEVICE void copy_bytes(uint8_t *dst, const uint8_t *src) { + static_assert(NUM_BYTES == 4 || NUM_BYTES == 8 || NUM_BYTES % 16 == 0); + if constexpr (NUM_BYTES == 4) { + *(uint32_t*)dst = *(const uint32_t*)src; + } else if constexpr (NUM_BYTES == 8) { + *(uint64_t*)dst = *(const uint64_t*)src; + } else { + CUTE_UNROLL + for (int i = 0; i < NUM_BYTES / 16; ++i) { + *((__int128_t*)dst + i) = *((const __int128_t*)src + i); + } + } +} + +// Dequantizes one KV block of B_TOPK tokens (this CTA's D dims, stored in format F) with one warpgroup: reads the raw rows +// gathered by TMA and the 1 B scales from shared memory, converts in registers, and stores the bf16 result into a tile in the +// canonical SW128 K-major layout. fp8 and fp4 share the code below and differ only in the bytes of raw data per step +// (RawWord), the loading of scales, and the conversion instructions. A block is entirely fp8 or entirely fp4, so the format +// is a compile-time parameter (see for_each_kv_block). +// +// The unit of work is a step: a thread turns one RawWord (8 B of fp8 / 4 B of fp4) into one 16 B chunk of 8 bf16. 8 threads +// per token cover one swizzle-atom column (ELEMS_PER_STEP = 64 elements) per step and write the 8 chunks of one 128 B +// swizzle-atom row, so the STS.128 of a wavefront (8 lanes) is conflict-free; 16 tokens per pass, B_TOPK / 16 passes. +// - Raw rows are RAW_TOKEN_SMEM_STRIDE apart. fp4 rows are read with LDS.32 and are padded to 8 mod 32 words (256 + 32, or +// 128 + 32 per CTA), so the 4 rows of a gather4 group start in 4 different quarters of the 32 banks and the LDS is +// conflict-free. fp8 rows are read with LDS.64 and have no padding (a 2-way bank conflict). +// - Scales are 1 B each (ue8m0 for fp8, e4m3 for fp4; V3.2's fp32 scales are converted by the index warp), SCALE_SMEM_STRIDE +// apart, this CTA's NUM_SCALES first. +template +struct KVBlockDequantizer { + static constexpr int GROUP_SIZE = 8, NUM_GROUPS = 128 / GROUP_SIZE, ROWS_PER_GROUP = B_TOPK / NUM_GROUPS; + static constexpr int ELEMS_PER_STEP = GROUP_SIZE * 8; // One swizzle-atom column + static constexpr int NUM_STEPS = D / ELEMS_PER_STEP; + static constexpr int RAW_BYTES_PER_STEP = F::IS_FP4 ? ELEMS_PER_STEP / 2 : ELEMS_PER_STEP; + static constexpr int RAW_BYTES_PER_THREAD_STEP = RAW_BYTES_PER_STEP / GROUP_SIZE; + static constexpr int NUM_SCALES = D / F::QUANT_TILE_SIZE; // 7 for the 448 fp8 dims of V4, 3 for CTA1's 192 of them + static constexpr int SCALE_LOAD_BYTES = (NUM_SCALES + 3) / 4 * 4; // Whole words (the scale rows are padded, see KVCacheFormat) + using RawWord = std::conditional_t; + static_assert(D % ELEMS_PER_STEP == 0 && B_TOPK % NUM_GROUPS == 0 && SCALE_LOAD_BYTES <= SCALE_SMEM_STRIDE); + static_assert(RAW_TOKEN_SMEM_STRIDE >= NUM_STEPS * RAW_BYTES_PER_STEP); // RAW_TOKEN_SMEM_STRIDE must hold a row's data + static_assert(!F::IS_FP4 || RAW_TOKEN_SMEM_STRIDE / 4 % 16 == 8); // fp4 rows start in 4 different bank quarters, see above + // The scale index of a step is a compile-time base plus this thread's part. fp8: the part is 0 or 1 (QUANT_TILE_SIZE >= 32, + // written as a comparison so that the byte array stays in registers as a select); fp4: the part is idx_in_group / 2, and the + // byte is extracted from a (compile-time) word of the row with one PRMT, replicated into bytes 0 and 1 + static_assert(F::IS_FP4 ? ELEMS_PER_STEP == 4 * F::QUANT_TILE_SIZE : F::QUANT_TILE_SIZE >= 32); + + int group_idx, idx_in_group; + uint32_t raw_offset; // Of this thread's first raw word within a block + uint32_t dst_offset; // Of this thread's first bf16 chunk within a tile (swizzle included: row % 8 is fixed for a thread since NUM_GROUPS % 8 == 0) + uint32_t scale_prmt_sel11; // fp4 only: the PRMT selector extracting this thread's scale of a step, see above + + CUTE_DEVICE explicit KVBlockDequantizer(int idx_in_warpgroup): + group_idx(idx_in_warpgroup / GROUP_SIZE), idx_in_group(idx_in_warpgroup % GROUP_SIZE) { + raw_offset = group_idx * RAW_TOKEN_SMEM_STRIDE + idx_in_group * RAW_BYTES_PER_THREAD_STEP; + // The thread's first chunk in the SW128 K-major tile: 8-row swizzle atoms are stacked along the rows first, and within an + // atom row the 16 B lane index is XORed with the row's position in the atom (the swizzle acts on byte addresses) + const int row_in_atom = group_idx % 8; + dst_offset = group_idx / 8 * (8 * 128) + row_in_atom * 128 + (idx_in_group ^ row_in_atom) * 16; + scale_prmt_sel11 = (uint32_t)(idx_in_group / 2) * 0x11; + } + + // raw / scales: the block's raw rows and scale rows in shared memory; dst: the shared memory address (cvta) of the bf16 tile. + // before_first_store() runs once, right before the first STS, so that waiting for the tile to be free overlaps with the first + // loads and conversions + template + CUTE_DEVICE void run(const uint8_t *raw, const uint8_t *scales, uint32_t dst, Fn &&before_first_store) const { + CUTE_UNROLL + for (int local_row_idx = 0; local_row_idx < ROWS_PER_GROUP; ++local_row_idx) { + const int row_idx = local_row_idx * NUM_GROUPS + group_idx; + alignas(16) uint8_t scales_row[SCALE_LOAD_BYTES]; + copy_bytes(scales_row, scales + row_idx * SCALE_SMEM_STRIDE); + const uint8_t *raw_row = raw + raw_offset + local_row_idx * NUM_GROUPS * RAW_TOKEN_SMEM_STRIDE; + RawWord cur_data = *(const RawWord*)raw_row; + CUTE_UNROLL + for (int local_col_idx = 0; local_col_idx < NUM_STEPS; ++local_col_idx) { + RawWord data = cur_data; + if (local_col_idx + 1 < NUM_STEPS) + cur_data = *(const RawWord*)(raw_row + (local_col_idx + 1) * RAW_BYTES_PER_STEP); + // Elements [local_col_idx*64 + idx_in_group*8, +8) of the row lie in this quant tile / word (see the ctor) + const int scale_idx_base = local_col_idx * ELEMS_PER_STEP / F::QUANT_TILE_SIZE; + ku::nvbf16x2 data_bf16[4]; + if constexpr (F::IS_FP4) { + fp4x8_to_bf16x2x4(data, data_bf16); + const uint32_t scale_word = *(const uint32_t*)(scales_row + scale_idx_base); // Constant offset: the array stays in registers + ku::nvbf16x2 scale = e4m3x2_to_bf16x2(__byte_perm(scale_word, 0, scale_prmt_sel11)); // (s, s) + CUTE_UNROLL + for (int i = 0; i < 4; ++i) + data_bf16[i] = __hmul2(data_bf16[i], scale); // Exact: e2m1 x e4m3 has at most 2 + 4 significant bits + } else { + const int scale_idx = scale_idx_base + (F::QUANT_TILE_SIZE == 32 ? idx_in_group >= GROUP_SIZE / 2 : 0); + CUTE_UNROLL + for (int i = 0; i < 4; ++i) { + data_bf16[i] = fp8x2_to_bf16x2_with_scale(((ku::nve4m3x2*)&data)[i], ((__nv_fp8_e8m0*)scales_row)[scale_idx]); + } + } + if (local_row_idx == 0 && local_col_idx == 0) { + before_first_store(); + } + asm volatile ("st.weak.shared::cta.b128 [%0], %1;\n" + : + : "r"(dst + dst_offset + local_row_idx * NUM_GROUPS * 128 + local_col_idx * B_TOPK * 128), "q"(*(__int128_t*)data_bf16) + ); + } + } + } +}; + +// Host: TMA tensor map of one CTA's share of the quantized part of a paged KV cache, for gather4. dim0 = DIM0_BYTES of a token +// starting at byte `cta_byte_offset` (as uint32), dim1 = tokens at F::TMA_K_STRIDE; box = {BOX_BYTES, 1 row}. BOX_BYTES may exceed +// DIM0_BYTES, the rest of the box is out of bounds and zero-filled by TMA (the padding of the fp4 raw rows, see +// KVBlockDequantizer) -- or be smaller, reading a box out of a wider view (head128). Each gather4 writes its 4 rows BOX_BYTES +// apart in shared memory +template +static CUtensorMap make_kv_quant_part_tensor_map(const char *name, void *kv, int num_blocks, int64_t block_stride_bytes, int row_stride_bytes, int cta_byte_offset) { + static_assert(DIM0_BYTES % 4 == 0 && BOX_BYTES % 16 == 0 && BOX_BYTES / 4 <= 256); + KU_ASSERT((int64_t)kv % 16 == 0, "The base address of %s (%p) must be 16B aligned", name, kv); + KU_ASSERT(row_stride_bytes == F::BYTES_PER_TOKEN, "%s.stride(-2) (%d) must be %d, i.e. each page block in the KV cache must be contiguous", name, row_stride_bytes, F::BYTES_PER_TOKEN); + KU_ASSERT(block_stride_bytes % F::TMA_K_STRIDE == 0, "%s.stride(0) (%ld) must be a multiple of %d. Padding might be necessary", name, block_stride_bytes, F::TMA_K_STRIDE); + KU_ASSERT((uint64_t)num_blocks * (uint64_t)(block_stride_bytes / F::TMA_K_STRIDE) <= INT32_MAX, "%s: too many rows for the int32 TMA coordinates", name); + return ku::make_tensor_map( + {(uint64_t)DIM0_BYTES / 4, (uint64_t)num_blocks * (uint64_t)(block_stride_bytes / F::TMA_K_STRIDE)}, + {(uint64_t)F::TMA_K_STRIDE}, + {BOX_BYTES / 4, 1}, + (uint8_t*)kv + cta_byte_offset, + CUtensorMapDataType::CU_TENSOR_MAP_DATA_TYPE_UINT32, + CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_NONE, + CUtensorMapL2promotion::CU_TENSOR_MAP_L2_PROMOTION_L2_128B + ); +} + +// Host: the TMA tensor map of the bf16 (RoPE) part of a paged KV cache (F::D_BF16 > 0), for gather4 into a SWIZZLE-byte swizzled tile; box = {BOX_ELEMS, 1 row} +template +static CUtensorMap make_kv_bf16_part_tensor_map(void *kv, int num_blocks, int64_t block_stride_bytes) { + static_assert(F::D_BF16 > 0 && (SWIZZLE == 64 || SWIZZLE == 128) && BOX_ELEMS * sizeof(bf16) <= SWIZZLE); + return ku::make_tensor_map( + {(uint64_t)F::D_BF16, (uint64_t)num_blocks * (uint64_t)(block_stride_bytes / F::TMA_K_STRIDE)}, + {(uint64_t)F::TMA_K_STRIDE}, + {BOX_ELEMS, 1}, + (uint8_t*)kv + (F::TMA_K_STRIDE - 2 * F::D_BF16), // The bf16 part is the tail of the token's data + CUtensorMapDataType::CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, + SWIZZLE == 64 ? CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_64B : CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_128B, + CUtensorMapL2promotion::CU_TENSOR_MAP_L2_PROMOTION_L2_128B + ); +} + +} diff --git a/csrc/kernels/sm100/helpers.h b/csrc/kernels/sm100/helpers.h new file mode 100644 index 00000000..0fd5e0ef --- /dev/null +++ b/csrc/kernels/sm100/helpers.h @@ -0,0 +1,167 @@ +#pragma once + +#include + +#include +#include + +namespace sm100 { + +using namespace cute; + +struct bf16x8 { + __nv_bfloat162 a01; + __nv_bfloat162 a23; + __nv_bfloat162 a45; + __nv_bfloat162 a67; +}; + +CUTE_DEVICE +int int4_max(int4 t) { + return max(max(t.x, t.y), max(t.z, t.w)); +} + +CUTE_DEVICE +int int4_min(int4 t) { + return min(min(t.x, t.y), min(t.z, t.w)); +} + +struct int32x8_t { + int a0, a1, a2, a3, a4, a5, a6, a7; +}; + +struct float8 { + float2 a01, a23, a45, a67; +}; + +// CUDA 13.2 (PTX ISA 9.2) adds a direct {e4m3x2, e2m1x2} -> bf16x2 conversion, i.e. +// `cvt.rn.bf16x2.e4m3x2` / `cvt.rn.bf16x2.e2m1x2`. Older toolchains only offer the .f16x2 +// destination, so the value has to be widened through the FP32 domain instead. +#if defined(__CUDACC_VER_MAJOR__) && (__CUDACC_VER_MAJOR__ > 13 || (__CUDACC_VER_MAJOR__ == 13 && __CUDACC_VER_MINOR__ >= 2)) +#define SM100_HAS_NATIVE_BF16X2_CVT 1 +#else +#define SM100_HAS_NATIVE_BF16X2_CVT 0 +#endif + +// 2x f16 -> 2x bf16 (exact: f16 is a subset of f32, and every f16 the KV cache can store is representable in bf16) +CUTE_DEVICE +ku::nvbf16x2 f16x2_to_bf16x2(uint32_t f16x2) { + return __float22bfloat162_rn(__half22float2(*(__half2*)&f16x2)); +} + +// 2x fp8_e4m3 -> 2x bf16, without scaling +CUTE_DEVICE +ku::nvbf16x2 fp8x2_to_bf16x2(ku::nve4m3x2 data) { + uint32_t out; +#if SM100_HAS_NATIVE_BF16X2_CVT + asm("cvt.rn.bf16x2.e4m3x2 %0, %1;" + : "=r"(out) + : "h"(*(uint16_t*)&data)); +#else + asm("cvt.rn.f16x2.e4m3x2 %0, %1;" + : "=r"(out) + : "h"(*(uint16_t*)&data)); + const ku::nvbf16x2 bf16 = f16x2_to_bf16x2(out); + out = *(uint32_t*)&bf16; +#endif + return *(ku::nvbf16x2*)&out; +} + +// 1x ue8m0 scale -> 2x bf16 (the same scale in both halves) +CUTE_DEVICE +ku::nvbf16x2 ue8m0_to_bf16x2(__nv_fp8_e8m0 scale_e8m0) { + uint16_t packed = (uint16_t)(*(uint8_t*)&scale_e8m0) * 0x0101; + uint32_t out; + asm("cvt.rn.bf16x2.ue8m0x2 %0, %1;" + : "=r"(out) + : "h"(packed)); + return *(ku::nvbf16x2*)&out; +} + +// 2x fp8_e4m3 -> 2x bf16, multiplied by an ue8m0 scale +CUTE_DEVICE +ku::nvbf16x2 fp8x2_to_bf16x2_with_scale(ku::nve4m3x2 data, __nv_fp8_e8m0 scale_e8m0) { + return __hmul2(fp8x2_to_bf16x2(data), ue8m0_to_bf16x2(scale_e8m0)); +} + +// 8x fp4_e2m1 (packed in 32 bits, 2 per byte) -> 4x bf16x2 +// Written as one asm block so that ptxas selects the source byte with the .B0-.B3 operand selector of F2FP instead of emitting PRMTs +CUTE_DEVICE +void fp4x8_to_bf16x2x4(uint32_t packed, ku::nvbf16x2 *out) { + uint32_t o0, o1, o2, o3; +#if SM100_HAS_NATIVE_BF16X2_CVT + asm( + "{\n" + ".reg .b8 b0, b1, b2, b3;\n" + "mov.b32 {b0, b1, b2, b3}, %4;\n" + "cvt.rn.bf16x2.e2m1x2 %0, b0;\n" + "cvt.rn.bf16x2.e2m1x2 %1, b1;\n" + "cvt.rn.bf16x2.e2m1x2 %2, b2;\n" + "cvt.rn.bf16x2.e2m1x2 %3, b3;\n" + "}\n" + : "=r"(o0), "=r"(o1), "=r"(o2), "=r"(o3) + : "r"(packed) + ); + out[0] = *(ku::nvbf16x2*)&o0; + out[1] = *(ku::nvbf16x2*)&o1; + out[2] = *(ku::nvbf16x2*)&o2; + out[3] = *(ku::nvbf16x2*)&o3; +#else + asm( + "{\n" + ".reg .b8 b0, b1, b2, b3;\n" + "mov.b32 {b0, b1, b2, b3}, %4;\n" + "cvt.rn.f16x2.e2m1x2 %0, b0;\n" + "cvt.rn.f16x2.e2m1x2 %1, b1;\n" + "cvt.rn.f16x2.e2m1x2 %2, b2;\n" + "cvt.rn.f16x2.e2m1x2 %3, b3;\n" + "}\n" + : "=r"(o0), "=r"(o1), "=r"(o2), "=r"(o3) + : "r"(packed) + ); + const ku::nvbf16x2 b0 = f16x2_to_bf16x2(o0), b1 = f16x2_to_bf16x2(o1); + const ku::nvbf16x2 b2 = f16x2_to_bf16x2(o2), b3 = f16x2_to_bf16x2(o3); + out[0] = b0; out[1] = b1; out[2] = b2; out[3] = b3; +#endif +} + +// 4x fp8_e4m3 -> 2x bf16x2 without scaling, used for the e4m3 scales of the fp4 KV cache +CUTE_DEVICE +void fp8x4_to_bf16x2x2(uint32_t packed, ku::nvbf16x2 *out) { + uint32_t o0, o1; +#if SM100_HAS_NATIVE_BF16X2_CVT + asm( + "{\n" + ".reg .b16 h0, h1;\n" + "mov.b32 {h0, h1}, %2;\n" + "cvt.rn.bf16x2.e4m3x2 %0, h0;\n" + "cvt.rn.bf16x2.e4m3x2 %1, h1;\n" + "}\n" + : "=r"(o0), "=r"(o1) + : "r"(packed) + ); + out[0] = *(ku::nvbf16x2*)&o0; + out[1] = *(ku::nvbf16x2*)&o1; +#else + asm( + "{\n" + ".reg .b16 h0, h1;\n" + "mov.b32 {h0, h1}, %2;\n" + "cvt.rn.f16x2.e4m3x2 %0, h0;\n" + "cvt.rn.f16x2.e4m3x2 %1, h1;\n" + "}\n" + : "=r"(o0), "=r"(o1) + : "r"(packed) + ); + const ku::nvbf16x2 b0 = f16x2_to_bf16x2(o0), b1 = f16x2_to_bf16x2(o1); + out[0] = b0; out[1] = b1; +#endif +} + +// 2x fp8_e4m3 -> bf16x2 without scaling, used for the e4m3 scale of one quant tile of the fp4 KV cache +CUTE_DEVICE +ku::nvbf16x2 e4m3x2_to_bf16x2(uint16_t packed) { + return fp8x2_to_bf16x2(*(ku::nve4m3x2*)&packed); +} + +} diff --git a/csrc/sm100/prefill/dense/collective/fmha_common.hpp b/csrc/kernels/sm100/prefill/dense/collective/fmha_common.hpp similarity index 100% rename from csrc/sm100/prefill/dense/collective/fmha_common.hpp rename to csrc/kernels/sm100/prefill/dense/collective/fmha_common.hpp diff --git a/csrc/sm100/prefill/dense/collective/fmha_fusion.hpp b/csrc/kernels/sm100/prefill/dense/collective/fmha_fusion.hpp similarity index 100% rename from csrc/sm100/prefill/dense/collective/fmha_fusion.hpp rename to csrc/kernels/sm100/prefill/dense/collective/fmha_fusion.hpp diff --git a/csrc/sm100/prefill/dense/collective/sm100_fmha_fwd_epilogue_tma_warpspecialized.hpp b/csrc/kernels/sm100/prefill/dense/collective/sm100_fmha_fwd_epilogue_tma_warpspecialized.hpp similarity index 100% rename from csrc/sm100/prefill/dense/collective/sm100_fmha_fwd_epilogue_tma_warpspecialized.hpp rename to csrc/kernels/sm100/prefill/dense/collective/sm100_fmha_fwd_epilogue_tma_warpspecialized.hpp diff --git a/csrc/sm100/prefill/dense/collective/sm100_fmha_fwd_mainloop_tma_warpspecialized.hpp b/csrc/kernels/sm100/prefill/dense/collective/sm100_fmha_fwd_mainloop_tma_warpspecialized.hpp similarity index 100% rename from csrc/sm100/prefill/dense/collective/sm100_fmha_fwd_mainloop_tma_warpspecialized.hpp rename to csrc/kernels/sm100/prefill/dense/collective/sm100_fmha_fwd_mainloop_tma_warpspecialized.hpp diff --git a/csrc/sm100/prefill/dense/collective/sm100_fmha_load_tma_warpspecialized.hpp b/csrc/kernels/sm100/prefill/dense/collective/sm100_fmha_load_tma_warpspecialized.hpp similarity index 100% rename from csrc/sm100/prefill/dense/collective/sm100_fmha_load_tma_warpspecialized.hpp rename to csrc/kernels/sm100/prefill/dense/collective/sm100_fmha_load_tma_warpspecialized.hpp diff --git a/csrc/sm100/prefill/dense/collective/sm100_fmha_mla_fwd_mainloop_tma_warpspecialized.hpp b/csrc/kernels/sm100/prefill/dense/collective/sm100_fmha_mla_fwd_mainloop_tma_warpspecialized.hpp similarity index 100% rename from csrc/sm100/prefill/dense/collective/sm100_fmha_mla_fwd_mainloop_tma_warpspecialized.hpp rename to csrc/kernels/sm100/prefill/dense/collective/sm100_fmha_mla_fwd_mainloop_tma_warpspecialized.hpp diff --git a/csrc/sm100/prefill/dense/collective/sm100_fmha_mla_load_tma_warpspecialized.hpp b/csrc/kernels/sm100/prefill/dense/collective/sm100_fmha_mla_load_tma_warpspecialized.hpp similarity index 100% rename from csrc/sm100/prefill/dense/collective/sm100_fmha_mla_load_tma_warpspecialized.hpp rename to csrc/kernels/sm100/prefill/dense/collective/sm100_fmha_mla_load_tma_warpspecialized.hpp diff --git a/csrc/sm100/prefill/dense/common/gather_tensor.hpp b/csrc/kernels/sm100/prefill/dense/common/gather_tensor.hpp similarity index 100% rename from csrc/sm100/prefill/dense/common/gather_tensor.hpp rename to csrc/kernels/sm100/prefill/dense/common/gather_tensor.hpp diff --git a/csrc/sm100/prefill/dense/common/helper.h b/csrc/kernels/sm100/prefill/dense/common/helper.h similarity index 100% rename from csrc/sm100/prefill/dense/common/helper.h rename to csrc/kernels/sm100/prefill/dense/common/helper.h diff --git a/csrc/sm100/prefill/dense/common/mask.cuh b/csrc/kernels/sm100/prefill/dense/common/mask.cuh similarity index 100% rename from csrc/sm100/prefill/dense/common/mask.cuh rename to csrc/kernels/sm100/prefill/dense/common/mask.cuh diff --git a/csrc/sm100/prefill/dense/common/pipeline_mla.hpp b/csrc/kernels/sm100/prefill/dense/common/pipeline_mla.hpp similarity index 100% rename from csrc/sm100/prefill/dense/common/pipeline_mla.hpp rename to csrc/kernels/sm100/prefill/dense/common/pipeline_mla.hpp diff --git a/csrc/sm100/prefill/dense/common/pow_2.hpp b/csrc/kernels/sm100/prefill/dense/common/pow_2.hpp similarity index 100% rename from csrc/sm100/prefill/dense/common/pow_2.hpp rename to csrc/kernels/sm100/prefill/dense/common/pow_2.hpp diff --git a/csrc/sm100/prefill/dense/common/utils.hpp b/csrc/kernels/sm100/prefill/dense/common/utils.hpp similarity index 100% rename from csrc/sm100/prefill/dense/common/utils.hpp rename to csrc/kernels/sm100/prefill/dense/common/utils.hpp diff --git a/csrc/sm100/prefill/dense/device/fmha.hpp b/csrc/kernels/sm100/prefill/dense/device/fmha.hpp similarity index 100% rename from csrc/sm100/prefill/dense/device/fmha.hpp rename to csrc/kernels/sm100/prefill/dense/device/fmha.hpp diff --git a/csrc/sm100/prefill/dense/device/fmha_device_bwd.hpp b/csrc/kernels/sm100/prefill/dense/device/fmha_device_bwd.hpp similarity index 100% rename from csrc/sm100/prefill/dense/device/fmha_device_bwd.hpp rename to csrc/kernels/sm100/prefill/dense/device/fmha_device_bwd.hpp diff --git a/csrc/sm100/prefill/dense/fmha_cutlass_bwd_sm100.cu b/csrc/kernels/sm100/prefill/dense/fmha_cutlass_bwd_sm100.cu similarity index 100% rename from csrc/sm100/prefill/dense/fmha_cutlass_bwd_sm100.cu rename to csrc/kernels/sm100/prefill/dense/fmha_cutlass_bwd_sm100.cu diff --git a/csrc/sm100/prefill/dense/fmha_cutlass_bwd_sm100.cuh b/csrc/kernels/sm100/prefill/dense/fmha_cutlass_bwd_sm100.cuh similarity index 100% rename from csrc/sm100/prefill/dense/fmha_cutlass_bwd_sm100.cuh rename to csrc/kernels/sm100/prefill/dense/fmha_cutlass_bwd_sm100.cuh diff --git a/csrc/sm100/prefill/dense/fmha_cutlass_fwd_sm100.cu b/csrc/kernels/sm100/prefill/dense/fmha_cutlass_fwd_sm100.cu similarity index 100% rename from csrc/sm100/prefill/dense/fmha_cutlass_fwd_sm100.cu rename to csrc/kernels/sm100/prefill/dense/fmha_cutlass_fwd_sm100.cu diff --git a/csrc/sm100/prefill/dense/fmha_cutlass_fwd_sm100.cuh b/csrc/kernels/sm100/prefill/dense/fmha_cutlass_fwd_sm100.cuh similarity index 100% rename from csrc/sm100/prefill/dense/fmha_cutlass_fwd_sm100.cuh rename to csrc/kernels/sm100/prefill/dense/fmha_cutlass_fwd_sm100.cuh diff --git a/csrc/sm100/prefill/dense/interface.h b/csrc/kernels/sm100/prefill/dense/interface.h similarity index 100% rename from csrc/sm100/prefill/dense/interface.h rename to csrc/kernels/sm100/prefill/dense/interface.h diff --git a/csrc/sm100/prefill/dense/kernel/fmha_causal_tile_scheduler.hpp b/csrc/kernels/sm100/prefill/dense/kernel/fmha_causal_tile_scheduler.hpp similarity index 100% rename from csrc/sm100/prefill/dense/kernel/fmha_causal_tile_scheduler.hpp rename to csrc/kernels/sm100/prefill/dense/kernel/fmha_causal_tile_scheduler.hpp diff --git a/csrc/sm100/prefill/dense/kernel/fmha_kernel_bwd_convert.hpp b/csrc/kernels/sm100/prefill/dense/kernel/fmha_kernel_bwd_convert.hpp similarity index 100% rename from csrc/sm100/prefill/dense/kernel/fmha_kernel_bwd_convert.hpp rename to csrc/kernels/sm100/prefill/dense/kernel/fmha_kernel_bwd_convert.hpp diff --git a/csrc/sm100/prefill/dense/kernel/fmha_kernel_bwd_sum_OdO.hpp b/csrc/kernels/sm100/prefill/dense/kernel/fmha_kernel_bwd_sum_OdO.hpp similarity index 100% rename from csrc/sm100/prefill/dense/kernel/fmha_kernel_bwd_sum_OdO.hpp rename to csrc/kernels/sm100/prefill/dense/kernel/fmha_kernel_bwd_sum_OdO.hpp diff --git a/csrc/sm100/prefill/dense/kernel/fmha_options.hpp b/csrc/kernels/sm100/prefill/dense/kernel/fmha_options.hpp similarity index 100% rename from csrc/sm100/prefill/dense/kernel/fmha_options.hpp rename to csrc/kernels/sm100/prefill/dense/kernel/fmha_options.hpp diff --git a/csrc/sm100/prefill/dense/kernel/fmha_tile_scheduler.hpp b/csrc/kernels/sm100/prefill/dense/kernel/fmha_tile_scheduler.hpp similarity index 100% rename from csrc/sm100/prefill/dense/kernel/fmha_tile_scheduler.hpp rename to csrc/kernels/sm100/prefill/dense/kernel/fmha_tile_scheduler.hpp diff --git a/csrc/sm100/prefill/dense/kernel/sm100_fmha_bwd_kernel_tma_warpspecialized.hpp b/csrc/kernels/sm100/prefill/dense/kernel/sm100_fmha_bwd_kernel_tma_warpspecialized.hpp similarity index 100% rename from csrc/sm100/prefill/dense/kernel/sm100_fmha_bwd_kernel_tma_warpspecialized.hpp rename to csrc/kernels/sm100/prefill/dense/kernel/sm100_fmha_bwd_kernel_tma_warpspecialized.hpp diff --git a/csrc/sm100/prefill/dense/kernel/sm100_fmha_bwd_mla_kernel_tma_warpspecialized.hpp b/csrc/kernels/sm100/prefill/dense/kernel/sm100_fmha_bwd_mla_kernel_tma_warpspecialized.hpp similarity index 100% rename from csrc/sm100/prefill/dense/kernel/sm100_fmha_bwd_mla_kernel_tma_warpspecialized.hpp rename to csrc/kernels/sm100/prefill/dense/kernel/sm100_fmha_bwd_mla_kernel_tma_warpspecialized.hpp diff --git a/csrc/sm100/prefill/dense/kernel/sm100_fmha_fwd_kernel_tma_warpspecialized.hpp b/csrc/kernels/sm100/prefill/dense/kernel/sm100_fmha_fwd_kernel_tma_warpspecialized.hpp similarity index 100% rename from csrc/sm100/prefill/dense/kernel/sm100_fmha_fwd_kernel_tma_warpspecialized.hpp rename to csrc/kernels/sm100/prefill/dense/kernel/sm100_fmha_fwd_kernel_tma_warpspecialized.hpp diff --git a/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/config.h b/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/config.h new file mode 100644 index 00000000..1b31a9ad --- /dev/null +++ b/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/config.h @@ -0,0 +1,310 @@ +/* +Fused "Q Norm + Q RoPE + Core Attention + O RoPE + O FP8 Cast" kernel for DeepSeek V4 / V4.1 +(d_qk = d_v = 512, h_kv = 1, token-level sparse (exact top-k) attention) on SM100f. + +Instead of launching separate kernels for RMSNorm, RoPE, attention, and output quantization, this +kernel fuses the whole "q_b_proj output -> wv_proj input" segment of the MLA block: +1. Q Norm (V4 only, MODEL_TYPE == ModelType::V4): computes the per-head RMSNorm denominator + rsqrt(sum(q^2)/d_qk + eps) on the fly while loading Q, and folds it into the softmax scale. + V4.1 (ModelType::V41) skips this step +2. Q RoPE: applies non-neox-style RoPE (rope_dim = 64) to the last D_ROPE dims of each Q head, + using `token_positions` and `cos_sin_cache` +3. Core attention: token-level sparse attention over the (at most) `topk` KV tokens selected by + `indices`. Invalid indices (< 0 or >= s_kv) and positions beyond `topk_length` are masked out. + Supports an optional per-head `attn_sink` (affects output but not lse / max_logits) +4. O RoPE: applies the conjugate RoPE to the last D_ROPE dims of each output head +5. O FP8 Cast: quantizes the output to fp8_e4m3 with per-32-element ue8m0 scale factors + (round_sf + packed ue8m0, TMA-aligned col-major sf layout) + +Template parameters: +- FWD_MODE: SparseAttnFwdMode::Prefill or SparseAttnFwdMode::Decode +- MODEL_TYPE: ModelType::V4 (Q norm enabled, V4 KV cache layout) or ModelType::V41 +- EXTRA_MODEL_TYPE: the format of the extra KV cache (decode only), MODEL_TYPE or ModelType::V41_FP4 + (V4.1 with fp4 KV cache, e4m3 per-16 scales) with MODEL_TYPE == V41 +- H_Q: number of Q heads, 64 or 128 + +I/O (see `csrc/api/fused_norm_rope_attn_rope_cast_fwd.cpp` and the Python docstrings in +`flash_mla/fused_norm_rope_attn_rope_cast.py` for the full parameter list): +- q: [s_q, h_q, d_qk], bf16, WITHOUT RoPE applied, in the PERMUTED layout produced by the + `permute_q_b_proj` kernel (16-element d-chunks interleaved across heads). Each token's + h_q*d_qk elements must be contiguous +- Prefill: kv [s_kv, 1, d_qk] (bf16, non-paged) + indices [s_q, 1, topk] + Decode: paged quantized KV cache(s) (V4 / V4.1 / V4.1 fp4 format, see below) + indices_in_kvcache [s_q, topk], + plus an optional secondary ("extra") KV cache with its own indices / topk_length +- out_fp8: [s_q, n_wv_group, wv_group_size * d_v], fp8_e4m3, in the permuted layout expected by + the `permute_wv_proj`-transformed weights; out_sf: packed ue8m0 scale factors (always per-32) +- lse / max_logits (prefill only for max_logits): [s_q, h_q], fp32 + +Execution structure: +- Persistent kernel scheduled via CLC (Cluster Launch Control); the grid is + (s_q * CLUSTER_SIZE, 1, 1) and each cluster processes one query token per job +- CLUSTER_SIZE = ceil(H_Q / 64): 1 CTA for h64, a 2-CTA cluster (dual-CTA UMMA, CTA0 owns + V[:, 0:256] and CTA1 owns V[:, 256:512]) for h128. Each CTA has 512 threads (4 warpgroups): + - WG0: Q fetching (q_sqr_sum + Q RoPE + store to TMEM) & O epilogue (TMEM load, O RoPE, + FP8 quant, store to gmem). Timeline: Q0 Q1 O0 Q2 O1 ... Qn O(n-1) On + - WG1: KV producer. Prefill: gathers the bf16 KV block via TMA gather4. Decode: loads the fp8 / fp4 + part from the paged cache and dequantizes it into the smem KV slots in registers + - WG2: MMA warp (warp 8, issues UMMAs on CTA0 only), CLC warp (warp 9), indices / validity + mask generator (warp 10), and (decode only) TMA warp for the bf16 KV part (warp 11) + - WG3: Scale & Exp: reduces P, maintains the online softmax state (mi / li), produces S +- KV tokens are processed in blocks of B_TOPK (96 for 2-CTA, 64 for 1-CTA) with NUM_KV_SLOTS-deep + software pipelining + +Multi-rail GeMM is always used: +- For cases where CLUSTER_SIZE is 1, we FOLD Q by FOLD_FACTOR and do a "batched gemm" in batch size = FOLD_FACTOR, and reduce P on shared memory +- For cases where CLUSTER_SIZE is 2, we also fold Q by FOLD_FACTOR and perform batched GeMM, and reduce P on shared memory + +Decoding mode (FWD_MODE == SparseAttnFwdMode::Decode): +- Batch size must be 1; the grid covers the s_q query tokens +- The KV cache is a paged FP8 cache with the same format as `sm100::decode::sparse::head64` (V4 layout). + The fp8 (D_FP8) part is loaded from global memory and dequantized in-place into the KV slots by + warpgroup 1 (no intermediate raw-fp8 smem buffer), while the bf16 (D_BF16, RoPE) part is loaded via + TMA gather4 by warp 11. An optional extra (secondary) KV cache is supported +- EXTRA_MODEL_TYPE == ModelType::V41_FP4 selects an fp4 extra KV cache: every token is 512 e2m1 + 32 e4m3 scales, + each page block stores [page_block_size x 256 B data rows] + [page_block_size x 32 B scale rows]. Warpgroup 1 then + dequantizes both caches with a common code path (see there) instead of the fp8-only one +- Split-KV is not supported (the fused RoPE + FP8-quant epilogue cannot be combined across splits) +*/ + +#pragma once + +#include +#include +#include + +#include "kernels/defines.h" +#include "kernels/params.h" +#include "kernels/kv_cache_format.h" + +#include "kernel.h" + +namespace sm100::prefill::fused_norm_rope_attn_rope_cast_fwd::core_attn { + +using namespace cute; + +template +struct Kernel { + +static constexpr SparseAttnFwdMode FWD_MODE = CONFIG.FWD_MODE; +static constexpr ModelType MODEL_TYPE = CONFIG.MODEL_TYPE; +static constexpr ModelType EXTRA_MODEL_TYPE = CONFIG.EXTRA_MODEL_TYPE; +static constexpr uint32_t H_Q = CONFIG.H_Q; + +using Params = ParamT; + +static_assert(FWD_MODE == SparseAttnFwdMode::Prefill || FWD_MODE == SparseAttnFwdMode::Decode); +static_assert(H_Q == 64 || H_Q == 128); + +static constexpr bool IS_DECODE = is_decode_v; + +// Model parameters +static constexpr uint32_t D_QK = 512; +static constexpr uint32_t D_VO = 512; +static constexpr uint32_t O_QUANT_TILE_SIZE = 32; +static constexpr uint32_t D_ROPE = 64; +static constexpr uint32_t D_NOPE = D_QK - D_ROPE; +static constexpr uint32_t WV_GROUP_SIZE = 8; +static constexpr bool ENABLE_Q_NORM = CONFIG.ENABLE_Q_NORM; + +// Cluster shape selection +static constexpr uint32_t CLUSTER_SIZE = ku::ceil_div((uint32_t)H_Q, 64u); +static constexpr uint32_t IS_2CTA = CLUSTER_SIZE == 2; + +// Tiling Shape Selection +static constexpr uint32_t B_TOPK = CLUSTER_SIZE == 2 ? 96 : 64; +static constexpr uint32_t H_Q_PER_CTA = 64; + +// Paged quantized KV cache format for decoding, plus the constants of this kernel's common dequant path. The members are re-exported as uint32_t +template +struct KVFormat { + using Base = KVCacheFormat; + static constexpr bool IS_FP4 = Base::IS_FP4; + static constexpr uint32_t D_FP4 = Base::D_FP4; + static constexpr uint32_t D_FP8 = Base::D_FP8; + static constexpr uint32_t D_BF16 = Base::D_BF16; + static constexpr uint32_t QUANT_TILE_SIZE = Base::QUANT_TILE_SIZE; + static constexpr uint32_t NUM_SCALES_EACH_TOKEN = Base::NUM_SCALES_EACH_TOKEN; + static constexpr uint32_t TMA_K_STRIDE = Base::TMA_K_STRIDE; + static constexpr uint32_t BYTES_PER_TOKEN = Base::BYTES_PER_TOKEN; + // The common dequant path of WG1 only (HAS_FP4_KV). A raw row is this CTA's part of the quantized data of a token; it is + // gathered with a box 16 B wider than its data (zero-filled by TMA) so that RAW_TOKEN_SMEM_STRIDE / 16 is odd, which keeps the + // LDS.128 of WG1 free of bank conflicts. A chunk is 16 B of raw data, by one LDS.128 + static constexpr uint32_t RAW_TOKEN_DATA_BYTES = Base::QUANT_BYTES / CLUSTER_SIZE; + static constexpr uint32_t RAW_TOKEN_SMEM_STRIDE = RAW_TOKEN_DATA_BYTES + 16; // 272 / 144 (fp4), 528 / 272 (V41 fp8) + static constexpr uint32_t NUM_CHUNKS_PER_ROW = RAW_TOKEN_DATA_BYTES / 16; + static constexpr uint32_t CHUNK_ELEMS = Base::IS_FP4 ? 32 : 16; +}; +using OrigKVFormat = KVFormat; // Format of `kv` +using ExtraKVFormat = KVFormat; // Format of `extra_kv` +static_assert(is_valid_kv_format_pair(MODEL_TYPE, EXTRA_MODEL_TYPE)); +static constexpr bool HAS_FP4_KV = ExtraKVFormat::IS_FP4; // Selects the common (fp8 + fp4) dequant path of WG1 instead of the fp8-only one +// Shorthands for the format of `kv`, which is also the format of `extra_kv` unless HAS_FP4_KV +static constexpr uint32_t D_FP8 = OrigKVFormat::D_FP8; +static constexpr uint32_t D_BF16 = OrigKVFormat::D_BF16; +static constexpr uint32_t KV_QUANT_TILE_SIZE = OrigKVFormat::QUANT_TILE_SIZE; +static constexpr uint32_t NUM_SCALES_EACH_TOKEN = OrigKVFormat::NUM_SCALES_EACH_TOKEN; +static constexpr uint32_t NUM_SCALES_EACH_TOKEN_PER_CTA = NUM_SCALES_EACH_TOKEN / CLUSTER_SIZE; +static constexpr uint32_t TMA_K_STRIDE = OrigKVFormat::TMA_K_STRIDE; +static constexpr uint32_t KV_CACHE_BYTES_PER_TOKEN = OrigKVFormat::BYTES_PER_TOKEN; +// The common dequant path gathers the raw rows of a KV block into the beginning of the KV slot (one 4-row group per gather4) and +// dequantizes them in place. A group is padded to 128 B since the destination of a gather4 must be 128 B aligned +static constexpr uint32_t RAW_KV_GROUP_BYTES = ku::ceil_div(4 * std::max(OrigKVFormat::RAW_TOKEN_SMEM_STRIDE, ExtraKVFormat::RAW_TOKEN_SMEM_STRIDE), 128u) * 128; +static_assert(!HAS_FP4_KV || B_TOPK / 4 * RAW_KV_GROUP_BYTES <= B_TOPK * D_QK / CLUSTER_SIZE * sizeof(bf16)); + +static constexpr uint32_t NUM_THREADS = 512; +static constexpr uint32_t NUM_WORKING_THREADS = + CLUSTER_SIZE == 1 ? ( + IS_DECODE ? + 128 + 128 + (1+1+32) + 128 : // WG0 + WG3 + (MMA + CLC + indices) + WG1 (dequant) + 128 + 128 + (1+1+32) + 4 // WG0 + WG3 + (MMA + CLC + indices) + WG1 (KV producer, 1 elected thread per warp) + ) : ( + IS_DECODE ? + 128*2 + 128*2 + (1+2+32*2) + 128*2 : + 128*2 + 128*2 + (1+2+32*2) + 4*2 + ); + +static constexpr uint32_t FOLD_FACTOR = 128 / H_Q_PER_CTA; +static constexpr uint32_t NUM_MRGEMM_RAILS = 2; // The number of "rails" (batch size) during multi-rail GeMM. Currently must be 2 +static constexpr uint32_t NUM_P_ELEMS_PER_THREAD = H_Q_PER_CTA * B_TOPK / 128; + +static constexpr uint32_t NUM_KV_SLOTS = 3; +static constexpr uint32_t NUM_INDICES_BUFS = 4; +static constexpr uint32_t NUM_P_BUFS = CLUSTER_SIZE == 2 ? 1 : 2; +static constexpr uint32_t NEED_TP_EMPTY_BAR = NUM_P_BUFS == 1; // Don't need to wait for P's emptiness as long as P has >= 2 buffers, since "we are issuing P[i]" <-- "O[i-2] has been issued" <-- "S[i-2] is ready" <-- "P[i-2] is free" + +struct tmem_cols { + static constexpr uint32_t O = 0; + static constexpr uint32_t Q = O + D_VO / FOLD_FACTOR; + static constexpr uint32_t P_0 = Q + D_QK / NUM_MRGEMM_RAILS / 2; // /2 since 2 bf16 is packed in 1 uint32 + static constexpr uint32_t P_1 = P_0 + B_TOPK*NUM_MRGEMM_RAILS/FOLD_FACTOR; + + static constexpr uint32_t get_p(const uint32_t &p_buf_idx) { + if constexpr (NUM_P_BUFS == 1) { + return P_0; + } else if constexpr (NUM_P_BUFS == 2) { + return p_buf_idx ? P_1 : P_0; + } else { + static_assert(NUM_P_BUFS == 1 || NUM_P_BUFS == 2); + } + } + static_assert(get_p(NUM_P_BUFS-1) + B_TOPK*NUM_MRGEMM_RAILS/FOLD_FACTOR <= 512); +}; + +using MMAAtom_QK = cute::conditional_t< + CLUSTER_SIZE == 2, + SM100_MMA_F16BF16_2x1SM_TS_NOELECT, + SM100_MMA_F16BF16_WS_TS_NOELECT +>; +using TiledMMA_QK = decltype(make_tiled_mma(MMAAtom_QK{})); +using TiledMMA_SV = cute::conditional_t< + CLUSTER_SIZE == 2, + decltype(make_tiled_mma( + SM100_MMA_F16BF16_2x1SM_SS_NOELECT{}, + Layout>{}, + Tile, Layout, Stride<_1, _256, _128>>, _16>{} // We use this permutation layout to let CTA0 takes V[:, 0:256] and CTA1 takes V[:, 256:512] + )), + decltype(make_tiled_mma(SM100_MMA_F16BF16_WS_SS_NOELECT{})) +>; + +struct SharedMemoryPlan { + CUTE_ALIGNAS(1024) bf16 kv_slots[NUM_KV_SLOTS][B_TOPK * D_QK / CLUSTER_SIZE]; // Cluster size = 1: the whole KV; cluster size = 2: half KV + CUTE_ALIGNAS(1024) bf16 s[H_Q_PER_CTA * B_TOPK]; + CUTE_ALIGNAS(1024) float p_exchange_buf[4][32*NUM_P_ELEMS_PER_THREAD]; + CUTE_ALIGNAS(1024) uint8_t is_k_valid[NUM_INDICES_BUFS][ku::find_next_power_of_2(B_TOPK/8)]; + // Decode: WG10 produces metadata once for the four dequant warps. 16 B aligned so that the fp4 path can read + // the coordinates of 4 consecutive rows with one LDS.128 + CUTE_ALIGNAS(16) int decode_tma_coords[IS_DECODE ? NUM_INDICES_BUFS : 0][B_TOPK]; + uint8_t decode_scales[IS_DECODE ? NUM_INDICES_BUFS : 0][B_TOPK * NUM_SCALES_EACH_TOKEN_PER_CTA]; // Scales of the fp8 tokens + // fp4 tokens: 4 buffers of 32 B scales per token do not fit into shared memory, so the dequant warps load the scales themselves + // and only get their global addresses here (nullptr for an invalid token) + const uint8_t *decode_scale_ptrs[(IS_DECODE && HAS_FP4_KV) ? NUM_INDICES_BUFS : 0][B_TOPK]; + float q_sqr_sum_buf[ENABLE_Q_NORM ? 2 : 0][128]; // We have 2 q_sqr_sum_buf to save some barriers, as Q[i+2] starts to fetch -> O[i] have finished -> Q[i]'s q_sqr_sum_buf is useless + float rowwise_max_buf[128]; + float rowwise_mi_buf[H_Q_PER_CTA]; + float rowwise_li_buf[128]; // 128: warpgroup size + + transac_bar_t bar_kv_slot_full[NUM_KV_SLOTS], bar_kv_slot_empty[NUM_KV_SLOTS]; + transac_bar_t bar_indices_full[NUM_INDICES_BUFS], bar_indices_empty[NUM_INDICES_BUFS]; + transac_bar_t bar_tQ_empty, bar_tQ_full; + transac_bar_t bar_q_sqr_sum_full; // Only used for 2-CTA + transac_bar_t bar_tO_empty, bar_tO_full; + transac_bar_t bar_tP_full[NUM_P_BUFS], bar_tP_empty[NEED_TP_EMPTY_BAR ? NUM_P_BUFS : 0]; + transac_bar_t bar_SO_full, bar_SO_empty; + transac_bar_t bar_clc_full, bar_clc_empty; + transac_bar_t bar_li_mi_full, bar_li_mi_empty; + transac_bar_t bar_raw_kv_full; + + ku::CLCResponseObj clc_response_obj; + array_aligned tmem_start_addr; +}; +static_assert(sizeof(SharedMemoryPlan) <= 227 * 1024); + +struct TMAParams { + // Prefill only + CUtensorMap tensor_map_kv; // the whole (bf16, non-paged) KV cache + // Decode only + CUtensorMap tensor_map_kv_fp8_part_cta0; + CUtensorMap tensor_map_kv_fp8_part_cta1; + CUtensorMap tensor_map_extra_kv_fp8_part_cta0; + CUtensorMap tensor_map_extra_kv_fp8_part_cta1; + CUtensorMap tensor_map_kv_bf16_part; // the bf16 (RoPE) part of the paged KV cache + CUtensorMap tensor_map_extra_kv_bf16_part; // the bf16 (RoPE) part of the extra paged KV cache. Invalid if extra_topk == 0 + CUtensorMap tensor_map_extra_kv_fp4_part_cta0; + CUtensorMap tensor_map_extra_kv_fp4_part_cta1; +}; +static constexpr uint32_t D_FP8_CTA0 = CLUSTER_SIZE == 1 ? D_FP8 : D_VO/2; +static constexpr uint32_t D_FP8_CTA1 = D_FP8 - D_FP8_CTA0; +static constexpr bool IS_CTA0_RAW_KV_PADDED = D_FP8_CTA0 % 128 == 0; +static constexpr bool IS_CTA1_RAW_KV_PADDED = D_FP8_CTA1 % 128 == 0; +// Bytes of one raw fp8 row in shared memory, i.e. the box of the fp8 tensor maps: the fp8-only dequant path pads a row by 64 B +// when needed, the common path (HAS_FP4_KV) by 16 B (KVFormat::RAW_TOKEN_SMEM_STRIDE) +static constexpr uint32_t RAW_FP8_TOKEN_SMEM_STRIDE_CTA0 = HAS_FP4_KV ? OrigKVFormat::RAW_TOKEN_SMEM_STRIDE : D_FP8_CTA0 + (IS_CTA0_RAW_KV_PADDED ? 64 : 0); +static constexpr uint32_t RAW_FP8_TOKEN_SMEM_STRIDE_CTA1 = HAS_FP4_KV ? OrigKVFormat::RAW_TOKEN_SMEM_STRIDE : D_FP8_CTA1 + (IS_CTA1_RAW_KV_PADDED ? 64 : 0); + +using AllocatorT = std::conditional_t; + +struct AuxParams { + cutlass::FastDivmod fast_divmod_page_block_size; + cutlass::FastDivmod fast_divmod_extra_page_block_size; +}; + +// Some helper functions for buffer arrival +static __device__ __forceinline__ void umma_arrive_on_every_cta(transac_bar_t &bar) { + // Perform UMMA arrive, possibly with multicast, to every CTA + if constexpr (IS_2CTA) { + ku::umma_arrive_multicast_2x1SM_noelect(bar, 1|2); + } else { + ku::umma_arrive_noelect(bar); + } +} +static __device__ __forceinline__ void umma_arrive_on_cta0(transac_bar_t &bar) { + // Perform UMMA arrive on CTA0 + if constexpr (IS_2CTA) { + ku::umma_arrive_2x1SM_noelect(bar); + } else { + ku::umma_arrive_noelect(bar); + } +} +static __device__ __forceinline__ void arrive_on_cta0_barrier(transac_bar_t &bar) { + if constexpr (IS_2CTA) { + bar.arrive(0u); + } else { + bar.arrive(); + } +} + +struct barrier_ids { + static constexpr int WG0_SYNC = 0; + static constexpr int WG3_SYNC = 1; + static constexpr int WG3_WARP02_SYNC = 2; + static constexpr int WG3_WARP13_SYNC = 3; +}; + +static __device__ __forceinline__ void +devfunc(const Params ¶ms, const TMAParams &tma_params, const AuxParams &aux_params); + +static void run(const Params ¶ms); + +}; + +} diff --git a/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41_h128_decode_nonorm.cu b/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41_h128_decode_nonorm.cu new file mode 100644 index 00000000..0518b910 --- /dev/null +++ b/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41_h128_decode_nonorm.cu @@ -0,0 +1,8 @@ +#include "../kernel.h" +#include "../kernel.cuh" + +namespace sm100::prefill::fused_norm_rope_attn_rope_cast_fwd::core_attn { + +template void run_fused_norm_rope_attn_rope_cast_fwd_kernel(const ParamT& params); + +} diff --git a/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41_h128_decode_norm.cu b/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41_h128_decode_norm.cu new file mode 100644 index 00000000..2db860f5 --- /dev/null +++ b/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41_h128_decode_norm.cu @@ -0,0 +1,8 @@ +#include "../kernel.h" +#include "../kernel.cuh" + +namespace sm100::prefill::fused_norm_rope_attn_rope_cast_fwd::core_attn { + +template void run_fused_norm_rope_attn_rope_cast_fwd_kernel(const ParamT& params); + +} diff --git a/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41_h64_decode_nonorm.cu b/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41_h64_decode_nonorm.cu new file mode 100644 index 00000000..43f765fd --- /dev/null +++ b/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41_h64_decode_nonorm.cu @@ -0,0 +1,8 @@ +#include "../kernel.h" +#include "../kernel.cuh" + +namespace sm100::prefill::fused_norm_rope_attn_rope_cast_fwd::core_attn { + +template void run_fused_norm_rope_attn_rope_cast_fwd_kernel(const ParamT& params); + +} diff --git a/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41_h64_decode_norm.cu b/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41_h64_decode_norm.cu new file mode 100644 index 00000000..2b5b36cc --- /dev/null +++ b/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41_h64_decode_norm.cu @@ -0,0 +1,8 @@ +#include "../kernel.h" +#include "../kernel.cuh" + +namespace sm100::prefill::fused_norm_rope_attn_rope_cast_fwd::core_attn { + +template void run_fused_norm_rope_attn_rope_cast_fwd_kernel(const ParamT& params); + +} diff --git a/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41fp4_h128_decode_nonorm.cu b/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41fp4_h128_decode_nonorm.cu new file mode 100644 index 00000000..97fd35a9 --- /dev/null +++ b/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41fp4_h128_decode_nonorm.cu @@ -0,0 +1,8 @@ +#include "../kernel.h" +#include "../kernel.cuh" + +namespace sm100::prefill::fused_norm_rope_attn_rope_cast_fwd::core_attn { + +template void run_fused_norm_rope_attn_rope_cast_fwd_kernel(const ParamT& params); + +} diff --git a/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41fp4_h128_decode_norm.cu b/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41fp4_h128_decode_norm.cu new file mode 100644 index 00000000..0d4e152c --- /dev/null +++ b/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41fp4_h128_decode_norm.cu @@ -0,0 +1,8 @@ +#include "../kernel.h" +#include "../kernel.cuh" + +namespace sm100::prefill::fused_norm_rope_attn_rope_cast_fwd::core_attn { + +template void run_fused_norm_rope_attn_rope_cast_fwd_kernel(const ParamT& params); + +} diff --git a/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41fp4_h64_decode_nonorm.cu b/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41fp4_h64_decode_nonorm.cu new file mode 100644 index 00000000..dca56874 --- /dev/null +++ b/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41fp4_h64_decode_nonorm.cu @@ -0,0 +1,8 @@ +#include "../kernel.h" +#include "../kernel.cuh" + +namespace sm100::prefill::fused_norm_rope_attn_rope_cast_fwd::core_attn { + +template void run_fused_norm_rope_attn_rope_cast_fwd_kernel(const ParamT& params); + +} diff --git a/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41fp4_h64_decode_norm.cu b/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41fp4_h64_decode_norm.cu new file mode 100644 index 00000000..d71f00a0 --- /dev/null +++ b/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41fp4_h64_decode_norm.cu @@ -0,0 +1,8 @@ +#include "../kernel.h" +#include "../kernel.cuh" + +namespace sm100::prefill::fused_norm_rope_attn_rope_cast_fwd::core_attn { + +template void run_fused_norm_rope_attn_rope_cast_fwd_kernel(const ParamT& params); + +} diff --git a/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h128_decode_nonorm.cu b/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h128_decode_nonorm.cu new file mode 100644 index 00000000..96070152 --- /dev/null +++ b/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h128_decode_nonorm.cu @@ -0,0 +1,8 @@ +#include "../kernel.h" +#include "../kernel.cuh" + +namespace sm100::prefill::fused_norm_rope_attn_rope_cast_fwd::core_attn { + +template void run_fused_norm_rope_attn_rope_cast_fwd_kernel(const ParamT& params); + +} diff --git a/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h128_decode_norm.cu b/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h128_decode_norm.cu new file mode 100644 index 00000000..e11b9234 --- /dev/null +++ b/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h128_decode_norm.cu @@ -0,0 +1,8 @@ +#include "../kernel.h" +#include "../kernel.cuh" + +namespace sm100::prefill::fused_norm_rope_attn_rope_cast_fwd::core_attn { + +template void run_fused_norm_rope_attn_rope_cast_fwd_kernel(const ParamT& params); + +} diff --git a/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h128_prefill_nonorm.cu b/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h128_prefill_nonorm.cu new file mode 100644 index 00000000..416fa279 --- /dev/null +++ b/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h128_prefill_nonorm.cu @@ -0,0 +1,8 @@ +#include "../kernel.h" +#include "../kernel.cuh" + +namespace sm100::prefill::fused_norm_rope_attn_rope_cast_fwd::core_attn { + +template void run_fused_norm_rope_attn_rope_cast_fwd_kernel(const ParamT& params); + +} diff --git a/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h128_prefill_norm.cu b/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h128_prefill_norm.cu new file mode 100644 index 00000000..c611d14d --- /dev/null +++ b/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h128_prefill_norm.cu @@ -0,0 +1,8 @@ +#include "../kernel.h" +#include "../kernel.cuh" + +namespace sm100::prefill::fused_norm_rope_attn_rope_cast_fwd::core_attn { + +template void run_fused_norm_rope_attn_rope_cast_fwd_kernel(const ParamT& params); + +} diff --git a/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h64_decode_nonorm.cu b/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h64_decode_nonorm.cu new file mode 100644 index 00000000..645eac49 --- /dev/null +++ b/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h64_decode_nonorm.cu @@ -0,0 +1,8 @@ +#include "../kernel.h" +#include "../kernel.cuh" + +namespace sm100::prefill::fused_norm_rope_attn_rope_cast_fwd::core_attn { + +template void run_fused_norm_rope_attn_rope_cast_fwd_kernel(const ParamT& params); + +} diff --git a/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h64_decode_norm.cu b/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h64_decode_norm.cu new file mode 100644 index 00000000..eb7289e4 --- /dev/null +++ b/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h64_decode_norm.cu @@ -0,0 +1,8 @@ +#include "../kernel.h" +#include "../kernel.cuh" + +namespace sm100::prefill::fused_norm_rope_attn_rope_cast_fwd::core_attn { + +template void run_fused_norm_rope_attn_rope_cast_fwd_kernel(const ParamT& params); + +} diff --git a/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h64_prefill_nonorm.cu b/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h64_prefill_nonorm.cu new file mode 100644 index 00000000..ba749ff4 --- /dev/null +++ b/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h64_prefill_nonorm.cu @@ -0,0 +1,8 @@ +#include "../kernel.h" +#include "../kernel.cuh" + +namespace sm100::prefill::fused_norm_rope_attn_rope_cast_fwd::core_attn { + +template void run_fused_norm_rope_attn_rope_cast_fwd_kernel(const ParamT& params); + +} diff --git a/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h64_prefill_norm.cu b/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h64_prefill_norm.cu new file mode 100644 index 00000000..6f86f6a1 --- /dev/null +++ b/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h64_prefill_norm.cu @@ -0,0 +1,8 @@ +#include "../kernel.h" +#include "../kernel.cuh" + +namespace sm100::prefill::fused_norm_rope_attn_rope_cast_fwd::core_attn { + +template void run_fused_norm_rope_attn_rope_cast_fwd_kernel(const ParamT& params); + +} diff --git a/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/kernel.cuh b/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/kernel.cuh new file mode 100644 index 00000000..b55c0855 --- /dev/null +++ b/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/kernel.cuh @@ -0,0 +1,1624 @@ +#pragma once + +#include +#include // CUDART_INF_F +#include +#include +#include +#include + +#include "kernels/utils.h" +#include "kernels/sm100/helpers.h" +#include "kernels/sm100/common_subroutine.h" + +#include "config.h" + +namespace sm100::prefill::fused_norm_rope_attn_rope_cast_fwd::core_attn { + +static constexpr float MAX_INIT_VAL = -1e30; +static constexpr float O_QUANT_CLAMP_MIN_VALUE = 1e-4; + +#ifdef KERUTILS_ENABLE_SM103A +static constexpr bool IS_TMEM_LD_WITH_RED_AVAILABLE = true; +#else +static constexpr bool IS_TMEM_LD_WITH_RED_AVAILABLE = false; +#endif + +__device__ __forceinline__ +float2 apply_rope(const float2 &x, const float &cur_cos, const float &cur_sin) { + float2 a = {x.x, x.x}; + float2 b = {cur_cos, cur_sin}; + float2 c = {-x.y * cur_sin, +x.y * cur_cos}; + float2 y = ku::float2_fma(a, b, c); + return y; +} + +// To achieve prefill - decoding alignment while using block_size = 96, decoding must reproduce +// prefill's "natural" blocking of the concatenated indices array ([topk orig slots; extra_topk +// extra slots], tiled by B_TOPK). Since topk (e.g. 128) may not be a multiple of B_TOPK (e.g. 96), +// one KV block may straddle the orig/extra boundary, i.e. the last (partial) block of the orig KV +// and the first tokens of the extra KV are "stitched" into one block. Each KV block is thus +// classified into one of the following three categories: +enum class KVLocation { + ORIG, // All slots of this block come from `kv`/`indices` + ORIG_AND_EXTRA, // This block straddles the boundary: slots < num_orig_slots come from `kv`/`indices`, the rest from `extra_kv`/`extra_indices` + EXTRA // All slots of this block come from `extra_kv`/`extra_indices` +}; + +template +__device__ __forceinline__ +void Kernel::devfunc(const Params ¶ms, const TMAParams &tma_params, const AuxParams &aux_params) { +#ifdef KERUTILS_ENABLE_SM100A + const uint32_t cta_idx = IS_2CTA ? blockIdx.x % 2 : 0; + const uint32_t warp_idx = cutlass::canonical_warp_idx_sync(); + const uint32_t warpgroup_idx = __shfl_sync(0xffffffff, threadIdx.x / 128, 0); + const uint32_t idx_in_warpgroup = threadIdx.x % 128; + const uint32_t lane_idx = threadIdx.x % 32; + + extern __shared__ char smem_buf[]; + SharedMemoryPlan &smem = *reinterpret_cast(smem_buf); + + if constexpr (IS_2CTA) { + ku::barrier_cluster_arrive_relaxed(); + ku::barrier_cluster_wait_acquire(); + } + + if (warp_idx == 0 && elect_one_sync()) { + // Prefetch TMA descriptors + if constexpr (IS_DECODE) { + if constexpr (D_BF16 > 0) { + cute::prefetch_tma_descriptor(&tma_params.tensor_map_kv_bf16_part); + } + } else { + cute::prefetch_tma_descriptor(&tma_params.tensor_map_kv); + } + } else if (warp_idx == 1 && elect_one_sync()) { + // Init barriers + CUTE_UNROLL + for (uint32_t i = 0; i < NUM_KV_SLOTS; ++i) { + // bar_kv_slot_full: + // Prefill: 1 arrive (arrive_and_expect_tx from the MMA warp) + TMA transactions of the whole KV block + // Decode: 128 arrives (the dequant warpgroup) from each CTA + 1 arrive (arrive_and_expect_tx from the MMA warp) from CTA0 + TMA transactions of the bf16 (RoPE) part + smem.bar_kv_slot_full[i].init(IS_DECODE ? 128*CLUSTER_SIZE + 1 : 1); // bar_kv_full: Every CTA -> CTA0 + smem.bar_kv_slot_empty[i].init(1); // bar_kv_empty: CTA0 -> Every CTA + } + CUTE_UNROLL + for (uint32_t i = 0; i < NUM_INDICES_BUFS; ++i) { + smem.bar_indices_full[i].init(32); // CTA-local + smem.bar_indices_empty[i].init(IS_DECODE ? 256 : 128); // CTA-local + } + CUTE_UNROLL + for (uint32_t i = 0; i < NUM_P_BUFS; ++i) { + smem.bar_tP_full[i].init(1); // CTA0 -> Every CTA + if constexpr (NEED_TP_EMPTY_BAR) { + smem.bar_tP_empty[i].init(128*CLUSTER_SIZE); // Every CTA -> CTA0 + } + } + if constexpr (IS_2CTA && ENABLE_Q_NORM) { + smem.bar_q_sqr_sum_full.init(128); // CTA-local + } + smem.bar_clc_full.init(1); // CTA0 -> Every CTA + smem.bar_clc_empty.init(cta_idx == 1 ? 1 : NUM_WORKING_THREADS); // Every CTA -> CTA0 + smem.bar_tQ_full.init(128*CLUSTER_SIZE); // Every CTA -> CTA0 + smem.bar_tQ_empty.init(1+128); // CTA0 -> Every CTA (arrive by MMA thread), as well as CTA-local (arrived by exp warpgroup) + smem.bar_tO_full.init(1); // CTA0 -> Every CTA + smem.bar_tO_empty.init(128*CLUSTER_SIZE); // Every CTA -> CTA0 + smem.bar_SO_full.init(128*CLUSTER_SIZE); // Every CTA -> CTA0 + smem.bar_SO_empty.init(1); // CTA0 -> Every CTA + smem.bar_li_mi_full.init(128); // CTA-local + smem.bar_li_mi_empty.init(128); // CTA-local + if constexpr (IS_DECODE) { + smem.bar_raw_kv_full.init(1); // CTA-local + } + fence_barrier_init(); + } else if (warp_idx == 3) { + // Allocate TMEM + AllocatorT().allocate(512, smem.tmem_start_addr.data()); + AllocatorT().release_allocation_lock(); + KU_TRAP_ONLY_DEVICE_ASSERT(smem.tmem_start_addr.data()[0] == 0); + } + + if constexpr (IS_2CTA) { + ku::barrier_cluster_arrive_relaxed(); + ku::barrier_cluster_wait_acquire(); + } else { + __syncthreads(); + } + + struct OuterloopArgs { + bool is_valid; + uint32_t s_q_idx; + uint32_t job_idx_mod_2; + uint32_t topk_length; + uint32_t num_kv_blocks; + // Decoding only: + uint32_t extra_topk_length; // Number of valid extra-topk entries of the current request + uint32_t num_orig_slots; // Number of slots occupied by the orig KV in the unified slot space, i.e. slots < num_orig_slots come from `kv`/`indices` (will be set to -1 if we don't have so many valid indices) while the others come from `extra_kv`/`extra_indices`. 0xFFFFFFFF when there is no extra KV (so that every slot belongs to the orig KV) + }; + + auto _make_outer_loop_args = [&](uint32_t job_idx_mod_2, uint32_t cta_x_idx) -> OuterloopArgs { + uint32_t s_q_idx = cta_x_idx / CLUSTER_SIZE; + if constexpr (IS_DECODE) { + uint32_t topk_length = params.topk_length ? (uint32_t)__ldg(params.topk_length + s_q_idx) : (uint32_t)params.topk; + uint32_t extra_topk_length = params.extra_topk_length ? (uint32_t)__ldg(params.extra_topk_length + s_q_idx) : (uint32_t)params.extra_topk; + bool have_extra_kv = params.extra_topk > 0; + uint32_t num_orig_slots, num_kv_blocks; + if (have_extra_kv) { + num_orig_slots = (uint32_t)params.topk; + num_kv_blocks = ku::ceil_div(num_orig_slots + extra_topk_length, (uint32_t)B_TOPK); // When extra_kv is present, always round topk up to the full cycle + } else { + num_orig_slots = 0xFFFFFFFFu; + num_kv_blocks = ku::ceil_div(topk_length, (uint32_t)B_TOPK); + } + num_kv_blocks = std::max(num_kv_blocks, 1u); + return { + true, + s_q_idx, + job_idx_mod_2, + topk_length, + num_kv_blocks, + extra_topk_length, + num_orig_slots + }; + } else { + uint32_t topk_length = params.topk_length ? __ldg(params.topk_length + s_q_idx) : params.topk; + uint32_t num_kv_blocks = std::max(ku::ceil_div(topk_length, (uint32_t)B_TOPK), 1u); + return { + true, + s_q_idx, + job_idx_mod_2, + topk_length, + num_kv_blocks + }; + } + }; + + // A handy function (decode only) to run along all KV blocks in the unified slot space. + // Should be provided with a template function, which will be invoked as callable(kv_block_idx), + // where LOC (a KVLocation) tells where the tokens of the current block come from. Only the (at most + // one) block straddling the orig/extra boundary is invoked with ORIG_AND_EXTRA, so that ORIG-only + // and EXTRA-only blocks stay on branch-free fast paths + auto run_along_kv_blocks = [&](const OuterloopArgs &cur_args, auto callable) { + uint32_t num_full_orig_blocks = std::min(cur_args.num_kv_blocks, cur_args.num_orig_slots / B_TOPK); + bool has_mixed_block = num_full_orig_blocks < cur_args.num_kv_blocks && cur_args.num_orig_slots % B_TOPK != 0; + CUTE_NO_UNROLL + for (uint32_t kv_block_idx = 0; kv_block_idx < num_full_orig_blocks; ++kv_block_idx) { + callable.template operator()(kv_block_idx); + } + if (has_mixed_block) { + callable.template operator()(num_full_orig_blocks); + } + CUTE_NO_UNROLL + for (uint32_t kv_block_idx = num_full_orig_blocks + has_mixed_block; kv_block_idx < cur_args.num_kv_blocks; ++kv_block_idx) { + callable.template operator()(kv_block_idx); + } + }; + + auto get_first_job = [&]() -> OuterloopArgs { + return _make_outer_loop_args(0, blockIdx.x); + }; + auto get_next_job = [&](const OuterloopArgs &cur_args) -> OuterloopArgs { + smem.bar_clc_full.wait(cur_args.job_idx_mod_2); + ku::CLCResult next_cta0_idx = ku::get_clc_query_response(smem.clc_response_obj); + arrive_on_cta0_barrier(smem.bar_clc_empty); + + if (!next_cta0_idx.is_valid) { + return OuterloopArgs {false}; + } else { + return _make_outer_loop_args( + cur_args.job_idx_mod_2^1, + next_cta0_idx.x + ); + } + }; + + if (warpgroup_idx == 0) { + /* + Q fetching & Epilogue warpgroup + + The timeline of this warpgroup is as follows: + Q0 Q1 O0 Q2 O1 Q3 O2 Q4 O3 ... Qn O(n-1) On + + Where + - Qi means loading the Q of the i-th request, computing the sum of squares of each head on the fly, + performing RoPE transformation, then writing Q to TMEM + - Oi means reading the O of the i-th request from TMEM, performing RoPE transformation, + quantizing to FP8, and writing back to global memory + + About cached_cos and cached_sin: + - D_ROPE is always 64, so the i-th lane caches the i-th cos and sin dimension of the + corresponding position + - We always cache the cos and sin of the i-th request and the (i-1)-th request, so cached_cos + and cached_sin have two slots + - Every time a new Q is loaded, cached_cos/sin[1] (the first slot) is assigned to the zeroth slot + (and "sin" is negated to prepare for the conjugate RoPE of O), + and the cos/sin of the new Q is saved to the 1st slot + - This way, O RoPE only needs to read from the 0th slot (except for the last O) + */ + cutlass::arch::warpgroup_reg_alloc<184>(); + + #pragma nv_diag_suppress 549 // Uninitialized variable. The following two variables are indeed initialized, but the compiler cannot prove it + float cached_cos[2], cached_sin[2]; + auto shift_cached_cos_and_cached_sin = [&]() { + cached_cos[0] = cached_cos[1]; + cached_sin[0] = -cached_sin[1]; + }; + + auto load_q_and_save_to_tmem = [&](const OuterloopArgs &cur_job) { + static constexpr uint32_t NUM_CACHED_BF16_PER_THREAD = H_Q_PER_CTA * D_QK / 128; + static constexpr uint32_t NUM_BF16_PER_LOAD = 256 / 16; // LDG 256 + bf16 cached_q[NUM_CACHED_BF16_PER_THREAD]; + float q_sqr_sum = 0.0f; + + uint32_t cur_q_position = __ldg(params.token_positions + cur_job.s_q_idx); + cached_cos[1] = __ldg(params.cos_sin_cache + cur_q_position * D_ROPE + lane_idx); + cached_sin[1] = __ldg(params.cos_sin_cache + cur_q_position * D_ROPE + D_ROPE / 2 + lane_idx); + + bf16 *q_token_base = params.q + (uint64_t)cur_job.s_q_idx * params.stride_q_s_q; + + static constexpr uint32_t TILE_SIZE = 64; + CUTE_UNROLL + for (uint32_t local_tile_idx = 3; local_tile_idx != 0xFFFFFFFF; --local_tile_idx) { + uint32_t tile_idx = + CLUSTER_SIZE == 1 ? + local_tile_idx * 2 + (warp_idx / (H_Q_PER_CTA / 32)) : // Don't use idx_in_warpgroup / H_Q_PER_CTA to hint the compiler that warps does not diverge here + (warp_idx / (H_Q_PER_CTA / 32)) * 4 + local_tile_idx; + CUTE_UNROLL + for (uint32_t i = 0; i < TILE_SIZE / NUM_BF16_PER_LOAD; ++i) { + static_assert(NUM_MRGEMM_RAILS == 2); + uint32_t h_q_idx = cta_idx * H_Q_PER_CTA + idx_in_warpgroup % H_Q_PER_CTA; + uint32_t d_q_idx = tile_idx * TILE_SIZE + i * NUM_BF16_PER_LOAD; + KU_LDG_256( + q_token_base + h_q_idx * NUM_BF16_PER_LOAD + d_q_idx * H_Q, + cached_q + local_tile_idx * TILE_SIZE + i * NUM_BF16_PER_LOAD, + ".nc", "no_allocate", "evict_first", "256B" + ); + } + // Perform RoPE + if (local_tile_idx == 3 && tile_idx == D_VO / TILE_SIZE - 1) { + float2 cur_q_sqr_sum = {0.0f, 0.0f}; + CUTE_UNROLL + for (uint32_t j = 0; j < TILE_SIZE; j += 2) { + float2 x = __bfloat1622float2(*(nv_bfloat162*)(cached_q+local_tile_idx*TILE_SIZE+j)); + if constexpr (ENABLE_Q_NORM) { + cur_q_sqr_sum = ku::float2_fma(x, x, cur_q_sqr_sum); + } + float cur_cos = __shfl_sync(0xFFFFFFFF, cached_cos[1], j/2); + float cur_sin = __shfl_sync(0xFFFFFFFF, cached_sin[1], j/2); + float2 y = apply_rope(x, cur_cos, cur_sin); + *(nv_bfloat162*)(cached_q+local_tile_idx*TILE_SIZE+j) = nv_bfloat162{__float2bfloat16_rn(y.x), __float2bfloat16_rn(y.y)}; + } + q_sqr_sum += cur_q_sqr_sum.x + cur_q_sqr_sum.y; + } else { + if constexpr (ENABLE_Q_NORM) { + // Accumulate \sum q_i^2 for RMS norm + CUTE_UNROLL + for (uint32_t i = 0; i < TILE_SIZE; ++i) + asm volatile ("fma.rn.f32.bf16 %0, %1, %1, %0;\n" : "+f"(q_sqr_sum) : "h"(*(uint16_t*)(cached_q+local_tile_idx*TILE_SIZE+i))); + } + } + } + + if constexpr (ENABLE_Q_NORM) { + smem.q_sqr_sum_buf[cur_job.job_idx_mod_2][idx_in_warpgroup] = q_sqr_sum; + } + + smem.bar_tQ_empty.wait(cur_job.job_idx_mod_2^1); + ku::tcgen05_after_thread_sync(); + + static constexpr uint32_t NUM_CACHED_UINT32 = NUM_CACHED_BF16_PER_THREAD/2; + ku::tmem_st_32dp32bNx(tmem_cols::Q, cached_q); + ku::tmem_st_32dp32bNx(tmem_cols::Q+NUM_CACHED_UINT32/2, cached_q+NUM_CACHED_UINT32/2*2); // We split the tmem_st into two parts, otherwise NVCC complains about insufficient registers. I suspect this is because PTXAS ignores the warpgroup_reg_alloc<168> above and uses 128 as the available register count per thread (with 512 total threads, each thread initially has only 128 registers) + cutlass::arch::fence_view_async_tmem_store(); + + ku::tcgen05_before_thread_sync(); + arrive_on_cta0_barrier(smem.bar_tQ_full); + if constexpr (IS_2CTA && ENABLE_Q_NORM) { + smem.bar_q_sqr_sum_full.arrive(); + } + }; + auto store_o = [&](const OuterloopArgs &cur_job, const bool &is_last_job) { + smem.bar_li_mi_full.wait(cur_job.job_idx_mod_2); + float li = 0.0f; + float mi = smem.rowwise_mi_buf[idx_in_warpgroup % H_Q_PER_CTA]; + if constexpr (FOLD_FACTOR == 2) { + li = smem.rowwise_li_buf[idx_in_warpgroup] + smem.rowwise_li_buf[idx_in_warpgroup^64]; + } else { + static_assert(FOLD_FACTOR == 4); + li = __fadd_rn( + __fadd_rn(smem.rowwise_li_buf[idx_in_warpgroup], smem.rowwise_li_buf[idx_in_warpgroup^64]), + __fadd_rn(smem.rowwise_li_buf[idx_in_warpgroup^32], smem.rowwise_li_buf[idx_in_warpgroup^96]) + ); + } + smem.bar_li_mi_empty.arrive(); + + if (idx_in_warpgroup < H_Q_PER_CTA) { + uint32_t global_index = cur_job.s_q_idx * H_Q + cta_idx * H_Q_PER_CTA + idx_in_warpgroup; + float cur_lse = fmaf(mi, CUDART_LN2_F, logf(li)); + cur_lse = cur_lse == -CUDART_INF_F ? +CUDART_INF_F : cur_lse; + params.lse[global_index] = cur_lse; + } + + float attn_sink = params.attn_sink == nullptr ? -CUDART_INF_F : __ldg(params.attn_sink + cta_idx * H_Q_PER_CTA + idx_in_warpgroup % H_Q_PER_CTA) * CUDART_L2E_F; + float output_scale = li == 0.0f ? 0.0f : __fdividef(1.0f, li + exp2f(attn_sink - mi)); + + smem.bar_tO_full.wait(cur_job.job_idx_mod_2); + ku::tcgen05_after_thread_sync(); + + static constexpr uint32_t MMA_ATOM_N = 256; + static constexpr uint32_t NUM_MMA_ATOMS = D_VO / MMA_ATOM_N; + static constexpr uint32_t NUM_O_TMEM_COLS_PER_ATOM = MMA_ATOM_N / FOLD_FACTOR; + static constexpr uint32_t EPILOGUE_TILE_SIZE = O_QUANT_TILE_SIZE; + static constexpr uint32_t NUM_EPILOGUE_TILES_PER_ATOM = NUM_O_TMEM_COLS_PER_ATOM / EPILOGUE_TILE_SIZE; + static_assert(NUM_O_TMEM_COLS_PER_ATOM % EPILOGUE_TILE_SIZE == 0); // TODO When FOLD_FACTOR is 4 and MODEL_TYPE is V4 (so NUM_O_TMEM_COLS_PER_ATOM is 128), this isn't hold + fp8_e4m3 output_fp8[NUM_MMA_ATOMS][NUM_O_TMEM_COLS_PER_ATOM]; + uint8_t output_sf[NUM_MMA_ATOMS][NUM_O_TMEM_COLS_PER_ATOM / O_QUANT_TILE_SIZE]; + + CUTE_UNROLL + for (uint32_t mma_atom_idx = 0; mma_atom_idx < NUM_MMA_ATOMS; mma_atom_idx += 1) { + CUTE_UNROLL + for (uint32_t epilogue_tile_idx_in_atom = 0; epilogue_tile_idx_in_atom < NUM_EPILOGUE_TILES_PER_ATOM; ++epilogue_tile_idx_in_atom) { + // Fetch output from TMEM + uint32_t tmem_col_base = tmem_cols::O + mma_atom_idx * NUM_O_TMEM_COLS_PER_ATOM + epilogue_tile_idx_in_atom * EPILOGUE_TILE_SIZE; + float output[EPILOGUE_TILE_SIZE]; + float reduce_result_by_tmem_ld; + if constexpr (IS_TMEM_LD_WITH_RED_AVAILABLE) { + ku::tmem_ld_red_32dp32bNx(tmem_col_base, output, reduce_result_by_tmem_ld); + } else { + ku::tmem_ld_32dp32bNx(tmem_col_base, output); + } + cutlass::arch::fence_view_async_tmem_load(); + + // Notify tO's emptyness + if (mma_atom_idx+1 == NUM_MMA_ATOMS && epilogue_tile_idx_in_atom+1 == NUM_EPILOGUE_TILES_PER_ATOM) { + ku::tcgen05_before_thread_sync(); + if (!is_last_job) { + // Don't arrive on the barrier if this job is the last job, to avoid "cluster target block not present" + arrive_on_cta0_barrier(smem.bar_tO_empty); + } + } + + // RoPE (conjugate) + bool should_perform_rope; + { + static_assert(FOLD_FACTOR == 2); + static_assert(D_ROPE % EPILOGUE_TILE_SIZE == 0); + should_perform_rope = + mma_atom_idx + 1 == NUM_MMA_ATOMS && + epilogue_tile_idx_in_atom >= NUM_EPILOGUE_TILES_PER_ATOM - D_ROPE/EPILOGUE_TILE_SIZE && + warp_idx >= 2; + if (should_perform_rope) { + CUTE_UNROLL + for (uint32_t j = 0; j < EPILOGUE_TILE_SIZE; j += 2) { + float2 x = *(float2*)(output + j); + uint32_t src_lane = j/2 + (epilogue_tile_idx_in_atom + 1 == NUM_EPILOGUE_TILES_PER_ATOM ? EPILOGUE_TILE_SIZE / 2 : 0); + float cur_cos = __shfl_sync(0xFFFFFFFF, cached_cos[0], src_lane); + float cur_sin = __shfl_sync(0xFFFFFFFF, cached_sin[0], src_lane); + float2 y = apply_rope(x, cur_cos, cur_sin); + *(float2*)(output + j) = y; + } + } + } + + // Cast to FP8, and save to global memory + float output_abs_max; + if (!IS_TMEM_LD_WITH_RED_AVAILABLE || should_perform_rope) { + output_abs_max = get_max(output) * output_scale; + } else { + output_abs_max = reduce_result_by_tmem_ld * output_scale; + } + output_abs_max = max(O_QUANT_CLAMP_MIN_VALUE, output_abs_max); + float sf = output_abs_max / 448.0f; + uint32_t sf_as_uint32 = *reinterpret_cast(&sf); + uint32_t exp_sf = (int32_t)((sf_as_uint32-1) >> 23) + (1 - 127); + uint32_t sf_inv_as_uint32 = (127 - exp_sf) << 23; + float sf_inv = *reinterpret_cast(&sf_inv_as_uint32); + float cur_multiplier = output_scale * sf_inv; + float2 cur_multiplier_float2 = float2(cur_multiplier, cur_multiplier); + + CUTE_UNROLL + for (uint32_t j = 0; j < EPILOGUE_TILE_SIZE; j += 2) { + float2 x = *(float2*)(output + j); + x = ku::float2_mul(x, cur_multiplier_float2); + *(__nv_fp8x2_storage_t*)(output_fp8[mma_atom_idx] + epilogue_tile_idx_in_atom * EPILOGUE_TILE_SIZE + j) = __nv_cvt_float2_to_fp8x2( + x, + __NV_SATFINITE, + __nv_fp8_interpretation_t::__NV_E4M3 + ); // NOTE. Here we don't use cvt.f8x4type.f32 since it only has .rs mode, which affects accuracy + } + output_sf[mma_atom_idx][epilogue_tile_idx_in_atom] = exp_sf + 127; + } + } + + uint32_t head_idx = cta_idx*H_Q_PER_CTA + idx_in_warpgroup%H_Q_PER_CTA; + uint32_t wv_group_idx = head_idx / WV_GROUP_SIZE; + uint32_t head_idx_in_wv_group = head_idx % WV_GROUP_SIZE; + CUTE_UNROLL + for (uint32_t mma_atom_idx = 0; mma_atom_idx < NUM_MMA_ATOMS; mma_atom_idx += 1) { + // Store SF + // 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 (numerically) per-128 scaled when num_per_channels is 128 + static constexpr uint32_t OUTPUT_SAVE_AS_SCALE_GRAN = 32; + // Layout of O in Tensor Memory: + // For HEAD_DIM_QK = 64 (head64): + // - Atom 0 computes O[:, 0:256]; Atom 1 computes O[:, 256:512] + // - Mapping to TMEM: + // O[0:128] -> TMEM[0:64, 0:128] + // O[128:256] -> TMEM[64:128, 0:128] + // O[256:384] -> TMEM[0:64, 128:256] + // O[384:512] -> TMEM[64:128, 128:256] + // - Visually (label the four 128-wide O chunks as 0..3): + // +---+---+ + // | 0 | 2 | + // +---+---+ + // | 1 | 3 | + // +---+---+ + // For HEAD_DIM_QK = 128 (head128): + // - CTA0 holds V[:, 0:256]; CTA1 holds V[:, 256:512] + // - Atom 0 computes O[:, 0:128] and O[:, 256:384] + // Atom 1 computes O[:, 128:256] and O[:, 384:512] + // - Mapping to TMEM: + // O[0:128] -> TMEM[0:64, 0:128] + // O[128:256] -> TMEM[0:64, 128:256] + // O[256:384] -> TMEM[64:128, 0:128] + // O[384:512] -> TMEM[64:128, 128:256] + // - Visually (label the four 128-wide O chunks as 0..3): + // +---+---+ + // | 0 | 1 | + // +---+---+ + // | 2 | 3 | + // +---+---+ + uint32_t head_dim_idx_base = + CLUSTER_SIZE == 1 ? + mma_atom_idx * MMA_ATOM_N + (warp_idx/(H_Q_PER_CTA/32)) * (MMA_ATOM_N/FOLD_FACTOR) : + mma_atom_idx * NUM_O_TMEM_COLS_PER_ATOM + (warp_idx/(H_Q_PER_CTA/32)) * MMA_ATOM_N; + CUTE_UNROLL + for (uint32_t i = 0; i < NUM_O_TMEM_COLS_PER_ATOM / OUTPUT_SAVE_AS_SCALE_GRAN; ++i) { + uint32_t sf_block_idx = head_idx_in_wv_group + (head_dim_idx_base/OUTPUT_SAVE_AS_SCALE_GRAN+i) * WV_GROUP_SIZE; + *((uint8_t*)(params.out_sf + cur_job.s_q_idx + wv_group_idx*params.stride_out_sf_wv_group + (sf_block_idx/4)*params.stride_out_sf_head_dim) + sf_block_idx%4) = output_sf[mma_atom_idx][i]; // TODO Optimize + } + // Store output + CUTE_UNROLL + for (uint32_t i = 0; i < NUM_O_TMEM_COLS_PER_ATOM; i += 32) { + uint32_t head_dim_idx = head_dim_idx_base + i; + KU_STG_256( + params.out_fp8 + (uint64_t)cur_job.s_q_idx*(H_Q*D_VO) + wv_group_idx*(WV_GROUP_SIZE*D_VO) + head_idx_in_wv_group*32 + head_dim_idx*WV_GROUP_SIZE, + output_fp8[mma_atom_idx] + i, + "no_allocate", + "evict_first" + ); + } + } + }; + + OuterloopArgs cur_job = get_first_job(); + load_q_and_save_to_tmem(cur_job); + do { + OuterloopArgs next_job = get_next_job(cur_job); + shift_cached_cos_and_cached_sin(); + if (next_job.is_valid) { + load_q_and_save_to_tmem(next_job); + } + store_o(cur_job, !next_job.is_valid); + cur_job = next_job; + } while (cur_job.is_valid); + + NamedBarrier::arrive_and_wait(128, barrier_ids::WG0_SYNC); + if (warp_idx == 0) { + AllocatorT().free(0, 512); + } + } else if (warpgroup_idx == 3) { + // Scale & Exp warpgroup + cutlass::arch::warpgroup_reg_alloc<128>(); + + OuterloopArgs cur_job = get_first_job(); + uint32_t local_warp_idx = warp_idx - 12; + static_assert(FOLD_FACTOR == 2); + bf16* sS_base = smem.s + (local_warp_idx >= 2 ? H_Q_PER_CTA * (B_TOPK/2) : 0) + (idx_in_warpgroup%H_Q_PER_CTA) * 8; + RingBufferState rs; + do { + // For definition and consistency about `mi`, `li`, and `real_mi`, plz refer to head64 prefill + static constexpr uint32_t NUM_ELEMS_PER_THREAD = B_TOPK * H_Q_PER_CTA / 128; + float mi = MAX_INIT_VAL; + float li = 0.0f; + float real_mi = -CUDART_INF_F; + + float score_multiplier; // qk_scale * rms_norm's denominator (if q norm is enabled) + if constexpr (ENABLE_Q_NORM) { + if constexpr (IS_2CTA) { + smem.bar_q_sqr_sum_full.wait(cur_job.job_idx_mod_2); + } else { + smem.bar_tQ_full.wait(cur_job.job_idx_mod_2); + } + + if constexpr (H_Q_PER_CTA == 64) { + score_multiplier = smem.q_sqr_sum_buf[cur_job.job_idx_mod_2][idx_in_warpgroup] + smem.q_sqr_sum_buf[cur_job.job_idx_mod_2][idx_in_warpgroup^64]; + } else { + static_assert(H_Q_PER_CTA == 32); + score_multiplier = __fadd_rn( + __fadd_rn(smem.q_sqr_sum_buf[cur_job.job_idx_mod_2][idx_in_warpgroup], smem.q_sqr_sum_buf[cur_job.job_idx_mod_2][idx_in_warpgroup^64]), + __fadd_rn(smem.q_sqr_sum_buf[cur_job.job_idx_mod_2][idx_in_warpgroup^32], smem.q_sqr_sum_buf[cur_job.job_idx_mod_2][idx_in_warpgroup^96]) + ); + } + score_multiplier = params.sm_scale_div_log2 * rsqrtf(score_multiplier / D_QK + params.rms_norm_eps); // rsqrt is translated to `MUFU.RSQ` + } else { + score_multiplier = params.sm_scale_div_log2; + } + + smem.bar_tQ_empty.arrive(); // Must arrive on the empty barrier here, to prevent smem.bar_tQ_full being phase-skipped + + CUTE_NO_UNROLL + for (uint32_t kv_block_idx = 0; kv_block_idx < cur_job.num_kv_blocks; ++kv_block_idx) { + auto [indices_buf_idx, indices_bar_phase] = rs.get(); + auto [p_buf_idx, p_bar_phase] = rs.get(); + smem.bar_tP_full[p_buf_idx].wait(p_bar_phase); + smem.bar_indices_full[indices_buf_idx].wait(indices_bar_phase); + ku::tcgen05_after_thread_sync(); + + float p[NUM_ELEMS_PER_THREAD]; + retrieve_mask_and_reduce_p< + NUM_ELEMS_PER_THREAD, + barrier_ids::WG3_WARP02_SYNC, + barrier_ids::WG3_WARP13_SYNC, + false + >( + tmem_cols::get_p(p_buf_idx), + (char*)&smem.is_k_valid[indices_buf_idx], + local_warp_idx, + lane_idx, + [&]() { + if constexpr (NEED_TP_EMPTY_BAR) { + arrive_on_cta0_barrier(smem.bar_tP_empty[p_buf_idx]); + } + }, + smem.p_exchange_buf, + p + ); + + float cur_pi_max = get_max(p); + cur_pi_max *= score_multiplier; + + smem.rowwise_max_buf[idx_in_warpgroup] = cur_pi_max; + NamedBarrier::arrive_and_wait(64, barrier_ids::WG3_WARP02_SYNC + (local_warp_idx&1)); + smem.bar_indices_empty[indices_buf_idx].arrive(); // Put it here to give the compiler more room for SASS code reordering + cur_pi_max = max(cur_pi_max, smem.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); + + 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_bfloat16 s[NUM_ELEMS_PER_THREAD]; + float cur_sum = get_s_from_p((nv_bfloat162*)s, p, score_multiplier, new_max); + li = fmaf(li, scale_for_old, cur_sum); + + // Store S + smem.bar_SO_empty.wait(rs.get<1>().second^1); + CUTE_UNROLL + for (int i = 0; i < NUM_ELEMS_PER_THREAD/8; ++i) { + ku::st_shared(sS_base + i*8*H_Q_PER_CTA, *(__int128_t*)(s + i*8)); + } + + // Rescale O + if (kv_block_idx > 0 && should_scale_o) { + ku::tcgen05_after_thread_sync(); + rescale_O(scale_for_old); + ku::tcgen05_before_thread_sync(); + } + + fence_view_async_shared(); + ku::tcgen05_before_thread_sync(); + arrive_on_cta0_barrier(smem.bar_SO_full); + rs.update(); + } + + if (real_mi == -CUDART_INF_F) { + // No valid TopK indices + li = 0.0f; + mi = -CUDART_INF_F; + } + + smem.bar_li_mi_empty.wait(cur_job.job_idx_mod_2^1); + static_assert(H_Q_PER_CTA % 32 == 0); + if (local_warp_idx < H_Q_PER_CTA / 32) { + if constexpr (!IS_DECODE) { + uint32_t global_index = cur_job.s_q_idx * H_Q + cta_idx * H_Q_PER_CTA + idx_in_warpgroup; + params.max_logits[global_index] = real_mi * CUDART_LN2_F; + } + smem.rowwise_mi_buf[idx_in_warpgroup] = mi; + } + smem.rowwise_li_buf[idx_in_warpgroup] = li; + smem.bar_li_mi_full.arrive(); + + cur_job = get_next_job(cur_job); + } while (cur_job.is_valid); + } else if (warpgroup_idx == 2) { + cutlass::arch::warpgroup_reg_dealloc<72>(); + if (warp_idx == 8 && cta_idx == 0 && elect_one_sync()) { + // MMA warp (CTA0 only) + auto tiled_mma_qk = TiledMMA_QK{}; + auto tiled_mma_sv = TiledMMA_SV{}; + Tensor tQ = tiled_mma_qk.get_slice(_0{}).make_fragment_A( + partition_shape_A(tiled_mma_qk, Shape, Int>{}) + ); + Tensor tP = partition_fragment_C(tiled_mma_qk, Shape, Int>{}); + Tensor sS = make_tensor( + make_smem_ptr(smem.s), + ku::make_umma_canonical_k_major_layout() + ); + Tensor tO = partition_fragment_C(tiled_mma_sv, Shape, Int>{}); + tQ.data().get() = tmem_cols::Q; + tO.data().get() = tmem_cols::O; + // tP.data() will be assigned in the loop since it has double buffers + + RingBufferState rs_qk, rs_sv; + auto run_qk_gemm = [&](const OuterloopArgs &job, uint32_t kv_block_idx) { + if (kv_block_idx == 0) { + smem.bar_tQ_full.wait(job.job_idx_mod_2); + } + auto [kv_slot_idx, kv_bar_phase] = rs_qk.get(); + Tensor sK = make_tensor( + make_smem_ptr(smem.kv_slots[kv_slot_idx]), + ku::make_umma_canonical_k_major_layout() + ); + // Expected TMA transaction bytes on bar_kv_slot_full: + // Prefill: the whole KV block; Decode: only the bf16 (RoPE) part (the fp8 part is dequantized by WG1, which arrives on the same barrier) + if constexpr (IS_DECODE && D_BF16 == 0) { + // No bf16 part, so no TMA transaction is expected (expect-tx count must not be 0) + smem.bar_kv_slot_full[kv_slot_idx].arrive(); + } else { + smem.bar_kv_slot_full[kv_slot_idx].arrive_and_expect_tx(IS_DECODE ? B_TOPK*D_BF16*sizeof(bf16) : B_TOPK*D_QK*sizeof(bf16)); + } + smem.bar_kv_slot_full[kv_slot_idx].wait(kv_bar_phase); + auto [p_buf_idx, p_bar_phase] = rs_qk.get(); + if constexpr (NEED_TP_EMPTY_BAR) { + smem.bar_tP_empty[p_buf_idx].wait(p_bar_phase ^ 1); + } + tP.data().get() = tmem_cols::get_p(p_buf_idx); + + ku::tcgen05_after_thread_sync(); + ku::utcmma_ts(tiled_mma_qk, tQ, sK, tP, true); + umma_arrive_on_every_cta(smem.bar_tP_full[p_buf_idx]); + + if (kv_block_idx == job.num_kv_blocks-1) { + umma_arrive_on_every_cta(smem.bar_tQ_empty); + } + rs_qk.update(); + }; + + auto run_sv_gemm = [&](const OuterloopArgs &job, uint32_t kv_block_idx) { + if (kv_block_idx == 0) { + smem.bar_tO_empty.wait(job.job_idx_mod_2^1); + } + auto [kv_slot_idx, _] = rs_sv.get(); + smem.bar_SO_full.wait(rs_sv.get<1>().second); + Tensor sV = make_tensor( + make_smem_ptr(smem.kv_slots[kv_slot_idx]), + ku::make_umma_canonical_mn_major_layout() + ); + ku::tcgen05_after_thread_sync(); + ku::utcmma_ss(tiled_mma_sv, sS, sV, tO, kv_block_idx == 0); + umma_arrive_on_every_cta(smem.bar_kv_slot_empty[kv_slot_idx]); + umma_arrive_on_every_cta(smem.bar_SO_empty); + if (kv_block_idx == job.num_kv_blocks-1) { + umma_arrive_on_every_cta(smem.bar_tO_full); + } + rs_sv.update(); + }; + + OuterloopArgs cur_job = get_first_job(); + run_qk_gemm(cur_job, 0); + do { + CUTE_NO_UNROLL + for (uint32_t kv_block_idx = 1; kv_block_idx < cur_job.num_kv_blocks; ++kv_block_idx) { + run_qk_gemm(cur_job, kv_block_idx); + run_sv_gemm(cur_job, kv_block_idx-1); + } + + OuterloopArgs next_job = get_next_job(cur_job); + if (next_job.is_valid) { + run_qk_gemm(next_job, 0); + } + run_sv_gemm(cur_job, cur_job.num_kv_blocks-1); + + cur_job = next_job; + } while (cur_job.is_valid); + } else if (warp_idx == 9 && elect_one_sync()) { + // CLC warp + bool phase = 0; + while (true) { + if (cta_idx == 0) { + smem.bar_clc_empty.wait(phase^1); + ku::issue_clc_query_multicast_cluster_all(smem.bar_clc_full, smem.clc_response_obj); + } + smem.bar_clc_full.arrive_and_expect_tx(sizeof(smem.clc_response_obj)); + + smem.bar_clc_full.wait(phase&1); + ku::CLCResult clc_result = ku::get_clc_query_response(smem.clc_response_obj); + arrive_on_cta0_barrier(smem.bar_clc_empty); + if (!clc_result.is_valid) + break; + + phase ^= 1; + } + if constexpr (IS_2CTA) { + if (cta_idx == 0) { + smem.bar_clc_empty.wait(phase); // Wait for all threads' arrival on `bar_clc_empty`, which means that there will be no further operations on distributed shared memory (including barrier arrive and MMA), avoiding the "cluster target block not present" error + smem.bar_clc_empty.arrive(1u); // Transfer the signal above to CTA1 + } else { + smem.bar_clc_empty.wait(0); + } + } + } else if (warp_idx == 10) { + // Indices generator + // (also generates TMA coords & scales for dequant warps) + OuterloopArgs cur_job = get_first_job(); + RingBufferState rs; + static constexpr uint32_t NUM_INDICES_PER_THREAD = B_TOPK / 32; + static_assert(B_TOPK % 32 == 0); + + do { + if constexpr (!IS_DECODE) { + auto body = [&]() { + CUTE_NO_UNROLL + for (uint32_t kv_block_idx = 0; kv_block_idx < cur_job.num_kv_blocks; ++kv_block_idx) { + auto [indices_buf_idx, indices_bar_phase] = rs.get(); + smem.bar_indices_empty[indices_buf_idx].wait(indices_bar_phase^1); + + CUTE_UNROLL + for (uint32_t i = 0; i < NUM_INDICES_PER_THREAD; ++i) { + uint32_t pos = kv_block_idx * B_TOPK + i * 32 + lane_idx; + int cur_index; + if constexpr (CHECK_TOPK_SUBSCRIPT) { + // Predicate the load on `pos < topk_length` to prevent IMA + cur_index = pos < cur_job.topk_length ? __ldg(params.indices + cur_job.s_q_idx * params.stride_indices_s_q + pos) : -1; + } else { + // topk_length % B_TOPK == 0, so every pos is within the row + cur_index = __ldg(params.indices + cur_job.s_q_idx * params.stride_indices_s_q + pos); + } + bool is_index_valid = (uint32_t)cur_index < (uint32_t)params.s_kv; // Don't need to check `index >= 0`, since if `index < 0` holds, `(uint32_t)index` must lies in 2147483648 ~ 4294967295, which is definitely greater than `params.s_kv` + uint32_t mask = __ballot_sync(0xFFFFFFFF, is_index_valid); + if (lane_idx == 0) { + *((uint32_t*)smem.is_k_valid[indices_buf_idx] + i) = mask; + } + } + + smem.bar_indices_full[indices_buf_idx].arrive(); + rs.update(); + } + }; + if (cur_job.topk_length % B_TOPK == 0 && cur_job.topk_length != 0) + body.template operator()(); + else + body.template operator()(); + } else { + int *indices_base = params.indices + (int64_t)cur_job.s_q_idx * params.stride_indices_s_q; + int *extra_indices_base = params.extra_indices + (int64_t)cur_job.s_q_idx * params.stride_extra_indices_s_q; + run_along_kv_blocks(cur_job, [&](uint32_t kv_block_idx) { + auto [indices_buf_idx, indices_bar_phase] = rs.get(); + smem.bar_indices_empty[indices_buf_idx].wait(indices_bar_phase^1); + + CUTE_UNROLL + for (uint32_t i = 0; i < NUM_INDICES_PER_THREAD; ++i) { + uint32_t pos = kv_block_idx * B_TOPK + i * 32 + lane_idx; + bool in_extra; + if constexpr (LOC == KVLocation::ORIG) { + in_extra = false; + } else if constexpr (LOC == KVLocation::EXTRA) { + in_extra = true; + } else { + in_extra = pos >= cur_job.num_orig_slots; + } + uint32_t local_pos = in_extra ? pos - cur_job.num_orig_slots : pos; + uint32_t valid_len = in_extra ? cur_job.extra_topk_length : cur_job.topk_length; + // Predicate the load on `local_pos < valid_len`: the last KV block may run + // beyond the end of the indices row (e.g. topk=128 with B_TOPK=96), so an + // unconditional load could read OOB of `indices` / `extra_indices`. Slots + // beyond `valid_len` are masked out anyway, so skip the load for them + int cur_index = local_pos < valid_len ? __ldg((in_extra ? extra_indices_base : indices_base) + local_pos) : -1; + bool is_index_valid = cur_index >= 0; + + // Share coord/scale generation instead of repeating it in all dequant warps. + int64_t src_block_stride = in_extra ? params.stride_extra_kv_block : params.stride_kv_block; + const auto &fast_divmod = in_extra + ? aux_params.fast_divmod_extra_page_block_size + : aux_params.fast_divmod_page_block_size; + uint32_t token_idx = is_index_valid ? (uint32_t)cur_index : 0; + int idx_in_block; + int block_idx = fast_divmod.divmod(idx_in_block, (int)token_idx); + uint32_t row = i * 32 + lane_idx; + // The extra KV cache may have another format (HAS_FP4_KV); ORIG / EXTRA blocks resolve it at compile time + uint32_t tma_k_stride = in_extra ? ExtraKVFormat::TMA_K_STRIDE : OrigKVFormat::TMA_K_STRIDE; + uint32_t num_scales_each_token = in_extra ? ExtraKVFormat::NUM_SCALES_EACH_TOKEN : OrigKVFormat::NUM_SCALES_EACH_TOKEN; + smem.decode_tma_coords[indices_buf_idx][row] = is_index_valid + ? (src_block_stride / tma_k_stride) * block_idx + idx_in_block + : -1; + + uint32_t page_block_size = in_extra ? params.extra_page_block_size : params.page_block_size; + uint8_t *kv_base = (uint8_t*)(in_extra ? params.extra_kv : params.kv); + uint8_t *scale_ptr = kv_base + page_block_size * tma_k_stride + + (int64_t)block_idx * src_block_stride + + idx_in_block * num_scales_each_token + + cta_idx * (num_scales_each_token / 2); + bool is_fp4_token = HAS_FP4_KV && in_extra; // The dequant warps load the scales of fp4 tokens themselves + if (!is_fp4_token) { + uint8_t *scale_dst = smem.decode_scales[indices_buf_idx] + + row * NUM_SCALES_EACH_TOKEN_PER_CTA; + if constexpr (NUM_SCALES_EACH_TOKEN_PER_CTA == 4) { + uint32_t scales; + asm volatile ( + "ld.global.nc.L1::no_allocate.b32 %0, [%1];" + : "=r"(scales) + : "l"((uint64_t)scale_ptr) + ); + *(uint32_t*)scale_dst = is_index_valid ? scales : 0; + } else if constexpr (NUM_SCALES_EACH_TOKEN_PER_CTA == 8) { + uint64_t scales; + asm volatile ( + "ld.global.nc.L1::no_allocate.b64 %0, [%1];" + : "=l"(scales) + : "l"((uint64_t)scale_ptr) + ); + *(uint64_t*)scale_dst = is_index_valid ? scales : 0; + } else { + static_assert(NUM_SCALES_EACH_TOKEN_PER_CTA == 16); + __int128_t scales; + asm volatile ( + "ld.global.nc.L1::no_allocate.b128 %0, [%1];" + : "=q"(scales) + : "l"((uint64_t)scale_ptr) + ); + *(__int128_t*)scale_dst = is_index_valid ? scales : 0; + } + } else { + smem.decode_scale_ptrs[indices_buf_idx][row] = is_index_valid ? scale_ptr : nullptr; + } + uint32_t mask = __ballot_sync(0xFFFFFFFF, is_index_valid); + if (lane_idx == 0) { + *((uint32_t*)smem.is_k_valid[indices_buf_idx] + i) = mask; + } + } + + smem.bar_indices_full[indices_buf_idx].arrive(); + rs.update(); + }); + } + cur_job = get_next_job(cur_job); + } while (cur_job.is_valid); + } + } else if (warpgroup_idx == 1) { + cutlass::arch::warpgroup_reg_alloc<128>(); + if constexpr (IS_DECODE) { + if constexpr (HAS_FP4_KV) { + // KV producer for an fp8 `kv` plus an fp4 `extra_kv`: gathers the quantized rows of every selected KV token + // into the beginning of the KV slot via TMA gather4, dequantizes them in registers, and stores the bf16 result + // into the KV slot inplace (SW128 K-major layout). fp8 tokens (from `kv`) and fp4 + // tokens (from `extra_kv`) share the code below and differ only in the tensor map, the number of elements per + // 16 B of raw data (CHUNK_ELEMS), loading of scales, and the conversion instructions. A KV block + // straddling the orig/extra boundary (KVLocation::ORIG_AND_EXTRA) mixes both kinds of tokens; the format + // is resolved per 8 rows, so `run()` asserts topk % 8 == 0. + // + // The unit of work is a chunk: 16 B of raw data (one LDS.128) = CHUNK_ELEMS elements = CHUNK_ELEMS * 2 B of the bf16 + // 128 B swizzle-atom row of the KV slot (CHUNK_ELEMS / 8 STS.128). Lane-to-token mapping: 4 lanes per row, each + // owning a quarter of the row's chunks, and 8 consecutive rows (one "token" of the thread) per 8 consecutive lanes, + // NUM_TOKENS_PER_THREAD tokens per thread. Unlike the fp8-only path, the 4 lanes of one row are spread over the 4 + // wavefronts of an LDS/STS.128 (8 consecutive lanes each), so that one wavefront covers 8 consecutive rows working + // on the same chunk position (with the fp8-only mapping, two of the 4 lanes of a row would own fp4 chunks of the + // same parity and their STS.128 would conflict). With C = chunks per lane, and rows relative to the token: + // + // lane 0 1 2 3 | 4 5 6 7 | 8 .. 11 | 12 .. 15 | 16 .. 19 | 20 .. 23 | 24 .. 27 | 28 .. 31 + // row 0 1 2 3 | 4 5 6 7 | 0 .. 3 | 4 .. 7 | 0 .. 3 | 4 .. 7 | 0 .. 3 | 4 .. 7 + // chunks [0, C) | rotated by 4| [C, 2C) | rotated | [2C, 3C) | rotated | [3C, 4C) | rotated + // + // "rotated by 4": the lanes of rows 4..7 process the same chunks as the lanes of rows 0..3 but in an order shifted + // by 4 chunks (= 64 B), see get_chunk_base. This makes every wavefront bank-conflict-free: + // - Raw LDS.128: the rows of a 4-row gather4 group are RAW_TOKEN_SMEM_STRIDE bytes apart with RAW_TOKEN_SMEM_STRIDE / 16 odd, + // so they hit 4 consecutive 16 B bank groups; the groups themselves start 128 B aligned, so the second group + // of the wavefront reads chunks rotated by 4 to hit the other 4 bank groups. + // - STS.128 into the SW128 K-major layout: a 16 B part of chunk c lands in the 16 B bank group + // ((c % CHUNKS_PER_ATOM_ROW) * NUM_STS_PER_CHUNK + j) ^ (row % 8) of its swizzle-atom row. The 8 rows of a + // wavefront have 8 distinct row % 8 and, the rotation being a multiple of CHUNKS_PER_ATOM_ROW, the same + // c % CHUNKS_PER_ATOM_ROW. + // + // Dequant pipeline per KV block: + // 1. Read the scales of the fp8 tokens from smem, issue gather4 for the raw rows (completing on bar_raw_kv_full), + // then load the scales of the fp4 tokens from global memory (their latency overlaps with the gather). + // 2. Wait for TMA, then read the raw data via LDS.128. Synchronize afterward to prevent the subsequent + // write-back from overwriting the read data. + // 3. Per chunk: fp8: 8x F2FP (e4m3x2 x ue8m0 -> bf16x2); fp4: 16x F2FP (e2m1x2 -> bf16x2) + 16x HMUL2.BF16 + // (x the bf16 of the e4m3 scale, exact: the product has at most 2 + 4 significant bits). Then the STS.128s. + static_assert(!OrigKVFormat::IS_FP4 && ExtraKVFormat::IS_FP4 && ExtraKVFormat::D_BF16 == 0); + uint32_t local_warp_idx = warp_idx - 4; + static constexpr uint32_t NUM_DEQUANT_WARPS = 4, NUM_LANES_PER_ROW = 4; + static constexpr uint32_t NUM_ROWS_PER_WAVEFRONT = 32 / NUM_LANES_PER_ROW; // Rows covered by 8 consecutive lanes (one LDS/STS.128 wavefront) + static constexpr uint32_t NUM_TOKENS_PER_THREAD = B_TOPK / (NUM_DEQUANT_WARPS * NUM_ROWS_PER_WAVEFRONT); + static_assert(B_TOPK % (NUM_DEQUANT_WARPS * NUM_ROWS_PER_WAVEFRONT) == 0); + static constexpr uint32_t NUM_ROWS_PER_WARP = NUM_TOKENS_PER_THREAD * NUM_ROWS_PER_WAVEFRONT; + static constexpr uint32_t MAX_CHUNKS_PER_LANE = OrigKVFormat::NUM_CHUNKS_PER_ROW / NUM_LANES_PER_ROW; // fp8 rows have the most: 8 / 4 + const uint32_t row_in_wavefront = lane_idx % NUM_ROWS_PER_WAVEFRONT, idx_in_row = lane_idx / NUM_ROWS_PER_WAVEFRONT; + const uint32_t chunk_rotation = row_in_wavefront / 4 * 4; + auto get_row_idx = [&](uint32_t token_idx) { + return local_warp_idx * NUM_ROWS_PER_WARP + token_idx * NUM_ROWS_PER_WAVEFRONT + row_in_wavefront; + }; + // The chunks this lane processes are numbered g = 0.. within the lane; chunk g of the row is chunk_base(g / 4) + g % 4, + // so that all addresses are a per-lane base plus an immediate. The lanes of rows 4..7 rotate the order by 4 chunks: + // as a rotation of the chunk indices of the row when a lane owns at most 4 chunks, as a swap of its two halves + // (g ^ 4) when it owns 8 + auto get_chunk_base = [&](uint32_t half) { + constexpr uint32_t NUM_CHUNKS_PER_LANE = F::NUM_CHUNKS_PER_ROW / NUM_LANES_PER_ROW; + static_assert(NUM_CHUNKS_PER_LANE == 2 || NUM_CHUNKS_PER_LANE == 4 || NUM_CHUNKS_PER_LANE == 8); + if constexpr (NUM_CHUNKS_PER_LANE <= 4) { + return (idx_in_row * NUM_CHUNKS_PER_LANE + chunk_rotation) % F::NUM_CHUNKS_PER_ROW; + } else { + return idx_in_row * NUM_CHUNKS_PER_LANE + (half * 4 ^ chunk_rotation); + } + }; + auto get_raw_row_offset = [&](uint32_t row) { + return row / 4 * RAW_KV_GROUP_BYTES; // + (row % 4) * F::RAW_TOKEN_SMEM_STRIDE, which depends on the format + }; + + // STS.128 offsets of the 8 (swizzled) 16 B parts of the first swizzle-atom row of this lane's first token. The other + // atom rows / tokens are whole swizzle-atom columns / multiples of 8 rows further, i.e. plain offsets + static constexpr uint32_t STS_ATOM_COL_STRIDE_BYTES = B_TOPK * 128; + static constexpr uint32_t STS_TOKEN_STRIDE_BYTES = NUM_ROWS_PER_WAVEFRONT * 128; + uint32_t sts_offsets[8]; + { + Tensor sKV = make_tensor(make_smem_ptr(smem.kv_slots[0]), ku::make_umma_canonical_k_major_layout()); + CUTE_UNROLL + for (uint32_t i = 0; i < 8; ++i) { + sts_offsets[i] = (uint32_t)((&sKV(get_row_idx(0), i * 8) - smem.kv_slots[0]) * sizeof(bf16)); + } + } + + OuterloopArgs cur_job = get_first_job(); + RingBufferState rs; + do { + // Processes one KV block whose first `num_orig_rows` rows come from `kv` (fp8) and the others from `extra_kv` (fp4), + // the first NUM_FP8_TOKENS tokens of this thread being the fp8 ones. NUM_FP8_TOKENS is a template parameter so that + // the format of every token, and with it the code operating on the token's registers, is fixed at compile time + auto process_block = [&](uint32_t num_orig_rows) { + // Runs `callable.template operator()(token_idx)` for every token of this thread, F being its format + auto for_each_token = [&](auto callable) { + cute::for_each(cute::make_int_sequence{}, [&](auto token_idx) { + if constexpr (token_idx < NUM_FP8_TOKENS) { + callable.template operator()(token_idx); + } else { + callable.template operator()(token_idx); + } + }); + }; + + auto [indices_buf_idx, indices_bar_phase] = rs.get(); + smem.bar_indices_full[indices_buf_idx].wait(indices_bar_phase); + // The scales of this lane's chunks, per token. fp8 tokens: one ue8m0 per QUANT_TILE_SIZE elements = per 2 chunks, + // read from smem here. fp4 tokens: two e4m3 per chunk, loaded from global memory after the gather below + static constexpr uint32_t FP8_SCALE_BYTES_PER_LANE = OrigKVFormat::NUM_CHUNKS_PER_ROW / NUM_LANES_PER_ROW / 2; // 4 / 2 + static constexpr uint32_t FP4_SCALE_BYTES_PER_LANE = ExtraKVFormat::NUM_CHUNKS_PER_ROW / NUM_LANES_PER_ROW * 2; // 8 / 4 + static_assert(FP8_SCALE_BYTES_PER_LANE <= FP4_SCALE_BYTES_PER_LANE && FP4_SCALE_BYTES_PER_LANE % 4 == 0); + uint32_t cached_scales[NUM_TOKENS_PER_THREAD][FP4_SCALE_BYTES_PER_LANE / 4]; + const uint8_t *scale_ptrs[NUM_TOKENS_PER_THREAD]; + for_each_token([&](uint32_t token_idx) { + uint32_t row = get_row_idx(token_idx); + if constexpr (F::IS_FP4) { + scale_ptrs[token_idx] = smem.decode_scale_ptrs[indices_buf_idx][row]; + } else { + // Arranged so that chunk g uses byte g / 2 of the loaded word: a lane owning 8 chunks swaps their two halves + // (see get_chunk_base), i.e. the two halves of its 4 scale bytes; a lane owning 4 chunks rotates the chunk + // indices, so its 2 scale bytes are read from the rotated position + const uint8_t *scale_src = smem.decode_scales[indices_buf_idx] + row * NUM_SCALES_EACH_TOKEN_PER_CTA; + if constexpr (FP8_SCALE_BYTES_PER_LANE == 4) { + uint32_t scales = *(uint32_t*)(scale_src + idx_in_row * 4); + cached_scales[token_idx][0] = chunk_rotation ? __byte_perm(scales, scales, 0x1032) : scales; + } else { + static_assert(FP8_SCALE_BYTES_PER_LANE == 2); + cached_scales[token_idx][0] = *(uint16_t*)(scale_src + get_chunk_base.template operator()(0) / 2); + } + } + }); + + auto [kv_slot_idx, kv_bar_phase] = rs.get(); + smem.bar_kv_slot_empty[kv_slot_idx].wait(kv_bar_phase^1); + uint8_t *slot_base = (uint8_t*)smem.kv_slots[kv_slot_idx]; + + cute::for_each(cute::make_int_sequence{}, [&](auto i) { + // Each tma_gather4 covers 4 consecutive rows of one token, which share one tensor map. Only the + // elected lane needs the coordinates, and the 4 rows are consecutive in `decode_tma_coords`, so it + // reads them with one LDS.128 + uint32_t row_start = local_warp_idx * NUM_ROWS_PER_WARP + i * 4; + constexpr bool is_fp4 = i / (NUM_ROWS_PER_WAVEFRONT / 4) >= NUM_FP8_TOKENS; + auto tensor_map = is_fp4 ? + (cta_idx == 0 ? &tma_params.tensor_map_extra_kv_fp4_part_cta0 : &tma_params.tensor_map_extra_kv_fp4_part_cta1) : + (cta_idx == 0 ? &tma_params.tensor_map_kv_fp8_part_cta0 : &tma_params.tensor_map_kv_fp8_part_cta1); + if (elect_one_sync()) { + int4 coords = *(int4*)(smem.decode_tma_coords[indices_buf_idx] + row_start); + ku::tma_gather4( + tensor_map, + smem.bar_raw_kv_full, + slot_base + get_raw_row_offset(row_start), + 0, + coords, + (int64_t)TMA::CacheHintSm90::EVICT_FIRST + ); + } + }); + smem.bar_indices_empty[indices_buf_idx].arrive(); + + // Invalid fp4 tokens (nullptr) are zero-filled by TMA and get scale 0, so they dequantize to 0. The "memory" + // clobber keeps the loads behind the gather4s above, otherwise ptxas hoists them and the gather4s wait for them + for_each_token([&](uint32_t token_idx) { + if constexpr (F::IS_FP4) { + uint64_t scale_addr = (uint64_t)(scale_ptrs[token_idx] + get_chunk_base.template operator()(0) * 2); + if constexpr (FP4_SCALE_BYTES_PER_LANE == 8) { + uint64_t scales = 0; + if (scale_ptrs[token_idx] != nullptr) { + asm volatile ( + "ld.global.nc.L1::no_allocate.b64 %0, [%1];" + : "=l"(scales) + : "l"(scale_addr) + : "memory" + ); + } + *(uint64_t*)cached_scales[token_idx] = scales; + } else { + static_assert(FP4_SCALE_BYTES_PER_LANE == 4); + uint32_t scales = 0; + if (scale_ptrs[token_idx] != nullptr) { + asm volatile ( + "ld.global.nc.L1::no_allocate.b32 %0, [%1];" + : "=r"(scales) + : "l"(scale_addr) + : "memory" + ); + } + cached_scales[token_idx][0] = scales; + } + } + }); + + if (idx_in_warpgroup == 0) { + smem.bar_raw_kv_full.arrive_and_expect_tx(num_orig_rows * OrigKVFormat::RAW_TOKEN_SMEM_STRIDE + (B_TOPK - num_orig_rows) * ExtraKVFormat::RAW_TOKEN_SMEM_STRIDE); + } + smem.bar_raw_kv_full.wait(rs.get<1>().second); + + uint32_t cached_input[NUM_TOKENS_PER_THREAD][MAX_CHUNKS_PER_LANE][4]; // fp4 tokens (half as many chunks) use the first half + for_each_token([&](uint32_t token_idx) { + uint32_t row = get_row_idx(token_idx); + uint8_t *row_base = slot_base + get_raw_row_offset(row) + row % 4 * F::RAW_TOKEN_SMEM_STRIDE; + CUTE_UNROLL + for (uint32_t g = 0; g < F::NUM_CHUNKS_PER_ROW / NUM_LANES_PER_ROW; ++g) { + *(__int128_t*)(cached_input[token_idx][g]) = ku::ld_shared( + row_base + get_chunk_base.template operator()(g / 4) * 16 + g % 4 * 16 + ); + } + }); + NamedBarrier::arrive_and_wait(128, 7); // Make sure everyone has finished reading + + for_each_token([&](uint32_t token_idx) { + static constexpr uint32_t C_LANE = F::NUM_CHUNKS_PER_ROW / NUM_LANES_PER_ROW; + static constexpr uint32_t NUM_STS_PER_CHUNK = F::CHUNK_ELEMS * sizeof(bf16) / 16; + ku::nvbf16x2 fp4_scales[2]; // fp4: the two scales of each of the current pair of chunks, as bf16 + CUTE_UNROLL + for (uint32_t g = 0; g < C_LANE; ++g) { + ku::nvbf16x2 data_bf16x2[F::CHUNK_ELEMS / 2]; + if constexpr (F::IS_FP4) { + if (g % 2 == 0) { + fp8x4_to_bf16x2x2(cached_scales[token_idx][g / 2], fp4_scales); + } + CUTE_UNROLL + for (uint32_t j = 0; j < 4; ++j) { + fp4x8_to_bf16x2x4(cached_input[token_idx][g][j], data_bf16x2 + j * 4); + } + ku::nvbf16x2 scale_lo = __low2bfloat162(fp4_scales[g % 2]), scale_hi = __high2bfloat162(fp4_scales[g % 2]); + CUTE_UNROLL + for (uint32_t k = 0; k < F::CHUNK_ELEMS / 2; ++k) { + data_bf16x2[k] = __hmul2(data_bf16x2[k], k < F::QUANT_TILE_SIZE / 2 ? scale_lo : scale_hi); + } + } else { + __nv_fp8_e8m0 scale = ((__nv_fp8_e8m0*)cached_scales[token_idx])[g * F::CHUNK_ELEMS / F::QUANT_TILE_SIZE]; + CUTE_UNROLL + for (uint32_t k = 0; k < F::CHUNK_ELEMS / 2; ++k) { + data_bf16x2[k] = fp8x2_to_bf16x2_with_scale(((ku::nve4m3x2*)cached_input[token_idx][g])[k], scale); + } + } + // The chunk covers parts [(g % CHUNKS_PER_ATOM_ROW) * NUM_STS_PER_CHUNK, +NUM_STS_PER_CHUNK) of the atom row + // chunk_base / CHUNKS_PER_ATOM_ROW + g % 4 / CHUNKS_PER_ATOM_ROW (exact: chunk_base is a multiple of + // CHUNKS_PER_ATOM_ROW), so the swizzled offset is selected at compile time + static constexpr uint32_t CHUNKS_PER_ATOM_ROW = 128 / (F::CHUNK_ELEMS * 2); + uint32_t atom_col = get_chunk_base.template operator()(g / 4) / CHUNKS_PER_ATOM_ROW + g % 4 / CHUNKS_PER_ATOM_ROW; + CUTE_UNROLL + for (uint32_t j = 0; j < NUM_STS_PER_CHUNK; ++j) { + ku::st_shared( + slot_base + sts_offsets[g % CHUNKS_PER_ATOM_ROW * NUM_STS_PER_CHUNK + j] + atom_col * STS_ATOM_COL_STRIDE_BYTES + token_idx * STS_TOKEN_STRIDE_BYTES, + *(__int128_t*)(data_bf16x2 + j * 4) + ); + } + } + }); + + fence_view_async_shared(); + arrive_on_cta0_barrier(smem.bar_kv_slot_full[kv_slot_idx]); + rs.update(); + }; + + run_along_kv_blocks(cur_job, [&](uint32_t kv_block_idx) { + if constexpr (LOC == KVLocation::ORIG) { + process_block.template operator()(B_TOPK); + } else if constexpr (LOC == KVLocation::EXTRA) { + process_block.template operator()<0>(0); + } else { + // The block straddles the orig/extra boundary. The number of fp8 tokens of this thread is warp-uniform since + // topk % 8 == 0 (asserted by `run()`); dispatch to the matching instantiation + uint32_t num_orig_rows = cur_job.num_orig_slots - kv_block_idx * B_TOPK; + uint32_t num_fp8_tokens = (uint32_t)std::clamp((int)num_orig_rows - (int)(local_warp_idx * NUM_ROWS_PER_WARP), 0, (int)NUM_ROWS_PER_WARP) / NUM_ROWS_PER_WAVEFRONT; + [&](std::integer_sequence) { + ((num_fp8_tokens == Ks ? process_block.template operator()(num_orig_rows) : void()), ...); + }(std::make_integer_sequence{}); + } + }); + + cur_job = get_next_job(cur_job); + } while (cur_job.is_valid); + } else { + // KV producer: loads the fp8 part of every selected KV token directly from global + // memory, dequantizes it in registers, and stores the bf16 result into the KV slot (SW128 + // K-major layout, the first D_FP8 columns). Also load the bf16 (RoPE) part via TMA gather4. + + // Maximize shared memory throughput by using LDS.128 and STS.128 with 8 tokens processed per warp. + // Lane-to-token mapping for the 8 tokens being processed: + // + // 0 1 2 3 + // 4 5 6 7 + // 8 9 10 11 + // 12 13 14 15 + // 16 17 18 19 + // 20 21 22 23 + // 24 25 26 27 + // 28 29 30 31 + // + // Dequant pipeline: + // 1. Load raw FP8 KV from global memory into the target shared-memory KV buffer via TMA gather4. + // 2. While TMA is in flight, load scale factors via plain global loads. + // 3. Wait for TMA, then read raw FP8 KV via LDS.128. Synchronize afterward to prevent the + // subsequent write-back from overwriting the read data. + // 4. Dequantize in registers, then write back bf16 result via STS.128. + // + // This pipeline is bank-conflict-free: + // - On TMA gather4 load: if the per-row FP8 count (D_FP8_CTA0/1) is a multiple of 128, + // the box size is padded by 64B so that lanes i..i+8 see no bank conflicts during LDS.128. + // - On STS.128 write-back: swizzling avoids bank conflicts. + using fp8_e8m0 = __nv_fp8_e8m0; + uint32_t local_warp_idx = warp_idx - 4; + static constexpr uint32_t NUM_DEQUANT_WARPS = 4, NUM_ROWS_PER_WARP = 8; + static constexpr uint32_t GROUP_SIZE = 4; + static constexpr uint32_t NUM_COLS_PER_GROUP = D_VO / CLUSTER_SIZE / (GROUP_SIZE*16); + static_assert((D_VO/CLUSTER_SIZE) % (GROUP_SIZE*16) == 0); + uint32_t group_idx = lane_idx / GROUP_SIZE, idx_in_group = lane_idx % GROUP_SIZE; + static constexpr uint32_t NUM_CHUNKS_PER_WARP = B_TOPK / (NUM_DEQUANT_WARPS*NUM_ROWS_PER_WARP); + static_assert(B_TOPK % (NUM_DEQUANT_WARPS*NUM_ROWS_PER_WARP) == 0); + static constexpr uint32_t NUM_TOKENS_PER_THREAD = NUM_CHUNKS_PER_WARP; + + auto get_row_idx = [&](uint32_t local_row_idx) { + return local_row_idx*NUM_DEQUANT_WARPS*NUM_ROWS_PER_WARP + local_warp_idx*NUM_ROWS_PER_WARP + group_idx; + }; + auto [sts_base_offset_0, sts_base_offset_1] = [&] { + Tensor sKV = make_tensor(make_smem_ptr(smem.kv_slots[0]), ku::make_umma_canonical_k_major_layout()); + return std::pair { + (uint32_t)(&sKV(local_warp_idx*NUM_ROWS_PER_WARP+group_idx, idx_in_group*16 + 0) - smem.kv_slots[0]), + (uint32_t)(&sKV(local_warp_idx*NUM_ROWS_PER_WARP+group_idx, idx_in_group*16 + 8) - smem.kv_slots[0]) + }; + }(); + + uint32_t d_fp8_this_cta = cta_idx == 0 ? D_FP8_CTA0 : D_FP8_CTA1; + uint32_t d_fp8_this_cta_padded = d_fp8_this_cta + ((cta_idx==0?IS_CTA0_RAW_KV_PADDED:IS_CTA1_RAW_KV_PADDED)?64:0); + + OuterloopArgs cur_job = get_first_job(); + RingBufferState rs; + do { + auto loop_body = [&](uint32_t kv_block_idx) { + // Tells whether the unified slot `pos` comes from the extra KV. For ORIG / EXTRA + // blocks this is a compile-time constant, so all the source selections below fold + // into branch-free code; only the (at most one) ORIG_AND_EXTRA block resolves the + // source at runtime, per row + auto is_pos_in_extra = [&](uint32_t pos) -> bool { + if constexpr (LOC == KVLocation::ORIG) { + return false; + } else if constexpr (LOC == KVLocation::EXTRA) { + return true; + } else { + return pos >= cur_job.num_orig_slots; + } + }; + + static constexpr uint32_t NUM_SCALED_EACH_TOKEN_LOCAL = NUM_SCALES_EACH_TOKEN / CLUSTER_SIZE; + fp8_e8m0 cached_scales[NUM_TOKENS_PER_THREAD][NUM_SCALED_EACH_TOKEN_LOCAL]; + int cached_tma_coord[NUM_TOKENS_PER_THREAD]; + auto [indices_buf_idx, indices_bar_phase] = rs.get(); + smem.bar_indices_full[indices_buf_idx].wait(indices_bar_phase); + CUTE_UNROLL + for (uint32_t local_row_idx = 0; local_row_idx < NUM_TOKENS_PER_THREAD; ++local_row_idx) { + uint32_t row = get_row_idx(local_row_idx); + cached_tma_coord[local_row_idx] = smem.decode_tma_coords[indices_buf_idx][row]; + uint8_t *scale_src = smem.decode_scales[indices_buf_idx] + + row * NUM_SCALES_EACH_TOKEN_PER_CTA; + if constexpr (NUM_SCALED_EACH_TOKEN_LOCAL == 4) { + *(uint32_t*)cached_scales[local_row_idx] = *(uint32_t*)scale_src; + } else if constexpr (NUM_SCALED_EACH_TOKEN_LOCAL == 8) { + *(uint64_t*)cached_scales[local_row_idx] = *(uint64_t*)scale_src; + } else { + static_assert(NUM_SCALED_EACH_TOKEN_LOCAL == 16); + *(__int128_t*)cached_scales[local_row_idx] = *(__int128_t*)scale_src; + } + } + smem.bar_indices_empty[indices_buf_idx].arrive(); + + auto [kv_slot_idx, kv_bar_phase] = rs.get(); + smem.bar_kv_slot_empty[kv_slot_idx].wait(kv_bar_phase^1); + + int4 collected_tma_coords[NUM_TOKENS_PER_THREAD][2]; + CUTE_UNROLL + for (uint32_t local_row_idx = 0; local_row_idx < NUM_TOKENS_PER_THREAD; ++local_row_idx) { + CUTE_UNROLL + for (uint32_t i = 0; i < 2; ++i) { + // Each tma_gather4 covers 4 consecutive rows, which must share one tensor + // map. `run()` asserts topk % 4 == 0 when the extra KV is present, so a + // 4-row group never straddles the orig/extra boundary + uint32_t row_start = local_row_idx*NUM_DEQUANT_WARPS*NUM_ROWS_PER_WARP + local_warp_idx*NUM_ROWS_PER_WARP + i*4; + bool group_in_extra = is_pos_in_extra(kv_block_idx*B_TOPK + row_start); + auto tensor_map = group_in_extra ? + (cta_idx == 0 ? &tma_params.tensor_map_extra_kv_fp8_part_cta0 : &tma_params.tensor_map_extra_kv_fp8_part_cta1) : + (cta_idx == 0 ? &tma_params.tensor_map_kv_fp8_part_cta0 : &tma_params.tensor_map_kv_fp8_part_cta1); + int4 coords; + coords.x = i == 0 ? cached_tma_coord[local_row_idx] : __shfl_sync(0xFFFFFFFF, cached_tma_coord[local_row_idx], i*16); // Since the thread being elected is always lane 0 on SM100 + coords.y = __shfl_sync(0xFFFFFFFF, cached_tma_coord[local_row_idx], i*16+4); + coords.z = __shfl_sync(0xFFFFFFFF, cached_tma_coord[local_row_idx], i*16+8); + coords.w = __shfl_sync(0xFFFFFFFF, cached_tma_coord[local_row_idx], i*16+12); + collected_tma_coords[local_row_idx][i] = coords; + if (elect_one_sync()) { + auto smem_ptr = (fp8_e4m3*)smem.kv_slots[kv_slot_idx] + row_start * d_fp8_this_cta_padded; + ku::tma_gather4( + tensor_map, + smem.bar_raw_kv_full, + smem_ptr, + 0, + coords, + (int64_t)TMA::CacheHintSm90::EVICT_FIRST + ); + } + } + } + + if (D_BF16 > 0 && cta_idx+1 == CLUSTER_SIZE && elect_one_sync()) { + CUTE_UNROLL + for (uint32_t local_row_idx = 0; local_row_idx < NUM_TOKENS_PER_THREAD; ++local_row_idx) { + CUTE_UNROLL + for (uint32_t i = 0; i < 2; ++i) { + uint32_t row_start = local_row_idx*NUM_DEQUANT_WARPS*NUM_ROWS_PER_WARP + local_warp_idx*NUM_ROWS_PER_WARP + i*4; + // Like the fp8 part above, a 4-row gather4 group never straddles the orig/extra boundary + auto tensor_map = is_pos_in_extra(kv_block_idx*B_TOPK + row_start) ? &tma_params.tensor_map_extra_kv_bf16_part : &tma_params.tensor_map_kv_bf16_part; + auto smem_ptr = smem.kv_slots[kv_slot_idx] + (IS_2CTA ? (D_FP8-D_VO/2)/64 : D_FP8/64) * B_TOPK * 64 + row_start * 64; + if constexpr (IS_2CTA) { + ku::tma_gather4_cta_group_2( + tensor_map, + smem.bar_kv_slot_full[kv_slot_idx], + smem_ptr, + 0, + collected_tma_coords[local_row_idx][i], + (int64_t)TMA::CacheHintSm90::EVICT_FIRST + ); + } else { + ku::tma_gather4( + tensor_map, + smem.bar_kv_slot_full[kv_slot_idx], + smem_ptr, + 0, + collected_tma_coords[local_row_idx][i], + (int64_t)TMA::CacheHintSm90::EVICT_FIRST + ); + } + } + } + } + + if (idx_in_warpgroup == 0) { + smem.bar_raw_kv_full.arrive_and_expect_tx(B_TOPK*d_fp8_this_cta_padded*sizeof(fp8_e4m3)); + } + smem.bar_raw_kv_full.wait(rs.get<1>().second); + + fp8_e4m3 cached_input[NUM_TOKENS_PER_THREAD][NUM_COLS_PER_GROUP][16]; + CUTE_UNROLL + for (uint32_t local_row_idx = 0; local_row_idx < NUM_TOKENS_PER_THREAD; ++local_row_idx) { + uint32_t row = get_row_idx(local_row_idx); + for (uint32_t local_col_idx = 0; local_col_idx < NUM_COLS_PER_GROUP; ++local_col_idx) { + if (MODEL_TYPE == ModelType::V4 && cta_idx+1 == CLUSTER_SIZE && local_col_idx+1 == NUM_COLS_PER_GROUP) { + // Skip the last K/V block for V4 + continue; + } + *(__int128_t*)(cached_input[local_row_idx][local_col_idx]) = ku::ld_shared( + (fp8_e4m3*)smem.kv_slots[kv_slot_idx] + + row*d_fp8_this_cta_padded + + local_col_idx*GROUP_SIZE*16 + + idx_in_group*16 + ); + } + } + NamedBarrier::arrive_and_wait(128, 7); // Make sure everyone has finished reading + + CUTE_UNROLL + for (uint32_t local_row_idx = 0; local_row_idx < NUM_TOKENS_PER_THREAD; ++local_row_idx) { + CUTE_UNROLL + for (uint32_t local_col_idx = 0; local_col_idx < NUM_COLS_PER_GROUP; ++local_col_idx) { + if (MODEL_TYPE == ModelType::V4 && cta_idx+1 == CLUSTER_SIZE && local_col_idx+1 == NUM_COLS_PER_GROUP) { + // Skip the last K/V block for V4 + continue; + } + ku::nve4m3x2 data_fp8x2[8]; + ku::nvbf16x2 data_bf16x2[8]; + *(__int128_t*)data_fp8x2 = *(__int128_t*)(cached_input[local_row_idx][local_col_idx]); + static_assert(KV_QUANT_TILE_SIZE == 64 || KV_QUANT_TILE_SIZE == 32); + uint32_t scale_idx = KV_QUANT_TILE_SIZE == 64 ? local_col_idx : local_col_idx*2 + (idx_in_group >= GROUP_SIZE/2); + CUTE_UNROLL + for (uint32_t j = 0; j < 8; ++j) { + data_bf16x2[j] = fp8x2_to_bf16x2_with_scale(data_fp8x2[j], cached_scales[local_row_idx][scale_idx]); + } + ku::st_shared( + smem.kv_slots[kv_slot_idx] + sts_base_offset_0 + local_row_idx*(NUM_DEQUANT_WARPS*NUM_ROWS_PER_WARP*64) + local_col_idx*(B_TOPK*GROUP_SIZE*16), + *(__int128_t*)(data_bf16x2 + 0) + ); + ku::st_shared( + smem.kv_slots[kv_slot_idx] + sts_base_offset_1 + local_row_idx*(NUM_DEQUANT_WARPS*NUM_ROWS_PER_WARP*64) + local_col_idx*(B_TOPK*GROUP_SIZE*16), + *(__int128_t*)(data_bf16x2 + 4) + ); + } + } + + fence_view_async_shared(); + arrive_on_cta0_barrier(smem.bar_kv_slot_full[kv_slot_idx]); + rs.update(); + }; + + run_along_kv_blocks(cur_job, loop_body); + + cur_job = get_next_job(cur_job); + } while (cur_job.is_valid); + } + } else { + // KV Producer (prefill): gathers the whole bf16 KV block via TMA gather4 + if (elect_one_sync()) { + OuterloopArgs cur_job = get_first_job(); + RingBufferState rs; + uint32_t local_warp_idx = warp_idx - 4; + do { + for (uint32_t i = 0; i < cur_job.num_kv_blocks; ++i) { + static_assert(B_TOPK % (4*4) == 0); + static constexpr uint32_t NUM_ROW_PER_WARP = B_TOPK / 4; + int4 topk_idxs[NUM_ROW_PER_WARP / 4]; + CUTE_UNROLL + for (uint32_t local_row = 0; local_row < NUM_ROW_PER_WARP / 4; local_row += 1) { + uint32_t row = local_row * 4 * 4 + local_warp_idx * 4; + uint32_t pos = i * B_TOPK + row; + // Predicate the load on `pos < topk_length` to avoid reading OOB of the + // indices row in the last (partial) KV block. Chunks with no slot below + // topk_length are masked out by the indices generator warp anyway, so feed + // them invalid indices (-1), which tma_gather4 bounds-checks into zero-fill + topk_idxs[local_row] = pos < cur_job.topk_length + ? __ldg((int4*)(params.indices + cur_job.s_q_idx * params.stride_indices_s_q + pos)) + : int4{-1, -1, -1, -1}; + } + + auto [kv_slot_idx, kv_bar_phase] = rs.get(); + smem.bar_kv_slot_empty[kv_slot_idx].wait(kv_bar_phase^1); + CUTE_UNROLL + for (uint32_t local_row = 0; local_row < NUM_ROW_PER_WARP / 4; local_row += 1) { + uint32_t row = local_row * 4 * 4 + local_warp_idx * 4; + CUTE_UNROLL + for (uint32_t tile_idx = 0; tile_idx < (CLUSTER_SIZE == 2 ? (D_QK/2/64) : (D_QK/64)); ++tile_idx) { + /* + For cases where CLUSTER_SIZE == 1, each CTA reads the full K/V + For cases where CLUSTER_SIZE == 2, CTA0 reads KV[:, :D_QK/2] and CTA1 reads KV[:, D_QK/2:], since dual GeM< is used + */ + if constexpr (CLUSTER_SIZE == 1) { + ku::tma_gather4( + &tma_params.tensor_map_kv, + smem.bar_kv_slot_full[kv_slot_idx], + smem.kv_slots[kv_slot_idx] + tile_idx * B_TOPK * 64 + row * 64, + tile_idx * 64, + topk_idxs[local_row], + (int64_t)TMA::CacheHintSm90::EVICT_LAST + ); + } else { + ku::tma_gather4_cta_group_2( + &tma_params.tensor_map_kv, + smem.bar_kv_slot_full[kv_slot_idx], + smem.kv_slots[kv_slot_idx] + tile_idx * B_TOPK * 64 + row * 64, + tile_idx * 64 + cta_idx * (D_QK/2), + topk_idxs[local_row], + (int64_t)TMA::CacheHintSm90::EVICT_LAST + ); + } + } + } + rs.update(); + } + cur_job = get_next_job(cur_job); + } while (cur_job.is_valid); + } + } + } + +#else + if (cute::thread0()) { + CUTE_INVALID_CONTROL_PATH("This kernel only supports sm100"); + } +#endif +} + + +template +__global__ void __launch_bounds__(Kernel::NUM_THREADS, 1, Kernel::CLUSTER_SIZE) +fwd_kernel(__grid_constant__ const typename Kernel::Params params, __grid_constant__ const typename Kernel::TMAParams tma_params, __grid_constant__ const typename Kernel::AuxParams aux_params) { + Kernel::devfunc(params, tma_params, aux_params); +} + + +template +void Kernel::run(const Params ¶ms) { + KU_ASSERT(params.h_q == H_Q); + KU_ASSERT(params.h_kv == 1); + KU_ASSERT(params.d_qk == D_QK); + KU_ASSERT(params.d_v == D_VO); + KU_ASSERT(params.stride_indices_s_q*sizeof(int) % 32 == 0, "indices.stride(0) must be 32B aligned, got %d elements (%d bytes)", params.stride_indices_s_q, (int)(params.stride_indices_s_q*sizeof(int))); + + KU_ASSERT(params.enable_q_norm == ENABLE_Q_NORM); + KU_ASSERT(params.is_rope_neox_style == false && params.rope_dim == 64); + KU_ASSERT(params.wv_group_size == WV_GROUP_SIZE); + KU_ASSERT(params.num_per_channels == O_QUANT_TILE_SIZE); + KU_ASSERT(params.use_tma_aligned_col_major_sf == true && params.round_sf == true && params.use_packed_ue8m0 == true); + + TMAParams tma_params = {}; + if constexpr (IS_DECODE) { + KU_ASSERT(params.b == 1, "Only batch size 1 is supported for the fused decoding kernel"); + KU_ASSERT(params.model_type == MODEL_TYPE && params.extra_model_type == EXTRA_MODEL_TYPE); + if (params.extra_topk > 0) { + // A KV block may straddle the orig/extra boundary (see KVLocation::ORIG_AND_EXTRA). Since + // one TMA gather4 covers 4 consecutive rows sharing one tensor map, the boundary (i.e. + // topk) must be aligned to 4 rows; the common dequant path resolves the format per 8 rows + KU_ASSERT(params.topk % (HAS_FP4_KV ? 8 : 4) == 0, "topk (%d) must be a multiple of %d when the extra KV cache is used", params.topk, HAS_FP4_KV ? 8 : 4); + } + auto make_fp4_kv_tensor_map = [](bool is_extra, void *kv_ptr, int num_blocks, int64_t block_stride_bytes, int row_stride_bytes) -> std::pair { + using F = ExtraKVFormat; + KU_ASSERT((int64_t)kv_ptr % 16 == 0, "The base address of %skv (%p) must be 16B aligned", is_extra?"extra_":"", kv_ptr); + KU_ASSERT(row_stride_bytes == (int)F::BYTES_PER_TOKEN, "%skv_cache.stride(-2) (%d) must be %d, i.e. each page block in the KV cache must be contiguous", is_extra?"extra_":"", row_stride_bytes, (int)F::BYTES_PER_TOKEN); + KU_ASSERT(block_stride_bytes % F::TMA_K_STRIDE == 0, "%skv_cache.stride(0) (%ld) must be a multiple of %d. Padding might be necessary", is_extra?"extra_":"", block_stride_bytes, (int)F::TMA_K_STRIDE); + KU_ASSERT((uint64_t)num_blocks * (uint64_t)(block_stride_bytes / F::TMA_K_STRIDE) <= INT32_MAX, "%skv: too many rows for the int32 TMA coordinates", is_extra?"extra_":""); + // One 2D view per CTA: rows of the CTA's D_FP4 / CLUSTER_SIZE dims, TMA_K_STRIDE bytes apart. The box is 16 B + // wider than the row (the extra bytes are out of bounds and zero-filled by TMA), see KVFormat::RAW_TOKEN_SMEM_STRIDE + static_assert(F::RAW_TOKEN_DATA_BYTES % 4 == 0); + auto make_tensor_map_for_cta = [&](uint32_t cta_idx) { + return ku::make_tensor_map( + {(uint64_t)F::RAW_TOKEN_DATA_BYTES/4, (uint64_t)num_blocks * (uint64_t)(block_stride_bytes / F::TMA_K_STRIDE)}, + {(uint64_t)F::TMA_K_STRIDE}, + {F::RAW_TOKEN_SMEM_STRIDE/4, 1}, // Use UINT32 as dtype and divide the box size by 4, like the fp8 part + (uint8_t*)kv_ptr + cta_idx * F::RAW_TOKEN_DATA_BYTES, + CUtensorMapDataType::CU_TENSOR_MAP_DATA_TYPE_UINT32, + CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_NONE, + CUtensorMapL2promotion::CU_TENSOR_MAP_L2_PROMOTION_L2_256B + ); + }; + CUtensorMap fp4_part_tensor_map_cta1 = {}; + if constexpr (IS_2CTA) { + fp4_part_tensor_map_cta1 = make_tensor_map_for_cta(1); + } + return {make_tensor_map_for_cta(0), fp4_part_tensor_map_cta1}; + }; + auto make_kv_tensor_map = [](bool is_extra, void *kv_ptr, int num_blocks, int64_t block_stride_bytes, int row_stride_bytes) -> std::tuple { + KU_ASSERT((int64_t)kv_ptr % 16 == 0, "The base address of %skv (%p) must be 16B aligned", is_extra?"extra_":"", kv_ptr); + KU_ASSERT(row_stride_bytes == (int)KV_CACHE_BYTES_PER_TOKEN, "%skv_cache.stride(-2) (%d) must be %d, i.e. each page block in the KV cache must be contiguous", is_extra?"extra_":"", row_stride_bytes, (int)KV_CACHE_BYTES_PER_TOKEN); + KU_ASSERT(block_stride_bytes % TMA_K_STRIDE == 0, "%skv_cache.stride(0) (%ld) must be a multiple of %d. Padding might be necessary", is_extra?"extra_":"", block_stride_bytes, (int)TMA_K_STRIDE); + KU_ASSERT((uint64_t)num_blocks * (uint64_t)(block_stride_bytes / TMA_K_STRIDE) <= INT32_MAX, "%skv: too many rows for the int32 TMA coordinates", is_extra?"extra_":""); + static_assert(D_FP8_CTA0%4 == 0); + auto fp8_part_tensor_map_cta0 = ku::make_tensor_map( + {(uint64_t)D_FP8_CTA0/4, (uint64_t)num_blocks * (uint64_t)(block_stride_bytes / TMA_K_STRIDE)}, + {(uint64_t)TMA_K_STRIDE}, + {RAW_FP8_TOKEN_SMEM_STRIDE_CTA0/4, 1}, // Use UINT32 as dtype and divide the box size by 4, to satisfy the requirement that TMA's box size must <= 256. Add 64 bytes to prevent bank conflict when loading from SMEM + (uint8_t*)kv_ptr, + CUtensorMapDataType::CU_TENSOR_MAP_DATA_TYPE_UINT32, + CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_NONE, + CUtensorMapL2promotion::CU_TENSOR_MAP_L2_PROMOTION_L2_256B + ); + CUtensorMap fp8_part_tensor_map_cta1 = {}; + if constexpr (IS_2CTA) { + static_assert(D_FP8_CTA1%4 == 0); + fp8_part_tensor_map_cta1 = ku::make_tensor_map( + {(uint64_t)D_FP8_CTA1/4, (uint64_t)num_blocks * (uint64_t)(block_stride_bytes / TMA_K_STRIDE)}, + {(uint64_t)TMA_K_STRIDE}, + {RAW_FP8_TOKEN_SMEM_STRIDE_CTA1/4, 1}, + (uint8_t*)kv_ptr + D_FP8_CTA0, + CUtensorMapDataType::CU_TENSOR_MAP_DATA_TYPE_UINT32, + CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_NONE, + CUtensorMapL2promotion::CU_TENSOR_MAP_L2_PROMOTION_L2_256B + ); + } + CUtensorMap bf16_part_tensor_map = {}; + if constexpr (D_BF16 > 0) { + bf16_part_tensor_map = ku::make_tensor_map( + {(uint64_t)D_BF16, (uint64_t)num_blocks * (uint64_t)(block_stride_bytes / TMA_K_STRIDE)}, + {(uint64_t)TMA_K_STRIDE}, + {D_BF16, 1}, + (uint8_t*)kv_ptr + D_FP8, + CUtensorMapDataType::CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, + CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_128B, + CUtensorMapL2promotion::CU_TENSOR_MAP_L2_PROMOTION_L2_128B + ); + } + return {fp8_part_tensor_map_cta0, fp8_part_tensor_map_cta1, bf16_part_tensor_map}; + }; + { + auto [fp8_part_cta0, fp8_part_cta1, bf16_part] = make_kv_tensor_map(false, params.kv, params.num_blocks, params.stride_kv_block, params.stride_kv_row); + tma_params.tensor_map_kv_fp8_part_cta0 = fp8_part_cta0; + tma_params.tensor_map_kv_fp8_part_cta1 = fp8_part_cta1; + tma_params.tensor_map_kv_bf16_part = bf16_part; + } + if (params.extra_topk > 0) { + if constexpr (HAS_FP4_KV) { + auto [fp4_part_cta0, fp4_part_cta1] = make_fp4_kv_tensor_map(true, params.extra_kv, params.extra_num_blocks, params.stride_extra_kv_block, params.stride_extra_kv_row); + tma_params.tensor_map_extra_kv_fp4_part_cta0 = fp4_part_cta0; + tma_params.tensor_map_extra_kv_fp4_part_cta1 = fp4_part_cta1; + } else { + auto [fp8_part_cta0, fp8_part_cta1, bf16_part] = make_kv_tensor_map(true, params.extra_kv, params.extra_num_blocks, params.stride_extra_kv_block, params.stride_extra_kv_row); + tma_params.tensor_map_extra_kv_fp8_part_cta0 = fp8_part_cta0; + tma_params.tensor_map_extra_kv_fp8_part_cta1 = fp8_part_cta1; + tma_params.tensor_map_extra_kv_bf16_part = bf16_part; + } + } + } else { + tma_params.tensor_map_kv = ku::make_tensor_map( + {(uint64_t)D_QK, (uint64_t)params.s_kv}, + {(uint64_t)params.stride_kv_s_kv * sizeof(bf16)}, + {64, 1}, + params.kv, + CUtensorMapDataType::CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, + CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_128B, + CUtensorMapL2promotion::CU_TENSOR_MAP_L2_PROMOTION_L2_256B + ); + } + + auto aux_params = AuxParams { + }; + if constexpr (IS_DECODE) { + aux_params.fast_divmod_page_block_size = cutlass::FastDivmod(params.page_block_size); + aux_params.fast_divmod_extra_page_block_size = cutlass::FastDivmod(params.extra_kv != nullptr ? params.extra_page_block_size : 1); + } + + auto kernel = &fwd_kernel; + constexpr size_t smem_size = sizeof(SharedMemoryPlan); + KU_CUDA_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); + + cutlass::ClusterLaunchParams launch_params = { + dim3(params.s_q * CLUSTER_SIZE, 1, 1), + dim3(NUM_THREADS, 1, 1), + dim3(CLUSTER_SIZE, 1, 1), + smem_size, + params.stream + }; + KU_CUTLASS_CHECK(cutlass::launch_kernel_on_cluster( + launch_params, (void*)kernel, params, tma_params, aux_params + )); +} + + +template +void run_fused_norm_rope_attn_rope_cast_fwd_kernel(const ParamT& params) { + using KernelType = Kernel; + KernelType::run(params); +} + + +} diff --git a/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/kernel.h b/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/kernel.h new file mode 100644 index 00000000..9c980c47 --- /dev/null +++ b/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/kernel.h @@ -0,0 +1,59 @@ +#pragma once + +#include + +#include "kernels/params.h" + +namespace sm100::prefill::fused_norm_rope_attn_rope_cast_fwd::core_attn { + +// Local alias: `kernels/defines.h` calls this type `fp8`; this header needs the explicit name +// (see `csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/config.h` for the same pattern) +using fp8_e4m3 = cutlass::float_e4m3_t; + +// Compile-time configuration for the fused_norm_rope_attn_rope_cast core attention kernel. +// Notes: +// - For Prefill mode, MODEL_TYPE is always V4 since V4 and V41 has no difference. +// - For Decode mode, MODEL_TYPE (the format of `kv`) can be V4 or V41, and EXTRA_MODEL_TYPE (the format of `extra_kv`) +// can be MODEL_TYPE or, for MODEL_TYPE == V41, V41_FP4 (fp4 KV cache). +struct Config { + SparseAttnFwdMode FWD_MODE; + ModelType MODEL_TYPE; // V4 only for prefill; V4 or V41 for decode + ModelType EXTRA_MODEL_TYPE; // Decode only, the format of the extra KV cache. Equals MODEL_TYPE for prefill + uint32_t H_Q; + bool ENABLE_Q_NORM; +}; + +// Parameters for the fused Q-b-norm + Q RoPE + Core Attention Forward (prefill/decoding) + O RoPE + O Cast kernel +template +struct ParamsTemplate : Base { + // Q Norm + bool enable_q_norm; + float rms_norm_eps; + + // Q/O RoPE + uint32_t* __restrict__ token_positions; // [s_q] + bool is_rope_neox_style; // Must be false + uint32_t rope_dim; // Must be 64 + float* __restrict__ cos_sin_cache; // [*, rope_dim], must be contiguous, from vllm.RotaryEmbedding.cos_sin_cache + + // O Cast + uint32_t n_wv_group; + uint32_t wv_group_size; + uint32_t num_per_channels; // Must be 128 (for v4) or 32 (for v4.1) + bool use_tma_aligned_col_major_sf; // Must be true + bool round_sf; // Must be true + bool use_packed_ue8m0; // Must be true + + // Output + fp8_e4m3* __restrict__ out_fp8; // [s_q, n_wv_group, wv_group_size * d_v] + uint32_t* __restrict__ out_sf; // [s_q, n_wv_group, (wv_group_size*d_v) / 32 / 4], contiguous on the first (s_q) dim + uint32_t stride_out_sf_wv_group, stride_out_sf_head_dim; +}; + +template +using ParamT = std::conditional_t, ParamsTemplate, ParamsTemplate>; + +template +void run_fused_norm_rope_attn_rope_cast_fwd_kernel(const ParamT& params); + +} diff --git a/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/permute_q_b_proj/kernel.cu b/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/permute_q_b_proj/kernel.cu new file mode 100644 index 00000000..c1e782ec --- /dev/null +++ b/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/permute_q_b_proj/kernel.cu @@ -0,0 +1,71 @@ +/* +Transform q_b_proj layout (include the weight and its scale factor) into layout required by "fused_norm_rope_attn_rope_cast_fwd" kernel + +"fused_norm_rope_attn_rope_cast_fwd" 那边，为了最优秀的性能，我们希望输入的 q 的 layout 长这样： + +[(h0d0 h0d1 ... h0d15) (h1d0 h1d1 ... h1d15) ... (hHd0 ... hHd15)] [(h0d16 ... h0d31) ... (hHd16 ... hHd31)] ... [(h0d496 ... h0d511) ... (hHd496 ... hHdD)] + +where H = (head of q) - 1, D = (headdim of q) - 1 + +因此，我们需要把 q_b_proj 的行（如果假设 q_b_proj 的 shape 是 (H*D) * q_lora_rank 的话）互换。这个 kernel 负责该互换。 + +block dim: 32，一个 warp 负责 permute q_b_proj 的一行 +grid dim: H*D +*/ + +#include "kernel.h" + +#include + +namespace sm100::prefill::fused_norm_rope_attn_rope_cast_fwd::permute_q_b_proj { + +__launch_bounds__(32) +__global__ void permute_q_b_proj_kernel(__grid_constant__ const Params params) { + uint32_t row_idx = blockIdx.x; + uint32_t head_idx = row_idx / params.d_q; + uint32_t head_dim_idx = row_idx % params.d_q; + uint32_t out_row_idx = head_dim_idx % 16 + head_idx * 16 + (head_dim_idx / 16u) * (params.h_q * 16); + uint32_t num_scales_per_row = params.q_lora_rank / params.gran / 4; + + __shared__ uint64_t bar_storage; + ku::transac_bar_t& bar = *(ku::transac_bar_t*)&bar_storage; + if (cute::elect_one_sync()) { + bar.init(1); + cutlass::arch::fence_barrier_init(); + } + __syncthreads(); + + CUTE_ALIGNAS(1024) extern __shared__ fp8_e4m3 row_buf[]; + if (cute::elect_one_sync()) { + cute::SM90_BULK_COPY_G2S::copy( + params.q_b_proj + row_idx * params.stride_q_b_dim0, + (uint64_t*)&bar, + row_buf, + params.q_lora_rank + ); + } + + for (uint32_t scale_idx = threadIdx.x; scale_idx < num_scales_per_row; scale_idx += 32) { + auto cur_scale = params.scale_factors[row_idx + scale_idx * params.stride_scale_factors_dim1]; + params.scale_factors_permuted[out_row_idx + scale_idx * params.stride_scale_factors_permuted_dim1] = cur_scale; + } + + if (cute::elect_one_sync()) { + bar.arrive_and_expect_tx(params.q_lora_rank); + bar.wait(0); + cutlass::arch::fence_view_async_shared(); + cute::SM90_BULK_COPY_S2G::copy( + row_buf, + params.q_b_proj_permuted + out_row_idx * params.stride_q_b_permuted_dim0, + params.q_lora_rank + ); + } +} + +void run_permute_q_b_proj_kernel(const Params& params) { + uint32_t smem_size = params.q_lora_rank; + KU_CUDA_CHECK(cudaFuncSetAttribute(permute_q_b_proj_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); + permute_q_b_proj_kernel<<>>(params); +} + +} diff --git a/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/permute_q_b_proj/kernel.h b/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/permute_q_b_proj/kernel.h new file mode 100644 index 00000000..ec2702db --- /dev/null +++ b/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/permute_q_b_proj/kernel.h @@ -0,0 +1,35 @@ +#pragma once + +#include + +#include "kernels/params.h" + +namespace sm100::prefill::fused_norm_rope_attn_rope_cast_fwd::permute_q_b_proj { + +// Local alias: `kernels/defines.h` calls this type `fp8`; this header needs the explicit name +using fp8_e4m3 = cutlass::float_e4m3_t; + +struct Params { + uint32_t h_q; // Number of q heads + uint32_t d_q; // Q head dimension + uint32_t q_lora_rank; // Q LoRA rank (1024 for V4 Flash, 1536 for V4 Pro) + uint32_t gran; // Scale granularity, 32 or 128 + + // Input tensors + fp8_e4m3* __restrict__ q_b_proj; // [h_q*d_q, q_lora_rank], contiguous on the last dim + uint64_t stride_q_b_dim0; + int32_t* __restrict__ scale_factors; // [h_q*d_q, q_lora_rank/gran/4], contiguous on the FIRST dim (DeepGeMM's format) + uint64_t stride_scale_factors_dim1; + + // Output tensors + fp8_e4m3* __restrict__ q_b_proj_permuted; // The same shape as q_b_proj + uint64_t stride_q_b_permuted_dim0; + int32_t* __restrict__ scale_factors_permuted; // The same shape as scale_factors + uint64_t stride_scale_factors_permuted_dim1; + + cudaStream_t stream; +}; + +void run_permute_q_b_proj_kernel(const Params& params); + +} diff --git a/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/permute_wv_proj/kernel.cu b/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/permute_wv_proj/kernel.cu new file mode 100644 index 00000000..beed958c --- /dev/null +++ b/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/permute_wv_proj/kernel.cu @@ -0,0 +1,106 @@ +/* +Transform wv_proj layout (include the weight and its scale factor) into layout required by "fused_norm_rope_attn_rope_cast_fwd" kernel + +"fused_norm_rope_attn_rope_cast_fwd" 那边，为了最优秀的性能，我们输出的 o 的 layout 长这样： + +- 首先，我们只关注一个 token 对应的 o，其 shape 为 [h_o, d_o] +- 将 head 按照 wv group 分组，其 shape 变成 [n_wv_group, wv_group_size, d_o] +- 对于每个 wv group（此时 shape 为 [wv_group_size, d_o]），输出为 + [(h0d0 h0d1 ... h0d31) (h1d0 h1d1 ... h1d31) ... (hGd0 ... hGd31)] [(h0d32 ... h0d63) ... (hGd32 ... hGd63)] ... [(h0d480 ... h0d511) ... (hGd480 ... hGdD)] +- 这个操作相当于把 head dim 上的每 32 个元素打包后，进行转置操作 + +where G = wv_group_size - 1, D = (headdim of o) - 1 + +因此，我们需要对 wv_proj 中的每个 wv_group 的权重分别处理。对于一个 wv group 的权重，把 wv_proj 的不同列互换。这个 kernel 负责该互换。 + +关于 scale factor：由于这一次我们在 permute input channel（而不是像 q_b 一样 permute output channel），且 permute 粒度为 32 + +block dim: 32，一个 warp 负责变换 wv proj weight 的某个 wv_group 的一整行（如果认为 wv_proj 的 shape 是 [n_wv_group, d_proj_out, wv_group_size * d_o] 的话） +grid dim: (d_proj_out, n_wv_group) +*/ + +#include "kernel.h" + +#include + +namespace sm100::prefill::fused_norm_rope_attn_rope_cast_fwd::permute_wv_proj { + +static constexpr uint32_t INPUT_GRAN = 32; +static constexpr uint32_t CHUNK_SIZE = 32; // 沿着 d_o 方向，每 32 个元素为一个 chunk +static constexpr uint32_t D_O = 512; + +__launch_bounds__(32) +__global__ void permute_wv_proj_kernel(__grid_constant__ const Params params) { + uint32_t row_idx = blockIdx.x; + uint32_t wv_group_idx = blockIdx.y; + uint32_t row_size = params.wv_group_size * D_O; + + __shared__ uint64_t bar_storage; + ku::transac_bar_t& bar = *(ku::transac_bar_t*)&bar_storage; + if (cute::elect_one_sync()) { + bar.init(1); + cutlass::arch::fence_barrier_init(); + } + __syncthreads(); + + // Load the row + CUTE_ALIGNAS(1024) extern __shared__ fp8_e4m3 smem_buf[]; + fp8_e4m3* row_buf = smem_buf; // [row_size] + fp8_e4m3* permuted_row_buf = smem_buf + row_size; // [row_size] + if (cute::elect_one_sync()) { + cute::SM90_BULK_COPY_G2S::copy( + params.wv_proj + wv_group_idx * params.stride_wv_proj_dim0 + row_idx * params.stride_wv_proj_dim1, + &bar_storage, + row_buf, + row_size + ); + bar.arrive_and_expect_tx(row_size); + } + + // Permute SF + uint32_t num_sf_per_head = D_O / 32; // The output granularity is fixed to 32 + uint32_t num_sf = params.wv_group_size * num_sf_per_head; + for (uint32_t i = threadIdx.x; i < num_sf; i += 32) { + uint32_t head_idx = i / num_sf_per_head; + uint32_t head_dim_chunk_idx = i % num_sf_per_head; + uint32_t input_sf_idx_in_row = i; + uint8_t cur_sf = *((uint8_t*)(params.scale_factors + wv_group_idx * params.stride_scale_factors_dim0 + row_idx + (input_sf_idx_in_row / 4) * params.stride_scale_factors_dim2) + input_sf_idx_in_row % 4); + uint32_t output_sf_idx_in_row = head_idx + head_dim_chunk_idx * params.wv_group_size; + *((uint8_t*)(params.scale_factors_permuted + wv_group_idx * params.stride_scale_factors_permuted_dim0 + row_idx + (output_sf_idx_in_row / 4) * params.stride_scale_factors_permuted_dim2) + output_sf_idx_in_row % 4) = cur_sf; + } + + // Wait for the row to be ready, and permute the row + bar.wait(0); + for (uint32_t i = threadIdx.x; i < params.wv_group_size * (D_O / CHUNK_SIZE); i += 32) { + fp8_e4m3 data[CHUNK_SIZE]; + *(__int128_t*)(data + 0) = ku::ld_shared(row_buf + i * CHUNK_SIZE); + *(__int128_t*)(data + 16) = ku::ld_shared(row_buf + i * CHUNK_SIZE + 16); + uint32_t head_idx = i / (D_O / CHUNK_SIZE); + uint32_t head_dim_chunk_idx = i % (D_O / CHUNK_SIZE); + uint32_t chunk_idx_in_permuted_row = head_dim_chunk_idx * params.wv_group_size + head_idx; + ku::st_shared(permuted_row_buf + chunk_idx_in_permuted_row * CHUNK_SIZE + 0, *(__int128_t*)(data + 0)); + ku::st_shared(permuted_row_buf + chunk_idx_in_permuted_row * CHUNK_SIZE + 16, *(__int128_t*)(data + 16)); + } + + // Store the row + cutlass::arch::fence_view_async_shared(); + __syncthreads(); + if (cute::elect_one_sync()) { + cute::SM90_BULK_COPY_S2G::copy( + permuted_row_buf, + params.wv_proj_permuted + wv_group_idx * params.stride_wv_proj_permuted_dim0 + row_idx * params.stride_wv_proj_permuted_dim1, + row_size + ); + } +} + +void run_permute_wv_proj_kernel(const Params& params) { + KU_ASSERT(params.d_o == D_O); + KU_ASSERT(params.input_gran == INPUT_GRAN); + KU_ASSERT(D_O % (params.input_gran * 4) == 0); + uint32_t smem_size = 2 * params.wv_group_size * params.d_o; + KU_CUDA_CHECK(cudaFuncSetAttribute(permute_wv_proj_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); + permute_wv_proj_kernel<<>>(params); +} + +} diff --git a/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/permute_wv_proj/kernel.h b/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/permute_wv_proj/kernel.h new file mode 100644 index 00000000..7a25f79e --- /dev/null +++ b/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/permute_wv_proj/kernel.h @@ -0,0 +1,36 @@ +#pragma once + +#include + +#include "kernels/params.h" + +namespace sm100::prefill::fused_norm_rope_attn_rope_cast_fwd::permute_wv_proj { + +// Local alias: `kernels/defines.h` calls this type `fp8`; this header needs the explicit name +using fp8_e4m3 = cutlass::float_e4m3_t; + +struct Params { + uint32_t d_o; // O head dimension + uint32_t input_gran; // Scale granularity for input, must be 32 + uint32_t wv_group_size; // Number of O heads per WV group. Always equal to 8 for V4 pro/flash + uint32_t n_wv_group; // = h_o / wv_group_size + uint32_t d_proj_out; // dimension of wv_proj's output. Equal to vLLM's V33Attention.o_head_dim + + // Input tensors + fp8_e4m3* __restrict__ wv_proj; // [n_wv_group, d_proj_out, wv_group_size*d_o], contiguous on the last dim + uint64_t stride_wv_proj_dim0, stride_wv_proj_dim1; + int32_t* __restrict__ scale_factors; // [n_wv_group, d_proj_out, wv_group_size*d_o/gran/4], contiguous on the SECOND dim (DeepGeMM's format) + uint64_t stride_scale_factors_dim0, stride_scale_factors_dim2; + + // Output tensors + fp8_e4m3* __restrict__ wv_proj_permuted; // The same shape as wv_proj + uint64_t stride_wv_proj_permuted_dim0, stride_wv_proj_permuted_dim1; + int32_t* __restrict__ scale_factors_permuted; // [n_wv_group, d_proj_out, wv_group_size * d_o / 32 / 4] - The granularity is fixed to 32 + uint64_t stride_scale_factors_permuted_dim0, stride_scale_factors_permuted_dim2; + + cudaStream_t stream; +}; + +void run_permute_wv_proj_kernel(const Params& params); + +} diff --git a/csrc/sm100/prefill/sparse/fwd/head128/config.h b/csrc/kernels/sm100/prefill/sparse/fwd/head128/config.h similarity index 97% rename from csrc/sm100/prefill/sparse/fwd/head128/config.h rename to csrc/kernels/sm100/prefill/sparse/fwd/head128/config.h index 6c846bb4..ae168063 100644 --- a/csrc/sm100/prefill/sparse/fwd/head128/config.h +++ b/csrc/kernels/sm100/prefill/sparse/fwd/head128/config.h @@ -4,10 +4,10 @@ #include #include -#include "params.h" -#include "defines.h" +#include "kernels/params.h" +#include "kernels/defines.h" -namespace sm100::fwd::head128 { +namespace sm100::prefill::sparse_fwd::head128 { using namespace cute; @@ -21,10 +21,6 @@ struct TmaParams { CUtensorMap tensor_map_kv; }; -struct float2x2 { - float2 lo, hi; -}; - template struct KernelTemplate { @@ -38,7 +34,6 @@ static constexpr int B_TOPK = 128; // For 2 CTAs static constexpr int NUM_BUFS = 2; static constexpr int NUM_THREADS = 256 + 128 + 128; // 128 scale & exp threads, 128x2 TMA threads, 32 UTCMMA threads - static constexpr int D_tQ = 384, NUM_tQ_TILES = D_tQ / 64; static constexpr int D_sQ = D_QK-D_tQ, NUM_sQ_TILES = D_sQ / 64; static_assert(D_sQ%64 == 0 && D_tQ%64 == 0 && D_sQ + D_tQ == D_Q); @@ -103,7 +98,6 @@ struct SharedMemoryPlan { array_aligned> o; } u; array_aligned>> s; - float p[(B_H/2)*B_TOPK]; char is_k_valid[NUM_BUFS][B_TOPK/8]; transac_bar_t bar_prologue_q, bar_prologue_utccp; transac_bar_t bar_qk_part_done[NUM_BUFS], bar_qk_done[NUM_BUFS]; // Pi = QKi^T done (i.e. Ki free) diff --git a/csrc/kernels/sm100/prefill/sparse/fwd/head128/instantiations/phase1_k512.cu b/csrc/kernels/sm100/prefill/sparse/fwd/head128/instantiations/phase1_k512.cu new file mode 100644 index 00000000..7b4ebb17 --- /dev/null +++ b/csrc/kernels/sm100/prefill/sparse/fwd/head128/instantiations/phase1_k512.cu @@ -0,0 +1,8 @@ +#include "../phase1.h" +#include "../phase1.cuh" + +namespace sm100::prefill::sparse_fwd::head128 { + +template void run_sparse_fwd_phase1_kernel<512>(const SparseAttnFwdParams& params); + +} diff --git a/csrc/kernels/sm100/prefill/sparse/fwd/head128/instantiations/phase1_k576.cu b/csrc/kernels/sm100/prefill/sparse/fwd/head128/instantiations/phase1_k576.cu new file mode 100644 index 00000000..697aac84 --- /dev/null +++ b/csrc/kernels/sm100/prefill/sparse/fwd/head128/instantiations/phase1_k576.cu @@ -0,0 +1,8 @@ +#include "../phase1.h" +#include "../phase1.cuh" + +namespace sm100::prefill::sparse_fwd::head128 { + +template void run_sparse_fwd_phase1_kernel<576>(const SparseAttnFwdParams& params); + +} diff --git a/csrc/sm100/prefill/sparse/fwd/head128/phase1.cuh b/csrc/kernels/sm100/prefill/sparse/fwd/head128/phase1.cuh similarity index 96% rename from csrc/sm100/prefill/sparse/fwd/head128/phase1.cuh rename to csrc/kernels/sm100/prefill/sparse/fwd/head128/phase1.cuh index ec1192b9..72d15ddc 100644 --- a/csrc/sm100/prefill/sparse/fwd/head128/phase1.cuh +++ b/csrc/kernels/sm100/prefill/sparse/fwd/head128/phase1.cuh @@ -1,3 +1,18 @@ +/* +Sparse Attention Forward Pass (Phase 1) — SM100, h_q=128 + +Forward attention kernel specialized for 128 query heads. Uses UTCMMA instructions +and Tensor Memory (TMEM). Two-CTA cluster design: B_H=128 split across 2 CTAs, +B_TOPK=128 split across 2 CTAs. The pipeline interleaves KV copy (TMA), Q*K^T MMA, +and softmax scale/exp operations. + +Template parameters: + D_QK — Head dimension for QK (512 or 576) + +Grid: [2*s_q, 1, 1], Cluster: [2, 1, 1] + +I/O: See SparseAttnFwdParams in kernels/params.h +*/ #pragma once #include "phase1.h" @@ -8,13 +23,14 @@ #include #include -#include "params.h" -#include "utils.h" -#include "sm100/helpers.h" +#include "kernels/params.h" +#include "kernels/utils.h" +#include "kernels/sm100/helpers.h" +#include "kernels/sm100/common_subroutine.h" #include "config.h" -namespace sm100::fwd::head128 { +namespace sm100::prefill::sparse_fwd::head128 { using namespace cute; @@ -127,7 +143,8 @@ KernelTemplate::sparse_attn_fwd_kernel_devfunc(const SparseAttnFwdParams & } } - cute::cluster_sync(); // We must add a cluster_sync() here, or TMA from CTA1 may launch before barrier initialization in CTA0 + ku::barrier_cluster_arrive_relaxed(); // We must add a cluster_sync() here, or TMA from CTA1 may launch before barrier initialization in CTA0 + ku::barrier_cluster_wait_acquire(); if (warp_idx == 0) { if (elect_one_sync()) { @@ -148,7 +165,7 @@ KernelTemplate::sparse_attn_fwd_kernel_devfunc(const SparseAttnFwdParams & __syncthreads(); // Wait for TMEM allocation if (warpgroup_idx == 0) { - cutlass::arch::warpgroup_reg_alloc<144>(); + cutlass::arch::warpgroup_reg_alloc<152>(); // Scale & Exp warps // The following three numbers are @@ -164,7 +181,6 @@ KernelTemplate::sparse_attn_fwd_kernel_devfunc(const SparseAttnFwdParams & const float2 scale = float2 {params.sm_scale_div_log2, params.sm_scale_div_log2}; uint128_t* sS_base = (uint128_t*)plan.s.data() + idx_in_warpgroup%64 + 64*((idx_in_warpgroup/64)*8); - float* sP_base = plan.p + idx_in_warpgroup%64*4 + (idx_in_warpgroup/64)*((B_H/2)*(B_TOPK/2)); CUTE_NO_UNROLL for (int k = 0; k < num_k_blocks; ++k) { @@ -200,11 +216,7 @@ KernelTemplate::sparse_attn_fwd_kernel_devfunc(const SparseAttnFwdParams & } // Get rowwise max of Pi - float cur_pi_max = -CUDART_INF_F; - CUTE_UNROLL - for (int i = 0; i < (B_TOPK/2); i += 1) { - cur_pi_max = max(cur_pi_max, p_float[i]); - } + float cur_pi_max = get_max(p_float); cur_pi_max *= params.sm_scale_div_log2; plan.bar_k_valid_free[k%NUM_BUFS].arrive(); @@ -219,7 +231,6 @@ KernelTemplate::sparse_attn_fwd_kernel_devfunc(const SparseAttnFwdParams & // - cur_pi_max, real_mi, and mi is identical within each row (i.e. thread 0+64, 1+65, ...) // - should_scale_o is identical among threads 0~31+64~95; and is identical among threads 32~63+96~127 - // Calc scale factor, and scale li float new_max, scale_for_old; if (!should_scale_o) { @@ -605,7 +616,6 @@ KernelTemplate::sparse_attn_fwd_kernel_devfunc(const SparseAttnFwdParams & } } - #else if (cute::thread0()) { CUTE_INVALID_CONTROL_PATH("This kernel only supports sm100"); @@ -620,7 +630,7 @@ sparse_attn_fwd_kernel(__grid_constant__ const SparseAttnFwdParams params, __gri } template -void run_fwd_phase1_kernel(const SparseAttnFwdParams& params) { +void run_sparse_fwd_phase1_kernel(const SparseAttnFwdParams& params) { static_assert(D_QK == 576 || D_QK == 512); using Kernel = KernelTemplate; @@ -639,7 +649,7 @@ void run_fwd_phase1_kernel(const SparseAttnFwdParams& params) { make_stride(params.stride_q_h_q, _1{}, params.stride_q_s_q) ) ), - (typename Kernel::template SmemLayoutQTiles){} + (typename Kernel::SmemLayoutQTiles){} ); auto shape_O = make_shape(params.h_q, params.d_v, params.s_q); @@ -652,7 +662,7 @@ void run_fwd_phase1_kernel(const SparseAttnFwdParams& params) { make_stride(params.d_v, _1{}, params.h_q*params.d_v) ) ), - (typename Kernel::template SmemLayoutOTiles<1>){} + (typename Kernel::SmemLayoutOTiles<1>){} ); CUtensorMap tensor_map_kv; diff --git a/csrc/kernels/sm100/prefill/sparse/fwd/head128/phase1.h b/csrc/kernels/sm100/prefill/sparse/fwd/head128/phase1.h new file mode 100644 index 00000000..332cd10d --- /dev/null +++ b/csrc/kernels/sm100/prefill/sparse/fwd/head128/phase1.h @@ -0,0 +1,10 @@ +#pragma once + +#include "kernels/params.h" + +namespace sm100::prefill::sparse_fwd::head128 { + +template +void run_sparse_fwd_phase1_kernel(const SparseAttnFwdParams& params); + +} diff --git a/csrc/sm100/prefill/sparse/fwd/head64/config.h b/csrc/kernels/sm100/prefill/sparse/fwd/head64/config.h similarity index 58% rename from csrc/sm100/prefill/sparse/fwd/head64/config.h rename to csrc/kernels/sm100/prefill/sparse/fwd/head64/config.h index 8d6eb77a..6199af03 100644 --- a/csrc/sm100/prefill/sparse/fwd/head64/config.h +++ b/csrc/kernels/sm100/prefill/sparse/fwd/head64/config.h @@ -3,12 +3,18 @@ #include #include -#include "defines.h" +#include "kernels/defines.h" +#include "kernels/params.h" -namespace sm100::fwd::head64 { +namespace sm100::prefill::sparse_fwd::head64 { using namespace cute; +template +struct KernelTemplate { + +static constexpr uint32_t B_H = 64; // Head block size. This kernel only supports h_q == B_H + template< typename Shape_Q_NoPE, typename TMA_Q_NoPE, typename Shape_Q_RoPE, typename TMA_Q_RoPE, @@ -21,31 +27,36 @@ struct TmaParams { CUtensorMap tensor_map_kv_nope; }; -struct float2x2 { - float2 lo, hi; -}; +static_assert(D_QK == 576 || D_QK == 512); -constexpr int D_Q = 576; -constexpr int D_K = 576; -constexpr int D_V = 512; -constexpr float MAX_INIT_VAL = -1e30; // We use this number as the initial value for mi (max logits) to avoid -inf - (-inf) = nan +static constexpr int D_Q = D_QK; +static constexpr int D_K = D_QK; +static constexpr int D_V = 512; +static constexpr float MAX_INIT_VAL = -1e30; // We use this number as the initial value for mi (max logits) to avoid -inf - (-inf) = nan +static constexpr int D_NOPE = D_V; +static constexpr int D_ROPE = D_QK - D_V; +static constexpr bool HAVE_ROPE = D_ROPE > 0; -constexpr int B_H = 64; -constexpr int B_TOPK = 64; -constexpr int NUM_BUFS = 3; -constexpr int NUM_THREADS = 128 + 128 + 128; // 128 scale & exp threads, 128 TMA threads, 32 UTCMMA threads +static constexpr int B_TOPK = 64; +static constexpr int NUM_BUFS = 3; +static constexpr int NUM_THREADS = 128 + 128 + 128; +static constexpr int NUM_WORKER_THREADS = 128 + 128 + 1 + B_TOPK/8+1 + (HAVE_ROPE?64:0); +static constexpr int B_EPI = 64; +static constexpr int B_EPI_SB = 256; // "SB" means SuperBlock +static_assert(D_V % B_EPI_SB == 0); +static_assert(B_EPI_SB % ((128/B_H)*B_EPI) == 0); // Tensor memory columns -namespace tmem_cols { +struct tmem_cols { // 0 ~ 256: output // 256 ~ 400: Q // 400 ~ 464: P - constexpr int O = 0; - constexpr int Q = 256; - constexpr int Q_RoPE = 256 + 128; - constexpr int P = 400; -} + static constexpr int O = 0; + static constexpr int Q = 256; + static constexpr int Q_RoPE = 256 + 128; + static constexpr int P = 400; +}; using SmemLayoutQNoPE = decltype(coalesce(tile_to_shape( UMMA::Layout_K_SW128_Atom{}, @@ -57,16 +68,16 @@ using SmemLayoutQRoPE = decltype(coalesce(tile_to_shape( UMMA::Layout_K_SW64_Atom{}, Shape, Int>{}, Step<_1, _2>{} -), Shape<_1, _1>{})); +), Shape<_1, _1>{})); // TODO Explain why SW64 template using SmemLayoutOTiles = decltype(coalesce(tile_to_shape( UMMA::Layout_K_SW128_Atom{}, - Shape, Int<64*NUM_TILES>>{}, + Shape, Int>{}, Step<_1, _2>{} ), Shape<_1, _1>{})); -using SmemLayoutO = SmemLayoutOTiles<8>; +using SmemLayoutO = SmemLayoutOTiles; template using SmemLayoutKTiles = decltype(coalesce(tile_to_shape( @@ -107,25 +118,14 @@ using SmemLayoutS = decltype(coalesce(tile_to_shape( Step<_1, _2>{} ), Shape<_1, _1>{})); - struct SharedMemoryPlan { union { - struct { - array_aligned> _k_rope_pad; - array_aligned> _k_pad[2]; // So that q_nope covers k[2] - array_aligned> q_nope; - } q_full; - struct { - array_aligned> k_rope; - array_aligned> k_nope[NUM_BUFS]; - } k; - array_aligned> o; + static_assert(B_H <= B_TOPK); + array_aligned qko_slots[NUM_BUFS]; } u; - float p_exchange_buf[4][32 * (B_TOPK/2)]; - union { - bf16 s[B_H*B_TOPK]; - array_aligned> q_rope; - } s_q_rope; + bf16 qk_rope_slot[B_TOPK*(D_Q-D_V)]; + float p_exchange_buf[4][32 * (B_TOPK/(128/B_H))]; + bf16 s[B_H*B_TOPK]; char is_k_valid[NUM_BUFS][B_TOPK/8]; transac_bar_t bar_prologue_q_nope, bar_prologue_q_rope, bar_prologue_utccp_nope, bar_prologue_utccp_rope; transac_bar_t bar_qk_nope_done[NUM_BUFS], bar_qk_rope_done; // Pi = QKi^T (the nope part) done @@ -133,13 +133,21 @@ struct SharedMemoryPlan { transac_bar_t bar_kv_nope_ready[NUM_BUFS][2], bar_kv_rope_ready; transac_bar_t bar_p_free; transac_bar_t bar_so_ready; // S and O are ready + transac_bar_t bar_o_write_back_done; + transac_bar_t bar_o_write_back_done_waited; // Whether the barrier above has been waited for. Used to prevent double arrive transac_bar_t bar_k_valid_ready[NUM_BUFS], bar_k_valid_free[NUM_BUFS]; + transac_bar_t bar_clc_full, bar_clc_empty; array_aligned tmem_start_addr; float rowwise_max_buf[128], rowwise_li_buf[128]; + ku::CLCResponseObj clc_response_obj; }; +static constexpr int QK_MRGEMM_N = (128/B_H)*B_TOPK; +static constexpr int QK_MRGEMM_K_NOPE = D_V / (128/B_H); +static constexpr int QK_MRGEMM_K_ROPE = D_ROPE / (128/B_H); + using TiledMMA_P = decltype(make_tiled_mma( - SM100_MMA_F16BF16_WS_TS_NOELECT{} // Here we use N = 128 = 2*B_TOPK since we're going to use implicit dual gemm: + SM100_MMA_F16BF16_WS_TS_NOELECT{} )); using TiledMMA_O = decltype(make_tiled_mma( @@ -150,8 +158,14 @@ enum NamedBarriers : int { wg0_sync = 0, wg0_warp02_sync = 1, wg0_warp13_sync = 2, - pepi_sync = 3, }; +template +static __device__ void +sparse_attn_fwd_kernel_devfunc(const SparseAttnFwdParams ¶ms, const TmaParam &tma_params); + +static void run(const SparseAttnFwdParams& params); + +}; } diff --git a/csrc/kernels/sm100/prefill/sparse/fwd/head64/instantiations/phase1_h64_k512.cu b/csrc/kernels/sm100/prefill/sparse/fwd/head64/instantiations/phase1_h64_k512.cu new file mode 100644 index 00000000..6ece2bf2 --- /dev/null +++ b/csrc/kernels/sm100/prefill/sparse/fwd/head64/instantiations/phase1_h64_k512.cu @@ -0,0 +1,8 @@ +#include "../phase1.h" +#include "../phase1.cuh" + +namespace sm100::prefill::sparse_fwd::head64 { + +template void run_sparse_fwd_phase1_kernel(const SparseAttnFwdParams& params); + +} diff --git a/csrc/kernels/sm100/prefill/sparse/fwd/head64/instantiations/phase1_h64_k576.cu b/csrc/kernels/sm100/prefill/sparse/fwd/head64/instantiations/phase1_h64_k576.cu new file mode 100644 index 00000000..4b4a999d --- /dev/null +++ b/csrc/kernels/sm100/prefill/sparse/fwd/head64/instantiations/phase1_h64_k576.cu @@ -0,0 +1,8 @@ +#include "../phase1.h" +#include "../phase1.cuh" + +namespace sm100::prefill::sparse_fwd::head64 { + +template void run_sparse_fwd_phase1_kernel(const SparseAttnFwdParams& params); + +} diff --git a/csrc/kernels/sm100/prefill/sparse/fwd/head64/phase1.cuh b/csrc/kernels/sm100/prefill/sparse/fwd/head64/phase1.cuh new file mode 100644 index 00000000..aa5b824b --- /dev/null +++ b/csrc/kernels/sm100/prefill/sparse/fwd/head64/phase1.cuh @@ -0,0 +1,798 @@ +/* +Sparse Attention Forward Pass (Phase 1) — SM100, h_q == 64 + +Forward attention kernel for h_q == 64 query heads. +Uses UTCMMA and TMEM. Single-CTA design with B_TOPK=64. Prefill mode only. + +Template parameters: + FWD_MODE — Forward mode (Prefill) + D_QK — Head dimension for QK (512 or 576) + +Grid: [s_q, 1, 1] + +I/O: See SparseAttnFwdParams in kernels/params.h +*/ +#pragma once +#include "phase1.h" + +#include +#include +#include +#include +#include + +#include + +#include "kernels/params.h" +#include "kernels/utils.h" +#include "kernels/sm100/helpers.h" +#include "kernels/sm100/common_subroutine.h" +#include "config.h" + +namespace sm100::prefill::sparse_fwd::head64 { + +using namespace cute; + +/* +Pipeline Overview: + +| Copy | MMA | Scale & Exp | + +KV0 +KV1 +KV2 + P0 = QK0^T + S0 = exp(P0) + scale(O) w.r.t P0 + P1 = QK1^T + S1 = exp(P1) + O += S0V0 +KV3 scale(O) w.r.t P1 + P2 = QK2^T + S2 = exp(P2) + O += S1V1 +KV4 scale(O) w.r.t P2 + P3 = QK3^T + S3 = exp(P3) + O += S2V2 +KV5 scale(O) w.r.t P3 + +... + + O += S(n-3)V(n-3) + scale(O) w.r.t P(n-2) + P(n-1) = QK(n-1)^T + S(n-1) = exp(P(n-1)) + O += S(n-2)V(n-2) + scale(O) w.r.t P(n-1) + O += S(n-1)V(n-1) +*/ + +using FwdMode = SparseAttnFwdMode; + +template +template +__device__ void +KernelTemplate::sparse_attn_fwd_kernel_devfunc(const SparseAttnFwdParams ¶ms, const TmaParam &tma_params) { +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000 && __CUDA_ARCH__ < 1200)) || (defined(__CLION_IDE__) || defined(__VSCODE_IDE__)) + // Grid shape: [s_q, 1, 1] + + const int warp_idx = cutlass::canonical_warp_idx_sync(); + const int lane_idx = threadIdx.x % 32; + const int warpgroup_idx = __shfl_sync(0xffffffff, threadIdx.x / 128, 0); + const int idx_in_warpgroup = threadIdx.x % 128; + + // Define shared tensors + extern __shared__ char wksp_buf[]; + SharedMemoryPlan &plan = *reinterpret_cast(wksp_buf); + + if (warp_idx == 0 && elect_one_sync()) { + if constexpr (HAVE_ROPE) { + cute::prefetch_tma_descriptor(tma_params.tma_Q_rope.get_tma_descriptor()); + } + cute::prefetch_tma_descriptor(tma_params.tma_Q_nope.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_params.tma_O.get_tma_descriptor()); + cute::prefetch_tma_descriptor(&(tma_params.tensor_map_kv_nope)); + + plan.bar_prologue_q_nope.init(1); + plan.bar_prologue_utccp_nope.init(1); + if constexpr (HAVE_ROPE) { + plan.bar_prologue_q_rope.init(1); + plan.bar_prologue_utccp_rope.init(1); + } + plan.bar_clc_full.init(1); + plan.bar_clc_empty.init(NUM_WORKER_THREADS); + fence_barrier_init(); + } else if (warp_idx == 1 && elect_one_sync()) { + // Initialize other barriers + CUTE_UNROLL + for (int i = 0; i < NUM_BUFS; ++i) { + plan.bar_qk_nope_done[i].init(1); + plan.bar_sv_done[i].init(1); + plan.bar_kv_nope_ready[i][0].init(1); + plan.bar_kv_nope_ready[i][1].init(1); + plan.bar_k_valid_ready[i].init(B_TOPK/8); + plan.bar_k_valid_free[i].init(128); + } + plan.bar_p_free.init(128); + plan.bar_so_ready.init(128); + if constexpr (HAVE_ROPE) { + plan.bar_qk_rope_done.init(1); + plan.bar_kv_rope_ready.init(64); + } + plan.bar_o_write_back_done.init(128); + plan.bar_o_write_back_done_waited.init(4); + fence_barrier_init(); + } else if (warp_idx == 2) { + // Initialize TMEM + cute::TMEM::Allocator1Sm().allocate(512, plan.tmem_start_addr.data()); + TRAP_ONLY_DEVICE_ASSERT(plan.tmem_start_addr.data()[0] == 0); + cute::TMEM::Allocator1Sm().release_allocation_lock(); + } + + __syncthreads(); + + struct OuterloopArgs { + bool outer_loop_phase; + int s_q_idx; + int num_k_blocks; + int topk_length; + }; + + auto issue_q_rope_tma = [&](int s_q_idx) { + if constexpr (HAVE_ROPE) { + Tensor gQ_rope = tma_params.tma_Q_rope.get_tma_tensor(tma_params.shape_Q_rope)(_, _, s_q_idx); + Tensor sQ_rope = make_tensor(make_smem_ptr(plan.qk_rope_slot), SmemLayoutQRoPE{}); + ku::launch_tma_copy(tma_params.tma_Q_rope, gQ_rope, sQ_rope, plan.bar_prologue_q_rope, TMA::CacheHintSm90::EVICT_FIRST); + } + }; + + auto issue_q_rope_utccp = [&](bool outer_loop_phase) { + if constexpr (HAVE_ROPE) { + plan.bar_prologue_q_rope.arrive_and_expect_tx(B_H*(D_Q-D_V)*sizeof(bf16)); + plan.bar_prologue_q_rope.wait(outer_loop_phase); + ku::tcgen05_after_thread_sync(); + + UMMA::SmemDescriptor sQ_rope_desc = UMMA::make_umma_desc( + make_tensor( + make_smem_ptr(plan.qk_rope_slot), + tile_to_shape( + UMMA::Layout_K_SW64_Atom{}, + Shape, Int<32>>{} + ) + ) + ); + + // Copy the RoPE tile: (2*B_H) rows * 32 cols (64B) (in UTCCP's view), or B_H rows * 64 cols (in our view) + // A subtile is (2*B_H) rows * 16 cols (256b, 32B) (in UTCCP's view), or B_H rows * 16 cols * 2 (in our view) + CUTE_UNROLL + for (int subtile_idx = 0; subtile_idx < 2; ++subtile_idx) { + SM100_UTCCP_128dp256bit_1cta::copy( + sQ_rope_desc + (subtile_idx*32) / 16, + tmem_cols::Q_RoPE + subtile_idx*8 + ); + } + ku::umma_arrive_noelect(plan.bar_prologue_utccp_rope); + } + }; + + auto issue_q_nope_tma = [&](int s_q_idx, int qko_slot_idx) { + Tensor gQ_nope = tma_params.tma_Q_nope.get_tma_tensor(tma_params.shape_Q_nope)(_, _, s_q_idx); + Tensor sQ_nope = make_tensor(make_smem_ptr(plan.u.qko_slots[qko_slot_idx].data()), SmemLayoutQNoPE{}); + ku::launch_tma_copy(tma_params.tma_Q_nope, gQ_nope, sQ_nope, plan.bar_prologue_q_nope, TMA::CacheHintSm90::EVICT_FIRST); + }; + + auto issue_q_nope_utccp = [&](bool outer_loop_phase, int qko_slot_idx) { + plan.bar_prologue_q_nope.arrive_and_expect_tx(B_H*D_V*sizeof(bf16)); + plan.bar_prologue_q_nope.wait(outer_loop_phase); + ku::tcgen05_after_thread_sync(); + UMMA::SmemDescriptor sQ_nope_desc = UMMA::make_umma_desc( + make_tensor( + make_smem_ptr(plan.u.qko_slots[qko_slot_idx].data()), + tile_to_shape( + UMMA::Layout_K_SW128_Atom{}, + Shape, Int<64>>{} // TODO Explain this layout and dual gemm + ) + ) + ); + + CUTE_UNROLL + for (int tile_idx = 0; tile_idx < D_V/64/2; ++tile_idx) { + // A tile is (2*B_H) rows * 64 cols (128B) (in UTCCP's view), or B_H rows * 128 cols (in our view) + CUTE_UNROLL + for (int subtile_idx = 0; subtile_idx < 4; ++subtile_idx) { + // A subtile is 128 rows * 16 cols (256b, 32B) (in UTCCP's view), or B_H rows * 16 cols * 2 (in our view) + SM100_UTCCP_128dp256bit_1cta::copy( + sQ_nope_desc + (tile_idx*(B_H*128*2) + subtile_idx*32) / 16, // Remember that 4 LSBs are not included + tmem_cols::Q + tile_idx*32 + subtile_idx*8 + ); + } + } + ku::umma_arrive_noelect(plan.bar_prologue_utccp_nope); + }; + + auto run_outer_loop = [&](auto loop_body) { + int outer_loop_phase = false; + ku::CLCResult next_job = {true, (int)blockIdx.x, 0, 0}; + CUTE_NO_UNROLL + while (next_job.is_valid) { + int s_q_idx = next_job.x; + int topk_length = params.topk_length != nullptr ? __ldg(params.topk_length + s_q_idx) : params.topk; + int num_k_blocks = max(cute::ceil_div(topk_length, (int)B_TOPK), 2); // num_k_blocks always >= 2 to simplify synchronizations across outer loop boundries + OuterloopArgs args = { + (bool)outer_loop_phase, + s_q_idx, + num_k_blocks, + topk_length + }; + loop_body(args); + + plan.bar_clc_full.wait(outer_loop_phase); + next_job = ku::get_clc_query_response(plan.clc_response_obj); + outer_loop_phase ^= 1; + + plan.bar_clc_empty.arrive(); + } + }; + + RingBufferState rs; + + if (warpgroup_idx == 0) { + // Scale & Exp warps + bf16* sS_base = plan.s + lane_idx*8 + (warp_idx&1)*(B_H/2)*8 + (warp_idx/2)*B_H*(B_TOPK/2); + static constexpr int NUM_ELEMS_PER_THREAD = B_TOPK / 2; + + run_outer_loop([&](const OuterloopArgs &args) { + // The following three numbers are + // - mi: max_logits used to scale Pi (i.e. O := exp2(Pi*scale - mi) @ V) + // - li: sumexp, i.e. li := sum(exp(Pi*scale - mi)) + // - real_mi: real max logits, i.e. real_mi := max(Pi*scale) + // where Pi is the i-th row of P, P := QK^T + // mi and real_mi are always consistent within the two threads that + // controls one row (i.e. thread 0+64, 1+65, 2+66, ...) after every update + float mi = MAX_INIT_VAL; + float li = 0.0f; + float real_mi = -CUDART_INF_F; + + cute::tma_store_wait<0>(); + plan.bar_o_write_back_done.arrive(); + + CUTE_NO_UNROLL + for (int k = 0; k < args.num_k_blocks; ++k) { + // Wait for P + NamedBarrier::arrive_and_wait(64, NamedBarriers::wg0_warp02_sync+(warp_idx&1)); + auto [buf_idx, bar_phase] = rs.get(); + plan.bar_qk_nope_done[buf_idx].wait(bar_phase); + plan.bar_k_valid_ready[buf_idx].wait(bar_phase); // Put the barrier wait here for more code reordering space + ku::tcgen05_after_thread_sync(); + + // Load P + float p[NUM_ELEMS_PER_THREAD]; + retrieve_mask_and_reduce_p< + NUM_ELEMS_PER_THREAD, + NamedBarriers::wg0_warp02_sync, + NamedBarriers::wg0_warp13_sync, + false // Prefill keeps P in registers (no store_back_p) + >( + tmem_cols::P, + plan.is_k_valid[buf_idx], + warp_idx, lane_idx, + [&]() {plan.bar_p_free.arrive();}, + plan.p_exchange_buf, + p + ); + plan.bar_k_valid_free[buf_idx].arrive(); + + // Get rowwise max of Pi + 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); + 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, ...) + // - should_scale_o is identical among every warp, and is identical among threads that controls the same row (i.e. among threads 0~31+64~95; and is identical among threads 32~63+96~127) + + // 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]; + float cur_sum = get_s_from_p(s, p, params.sm_scale_div_log2, new_max); + li = fma(li, scale_for_old, cur_sum); + + // Wait for last SV gemm, write S + if (k > 0) { + auto [last_sv_buf_idx, last_sv_bar_phase] = rs.offset_by(-1).get(); + plan.bar_sv_done[last_sv_buf_idx].wait(last_sv_bar_phase); + } + 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 (k > 0 && should_scale_o) { + // plan.bar_sv_done[(k-1)%NUM_BUFS].wait(((k-1)/NUM_BUFS)&1); // NOTE We have waited for last SV gemm before + ku::tcgen05_after_thread_sync(); + rescale_O(scale_for_old); + ku::tcgen05_before_thread_sync(); + } + + fence_view_async_shared(); + plan.bar_so_ready.arrive(); + + rs.update(); + } + + plan.bar_o_write_back_done_waited.wait(args.outer_loop_phase); + + // Epilogue + 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_li_buf[idx_in_warpgroup] = li; + NamedBarrier::arrive_and_wait(128, NamedBarriers::wg0_sync); + li += plan.rowwise_li_buf[idx_in_warpgroup^64]; + + // Store mi and li + if (idx_in_warpgroup < B_H) { + bool is_padding_row = idx_in_warpgroup >= params.h_q; + float cur_lse = fmaf(mi, CUDART_LN2_F, logf(li)); + cur_lse = cur_lse == -CUDART_INF_F ? +CUDART_INF_F : cur_lse; + if (!is_padding_row) { + int global_index = args.s_q_idx*params.h_q + idx_in_warpgroup; + params.max_logits[global_index] = real_mi*CUDART_LN2_F; + params.lse[global_index] = cur_lse; + } + } + + auto [o_slot_idx, __] = rs.offset_by(+2).get(); + + // Store O + float attn_sink = (params.attn_sink == nullptr || idx_in_warpgroup%B_H >= params.h_q) + ? -CUDART_INF_F : __ldg(params.attn_sink + (idx_in_warpgroup%B_H))*CUDART_L2E_F; + float output_scale = __fdividef(1.0f, li + exp2f(attn_sink - mi)); + Tensor sO = make_tensor(make_smem_ptr(plan.u.qko_slots[o_slot_idx].data()), SmemLayoutO{}); + Tensor tma_gO = flat_divide( + tma_params.tma_O.get_tma_tensor(tma_params.shape_O)(_, _, args.s_q_idx), + Shape, Int>{} + )(_, _, _0{}, _); + Tensor sO_divided = flat_divide( + sO, + Shape, Int>{} + )(_, _, _0{}, _); + auto thr_tma = tma_params.tma_O.get_slice(_0{}); + + float2 o[B_EPI/2]; + bool have_valid_indices = __any_sync(0xffffffff, li != 0); // Prevent some threads' li == 0 and some threads' li != 0 which lead to deadlock during ku::tmem_ld + if (!have_valid_indices) { + // If there are no valid indices, we set o[i] to 0 and don't load from TMEM + CUTE_UNROLL + for (int i = 0; i < B_EPI/2; ++i) + o[i].x = o[i].y = 0.0f; + output_scale = 1.0f; + } + + float2 output_scale_float2 = make_float2(output_scale, output_scale); + + bf16* sO_addrs[8]; + CUTE_UNROLL + for (int i = 0; i < B_EPI/8; ++i) { + sO_addrs[i] = &sO(idx_in_warpgroup%B_H, i*8); + } + + // Wait for the last GEMM + { + auto [last_sv_buf_idx, last_sv_bar_phase] = rs.offset_by(-1).get(); + plan.bar_sv_done[last_sv_buf_idx].wait(last_sv_bar_phase); + ku::tcgen05_after_thread_sync(); + } + + static constexpr int NUM_EPI_SB = D_V/B_EPI_SB; + static constexpr int NUM_TMA_PARTS = 2; + CUTE_UNROLL + for (int c = 0; c < NUM_EPI_SB; ++c) { + // Each tile: B_H x B_EPI_SB + CUTE_UNROLL + for (int k = 0; k < B_EPI_SB/B_EPI/NUM_TMA_PARTS; ++k) { + // Load O from tO + if (have_valid_indices) { + ku::tmem_ld_32dp32bNx(tmem_cols::O + c*(B_EPI_SB/NUM_TMA_PARTS) + k*B_EPI, o); + cutlass::arch::fence_view_async_tmem_load(); + } + // NOTE. We neither signal any barrier after tmem_O is free, nor do we wait for any barrier in the UTCMMA warp, since the first O gemm in the next round depends on S, which depends on this warpgroup + + // Convert and store + CUTE_UNROLL + for (int i = 0; i < B_EPI/8; ++i) { + nv_bfloat162 o_bf16[4]; + CUTE_UNROLL + for (int j = 0; j < 4; ++j) { + o[i*4+j] = ku::float2_mul(o[i*4+j], output_scale_float2); + o_bf16[j] = __float22bfloat162_rn(o[i*4+j]); + } + bf16* o_smem_ptr = sO_addrs[i] + (c*B_EPI_SB + (idx_in_warpgroup/B_H)*(B_EPI_SB/(128/B_H)) + k*B_EPI)*B_H; + ku::st_shared(o_smem_ptr, *(__int128_t*)(o_bf16)); + } + + // Sync + fence_view_async_shared(); + NamedBarrier::arrive_and_wait(128, NamedBarriers::wg0_sync); + + // Store into global memory + if (warp_idx < NUM_TMA_PARTS && elect_one_sync()) { + int epi_chunk_idx = c*(B_EPI_SB/B_EPI) + ((B_EPI_SB/B_EPI)/NUM_TMA_PARTS)*warp_idx + k; + cute::copy( + tma_params.tma_O, + thr_tma.partition_S(sO_divided(_, _, epi_chunk_idx)), + thr_tma.partition_D(tma_gO(_, _, epi_chunk_idx)) + ); + } + } + } + cute::tma_store_arrive(); + }); + + if (warp_idx == 3) { + cute::TMEM::Allocator1Sm().free(0, 512); + } + } else if (warpgroup_idx == 1) { + // Producer warp for KV + int warp_idx = cutlass::canonical_warp_idx_sync() - 4; + constexpr int NUM_WARPS = 4, NUM_LOCAL_ROWS_PER_WARP = (B_TOPK/4)/NUM_WARPS; + run_outer_loop([&](const OuterloopArgs &args) { + if (elect_one_sync()) { + int* gIndices = params.indices + args.s_q_idx*params.stride_indices_s_q; // [topk] + CUTE_NO_UNROLL + for (int k = 0; k < args.num_k_blocks; ++k) { + int4 indices[NUM_LOCAL_ROWS_PER_WARP]; + int max_indices = -1, min_indices = params.s_kv; + CUTE_UNROLL + for (int local_row = 0; local_row < NUM_LOCAL_ROWS_PER_WARP; ++local_row) { + indices[local_row] = __ldg((int4*)(gIndices + k*B_TOPK) + local_row*NUM_WARPS + warp_idx); + max_indices = max(max_indices, int4_max(indices[local_row])); + min_indices = min(min_indices, int4_min(indices[local_row])); + } + bool is_all_rows_invalid = min_indices == params.s_kv || max_indices == -1; + bool should_skip_tma = is_all_rows_invalid && k >= NUM_BUFS; // Don't skip TMA for the first NUM_BUFS turns to swipe out invalid values in kv buffer + + if (k == 1) { + // Since q_nope coincidences with k["buffer idx of the 1st block"] + plan.bar_prologue_utccp_nope.wait(args.outer_loop_phase); + } else if (k == 2) { + // Since o_buf coincidences with k["buffer idx of the 2nd block"] + plan.bar_o_write_back_done.wait(args.outer_loop_phase); + } + + // Copy NoPE + auto [buf_idx, bar_phase] = rs.get(); + plan.bar_sv_done[buf_idx].wait(bar_phase^1); + bf16* sK_nope_base = plan.u.qko_slots[buf_idx].data() + warp_idx*4*64; + + auto load_kv_nope_part = [&](int part_idx) { + CUTE_UNROLL + for (int local_row = 0; local_row < NUM_LOCAL_ROWS_PER_WARP; ++local_row) { + CUTE_UNROLL + for (int local_col = part_idx*(D_V/2/64); local_col < (part_idx+1)*(D_V/2/64); ++local_col) { + ku::tma_gather4( + &(tma_params.tensor_map_kv_nope), + plan.bar_kv_nope_ready[buf_idx][part_idx], + sK_nope_base + local_row*(4*NUM_WARPS)*64 + local_col*(B_TOPK*64), + local_col*64, + indices[local_row], + (int64_t)TMA::CacheHintSm90::EVICT_LAST + ); + } + } + }; + + if (!should_skip_tma) { + load_kv_nope_part(0); + load_kv_nope_part(1); + } else { + // NOTE See head128/phase1.cuh for this TMA skipping technique + CUTE_UNROLL + for (int part_idx = 0; part_idx < 2; ++part_idx) + plan.bar_kv_nope_ready[buf_idx][part_idx].complete_transaction(NUM_LOCAL_ROWS_PER_WARP*4*D_V/2*sizeof(bf16)); + } + + rs.update(); + } + if (args.num_k_blocks <= 2) { + plan.bar_o_write_back_done.wait(args.outer_loop_phase); + } + plan.bar_o_write_back_done_waited.arrive(); + } + __syncwarp(); + }); + } else { + // MMA warp + if (warp_idx == 8 && elect_one_sync()) { + // Allocate tmem tensors + TiledMMA tiled_mma_P = TiledMMA_P{}; + TiledMMA tiled_mma_O = TiledMMA_O{}; + // NOTE These tXXX tensors are only for a forged layout (so that CuTe is able to generate correct address in cute::gemm) + Tensor tP = partition_fragment_C(tiled_mma_P, Shape, Int>{}); + Tensor tQ_nope_part0 = tiled_mma_P.get_slice(_0{}).make_fragment_A( + partition_shape_A(tiled_mma_P, Shape, Int>{}) + ); + Tensor tQ_nope_part1 = tiled_mma_P.get_slice(_0{}).make_fragment_A( + partition_shape_A(tiled_mma_P, Shape, Int>{}) + ); + Tensor tQ_rope = tiled_mma_P.get_slice(_0{}).make_fragment_A( + partition_shape_A(tiled_mma_P, Shape, Int>{}) + ); + Tensor tO = partition_fragment_C(tiled_mma_O, Shape, Int>{}); + tP.data().get() = tmem_cols::P; + tQ_nope_part0.data().get() = tmem_cols::Q; + tQ_nope_part1.data().get() = tmem_cols::Q + 64; + tQ_rope.data().get() = tmem_cols::Q_RoPE; + tO.data().get() = tmem_cols::O; + + run_outer_loop([&](const OuterloopArgs &args) { + // Copy Q into k["buffer idx of the 1st block"+1] + // NOTE. As we reach here, we must have already issued the last O gemm of the last round, which means that the penultimate O gemm must be finished (since the last O gemm |-> bar_so_ready (the last S is ready) |-> the penultimate O gemm must be finished). So the corresponding K buffer must be free. + // For the RoPE part, as we reach here, the last S from the last round must be ready, which means that the corresponding K RoPE buffer must be free. + int q_slot_idx = rs.offset_by(+1).get().first; + issue_q_rope_tma(args.s_q_idx); + issue_q_nope_tma(args.s_q_idx, q_slot_idx); + issue_q_rope_utccp(args.outer_loop_phase); + issue_q_nope_utccp(args.outer_loop_phase, q_slot_idx); + // NOTE Here we don't need to wait for bar_prologue_utccp_rope, since the copy-in of the first RoPE relies on bar_prologue_utccp_rope + + CUTE_NO_UNROLL + for (int k = 0; k < args.num_k_blocks+1; ++k) { + if (k < args.num_k_blocks) { + // Pi = QKi^T + auto [buf_idx, bar_phase] = rs.get(); + Tensor sK_nope = make_tensor(make_smem_ptr(plan.u.qko_slots[buf_idx].data()), SmemLayoutKNoPE_TiledMMA{}); + Tensor sK_rope = make_tensor(make_smem_ptr(plan.qk_rope_slot), SmemLayoutKRoPE_TiledMMA{}); + + auto [__, binary_bar_phase] = rs.get<1>(); + plan.bar_p_free.wait(binary_bar_phase^1); + ku::tcgen05_after_thread_sync(); + + // Wait for K (RoPE) + // P = Q(rope) @ K(rope)^T + if constexpr (HAVE_ROPE) { + plan.bar_kv_rope_ready.wait(binary_bar_phase); + ku::tcgen05_after_thread_sync(); + ku::utcmma_ts(tiled_mma_P, tQ_rope, sK_rope, tP, true); + ku::umma_arrive_noelect(plan.bar_qk_rope_done); + } + + // Wait for K (NoPE) + if (k == 0) { + plan.bar_prologue_utccp_nope.wait(args.outer_loop_phase); + } + Tensor sK_nope_divided = flat_divide(sK_nope, Tile, Int>{})(_, _, _0{}, _); + CUTE_UNROLL + for (int kv_nope_part_idx = 0; kv_nope_part_idx < 2; ++kv_nope_part_idx) { + plan.bar_kv_nope_ready[buf_idx][kv_nope_part_idx].arrive_and_expect_tx(B_TOPK*D_V/2*sizeof(bf16)); + plan.bar_kv_nope_ready[buf_idx][kv_nope_part_idx].wait(bar_phase); + ku::tcgen05_after_thread_sync(); + + // P += Q(nope) @ K(nope)^T + bool clear_accum = (!HAVE_ROPE) && kv_nope_part_idx == 0; + ku::utcmma_ts(tiled_mma_P, kv_nope_part_idx ? tQ_nope_part1 : tQ_nope_part0, sK_nope_divided(_, _, kv_nope_part_idx), tP, clear_accum); + } + ku::umma_arrive_noelect(plan.bar_qk_nope_done[buf_idx]); + } + if (k > 0) { + // O += S(i-1)V(i-1) + auto [buf_idx, bar_phase] = rs.offset_by(-1).get(); + + Tensor sS = make_tensor(make_smem_ptr(plan.s), SmemLayoutS{}); + Tensor sV = make_tensor(make_smem_ptr(plan.u.qko_slots[buf_idx].data()), SmemLayoutV{}); + + // Wait for S(i-1) and O to be scaled + auto [__, binary_bar_phase] = rs.offset_by(-1).get<1>(); + plan.bar_so_ready.wait(binary_bar_phase); + ku::tcgen05_after_thread_sync(); + + // O += sS @ sV + ku::utcmma_ss(tiled_mma_O, sS, sV, tO, k == 1); + ku::umma_arrive_noelect(plan.bar_sv_done[buf_idx]); + } + + rs.update(); + } + rs = rs.offset_by(-1); + }); + } else if (warp_idx == 9) { + // KV valid loading + CLC producer warp + if (lane_idx < B_TOPK/8) { + run_outer_loop([&](const OuterloopArgs &args) { + int* gIndices = params.indices + args.s_q_idx*params.stride_indices_s_q; // [topk] + CUTE_NO_UNROLL + for (int k = 0; k < args.num_k_blocks; ++k) { + char k_validness_mask = load_indices_and_generate_mask( + lane_idx, + gIndices + k*B_TOPK, + params.s_kv, + k*B_TOPK, + args.topk_length + ); + + auto [buf_idx, bar_phase] = rs.get(); + plan.bar_k_valid_free[buf_idx].wait(bar_phase^1); + plan.is_k_valid[buf_idx][lane_idx] = k_validness_mask; + plan.bar_k_valid_ready[buf_idx].arrive(); + + rs.update(); + } + }); + } else if (lane_idx == B_TOPK/8) { + run_outer_loop([&](const OuterloopArgs &args) { + plan.bar_clc_empty.wait(args.outer_loop_phase^1); + ku::issue_clc_query(plan.bar_clc_full, plan.clc_response_obj); + plan.bar_clc_full.arrive_and_expect_tx(sizeof(plan.clc_response_obj)); + }); + } + } else if (warp_idx == 10 || warp_idx == 11) { + // RoPE loading warp + if constexpr (HAVE_ROPE) { + int thread_idx = threadIdx.x - 10*32; + constexpr int GROUP_SIZE = 8, NUM_GROUPS = B_H/GROUP_SIZE, ROWS_PER_THREAD = B_TOPK/NUM_GROUPS; + int group_idx = thread_idx / GROUP_SIZE, idx_in_group = thread_idx % GROUP_SIZE; + Tensor sK_rope = make_tensor(make_smem_ptr(plan.qk_rope_slot), SmemLayoutKRoPE{}); + bf16* sK_rope_base = &sK_rope(group_idx, idx_in_group*8); + run_outer_loop([&](const OuterloopArgs &args) { + int* gIndices = params.indices + args.s_q_idx*params.stride_indices_s_q; // [topk] + CUTE_NO_UNROLL + for (int k = 0; k < args.num_k_blocks; ++k) { + auto [_, binary_bar_phase] = rs.get<1>(); + int indices[ROWS_PER_THREAD]; + CUTE_UNROLL + for (int local_row = 0; local_row < ROWS_PER_THREAD; ++local_row) + indices[local_row] = __ldg(gIndices + k*B_TOPK + group_idx + local_row*NUM_GROUPS); + plan.bar_qk_rope_done.wait(binary_bar_phase^1); + if (k == 0) { + plan.bar_prologue_utccp_rope.wait(args.outer_loop_phase); // Wait for Q RoPE's UTCCP so that qk_rope_slot is empty + } + CUTE_UNROLL + for (int local_row = 0; local_row < ROWS_PER_THREAD; ++local_row) { + int index = indices[local_row]; + ku::cp_async_cacheglobal( + params.kv + (int64_t)index*params.stride_kv_s_kv + 512 + idx_in_group*8, + sK_rope_base + local_row*NUM_GROUPS*32, + index >= 0 && index < params.s_kv + ); // NOTE Using cp.async instead of TMA is faster here + // NOTE Here we only consider the range of `index` instead of also checking against topk_length, as it's noted that under this scenario (i.e. there exists a valid index among indices[topk_length: ] that points to a token who has NaN inside) + } + cutlass::arch::cpasync_barrier_arrive_noinc((uint64_t*)&(plan.bar_kv_rope_ready)); + rs.update(); + } + }); + } + } + } + +#else + if (cute::thread0()) { + CUTE_INVALID_CONTROL_PATH("This kernel only supports sm100"); + } +#endif +} + +template +__global__ void __launch_bounds__(Kernel::NUM_THREADS, 1, 1) +sparse_attn_fwd_kernel(__grid_constant__ const SparseAttnFwdParams params, __grid_constant__ const TmaParams tma_params) { + Kernel::sparse_attn_fwd_kernel_devfunc(params, tma_params); +} + +template +void KernelTemplate::run(const SparseAttnFwdParams& params) { + KU_ASSERT(params.h_kv == 1); + KU_ASSERT(params.topk % B_TOPK == 0); // To save some boundry checkings + KU_ASSERT(params.topk >= 128); // To simplify synchronizations between outer loop boundries + KU_ASSERT(params.h_q == 64); + KU_ASSERT(params.d_qk == D_QK); + static_assert(D_QK == 576 || D_QK == 512); + + auto shape_Q_nope = make_shape(params.h_q, D_V, params.s_q); + auto tma_Q_nope = cute::make_tma_copy( + SM90_TMA_LOAD{}, + make_tensor( + make_gmem_ptr((bf16*)params.q), + make_layout( + shape_Q_nope, + make_stride(params.stride_q_h_q, _1{}, params.stride_q_s_q) + ) + ), + SmemLayoutQNoPE{} + ); + + auto shape_Q_rope = make_shape(params.h_q, D_Q-D_V == 0 ? 64 : D_Q-D_V, params.s_q); // If + auto tma_Q_rope = cute::make_tma_copy( + SM90_TMA_LOAD{}, + make_tensor( + make_gmem_ptr((bf16*)params.q + D_V), + make_layout( + shape_Q_rope, + make_stride(params.stride_q_h_q, _1{}, params.stride_q_s_q) + ) + ), + SmemLayoutQRoPE{} + ); + + auto shape_O = make_shape(params.h_q, params.d_v, params.s_q); + auto tma_O = cute::make_tma_copy( + SM90_TMA_STORE{}, + make_tensor( + make_gmem_ptr((bf16*)params.out), + make_layout( + shape_O, + make_stride(params.d_v, _1{}, params.h_q*params.d_v) + ) + ), + SmemLayoutOTiles<1>{} + ); + + CUtensorMap tensor_map_kv_nope; + { + uint64_t size[2] = {D_V, (unsigned long)params.s_kv}; + uint64_t stride[1] = {params.stride_kv_s_kv*sizeof(bf16)}; + uint32_t box_size[2] = {64, 1}; + uint32_t elem_stride[2] = {1, 1}; + CUresult res = CUTLASS_CUDA_DRIVER_WRAPPER_CALL(cuTensorMapEncodeTiled)( + &tensor_map_kv_nope, + CUtensorMapDataType::CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, + 2, + params.kv, + size, + stride, + box_size, + elem_stride, + CUtensorMapInterleave::CU_TENSOR_MAP_INTERLEAVE_NONE, + CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_128B, + CUtensorMapL2promotion::CU_TENSOR_MAP_L2_PROMOTION_L2_256B, + CUtensorMapFloatOOBfill::CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE + ); + KU_ASSERT(res == CUresult::CUDA_SUCCESS); + } + + TmaParams< + decltype(shape_Q_nope), decltype(tma_Q_nope), + decltype(shape_Q_rope), decltype(tma_Q_rope), + decltype(shape_O), decltype(tma_O) + > tma_params = { + shape_Q_nope, tma_Q_nope, + shape_Q_rope, tma_Q_rope, + shape_O, tma_O, + tensor_map_kv_nope + }; + auto kernel = &sparse_attn_fwd_kernel, decltype(tma_params)>; + + constexpr size_t smem_size = sizeof(SharedMemoryPlan); + KU_CUDA_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); + + kernel<<>>(params, tma_params); + KU_CHECK_KERNEL_LAUNCH(); +} + +template +void run_sparse_fwd_phase1_kernel(const SparseAttnFwdParams& params) { + KernelTemplate::run(params); +} + +} + diff --git a/csrc/kernels/sm100/prefill/sparse/fwd/head64/phase1.h b/csrc/kernels/sm100/prefill/sparse/fwd/head64/phase1.h new file mode 100644 index 00000000..dcc9da0d --- /dev/null +++ b/csrc/kernels/sm100/prefill/sparse/fwd/head64/phase1.h @@ -0,0 +1,10 @@ +#pragma once + +#include "kernels/params.h" + +namespace sm100::prefill::sparse_fwd::head64 { + +template +void run_sparse_fwd_phase1_kernel(const SparseAttnFwdParams& params); + +} diff --git a/csrc/sm100/prefill/sparse/fwd_for_small_topk/head128/config.h b/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/config.h similarity index 53% rename from csrc/sm100/prefill/sparse/fwd_for_small_topk/head128/config.h rename to csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/config.h index e4880078..e7abe49c 100644 --- a/csrc/sm100/prefill/sparse/fwd_for_small_topk/head128/config.h +++ b/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/config.h @@ -6,14 +6,16 @@ #include #include -#include "defines.h" -#include "params.h" +#include "kernels/defines.h" +#include "kernels/params.h" +#include "kernels/kv_cache_format.h" +#include "kernels/sm100/dequant_utils.cuh" -namespace sm100::fwd_for_small_topk::head128 { +namespace sm100::prefill::sparse_fwd_for_small_topk::head128 { using namespace cute; -template +template struct KernelTemplate { using ArgT = SparseFwdArgT; @@ -32,10 +34,10 @@ struct TmaParamsForDecode { CUtensorMap tensor_map_q; CUtensorMap tensor_map_o; CUtensorMap tensor_map_o_accum; - CUtensorMap tensor_map_kv_nope; - CUtensorMap tensor_map_kv_rope; - CUtensorMap tensor_map_extra_kv_nope; // Only available if extra_kv is enabled - CUtensorMap tensor_map_extra_kv_rope; + CUtensorMap tensor_map_kv_quant_part[2]; // One map per CTA: each CTA of the pair gathers its half of the token + CUtensorMap tensor_map_kv_bf16_part; + CUtensorMap tensor_map_extra_kv_quant_part[2]; // Only available if extra_kv is enabled + CUtensorMap tensor_map_extra_kv_bf16_part; }; using TmaParams = std::conditional_t< @@ -46,15 +48,25 @@ using TmaParams = std::conditional_t< static_assert(D_QK == 512); +using OrigKVFormat = KVCacheFormat; // Format of the paged `kv`, decode only +using ExtraKVFormat = KVCacheFormat; // Format of the paged `extra_kv`, decode only +static_assert(is_valid_kv_format_pair(MODEL_TYPE, EXTRA_MODEL_TYPE)); +static_assert(OrigKVFormat::D_QK == D_QK); +static_assert(ExtraKVFormat::D_FP8 + ExtraKVFormat::D_FP4 == OrigKVFormat::D_FP8 && ExtraKVFormat::D_BF16 == OrigKVFormat::D_BF16); + static constexpr int D_Q = D_QK; static constexpr int D_K = D_QK; static constexpr int D_V = 512; static constexpr float MAX_INIT_VAL = -1e30; // We use this number as the initial value for mi (max logits) to avoid -inf - (-inf) = nan +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 H_Q = 128; // For 2 CTAs static constexpr int B_TOPK = 64; // For 2 CTAs static constexpr int NUM_THREADS = 128*4; -static constexpr int NUM_WORKER_THREADS = IS_PREFILL ? (128 + 4 + (B_TOPK/8) + 1 + 128)*2 + 1 : (128 + 128 + 1 + 32 + 2 + 128)*2; +static constexpr int NUM_WORKER_THREADS = IS_PREFILL ? (128 + 4 + (B_TOPK/8) + 1 + 128)*2 + 1 : (128 + 128 + 1 + 32 + 2 + 128)*2 - (D_BF16 == 0); // For non-decode mode, we have 4 (half-)KV buffers // For decode mode, we have 3 (half-)KV buffers with two raw KV buffers @@ -62,10 +74,25 @@ static constexpr int NUM_K_BUFS = IS_DECODE ? 3 : 4; static constexpr int NUM_RAW_K_BUFS = IS_DECODE ? 2 : 0; static constexpr int NUM_INDEX_BUFS = IS_DECODE ? 4 : 4; -static constexpr int D_NOPE = 448; -static constexpr int D_ROPE = 64; -static constexpr int TMA_K_STRIDE_FOR_DECODING = D_NOPE + 2*D_ROPE; -static constexpr int NUM_SCALES_EACH_TOKEN = 8; // 7 scales + 1 padding +static constexpr int QUANT_TILE_SIZE = OrigKVFormat::QUANT_TILE_SIZE; +static constexpr int NUM_SCALES_EACH_TOKEN = OrigKVFormat::NUM_SCALES_EACH_TOKEN; +static constexpr int TMA_K_STRIDE_FOR_DECODING = OrigKVFormat::TMA_K_STRIDE; + +// Decode only. Each CTA of the pair gathers and dequantizes its half ([cta*256, cta*256 + 256)) of every selected token. +template static constexpr int RAW_TOKEN_SMEM_STRIDE = F::IS_FP4 ? F::QUANT_BYTES/2 + 32 : D_K/2; +// 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); +// Byte offset of CTA1's share within a token's data, and dim0 of each CTA's tensor map: for fp8, dim0 runs from the CTA's offset to +// the end of the token's data rather than to the end of its share -- a box reaching beyond the fp8 part then reads the raw bytes of +// the bf16 (RoPE) tail (CTA1, V4), which is actually faster, probably because it prefetches part of the BF16 part of the selected +// tokens. For fp4, dim0 is exactly the CTA's share and the rest of the box is zero-filled padding +template static constexpr int QUANT_PART_CTA_OFFSET = F::IS_FP4 ? F::QUANT_BYTES/2 : D_K/2; +template static constexpr int QUANT_PART_MAP_DIM0 = F::IS_FP4 ? F::QUANT_BYTES/2 : F::TMA_K_STRIDE - CTA*(D_K/2); +// One row of this CTA's 1 B scales per token. In a kernel with an fp4 extra_kv the rows are 16 B and the fp8 tokens use the first 8 B +static constexpr int SCALE_SMEM_STRIDE_PER_CTA = std::max(OrigKVFormat::NUM_SCALES_EACH_TOKEN, ExtraKVFormat::NUM_SCALES_EACH_TOKEN) / 2; +// The dequantized dims of this CTA: its half of the token, minus (CTA1) the bf16 RoPE tail +template using DequantizerT = KVBlockDequantizer, SCALE_SMEM_STRIDE_PER_CTA>; +static constexpr int K_ROPE_SW = MODEL_TYPE == ModelType::V41 ? 0 : 128; // RoPE part stored in SW128, in bytes. 0 for V41 since RoPE is fp8 static constexpr int B_EPI = 64; // Epilogue block size for normal case (i.e. prefill or non-splitkv decoding) static constexpr int B_EPI_SPLITKV = 32; // Epilogue block size for splitkv decoding @@ -83,16 +110,17 @@ struct tmem_cols { }; struct SharedMemoryPlan { - array_aligned Q; // Will be output for epilogue array_aligned K[NUM_K_BUFS]; - array_aligned K_raw[NUM_RAW_K_BUFS]; array_aligned S; + array_aligned Q; // Will be output for epilogue + array_aligned K_raw[NUM_RAW_K_BUFS]; // 128 B aligned for gather4 + static_assert(!IS_DECODE || B_TOPK * RAW_TOKEN_SMEM_STRIDE <= B_TOPK * (D_K/2)); float P_exchange[4][(H_Q/2/2)*(B_TOPK/2)]; float rowwise_max_buf[128], rowwise_li_buf[128]; CUTE_ALIGNAS(16) char is_k_valid[NUM_INDEX_BUFS][B_TOPK/8]; CUTE_ALIGNAS(16) int tma_coord[NUM_INDEX_BUFS][B_TOPK]; - CUTE_ALIGNAS(16) fp8_e8m0 scales[NUM_INDEX_BUFS][B_TOPK][NUM_SCALES_EACH_TOKEN/2]; + CUTE_ALIGNAS(16) uint8_t scales[NUM_INDEX_BUFS][B_TOPK][IS_DECODE ? SCALE_SMEM_STRIDE_PER_CTA : 0]; transac_bar_t bar_sQ_full, bar_tQ_empty, bar_tQ_full; transac_bar_t bar_tOut_full, bar_tOut_empty; diff --git a/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512.cu b/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512.cu new file mode 100644 index 00000000..9e8e2e1b --- /dev/null +++ b/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512.cu @@ -0,0 +1,8 @@ +#include "../phase1.h" +#include "../phase1.cuh" + +namespace sm100::prefill::sparse_fwd_for_small_topk::head128 { + +template void run_sparse_fwd_for_small_topk_phase1_kernel(const SparseAttnDecodeParams& params); + +} diff --git a/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512_splitkv.cu b/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512_splitkv.cu new file mode 100644 index 00000000..31ab52ee --- /dev/null +++ b/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512_splitkv.cu @@ -0,0 +1,8 @@ +#include "../phase1.h" +#include "../phase1.cuh" + +namespace sm100::prefill::sparse_fwd_for_small_topk::head128 { + +template void run_sparse_fwd_for_small_topk_phase1_kernel(const SparseAttnDecodeParams& params); + +} diff --git a/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512_v41.cu b/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512_v41.cu new file mode 100644 index 00000000..22d50c6c --- /dev/null +++ b/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512_v41.cu @@ -0,0 +1,8 @@ +#include "../phase1.h" +#include "../phase1.cuh" + +namespace sm100::prefill::sparse_fwd_for_small_topk::head128 { + +template void run_sparse_fwd_for_small_topk_phase1_kernel(const SparseAttnDecodeParams& params); + +} diff --git a/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512_v41_splitkv.cu b/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512_v41_splitkv.cu new file mode 100644 index 00000000..88e0ad60 --- /dev/null +++ b/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512_v41_splitkv.cu @@ -0,0 +1,8 @@ +#include "../phase1.h" +#include "../phase1.cuh" + +namespace sm100::prefill::sparse_fwd_for_small_topk::head128 { + +template void run_sparse_fwd_for_small_topk_phase1_kernel(const SparseAttnDecodeParams& params); + +} diff --git a/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512_v41fp4.cu b/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512_v41fp4.cu new file mode 100644 index 00000000..b7323ffb --- /dev/null +++ b/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512_v41fp4.cu @@ -0,0 +1,8 @@ +#include "../phase1.h" +#include "../phase1.cuh" + +namespace sm100::prefill::sparse_fwd_for_small_topk::head128 { + +template void run_sparse_fwd_for_small_topk_phase1_kernel(const SparseAttnDecodeParams& params); + +} diff --git a/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512_v41fp4_splitkv.cu b/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512_v41fp4_splitkv.cu new file mode 100644 index 00000000..50f8e420 --- /dev/null +++ b/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512_v41fp4_splitkv.cu @@ -0,0 +1,8 @@ +#include "../phase1.h" +#include "../phase1.cuh" + +namespace sm100::prefill::sparse_fwd_for_small_topk::head128 { + +template void run_sparse_fwd_for_small_topk_phase1_kernel(const SparseAttnDecodeParams& params); + +} diff --git a/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_k512.cu b/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_k512.cu new file mode 100644 index 00000000..53aa1040 --- /dev/null +++ b/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_k512.cu @@ -0,0 +1,8 @@ +#include "../phase1.h" +#include "../phase1.cuh" + +namespace sm100::prefill::sparse_fwd_for_small_topk::head128 { + +template void run_sparse_fwd_for_small_topk_phase1_kernel(const SparseAttnFwdParams& params); + +} diff --git a/csrc/sm100/prefill/sparse/fwd_for_small_topk/head128/phase1.cuh b/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/phase1.cuh similarity index 75% rename from csrc/sm100/prefill/sparse/fwd_for_small_topk/head128/phase1.cuh rename to csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/phase1.cuh index 388abc87..e9bdcc8e 100644 --- a/csrc/sm100/prefill/sparse/fwd_for_small_topk/head128/phase1.cuh +++ b/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/phase1.cuh @@ -1,3 +1,20 @@ +/* +Sparse Attention Forward Pass (Small TopK) — SM100, h_q=128 + +Specialized forward kernel optimized for small topk values (topk <= 1280), with 128 +query heads. Uses a different tiling strategy than the standard forward kernel. +Supports both prefill and decode modes (including split-KV decoding). Uses FP8 KV +cache for decoding mode with dequantization and RoPE/NoPE separation. + +Template parameters: + FWD_MODE — Forward mode (Prefill / Decode / DecodeWithSplitKV) + D_QK — Head dimension for QK (only 512 supported) + +Grid: [2*s_q, 1, 1] (prefill) or [2*b, 1, 1] (decode), Cluster: [2, 1, 1] +NUM_THREADS=512 (4 warpgroups) + +I/O: See SparseAttnFwdParams or SparseAttnDecodeParams in kernels/params.h +*/ #pragma once #include "phase1.h" @@ -7,39 +24,50 @@ #include #include -#include "params.h" -#include "utils.h" -#include "sm100/prefill/sparse/common_subroutine.h" -#include "sm100/helpers.h" +#include "kernels/params.h" +#include "kernels/utils.h" +#include "kernels/sm100/common_subroutine.h" +#include "kernels/sm100/helpers.h" #include "config.h" -namespace sm100::fwd_for_small_topk::head128 { +namespace sm100::prefill::sparse_fwd_for_small_topk::head128 { using namespace cute; using FwdMode = SparseAttnFwdMode; -template +// NOTES. We found out that, if we use kerutils::st_shared, warpgroup 0 suffers from register spilling. However, if we change to this dump implementation with `__cvta_generic_to_shared`, there is no register spilling. +CUTE_DEVICE +void st_shared_with_cvta(void* ptr, __int128_t val) { + asm volatile("st.shared.b128 [%0], %1;" :: "l"(__cvta_generic_to_shared(ptr)), "q"(val)); +} + +template __device__ void -KernelTemplate::sparse_attn_fwd_kernel_devfunc(const ArgT ¶ms, const TmaParams &tma_params) { +KernelTemplate::sparse_attn_fwd_kernel_devfunc(const ArgT ¶ms, const TmaParams &tma_params) { #ifdef KERUTILS_ENABLE_SM100A - // Grid shape: [2*s_q, 1, 1] for prefilling, [2*s_q, num_sm_parts, 1] for decoding + // Grid shape: [2*s_q, 1, 1] for prefilling, [2*s_q, num_sm_parts, 1] for decoding w/ splitKV, [2*s_q, batch, 1] for decoding w/o splitKV // Cluster shape: [2, 1, 1] const int warp_idx = cutlass::canonical_warp_idx_sync(); const int lane_idx = threadIdx.x % 32; const int warpgroup_idx = cutlass::canonical_warp_group_idx(); const int idx_in_warpgroup = threadIdx.x % 128; - const int cta_idx = block_id_in_cluster().x; + int cta_idx = blockIdx.x % 2; extern __shared__ char wksp_buf[]; SharedMemoryPlan &smem = *reinterpret_cast(wksp_buf); + ku::barrier_cluster_arrive_relaxed(); + ku::barrier_cluster_wait_acquire(); + if (warp_idx == 0 && elect_one_sync()) { cute::prefetch_tma_descriptor(&tma_params.tensor_map_q); cute::prefetch_tma_descriptor(&tma_params.tensor_map_o); if constexpr (IS_DECODE) { - cute::prefetch_tma_descriptor(&tma_params.tensor_map_kv_nope); - cute::prefetch_tma_descriptor(&tma_params.tensor_map_kv_rope); + cute::prefetch_tma_descriptor(&tma_params.tensor_map_kv_quant_part[cta_idx]); + if constexpr (D_BF16 > 0) { + cute::prefetch_tma_descriptor(&tma_params.tensor_map_kv_bf16_part); + } } else { cute::prefetch_tma_descriptor(&tma_params.tensor_map_kv); } @@ -73,7 +101,7 @@ KernelTemplate::sparse_attn_fwd_kernel_devfunc(const ArgT ¶m CUTE_UNROLL for (int i = 0; i < NUM_INDEX_BUFS; ++i) { smem.bar_valid_coord_scales_full[i].init(IS_PREFILL ? B_TOPK/8 : 32); - smem.bar_valid_coord_scales_empty[i].init(IS_PREFILL ? 128 : (128 + (cta_idx==1) + 2 + 128)); + smem.bar_valid_coord_scales_empty[i].init(IS_PREFILL ? 128 : (128 + (D_BF16 > 0 && block_id_in_cluster().x==1) + 2 + 128)); // We use `block_id_in_cluster().x` instead of cta_idx to prevent reg spilling } if constexpr (IS_DECODE) { CUTE_UNROLL @@ -126,11 +154,6 @@ KernelTemplate::sparse_attn_fwd_kernel_devfunc(const ArgT ¶m bool 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); int 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); - // start_block_idx = 0; - // end_block_idx = total_topk_padded / B_TOPK; - // is_split = false; - // n_split_idx = 0; - OuterloopArgs args = { (bool)outer_loop_phase, batch_idx, s_q_idx, @@ -145,7 +168,7 @@ KernelTemplate::sparse_attn_fwd_kernel_devfunc(const ArgT ¶m outer_loop_phase ^= 1; } } else { - // Prefill mode. Use CLC to allocate different s_q (for decoding, different batches + s_q) to different workers + // Prefill (or decoding w/o splitKV) mode. Use CLC to allocate different s_q (for decoding, different batches + s_q) to different workers ku::CLCResult next_job = {true, (int)blockIdx.x, IS_PREFILL ? 0 : (int)blockIdx.y, 0}; CUTE_NO_UNROLL while (next_job.is_valid) { @@ -220,7 +243,7 @@ KernelTemplate::sparse_attn_fwd_kernel_devfunc(const ArgT ¶m float2 output_scale_float2 = float2 {output_scale, output_scale}; smem.bar_li_empty.arrive(); - // Retrieve and store O, and calculate delta := sum(O*dO, dim=-1) if FWD_MODE is Recompute + // Retrieve and store O smem.bar_tOut_full.wait(args.outer_loop_phase); if (is_last_o && elect_one_sync()) { cudaTriggerProgrammaticLaunchCompletion(); @@ -243,11 +266,11 @@ KernelTemplate::sparse_attn_fwd_kernel_devfunc(const ArgT ¶m o[i*4+j] = ku::float2_mul(o[i*4+j], output_scale_float2); o_bf16[j] = __float22bfloat162_rn(o[i*4+j]); } - bf16* o_do_addr = sO_addrs[i] + k*B_EPI*(H_Q/2); - if (k == 0 && i == 0) { - smem.bar_tQ_full.wait(args.outer_loop_phase^1^is_last_o); // Wait for sQ's availability - } - ku::st_shared(o_do_addr, *(__int128_t*)o_bf16); + bf16* o_addr = sO_addrs[i] + k*B_EPI*(H_Q/2); + if (k == 0 && i == 0) { + smem.bar_tQ_full.wait(args.outer_loop_phase^1^is_last_o); // Wait for sQ's availability + } + st_shared_with_cvta(o_addr, *(__int128_t*)o_bf16); } } @@ -287,7 +310,7 @@ KernelTemplate::sparse_attn_fwd_kernel_devfunc(const ArgT ¶m if (k == 0 && i == 0) { smem.bar_tQ_full.wait(args.outer_loop_phase^1^is_last_o); // Wait for sQ's availability } - ku::st_shared( + st_shared_with_cvta( sO_accum_addrs[i] + cur_buf_idx*((H_Q/2)*B_EPI_SPLITKV*2), *(__int128_t*)(o + i*4) ); @@ -365,7 +388,8 @@ KernelTemplate::sparse_attn_fwd_kernel_devfunc(const ArgT ¶m // A subtile is 128 rows * 16 cols (256b, 32B) (in UTCCP's view), or 64 rows * 16 cols * 2 (in our view) // NOTE Using `sQ_desc+((tile_idx*((H_Q/2)*128*2) + subtile_idx*32) >> 4)` leads to IMA, doesn't know why UMMA::SmemDescriptor cur_sQ_desc = sQ_desc; - cur_sQ_desc.lo += ((tile_idx*((H_Q/2)*128*2) + subtile_idx*32) >> 4); + // cur_sQ_desc.lo += ((tile_idx*((H_Q/2)*128*2) + subtile_idx*32) >> 4); + asm volatile ("add.u32 %0, %0, %1;" : "+r"(cur_sQ_desc.lo) : "r"((tile_idx*((H_Q/2)*128*2) + subtile_idx*32) >> 4)); // uint64_t cur_sQ_desc = sQ_desc; // cur_sQ_desc += ((tile_idx*((H_Q/2)*128*2) + subtile_idx*32) >> 4); SM100_UTCCP_128dp256bit_2cta::copy( @@ -451,67 +475,33 @@ KernelTemplate::sparse_attn_fwd_kernel_devfunc(const ArgT ¶m } } else { - // 8 threads per token struct IsCTA0 {}; struct IsCTA1 {}; auto launch_dequant_wg = [&](auto cta_id_t) { static constexpr bool IS_CTA1 = std::is_same::value; - constexpr int GROUP_SIZE = 8, NUM_GROUPS = 128/8, ROWS_PER_GROUP = B_TOPK / NUM_GROUPS, COLS_PER_GROUP = (IS_CTA1 ? 256-64 : 256) / (GROUP_SIZE*8); - int group_idx = idx_in_warpgroup/GROUP_SIZE, idx_in_group = idx_in_warpgroup%GROUP_SIZE; - Tensor nope0 = make_tensor(make_smem_ptr(smem.K[0].data()), ku::make_umma_canonical_k_major_layout()); - bf16* nope0_base = &nope0(group_idx, idx_in_group*8); - fp8_e4m3* raw_nope0_base = smem.K_raw[0].data() + group_idx*(D_K/2) + idx_in_group*8; + const DequantizerT dequant_orig(idx_in_warpgroup); + const DequantizerT dequant_extra(idx_in_warpgroup); // Identical to dequant_orig unless the extra cache is fp4 + const uint8_t *raw_0 = smem.K_raw[0].data(), *raw_1 = smem.K_raw[1].data(); + static_assert(NUM_RAW_K_BUFS == 2); run_outer_loop([&](const OuterloopArgs &args) { - CUTE_NO_UNROLL - for (int block_idx = args.start_block_idx; block_idx < args.end_block_idx; ++block_idx) { + for_each_kv_block(args.start_block_idx, args.end_block_idx, args.num_orig_kv_blocks, [&](int, bool) { auto [k_buf_idx, k_bar_phase] = rs.get(); auto [raw_k_buf_idx, raw_k_bar_phase] = rs.get(); auto [index_buf_idx, index_bar_phase] = rs.get(); - fp8_e4m3* raw_nope_base = raw_nope0_base + raw_k_buf_idx * (B_TOPK*(D_K/2)); - auto get_raw_fp8 = [&](int local_row_idx, int local_col_idx) -> uint64_t { - return *(uint64_t*)(raw_nope_base + local_row_idx*NUM_GROUPS*(D_K/2) + local_col_idx*(GROUP_SIZE*8)); - }; - bf16* nope_base = nope0_base + k_buf_idx * (B_TOPK*(D_K/2)); - uint32_t cur_nope_base_uint_addr = cute::cast_smem_ptr_to_uint(nope_base); - auto st_128b = [&](int local_row_idx, int local_col_idx, __int128_t &data) { - asm volatile ("st.weak.shared::cta.b128 [%0], %1;\n" - : - : "r"(cur_nope_base_uint_addr + 2*(local_row_idx*NUM_GROUPS*64 + local_col_idx*B_TOPK*64)), "q"(data) // 2 for sizeof(bf16) - ); // We have this `asm volatile` here, otherwise the compiler generates ST.E instead of STS - }; smem.bar_valid_coord_scales_full[index_buf_idx].wait(index_bar_phase); smem.bar_raw_KV_full[raw_k_buf_idx].wait(raw_k_bar_phase); - CUTE_UNROLL - for (int local_row_idx = 0; local_row_idx < ROWS_PER_GROUP; ++local_row_idx) { - int row_idx = local_row_idx*NUM_GROUPS + group_idx; - bf16 scales[4]; - fp8_e8m0 scales_e8m0[4]; - *(uint32_t*)scales_e8m0 = *(uint32_t*)(smem.scales[index_buf_idx][row_idx]); - *(__nv_bfloat162_raw*)(scales+0) = __nv_cvt_e8m0x2_to_bf162raw(*(unsigned short*)(scales_e8m0+0)); - *(__nv_bfloat162_raw*)(scales+2) = __nv_cvt_e8m0x2_to_bf162raw(*(unsigned short*)(scales_e8m0+2)); - - uint64_t cur_data_fp8x8 = get_raw_fp8(local_row_idx, 0); - CUTE_UNROLL - for (int local_col_idx = 0; local_col_idx < COLS_PER_GROUP; ++local_col_idx) { - ku::nve4m3x2 data_fp8[4]; - ku::nvbf16x2 data_bf16[4]; - *(uint64_t*)data_fp8 = cur_data_fp8x8; - if (local_col_idx+1 < COLS_PER_GROUP) - cur_data_fp8x8 = get_raw_fp8(local_row_idx, local_col_idx+1); - bf16 scale = scales[local_col_idx]; - CUTE_UNROLL - for (int i = 0; i < 4; ++i) { - data_bf16[i] = fp8x2_to_bf16x2_with_scale(data_fp8[i], *(ku::nvbf16*)(&scale)); - } - if (local_row_idx == 0 && local_col_idx == 0) { - smem.bar_KV_empty[k_buf_idx].wait(k_bar_phase^1); - } - st_128b(local_row_idx, local_col_idx, *(__int128_t*)data_bf16); - } - } + const auto *dequant = [&] { + if constexpr (std::is_same_v) return &dequant_orig; else return &dequant_extra; + }(); + dequant->run( + raw_k_buf_idx == 0 ? raw_0 : raw_1, + smem.scales[index_buf_idx][0], + cute::cast_smem_ptr_to_uint(smem.K[k_buf_idx].data()), + [&] { smem.bar_KV_empty[k_buf_idx].wait(k_bar_phase^1); } + ); fence_view_async_shared(); // NOTE Should we use shared::cluster here? __syncwarp(); @@ -521,7 +511,7 @@ KernelTemplate::sparse_attn_fwd_kernel_devfunc(const ArgT ¶m smem.bar_KV_full[k_buf_idx].arrive(0u); } rs.update(); - } + }); }); }; if (cta_idx == 0) { @@ -558,8 +548,11 @@ KernelTemplate::sparse_attn_fwd_kernel_devfunc(const ArgT ¶m if constexpr (IS_PREFILL) { smem.bar_KV_full[k_buf_idx].arrive_and_expect_tx(B_TOPK*D_K*sizeof(bf16)); } else { - // RoPE only - smem.bar_KV_full[k_buf_idx].arrive_and_expect_tx(B_TOPK*D_ROPE*sizeof(bf16)); + if constexpr (D_BF16 > 0) { + smem.bar_KV_full[k_buf_idx].arrive_and_expect_tx(B_TOPK*D_BF16*sizeof(bf16)); + } else { + smem.bar_KV_full[k_buf_idx].arrive(); + } } smem.bar_KV_full[k_buf_idx].wait(k_bar_phase); ku::tcgen05_after_thread_sync(); @@ -614,8 +607,8 @@ KernelTemplate::sparse_attn_fwd_kernel_devfunc(const ArgT ¶m ku::umma_arrive_multicast_2x1SM_noelect(smem.bar_tOut_full, 1|2); }); } else if (warp_idx == 8 && cta_idx == 1 && elect_one_sync()) { - if constexpr (IS_DECODE) { - // KV RoPE fetching warp + if constexpr (IS_DECODE && D_BF16 > 0) { + // KV BF16 part fetching warp run_outer_loop([&](const OuterloopArgs &args) { CUTE_NO_UNROLL for (int block_idx = args.start_block_idx; block_idx < args.end_block_idx; ++block_idx) { @@ -627,12 +620,12 @@ KernelTemplate::sparse_attn_fwd_kernel_devfunc(const ArgT ¶m for (int row = 0; row < B_TOPK; row += 4) { int4 cur_indices = *(int4*)(smem.tma_coord[index_buf_idx] + row); ku::tma_gather4_cta_group_2( - block_idx >= args.num_orig_kv_blocks ? &tma_params.tensor_map_extra_kv_rope : &tma_params.tensor_map_kv_rope, + block_idx >= args.num_orig_kv_blocks ? &tma_params.tensor_map_extra_kv_bf16_part : &tma_params.tensor_map_kv_bf16_part, smem.bar_KV_full[k_buf_idx], smem.K[k_buf_idx].data() + (D_NOPE-D_K/2)*B_TOPK + row*D_ROPE, 0, cur_indices, - (int64_t)TMA::CacheHintSm90::EVICT_LAST + (int64_t)TMA::CacheHintSm90::EVICT_FIRST ); } smem.bar_valid_coord_scales_empty[index_buf_idx].arrive(); @@ -668,11 +661,11 @@ KernelTemplate::sparse_attn_fwd_kernel_devfunc(const ArgT ¶m } else { static_assert(B_TOPK == 64); // Each thread is responsible for 2 tokens - static constexpr int tma_coords_step_per_token = 576/TMA_K_STRIDE_FOR_DECODING; + static constexpr int tma_coords_step_per_token = 1; // A token's data is exactly one TMA_K_STRIDE row, for all formats int tma_coords_step_per_block = params.stride_kv_block / TMA_K_STRIDE_FOR_DECODING; // must < 2G since k_batch_stride < 1T and TMA_K_STRIDE_FOR_DECODING > 512 - int tma_coords_step_per_extra_block = params.stride_extra_kv_block / TMA_K_STRIDE_FOR_DECODING; - uint8_t* k_scales_ptr = (uint8_t*)params.kv + params.page_block_size*(D_NOPE+2*D_ROPE); - uint8_t* extra_k_scales_ptr = (uint8_t*)params.extra_kv + params.extra_page_block_size*(D_NOPE+2*D_ROPE); + int tma_coords_step_per_extra_block = params.stride_extra_kv_block / ExtraKVFormat::TMA_K_STRIDE; + uint8_t* k_scales_ptr = (uint8_t*)params.kv + params.page_block_size*TMA_K_STRIDE_FOR_DECODING; + uint8_t* extra_k_scales_ptr = (uint8_t*)params.extra_kv + params.extra_page_block_size*ExtraKVFormat::TMA_K_STRIDE; run_outer_loop([&](const OuterloopArgs &args) { int* indices = (int*)params.indices + params.stride_indices_b*args.batch_idx + params.stride_indices_s_q*args.s_q_idx; @@ -680,61 +673,93 @@ KernelTemplate::sparse_attn_fwd_kernel_devfunc(const ArgT ¶m struct IsOrigBlock {}; struct IsExtraBlock {}; - auto process_one_block = [&](int block_idx, auto is_extra_block_t) { + // Prefetch the next block's indices while processing the current one, so that the + // index LDG latency overlaps with the scale LDG and the computation of the current block. + auto load_block_indices = [&](int block_idx, auto is_extra_block_t) -> int2 { + static constexpr bool IS_EXTRA_BLOCK = std::is_same_v; + if constexpr (!IS_EXTRA_BLOCK) { + return __ldg((int2*)(indices + block_idx*B_TOPK + lane_idx*2)); + } else { + return __ldg((int2*)(extra_indices + (block_idx-args.num_orig_kv_blocks)*B_TOPK + lane_idx*2)); + } + }; + auto process_one_block = [&](int block_idx, int2 my_indices, auto is_extra_block_t) { auto [index_buf_idx, index_bar_phase] = rs.get(); static constexpr bool IS_EXTRA_BLOCK = std::is_same_v; + using F = std::conditional_t; int cur_block_size = IS_EXTRA_BLOCK ? params.extra_page_block_size : params.page_block_size; int64_t cur_k_block_stride = IS_EXTRA_BLOCK ? params.stride_extra_kv_block : params.stride_kv_block; [[maybe_unused]] int cur_k_row_stride = IS_EXTRA_BLOCK ? params.stride_extra_kv_row : params.stride_kv_row; uint8_t* cur_k_scales_ptr = IS_EXTRA_BLOCK ? extra_k_scales_ptr : k_scales_ptr; int cur_tma_coords_step_per_block = IS_EXTRA_BLOCK ? tma_coords_step_per_extra_block : tma_coords_step_per_block; - int abs_pos, my_indices[2]; - if (!IS_EXTRA_BLOCK) { - abs_pos = block_idx*B_TOPK + lane_idx*2; - *(int2*)my_indices = __ldg((int2*)(indices + abs_pos)); - } else { - abs_pos = (block_idx-args.num_orig_kv_blocks)*B_TOPK + lane_idx*2; - *(int2*)my_indices = __ldg((int2*)(extra_indices + abs_pos)); + int abs_pos = IS_EXTRA_BLOCK ? (block_idx-args.num_orig_kv_blocks)*B_TOPK + lane_idx*2 : block_idx*B_TOPK + lane_idx*2; + + // Issue the scale LDGs before waiting for the index buffer, so their latency + // overlaps with the (usually long) empty-barrier wait. + alignas(16) uint8_t scales[2][SCALE_SMEM_STRIDE_PER_CTA]; + int kv_block_idx_arr[2], idx_in_block_arr[2]; + CUTE_UNROLL + for (int i = 0; i < 2; ++i) { + int cur_idx = i == 0 ? my_indices.x : my_indices.y; + int kv_block_idx = (unsigned int)cur_idx / cur_block_size; + int idx_in_block = (unsigned int)cur_idx % cur_block_size; + kv_block_idx_arr[i] = kv_block_idx; + idx_in_block_arr[i] = idx_in_block; + // This CTA's half of the token's row of 1 B scales + int64_t offset = kv_block_idx*cur_k_block_stride + (idx_in_block*F::NUM_SCALES_EACH_TOKEN + (cta_idx == 1 ? F::NUM_SCALES_EACH_TOKEN/2 : 0)); + bool is_token_valid = cur_idx != -1 && (abs_pos+i < (IS_EXTRA_BLOCK?args.extra_topk_length:args.topk_length)); + ldg_or_zero(scales[i], cur_k_scales_ptr + offset, is_token_valid); } smem.bar_valid_coord_scales_empty[index_buf_idx].wait(index_bar_phase^1); int tma_coords[2]; - fp8_e8m0 scales[2*(NUM_SCALES_EACH_TOKEN/2)]; char valid_mask = 0; CUTE_UNROLL for (int i = 0; i < 2; ++i) { - int block_idx, idx_in_block; - block_idx = (unsigned int)my_indices[i] / cur_block_size; - idx_in_block = (unsigned int)my_indices[i] % cur_block_size; - bool is_token_valid = my_indices[i] != -1 && (abs_pos+i < (IS_EXTRA_BLOCK?args.extra_topk_length:args.topk_length)); + int cur_idx = i == 0 ? my_indices.x : my_indices.y; + bool is_token_valid = cur_idx != -1 && (abs_pos+i < (IS_EXTRA_BLOCK?args.extra_topk_length:args.topk_length)); valid_mask |= is_token_valid << i; - tma_coords[i] = is_token_valid ? block_idx*cur_tma_coords_step_per_block + idx_in_block*tma_coords_step_per_token : -1; // If the token is invalid because it topk position exceeds topk_length, we must manually fill tma_coords with -1 to avoid copying-in NaN. - - int64_t offset = block_idx*cur_k_block_stride + (idx_in_block*8 + (cta_idx == 1 ? 4 : 0)); // Each token has 7 scale factors with an extra 1B padding - uint32_t scalesx4 = is_token_valid ? __ldg((uint32_t*)(cur_k_scales_ptr + offset)) : 0; - *(uint32_t*)(scales+i*(NUM_SCALES_EACH_TOKEN/2)) = scalesx4; + tma_coords[i] = is_token_valid ? kv_block_idx_arr[i]*cur_tma_coords_step_per_block + idx_in_block_arr[i]*tma_coords_step_per_token : -1; // If the token is invalid because it topk position exceeds topk_length, we must manually fill tma_coords with -1 to avoid copying-in NaN. } valid_mask <<= lane_idx%4*2; valid_mask |= __shfl_xor_sync(0xFFFFFFFF, valid_mask, 0x1); valid_mask |= __shfl_xor_sync(0xFFFFFFFF, valid_mask, 0x2); - *(uint64_t*)(smem.scales[index_buf_idx] + lane_idx*2) = *(uint64_t*)scales; + + if constexpr (SCALE_SMEM_STRIDE_PER_CTA == F::NUM_SCALES_EACH_TOKEN/2) { + copy_bytes<2*SCALE_SMEM_STRIDE_PER_CTA>(smem.scales[index_buf_idx][lane_idx*2], scales[0]); + } else { + // Mixed-format kernel, fp8 tokens: each token uses the first bytes of its 16 B row + CUTE_UNROLL + for (int i = 0; i < 2; ++i) + copy_bytes(smem.scales[index_buf_idx][lane_idx*2 + i], scales[i]); + } *(int2*)(smem.tma_coord[index_buf_idx] + lane_idx*2) = *(int2*)tma_coords; if (lane_idx%4 == 0) smem.is_k_valid[index_buf_idx][lane_idx/4] = valid_mask; - + smem.bar_valid_coord_scales_full[index_buf_idx].arrive(); rs.update(); }; + const int orig_end = min(args.num_orig_kv_blocks, args.end_block_idx); + int2 my_indices = args.start_block_idx < orig_end ? load_block_indices(args.start_block_idx, IsOrigBlock{}) : int2{}; CUTE_NO_UNROLL - for (int block_idx = args.start_block_idx; block_idx < min(args.num_orig_kv_blocks, args.end_block_idx); ++block_idx) { - process_one_block(block_idx, IsOrigBlock{}); + for (int block_idx = args.start_block_idx; block_idx < orig_end; ++block_idx) { + bool has_next = block_idx+1 < orig_end; + int2 next_indices = has_next ? load_block_indices(block_idx+1, IsOrigBlock{}) : int2{}; + process_one_block(block_idx, my_indices, IsOrigBlock{}); + if (has_next) my_indices = next_indices; } + const int extra_start = max(args.start_block_idx, args.num_orig_kv_blocks); + my_indices = extra_start < args.end_block_idx ? load_block_indices(extra_start, IsExtraBlock{}) : int2{}; CUTE_NO_UNROLL - for (int block_idx = max(args.start_block_idx, args.num_orig_kv_blocks); block_idx < args.end_block_idx; ++block_idx) { - process_one_block(block_idx, IsExtraBlock{}); + for (int block_idx = extra_start; block_idx < args.end_block_idx; ++block_idx) { + bool has_next = block_idx+1 < args.end_block_idx; + int2 next_indices = has_next ? load_block_indices(block_idx+1, IsExtraBlock{}) : int2{}; + process_one_block(block_idx, my_indices, IsExtraBlock{}); + if (has_next) my_indices = next_indices; } }); } @@ -751,14 +776,21 @@ KernelTemplate::sparse_attn_fwd_kernel_devfunc(const ArgT ¶m }); } } else { - // Raw KV NoPE Producer thread + // Raw (quantized) KV part Producer thread (+CLC Producer if splitKV is not enabled) run_outer_loop([&](const OuterloopArgs &args) { - CUTE_NO_UNROLL - for (int block_idx = args.start_block_idx; block_idx < args.end_block_idx; ++block_idx) { + if (warp_idx == 10 && FWD_MODE == FwdMode::Decode) { + if (cta_idx == 0) { + smem.bar_clc_empty.wait(args.outer_loop_phase^1); + ku::issue_clc_query_multicast_cluster_all(smem.bar_clc_full, smem.clc_response_obj); + } + smem.bar_clc_full.arrive_and_expect_tx(sizeof(smem.clc_response_obj)); + } + for_each_kv_block(args.start_block_idx, args.end_block_idx, args.num_orig_kv_blocks, [&](int block_idx, bool is_extra_block) { auto [raw_k_buf_idx, raw_k_bar_phase] = rs.get(); auto [index_buf_idx, index_bar_phase] = rs.get(); smem.bar_valid_coord_scales_full[index_buf_idx].wait(index_bar_phase); smem.bar_raw_KV_empty[raw_k_buf_idx].wait(raw_k_bar_phase^1); + const CUtensorMap *tensor_map = is_extra_block ? &tma_params.tensor_map_extra_kv_quant_part[cta_idx] : &tma_params.tensor_map_kv_quant_part[cta_idx]; int4 nxt_indices = *(int4*)(smem.tma_coord[index_buf_idx] + (warp_idx == 10 ? 0 : 4)); CUTE_UNROLL @@ -767,20 +799,20 @@ KernelTemplate::sparse_attn_fwd_kernel_devfunc(const ArgT ¶m if (row+8 < B_TOPK) nxt_indices = *(int4*)(smem.tma_coord[index_buf_idx] + row + 8); ku::tma_gather4( - block_idx >= args.num_orig_kv_blocks ? &tma_params.tensor_map_extra_kv_nope : &tma_params.tensor_map_kv_nope, + tensor_map, smem.bar_raw_KV_full[raw_k_buf_idx], - smem.K_raw[raw_k_buf_idx].data() + row*(D_K/2), - cta_idx*(D_K/2), + smem.K_raw[raw_k_buf_idx].data() + row*RAW_TOKEN_SMEM_STRIDE, + 0, cur_indices, - (int64_t)TMA::CacheHintSm90::EVICT_LAST + (int64_t)TMA::CacheHintSm90::EVICT_FIRST ); } if (warp_idx == 10) { - smem.bar_raw_KV_full[raw_k_buf_idx].arrive_and_expect_tx(B_TOPK*(D_K/2)*sizeof(fp8_e4m3)); + smem.bar_raw_KV_full[raw_k_buf_idx].arrive_and_expect_tx(B_TOPK*RAW_TOKEN_SMEM_STRIDE); } smem.bar_valid_coord_scales_empty[index_buf_idx].arrive(); rs.update(); - } + }); }); } } @@ -813,11 +845,11 @@ KernelTemplate::sparse_attn_fwd_kernel_devfunc(const ArgT ¶m ku::tcgen05_after_thread_sync(); retrieve_mask_and_reduce_p< NUM_ELEMS_PER_THREAD, - tmem_cols::P, barrier_ids::WG2_WARP02_SYNC, barrier_ids::WG2_WARP13_SYNC, - false + false // We never write P back through the exchange buffer >( + tmem_cols::P, smem.is_k_valid[indices_buf_idx], local_warp_idx, lane_idx, @@ -836,7 +868,6 @@ KernelTemplate::sparse_attn_fwd_kernel_devfunc(const ArgT ¶m real_mi = max(real_mi, cur_pi_max); bool should_scale_o = __any_sync(0xffffffff, cur_pi_max - mi > 6.0f); - // Calc scale factor, and scale li float new_max, scale_for_old; if (!should_scale_o) { @@ -864,7 +895,7 @@ KernelTemplate::sparse_attn_fwd_kernel_devfunc(const ArgT ¶m // Rescale O if (k > 0 && should_scale_o) { ku::tcgen05_after_thread_sync(); - rescale_O(scale_for_old); + rescale_O(scale_for_old); ku::tcgen05_before_thread_sync(); } @@ -915,7 +946,6 @@ KernelTemplate::sparse_attn_fwd_kernel_devfunc(const ArgT ¶m params.lse_accum[args.n_split_idx*params.stride_lse_accum_split + args.s_q_idx*params.stride_lse_accum_s_q + head_idx] = cur_lse_2base; } } - } }); } @@ -944,28 +974,41 @@ flash_fwd_splitkv_mla_fp8_sparse_kernel(__grid_constant__ const typename Kernel: Kernel::sparse_attn_fwd_kernel_devfunc(params, tma_params); } -template -void KernelTemplate::run(const ArgT& params) { - static_assert(D_QK == 576 || D_QK == 512); - +template +void KernelTemplate::run(const ArgT& params) { KU_ASSERT(params.h_kv == 1); KU_ASSERT(params.topk % B_TOPK == 0); // To save some boundry checkings KU_ASSERT(params.h_q == H_Q); // To save some calculation KU_ASSERT(params.d_qk == D_QK); + if constexpr (IS_DECODE) { + KU_ASSERT(params.model_type == MODEL_TYPE && params.extra_model_type == EXTRA_MODEL_TYPE); + } static_assert(D_Q == 512); CUtensorMap tensor_map_q; if constexpr (IS_DECODE) { - KU_ASSERT(params.stride_q_b % params.stride_q_s_q == 0, "In decode mode for MODEL1 sparse fp8 decoding on sm100f, q.stride(0) (on the batch dimension) must be divisible by q.stride(1) (on the sequence dimension)."); - tensor_map_q = ku::make_tensor_map( - {64ul, H_Q, 2ul, (D_Q/64ul)/2ul, (unsigned long)params.b * (params.stride_q_b / params.stride_q_s_q)}, - ku::make_stride_helper({params.stride_q_h_q, D_Q/2, 64, params.stride_q_s_q}, sizeof(bf16)), - {64, H_Q/2, 2, (D_Q/64)/2, 1}, - params.q, - CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, - CU_TENSOR_MAP_SWIZZLE_128B, - CU_TENSOR_MAP_L2_PROMOTION_L2_256B - ); + if (params.b > 1) { + KU_ASSERT(params.stride_q_b % params.stride_q_s_q == 0, "In decode mode for V4 sparse fp8 decoding on sm100f, q.stride(0) (on the batch dimension) must be divisible by q.stride(1) (on the sequence dimension)."); + tensor_map_q = ku::make_tensor_map( + {64ul, H_Q, 2ul, (D_Q/64ul)/2ul, (unsigned long)params.b * (params.stride_q_b / params.stride_q_s_q)}, + ku::make_stride_helper({params.stride_q_h_q, D_Q/2, 64, params.stride_q_s_q}, sizeof(bf16)), + {64, H_Q/2, 2, (D_Q/64)/2, 1}, + params.q, + CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, + CU_TENSOR_MAP_SWIZZLE_128B, + CU_TENSOR_MAP_L2_PROMOTION_L2_256B + ); + } else { + tensor_map_q = ku::make_tensor_map( + {64ul, H_Q, 2ul, (D_Q/64ul)/2ul, (unsigned long)params.s_q}, + ku::make_stride_helper({params.stride_q_h_q, D_Q/2, 64, params.s_q == 1 ? 0 : params.stride_q_s_q}, sizeof(bf16)), + {64, H_Q/2, 2, (D_Q/64)/2, 1}, + params.q, + CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, + CU_TENSOR_MAP_SWIZZLE_128B, + CU_TENSOR_MAP_L2_PROMOTION_L2_256B + ); + } } else { tensor_map_q = ku::make_tensor_map( {64ul, H_Q, 2ul, (D_Q/64ul)/2ul, (unsigned long)params.s_q}, @@ -979,34 +1022,26 @@ void KernelTemplate::run(const ArgT& params) { } CUtensorMap tensor_map_kv; - CUtensorMap tensor_map_kv_nope, tensor_map_kv_rope, tensor_map_extra_kv_nope = {}, tensor_map_extra_kv_rope = {}; + CUtensorMap tensor_map_kv_quant_part[2] = {}, tensor_map_extra_kv_quant_part[2] = {}; + CUtensorMap tensor_map_kv_bf16_part = {}, tensor_map_extra_kv_bf16_part = {}; if constexpr (IS_DECODE) { - auto get_kv_tensormap = [&](bool is_extra, void* k_ptr, int num_blocks, int64_t stride_kv_block, int64_t stride_kv_row) -> std::pair { - KU_ASSERT((int64_t)k_ptr % 16 == 0, "The base address of %sk_ptr (%p) must be 16B aligned for sparse fp8 attention on sm100f", is_extra?"extra_":"", k_ptr); - KU_ASSERT(stride_kv_block % TMA_K_STRIDE_FOR_DECODING == 0, "%sk_cache.stride(0) (%ld) must be a multiple of %d. Padding might be necessary", is_extra?"extra_":"", stride_kv_block, TMA_K_STRIDE_FOR_DECODING); - CUtensorMap tensor_map_kv_nope = ku::make_tensor_map( - {D_NOPE + D_ROPE*2, (uint64_t)num_blocks * (stride_kv_block/TMA_K_STRIDE_FOR_DECODING)}, - {TMA_K_STRIDE_FOR_DECODING}, - {D_K/2, 1}, - k_ptr, - CUtensorMapDataType::CU_TENSOR_MAP_DATA_TYPE_UINT8, - CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_NONE, - CUtensorMapL2promotion::CU_TENSOR_MAP_L2_PROMOTION_L2_128B - ); // NOTE: Here we use `D_NOPE+D_ROPE*2` as the box shape instead of D_NOPE because it's actually faster. I think that's because, if we use `D_NOPE+D_ROPE*2`, we can prefetch part of the RoPE part of the selected tokens. - CUtensorMap tensor_map_kv_rope = ku::make_tensor_map( - {D_ROPE, (uint64_t)num_blocks * (stride_kv_block/TMA_K_STRIDE_FOR_DECODING)}, - {TMA_K_STRIDE_FOR_DECODING}, - {64, 1}, - (uint8_t*)k_ptr + D_NOPE, - CUtensorMapDataType::CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, - CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_128B, - CUtensorMapL2promotion::CU_TENSOR_MAP_L2_PROMOTION_L2_128B - ); - return {tensor_map_kv_nope, tensor_map_kv_rope}; + // Two maps per cache, one per CTA of the pair + auto make_quant_part_maps = [&](F, const char *name, void* k_ptr, int num_blocks, int64_t stride_kv_block, int64_t stride_kv_row, CUtensorMap (&maps)[2]) { + maps[0] = make_kv_quant_part_tensor_map, RAW_TOKEN_SMEM_STRIDE>( + name, k_ptr, num_blocks, stride_kv_block, stride_kv_row, 0); + maps[1] = make_kv_quant_part_tensor_map, RAW_TOKEN_SMEM_STRIDE>( + name, k_ptr, num_blocks, stride_kv_block, stride_kv_row, QUANT_PART_CTA_OFFSET); }; - std::tie(tensor_map_kv_nope, tensor_map_kv_rope) = get_kv_tensormap(false, params.kv, params.num_blocks, params.stride_kv_block, params.stride_kv_row); - if (params.extra_topk > 0) - std::tie(tensor_map_extra_kv_nope, tensor_map_extra_kv_rope) = get_kv_tensormap(true, params.extra_kv, params.extra_num_blocks, params.stride_extra_kv_block, params.stride_extra_kv_row); + make_quant_part_maps(OrigKVFormat{}, "k_cache", params.kv, params.num_blocks, params.stride_kv_block, params.stride_kv_row, tensor_map_kv_quant_part); + if constexpr (D_BF16 > 0) { + tensor_map_kv_bf16_part = make_kv_bf16_part_tensor_map(params.kv, params.num_blocks, params.stride_kv_block); + } + if (params.extra_topk > 0) { + make_quant_part_maps(ExtraKVFormat{}, "extra_k_cache", params.extra_kv, params.extra_num_blocks, params.stride_extra_kv_block, params.stride_extra_kv_row, tensor_map_extra_kv_quant_part); + if constexpr (ExtraKVFormat::D_BF16 > 0) { + tensor_map_extra_kv_bf16_part = make_kv_bf16_part_tensor_map(params.extra_kv, params.extra_num_blocks, params.stride_extra_kv_block); + } + } } else { tensor_map_kv = ku::make_tensor_map( {D_QK, (unsigned long)params.s_kv}, @@ -1042,7 +1077,6 @@ void KernelTemplate::run(const ArgT& params) { ); } - CUtensorMap tensor_map_o_accum = {}; if constexpr (FWD_MODE == FwdMode::DecodeWithSplitKV) { tensor_map_o_accum = ku::make_tensor_map( @@ -1062,10 +1096,10 @@ void KernelTemplate::run(const ArgT& params) { tensor_map_q, tensor_map_o, tensor_map_o_accum, - tensor_map_kv_nope, - tensor_map_kv_rope, - tensor_map_extra_kv_nope, - tensor_map_extra_kv_rope + {tensor_map_kv_quant_part[0], tensor_map_kv_quant_part[1]}, + tensor_map_kv_bf16_part, + {tensor_map_extra_kv_quant_part[0], tensor_map_extra_kv_quant_part[1]}, + tensor_map_extra_kv_bf16_part }; } else { tma_params = { @@ -1075,7 +1109,13 @@ void KernelTemplate::run(const ArgT& params) { }; } - auto kernel = IS_PREFILL ? &sparse_attn_fwd_for_small_topk_kernel> : &flash_fwd_splitkv_mla_fp8_sparse_kernel>; + using KT = KernelTemplate; + void (*kernel)(__grid_constant__ const typename KT::ArgT, __grid_constant__ const typename KT::TmaParams); + if constexpr (IS_PREFILL) { + kernel = &sparse_attn_fwd_for_small_topk_kernel; + } else { + kernel = &flash_fwd_splitkv_mla_fp8_sparse_kernel; + } constexpr size_t smem_size = sizeof(SharedMemoryPlan); KU_CUDA_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); @@ -1098,9 +1138,9 @@ void KernelTemplate::run(const ArgT& params) { )); } -template -void run_fwd_for_small_topk_phase1_kernel(const SparseFwdArgT& params) { - using Kernel = KernelTemplate; +template +void run_sparse_fwd_for_small_topk_phase1_kernel(const SparseFwdArgT& params) { + using Kernel = KernelTemplate; Kernel::run(params); } diff --git a/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/phase1.h b/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/phase1.h new file mode 100644 index 00000000..7ce4a3ee --- /dev/null +++ b/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/phase1.h @@ -0,0 +1,10 @@ +#pragma once + +#include "kernels/params.h" + +namespace sm100::prefill::sparse_fwd_for_small_topk::head128 { + +template +void run_sparse_fwd_for_small_topk_phase1_kernel(const SparseFwdArgT& params); + +} diff --git a/csrc/sm90/decode/dense/config.h b/csrc/kernels/sm90/decode/dense/config.h similarity index 84% rename from csrc/sm90/decode/dense/config.h rename to csrc/kernels/sm90/decode/dense/config.h index e97e0bca..76f8606b 100644 --- a/csrc/sm90/decode/dense/config.h +++ b/csrc/kernels/sm90/decode/dense/config.h @@ -1,5 +1,7 @@ #pragma once +namespace sm90::decode::dense { + namespace Config { static constexpr int BLOCK_SIZE_M = 64; @@ -9,3 +11,5 @@ static constexpr int HEAD_DIM_K = 576; static constexpr int HEAD_DIM_V = 512; } + +} diff --git a/csrc/sm90/decode/dense/instantiations/bf16.cu b/csrc/kernels/sm90/decode/dense/instantiations/bf16.cu similarity index 83% rename from csrc/sm90/decode/dense/instantiations/bf16.cu rename to csrc/kernels/sm90/decode/dense/instantiations/bf16.cu index 3a1dce9b..0a819362 100644 --- a/csrc/sm90/decode/dense/instantiations/bf16.cu +++ b/csrc/kernels/sm90/decode/dense/instantiations/bf16.cu @@ -1,7 +1,7 @@ #include "../splitkv_mla.cuh" #include "../splitkv_mla.h" -namespace sm90 { +namespace sm90::decode::dense { template void run_flash_splitkv_mla_kernel(DenseAttnDecodeParams ¶ms); diff --git a/csrc/sm90/decode/dense/instantiations/fp16.cu b/csrc/kernels/sm90/decode/dense/instantiations/fp16.cu similarity index 85% rename from csrc/sm90/decode/dense/instantiations/fp16.cu rename to csrc/kernels/sm90/decode/dense/instantiations/fp16.cu index bc6cd648..43adc644 100644 --- a/csrc/sm90/decode/dense/instantiations/fp16.cu +++ b/csrc/kernels/sm90/decode/dense/instantiations/fp16.cu @@ -1,7 +1,7 @@ #include "../splitkv_mla.cuh" #include "../splitkv_mla.h" -namespace sm90 { +namespace sm90::decode::dense { #ifndef FLASH_MLA_DISABLE_FP16 template void run_flash_splitkv_mla_kernel(DenseAttnDecodeParams ¶ms); diff --git a/csrc/sm90/decode/dense/splitkv_mla.cuh b/csrc/kernels/sm90/decode/dense/splitkv_mla.cuh similarity index 99% rename from csrc/sm90/decode/dense/splitkv_mla.cuh rename to csrc/kernels/sm90/decode/dense/splitkv_mla.cuh index cdd54413..bdd3c789 100644 --- a/csrc/sm90/decode/dense/splitkv_mla.cuh +++ b/csrc/kernels/sm90/decode/dense/splitkv_mla.cuh @@ -1,15 +1,15 @@ #include -#include "utils.h" +#include "kernels/utils.h" -#include "params.h" +#include "kernels/params.h" #include "config.h" #include "traits.h" using namespace cute; using cutlass::arch::NamedBarrier; -namespace sm90 { +namespace sm90::decode::dense { // Here we use MAX_INIT_VAL_SM to initialize sM, and MAX_INIT_VAL for masking // The reason is that, we need to calculate new_max = max(sM(row_idx), cur_max*scale_softmax_log2) diff --git a/csrc/sm90/decode/dense/splitkv_mla.h b/csrc/kernels/sm90/decode/dense/splitkv_mla.h similarity index 64% rename from csrc/sm90/decode/dense/splitkv_mla.h rename to csrc/kernels/sm90/decode/dense/splitkv_mla.h index b2c50c8a..d1f89047 100644 --- a/csrc/sm90/decode/dense/splitkv_mla.h +++ b/csrc/kernels/sm90/decode/dense/splitkv_mla.h @@ -1,8 +1,8 @@ #pragma once -#include "params.h" +#include "kernels/params.h" -namespace sm90 { +namespace sm90::decode::dense { template void run_flash_splitkv_mla_kernel(DenseAttnDecodeParams ¶ms); diff --git a/csrc/sm90/decode/dense/traits.h b/csrc/kernels/sm90/decode/dense/traits.h similarity index 99% rename from csrc/sm90/decode/dense/traits.h rename to csrc/kernels/sm90/decode/dense/traits.h index 5f915a68..1f3eaa95 100644 --- a/csrc/sm90/decode/dense/traits.h +++ b/csrc/kernels/sm90/decode/dense/traits.h @@ -7,6 +7,8 @@ #include "config.h" +namespace sm90::decode::dense { + using TMABarrier = cutlass::arch::ClusterTransactionBarrier; using namespace cute; @@ -105,3 +107,5 @@ enum NamedBarriers : int { rO1sP0sV0RIssued = 3, sMInitialized = 4, }; + +} diff --git a/csrc/sm90/decode/sparse_fp8/components/config.h b/csrc/kernels/sm90/decode/sparse/components/config.h similarity index 90% rename from csrc/sm90/decode/sparse_fp8/components/config.h rename to csrc/kernels/sm90/decode/sparse/components/config.h index f38915ba..897cb258 100644 --- a/csrc/sm90/decode/sparse_fp8/components/config.h +++ b/csrc/kernels/sm90/decode/sparse/components/config.h @@ -3,11 +3,11 @@ #include #include #include -#include "defines.h" +#include "kernels/defines.h" using namespace cute; -namespace sm90::decode::sparse_fp8 { +namespace sm90::decode::sparse { static constexpr int HEAD_DIM_K = 576; static constexpr int HEAD_DIM_V = 512; diff --git a/csrc/sm90/decode/sparse_fp8/components/dequant.h b/csrc/kernels/sm90/decode/sparse/components/dequant.h similarity index 98% rename from csrc/sm90/decode/sparse_fp8/components/dequant.h rename to csrc/kernels/sm90/decode/sparse/components/dequant.h index 0c4022d6..13240ea7 100644 --- a/csrc/sm90/decode/sparse_fp8/components/dequant.h +++ b/csrc/kernels/sm90/decode/sparse/components/dequant.h @@ -3,9 +3,9 @@ #include #include -#include "defines.h" +#include "kernels/defines.h" -namespace sm90::decode::sparse_fp8 { +namespace sm90::decode::sparse { struct fp8x8 { __nv_fp8x4_e4m3 lo; diff --git a/csrc/sm90/decode/sparse_fp8/components/helpers.h b/csrc/kernels/sm90/decode/sparse/components/helpers.h similarity index 99% rename from csrc/sm90/decode/sparse_fp8/components/helpers.h rename to csrc/kernels/sm90/decode/sparse/components/helpers.h index d47e4922..5d1c1d29 100644 --- a/csrc/sm90/decode/sparse_fp8/components/helpers.h +++ b/csrc/kernels/sm90/decode/sparse/components/helpers.h @@ -7,7 +7,7 @@ using namespace cute; -namespace sm90::decode::sparse_fp8 { +namespace sm90::decode::sparse { // In the layout of fragment A and fragment C during WGMMA, data each thread holds resides in two particular rows. This function converts the local_row_idx (0~1) to the actual row_idx // You may refer to this link for the detailed layout: https://docs.nvidia.com/cuda/parallel-thread-execution/#wgmma-64n16-a diff --git a/csrc/sm90/decode/sparse_fp8/config.h b/csrc/kernels/sm90/decode/sparse/config.h similarity index 98% rename from csrc/sm90/decode/sparse_fp8/config.h rename to csrc/kernels/sm90/decode/sparse/config.h index e5631f31..2fe2f2c3 100644 --- a/csrc/sm90/decode/sparse_fp8/config.h +++ b/csrc/kernels/sm90/decode/sparse/config.h @@ -5,12 +5,12 @@ #include #include -#include "defines.h" -#include "params.h" +#include "kernels/defines.h" +#include "kernels/params.h" using namespace cute; -namespace sm90::decode::sparse_fp8 { +namespace sm90::decode::sparse { template class KernelTemplate { @@ -26,7 +26,7 @@ static constexpr int HEAD_DIM_ROPE = 64; static constexpr int HEAD_DIM_NOPE = HEAD_DIM_K - HEAD_DIM_ROPE; static constexpr int QUANT_TILE_SIZE = MODEL_TYPE == ModelType::V32 ? 128 : 64; -static constexpr int NUM_SCALES = MODEL_TYPE == ModelType::V32 ? 4 : 8; // For MODEL1: 7 fp8_e4m3 + 1 padding +static constexpr int NUM_SCALES = MODEL_TYPE == ModelType::V32 ? 4 : 8; // For DeepSeek-V4: 7 fp8_e4m3 + 1 padding static constexpr int NUM_THREADS = 128*3; static constexpr int BLOCK_M = 64; diff --git a/csrc/sm90/decode/sparse_fp8/instantiations/v32_persistent_h128.cu b/csrc/kernels/sm90/decode/sparse/instantiations/v32_persistent_h128.cu similarity index 80% rename from csrc/sm90/decode/sparse_fp8/instantiations/v32_persistent_h128.cu rename to csrc/kernels/sm90/decode/sparse/instantiations/v32_persistent_h128.cu index 97276426..44ce6c2c 100644 --- a/csrc/sm90/decode/sparse_fp8/instantiations/v32_persistent_h128.cu +++ b/csrc/kernels/sm90/decode/sparse/instantiations/v32_persistent_h128.cu @@ -1,6 +1,6 @@ #include "../splitkv_mla.cuh" -namespace sm90::decode::sparse_fp8 { +namespace sm90::decode::sparse { template void run_flash_splitkv_mla_fp8_sparse_kernel(const SparseAttnDecodeParams ¶ms); diff --git a/csrc/sm90/decode/sparse_fp8/instantiations/v32_persistent_h64.cu b/csrc/kernels/sm90/decode/sparse/instantiations/v32_persistent_h64.cu similarity index 80% rename from csrc/sm90/decode/sparse_fp8/instantiations/v32_persistent_h64.cu rename to csrc/kernels/sm90/decode/sparse/instantiations/v32_persistent_h64.cu index f7a3f19a..a0fb6ed1 100644 --- a/csrc/sm90/decode/sparse_fp8/instantiations/v32_persistent_h64.cu +++ b/csrc/kernels/sm90/decode/sparse/instantiations/v32_persistent_h64.cu @@ -1,6 +1,6 @@ #include "../splitkv_mla.cuh" -namespace sm90::decode::sparse_fp8 { +namespace sm90::decode::sparse { template void run_flash_splitkv_mla_fp8_sparse_kernel(const SparseAttnDecodeParams ¶ms); diff --git a/csrc/sm90/decode/sparse_fp8/instantiations/model1_persistent_h128.cu b/csrc/kernels/sm90/decode/sparse/instantiations/v4_persistent_h128.cu similarity index 52% rename from csrc/sm90/decode/sparse_fp8/instantiations/model1_persistent_h128.cu rename to csrc/kernels/sm90/decode/sparse/instantiations/v4_persistent_h128.cu index af5058f7..1e0267ca 100644 --- a/csrc/sm90/decode/sparse_fp8/instantiations/model1_persistent_h128.cu +++ b/csrc/kernels/sm90/decode/sparse/instantiations/v4_persistent_h128.cu @@ -1,7 +1,7 @@ #include "../splitkv_mla.cuh" -namespace sm90::decode::sparse_fp8 { +namespace sm90::decode::sparse { -template void run_flash_splitkv_mla_fp8_sparse_kernel(const SparseAttnDecodeParams ¶ms); +template void run_flash_splitkv_mla_fp8_sparse_kernel(const SparseAttnDecodeParams ¶ms); } diff --git a/csrc/sm90/decode/sparse_fp8/instantiations/model1_persistent_h64.cu b/csrc/kernels/sm90/decode/sparse/instantiations/v4_persistent_h64.cu similarity index 52% rename from csrc/sm90/decode/sparse_fp8/instantiations/model1_persistent_h64.cu rename to csrc/kernels/sm90/decode/sparse/instantiations/v4_persistent_h64.cu index 902a5910..3e63a575 100644 --- a/csrc/sm90/decode/sparse_fp8/instantiations/model1_persistent_h64.cu +++ b/csrc/kernels/sm90/decode/sparse/instantiations/v4_persistent_h64.cu @@ -1,8 +1,8 @@ #include "../splitkv_mla.cuh" -namespace sm90::decode::sparse_fp8 { +namespace sm90::decode::sparse { -template void run_flash_splitkv_mla_fp8_sparse_kernel(const SparseAttnDecodeParams ¶ms); +template void run_flash_splitkv_mla_fp8_sparse_kernel(const SparseAttnDecodeParams ¶ms); } diff --git a/csrc/sm90/decode/sparse_fp8/splitkv_mla.cuh b/csrc/kernels/sm90/decode/sparse/splitkv_mla.cuh similarity index 98% rename from csrc/sm90/decode/sparse_fp8/splitkv_mla.cuh rename to csrc/kernels/sm90/decode/sparse/splitkv_mla.cuh index 99945689..523afab4 100644 --- a/csrc/sm90/decode/sparse_fp8/splitkv_mla.cuh +++ b/csrc/kernels/sm90/decode/sparse/splitkv_mla.cuh @@ -11,13 +11,13 @@ #include -#include "utils.h" +#include "kernels/utils.h" #include "components/dequant.h" #include "components/helpers.h" #include "config.h" using namespace cute; -namespace sm90::decode::sparse_fp8 { +namespace sm90::decode::sparse { static constexpr float MAX_INIT_VAL = -1e30; // Prevent (-inf) - (-inf) = nan using cutlass::arch::fence_view_async_shared; @@ -158,7 +158,7 @@ __device__ void KernelTemplate::devfunc(const SparseAttnD int start_block_idx, end_block_idx; bool is_no_split; - // The following fields are only valid for MODEL1 + // The following fields are only valid for DeepSeek-V4 int topk_length, extra_topk_length, num_orig_kv_blocks; }; auto get_cur_req_info = [&](int batch_idx) -> MainloopArgs { @@ -528,8 +528,8 @@ __device__ void KernelTemplate::devfunc(const SparseAttnD nxt_token_indexs[round] = __ldg(gExtraIndices + (block_idx+1-args.num_orig_kv_blocks)*TOPK_BLOCK_SIZE + idx_in_cluster*(TOPK_BLOCK_SIZE/2) + my_token_idx); } - if constexpr (MODEL_TYPE == ModelType::MODEL1) { - // For MODEL1, we need to check whether the token_index is within topk_length + if constexpr (MODEL_TYPE == ModelType::V4) { + // For DeepSeek-V4, we need to check whether the token_index is within topk_length if (rel_block_idx*TOPK_BLOCK_SIZE + idx_in_cluster*(TOPK_BLOCK_SIZE/2) + my_token_idx >= topk_length) { token_index = -1; // To prevent IMA when we have invalid (e.g. INT_MAX) topk indexes outside topk_length } @@ -581,7 +581,7 @@ __device__ void KernelTemplate::devfunc(const SparseAttnD } CUTE_UNROLL for (int dim_idx = 0; dim_idx < HEAD_DIM_NOPE/64; dim_idx += 1) { - fp8x16 cur_fp8x16 = load_128b_from_gmem(gK_nope + dim_idx*64); // We use EVICT_LAST here since gK_base may not be aligned to 32B (for V3.2) and the performance is the best among all cache hints (for MODEL1) + fp8x16 cur_fp8x16 = load_128b_from_gmem(gK_nope + dim_idx*64); // We use EVICT_LAST here since gK_base may not be aligned to 32B (for V3.2) and the performance is the best among all cache hints (for DeepSeek-V4) bf16 scale = scales[MODEL_TYPE == ModelType::V32 ? dim_idx/2 : dim_idx]; auto dequant_and_save_bf16x8 = [&](const fp8x8 &data, int offset) { int smem_offset = (dim_idx*64 + offset) * TOPK_BLOCK_SIZE; @@ -690,11 +690,11 @@ void KernelTemplate::run(const SparseAttnDecodeParams &pa KU_ASSERT(params.d_qk == HEAD_DIM_K); KU_ASSERT(params.d_v == HEAD_DIM_V); KU_ASSERT(params.h_q % BLOCK_M == 0); - if constexpr (MODEL_TYPE == ModelType::MODEL1) { + if constexpr (MODEL_TYPE == ModelType::V4) { constexpr int BYTES_PER_TOKEN = HEAD_DIM_NOPE + 2*HEAD_DIM_ROPE + 8; - KU_ASSERT(params.stride_kv_row == BYTES_PER_TOKEN, "Each page block in KV cache must be contiguous for head64 sparse fp8 decoding attention in MODEL1"); // Each block must be contiguous + KU_ASSERT(params.stride_kv_row == BYTES_PER_TOKEN, "Each page block in KV cache must be contiguous for head64 sparse fp8 decoding attention in DeepSeek-V4"); // Each block must be contiguous if (params.extra_kv != nullptr) { - KU_ASSERT(params.stride_extra_kv_row == BYTES_PER_TOKEN, "Each page block in extra KV cache must be contiguous for head64 sparse fp8 decoding attention in MODEL1"); // Each block must be contiguous + KU_ASSERT(params.stride_extra_kv_row == BYTES_PER_TOKEN, "Each page block in extra KV cache must be contiguous for head64 sparse fp8 decoding attention in DeepSeek-V4"); // Each block must be contiguous } } else { KU_ASSERT(params.extra_kv == nullptr, "V3.2 does not support extra KV cache"); diff --git a/csrc/sm90/decode/sparse_fp8/splitkv_mla.h b/csrc/kernels/sm90/decode/sparse/splitkv_mla.h similarity index 71% rename from csrc/sm90/decode/sparse_fp8/splitkv_mla.h rename to csrc/kernels/sm90/decode/sparse/splitkv_mla.h index 13b659be..f69f257c 100644 --- a/csrc/sm90/decode/sparse_fp8/splitkv_mla.h +++ b/csrc/kernels/sm90/decode/sparse/splitkv_mla.h @@ -1,8 +1,8 @@ #pragma once -#include "params.h" +#include "kernels/params.h" -namespace sm90::decode::sparse_fp8 { +namespace sm90::decode::sparse { template void run_flash_splitkv_mla_fp8_sparse_kernel(const SparseAttnDecodeParams ¶ms); diff --git a/csrc/sm90/helpers.h b/csrc/kernels/sm90/helpers.h similarity index 100% rename from csrc/sm90/helpers.h rename to csrc/kernels/sm90/helpers.h diff --git a/csrc/sm90/prefill/sparse/config.h b/csrc/kernels/sm90/prefill/sparse/config.h similarity index 95% rename from csrc/sm90/prefill/sparse/config.h rename to csrc/kernels/sm90/prefill/sparse/config.h index 75005664..19deccbe 100644 --- a/csrc/sm90/prefill/sparse/config.h +++ b/csrc/kernels/sm90/prefill/sparse/config.h @@ -8,10 +8,10 @@ #include #include -#include "defines.h" -#include "params.h" +#include "kernels/defines.h" +#include "kernels/params.h" -namespace sm90::fwd { +namespace sm90::prefill::sparse_fwd { using namespace cute; @@ -72,7 +72,7 @@ struct SharedMemoryPlan { array_aligned> o; } q_o; array_aligned> k[2]; - array_aligned> s[D_QK == 576 ? 1 : 2]; // For V3.2 (whose D_QK is 576), we overlap sS[0] with k's RoPE part to save shared memory; For MODEL1 (whose D_QK is 512), we allocate two buffers + array_aligned> s[D_QK == 576 ? 1 : 2]; // For V3.2 (whose D_QK is 576), we overlap sS[0] with k's RoPE part to save shared memory; For DeepSeek-V4 (whose D_QK is 512), we allocate two buffers bool is_kv_valid[2][B_TOPK]; float2 sM[32]; diff --git a/csrc/sm90/prefill/sparse/instantiations/phase1_k512.cu b/csrc/kernels/sm90/prefill/sparse/instantiations/phase1_k512.cu similarity index 88% rename from csrc/sm90/prefill/sparse/instantiations/phase1_k512.cu rename to csrc/kernels/sm90/prefill/sparse/instantiations/phase1_k512.cu index 046cfb39..dd0ba23a 100644 --- a/csrc/sm90/prefill/sparse/instantiations/phase1_k512.cu +++ b/csrc/kernels/sm90/prefill/sparse/instantiations/phase1_k512.cu @@ -1,7 +1,7 @@ #include "../phase1.h" #include "../phase1.cuh" -namespace sm90::fwd { +namespace sm90::prefill::sparse_fwd { // NOTE (intlsy): We instantiate run_fwd_phase1_kernel in two .cu files as functions with HAVE_TOPK_LENGTH // = true / false respectively, to compile them in parallel. diff --git a/csrc/sm90/prefill/sparse/instantiations/phase1_k512_topklen.cu b/csrc/kernels/sm90/prefill/sparse/instantiations/phase1_k512_topklen.cu similarity index 88% rename from csrc/sm90/prefill/sparse/instantiations/phase1_k512_topklen.cu rename to csrc/kernels/sm90/prefill/sparse/instantiations/phase1_k512_topklen.cu index 45da9955..ef0da757 100644 --- a/csrc/sm90/prefill/sparse/instantiations/phase1_k512_topklen.cu +++ b/csrc/kernels/sm90/prefill/sparse/instantiations/phase1_k512_topklen.cu @@ -1,7 +1,7 @@ #include "../phase1.h" #include "../phase1.cuh" -namespace sm90::fwd { +namespace sm90::prefill::sparse_fwd { // NOTE (intlsy): We instantiate run_fwd_phase1_kernel in two .cu files as functions with HAVE_TOPK_LENGTH // = true / false respectively, to compile them in parallel. diff --git a/csrc/sm90/prefill/sparse/instantiations/phase1_k576.cu b/csrc/kernels/sm90/prefill/sparse/instantiations/phase1_k576.cu similarity index 78% rename from csrc/sm90/prefill/sparse/instantiations/phase1_k576.cu rename to csrc/kernels/sm90/prefill/sparse/instantiations/phase1_k576.cu index f35db006..72d7866b 100644 --- a/csrc/sm90/prefill/sparse/instantiations/phase1_k576.cu +++ b/csrc/kernels/sm90/prefill/sparse/instantiations/phase1_k576.cu @@ -1,7 +1,7 @@ #include "../phase1.h" #include "../phase1.cuh" -namespace sm90::fwd { +namespace sm90::prefill::sparse_fwd { template void run_fwd_phase1_kernel<576, false>(const SparseAttnFwdParams& params); diff --git a/csrc/sm90/prefill/sparse/instantiations/phase1_k576_topklen.cu b/csrc/kernels/sm90/prefill/sparse/instantiations/phase1_k576_topklen.cu similarity index 78% rename from csrc/sm90/prefill/sparse/instantiations/phase1_k576_topklen.cu rename to csrc/kernels/sm90/prefill/sparse/instantiations/phase1_k576_topklen.cu index bd0f0ca6..47d4d6eb 100644 --- a/csrc/sm90/prefill/sparse/instantiations/phase1_k576_topklen.cu +++ b/csrc/kernels/sm90/prefill/sparse/instantiations/phase1_k576_topklen.cu @@ -1,7 +1,7 @@ #include "../phase1.h" #include "../phase1.cuh" -namespace sm90::fwd { +namespace sm90::prefill::sparse_fwd { template void run_fwd_phase1_kernel<576, true>(const SparseAttnFwdParams& params); diff --git a/csrc/sm90/prefill/sparse/phase1.cuh b/csrc/kernels/sm90/prefill/sparse/phase1.cuh similarity index 99% rename from csrc/sm90/prefill/sparse/phase1.cuh rename to csrc/kernels/sm90/prefill/sparse/phase1.cuh index bf2fff84..201a59f1 100644 --- a/csrc/sm90/prefill/sparse/phase1.cuh +++ b/csrc/kernels/sm90/prefill/sparse/phase1.cuh @@ -2,10 +2,10 @@ #include "config.h" -#include "utils.h" +#include "kernels/utils.h" #include "../../helpers.h" -namespace sm90::fwd { +namespace sm90::prefill::sparse_fwd { using namespace cute; diff --git a/csrc/sm90/prefill/sparse/phase1.h b/csrc/kernels/sm90/prefill/sparse/phase1.h similarity index 65% rename from csrc/sm90/prefill/sparse/phase1.h rename to csrc/kernels/sm90/prefill/sparse/phase1.h index c315b2b6..c1a60326 100644 --- a/csrc/sm90/prefill/sparse/phase1.h +++ b/csrc/kernels/sm90/prefill/sparse/phase1.h @@ -1,8 +1,8 @@ #pragma once -#include "../../../params.h" +#include "kernels/params.h" -namespace sm90::fwd { +namespace sm90::prefill::sparse_fwd { template void run_fwd_phase1_kernel(const SparseAttnFwdParams& params); diff --git a/csrc/smxx/decode/combine/combine.cu b/csrc/kernels/smxx/decode/combine/combine.cu similarity index 99% rename from csrc/smxx/decode/combine/combine.cu rename to csrc/kernels/smxx/decode/combine/combine.cu index cb6531a6..7bd34da3 100644 --- a/csrc/smxx/decode/combine/combine.cu +++ b/csrc/kernels/smxx/decode/combine/combine.cu @@ -8,8 +8,8 @@ #include -#include "params.h" -#include "utils.h" +#include "kernels/params.h" +#include "kernels/utils.h" using namespace cute; diff --git a/csrc/smxx/decode/combine/combine.h b/csrc/kernels/smxx/decode/combine/combine.h similarity index 82% rename from csrc/smxx/decode/combine/combine.h rename to csrc/kernels/smxx/decode/combine/combine.h index 0ea21fde..93a543b6 100644 --- a/csrc/smxx/decode/combine/combine.h +++ b/csrc/kernels/smxx/decode/combine/combine.h @@ -1,6 +1,6 @@ #pragma once -#include "params.h" +#include "kernels/params.h" namespace smxx::decode { diff --git a/csrc/smxx/decode/get_decoding_sched_meta/get_decoding_sched_meta.cu b/csrc/kernels/smxx/decode/get_decoding_sched_meta/get_decoding_sched_meta.cu similarity index 99% rename from csrc/smxx/decode/get_decoding_sched_meta/get_decoding_sched_meta.cu rename to csrc/kernels/smxx/decode/get_decoding_sched_meta/get_decoding_sched_meta.cu index 083da60c..27c141fe 100644 --- a/csrc/smxx/decode/get_decoding_sched_meta/get_decoding_sched_meta.cu +++ b/csrc/kernels/smxx/decode/get_decoding_sched_meta/get_decoding_sched_meta.cu @@ -4,7 +4,7 @@ #include #include -#include "utils.h" +#include "kernels/utils.h" namespace smxx::decode { diff --git a/csrc/smxx/decode/get_decoding_sched_meta/get_decoding_sched_meta.h b/csrc/kernels/smxx/decode/get_decoding_sched_meta/get_decoding_sched_meta.h similarity index 80% rename from csrc/smxx/decode/get_decoding_sched_meta/get_decoding_sched_meta.h rename to csrc/kernels/smxx/decode/get_decoding_sched_meta/get_decoding_sched_meta.h index 0b1c288e..1c183211 100644 --- a/csrc/smxx/decode/get_decoding_sched_meta/get_decoding_sched_meta.h +++ b/csrc/kernels/smxx/decode/get_decoding_sched_meta/get_decoding_sched_meta.h @@ -1,6 +1,6 @@ #pragma once -#include "params.h" +#include "kernels/params.h" namespace smxx::decode { diff --git a/csrc/utils.h b/csrc/kernels/utils.h similarity index 100% rename from csrc/utils.h rename to csrc/kernels/utils.h diff --git a/csrc/kerutils/include/kerutils/device/common.h b/csrc/kerutils/include/kerutils/device/common.h index d13e7a90..05d671ee 100644 --- a/csrc/kerutils/include/kerutils/device/common.h +++ b/csrc/kerutils/include/kerutils/device/common.h @@ -61,10 +61,15 @@ static_assert(false, "kerutils doesn't support SM architectures below SM80"); #define KERUTILS_ENABLE_SM100A #endif +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1030 && __CUDA_ARCH__ < 1200)) +#define KERUTILS_ENABLE_SM103A +#endif + #if (defined(__CLION_IDE__) || defined(__VSCODE_IDE__)) #define KERUTILS_ENABLE_SM80 #define KERUTILS_ENABLE_SM90 #define KERUTILS_ENABLE_SM90A #define KERUTILS_ENABLE_SM100 #define KERUTILS_ENABLE_SM100A +#define KERUTILS_ENABLE_SM103A #endif diff --git a/csrc/kerutils/include/kerutils/device/sm100/gemm.cuh b/csrc/kerutils/include/kerutils/device/sm100/gemm.cuh index 8af4edcf..2500d80c 100644 --- a/csrc/kerutils/include/kerutils/device/sm100/gemm.cuh +++ b/csrc/kerutils/include/kerutils/device/sm100/gemm.cuh @@ -10,14 +10,82 @@ namespace cute { // CuTe don't support UTCMMA with .ws, so we add it here // Besides, CuTe's UTCMMA has an `elect_one_sync()` inside which is really disgusting, so we have our own variant without `elect_one_sync()` here +namespace UMMA { + +// 在双轨矩乘（或者其他用况）中，我们有时会想自己指定 tensor memory 的“折叠度”，也即，如果 M < 128，那么是在不同的 lane 上将数据进行多次复制，还是将 N 切成 128/M 份，使用 N/(128/M) 的 lane，还是某种介于二者之间的做法。 +// 因此我在这里实现了这个 `tmem_frg_elastic`，其支持自定义 FOLD_DEGREE。最终这个 fragment 会占用 N/FOLD_DEGREE 个 TMEM col（等价于将数据在 TMEM 中复制 (128/M) / FOLD_DEGREE 份。 +// 在正常情况（NV 官方的用法）中，由于 A 总是在不同的 Tensor Memory Row 之间复制，因此 FOLD_DEGREE = 1；在多轨矩乘情况下，FOLD_DEGREE 为轨道的数量。 +template +struct tmem_frg_elastic : tmem_frg_base { + static_assert(sizeof_bits_v <= sizeof_bits_v, "TMEM MMA allocations require StorageType big enough for ValueType."); + + // UMMA TMEM Allocator + // Each UMMA expects a specific MxN layout of TMEM for accumulators + // and sometimes a specific MxK layout of TMEM for A-values. + // @tparam ValueType The value type of the TMEM Tensor to allocate. + // @tparam StorageType The storage type of the TMEM Tensor to allocate. + // "Sparse" allocations often allocate ValueType=half_t within StorageType=uint32_t. + // "Dense" allocations often allocate ValueType=half_t within StorageType=half_t. + // @param tmem_shape ((M_MMA_SM,N_MMA_SM),MMA_M,MMA_N,...) + // The post-MMA-partitioned shape of TMEM to allocate. + // Note for UMMA_2SM_128xNx16, that M_MMA_SM will be 64, for example. + template + CUTE_HOST_DEVICE constexpr static auto + make(TmemShape const& tmem_shape) + { + CUTE_STATIC_ASSERT_V(size(tmem_shape)*Int)>{} <= TMEM::MAX_CAPACITY_BITS{}, + "Requesting more TMEM than is available."); + CUTE_STATIC_ASSERT_V(rank<0>(tmem_shape) == Int<2>{}, "Expected post-partitioned shape ((M_MMA,N_MMA),...)."); + constexpr int R = decltype(rank(tmem_shape))::value; + constexpr int M_MMA = decltype(size<0,0>(tmem_shape))::value; + constexpr int N_MMA = decltype(size<0,1>(tmem_shape))::value; + + // It's convenient to use "virtual tensor memory addressing" + // with DP_STRIDE=1, COL_STRIDE=128 to define the tmem_atom, + // then convert to "logical tensor memory addressing" on return. + using COL_ADDR = C::value / sizeof_bits::value>; + Layout tmem_restride = Layout, + Stride, COL_ADDR>>{}; + Layout tmem_atom = Layout, Shape, Int>>, + Stride<_1, Stride<_128, Int<128/FOLD_DEGREE>> >>{}; + constexpr int tile_stride = (128/M_MMA) / FOLD_DEGREE; + Layout tmem_logical_layout = tiled_product(tmem_atom, make_layout(take<1, R>(tmem_shape), + compact_col_major(take<1,R>(tmem_shape),Int{}))); + + return make_tensor(make_tmem_ptr(), composition(tmem_restride, tmem_logical_layout)); + } +}; + +} + +template +struct MakeTensor> { + template + CUTE_HOST_DEVICE constexpr auto + operator()(Shape const& tmem_shape) { + return UMMA::tmem_frg_elastic::make(shape(tmem_shape)); + } +}; + +template +struct MakeTensor> { + template + CUTE_HOST_DEVICE constexpr auto + operator()(Shape const& tmem_shape) { + return UMMA::tmem_frg_elastic::make(shape(tmem_shape)); + } +}; + + template struct SM100_MMA_F16BF16_WS_TS_NOELECT { static_assert(M == 32 || M == 64 || M == 128, "SM100_MMA_F16BF16_WS_TS_NOELECT M-mode size should be 32, 64 or 128 for 1 CTA cluster MMA."); static_assert(N == 64 || N == 128 || N == 256, - "SM100_MMA_F16BF16_WS_TS_NOELECT N-mode size should be 32, 64 or 128"); + "SM100_MMA_F16BF16_WS_TS_NOELECT N-mode size should be 64, 128 or 256"); using DRegisters = void; using ARegisters = uint64_t[1]; @@ -44,10 +112,12 @@ struct SM100_MMA_F16BF16_WS_TS_NOELECT template struct MMA_Traits> { using ValTypeD = c_type; @@ -56,9 +126,10 @@ struct MMA_Traits == cute::sizeof_bits_v && cute::sizeof_bits_v == 16, "SM100_MMA_F16BF16_WS_TS_NOELECT supports 16bit types"); - using FrgTypeA = UMMA::tmem_frg_1sm; // Actually this should be "duplicated", however, our great CuTe doesn't allow us to set it to "duplicated", so we just set it to NonInterleaved for a correct address calculation + // using FrgTypeA = UMMA::tmem_frg_1sm; // Actually this should be "duplicated", however, our great CuTe doesn't allow us to set it to "duplicated", so we just set it to NonInterleaved for a correct address calculation + using FrgTypeA = UMMA::tmem_frg_elastic; using FrgTypeB = UMMA::smem_desc; - using FrgTypeC = UMMA::tmem_frg_ws_1sm; + using FrgTypeC = UMMA::tmem_frg_elastic; // Logical shape-K is always 256 bits; transform to units of elements static constexpr int K = 256 / cute::sizeof_bits::value; @@ -103,6 +174,7 @@ struct MMA_Traits::fma(tmem_a, desc_b, tmem_c, uint32_t(traits.accumulate_), idesc); } }; @@ -114,7 +186,7 @@ struct SM100_MMA_F16BF16_WS_SS_NOELECT { static_assert(M == 32 || M == 64 || M == 128, "SM100_MMA_F16BF16_WS_SS_NOELECT M-mode size should be 32, 64 or 128 for 1 CTA cluster MMA."); static_assert(N == 64 || N == 128 || N == 256, - "SM100_MMA_F16BF16_WS_SS_NOELECT N-mode size should be 32, 64 or 128"); + "SM100_MMA_F16BF16_WS_SS_NOELECT N-mode size should be 64, 128 or 256"); using DRegisters = void; using ARegisters = uint64_t[1]; @@ -155,7 +227,7 @@ struct MMA_Traits; using FrgTypeB = UMMA::smem_desc; - using FrgTypeC = UMMA::tmem_frg_ws_1sm; + using FrgTypeC = UMMA::tmem_frg_elastic; // Logical shape-K is always 256bits, transform to units of elements static constexpr int K = 256 / cute::sizeof_bits::value; @@ -425,10 +497,8 @@ template + #include "kerutils/device/common.h" namespace kerutils { @@ -298,6 +300,230 @@ void tmem_ld_32dp32bNx(uint32_t tmem_start, void* data_) { #endif } +// Load from tensor memory with column-wise reduction, 32 data path lanes, 32-bit pattern, repeated N times. +// (https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-instructions-tcgen05-ld) +// USE_MAX: true = .max reduction, false = .min reduction +// USE_ABS: true = .abs qualifier (use absolute values for reduction) +// USE_NAN: true = .NaN qualifier (propagate NaN in reduction) +template +__device__ __forceinline__ +void tmem_ld_red_32dp32bNx(uint32_t tmem_start, void* data_, float& redval) { + uint32_t* data = (uint32_t*)data_; + static_assert(kNumElements == 2 || kNumElements == 4 || kNumElements == 8 || + kNumElements == 16 || kNumElements == 32 || kNumElements == 64 || + kNumElements == 128, + "Invalid kNumElements for tcgen05.ld.red (must be power of 2, at least 2)"); +#ifndef __VSCODE_IDE__ + static constexpr char s_max[] = ".max"; + static constexpr char s_min[] = ".min"; + static constexpr char s_abs[] = ".abs"; + static constexpr char s_nan[] = ".NaN"; + static constexpr char s_empty[] = ""; + // Operand numbering for each branch: + // outputs: %0..%{N-1} = data[0..N-1] ("=r"), %N = redval ("=f") + // inputs: %{N+1} = tmem_start ("r"), %{N+2} = redop ("C"), %{N+3} = abs ("C"), %{N+4} = nan ("C") + if constexpr (kNumElements == 2) { + // outputs: %0..%1 = data, %2 = redval; inputs: %3 = taddr, %4 = redop, %5 = abs, %6 = nan + asm volatile( + "tcgen05.ld.red.sync.aligned.32x32b.x2%4%5%6.f32" + " {%0, %1}, %2, [%3];\n" + : "=r"(data[0]), "=r"(data[1]), + "=f"(redval) + : "r"(tmem_start), + "C"(USE_MAX ? s_max : s_min), + "C"(USE_ABS ? s_abs : s_empty), + "C"(USE_NAN ? s_nan : s_empty) + ); + } else if constexpr (kNumElements == 4) { + // outputs: %0..%3 = data, %4 = redval; inputs: %5 = taddr, %6..%8 = C + asm volatile( + "tcgen05.ld.red.sync.aligned.32x32b.x4%6%7%8.f32" + " {%0, %1, %2, %3}, %4, [%5];\n" + : "=r"(data[0]), "=r"(data[1]), "=r"(data[2]), "=r"(data[3]), + "=f"(redval) + : "r"(tmem_start), + "C"(USE_MAX ? s_max : s_min), + "C"(USE_ABS ? s_abs : s_empty), + "C"(USE_NAN ? s_nan : s_empty) + ); + } else if constexpr (kNumElements == 8) { + // outputs: %0..%7 = data, %8 = redval; inputs: %9 = taddr, %10..%12 = C + asm volatile( + "tcgen05.ld.red.sync.aligned.32x32b.x8%10%11%12.f32" + " {%0, %1, %2, %3," + " %4, %5, %6, %7}, %8, [%9];\n" + : "=r"(data[0]), "=r"(data[1]), "=r"(data[2]), "=r"(data[3]), + "=r"(data[4]), "=r"(data[5]), "=r"(data[6]), "=r"(data[7]), + "=f"(redval) + : "r"(tmem_start), + "C"(USE_MAX ? s_max : s_min), + "C"(USE_ABS ? s_abs : s_empty), + "C"(USE_NAN ? s_nan : s_empty) + ); + } else if constexpr (kNumElements == 16) { + // outputs: %0..%15 = data, %16 = redval; inputs: %17 = taddr, %18..%20 = C + asm volatile( + "tcgen05.ld.red.sync.aligned.32x32b.x16%18%19%20.f32" + " {%0, %1, %2, %3," + " %4, %5, %6, %7," + " %8, %9, %10, %11," + " %12, %13, %14, %15}, %16, [%17];\n" + : "=r"(data[0]), "=r"(data[1]), "=r"(data[2]), "=r"(data[3]), + "=r"(data[4]), "=r"(data[5]), "=r"(data[6]), "=r"(data[7]), + "=r"(data[8]), "=r"(data[9]), "=r"(data[10]), "=r"(data[11]), + "=r"(data[12]), "=r"(data[13]), "=r"(data[14]), "=r"(data[15]), + "=f"(redval) + : "r"(tmem_start), + "C"(USE_MAX ? s_max : s_min), + "C"(USE_ABS ? s_abs : s_empty), + "C"(USE_NAN ? s_nan : s_empty) + ); + } else if constexpr (kNumElements == 32) { + // outputs: %0..%31 = data, %32 = redval; inputs: %33 = taddr, %34..%36 = C + asm volatile( + "tcgen05.ld.red.sync.aligned.32x32b.x32%34%35%36.f32" + " {%0, %1, %2, %3," + " %4, %5, %6, %7," + " %8, %9, %10, %11," + " %12, %13, %14, %15," + " %16, %17, %18, %19," + " %20, %21, %22, %23," + " %24, %25, %26, %27," + " %28, %29, %30, %31}, %32, [%33];\n" + : "=r"(data[0]), "=r"(data[1]), "=r"(data[2]), "=r"(data[3]), + "=r"(data[4]), "=r"(data[5]), "=r"(data[6]), "=r"(data[7]), + "=r"(data[8]), "=r"(data[9]), "=r"(data[10]), "=r"(data[11]), + "=r"(data[12]), "=r"(data[13]), "=r"(data[14]), "=r"(data[15]), + "=r"(data[16]), "=r"(data[17]), "=r"(data[18]), "=r"(data[19]), + "=r"(data[20]), "=r"(data[21]), "=r"(data[22]), "=r"(data[23]), + "=r"(data[24]), "=r"(data[25]), "=r"(data[26]), "=r"(data[27]), + "=r"(data[28]), "=r"(data[29]), "=r"(data[30]), "=r"(data[31]), + "=f"(redval) + : "r"(tmem_start), + "C"(USE_MAX ? s_max : s_min), + "C"(USE_ABS ? s_abs : s_empty), + "C"(USE_NAN ? s_nan : s_empty) + ); + } else if constexpr (kNumElements == 64) { + // outputs: %0..%63 = data, %64 = redval; inputs: %65 = taddr, %66..%68 = C + asm volatile( + "tcgen05.ld.red.sync.aligned.32x32b.x64%66%67%68.f32" + " {%0, %1, %2, %3," + " %4, %5, %6, %7," + " %8, %9, %10, %11," + " %12, %13, %14, %15," + " %16, %17, %18, %19," + " %20, %21, %22, %23," + " %24, %25, %26, %27," + " %28, %29, %30, %31," + " %32, %33, %34, %35," + " %36, %37, %38, %39," + " %40, %41, %42, %43," + " %44, %45, %46, %47," + " %48, %49, %50, %51," + " %52, %53, %54, %55," + " %56, %57, %58, %59," + " %60, %61, %62, %63}, %64, [%65];\n" + : "=r"(data[0]), "=r"(data[1]), "=r"(data[2]), "=r"(data[3]), + "=r"(data[4]), "=r"(data[5]), "=r"(data[6]), "=r"(data[7]), + "=r"(data[8]), "=r"(data[9]), "=r"(data[10]), "=r"(data[11]), + "=r"(data[12]), "=r"(data[13]), "=r"(data[14]), "=r"(data[15]), + "=r"(data[16]), "=r"(data[17]), "=r"(data[18]), "=r"(data[19]), + "=r"(data[20]), "=r"(data[21]), "=r"(data[22]), "=r"(data[23]), + "=r"(data[24]), "=r"(data[25]), "=r"(data[26]), "=r"(data[27]), + "=r"(data[28]), "=r"(data[29]), "=r"(data[30]), "=r"(data[31]), + "=r"(data[32]), "=r"(data[33]), "=r"(data[34]), "=r"(data[35]), + "=r"(data[36]), "=r"(data[37]), "=r"(data[38]), "=r"(data[39]), + "=r"(data[40]), "=r"(data[41]), "=r"(data[42]), "=r"(data[43]), + "=r"(data[44]), "=r"(data[45]), "=r"(data[46]), "=r"(data[47]), + "=r"(data[48]), "=r"(data[49]), "=r"(data[50]), "=r"(data[51]), + "=r"(data[52]), "=r"(data[53]), "=r"(data[54]), "=r"(data[55]), + "=r"(data[56]), "=r"(data[57]), "=r"(data[58]), "=r"(data[59]), + "=r"(data[60]), "=r"(data[61]), "=r"(data[62]), "=r"(data[63]), + "=f"(redval) + : "r"(tmem_start), + "C"(USE_MAX ? s_max : s_min), + "C"(USE_ABS ? s_abs : s_empty), + "C"(USE_NAN ? s_nan : s_empty) + ); + } else if constexpr (kNumElements == 128) { + // outputs: %0..%127 = data, %128 = redval; inputs: %129 = taddr, %130..%132 = C + asm volatile( + "tcgen05.ld.red.sync.aligned.32x32b.x128%130%131%132.f32" + " {%0, %1, %2, %3," + " %4, %5, %6, %7," + " %8, %9, %10, %11," + " %12, %13, %14, %15," + " %16, %17, %18, %19," + " %20, %21, %22, %23," + " %24, %25, %26, %27," + " %28, %29, %30, %31," + " %32, %33, %34, %35," + " %36, %37, %38, %39," + " %40, %41, %42, %43," + " %44, %45, %46, %47," + " %48, %49, %50, %51," + " %52, %53, %54, %55," + " %56, %57, %58, %59," + " %60, %61, %62, %63," + " %64, %65, %66, %67," + " %68, %69, %70, %71," + " %72, %73, %74, %75," + " %76, %77, %78, %79," + " %80, %81, %82, %83," + " %84, %85, %86, %87," + " %88, %89, %90, %91," + " %92, %93, %94, %95," + " %96, %97, %98, %99," + " %100, %101, %102, %103," + " %104, %105, %106, %107," + " %108, %109, %110, %111," + " %112, %113, %114, %115," + " %116, %117, %118, %119," + " %120, %121, %122, %123," + " %124, %125, %126, %127}, %128, [%129];\n" + : "=r"(data[0]), "=r"(data[1]), "=r"(data[2]), "=r"(data[3]), + "=r"(data[4]), "=r"(data[5]), "=r"(data[6]), "=r"(data[7]), + "=r"(data[8]), "=r"(data[9]), "=r"(data[10]), "=r"(data[11]), + "=r"(data[12]), "=r"(data[13]), "=r"(data[14]), "=r"(data[15]), + "=r"(data[16]), "=r"(data[17]), "=r"(data[18]), "=r"(data[19]), + "=r"(data[20]), "=r"(data[21]), "=r"(data[22]), "=r"(data[23]), + "=r"(data[24]), "=r"(data[25]), "=r"(data[26]), "=r"(data[27]), + "=r"(data[28]), "=r"(data[29]), "=r"(data[30]), "=r"(data[31]), + "=r"(data[32]), "=r"(data[33]), "=r"(data[34]), "=r"(data[35]), + "=r"(data[36]), "=r"(data[37]), "=r"(data[38]), "=r"(data[39]), + "=r"(data[40]), "=r"(data[41]), "=r"(data[42]), "=r"(data[43]), + "=r"(data[44]), "=r"(data[45]), "=r"(data[46]), "=r"(data[47]), + "=r"(data[48]), "=r"(data[49]), "=r"(data[50]), "=r"(data[51]), + "=r"(data[52]), "=r"(data[53]), "=r"(data[54]), "=r"(data[55]), + "=r"(data[56]), "=r"(data[57]), "=r"(data[58]), "=r"(data[59]), + "=r"(data[60]), "=r"(data[61]), "=r"(data[62]), "=r"(data[63]), + "=r"(data[64]), "=r"(data[65]), "=r"(data[66]), "=r"(data[67]), + "=r"(data[68]), "=r"(data[69]), "=r"(data[70]), "=r"(data[71]), + "=r"(data[72]), "=r"(data[73]), "=r"(data[74]), "=r"(data[75]), + "=r"(data[76]), "=r"(data[77]), "=r"(data[78]), "=r"(data[79]), + "=r"(data[80]), "=r"(data[81]), "=r"(data[82]), "=r"(data[83]), + "=r"(data[84]), "=r"(data[85]), "=r"(data[86]), "=r"(data[87]), + "=r"(data[88]), "=r"(data[89]), "=r"(data[90]), "=r"(data[91]), + "=r"(data[92]), "=r"(data[93]), "=r"(data[94]), "=r"(data[95]), + "=r"(data[96]), "=r"(data[97]), "=r"(data[98]), "=r"(data[99]), + "=r"(data[100]), "=r"(data[101]), "=r"(data[102]), "=r"(data[103]), + "=r"(data[104]), "=r"(data[105]), "=r"(data[106]), "=r"(data[107]), + "=r"(data[108]), "=r"(data[109]), "=r"(data[110]), "=r"(data[111]), + "=r"(data[112]), "=r"(data[113]), "=r"(data[114]), "=r"(data[115]), + "=r"(data[116]), "=r"(data[117]), "=r"(data[118]), "=r"(data[119]), + "=r"(data[120]), "=r"(data[121]), "=r"(data[122]), "=r"(data[123]), + "=r"(data[124]), "=r"(data[125]), "=r"(data[126]), "=r"(data[127]), + "=f"(redval) + : "r"(tmem_start), + "C"(USE_MAX ? s_max : s_min), + "C"(USE_ABS ? s_abs : s_empty), + "C"(USE_NAN ? s_nan : s_empty) + ); + } +#endif +} + // Load from tensor memory, 16 data path lanes, 128-bit pattern, repeated N times. (https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-instructions-tcgen05-ld) template __device__ __forceinline__ @@ -379,4 +605,45 @@ void tmem_st_32dp32bNx(uint32_t tmem_start, void const* data_) { #endif } +// Decompose a non-power-of-2 replication count into a sequence of power-of-2 calls. +// +// This helper splits kNumReplications into its constituent powers of 2 (from highest to lowest) +// and issues one call to tmem_ld_st_fn.template operator()() for each, automatically +// advancing the tmem address by kNumElemsPerUnit columns and the data pointer by kNumElemsPerUnit +// uint32_t elements per unit. +// +// Template parameters: +// kNumReplications — total number of replications (may be non-power-of-2) +// kNumElemsPerUnit — number of tmem columns (and uint32_t data elements) consumed per unit +// (e.g. 1 for 32dp32b, 2 for 16dp128b, 4 for 16dp256b) +// F — a callable with a template operator(): f.template operator()(uint32_t, void*) +// +// Usage example (with a template lambda): +// tmem_ld_st_decomposed<6, 1>(tmem_start, data_ptr, +// [](uint32_t ts, void* d) { tmem_ld_32dp32bNx(ts, d); }); +// // Equivalent to: +// // tmem_ld_32dp32bNx<4>(tmem_start, data_ptr); +// // tmem_ld_32dp32bNx<2>(tmem_start + 4, (void*)((uint32_t*)data_ptr + 4)); +template +__device__ __forceinline__ +void tmem_ld_st_decomposed(uint32_t tmem_start, void* data, F&& tmem_ld_st_fn) { + if constexpr (kNumReplications == 0) { + // Base case: nothing to do + return; + } else { + // Largest power-of-2 <= kNumReplications + constexpr int kHighBit = 1 << (31 - __builtin_clz(kNumReplications)); + tmem_ld_st_fn.template operator()(tmem_start, data); + // Recurse on the remainder + constexpr int kRemaining = kNumReplications - kHighBit; + if constexpr (kRemaining > 0) { + tmem_ld_st_decomposed( + tmem_start + kHighBit * kNumElemsPerUnit, + (void*)((uint32_t*)data + kHighBit * kNumElemsPerUnit), + tmem_ld_st_fn + ); + } + } +} + } diff --git a/csrc/kerutils/include/kerutils/device/sm80/helpers.cuh b/csrc/kerutils/include/kerutils/device/sm80/helpers.cuh index 551b0f13..be8aec60 100644 --- a/csrc/kerutils/include/kerutils/device/sm80/helpers.cuh +++ b/csrc/kerutils/include/kerutils/device/sm80/helpers.cuh @@ -31,7 +31,7 @@ float2 float2float2(const float &x) { CUTE_DEVICE void st_shared(void* ptr, __int128_t val) { - asm volatile("st.shared.b128 [%0], %1;" :: "l"(__cvta_generic_to_shared(ptr)), "q"(val)); + asm volatile("st.shared.b128 [%0], %1;" :: "r"(cute::cast_smem_ptr_to_uint(ptr)), "q"(val)); } CUTE_DEVICE @@ -40,14 +40,14 @@ void st_shared(void* ptr, float4 val) { } CUTE_DEVICE -__int128_t ld_shared(void* ptr) { +__int128_t ld_shared(const void* ptr) { __int128_t val; - asm volatile("ld.shared.b128 %0, [%1];" : "=q"(val) : "l"(__cvta_generic_to_shared(ptr))); + asm volatile("ld.shared.b128 %0, [%1];" : "=q"(val) : "r"(cute::cast_smem_ptr_to_uint(ptr))); return val; } CUTE_DEVICE -float4 ld_shared_float4(void* ptr) { +float4 ld_shared_float4(const void* ptr) { __int128_t temp = ld_shared(ptr); return *(float4*)&temp; } diff --git a/csrc/kerutils/include/kerutils/device/sm80/intrinsics.cuh b/csrc/kerutils/include/kerutils/device/sm80/intrinsics.cuh index 7039a0bf..c6238aff 100644 --- a/csrc/kerutils/include/kerutils/device/sm80/intrinsics.cuh +++ b/csrc/kerutils/include/kerutils/device/sm80/intrinsics.cuh @@ -110,8 +110,8 @@ void atomicadd_f32_with_policy_and_pred(void* global_addr, const float &data, in } // Get the id of the current SM -// About %smid (https://docs.nvidia.com/cuda/parallel-thread-execution/#special-registers-smid): PTX document says that %smid ranges from 0 to %nsmid-1, while "The SM identifier numbering is not guaranteed to be contiguous, so %nsmid may be larger than the physical number of SMs in the device.". However, result shows that, at least for sm90 and sm100f, %nsmid is the number of physical SMs - 1. For the sake of safety, I recommend you to check the return of get_sm_id manually or call `get_sm_id_with_range_check()` defined in `device/sm80/helpers.cuh`. -// Besides, PTX document also says that this number may change due to preemption, but currently this never happens according to [DATEN GELÖSCHT] +// About %smid (https://docs.nvidia.com/cuda/parallel-thread-execution/#special-registers-smid): PTX document says that %smid ranges from 0 to %nsmid-1, while "The SM identifier numbering is not guaranteed to be contiguous, so %nsmid may be larger than the physical number of SMs in the device.". However, result shows that, at least for sm90 and sm100, %nsmid is the number of physical SMs - 1. For the sake of safety, I recommend you to check the return of get_sm_id manually or call `get_sm_id_with_range_check()` defined in `device/sm80/helpers.cuh`. +// Besides, PTX document also says that this number may change due to preemption, but currently this never happens in practice CUTE_DEVICE uint32_t get_sm_id() { uint32_t ret; @@ -143,4 +143,18 @@ void trap() { ); \ } +// STG.128 (https://docs.nvidia.com/cuda/parallel-thread-execution/#data-movement-and-conversion-instructions-st) +// L2_CACHE_HINT_STR should be ".L2::evict_XXX" (only available on sm100+) or "" +#define KU_STG_128(global_addr, src, L1_CACHE_HINT_STR, L2_CACHE_HINT_STR) \ + { \ + static_assert(std::is_pointer_v || std::is_array_v, "`global_addr` must be a pointer"); \ + static_assert(std::is_pointer_v || std::is_array_v, "`src` must be a pointer"); \ + uint64_t const* src_as_uint64_ptr = (uint64_t const*)(src); \ + asm volatile( \ + "st.global.L1::" L1_CACHE_HINT_STR L2_CACHE_HINT_STR ".v2.u64 [%0], {%1, %2};\n" \ + : \ + : "l"(global_addr), "l"(src_as_uint64_ptr[0]), "l"(src_as_uint64_ptr[1]) \ + ); \ + } + } diff --git a/csrc/kerutils/include/kerutils/host/host.h b/csrc/kerutils/include/kerutils/host/host.h index 3bdd1249..bd7b05dc 100644 --- a/csrc/kerutils/include/kerutils/host/host.h +++ b/csrc/kerutils/include/kerutils/host/host.h @@ -1,14 +1,13 @@ #pragma once +#include #include #include #include #include #include -#include - -#include +#include #include "kerutils/common/common.h" @@ -35,35 +34,42 @@ class KUException final : public std::exception { #define THROW_KU_EXCEPTION(name, ...) \ throw kerutils::KUException(name, __FILE__, __LINE__, __VA_ARGS__) -#define KU_CUDA_CHECK(call) \ -do { \ - cudaError_t status_ = call; \ - if (status_ != cudaSuccess) { \ - fprintf(stderr, "CUDA error (%s:%d): %s\n", __FILE__, __LINE__, cudaGetErrorString(status_)); \ - THROW_KU_EXCEPTION("CUDA", "CUDA error: ", cudaGetErrorString(status_)); \ - } \ +#define KU_CUDA_CHECK(call) \ +do { \ + cudaError_t status_ = (call); \ + if (status_ != cudaSuccess) { \ + char _ku_buf[1024]; \ + snprintf(_ku_buf, sizeof(_ku_buf), "CUDA error (%s:%d): %s", __FILE__, __LINE__, cudaGetErrorString(status_)); \ + fprintf(stderr, "%s\n", _ku_buf); \ + THROW_KU_EXCEPTION("CUDA", _ku_buf); \ + } \ } while(0) -#define KU_CUTLASS_CHECK(call) \ -do { \ - cutlass::Status status_ = call; \ - if (status_ != cutlass::Status::kSuccess) { \ - fprintf(stderr, "CUTLASS error (%s:%d): %d\n", __FILE__, __LINE__, static_cast(status_)); \ - THROW_KU_EXCEPTION("CUTLASS", "CUTLASS error: ", static_cast(status_)); \ - } \ +#define KU_CUTLASS_CHECK(call) \ +do { \ + cutlass::Status status_ = (call); \ + if (status_ != cutlass::Status::kSuccess) { \ + char _ku_buf[1024]; \ + snprintf(_ku_buf, sizeof(_ku_buf), "CUTLASS error (%s:%d): %d", __FILE__, __LINE__, static_cast(status_)); \ + fprintf(stderr, "%s\n", _ku_buf); \ + THROW_KU_EXCEPTION("CUTLASS", _ku_buf); \ + } \ } while(0) // This `KU_ASSERT` is triggered no matter if the code is compiled with `-DNDEBUG` or not. -#define KU_ASSERT(cond, ...) \ - do { \ - if (not (cond)) { \ - fprintf(stderr, "Assertion `%s` failed (%s:%d): ", #cond, __FILE__, __LINE__); \ - if constexpr (sizeof(#__VA_ARGS__) > 1) { \ - fprintf(stderr, ", " __VA_ARGS__); \ - } \ - fprintf(stderr, "\n"); \ - THROW_KU_EXCEPTION("Assertion", "Assertion `", #cond, "` failed."); \ - } \ +#define KU_ASSERT(cond, ...) \ + do { \ + if (not (cond)) { \ + char _ku_buf[1024]; \ + int _ku_len = snprintf(_ku_buf, sizeof(_ku_buf), \ + "Assertion `%s` failed (%s:%d)", #cond, __FILE__, __LINE__); \ + __VA_OPT__( \ + _ku_len += snprintf(_ku_buf + _ku_len, sizeof(_ku_buf) - _ku_len, \ + ": " __VA_ARGS__); \ + ) \ + fprintf(stderr, "%s\n", _ku_buf); \ + THROW_KU_EXCEPTION("Assertion", _ku_buf); \ + } \ } while(0) #define KU_CHECK_KERNEL_LAUNCH() KU_CUDA_CHECK(cudaGetLastError()) @@ -78,6 +84,13 @@ inline __host__ __device__ constexpr T ceil(const T &a, const T &b) { return (a + b - 1) / b * b; } +template +inline __host__ __device__ constexpr T find_next_power_of_2(const T& x) { + if (x <= LOWER_BOUND) + return LOWER_BOUND; + return find_next_power_of_2(x); +} + // A wrapper for make_tensor_map static inline CUtensorMap make_tensor_map( const std::vector &size, @@ -103,8 +116,30 @@ static inline CUtensorMap make_tensor_map( } KU_ASSERT(strides.size() == (uint32_t)dim-1 && box_size.size() == (uint32_t)dim && element_strides.size() == (uint32_t)dim); + auto call_cuTensorMapEncodeTiled = [&](Args... args) { + cudaDriverEntryPointQueryResult cuda_status; + void* pfn = nullptr; +#if (__CUDACC_VER_MAJOR__ > 12) + KU_CUDA_CHECK(cudaGetDriverEntryPointByVersion( + "cuTensorMapEncodeTiled", + &pfn, 12000, + cudaEnableDefault, + &cuda_status)); +#else + KU_CUDA_CHECK(cudaGetDriverEntryPoint( + "cuTensorMapEncodeTiled", + &pfn, + cudaEnableDefault, + &cuda_status)); +#endif + if (cuda_status != cudaDriverEntryPointSuccess) { + KU_ASSERT(false, "Failed to load `cuTensorMapEncodeTiled`. cuda_status = %d", cuda_status); + } + return reinterpret_cast(pfn)(args...); \ + }; + CUtensorMap result; - CUresult ret_code = CUTLASS_CUDA_DRIVER_WRAPPER_CALL(cuTensorMapEncodeTiled)( + CUresult ret_code = call_cuTensorMapEncodeTiled( &result, data_type, dim, diff --git a/csrc/kerutils/include/kerutils/supplemental/torch_tensors.h b/csrc/kerutils/include/kerutils/supplemental/torch_tensors.h index 5b4e564c..cfe00da8 100644 --- a/csrc/kerutils/include/kerutils/supplemental/torch_tensors.h +++ b/csrc/kerutils/include/kerutils/supplemental/torch_tensors.h @@ -68,4 +68,4 @@ static inline PtrT* get_optional_tensor_ptr(const T& tensor_or_opt) { #define KU_CHECK_LAST_DIM_CONTIGUOUS(tensor) TORCH_CHECK(ku::_check_optional_tensor(tensor, [](const at::Tensor& t) { return t.size(-1) == 1 || t.stride(-1) == 1; }), #tensor " must have contiguous last dimension") // Check whether the given tensor (or optional) has the specified dtype -#define KU_CHECK_DTYPE(tensor, target_dtype) TORCH_CHECK(ku::_check_optional_tensor(tensor, [](const at::Tensor& t) { return t.dtype() == (target_dtype); }), #tensor " must have dtype " #target_dtype) +#define KU_CHECK_DTYPE(tensor, target_dtype) TORCH_CHECK(ku::_check_optional_tensor(tensor, [&](const at::Tensor& t) { return t.dtype() == (target_dtype); }), #tensor " must have dtype " #target_dtype) diff --git a/csrc/sm100/decode/head128/README.md b/csrc/sm100/decode/head128/README.md deleted file mode 100644 index 6cd90624..00000000 --- a/csrc/sm100/decode/head128/README.md +++ /dev/null @@ -1 +0,0 @@ -Head128 decoding kernels are located at `csrc/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512.cu` (for k_dim = 512) or simulated using 2x head64 kernel \ No newline at end of file diff --git a/csrc/sm100/decode/head64/instantiations/model1.cu b/csrc/sm100/decode/head64/instantiations/model1.cu deleted file mode 100644 index 868ff0c8..00000000 --- a/csrc/sm100/decode/head64/instantiations/model1.cu +++ /dev/null @@ -1,8 +0,0 @@ -#include "../kernel.cuh" - -namespace sm100::decode::head64 { - -template -void run_flash_splitkv_mla_fp8_sparse_kernel(const SparseAttnDecodeParams ¶ms); - -} diff --git a/csrc/sm100/decode/head64/instantiations/v32.cu b/csrc/sm100/decode/head64/instantiations/v32.cu deleted file mode 100644 index 08ce093f..00000000 --- a/csrc/sm100/decode/head64/instantiations/v32.cu +++ /dev/null @@ -1,8 +0,0 @@ -#include "../kernel.cuh" - -namespace sm100::decode::head64 { - -template -void run_flash_splitkv_mla_fp8_sparse_kernel(const SparseAttnDecodeParams ¶ms); - -} diff --git a/csrc/sm100/decode/head64/kernel.h b/csrc/sm100/decode/head64/kernel.h deleted file mode 100644 index 0b3c63c0..00000000 --- a/csrc/sm100/decode/head64/kernel.h +++ /dev/null @@ -1,11 +0,0 @@ -#pragma once - -#include "params.h" - -namespace sm100::decode::head64 { - -template -void run_flash_splitkv_mla_fp8_sparse_kernel(const SparseAttnDecodeParams ¶ms); - -} - diff --git a/csrc/sm100/helpers.h b/csrc/sm100/helpers.h deleted file mode 100644 index a5666257..00000000 --- a/csrc/sm100/helpers.h +++ /dev/null @@ -1,35 +0,0 @@ -#pragma once - -#include -#include -#include - -#include "defines.h" - -namespace sm100 { - -using namespace cute; - -CUTE_DEVICE -int int4_max(int4 t) { - return max(max(t.x, t.y), max(t.z, t.w)); -} - -CUTE_DEVICE -int int4_min(int4 t) { - return min(min(t.x, t.y), min(t.z, t.w)); -} - -// Convert 2x fp8_e4m3 to 2x bf16 with scaling -CUTE_DEVICE -nv_bfloat162 fp8x2_to_bf16x2_with_scale(__nv_fp8x2_e4m3 data, nv_bfloat16 scale) { - // TODO Use native conversion for CUDA >= 13.1 - float2 data_float2 = (float2)data; - nv_bfloat162 data_bf16x2 = __float22bfloat162_rn(data_float2); - return nv_bfloat162 { - data_bf16x2.x * scale, - data_bf16x2.y * scale - }; -} - -} diff --git a/csrc/sm100/prefill/sparse/fwd/head128/instantiations/phase1_k512.cu b/csrc/sm100/prefill/sparse/fwd/head128/instantiations/phase1_k512.cu deleted file mode 100644 index 5dcec83b..00000000 --- a/csrc/sm100/prefill/sparse/fwd/head128/instantiations/phase1_k512.cu +++ /dev/null @@ -1,8 +0,0 @@ -#include "../phase1.h" -#include "../phase1.cuh" - -namespace sm100::fwd::head128 { - -template void run_fwd_phase1_kernel<512>(const SparseAttnFwdParams& params); - -} diff --git a/csrc/sm100/prefill/sparse/fwd/head128/instantiations/phase1_k576.cu b/csrc/sm100/prefill/sparse/fwd/head128/instantiations/phase1_k576.cu deleted file mode 100644 index bfd01ced..00000000 --- a/csrc/sm100/prefill/sparse/fwd/head128/instantiations/phase1_k576.cu +++ /dev/null @@ -1,8 +0,0 @@ -#include "../phase1.h" -#include "../phase1.cuh" - -namespace sm100::fwd::head128 { - -template void run_fwd_phase1_kernel<576>(const SparseAttnFwdParams& params); - -} diff --git a/csrc/sm100/prefill/sparse/fwd/head128/phase1.h b/csrc/sm100/prefill/sparse/fwd/head128/phase1.h deleted file mode 100644 index b1057809..00000000 --- a/csrc/sm100/prefill/sparse/fwd/head128/phase1.h +++ /dev/null @@ -1,10 +0,0 @@ -#pragma once - -#include "params.h" - -namespace sm100::fwd::head128 { - -template -void run_fwd_phase1_kernel(const SparseAttnFwdParams& params); - -} diff --git a/csrc/sm100/prefill/sparse/fwd/head64/instantiations/phase1_k512.cu b/csrc/sm100/prefill/sparse/fwd/head64/instantiations/phase1_k512.cu deleted file mode 100644 index e1c87be9..00000000 --- a/csrc/sm100/prefill/sparse/fwd/head64/instantiations/phase1_k512.cu +++ /dev/null @@ -1,8 +0,0 @@ -#include "../phase1.h" -#include "../phase1.cuh" - -namespace sm100::fwd::head64 { - -template void run_fwd_phase1_kernel<512>(const SparseAttnFwdParams& params); - -} diff --git a/csrc/sm100/prefill/sparse/fwd/head64/instantiations/phase1_k576.cu b/csrc/sm100/prefill/sparse/fwd/head64/instantiations/phase1_k576.cu deleted file mode 100644 index 1bd214ee..00000000 --- a/csrc/sm100/prefill/sparse/fwd/head64/instantiations/phase1_k576.cu +++ /dev/null @@ -1,8 +0,0 @@ -#include "../phase1.h" -#include "../phase1.cuh" - -namespace sm100::fwd::head64 { - -template void run_fwd_phase1_kernel<576>(const SparseAttnFwdParams& params); - -} diff --git a/csrc/sm100/prefill/sparse/fwd/head64/phase1.cuh b/csrc/sm100/prefill/sparse/fwd/head64/phase1.cuh deleted file mode 100644 index b510b27f..00000000 --- a/csrc/sm100/prefill/sparse/fwd/head64/phase1.cuh +++ /dev/null @@ -1,673 +0,0 @@ -#pragma once -#include "phase1.h" - -#include -#include -#include -#include -#include - -#include - -#include "params.h" -#include "utils.h" -#include "sm100/helpers.h" -#include "sm100/prefill/sparse/common_subroutine.h" -#include "config.h" - -namespace sm100::fwd::head64 { - -using namespace cute; - -/* -Pipeline Overview: - -| Copy | MMA | Scale & Exp | - -KV0 -KV1 -KV2 - P0 = QK0^T - S0 = exp(P0) - scale(O) w.r.t P0 - P1 = QK1^T - S1 = exp(P1) - O += S0V0 -KV3 scale(O) w.r.t P1 - P2 = QK2^T - S2 = exp(P2) - O += S1V1 -KV4 scale(O) w.r.t P2 - P3 = QK3^T - S3 = exp(P3) - O += S2V2 -KV5 scale(O) w.r.t P3 - -... - - O += S(n-3)V(n-3) - scale(O) w.r.t P(n-2) - P(n-1) = QK(n-1)^T - S(n-1) = exp(P(n-1)) - O += S(n-2)V(n-2) - scale(O) w.r.t P(n-1) - O += S(n-1)V(n-1) -*/ - -using FwdMode = SparseAttnFwdMode; - -template -__global__ void __launch_bounds__(NUM_THREADS, 1, 1) -sparse_attn_fwd_kernel(__grid_constant__ const SparseAttnFwdParams params, __grid_constant__ const TmaParams tma_params) { -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000 && __CUDA_ARCH__ < 1200)) || (defined(__CLION_IDE__) || defined(__VSCODE_IDE__)) - // Grid shape: [s_q, 1, 1] - - const int s_q_idx = blockIdx.x; - const int warp_idx = cutlass::canonical_warp_idx_sync(); - const int lane_idx = threadIdx.x % 32; - const int warpgroup_idx = __shfl_sync(0xffffffff, threadIdx.x / 128, 0); - const int idx_in_warpgroup = threadIdx.x % 128; - const int topk_length = params.topk_length != nullptr ? __ldg(params.topk_length + s_q_idx) : params.topk; - const int num_k_blocks = max(cute::ceil_div(topk_length, (int)B_TOPK), 1); // num_k_blocks always >= 1 - - // Define shared tensors - extern __shared__ char wksp_buf[]; - SharedMemoryPlan &plan = *reinterpret_cast(wksp_buf); - - int* gIndices = params.indices + s_q_idx*params.stride_indices_s_q; // [topk] - - // Allocate tmem tensors - TiledMMA tiled_mma_P = TiledMMA_P{}; - TiledMMA tiled_mma_O = TiledMMA_O{}; - // NOTE These tXXX tensors are only for a forged layout (so that CuTe is able to generate correct address in cute::gemm) - Tensor tP = partition_fragment_C(tiled_mma_P, Shape, _128>{}); - Tensor tQ_nope_part0 = tiled_mma_P.get_slice(_0{}).make_fragment_A( - partition_shape_A(tiled_mma_P, Shape, Int<(D_V/2)/2>>{}) - ); - Tensor tQ_nope_part1 = tiled_mma_P.get_slice(_0{}).make_fragment_A( - partition_shape_A(tiled_mma_P, Shape, Int<(D_V/2)/2>>{}) - ); - Tensor tQ_rope = tiled_mma_P.get_slice(_0{}).make_fragment_A( - partition_shape_A(tiled_mma_P, Shape, Int<64/2>>{}) - ); - Tensor tO = partition_fragment_C(tiled_mma_O, Shape, Int>{}); - tP.data().get() = tmem_cols::P; - tQ_nope_part0.data().get() = tmem_cols::Q; - tQ_nope_part1.data().get() = tmem_cols::Q + 64; - tQ_rope.data().get() = tmem_cols::Q_RoPE; - tO.data().get() = tmem_cols::O; - - if (warp_idx == 0) { - if (elect_one_sync()) { - // Copy Q - if constexpr (HAVE_ROPE) { - cute::prefetch_tma_descriptor(tma_params.tma_Q_rope.get_tma_descriptor()); - } - cute::prefetch_tma_descriptor(tma_params.tma_Q_nope.get_tma_descriptor()); - - plan.bar_prologue_q_nope.init(1); - plan.bar_prologue_q_rope.init(1); - fence_barrier_init(); - - if constexpr (HAVE_ROPE) { - Tensor gQ_rope = tma_params.tma_Q_rope.get_tma_tensor(tma_params.shape_Q_rope)(_, _, s_q_idx); - Tensor sQ_rope = make_tensor(make_smem_ptr(plan.s_q_rope.q_rope.data()), SmemLayoutQRoPE{}); - ku::launch_tma_copy(tma_params.tma_Q_rope, gQ_rope, sQ_rope, plan.bar_prologue_q_rope, TMA::CacheHintSm90::EVICT_FIRST); - } - - Tensor gQ_nope = tma_params.tma_Q_nope.get_tma_tensor(tma_params.shape_Q_nope)(_, _, s_q_idx); - Tensor sQ_nope = make_tensor(make_smem_ptr(plan.u.q_full.q_nope.data()), SmemLayoutQNoPE{}); - ku::launch_tma_copy(tma_params.tma_Q_nope, gQ_nope, sQ_nope, plan.bar_prologue_q_nope, TMA::CacheHintSm90::EVICT_FIRST); - - cute::prefetch_tma_descriptor(tma_params.tma_O.get_tma_descriptor()); - cute::prefetch_tma_descriptor(&(tma_params.tensor_map_kv_nope)); - - // Initialize other barriers - plan.bar_prologue_utccp_rope.init(1); - plan.bar_prologue_utccp_nope.init(1); - CUTE_UNROLL - for (int i = 0; i < NUM_BUFS; ++i) { - plan.bar_qk_nope_done[i].init(1); - plan.bar_sv_done[i].init(1); - plan.bar_kv_nope_ready[i][0].init(1); - plan.bar_kv_nope_ready[i][1].init(1); - plan.bar_k_valid_ready[i].init(B_TOPK/8); - plan.bar_k_valid_free[i].init(128); - } - plan.bar_p_free.init(128); - plan.bar_so_ready.init(128); - plan.bar_qk_rope_done.init(1); - plan.bar_kv_rope_ready.init(64); - fence_barrier_init(); - } - - // Initialize TMEM - cute::TMEM::Allocator1Sm().allocate(512, plan.tmem_start_addr.data()); - TRAP_ONLY_DEVICE_ASSERT(plan.tmem_start_addr.data()[0] == 0); - cute::TMEM::Allocator1Sm().release_allocation_lock(); - } - - __syncthreads(); - - if (warpgroup_idx == 0) { - // Scale & Exp warps - - // The following three numbers are - // - mi: max_logits used to scale Pi (i.e. O := exp2(Pi*scale - mi) @ V) - // - li: sumexp, i.e. li := sum(exp(Pi*scale - mi)) - // - real_mi: real max logits, i.e. real_mi := max(Pi*scale) - // where Pi is the i-th row of P, P := QK^T - // mi and real_mi are always consistent within the two threads that - // controls one row (i.e. thread 0+64, 1+65, 2+66, ...) after every update - float mi = MAX_INIT_VAL; - float li = 0.0f; - float real_mi = -CUDART_INF_F; - - bf16* sS_base = plan.s_q_rope.s + lane_idx*8 + (warp_idx&1)*(B_H/2)*8 + (warp_idx/2)*B_H*(B_TOPK/2); - static constexpr int NUM_ELEMS_PER_THREAD = B_TOPK / 2; - - CUTE_NO_UNROLL - for (int k = 0; k < num_k_blocks; ++k) { - // Wait for P - NamedBarrier::arrive_and_wait(64, NamedBarriers::wg0_warp02_sync+(warp_idx&1)); - plan.bar_qk_nope_done[k%NUM_BUFS].wait((k/NUM_BUFS)&1); - plan.bar_k_valid_ready[k%NUM_BUFS].wait((k/NUM_BUFS)&1); // Put the barrier wait here for more code reordering space - ku::tcgen05_after_thread_sync(); - - // Load P - float p[NUM_ELEMS_PER_THREAD]; - retrieve_mask_and_reduce_p< - NUM_ELEMS_PER_THREAD, - tmem_cols::P, - NamedBarriers::wg0_warp02_sync, - NamedBarriers::wg0_warp13_sync, - false - >( - plan.is_k_valid[k%NUM_BUFS], - warp_idx, lane_idx, - [&]() {plan.bar_p_free.arrive();}, - plan.p_exchange_buf, - p - ); - plan.bar_k_valid_free[k%NUM_BUFS].arrive(); - - // Get rowwise max of Pi - 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); - 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, ...) - // - should_scale_o is identical among every warp, and is identical among threads that controls the same row (i.e. among threads 0~31+64~95; and is identical among threads 32~63+96~127) - - - // 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]; - float cur_sum = get_s_from_p(s, p, params.sm_scale_div_log2, new_max); - li = fma(li, scale_for_old, cur_sum); - - // Wait for last SV gemm, write S - if (k > 0) { - plan.bar_sv_done[(k-1)%NUM_BUFS].wait(((k-1)/NUM_BUFS)&1); - } - 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 (k > 0 && should_scale_o) { - // plan.bar_sv_done[(k-1)%NUM_BUFS].wait(((k-1)/NUM_BUFS)&1); // NOTE We have waited for last SV gemm before - ku::tcgen05_after_thread_sync(); - rescale_O(scale_for_old); - ku::tcgen05_before_thread_sync(); - } - - fence_view_async_shared(); - plan.bar_so_ready.arrive(); - } - - // Epilogue - - 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_li_buf[idx_in_warpgroup] = li; - NamedBarrier::arrive_and_wait(128, NamedBarriers::wg0_sync); - li += plan.rowwise_li_buf[idx_in_warpgroup^64]; - - // Store mi and li - if (idx_in_warpgroup < 64) { - int global_index = s_q_idx*params.h_q + idx_in_warpgroup; - float cur_lse = fmaf(mi, CUDART_LN2_F, logf(li)); - cur_lse = cur_lse == -CUDART_INF_F ? +CUDART_INF_F : cur_lse; - params.max_logits[global_index] = real_mi*CUDART_LN2_F; - params.lse[global_index] = cur_lse; - } - - // Wait for the last GEMM - plan.bar_sv_done[(num_k_blocks-1)%NUM_BUFS].wait(((num_k_blocks-1)/NUM_BUFS)&1); - ku::tcgen05_after_thread_sync(); - - // Fetch dO if necessary - - // Store O - float attn_sink = params.attn_sink == nullptr ? -CUDART_INF_F : __ldg(params.attn_sink + (idx_in_warpgroup%64))*CUDART_L2E_F; - float output_scale = __fdividef(1.0f, li + exp2f(attn_sink - mi)); - Tensor sO = make_tensor(make_smem_ptr(plan.u.o.data()), SmemLayoutO{}); - constexpr int B_EPI = 64; - Tensor tma_gO = flat_divide( - tma_params.tma_O.get_tma_tensor(tma_params.shape_O)(_, _, s_q_idx), - Shape, Int>{} - )(_, _, _0{}, _); - Tensor sO_divided = flat_divide( - sO, - Shape, Int>{} - )(_, _, _0{}, _); - auto thr_tma = tma_params.tma_O.get_slice(_0{}); - - float2 o[B_EPI/2]; - bool have_valid_indices = __any_sync(0xffffffff, li != 0); // Prevent some threads' li == 0 and some threads' li != 0 which lead to deadlock during ku::tmem_ld - if (!have_valid_indices) { - // If there are no valid indices, we set o[i] to 0 and don't load from TMEM - CUTE_UNROLL - for (int i = 0; i < B_EPI/2; ++i) - o[i].x = o[i].y = 0.0f; - output_scale = 1.0f; - } - - float2 output_scale_float2 = make_float2(output_scale, output_scale); - - bf16* sO_addrs[8]; - CUTE_UNROLL - for (int i = 0; i < B_EPI/8; ++i) { - sO_addrs[i] = &sO(idx_in_warpgroup%64, i*8); - } - - CUTE_UNROLL - for (int c = 0; c < 2; ++c) { - // Each tile: 64 x 256 - CUTE_UNROLL - for (int k = 0; k < (D_V/4)/B_EPI; ++k) { - // Load O from tO - if (have_valid_indices) { - ku::tmem_ld_32dp32bNx(tmem_cols::O + c*128 + k*B_EPI, o); - cutlass::arch::fence_view_async_tmem_load(); - } - - // Convert and store - CUTE_UNROLL - for (int i = 0; i < B_EPI/8; ++i) { - nv_bfloat162 o_bf16[4]; - CUTE_UNROLL - for (int j = 0; j < 4; ++j) { - o[i*4+j] = ku::float2_mul(o[i*4+j], output_scale_float2); - o_bf16[j] = __float22bfloat162_rn(o[i*4+j]); - } - *(uint128_t*)(sO_addrs[i] + (c*(D_V/2) + (idx_in_warpgroup/64)*(D_V/4) + k*B_EPI)*64) = *(uint128_t*)(o_bf16); - } - - // Sync - fence_view_async_shared(); - NamedBarrier::arrive_and_wait(128, NamedBarriers::wg0_sync); - - if (warp_idx == 0 && elect_one_sync()) { - int epi_chunk_idx = c*(D_V/2/B_EPI) + k; - cute::copy( - tma_params.tma_O, - thr_tma.partition_S(sO_divided(_, _, epi_chunk_idx)), - thr_tma.partition_D(tma_gO(_, _, epi_chunk_idx)) - ); - } - if (warp_idx == 1 && elect_one_sync()) { - int epi_chunk_idx = c*(D_V/2/B_EPI) + (D_V/B_EPI/4) + k; - cute::copy( - tma_params.tma_O, - thr_tma.partition_S(sO_divided(_, _, epi_chunk_idx)), - thr_tma.partition_D(tma_gO(_, _, epi_chunk_idx)) - ); - } - } - } - - - if (warp_idx == 0) { - cute::TMEM::Allocator1Sm().free(0, 512); - } - } else if (warpgroup_idx == 1) { - // Producer warp for KV - int warp_idx = cutlass::canonical_warp_idx_sync() - 4; - constexpr int NUM_WARPS = 4, NUM_LOCAL_ROWS_PER_WARP = (B_TOPK/4)/NUM_WARPS; - if (elect_one_sync()) { - CUTE_NO_UNROLL - for (int k = 0; k < num_k_blocks; ++k) { - int4 indices[NUM_LOCAL_ROWS_PER_WARP]; - int max_indices = -1, min_indices = params.s_kv; - CUTE_UNROLL - for (int local_row = 0; local_row < NUM_LOCAL_ROWS_PER_WARP; ++local_row) { - indices[local_row] = __ldg((int4*)(gIndices + k*B_TOPK) + local_row*NUM_WARPS + warp_idx); - max_indices = max(max_indices, int4_max(indices[local_row])); - min_indices = min(min_indices, int4_min(indices[local_row])); - } - bool is_all_rows_invalid = min_indices == params.s_kv || max_indices == -1; - bool should_skip_tma = is_all_rows_invalid && k >= NUM_BUFS; - - if (k == 2) { - plan.bar_prologue_utccp_nope.wait(0); // Since q_nope coincidences with k[2] - } - - // Copy NoPE - int cur_buf = k%NUM_BUFS; - plan.bar_sv_done[cur_buf].wait((k/NUM_BUFS)&1^1); - bf16* sK_nope_base = plan.u.k.k_nope[cur_buf].data() + warp_idx*4*64; - - auto load_kv_nope_part = [&](int part_idx) { - CUTE_UNROLL - for (int local_row = 0; local_row < NUM_LOCAL_ROWS_PER_WARP; ++local_row) { - CUTE_UNROLL - for (int local_col = part_idx*(D_V/2/64); local_col < (part_idx+1)*(D_V/2/64); ++local_col) { - ku::tma_gather4( - &(tma_params.tensor_map_kv_nope), - plan.bar_kv_nope_ready[cur_buf][part_idx], - sK_nope_base + local_row*(4*NUM_WARPS)*64 + local_col*(B_TOPK*64), - local_col*64, - indices[local_row], - (int64_t)TMA::CacheHintSm90::EVICT_LAST - ); - } - } - }; - - if (!should_skip_tma) { - load_kv_nope_part(0); - load_kv_nope_part(1); - } else { - // NOTE See head128/phase1.cuh for this TMA skipping technique - CUTE_UNROLL - for (int part_idx = 0; part_idx < 2; ++part_idx) - plan.bar_kv_nope_ready[cur_buf][part_idx].complete_transaction(NUM_LOCAL_ROWS_PER_WARP*4*D_V/2*sizeof(bf16)); - } - } - } - } else { - // MMA warp - if (warp_idx == 8 && elect_one_sync()) { - // S -> T copy for Q - UMMA::SmemDescriptor sQ_nope_desc = UMMA::make_umma_desc( - make_tensor( - make_smem_ptr(plan.u.q_full.q_nope.data()), - tile_to_shape( - UMMA::Layout_K_SW128_Atom{}, - Shape, Int<64>>{} // We use this shape for dual gemm (TODO Link) - ) - ) - ); - UMMA::SmemDescriptor sQ_rope_desc = UMMA::make_umma_desc( - make_tensor( - make_smem_ptr(plan.s_q_rope.q_rope.data()), - tile_to_shape( - UMMA::Layout_K_SW64_Atom{}, - Shape, Int<32>>{} - ) - ) - ); - - if constexpr (HAVE_ROPE) { - // Copy the RoPE tile: 128 rows * 32 cols (64B) (in UTCCP's view), or 64 rows * 64 cols (in our view) - plan.bar_prologue_q_rope.arrive_and_expect_tx(B_H*(D_Q-D_V)*sizeof(bf16)); - plan.bar_prologue_q_rope.wait(0); - ku::tcgen05_after_thread_sync(); - CUTE_UNROLL - for (int subtile_idx = 0; subtile_idx < 2; ++subtile_idx) { - // A subtile is 128 rows * 16 cols (256b, 32B) (in UTCCP's view), or 64 rows * 16 cols * 2 (in our view) - SM100_UTCCP_128dp256bit_1cta::copy( - sQ_rope_desc + (subtile_idx*32) / 16, - tmem_cols::Q_RoPE + subtile_idx*8 - ); - } - ku::umma_arrive_noelect(plan.bar_prologue_utccp_rope); - } - - plan.bar_prologue_q_nope.arrive_and_expect_tx(B_H*D_V*sizeof(bf16)); - plan.bar_prologue_q_nope.wait(0); - ku::tcgen05_after_thread_sync(); - CUTE_UNROLL - for (int tile_idx = 0; tile_idx < D_V/64/2; ++tile_idx) { - // A tile is 128 rows * 64 cols (128B) (in UTCCP's view), or 64 rows * 128 cols (in our view) - CUTE_UNROLL - for (int subtile_idx = 0; subtile_idx < 4; ++subtile_idx) { - // A subtile is 128 rows * 16 cols (256b, 32B) (in UTCCP's view), or 64 rows * 16 cols * 2 (in our view) - SM100_UTCCP_128dp256bit_1cta::copy( - sQ_nope_desc + (tile_idx*(B_H*128*2) + subtile_idx*32) / 16, // Remember that 4 LSBs are not included - tmem_cols::Q + tile_idx*32 + subtile_idx*8 - ); - } - } - ku::umma_arrive_noelect(plan.bar_prologue_utccp_nope); - - if constexpr (HAVE_ROPE) { - plan.bar_prologue_utccp_rope.wait(0); - } - - CUTE_NO_UNROLL - for (int k = 0; k < num_k_blocks+1; ++k) { - if (k < num_k_blocks) { - // Pi = QKi^T - int cur_buf = k%NUM_BUFS; - Tensor sK_nope = make_tensor(make_smem_ptr(plan.u.k.k_nope[cur_buf].data()), SmemLayoutKNoPE_TiledMMA{}); - Tensor sK_rope = make_tensor(make_smem_ptr(plan.u.k.k_rope.data()), SmemLayoutKRoPE_TiledMMA{}); - - plan.bar_p_free.wait(k&1^1); - ku::tcgen05_after_thread_sync(); - - // Wait for K (RoPE) - // P = Q(rope) @ K(rope)^T - if constexpr (HAVE_ROPE) { - plan.bar_kv_rope_ready.wait(k&1); - ku::tcgen05_after_thread_sync(); - ku::utcmma_ts(tiled_mma_P, tQ_rope, sK_rope, tP, true); - ku::umma_arrive_noelect(plan.bar_qk_rope_done); - } - - // Wait for K (NoPE) - if (k == 0) { - plan.bar_prologue_utccp_nope.wait(0); - } - Tensor sK_nope_divided = flat_divide(sK_nope, Tile, Int>{})(_, _, _0{}, _); - CUTE_UNROLL - for (int kv_nope_part_idx = 0; kv_nope_part_idx < 2; ++kv_nope_part_idx) { - plan.bar_kv_nope_ready[cur_buf][kv_nope_part_idx].arrive_and_expect_tx(B_TOPK*D_V/2*sizeof(bf16)); - plan.bar_kv_nope_ready[cur_buf][kv_nope_part_idx].wait((k/NUM_BUFS)&1); - ku::tcgen05_after_thread_sync(); - - // P += Q(nope) @ K(nope)^T - bool clear_accum = (!HAVE_ROPE) && kv_nope_part_idx == 0; - ku::utcmma_ts(tiled_mma_P, kv_nope_part_idx ? tQ_nope_part1 : tQ_nope_part0, sK_nope_divided(_, _, kv_nope_part_idx), tP, clear_accum); - } - ku::umma_arrive_noelect(plan.bar_qk_nope_done[cur_buf]); - } - if (k > 0) { - // O += S(i-1)V(i-1) - int cur_buf = (k-1)%NUM_BUFS; - - Tensor sS = make_tensor(make_smem_ptr(plan.s_q_rope.s), SmemLayoutS{}); - Tensor sV = make_tensor(make_smem_ptr(plan.u.k.k_nope[cur_buf].data()), SmemLayoutV{}); - - // Wait for S(i-1) and O to be scaled - plan.bar_so_ready.wait((k-1)&1); - ku::tcgen05_after_thread_sync(); - - // O += sS @ sV - ku::utcmma_ss(tiled_mma_O, sS, sV, tO, k == 1); - ku::umma_arrive_noelect(plan.bar_sv_done[cur_buf]); - } - } - } else if (warp_idx == 9) { - // KV valid loading warp - if (lane_idx < B_TOPK/8) { - CUTE_NO_UNROLL - for (int k = 0; k < num_k_blocks; ++k) { - char k_validness_mask = load_indices_and_generate_mask( - lane_idx, - gIndices + k*B_TOPK, - params.s_kv, - k*B_TOPK, - topk_length - ); - - int cur_buf = k%NUM_BUFS; - plan.bar_k_valid_free[cur_buf].wait((k/NUM_BUFS)&1^1); - plan.is_k_valid[cur_buf][lane_idx] = k_validness_mask; - plan.bar_k_valid_ready[cur_buf].arrive(); - } - } - } else if (warp_idx == 10 || warp_idx == 11) { - if constexpr (HAVE_ROPE) { - int thread_idx = threadIdx.x - 10*32; - constexpr int GROUP_SIZE = 8, NUM_GROUPS = 64/GROUP_SIZE, ROWS_PER_THREAD = B_TOPK/NUM_GROUPS; - int group_idx = thread_idx / GROUP_SIZE, idx_in_group = thread_idx % GROUP_SIZE; - Tensor sK_rope = make_tensor(make_smem_ptr(plan.u.k.k_rope.data()), SmemLayoutKRoPE{}); - bf16* sK_rope_base = &sK_rope(group_idx, idx_in_group*8); - CUTE_NO_UNROLL - for (int k = 0; k < num_k_blocks; ++k) { - int indices[ROWS_PER_THREAD]; - CUTE_UNROLL - for (int local_row = 0; local_row < ROWS_PER_THREAD; ++local_row) - indices[local_row] = __ldg(gIndices + k*B_TOPK + group_idx + local_row*NUM_GROUPS); - plan.bar_qk_rope_done.wait(k&1^1); - CUTE_UNROLL - for (int local_row = 0; local_row < ROWS_PER_THREAD; ++local_row) { - int index = indices[local_row]; - ku::cp_async_cacheglobal( - params.kv + (int64_t)index*params.stride_kv_s_kv + 512 + idx_in_group*8, - sK_rope_base + local_row*NUM_GROUPS*32, - index >= 0 && index < params.s_kv - ); // NOTE Using cp.async instead of TMA is faster here - // NOTE Here we only consider the range of `index` instead of also checking against topk_length, as it's noted that under this scenario (i.e. there exists a valid index among indices[topk_length: ] that points to a token who has NaN inside) - } - cutlass::arch::cpasync_barrier_arrive_noinc((uint64_t*)&(plan.bar_kv_rope_ready)); - } - } - } - } - - -#else - if (cute::thread0()) { - CUTE_INVALID_CONTROL_PATH("This kernel only supports sm100"); - } -#endif -} - -template -void run_fwd_phase1_kernel(const SparseAttnFwdParams& params) { - KU_ASSERT(params.h_kv == 1); - KU_ASSERT(params.topk % B_TOPK == 0); // To save some boundry checkings - KU_ASSERT(params.h_q == B_H); // To save some calculation - KU_ASSERT(params.d_qk == D_QK); - static_assert(D_QK == 576 || D_QK == 512); - - auto shape_Q_nope = make_shape(params.h_q, D_V, params.s_q); - auto tma_Q_nope = cute::make_tma_copy( - SM90_TMA_LOAD{}, - make_tensor( - make_gmem_ptr((bf16*)params.q), - make_layout( - shape_Q_nope, - make_stride(params.stride_q_h_q, _1{}, params.stride_q_s_q) - ) - ), - SmemLayoutQNoPE{} - ); - - auto shape_Q_rope = make_shape(params.h_q, D_Q-D_V, params.s_q); - auto tma_Q_rope = cute::make_tma_copy( - SM90_TMA_LOAD{}, - make_tensor( - make_gmem_ptr((bf16*)params.q + D_V), - make_layout( - shape_Q_rope, - make_stride(params.stride_q_h_q, _1{}, params.stride_q_s_q) - ) - ), - SmemLayoutQRoPE{} - ); - - auto shape_O = make_shape(params.h_q, params.d_v, params.s_q); - auto tma_O = cute::make_tma_copy( - SM90_TMA_STORE{}, - make_tensor( - make_gmem_ptr((bf16*)params.out), - make_layout( - shape_O, - make_stride(params.d_v, _1{}, params.h_q*params.d_v) - ) - ), - SmemLayoutOTiles<1>{} - ); - - - CUtensorMap tensor_map_kv_nope; - { - uint64_t size[2] = {D_V, (unsigned long)params.s_kv}; - uint64_t stride[1] = {params.stride_kv_s_kv*sizeof(bf16)}; - uint32_t box_size[2] = {64, 1}; - uint32_t elem_stride[2] = {1, 1}; - CUresult res = CUTLASS_CUDA_DRIVER_WRAPPER_CALL(cuTensorMapEncodeTiled)( - &tensor_map_kv_nope, - CUtensorMapDataType::CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, - 2, - params.kv, - size, - stride, - box_size, - elem_stride, - CUtensorMapInterleave::CU_TENSOR_MAP_INTERLEAVE_NONE, - CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_128B, - CUtensorMapL2promotion::CU_TENSOR_MAP_L2_PROMOTION_L2_256B, - CUtensorMapFloatOOBfill::CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE - ); - KU_ASSERT(res == CUresult::CUDA_SUCCESS); - } - - TmaParams< - decltype(shape_Q_nope), decltype(tma_Q_nope), - decltype(shape_Q_rope), decltype(tma_Q_rope), - decltype(shape_O), decltype(tma_O) - > tma_params = { - shape_Q_nope, tma_Q_nope, - shape_Q_rope, tma_Q_rope, - shape_O, tma_O, - tensor_map_kv_nope - }; - auto kernel = &sparse_attn_fwd_kernel; - - constexpr size_t smem_size = sizeof(SharedMemoryPlan); - KU_CUDA_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); - - kernel<<>>(params, tma_params); - KU_CHECK_KERNEL_LAUNCH(); -} - -} diff --git a/csrc/sm100/prefill/sparse/fwd/head64/phase1.h b/csrc/sm100/prefill/sparse/fwd/head64/phase1.h deleted file mode 100644 index 2962389d..00000000 --- a/csrc/sm100/prefill/sparse/fwd/head64/phase1.h +++ /dev/null @@ -1,10 +0,0 @@ -#pragma once - -#include "params.h" - -namespace sm100::fwd::head64 { - -template -void run_fwd_phase1_kernel(const SparseAttnFwdParams& params); - -} diff --git a/csrc/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512.cu b/csrc/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512.cu deleted file mode 100644 index a8b48956..00000000 --- a/csrc/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512.cu +++ /dev/null @@ -1,8 +0,0 @@ -#include "../phase1.h" -#include "../phase1.cuh" - -namespace sm100::fwd_for_small_topk::head128 { - -template void run_fwd_for_small_topk_phase1_kernel(const SparseAttnDecodeParams& params); - -} diff --git a/csrc/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_prefill_k512.cu b/csrc/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_prefill_k512.cu deleted file mode 100644 index 2f17fed9..00000000 --- a/csrc/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_prefill_k512.cu +++ /dev/null @@ -1,8 +0,0 @@ -#include "../phase1.h" -#include "../phase1.cuh" - -namespace sm100::fwd_for_small_topk::head128 { - -template void run_fwd_for_small_topk_phase1_kernel(const SparseAttnFwdParams& params); - -} diff --git a/csrc/sm100/prefill/sparse/fwd_for_small_topk/head128/phase1.h b/csrc/sm100/prefill/sparse/fwd_for_small_topk/head128/phase1.h deleted file mode 100644 index d1a092a3..00000000 --- a/csrc/sm100/prefill/sparse/fwd_for_small_topk/head128/phase1.h +++ /dev/null @@ -1,10 +0,0 @@ -#pragma once - -#include "params.h" - -namespace sm100::fwd_for_small_topk::head128 { - -template -void run_fwd_for_small_topk_phase1_kernel(const SparseFwdArgT& params); - -} diff --git a/csrc/sm90/prefill/sparse/fwd.cu b/csrc/sm90/prefill/sparse/fwd.cu deleted file mode 100644 index eb62b5d3..00000000 --- a/csrc/sm90/prefill/sparse/fwd.cu +++ /dev/null @@ -1,30 +0,0 @@ -#include "fwd.h" - -#include - -#include "phase1.h" - -namespace sm90 { - -void run_fwd_kernel(const SparseAttnFwdParams& params) { - const bool have_topk_length = params.topk_length != nullptr; - - // Dispatch based on d_qk dimension and presence of topk_length - if (params.d_qk == 512) { - if (have_topk_length) { - sm90::fwd::run_fwd_phase1_kernel<512, true>(params); - } else { - sm90::fwd::run_fwd_phase1_kernel<512, false>(params); - } - } else if (params.d_qk == 576) { - if (have_topk_length) { - sm90::fwd::run_fwd_phase1_kernel<576, true>(params); - } else { - sm90::fwd::run_fwd_phase1_kernel<576, false>(params); - } - } else { - throw std::runtime_error("Unsupported d_qk value in sparse attention fwd kernel"); - } -} - -} // namespace sm90 diff --git a/csrc/sm90/prefill/sparse/fwd.h b/csrc/sm90/prefill/sparse/fwd.h deleted file mode 100644 index 1c26d688..00000000 --- a/csrc/sm90/prefill/sparse/fwd.h +++ /dev/null @@ -1,9 +0,0 @@ -#pragma once - -#include "params.h" - -namespace sm90 { - -void run_fwd_kernel(const SparseAttnFwdParams& params); - -} diff --git a/flash_mla/__init__.py b/flash_mla/__init__.py index 02a8bbab..be366899 100644 --- a/flash_mla/__init__.py +++ b/flash_mla/__init__.py @@ -9,11 +9,14 @@ flash_mla_sparse_fwd ) +from . import fused_norm_rope_attn_rope_cast + __all__ = [ "get_mla_metadata", "flash_mla_with_kvcache", "flash_attn_varlen_func", "flash_attn_varlen_qkvpacked_func", "flash_attn_varlen_kvpacked_func", - "flash_mla_sparse_fwd" + "flash_mla_sparse_fwd", + "fused_norm_rope_attn_rope_cast" ] diff --git a/flash_mla/flash_mla_interface.py b/flash_mla/flash_mla_interface.py index a3740b0f..a9e8909f 100644 --- a/flash_mla/flash_mla_interface.py +++ b/flash_mla/flash_mla_interface.py @@ -97,6 +97,15 @@ def flash_mla_with_kvcache( - 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 and DeepSeek V4.1: + head_dim should be 512 while head_dim_v should be 512. + In FP8+sparse mode, 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 with an fp4 extra cache). + In all three, a page block stores `page_block_size` data rows first and `page_block_size` scale rows afterwards, so the scale factors are not interleaved into the data rows: + - V4 (584 Bytes per token): the data row is 448 float8_e4m3 NoPE values followed by 64 bfloat16 RoPE values (not quantized); the scale row is 8 Bytes, of which the first 7 are float8_e8m0 scales (the 8th byte is padding), each covering 64 consecutive float8_e4m3 values of the NoPE part. + - V4.1 (528 Bytes per token): the data row is 512 float8_e4m3 values (the 64 RoPE dimensions are quantized as well, so there is no bfloat16 part); the scale row is 16 Bytes of float8_e8m0 scales, each covering 32 consecutive float8_e4m3 values. + - V4.1 fp4 (288 Bytes per token): only valid for `extra_k_cache`, and only when `k_cache` is in the V4.1 format; the data row is 256 Bytes containing 512 e2m1 values (2 values per byte, the even-indexed one in the low nibble), and the scale row is 32 Bytes of float8_e4m3 scales, each covering 16 consecutive e2m1 values. + See tests/quant.py for quantization and dequantization details. + Return: out: (batch_size, seq_len_q, num_heads_q, head_dim_v). softmax_lse: (batch_size, num_heads_q, seq_len_q), torch.float32. diff --git a/flash_mla/fused_norm_rope_attn_rope_cast.py b/flash_mla/fused_norm_rope_attn_rope_cast.py new file mode 100644 index 00000000..0cb02ca2 --- /dev/null +++ b/flash_mla/fused_norm_rope_attn_rope_cast.py @@ -0,0 +1,203 @@ +from typing import Optional, Tuple + +import torch + +import flash_mla.cuda as flash_mla_cuda + + +def prefill( + enable_q_norm: bool, + rms_norm_eps: float, + + token_positions: torch.Tensor, + is_rope_neox_style: bool, + rope_dim: int, + cos_sin_cache: torch.Tensor, + + n_wv_group: int, + num_per_channels: int, + use_tma_aligned_col_major_sf: bool, + round_sf: bool, + use_packed_ue8m0: bool, + + q: torch.Tensor, + kv: torch.Tensor, + indices: torch.Tensor, + sm_scale: float, + d_v: int = 512, + attn_sink: Optional[torch.Tensor] = None, + topk_length: Optional[torch.Tensor] = None, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """ + A fused kernel for Q Norm + Q RoPE + Core Attn (sparse attention) + O RoPE + O cast to FP8, for DeepSeek-V4 & DeepSeek-V4.1 + Only support sm100 / sm103 GPU architecture. + + Args (Norm): + enable_q_norm: bool, whether to apply Q RMSNorm + rms_norm_eps: float. EPS for RMSNorm + + Args (RoPE): + token_positions: [s_q], int32 + is_rope_neox_style: must be False now + rope_dim: int32, must be 64 now + cos_sin_cache: [*, rope_dim], float32 + + Args (Cast): + n_wv_group: n_wv_group during o projection + num_per_channels: quantization granularity, must be 32 + use_tma_aligned_col_major_sf: must be True + round_sf: must be True + use_packed_ue8m0: must be True + + Args (Core Attn): + q: [s_q, h_q, d_qk], bfloat16 + kv: [s_kv, h_kv, d_qk], bfloat16 + indices: [s_q, h_kv, topk], int32. Invalid indices should be set to < 0, or >= s_kv + sm_scale: float, scaling factor for attention scores + d_v: value dimension, default (and only) is 512 + attn_sink: optional, [h_q], float32. + If attn_sink is provided, when computing output, output will be additionally multiplied by exp(lse) / (exp(lse) + exp(attn_sink)). + +-inf in attn_sink will be handled normally (i.e., -inf has no effect, +inf will make corresponding output all zeros). This has no effect on lse and max_logits. + topk_length: optional, [s_q], int32. If provided, the i-th q token will only attend to k tokens specified by indices[i, :, :topk_length[i]], ignoring later k tokens (even if provided in indices). This parameter is mainly used for variable-length topk attention scenarios, such as using sparse attention to simulate causal attention. + In extremely rare cases (topk_length provided, there is a valid topk index between topk_length[i] ~ s_kv, and that topk index points to a k token containing NaN), operator output will contain NaN, so please avoid this situation. + + Returns: + - out_fp8: [s_q, n_wv_group, wv_group_size * d_v], fp8_e4m3, quantized attention result + - out_sf: [s_q, n_wv_group, wv_group_size * d_v / (32*4)], int32_t, scaling factor. This scaling factor is ALWAYS stored in the per-32 scaled format, even if num_per_channels is 128 + - max_logits: [s_q, h_q], float + - lse: [s_q, h_q], float + If a q token does not attend to any k token, then max_logits is -inf, lse is +inf, out is all zeros. + """ + results = flash_mla_cuda.fused_norm_rope_attn_rope_cast_fwd( + q, kv, indices, sm_scale, d_v, attn_sink, topk_length, + + enable_q_norm, rms_norm_eps, token_positions, is_rope_neox_style, rope_dim, cos_sin_cache, + + n_wv_group, num_per_channels, use_tma_aligned_col_major_sf, round_sf, use_packed_ue8m0 + ) + out_fp8, out_sf, max_logits, lse = results + return out_fp8, out_sf, max_logits, lse + + +def decode( + enable_q_norm: bool, + rms_norm_eps: float, + + token_positions: torch.Tensor, + is_rope_neox_style: bool, + rope_dim: int, + cos_sin_cache: torch.Tensor, + + n_wv_group: int, + num_per_channels: int, + use_tma_aligned_col_major_sf: bool, + round_sf: bool, + use_packed_ue8m0: bool, + + q: torch.Tensor, + k_cache: torch.Tensor, + indices_in_kvcache: torch.Tensor, + sm_scale: float, + d_v: int = 512, + attn_sink: Optional[torch.Tensor] = None, + topk_length: Optional[torch.Tensor] = None, + extra_k_cache: Optional[torch.Tensor] = None, + extra_indices_in_kvcache: Optional[torch.Tensor] = None, + extra_topk_length: Optional[torch.Tensor] = None, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """ + Fused Decoding kernel: Q Norm + Q RoPE + Core Attn (decode, with paged FP8 KV cache) + O RoPE + O cast to FP8, for DeepSeek-V4 & DeepSeek-V4.1 + Only support sm100 / sm103 GPU architecture. + + The batch size is always 1, and the batch dimension is not included in any tensor shape below. + + Args (Norm): + enable_q_norm: bool, whether to apply Q RMSNorm + rms_norm_eps: float. EPS for RMSNorm + + Args (RoPE): + token_positions: [s_q], int32 + is_rope_neox_style: must be False now + rope_dim: int32, must be 64 now + cos_sin_cache: [*, rope_dim], float32 + + Args (Cast): + n_wv_group: n_wv_group during o projection + num_per_channels: quantization granularity, must be 32 + use_tma_aligned_col_major_sf: must be True + round_sf: must be True + use_packed_ue8m0: must be True + + Args (Core Attn): + q: [s_q, h_q, d_qk], bfloat16. Already permuted by permute_q_b_proj + k_cache: [num_blocks, page_block_size, h_kv, bytes_per_token], fp8_e4m3 (or int8/uint8). Paged quantized KV cache. The format is + detected from bytes_per_token: 584 (V4), 528 (V4.1) or 288 (V4.1 with fp4 e2m1 + per-16 e4m3 scales). See tests/quant.py for the layouts + indices_in_kvcache: [s_q, topk], int32. Page-relative KV token indices + sm_scale: float, scaling factor for attention scores + d_v: value dimension, default (and only) is 512 + attn_sink: optional, [h_q], float32. Per-head attention sink bias + topk_length: optional, [s_q], int32. Actual valid topk count of the request + extra_k_cache: optional, [extra_num_blocks, extra_page_block_size, h_kv, bytes_per_token]. Secondary paged FP8 KV cache + extra_indices_in_kvcache: optional, [s_q, extra_topk], int32. Indices into the extra KV cache + extra_topk_length: optional, [s_q], int32. Actual valid extra topk count of the request + + Returns: + - out_fp8: [s_q, n_wv_group, wv_group_size * d_v], fp8_e4m3, quantized attention result + - out_sf: [s_q, n_wv_group, wv_group_size * d_v / (32*4)], int32, scaling factor + - lse: [s_q, h_q], float + """ + out_fp8, out_sf, lse = flash_mla_cuda.fused_norm_rope_attn_rope_cast_decode( + q, k_cache, indices_in_kvcache, sm_scale, d_v, + attn_sink, topk_length, + extra_k_cache, extra_indices_in_kvcache, extra_topk_length, + enable_q_norm, rms_norm_eps, token_positions, is_rope_neox_style, rope_dim, cos_sin_cache, + n_wv_group, num_per_channels, use_tma_aligned_col_major_sf, round_sf, use_packed_ue8m0 + ) + return out_fp8, out_sf, lse + + +def permute_q_b_proj( + weight_and_sf: Tuple[torch.Tensor, torch.Tensor], + h_q: int, + d_q: int, +) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Permute the layout of the q_b_proj weight and its scale factors into the layout required by + fused_norm_rope_attn_rope_cast_fwd / fused_norm_rope_attn_rope_cast_decode. + + Args: + weight_and_sf: (q_b_proj, q_b_sf), where + - q_b_proj: [h_q*d_q, q_lora_rank], fp8_e4m3, the weight in its original layout + - q_b_sf: [h_q*d_q, q_lora_rank/gran/4], int32, the scale factors in their original layout (DeepGeMM format, i.e. contiguous on the first dim) + h_q: int, the number of Q heads + d_q: int, the Q head dimension + + Returns: + - q_b_proj_permuted: [h_q*d_q, q_lora_rank], fp8_e4m3 + - scale_factors_permuted: [h_q*d_q, q_lora_rank/gran/4], int32 + """ + return flash_mla_cuda.permute_q_b_proj(weight_and_sf[0], weight_and_sf[1], h_q, d_q) + + +def permute_wv_proj( + weight_and_sf: Tuple[torch.Tensor, torch.Tensor], + wv_group_size: int, + d_o: int, +) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Permute the layout of the wv_proj weight and its scale factors into the layout required by + fused_norm_rope_attn_rope_cast_fwd / fused_norm_rope_attn_rope_cast_decode. + Note: regardless of the input quantization granularity, the output granularity is always 32. + + Args: + weight_and_sf: (wv_proj, wv_sf), where + - wv_proj: [n_wv_groups, d_proj_out, wv_group_size * d_o], fp8_e4m3, the weight in its original layout + - wv_sf: [n_wv_groups, d_proj_out, wv_group_size * d_o/gran/4], int32, the scale factors in their original layout (DeepGeMM format, i.e. contiguous on the second dim) + wv_group_size: int + d_o: int, the O head dimension + + Returns: + - wv_proj_permuted: the same shape as wv_proj, fp8_e4m3 + - scale_factors_permuted: [n_wv_groups, d_proj_out, wv_group_size * d_o/32/4], int32 + """ + return flash_mla_cuda.permute_wv_proj(weight_and_sf[0], weight_and_sf[1], wv_group_size, d_o) diff --git a/setup.py b/setup.py index 513b4355..f09444ce 100644 --- a/setup.py +++ b/setup.py @@ -40,7 +40,9 @@ def get_arch_flags(): arch_flags = [] if not DISABLE_SM100: - arch_flags.extend(["-gencode", "arch=compute_100f,code=sm_100f"]) + # We use architecture-specific (sm_100a / sm_103a) targets instead of the family-specific one (sm_100f) for better SASS code generation + arch_flags.extend(["-gencode", "arch=compute_100a,code=sm_100a"]) + arch_flags.extend(["-gencode", "arch=compute_103a,code=sm_103a"]) if not DISABLE_SM90: arch_flags.extend(["-gencode", "arch=compute_90a,code=sm_90a"]) return arch_flags @@ -65,43 +67,79 @@ def get_nvcc_thread_args(): sources=[ # API "csrc/api/api.cpp", + "csrc/api/sparse_prefill.cpp", + "csrc/api/sparse_decode.cpp", + "csrc/api/dense_fwd.cpp", + "csrc/api/dense_bwd.cpp", + "csrc/api/dense_decode.cpp", + "csrc/api/fused_norm_rope_attn_rope_cast_fwd.cpp", # Misc kernels for decoding - "csrc/smxx/decode/get_decoding_sched_meta/get_decoding_sched_meta.cu", - "csrc/smxx/decode/combine/combine.cu", + "csrc/kernels/smxx/decode/get_decoding_sched_meta/get_decoding_sched_meta.cu", + "csrc/kernels/smxx/decode/combine/combine.cu", # sm90 dense decode - "csrc/sm90/decode/dense/instantiations/fp16.cu", - "csrc/sm90/decode/dense/instantiations/bf16.cu", + "csrc/kernels/sm90/decode/dense/instantiations/fp16.cu", + "csrc/kernels/sm90/decode/dense/instantiations/bf16.cu", # sm90 sparse decode - "csrc/sm90/decode/sparse_fp8/instantiations/model1_persistent_h64.cu", - "csrc/sm90/decode/sparse_fp8/instantiations/model1_persistent_h128.cu", - "csrc/sm90/decode/sparse_fp8/instantiations/v32_persistent_h64.cu", - "csrc/sm90/decode/sparse_fp8/instantiations/v32_persistent_h128.cu", + "csrc/kernels/sm90/decode/sparse/instantiations/v4_persistent_h64.cu", + "csrc/kernels/sm90/decode/sparse/instantiations/v4_persistent_h128.cu", + "csrc/kernels/sm90/decode/sparse/instantiations/v32_persistent_h64.cu", + "csrc/kernels/sm90/decode/sparse/instantiations/v32_persistent_h128.cu", # sm90 sparse prefill - "csrc/sm90/prefill/sparse/fwd.cu", - "csrc/sm90/prefill/sparse/instantiations/phase1_k512.cu", - "csrc/sm90/prefill/sparse/instantiations/phase1_k512_topklen.cu", - "csrc/sm90/prefill/sparse/instantiations/phase1_k576.cu", - "csrc/sm90/prefill/sparse/instantiations/phase1_k576_topklen.cu", + "csrc/kernels/sm90/prefill/sparse/instantiations/phase1_k512.cu", + "csrc/kernels/sm90/prefill/sparse/instantiations/phase1_k512_topklen.cu", + "csrc/kernels/sm90/prefill/sparse/instantiations/phase1_k576.cu", + "csrc/kernels/sm90/prefill/sparse/instantiations/phase1_k576_topklen.cu", # sm100 dense prefill & backward - "csrc/sm100/prefill/dense/fmha_cutlass_fwd_sm100.cu", - "csrc/sm100/prefill/dense/fmha_cutlass_bwd_sm100.cu", + "csrc/kernels/sm100/prefill/dense/fmha_cutlass_fwd_sm100.cu", + "csrc/kernels/sm100/prefill/dense/fmha_cutlass_bwd_sm100.cu", # sm100 sparse prefill - "csrc/sm100/prefill/sparse/fwd/head64/instantiations/phase1_k512.cu", - "csrc/sm100/prefill/sparse/fwd/head64/instantiations/phase1_k576.cu", - "csrc/sm100/prefill/sparse/fwd/head128/instantiations/phase1_k512.cu", - "csrc/sm100/prefill/sparse/fwd/head128/instantiations/phase1_k576.cu", - "csrc/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_prefill_k512.cu", + "csrc/kernels/sm100/prefill/sparse/fwd/head64/instantiations/phase1_h64_k512.cu", + "csrc/kernels/sm100/prefill/sparse/fwd/head64/instantiations/phase1_h64_k576.cu", + "csrc/kernels/sm100/prefill/sparse/fwd/head128/instantiations/phase1_k512.cu", + "csrc/kernels/sm100/prefill/sparse/fwd/head128/instantiations/phase1_k576.cu", + "csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_k512.cu", + + # sm100 fused norm + rope + attn + rope + cast + "csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h64_prefill_norm.cu", + "csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h64_prefill_nonorm.cu", + "csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h128_prefill_norm.cu", + "csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h128_prefill_nonorm.cu", + "csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h64_decode_norm.cu", + "csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h64_decode_nonorm.cu", + "csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h128_decode_norm.cu", + "csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h128_decode_nonorm.cu", + "csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41_h64_decode_norm.cu", + "csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41_h64_decode_nonorm.cu", + "csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41_h128_decode_norm.cu", + "csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41_h128_decode_nonorm.cu", + "csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41fp4_h64_decode_norm.cu", + "csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41fp4_h64_decode_nonorm.cu", + "csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41fp4_h128_decode_norm.cu", + "csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41fp4_h128_decode_nonorm.cu", + "csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/permute_q_b_proj/kernel.cu", + "csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/permute_wv_proj/kernel.cu", # sm100 sparse decode - "csrc/sm100/decode/head64/instantiations/v32.cu", - "csrc/sm100/decode/head64/instantiations/model1.cu", - "csrc/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512.cu", + "csrc/kernels/sm100/decode/sparse/head64/instantiations/v32_h64.cu", + "csrc/kernels/sm100/decode/sparse/head64/instantiations/v32_h64_no_split.cu", + "csrc/kernels/sm100/decode/sparse/head64/instantiations/v4_h64.cu", + "csrc/kernels/sm100/decode/sparse/head64/instantiations/v4_h64_no_split.cu", + "csrc/kernels/sm100/decode/sparse/head64/instantiations/v41_h64.cu", + "csrc/kernels/sm100/decode/sparse/head64/instantiations/v41_h64_no_split.cu", + "csrc/kernels/sm100/decode/sparse/head64/instantiations/v41fp4_h64.cu", + "csrc/kernels/sm100/decode/sparse/head64/instantiations/v41fp4_h64_no_split.cu", + "csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512.cu", + "csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512_splitkv.cu", + "csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512_v41.cu", + "csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512_v41_splitkv.cu", + "csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512_v41fp4.cu", + "csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512_v41fp4_splitkv.cu", ], extra_compile_args={ "cxx": cxx_args + get_features_args(), @@ -125,10 +163,11 @@ def get_nvcc_thread_args(): }, include_dirs=[ Path(this_dir) / "csrc", - Path(this_dir) / "csrc" / "kerutils" / "include", # TODO Remove me - Path(this_dir) / "csrc" / "sm90", + Path(this_dir) / "csrc" / "kerutils" / "include", Path(this_dir) / "csrc" / "cutlass" / "include", Path(this_dir) / "csrc" / "cutlass" / "tools" / "util" / "include", + Path(CUDA_HOME) / "targets" / "x86_64-linux" / "include" / "cccl", # for cuda/std headers in CUDA 13+ + Path(CUDA_HOME) / "targets" / "sbsa-linux" / "include" / "cccl", ], ) ) diff --git a/tests/lib.py b/tests/lib.py index 139e130d..ed6fad38 100644 --- a/tests/lib.py +++ b/tests/lib.py @@ -4,6 +4,7 @@ from typing import List, Optional import random +import argparse import torch import kernelkit as kk import flash_mla @@ -24,7 +25,9 @@ class ExtraTestParamForDecode: block_size: int = 64 extra_block_size: Optional[int] = None have_extra_topk_length: bool = False - + kvcache_layout: Optional["quant.KVCacheLayout"] = None # Must be specified for d_qk == 512 to distinguish V4 / V41 / V41_FP4 + extra_kvcache_layout: Optional["quant.KVCacheLayout"] = None # None: same as kvcache_layout + @dataclasses.dataclass class TestParam: s_q: int @@ -41,6 +44,8 @@ class TestParam: have_attn_sink: bool = False have_topk_length: bool = False decode: Optional[ExtraTestParamForDecode] = None + k_amplifier_portion: float = 0.0 # Amplify a portion of the KV tokens to create a more skewed attention distribution + k_amplifier_ratio: float = 1.0 # Amplification ratio for the amplified KV tokens @dataclasses.dataclass class RawTestParamForDecode: @@ -65,6 +70,8 @@ class RawTestParamForDecode: block_size: int = 64 extra_block_size: Optional[int] = None have_extra_topk_length: bool = False + kvcache_layout: Optional["quant.KVCacheLayout"] = None + extra_kvcache_layout: Optional["quant.KVCacheLayout"] = None d_qk: int = 576 # Q/K head dim (= dv + RoPE dim) d_v: int = 512 # V head dim check_correctness: bool = True @@ -82,7 +89,8 @@ def to_test_param(self) -> TestParam: decode = ExtraTestParamForDecode( self.b, self.is_varlen, self.have_zero_seqlen_k, self.extra_s_k, self.extra_topk, - self.block_size, self.extra_block_size, self.have_extra_topk_length + self.block_size, self.extra_block_size, self.have_extra_topk_length, + self.kvcache_layout, self.extra_kvcache_layout ) ) @@ -147,6 +155,11 @@ def generate_testcase(t: TestParam) -> Testcase: if t.have_topk_length: topk_length = torch.randint(0, max(t.topk + 1, 64), (t.s_q, ), dtype=torch.int32, device=q.device).clamp_max(t.topk) + if t.k_amplifier_portion > 0.0: + selected_indices = torch.randint(0, t.s_kv, (int(t.s_kv * t.k_amplifier_portion), ), device=kv.device) + amplifier_coeffs = torch.rand((selected_indices.size(0), ), device=kv.device) * (t.k_amplifier_ratio - 1) + 1 + kv[selected_indices] *= amplifier_coeffs.unsqueeze(-1).unsqueeze(-1) + q = kk.non_contiguousify(q) kv = kk.non_contiguousify(kv) do = kk.non_contiguousify(do) @@ -173,6 +186,7 @@ class KVScope: abs_indices: torch.Tensor indices_in_kvcache: torch.Tensor topk_length: Optional[torch.Tensor] + kvcache_layout: Optional["quant.KVCacheLayout"] = None blocked_k_quantized: Optional[torch.Tensor] = None def quant_and_dequant_(self): @@ -180,16 +194,18 @@ def quant_and_dequant_(self): For FP8 cases, we need to quantize the KV cache for Flash MLA. Besides, the quantization error may be too large to be distinguished from wrong kernels, so we de-quantize kvcache here to mitigate quantization error """ - fp8_kvcache_layout = None - if self.t.d_qk == 576: - fp8_kvcache_layout = quant.FP8KVCacheLayout.V32_FP8Sparse - elif self.t.d_qk == 512: - assert self.abs_indices is not None - fp8_kvcache_layout = quant.FP8KVCacheLayout.MODEL1_FP8Sparse - else: - assert False - self.blocked_k_quantized = quant.quantize_k_cache(self.blocked_k, fp8_kvcache_layout) - blocked_k_dequantized = quant.dequantize_k_cache(self.blocked_k_quantized, fp8_kvcache_layout) + kvcache_layout = self.kvcache_layout + if kvcache_layout is None: + if self.t.d_qk == 576: + kvcache_layout = quant.KVCacheLayout.V32_FP8Sparse + elif self.t.d_qk == 512: + assert self.abs_indices is not None + kvcache_layout = quant.KVCacheLayout.V4_FP8Sparse + else: + assert False + self.kvcache_layout = kvcache_layout + self.blocked_k_quantized = quant.quantize_k_cache(self.blocked_k, kvcache_layout) + blocked_k_dequantized = quant.dequantize_k_cache(self.blocked_k_quantized, kvcache_layout) self.blocked_k = blocked_k_dequantized def get_kvcache_for_flash_mla(self) -> torch.Tensor: @@ -211,6 +227,7 @@ def apply_perm(self, perm: torch.Tensor) -> "KVScope": self.abs_indices[perm], self.indices_in_kvcache[perm], self.topk_length[perm] if self.topk_length is not None else None, + self.kvcache_layout, self.blocked_k_quantized ) return new_kvscope @@ -239,7 +256,7 @@ def generate_testcase_for_decode(t: TestParam) -> TestcaseForDecode: attn_sink[inf_mask > 0.5] = float("inf") attn_sink[inf_mask < -0.5] = float("-inf") - def generate_one_k_scope(s_k: int, block_size: int, topk: int, is_varlen: bool, have_zero_seqlen: bool, is_all_indices_invalid: bool, have_topk_length: bool) -> KVScope: + def generate_one_k_scope(s_k: int, block_size: int, topk: int, is_varlen: bool, have_zero_seqlen: bool, is_all_indices_invalid: bool, have_topk_length: bool, kvcache_layout: Optional[quant.KVCacheLayout] = None) -> KVScope: b = t.decode.b # type: ignore cache_seqlens_cpu = torch.full((b,), s_k, dtype=torch.int32, device='cpu') if is_varlen: @@ -286,16 +303,17 @@ def generate_one_k_scope(s_k: int, block_size: int, topk: int, is_varlen: bool, block_table = kk.non_contiguousify(block_table) abs_indices = kk.non_contiguousify(abs_indices) indices_in_kvcache = kk.non_contiguousify(indices_in_kvcache) - return KVScope(t, cache_seqlens, block_table, blocked_k, abs_indices, indices_in_kvcache, topk_length) + return KVScope(t, cache_seqlens, block_table, blocked_k, abs_indices, indices_in_kvcache, topk_length, kvcache_layout) - kv_scope0 = generate_one_k_scope(t.s_kv, t.decode.block_size, t.topk, t.decode.is_varlen, t.decode.have_zero_seqlen_k, t.is_all_indices_invalid, t.have_topk_length) + kv_scope0 = generate_one_k_scope(t.s_kv, t.decode.block_size, t.topk, t.decode.is_varlen, t.decode.have_zero_seqlen_k, t.is_all_indices_invalid, t.have_topk_length, t.decode.kvcache_layout) kv_scope0.quant_and_dequant_() if t.decode.extra_topk is not None: if t.decode.extra_s_k is None: t.decode.extra_s_k = t.decode.extra_topk*2 if t.decode.extra_block_size is None: t.decode.extra_block_size = t.decode.block_size - kv_scope1 = generate_one_k_scope(t.decode.extra_s_k, t.decode.extra_block_size, t.decode.extra_topk, t.decode.is_varlen, t.decode.have_zero_seqlen_k, t.is_all_indices_invalid, t.decode.have_extra_topk_length) + extra_layout = t.decode.extra_kvcache_layout if t.decode.extra_kvcache_layout is not None else t.decode.kvcache_layout + kv_scope1 = generate_one_k_scope(t.decode.extra_s_k, t.decode.extra_block_size, t.decode.extra_topk, t.decode.is_varlen, t.decode.have_zero_seqlen_k, t.is_all_indices_invalid, t.decode.have_extra_topk_length, extra_layout) kv_scope1.quant_and_dequant_() else: assert t.decode.extra_block_size is None and t.decode.extra_s_k is None and not t.decode.have_extra_topk_length @@ -307,8 +325,7 @@ def generate_one_k_scope(s_k: int, block_size: int, topk: int, is_varlen: bool, return TestcaseForDecode(t, q, attn_sink, sm_scale, kv_scope0, kv_scope1) -def run_flash_mla_sparse_fwd(p: TestParam, t: Testcase, return_p_sum: bool): - assert not return_p_sum +def run_flash_mla_sparse_fwd(p: TestParam, t: Testcase): return flash_mla.flash_mla_sparse_fwd( t.q, t.kv, t.indices, sm_scale=t.sm_scale, @@ -341,6 +358,7 @@ class FlopsAndMemVolStatistics: """ fwd_flop: float fwd_mem_vol: float + fwd_prefill_with_fp8_out_mem_vol: float = 0.0 # Like `fwd_mem_vol`, but with the output stored as FP8 instead of bf16 def count_flop_and_mem_vol(p: TestParam, t: Testcase) -> FlopsAndMemVolStatistics: total_topk = (p.s_q*p.topk) if t.topk_length is None else t.topk_length.sum().item() @@ -354,6 +372,7 @@ def count_flop_and_mem_vol(p: TestParam, t: Testcase) -> FlopsAndMemVolStatistic return FlopsAndMemVolStatistics( fwd_flop, fwd_mem_vol, + fwd_mem_vol - p.s_q*p.h_q*p.d_v, # The FP8 output only stores d_v bytes per element (and no separate SF traffic is counted) ) @dataclasses.dataclass @@ -403,3 +422,7 @@ def get_num_retrieved_tokens(kv_scope: KVScope) -> int: def is_no_cooldown() -> bool: return os.environ.get('NO_COOLDOWN', '').lower() in ['1', 'yes', 'y'] + +def stick_unit_test_args(parser: argparse.ArgumentParser): + parser.add_argument("-nc", "--no-cooldown", action="store_true", help="Don't call time.sleep() before performance testcases") + parser.add_argument("-rf", "--run-to-finish", action="store_true", help="Don't exit when a testcase is failed") diff --git a/tests/quant.py b/tests/quant.py index b92b5399..9db53f02 100644 --- a/tests/quant.py +++ b/tests/quant.py @@ -3,23 +3,70 @@ import torch -class FP8KVCacheLayout(enum.Enum): +import kernelkit as kk + +class KVCacheLayout(enum.Enum): + V32_FP8 = 0 V32_FP8Sparse = 1 - MODEL1_FP8Sparse = 2 + V4_FP8Sparse = 2 + V41_FP8Sparse = 3 + V41_FP4 = 4 def get_meta(self) -> Tuple[int, int, int, int, int]: # Return: (d, d_nope, d_rope, tile_size, num_tiles) return { - FP8KVCacheLayout.V32_FP8Sparse: (576, 512, 64, 128, 4), - FP8KVCacheLayout.MODEL1_FP8Sparse: (512, 448, 64, 64, 7) + KVCacheLayout.V32_FP8: (576, 512, 64, 128, 4), + KVCacheLayout.V32_FP8Sparse: (576, 512, 64, 128, 4), + KVCacheLayout.V4_FP8Sparse: (512, 448, 64, 64, 7), + KVCacheLayout.V41_FP8Sparse: (512, 448, 64, 32, 16), # 14 NoPE + 2 RoPE tiles + KVCacheLayout.V41_FP4: (512, 448, 64, 16, 32), # 28 NoPE + 4 RoPE tiles, all fp4 + }[self] + + def get_bytes_per_token(self) -> int: + d, d_nope, d_rope, tile_size, num_tiles = self.get_meta() + return { + KVCacheLayout.V32_FP8: d_nope + num_tiles*4 + 2*d_rope, + KVCacheLayout.V32_FP8Sparse: d_nope + num_tiles*4 + 2*d_rope, + KVCacheLayout.V4_FP8Sparse: d_nope + 2*d_rope + num_tiles + 1, + KVCacheLayout.V41_FP8Sparse: d_nope + d_rope + num_tiles, + KVCacheLayout.V41_FP4: d // 2 + num_tiles, }[self] def _cast_scale_inv_to_ue8m0(scales_inv: torch.Tensor, out_dtype = torch.float32) -> torch.Tensor: - return torch.pow(2, torch.clamp_min(scales_inv, 1e-4).log2().ceil()).to(out_dtype) + return torch.pow(2, torch.clamp_min(scales_inv, 1e-4).log2().ceil()).to(out_dtype) # This 1e-4 align with Tile Kernel's FP8_AMAX_MARGIN + +_E2M1_MAGNITUDES = torch.tensor([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0], dtype=torch.float32) # Indexed by the 3 low bits of the code, bit 3 is the sign + +def _quantize_to_e2m1(x: torch.Tensor) -> torch.Tensor: + """ + Round to the nearest fp4_e2m1 value with the semantics of PTX `cvt.rn.satfinite.e2m1x2.f32` (ties to even, saturating to +-6) + and return the 4-bit codes as uint8. NaN is mapped to code 0 (fp4 has no NaN; the caller keeps the NaN in the scale instead) + """ + x = x.float() + if x.numel() > (1 << 26): + # The tie detection below broadcasts an elements x 7 intermediate; quantize big caches in chunks + out = torch.empty(x.shape, dtype=torch.uint8, device=x.device) + for i in range(0, x.numel(), 1 << 26): + out.reshape(-1)[i:i + (1 << 26)] = _quantize_to_e2m1(x.reshape(-1)[i:i + (1 << 26)]) + return out + mags = _E2M1_MAGNITUDES.to(x.device) + sign = (torch.signbit(x)).to(torch.uint8) << 3 + a = torch.nan_to_num(x.abs(), nan=0.0, posinf=6.0).clamp_max(6.0) + mids = (mags[:-1] + mags[1:]) / 2 + code = torch.bucketize(a, mids, right=True) + on_tie = (a.unsqueeze(-1) == mids).any(dim=-1) + tie_code = torch.bucketize(a, mids, right=False) # On a tie: the lower of the two candidate codes + code = torch.where(on_tie, tie_code + (tie_code & 1), code) + return (sign | code.to(torch.uint8)).to(torch.uint8) + +def _dequantize_e2m1(codes: torch.Tensor) -> torch.Tensor: + mags = _E2M1_MAGNITUDES.to(codes.device) + val = mags[(codes & 7).long()] + return torch.where((codes & 8) != 0, -val, val) def quantize_k_cache( input_k_cache: torch.Tensor, # (num_blocks, block_size, h_k, d) - kvcache_layout: FP8KVCacheLayout, + kvcache_layout: KVCacheLayout, ) -> torch.Tensor: """ Quantize the k-cache @@ -32,7 +79,28 @@ def quantize_k_cache( input_k_cache = input_k_cache.squeeze(2) # [num_blocks, block_size, d] input_elem_size = input_k_cache.element_size() - if kvcache_layout == FP8KVCacheLayout.V32_FP8Sparse: + if kvcache_layout == KVCacheLayout.V32_FP8: + bytes_per_block = block_size*d_nope + block_size*num_tiles*4 + block_size*input_elem_size*d_rope + result = kk.gen_non_contiguous_tensor((num_blocks, bytes_per_block), dtype=torch.float8_e4m3fn, device=input_k_cache.device) + result_k_nope_part = result[..., :block_size*d_nope].view(num_blocks, 512//16, block_size, 16) + result_k_scale_factor = result[..., block_size*d_nope:block_size*(d_nope+4*num_tiles)].view(torch.float32).view(num_blocks, num_tiles, block_size).permute(0, 2, 1) # [num_blocks, block_size, num_tiles] + result_k_rope_part = result[..., block_size*(d_nope+4*num_tiles):].view(input_k_cache.dtype).view(num_blocks, block_size, d_rope) + + result_k_rope_part[:] = input_k_cache[..., d_nope:] + + for tile_idx in range(0, num_tiles): + cur_scale_factors_inv = torch.abs(input_k_cache[..., tile_idx*tile_size:(tile_idx+1)*tile_size]).max(dim=-1).values.float() / 448.0 # [num_blocks, block_size] + cur_scale_factors_inv = _cast_scale_inv_to_ue8m0(cur_scale_factors_inv) + result_k_scale_factor[:, :, tile_idx] = cur_scale_factors_inv + + cur_scale_factors_inv = cur_scale_factors_inv.view(num_blocks, block_size, 1) + cur_quantized_nope = (input_k_cache[..., tile_idx*tile_size:(tile_idx+1)*tile_size].float() / cur_scale_factors_inv.float()).to(torch.float8_e4m3fn) + result_k_nope_part[:, tile_idx*tile_size//16:(tile_idx+1)*tile_size//16, :, :].permute(0, 2, 1, 3)[:] = cur_quantized_nope.view(num_blocks, block_size, tile_size//16, 16) + + result = result.view(num_blocks, block_size, 1, -1) + return result + + elif kvcache_layout == KVCacheLayout.V32_FP8Sparse: bytes_per_token = d_nope + num_tiles*4 + input_elem_size*d_rope result = torch.empty((num_blocks, block_size+1, bytes_per_token), dtype=torch.float8_e4m3fn, device=input_k_cache.device)[:, :block_size, :] result_k_nope_part = result[..., :d_nope] @@ -52,7 +120,7 @@ def quantize_k_cache( result = result.view(num_blocks, block_size, 1, -1) return result - elif kvcache_layout == FP8KVCacheLayout.MODEL1_FP8Sparse: + elif kvcache_layout == KVCacheLayout.V4_FP8Sparse: bytes_per_token = d_nope + 2*d_rope + num_tiles + 1 size_per_block_padded = (block_size*bytes_per_token + 576-1) // 576 * 576 result = torch.empty((num_blocks, size_per_block_padded), dtype=torch.float8_e4m3fn, device=input_k_cache.device)[:, :block_size*bytes_per_token] @@ -70,7 +138,46 @@ def quantize_k_cache( cur_scale_factors_inv = cur_scale_factors_inv.view(num_blocks, block_size, 1) cur_quantized_nope = (input_k_cache[..., tile_idx*tile_size:(tile_idx+1)*tile_size].float() / cur_scale_factors_inv.float()).to(torch.float8_e4m3fn) result_k_nope[:, :, tile_idx*tile_size:(tile_idx+1)*tile_size] = cur_quantized_nope - + + result = result.view(num_blocks, block_size, 1, -1) + return result + + elif kvcache_layout == KVCacheLayout.V41_FP8Sparse: + bytes_per_token = d_nope + d_rope + num_tiles + size_per_block_padded = (block_size*bytes_per_token + 512-1) // 512 * 512 + result = torch.empty((num_blocks, size_per_block_padded), dtype=torch.float8_e4m3fn, device=input_k_cache.device)[:, :block_size*bytes_per_token] + result_k_nope_rope_part = result[:, :block_size*(d_nope+d_rope)].view(num_blocks, block_size, d_nope + d_rope) + result_k_scale_factor = result[:, block_size*(d_nope+d_rope):].view(num_blocks, block_size, num_tiles).view(torch.float8_e8m0fnu) # [num_blocks, block_size, 16] + + for tile_idx in range(0, 16): + cur_scale_factors_inv = torch.abs(input_k_cache[..., tile_idx*tile_size:(tile_idx+1)*tile_size]).max(dim=-1).values.float() / 448.0 + cur_scale_factors_inv = _cast_scale_inv_to_ue8m0(cur_scale_factors_inv) + result_k_scale_factor[:, :, tile_idx] = cur_scale_factors_inv.to(torch.float8_e8m0fnu) + cur_scale_factors_inv = cur_scale_factors_inv.view(num_blocks, block_size, 1) + cur_quantized_nope = (input_k_cache[..., tile_idx*tile_size:(tile_idx+1)*tile_size].float() / cur_scale_factors_inv.float()).to(torch.float8_e4m3fn) + result_k_nope_rope_part[:, :, tile_idx*tile_size:(tile_idx+1)*tile_size] = cur_quantized_nope + + result = result.view(num_blocks, block_size, 1, -1) + return result + + elif kvcache_layout == KVCacheLayout.V41_FP4: + # Block layout: [block_size x 256 B fp4 rows][block_size x 32 B e4m3 scale rows]. Element i of a row is in byte i//2, even + # elements in the low nibble. The scale is amax / 6 (6 = max magnitude of e2m1) rounded to e4m3, without a per-tensor scale + bytes_per_token = d // 2 + num_tiles + size_per_block_padded = (block_size*bytes_per_token + 512-1) // 512 * 512 + result = torch.empty((num_blocks, size_per_block_padded), dtype=torch.float8_e4m3fn, device=input_k_cache.device)[:, :block_size*bytes_per_token] + result_k_data = result[:, :block_size*(d//2)].view(torch.uint8).view(num_blocks, block_size, d//2) + result_k_scale = result[:, block_size*(d//2):].view(num_blocks, block_size, num_tiles) + + x = input_k_cache.float() + amax = torch.nan_to_num(x.abs(), nan=float("inf")).view(num_blocks, block_size, num_tiles, tile_size).amax(dim=-1) # A NaN element poisons the whole tile + scale = torch.clamp(amax / 6.0, 2.0**-9, 448.0).to(torch.float8_e4m3fn) # Clamp to the e4m3 range first: torch maps overflow to NaN + scale = torch.where(torch.isinf(amax), torch.full_like(scale, float("nan")), scale) + codes = _quantize_to_e2m1(x.view(num_blocks, block_size, num_tiles, tile_size) / scale.float().unsqueeze(-1)) + codes = codes.view(num_blocks, block_size, d) + result_k_data[:] = codes[..., 0::2] | (codes[..., 1::2] << 4) + result_k_scale[:] = scale + result = result.view(num_blocks, block_size, 1, -1) return result @@ -80,7 +187,7 @@ def quantize_k_cache( def dequantize_k_cache( quant_k_cache: torch.Tensor, # (num_blocks, block_size, 1, bytes_per_token) - kvcache_layout: FP8KVCacheLayout, + kvcache_layout: KVCacheLayout, ) -> torch.Tensor: """ De-quantize the k-cache @@ -90,10 +197,22 @@ def dequantize_k_cache( assert h_k == 1 result = torch.empty((num_blocks, block_size, d), dtype=torch.bfloat16, device=quant_k_cache.device) - if kvcache_layout == FP8KVCacheLayout.V32_FP8Sparse: + if kvcache_layout == KVCacheLayout.V32_FP8: + quant_k_cache = quant_k_cache.view(num_blocks, -1) # [num_blocks, ...] + input_nope = quant_k_cache[..., :block_size*d_nope].view(num_blocks, 512//16, block_size, 16).view(torch.float8_e4m3fn) + input_scale = quant_k_cache[..., block_size*d_nope:block_size*(d_nope+4*num_tiles)].view(torch.float32).view(num_blocks, num_tiles, block_size).permute(0, 2, 1).contiguous() # [num_blocks, block_size, num_tiles] + input_rope = quant_k_cache[..., block_size*(d_nope+4*num_tiles):].view(torch.bfloat16).view(num_blocks, block_size, d_rope) + + result[..., d_nope:] = input_rope + for tile_idx in range(0, num_tiles): + cur_nope = input_nope[:, tile_idx*tile_size//16:(tile_idx+1)*tile_size//16, :, :].to(torch.float32).permute(0, 2, 1, 3).contiguous().view(num_blocks, block_size, tile_size) + cur_scales = input_scale[:, :, tile_idx].unsqueeze(-1) + result[..., tile_idx*tile_size:(tile_idx+1)*tile_size] = cur_nope * cur_scales + + elif kvcache_layout == KVCacheLayout.V32_FP8Sparse: quant_k_cache = quant_k_cache.view(num_blocks, block_size, -1) - input_nope = quant_k_cache[..., :d_nope] + input_nope = quant_k_cache[..., :d_nope].view(torch.float8_e4m3fn) input_scale = quant_k_cache[..., d_nope:d_nope + num_tiles*4].view(torch.float32) input_rope = quant_k_cache[..., d_nope + num_tiles*4:].view(torch.bfloat16) result[..., d_nope:] = input_rope @@ -103,10 +222,10 @@ def dequantize_k_cache( cur_scales = input_scale[..., tile_idx].unsqueeze(-1) result[..., tile_idx*tile_size:(tile_idx+1)*tile_size] = cur_nope * cur_scales - elif kvcache_layout == FP8KVCacheLayout.MODEL1_FP8Sparse: + elif kvcache_layout == KVCacheLayout.V4_FP8Sparse: quant_k_cache = quant_k_cache.view(num_blocks, -1) # [num_blocks, ...] input_nope_rope = quant_k_cache[:, :block_size*(d_nope+2*d_rope)].view(num_blocks, block_size, d_nope + 2*d_rope) - input_nope = input_nope_rope[:, :, :d_nope] + input_nope = input_nope_rope[:, :, :d_nope].view(torch.float8_e4m3fn) input_rope = input_nope_rope[:, :, d_nope:].view(torch.bfloat16) input_scale = quant_k_cache[:, block_size*(d_nope+2*d_rope):].view(num_blocks, block_size, 8)[:, :, :7].view(torch.float8_e8m0fnu) # [num_blocks, block_size, num_tiles] @@ -115,7 +234,34 @@ def dequantize_k_cache( cur_nope = input_nope[..., tile_idx*tile_size:(tile_idx+1)*tile_size].to(torch.bfloat16) cur_scales = input_scale[:, :, tile_idx].to(torch.bfloat16).unsqueeze(-1) result[..., tile_idx*tile_size: (tile_idx+1)*tile_size] = cur_nope * cur_scales + + elif kvcache_layout == KVCacheLayout.V41_FP8Sparse: + quant_k_cache = quant_k_cache.view(num_blocks, -1) # [num_blocks, ...] + input_nope_rope = quant_k_cache[:, :block_size*(d_nope+d_rope)].view(num_blocks, block_size, d_nope + d_rope) + input_scale = quant_k_cache[:, block_size*(d_nope+d_rope):].view(num_blocks, block_size, num_tiles).view(torch.float8_e8m0fnu) # [num_blocks, block_size, 16] + + # Dequant NoPE (tiles 0-13, each tile size 32) + for tile_idx in range(0, 16): + cur_nope_rope = input_nope_rope[..., tile_idx*tile_size:(tile_idx+1)*tile_size].to(torch.bfloat16) + cur_scales = input_scale[:, :, tile_idx].to(torch.bfloat16).unsqueeze(-1) + result[..., tile_idx*tile_size:(tile_idx+1)*tile_size] = cur_nope_rope * cur_scales + elif kvcache_layout == KVCacheLayout.V41_FP4: + quant_k_cache = quant_k_cache.view(num_blocks, -1) + input_data = quant_k_cache[:, :block_size*(d//2)].view(torch.uint8).view(num_blocks, block_size, d//2) + input_scale = quant_k_cache[:, block_size*(d//2):block_size*(d//2 + num_tiles)].view(torch.float8_e4m3fn).view(num_blocks, block_size, num_tiles) + + # Chunked along the blocks: the fp32 intermediates are ~12x the size of the cache + blocks_per_chunk = max(1, (1 << 26) // (block_size * d)) + for b0 in range(0, num_blocks, blocks_per_chunk): + b1 = min(b0 + blocks_per_chunk, num_blocks) + codes = torch.empty((b1 - b0, block_size, d), dtype=torch.uint8, device=quant_k_cache.device) + codes[..., 0::2] = input_data[b0:b1] & 0xF + codes[..., 1::2] = input_data[b0:b1] >> 4 + values = _dequantize_e2m1(codes).view(b1 - b0, block_size, num_tiles, tile_size) + # e2m1 x e4m3 has at most 2 + 4 significant bits, so the product is exact in bf16, as in the kernel + result[b0:b1] = (values * input_scale[b0:b1].float().unsqueeze(-1)).view(b1 - b0, block_size, d).to(torch.bfloat16) + else: raise NotImplementedError(f"Unsupported kvcache_layout: {kvcache_layout}") diff --git a/tests/ref.py b/tests/ref.py index e5a14b3e..4f0a39e6 100644 --- a/tests/ref.py +++ b/tests/ref.py @@ -16,7 +16,7 @@ def _merge_two_lse(lse0: torch.Tensor, lse1: Optional[torch.Tensor], s_q: int, h dim=0 ) -def ref_sparse_attn_fwd(p: TestParam, t: Testcase) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: +def ref_sparse_attn_fwd(p: TestParam, t: Testcase, rms_norm_scale_factor: Optional[torch.Tensor] = None) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: """ Returns: - o: [s_q, h_q, dv] @@ -34,7 +34,10 @@ def ref_sparse_attn_fwd(p: TestParam, t: Testcase) -> Tuple[torch.Tensor, torch. q = t.q.float() gathered_kv = t.kv.index_select(dim=0, index=indices.flatten()).reshape(p.s_q, p.topk, p.d_qk).float() # [s_q, topk, d_qk] P = (q @ gathered_kv.transpose(1, 2)) # [s_q, h_q, topk] - P *= t.sm_scale + if rms_norm_scale_factor is not None: + P *= t.sm_scale * rms_norm_scale_factor.unsqueeze(-1) + else: + P *= t.sm_scale P[invalid_mask.unsqueeze(1).broadcast_to(P.shape)] = float("-inf") orig_lse = torch.logsumexp(P, dim=-1) # [s_q, h_q] @@ -54,7 +57,8 @@ def ref_sparse_attn_fwd(p: TestParam, t: Testcase) -> Tuple[torch.Tensor, torch. def ref_sparse_attn_decode( p: TestParam, - t: TestcaseForDecode + t: TestcaseForDecode, + rms_norm_scale_factor: Optional[torch.Tensor] = None # [b, s_q, h_q] ) -> Tuple[torch.Tensor, torch.Tensor]: """ A reference implementation of sparse decoding attention in PyTorch @@ -83,7 +87,10 @@ def process_kv_scope(kv_scope: KVScope) -> Tuple[torch.Tensor, torch.Tensor]: gathered_kv[gathered_kv != gathered_kv] = 0.0 q = t.q.float().view(b*p.s_q, p.h_q, p.d_qk) attn_weight = q @ gathered_kv.transpose(-1, -2) # [t.b*t.s_q, t.h_q, topk+extra_topk] - attn_weight *= t.sm_scale + if rms_norm_scale_factor is not None: + attn_weight *= t.sm_scale * rms_norm_scale_factor.view(b*p.s_q, p.h_q).unsqueeze(-1) + else: + attn_weight *= t.sm_scale attn_weight[invalid_mask.view(b*p.s_q, 1, -1).broadcast_to(b*p.s_q, p.h_q, invalid_mask.size(-1))] = float("-inf") lse = attn_weight.logsumexp(dim=-1) # [t.b*t.s_q, t.h_q] attn_weight = torch.exp(attn_weight - lse.unsqueeze(-1)) diff --git a/tests/test_flash_mla_sparse_decoding.py b/tests/test_flash_mla_sparse_decoding.py index 31aac486..f07ea3e3 100644 --- a/tests/test_flash_mla_sparse_decoding.py +++ b/tests/test_flash_mla_sparse_decoding.py @@ -12,6 +12,7 @@ import flash_mla import lib +import quant from lib import TestParam from lib import RawTestParamForDecode as RawTestParam import ref @@ -21,6 +22,8 @@ """ def gen_testcase() -> List[RawTestParam]: + # The DeepSeek-V4.1 KV cache formats (V41 / V41_FP4) are only supported on SM100f + supports_v41 = torch.cuda.get_device_capability()[0] >= 10 correctness_cases = [] corner_cases = [] for d_qk in [576, 512]: @@ -99,18 +102,53 @@ def gen_testcase() -> List[RawTestParam]: ] corner_cases.extend(cur_corner_cases) + # DeepSeek-V4.1: fp8 (V41) KV cache, optionally with an fp4 (V41_FP4) extra KV cache + if supports_v41: + for extra_fp4 in [False, True]: + correctness_cases.extend([ + RawTestParam(b, h_q, s_q, 1, s_k, is_varlen, topk, + have_topk_length=have_topk_len, + enable_attn_sink=True, + extra_s_k=extra_s_k, + extra_topk=extra_topk, + block_size=block_size, + extra_block_size=extra_block_size, + have_extra_topk_length=have_extra_topk_len, + d_qk=512, + kvcache_layout=quant.KVCacheLayout.V41_FP8Sparse, + extra_kvcache_layout=quant.KVCacheLayout.V41_FP4 if extra_fp4 else None, + check_correctness=True, + num_runs=0) + for h_q in [64, 128] + for have_extra_topk_len in [False, True] + for have_topk_len in [False] + for (s_k, topk, block_size) in [(512, 64, 64), (1024, 576, 61)] + for (extra_s_k, extra_topk, extra_block_size) in [(512, 64, 64), (650, 576, 53)] + for b in [4] + for s_q in [1, 3] + for is_varlen in [True] + ]) + base_and_bszs = [ # V3.2 (RawTestParam(0, 128, 2, 1, 32768, True, topk=2048, d_qk=576), [2, 64, 74, 128]), - # MODEL1 CONFIG1 + # DeepSeek-V4 CONFIG1 (RawTestParam(0, 64, 2, 1, 16384, True, topk=128, d_qk=512, extra_s_k=16384, extra_topk=512, block_size=256, extra_block_size=64), [2, 64, 74, 128, 74*2, 256]), - # MODEL1 CONFIG2 + # DeepSeek-V4 CONFIG2 (RawTestParam(0, 128, 2, 1, 16384, True, topk=128, d_qk=512, extra_s_k=16384, extra_topk=1024, block_size=256, extra_block_size=64), [2, 64, 74, 128, 74*2, 256]), - # MODEL1 CONFIG3 + # DeepSeek-V4 CONFIG3 (RawTestParam(0, 64, 2, 1, 16384, True, topk=128, d_qk=512, extra_s_k=16384, extra_topk=1024, block_size=256, extra_block_size=2, have_extra_topk_length=True), [2, 64, 74, 128, 74*2, 256]), - # MODEL1 CONFIG4 + # DeepSeek-V4 CONFIG4 (RawTestParam(0, 128, 2, 1, 16384, True, topk=128, d_qk=512, extra_s_k=16384, extra_topk=1024, block_size=256, extra_block_size=2, have_extra_topk_length=True), [2, 64, 74, 128, 74*2, 256]), ] + if supports_v41: + base_and_bszs += [ + # DeepSeek-V4.1 CONFIG1 (fp8 V41 KV cache + fp4 V41_FP4 extra KV cache) + (RawTestParam(0, 64, 2, 1, 16384, True, topk=128, d_qk=512, extra_s_k=16384, extra_topk=512, block_size=256, extra_block_size=64, + kvcache_layout=quant.KVCacheLayout.V41_FP8Sparse, extra_kvcache_layout=quant.KVCacheLayout.V41_FP4), [2, 64, 74, 128, 74*2, 256]), + (RawTestParam(0, 128, 2, 1, 16384, True, topk=128, d_qk=512, extra_s_k=16384, extra_topk=512, block_size=256, extra_block_size=64, + kvcache_layout=quant.KVCacheLayout.V41_FP8Sparse, extra_kvcache_layout=quant.KVCacheLayout.V41_FP4), [2, 64, 74, 128, 74*2, 256]) + ] performance_cases = [ # Production cases dataclasses.replace(base, b=b) diff --git a/tests/test_flash_mla_sparse_prefill.py b/tests/test_flash_mla_sparse_prefill.py index f4d39dd7..f9dea0f6 100644 --- a/tests/test_flash_mla_sparse_prefill.py +++ b/tests/test_flash_mla_sparse_prefill.py @@ -24,7 +24,7 @@ def run_test(p: TestParam) -> bool: torch.cuda.synchronize() def run_prefill(): - return lib.run_flash_mla_sparse_fwd(p, t, False) + return lib.run_flash_mla_sparse_fwd(p, t) prefill_ans_out, prefill_ans_max_logits, prefill_ans_lse = run_prefill() torch.cuda.synchronize() @@ -105,7 +105,6 @@ def run_prefill(): (114, 384), ] for s_q in [62, 213] - for have_sink_lse in [False, True] for have_attn_sink in [False, True] for have_topk_length in [False, True] ] @@ -146,9 +145,9 @@ def run_prefill(): performance_case_templates = [ # V3.2 (576, 128, 2048, [8192, 32768, 65536, 98304, 131072]), - # MODEL1 CONFIG1 + # DeepSeek-V4 CONFIG1 (512, 64, 512, [8192, 32768, 49152, 65536]), - # MODEL1 CONFIG2 + # DeepSeek-V4 CONFIG2 (512, 128, 1024, [8192, 32768, 49152, 65536]), ] diff --git a/tests/test_fused_norm_rope_attn_rope_cast.py b/tests/test_fused_norm_rope_attn_rope_cast.py new file mode 100644 index 00000000..99769567 --- /dev/null +++ b/tests/test_fused_norm_rope_attn_rope_cast.py @@ -0,0 +1,542 @@ +import math +import time +import sys +from typing import Tuple, Optional +import random +import functools + +import argparse +import torch +import tile_kernels +import deep_gemm +import kernelkit as kk + +from lib import TestParam, Testcase, RawTestParamForDecode, TestcaseForDecode, ExtraTestParamForDecode +import lib +import quant +import ref + +import flash_mla + +# FMA with double precision is necessary for numerical consistency of fused ref +_GET_FMA_KERNEL = None +def _fma(a: torch.Tensor, b: torch.Tensor, c: torch.Tensor) -> torch.Tensor: + global _GET_FMA_KERNEL + assert a.numel() == b.numel() == c.numel() + if _GET_FMA_KERNEL is None: + import tilelang + from tilelang import language as T + @tilelang.jit( + pass_configs={ + tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True, + tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True, + }, + ) + def _get_fma_kernel(): + num_threads = 256 + numel_per_cta = 1024 + numel = T.symbolic("numel") + @T.prim_func + def fma_kernel( + out: T.Tensor[(numel,), torch.float32], + a: T.Tensor[(numel,), torch.float32], + b: T.Tensor[(numel,), torch.float32], + c: T.Tensor[(numel,), torch.float32] + ): + with T.Kernel(numel // numel_per_cta, threads=num_threads) as (tid_x, ): + a_fragment = T.alloc_fragment((numel_per_cta, ), torch.float32) + b_fragment = T.alloc_fragment((numel_per_cta, ), torch.float32) + c_fragment = T.alloc_fragment((numel_per_cta, ), torch.float32) + out_fragment = T.alloc_fragment((numel_per_cta, ), torch.float32) + + T.copy(a[tid_x * numel_per_cta], a_fragment) + T.copy(b[tid_x * numel_per_cta], b_fragment) + T.copy(c[tid_x * numel_per_cta], c_fragment) + for i in T.Parallel(numel_per_cta): + out_fragment[i] = a_fragment[i] * b_fragment[i] + c_fragment[i] + T.copy(out_fragment, out[tid_x * numel_per_cta]) + return fma_kernel + _GET_FMA_KERNEL = _get_fma_kernel + out = torch.empty_like(a).flatten() + _GET_FMA_KERNEL()(out, a.flatten(), b.flatten(), c.flatten()) + return out.view_as(a) + + +def _rope_inplace(q_or_o: torch.Tensor, conjugate: bool, cos_sin_table: torch.Tensor, token_positions: torch.Tensor, rope_dim: int = 64): + # cos_sin_table: [*, rope_dim]; token_positions: [s_q] (prefill) or [b, s_q] (decode) + # q_or_o: [s_q, h_q, d_qk] (prefill) or [b, s_q, h_q, d_qk] (decode) + # The unsqueeze dim depends on q_or_o's dimensionality: + # 3D => dim=1, 4D => dim=2, i.e., unsqueeze_dim = q_or_o.ndim - 2 + unsqueeze_dim = q_or_o.ndim - 2 + cos = cos_sin_table[token_positions, :rope_dim//2].unsqueeze(unsqueeze_dim) + sin = cos_sin_table[token_positions, rope_dim//2:].unsqueeze(unsqueeze_dim) + if conjugate: + sin = -sin + rope_part = q_or_o[..., -rope_dim:] + rope_part_x0 = rope_part[..., 0::2] + rope_part_x1 = rope_part[..., 1::2] + new_rope_part_x0 = _fma(rope_part_x0, cos.expand_as(rope_part_x0), -rope_part_x1 * sin) + new_rope_part_x1 = _fma(rope_part_x0, sin.expand_as(rope_part_x0), rope_part_x1 * cos) + rope_part[..., 0::2] = new_rope_part_x0 + rope_part[..., 1::2] = new_rope_part_x1 + +def ref_fused_norm_rope_attn_rope_cast_fwd(p: TestParam, t: Testcase, enable_q_norm: bool, cos_sin_table: torch.Tensor, token_positions: torch.Tensor): + # Q norm + if enable_q_norm: + q_float32 = t.q.float() + rms_norm_scale_factor = torch.rsqrt(torch.sum(q_float32*q_float32, dim=-1) / p.d_qk + rms_norm_eps) + else: + rms_norm_scale_factor = None + + # Q RoPE + q = t.q.clone().float() + _rope_inplace(q, False, cos_sin_table, token_positions) + q = q.to(torch.bfloat16) + + # Core attention + old_q = t.q + t.q = q + _, out, max_logits, lse = ref.ref_sparse_attn_fwd(p, t, rms_norm_scale_factor) + t.q = old_q + + # O RoPE + _rope_inplace(out, True, cos_sin_table, token_positions) + + return out, max_logits, lse + +def ref_fused_norm_rope_attn_rope_cast_decode( + p: TestParam, + t: TestcaseForDecode, + enable_q_norm: bool, + cos_sin_table: torch.Tensor, + token_positions: torch.Tensor # [b*s_q] +): + """ + Reference implementation for the fused Q norm + RoPE + Core Attn (decode) + O RoPE + Returns: (out_bf16, lse) + """ + b, s_q, h_q, d_qk = t.q.shape + + # Q norm + if enable_q_norm: + q_float32 = t.q.float() + rms_norm_scale_factor = torch.rsqrt(torch.sum(q_float32*q_float32, dim=-1) / p.d_qk + rms_norm_eps) + else: + rms_norm_scale_factor = None + + # Q RoPE + q = t.q.clone().reshape(b*s_q, h_q, d_qk).float() + _rope_inplace(q, False, cos_sin_table, token_positions) + q = q.to(torch.bfloat16).reshape(b, s_q, h_q, d_qk) + + # Core attention (decode) + old_q = t.q + t.q = q + out, lse = ref.ref_sparse_attn_decode(p, t, rms_norm_scale_factor) + t.q = old_q + + # O RoPE + out = out.float().reshape(b*s_q, h_q, p.d_v) + _rope_inplace(out, True, cos_sin_table, token_positions) + + return out.to(torch.bfloat16), lse.transpose(1, 2).reshape(b*s_q, h_q) + +def build_cos_sin_cache(max_token_position): + # Copied from vLLM + base = 100000 + rope_dim = 64 + mscale = 1.0 + inv_freq = 1.0 / (base ** (torch.arange(0, rope_dim, 2, dtype=torch.float, device="cuda") / rope_dim)) + freqs = torch.outer( + torch.arange(0, max_token_position, dtype=torch.float, device="cuda"), + inv_freq + ) + freqs_cis = torch.polar(torch.ones_like(freqs), freqs) + cos = freqs_cis.real.cuda() * mscale + sin = freqs_cis.imag.cuda() * mscale + return torch.cat((cos, sin), dim=-1) +max_token_position = 131072 +cos_sin_cache = build_cos_sin_cache(max_token_position) +rms_norm_eps = 1e-4 + +@functools.lru_cache(maxsize=None) # To avoid generating q weight and permuted q weight for multiple times, we cache previously generated q_weight and use them later +def get_q_b_weight(h_q: int, d_q: int, q_lora_rank: int, scale_gran: int): + kk.utils.set_random_seed(0) + q_weight = torch.randn((h_q*d_q, q_lora_rank), dtype=torch.bfloat16, device='cuda') / 10 + q_weight_quanted = tile_kernels.quant.per_token_cast(q_weight, 'e4m3', scale_gran, use_tma_aligned_col_major_sf=True, round_sf=True, use_packed_ue8m0=True) + q_weight_quanted_permuted = flash_mla.fused_norm_rope_attn_rope_cast.permute_q_b_proj(q_weight_quanted, h_q, d_q) + return q_weight_quanted, q_weight_quanted_permuted + +@functools.lru_cache(maxsize=None) +def get_wb_weight(n_wv_group: int, wv_group_size: int, d_o: int, scale_gran: int, wv_proj_out_dim: int): + kk.utils.set_random_seed(0) + o_weight = torch.randn((n_wv_group*wv_proj_out_dim, wv_group_size * d_o), dtype=torch.bfloat16, device='cuda') / 10 + o_weight_quanted = tile_kernels.quant.per_token_cast(o_weight, 'e4m3', scale_gran, use_tma_aligned_col_major_sf=False, round_sf=True, use_packed_ue8m0=False) + o_weight_quanted_sf = deep_gemm.transform_sf_into_required_layout( + o_weight_quanted[1].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, + ) + o_weight_quanted = (o_weight_quanted[0].view(n_wv_group, wv_proj_out_dim, wv_group_size * d_o), o_weight_quanted_sf) + o_weight_quanted_permuted = flash_mla.fused_norm_rope_attn_rope_cast.permute_wv_proj(o_weight_quanted, wv_group_size, d_o) + return o_weight_quanted, o_weight_quanted_permuted + +_counter = kk.Counter() + +@torch.inference_mode() +def run_test(p: TestParam) -> bool: + if p.seed == -1: + global _counter + p.seed = _counter.next() + + print("================") + print(f"Running on {p}") + + t = lib.generate_testcase(p) + torch.cuda.synchronize() + + q_lora_rank = random.choice([1024, 1536]) + q_b_proj_scale_gran = random.choice([32, 128]) + q_lora = torch.randn((p.s_q, q_lora_rank)) / 10 + q_lora = tile_kernels.quant.per_token_cast(q_lora, 'e4m3', q_b_proj_scale_gran, None, use_tma_aligned_col_major_sf=True, round_sf=True, use_packed_ue8m0=True) + q_b_proj, q_b_proj_permuted = get_q_b_weight(p.h_q, p.d_qk, q_lora_rank, q_b_proj_scale_gran) + # TODO Make q weight non-contiguous + + o_scale_gran = 32 + enable_q_norm = random.choice([False, True]) + token_positions = torch.randint(0, max_token_position, (p.s_q, ), device='cuda', dtype=torch.int32) + + wv_group_size = 8 + n_wv_group = p.h_q // 8 + wv_proj_out_dim = random.choice([256, 512, 1024]) + wv_proj_scale_gran = 32 + wv_proj, wv_proj_permuted = get_wb_weight(n_wv_group, wv_group_size, p.d_v, wv_proj_scale_gran, wv_proj_out_dim) + + def calculate_q_b_proj(weight) -> torch.Tensor: + q = torch.empty((p.s_q, p.h_q*p.d_qk), dtype=torch.bfloat16, device='cuda') + deep_gemm.fp8_gemm_nt(q_lora, weight, q, recipe_a=(1, q_b_proj_scale_gran), recipe_b=(1, q_b_proj_scale_gran)) + return q.view(p.s_q, p.h_q, p.d_qk) + q_for_fused = calculate_q_b_proj(q_b_proj_permuted) + + def calculate_wv_proj(output, weight) -> torch.Tensor: + wv_proj_out = torch.empty((p.s_q, n_wv_group, wv_proj_out_dim), dtype=torch.bfloat16, device='cuda') + deep_gemm.fp8_einsum( + "bhr,hdr->bhd", + output, + weight, + wv_proj_out, + recipe=(1, 1, 32), + ) + return wv_proj_out + + def run_fused_norm_rope_attn_rope_cast_fwd(): + return flash_mla.fused_norm_rope_attn_rope_cast.prefill( + enable_q_norm, + rms_norm_eps, + token_positions, + False, 64, cos_sin_cache, + n_wv_group, o_scale_gran, True, True, True, + q_for_fused, t.kv, t.indices, + sm_scale=t.sm_scale, + attn_sink=t.attn_sink, + topk_length=t.topk_length, + ) + + torch.cuda.synchronize() + ans_out_fp8, ans_out_sf, ans_max_logits, ans_lse = run_fused_norm_rope_attn_rope_cast_fwd() + torch.cuda.synchronize() + ans_wv_proj_out = calculate_wv_proj((ans_out_fp8, ans_out_sf), wv_proj_permuted) + + if p.num_runs > 0: + flops_and_mem_vol = lib.count_flop_and_mem_vol(p, t) + fused_time = kk.bench_kineto(run_fused_norm_rope_attn_rope_cast_fwd, num_tests=p.num_runs).get_kernel_time("fused_norm_rope_attn_rope_cast_fwd") + fused_flops = flops_and_mem_vol.fwd_flop/fused_time/1e12 + fused_mem_bw = flops_and_mem_vol.fwd_prefill_with_fp8_out_mem_vol/fused_time/1e12 + print(f"Fused: {fused_time*1e6:4.0f} us, {fused_flops:6.1f} TFlops, {fused_mem_bw:4.2f} TBps") + + if p.check_correctness: + out_criteria = {'abs_tol': 1.1e-3, 'rel_tol': 1.01/8, 'cos_diff_tol': 1e-3} if p.k_amplifier_portion == 0.0 else {'abs_tol': 1.0, 'rel_tol': 1.0, 'cos_diff_tol': 1e-3} + wv_proj_out_criteria = {'abs_tol': 1.0, 'rel_tol': 1.0, 'cos_diff_tol': 1e-3} if p.k_amplifier_portion == 0.0 else {'abs_tol': 100.0, 'rel_tol': 100.0, 'cos_diff_tol': 1e-3} + max_logits_criteria = {'abs_tol': 1e-5, 'rel_tol': 4.01/65536} if p.k_amplifier_portion == 0.0 else {'abs_tol': 1.0, 'rel_tol': 1.0} + lse_criteria = {'abs_tol': 1e-5, 'rel_tol': 4.01/65536} if p.k_amplifier_portion == 0 else {'abs_tol': 1e-4, 'rel_tol': 8.01/65536} + t.q = calculate_q_b_proj(q_b_proj) + ref_out, ref_max_logits, ref_lse = ref_fused_norm_rope_attn_rope_cast_fwd(p, t, enable_q_norm, cos_sin_cache, token_positions) + ref_lse[ref_lse == float("-inf")] = float("+inf") + + ans_out_sf_2d = ans_out_sf.view(p.s_q, -1) + if ans_out_sf_2d.stride(0) != 1: + # PyTorch normalizes strides of size-1 dims during view(), which breaks + # tile_kernels' col-major sf layout detection (`sf.stride(0) == 1`) when s_q == 1, + # so we restore the token-dim stride to 1 manually here + assert p.s_q == 1 + ans_out_sf_2d = ans_out_sf_2d.as_strided(ans_out_sf_2d.shape, (1, ans_out_sf_2d.stride(1))) + ans_wv_out_dequantized = tile_kernels.quant.per_token_cast_back((ans_out_fp8.view(p.s_q, -1), ans_out_sf_2d), 'fp32', 32) + ans_wv_out_dequantized = ans_wv_out_dequantized \ + .view(p.s_q, n_wv_group, p.d_v // o_scale_gran, wv_group_size, o_scale_gran) \ + .transpose(2, 3) \ + .reshape(p.s_q, p.h_q, p.d_v) + + assert o_scale_gran == 32 + ref_out_fp8, ref_out_sf = tile_kernels.quant.per_token_cast(ref_out.view(p.s_q, -1), 'e4m3', 32, round_sf=True, use_tma_aligned_col_major_sf=True, use_packed_ue8m0=True) + ref_wv_proj_out = calculate_wv_proj( + (ref_out_fp8.view(p.s_q, n_wv_group, wv_group_size * p.d_v), ref_out_sf.view(p.s_q, n_wv_group, wv_group_size * p.d_v // (32*4))), + wv_proj + ) + + is_correct = True + is_correct &= kk.check_is_allclose("out", ans_wv_out_dequantized, ref_out, **out_criteria) + is_correct &= kk.check_is_allclose("wv_proj_out", ans_wv_proj_out, ref_wv_proj_out, **wv_proj_out_criteria) + is_correct &= kk.check_is_allclose("max_logits", ans_max_logits, ref_max_logits, **max_logits_criteria) + is_correct &= kk.check_is_allclose("lse", ans_lse, ref_lse, **lse_criteria) + + return is_correct + else: + return True + + +@torch.inference_mode() +def run_decode_test(p: TestParam) -> bool: + assert p.decode is not None + if p.seed == -1: + global _counter + p.seed = _counter.next() + + print("================") + print(f"Running on {p}") + + t = lib.generate_testcase_for_decode(p) + b = p.decode.b + s_q = p.s_q + + q_lora_rank = random.choice([1024, 1536]) + q_b_proj_scale_gran = random.choice([32, 128]) + q_lora = torch.randn((b*s_q, q_lora_rank), device='cuda') / 10 + q_lora = tile_kernels.quant.per_token_cast(q_lora, 'e4m3', q_b_proj_scale_gran, None, use_tma_aligned_col_major_sf=True, round_sf=True, use_packed_ue8m0=True) + q_b_proj, q_b_proj_permuted = get_q_b_weight(p.h_q, p.d_qk, q_lora_rank, q_b_proj_scale_gran) + + o_scale_gran = 32 + enable_q_norm = random.choice([False, True]) + wv_group_size = 8 + n_wv_group = p.h_q // 8 + + def calculate_q_b_proj(weight) -> torch.Tensor: + q = torch.empty((b*s_q, p.h_q*p.d_qk), dtype=torch.bfloat16, device='cuda') + deep_gemm.fp8_gemm_nt(q_lora, weight, q, recipe_a=(1, q_b_proj_scale_gran), recipe_b=(1, q_b_proj_scale_gran)) + return q.view(b*s_q, p.h_q, p.d_qk) + q_for_fused = calculate_q_b_proj(q_b_proj_permuted) + token_positions = torch.randint(0, max_token_position, (b*s_q, ), device='cuda', dtype=torch.int32) + + def run_fused_norm_rope_attn_rope_cast_decode(): + # NOTE The kernel interface does not have the batch dimension (batch size is always 1), + # so we squeeze the batch dimension out of every per-request tensor here + return flash_mla.fused_norm_rope_attn_rope_cast.decode( + enable_q_norm=enable_q_norm, + rms_norm_eps=rms_norm_eps, + token_positions=token_positions, + is_rope_neox_style=False, + rope_dim=64, + cos_sin_cache=cos_sin_cache, + n_wv_group=n_wv_group, + num_per_channels=o_scale_gran, + use_tma_aligned_col_major_sf=True, + round_sf=True, + use_packed_ue8m0=True, + q=q_for_fused, + k_cache=t.kv_scope.get_kvcache_for_flash_mla(), + indices_in_kvcache=t.kv_scope.indices_in_kvcache.reshape(b*s_q, p.topk), + sm_scale=t.sm_scale, + d_v=p.d_v, + attn_sink=t.attn_sink, + topk_length=t.kv_scope.topk_length.repeat_interleave(s_q, 0) if t.kv_scope.topk_length is not None else None, + extra_k_cache=t.extra_kv_scope.get_kvcache_for_flash_mla() if t.extra_kv_scope is not None else None, + extra_indices_in_kvcache=t.extra_kv_scope.indices_in_kvcache.reshape(b*s_q, p.decode.extra_topk) if t.extra_kv_scope is not None else None, + extra_topk_length=t.extra_kv_scope.topk_length.repeat_interleave(s_q, 0) if t.extra_kv_scope is not None and t.extra_kv_scope.topk_length is not None else None, + ) + + torch.cuda.synchronize() + ans_out_fp8, ans_out_sf, ans_lse = run_fused_norm_rope_attn_rope_cast_decode() + torch.cuda.synchronize() + + if p.num_runs > 0: + flops_and_mem_vol = lib.count_flop_and_mem_vol_for_decode(p, t) + fused_time = kk.bench_kineto(run_fused_norm_rope_attn_rope_cast_decode, num_tests=p.num_runs).get_kernel_time("fused_norm_rope_attn_rope_cast_fwd") + fused_flops = flops_and_mem_vol.flop/fused_time/1e12 + fused_mem_bw = flops_and_mem_vol.mem_vol/fused_time/1e12 + print(f"Fused decode: {fused_time*1e6:4.0f} us, {fused_flops:6.1f} TFlops, {fused_mem_bw:4.2f} TBps") + + is_correct = True + if p.check_correctness: + out_criteria = {'abs_tol': 1.1e-3, 'rel_tol': 1.01/8, 'cos_diff_tol': 1e-3} if p.k_amplifier_portion == 0.0 else {'abs_tol': 1.0, 'rel_tol': 1.0, 'cos_diff_tol': 1e-3} + lse_criteria = {'abs_tol': 1e-5, 'rel_tol': 4.01/65536} if p.k_amplifier_portion == 0 else {'abs_tol': 1e-5, 'rel_tol': 4.01/65536} + + t.q = calculate_q_b_proj(q_b_proj).view(b, s_q, p.h_q, p.d_qk) + ref_out, ref_lse = ref_fused_norm_rope_attn_rope_cast_decode(p, t, enable_q_norm, cos_sin_cache, token_positions) + + # NOTE Tensors returned by the fused decode kernel do not have the batch dimension (batch size is always 1) + ans_out_sf_2d = ans_out_sf.view(b*s_q, -1) + if ans_out_sf_2d.stride(0) != 1: + assert s_q == 1 + ans_out_sf_2d = ans_out_sf_2d.as_strided(ans_out_sf_2d.shape, (1, ans_out_sf_2d.stride(1))) + ans_wv_out_dequantized = tile_kernels.quant.per_token_cast_back( + (ans_out_fp8.view(b*s_q, -1), ans_out_sf_2d), 'fp32', 32 + ) + ans_wv_out_dequantized = ans_wv_out_dequantized \ + .view(b*s_q, n_wv_group, p.d_v // 32, wv_group_size, 32) \ + .transpose(2, 3) \ + .reshape(b*s_q, p.h_q, p.d_v) + + is_correct &= kk.check_is_allclose("out", ans_wv_out_dequantized, ref_out.to(torch.float32), **out_criteria) + is_correct &= kk.check_is_allclose("lse", ans_lse, ref_lse, **lse_criteria) + + + return is_correct + + +if __name__ == '__main__': + device = torch.device("cuda:0") + torch.set_default_dtype(torch.bfloat16) + torch.set_default_device(device) + torch.cuda.set_device(device) + torch.set_float32_matmul_precision('high') + + parser = argparse.ArgumentParser() + lib.stick_unit_test_args(parser) + args = parser.parse_args() + + correctness_cases_prefill = [] + correctness_cases_decode = [] + + for h_q in [128, 64]: + for s_kv, topk in [ + # Regular shapes + (64, 64), + (128, 128), + (256, 256), + (512, 512), + + # Irregular shapes + (592, 120), + (1840, 240), + (1521, 600), + (3412, 2896), + + # Irregular shapes with OOB TopK + (95, 152), + (153, 264), + (2345, 5136), + + (32, 2048), # Some block may be fully invalid + ]: + for (have_attn_sink, is_all_indices_invalid, have_topk_length) in [ + (False, False, False), + (True, False, False), + (True, True, False), + (True, False, True), + (False, True, True) + ]: + for s_q in [1, 184, 2123]: + # Prefill + correctness_cases_prefill.extend([TestParam( + s_q, s_kv, topk, h_q, + is_all_indices_invalid=is_all_indices_invalid, + have_topk_length=have_topk_length, + k_amplifier_portion=k_amplifier_portion, + k_amplifier_ratio=k_amplifier_ratio, + num_runs=0 + ) + for (k_amplifier_portion, k_amplifier_ratio) in [ + (0.0, 1.0), + (0.02, 2**8) + ] + ]) + for (b, s_q) in [(1, 1), (69, 3), (19, 80)]: + for extra_s_kv, extra_topk, kvcache_layout, extra_kvcache_layout in [ + (None, None, quant.KVCacheLayout.V4_FP8Sparse, None), + (None, None, quant.KVCacheLayout.V41_FP8Sparse, None), + (64, 64, quant.KVCacheLayout.V4_FP8Sparse, None), + (32, 2048, quant.KVCacheLayout.V4_FP8Sparse, None), + (32, 2048, quant.KVCacheLayout.V41_FP8Sparse, None), + (32, 2048, quant.KVCacheLayout.V41_FP8Sparse, quant.KVCacheLayout.V41_FP4), + (592, 120, quant.KVCacheLayout.V4_FP8Sparse, None), + (512, 512, quant.KVCacheLayout.V41_FP8Sparse, None), + (512, 512, quant.KVCacheLayout.V41_FP8Sparse, quant.KVCacheLayout.V41_FP4), + (334, 432, quant.KVCacheLayout.V4_FP8Sparse, None), + (3412, 2896, quant.KVCacheLayout.V41_FP8Sparse, None), + (3412, 2896, quant.KVCacheLayout.V41_FP8Sparse, quant.KVCacheLayout.V41_FP4), + ]: + # The fp4 extra KV cache requires topk % 8 == 0 + if extra_kvcache_layout == quant.KVCacheLayout.V41_FP4 and topk % 8 != 0: + continue + for have_extra_topk_length in [False, True]: + if have_extra_topk_length and extra_s_kv is None: + continue + correctness_cases_decode.append(RawTestParamForDecode( + b, h_q, s_q, 1, s_kv, False, topk, + is_all_indices_invalid=is_all_indices_invalid, + have_zero_seqlen_k=False, + have_topk_length=have_topk_length, + enable_attn_sink=have_attn_sink, + extra_s_k=extra_s_kv, + extra_topk=extra_topk, + block_size=(64+topk%8), + extra_block_size=(128+(extra_topk%8) if extra_topk is not None else None), + have_extra_topk_length=have_extra_topk_length, + d_qk=512, + kvcache_layout=kvcache_layout, + extra_kvcache_layout=extra_kvcache_layout, + num_runs=0) + ) + + performance_case_templates = [ + # V4 small + # (512, 32, 512+128, [8192, 32768]), + # V4 / V4.1 + (512, 64, 512, 128, [8192, 32768], [(quant.KVCacheLayout.V4_FP8Sparse, None), (quant.KVCacheLayout.V41_FP8Sparse, None), (quant.KVCacheLayout.V41_FP8Sparse, quant.KVCacheLayout.V41_FP4)], [(140, 4), (320, 3)]), + # V4 / V4.1 teachar + (512, 128, 1024, 128, [8192, 32768], [(quant.KVCacheLayout.V4_FP8Sparse, None), (quant.KVCacheLayout.V41_FP8Sparse, None), (quant.KVCacheLayout.V41_FP8Sparse, quant.KVCacheLayout.V41_FP4)], [(80, 4)]), + ] + + performance_cases_prefill = [] + performance_cases_decode = [] + for (d_qk, h_q, extra_topk, topk, s_kv_list, kvcache_layouts, decoding_bsz_and_s_q) in performance_case_templates: + for s_kv in s_kv_list: + for prefill_s_q in [4096]: + performance_cases_prefill.append(TestParam(prefill_s_q, s_kv, topk+extra_topk, h_q, d_qk=d_qk, have_attn_sink=True, check_correctness=True)) + for (decoding_bsz, decoding_s_q) in decoding_bsz_and_s_q: + for kvcache_layout, extra_kvcache_layout in kvcache_layouts: + performance_cases_decode.append(RawTestParamForDecode( + decoding_bsz, h_q, decoding_s_q, 1, topk, False, topk, have_topk_length=False, + extra_s_k=s_kv, extra_topk=extra_topk, extra_block_size=128, + d_qk=d_qk, kvcache_layout=kvcache_layout, extra_kvcache_layout=extra_kvcache_layout)) + + # Prefill testcases + testcases = correctness_cases_prefill + correctness_cases_decode + performance_cases_prefill + performance_cases_decode + testcases = [ + (t.to_test_param() if isinstance(t, RawTestParamForDecode) else t) + for t in testcases + ] + + print(f"{kk.colors['CYAN_BG']}{len(testcases)} testcases to run{kk.colors['CLEAR']}") + + failed_cases = [] + for test_idx, test in enumerate(testcases): + if test != testcases[0] and test.num_runs > 0 and not args.no_cooldown: + time.sleep(0.3) + print(f'[{test_idx:5d}/{len(testcases):5d} ({test_idx/len(testcases)*100:5.1f}%)] ', end='') + is_correct = run_test(test) if test.decode is None else run_decode_test(test) + if not is_correct: + failed_cases.append(test) + if not args.run_to_finish: + sys.exit(1) + + total = len(testcases) + if len(failed_cases) > 0: + print(f"\033[31m\033[1m{len(failed_cases)} / {total} cases failed:\033[0m") + for case in failed_cases: + print(f" {case}") + sys.exit(1) + else: + print(f"\033[32m\033[1mAll {total} cases passed!\033[0m")