diff --git a/sgl-kernel/cmake/flashmla.cmake b/sgl-kernel/cmake/flashmla.cmake index b83fe0d0b..f1dde152e 100644 --- a/sgl-kernel/cmake/flashmla.cmake +++ b/sgl-kernel/cmake/flashmla.cmake @@ -4,7 +4,7 @@ include(FetchContent) FetchContent_Declare( repo-flashmla GIT_REPOSITORY https://github.com/sgl-project/FlashMLA - GIT_TAG bc8576abc3e507425cf6498f3d3393df7733ce37 + GIT_TAG 6fde8086398bc5ed144c1e9ed0de30cd43be2cda GIT_SHALLOW OFF ) FetchContent_Populate(repo-flashmla) @@ -44,6 +44,10 @@ set(FlashMLA_SOURCES ${repo-flashmla_SOURCE_DIR}/csrc/sm100/prefill/dense/fmha_cutlass_fwd_sm100.cu ${repo-flashmla_SOURCE_DIR}/csrc/sm100/prefill/dense/fmha_cutlass_bwd_sm100.cu ${repo-flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd.cu + + ${repo-flashmla_SOURCE_DIR}/csrc/extension/sm90/dense_fp8/dense_fp8_python_api.cpp + ${repo-flashmla_SOURCE_DIR}/csrc/extension/sm90/dense_fp8/flash_fwd_mla_fp8_sm90.cu + ${repo-flashmla_SOURCE_DIR}/csrc/extension/sm90/dense_fp8/flash_fwd_mla_metadata.cu ) Python_add_library(flashmla_ops MODULE USE_SABI ${SKBUILD_SABI_VERSION} WITH_SOABI ${FlashMLA_SOURCES}) @@ -51,6 +55,7 @@ target_compile_options(flashmla_ops PRIVATE $<$:${FLASHML target_include_directories(flashmla_ops PRIVATE ${repo-flashmla_SOURCE_DIR}/csrc ${repo-flashmla_SOURCE_DIR}/csrc/sm90 + ${repo-flashmla_SOURCE_DIR}/csrc/extension/sm90/dense_fp8/ ${repo-flashmla_SOURCE_DIR}/csrc/cutlass/include ${repo-flashmla_SOURCE_DIR}/csrc/cutlass/tools/util/include ) diff --git a/sgl-kernel/csrc/flashmla_extension.cc b/sgl-kernel/csrc/flashmla_extension.cc index e72cfa11c..8236c8867 100644 --- a/sgl-kernel/csrc/flashmla_extension.cc +++ b/sgl-kernel/csrc/flashmla_extension.cc @@ -27,6 +27,9 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) { "is_fp8_kvcache, int? topk) -> Tensor[]"); m.impl("get_mla_decoding_metadata", torch::kCUDA, &get_mla_decoding_metadata); + m.def("get_mla_decoding_metadata_dense_fp8(Tensor seqlens_k, int num_heads_per_head_k, int num_heads_k) -> Tensor[]"); + m.impl("get_mla_decoding_metadata_dense_fp8", torch::kCUDA, &get_mla_decoding_metadata_dense_fp8); + m.def( "fwd_kvcache_mla(Tensor q, Tensor kv_cache, int head_size_v, Tensor seqlens_k, Tensor block_table, float " "softmax_scale, bool is_causal, Tensor tile_scheduler_metadata, Tensor num_splits, bool is_fp8, Tensor? indices) " @@ -41,6 +44,12 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) { m.def("sparse_prefill_fwd(Tensor q, Tensor kv, Tensor indices, float sm_scale, int d_v) -> Tensor[]"); m.impl("sparse_prefill_fwd", torch::kCUDA, &sparse_prefill_fwd); + + m.def( + "fwd_kvcache_mla_fp8(Tensor q, Tensor kcache, int head_size_v, Tensor seqlens_k, Tensor block_table, float " + "softmax_scale, bool is_causal, Tensor tile_scheduler_metadata, Tensor num_splits, Tensor? descale_q, Tensor? " + "descale_k) -> Tensor[]"); + m.impl("fwd_kvcache_mla_fp8", torch::kCUDA, &fwd_kvcache_mla_fp8); } REGISTER_EXTENSION(flashmla_ops) diff --git a/sgl-kernel/include/sgl_kernel_ops.h b/sgl-kernel/include/sgl_kernel_ops.h index e95bcc2ff..a3240e783 100644 --- a/sgl-kernel/include/sgl_kernel_ops.h +++ b/sgl-kernel/include/sgl_kernel_ops.h @@ -935,3 +935,21 @@ void FMHACutlassSM100FwdRun( std::vector sparse_prefill_fwd(const at::Tensor& q, const at::Tensor& kv, const at::Tensor& indices, double sm_scale, int64_t d_v); + +std::vector fwd_kvcache_mla_fp8( + at::Tensor& q, // batch_size x seqlen_q x num_heads x head_size + const at::Tensor& kcache, // num_blocks x page_block_size x num_heads_k x head_size (when is_fp8 is False) or + // num_blocks x num_heads_k x (page_block_size*656) (when is_fp8 is True) + const int64_t head_size_v, + const at::Tensor& seqlens_k, // batch_size + const at::Tensor& block_table, // batch_size x max_num_blocks_per_seq + const double softmax_scale, + bool is_causal, + const at::Tensor& tile_scheduler_metadata, // num_sm_parts x TileSchedulerMetaDataSize + const at::Tensor& num_splits, // batch_size + 1 + const std::optional& descale_q, // None or batch_size + const std::optional& descale_k // None or batch_size +); + +std::vector get_mla_decoding_metadata_dense_fp8( + at::Tensor& seqlens_k, const int64_t num_heads_per_head_k, const int64_t num_heads_k); diff --git a/sgl-kernel/python/sgl_kernel/flash_mla.py b/sgl-kernel/python/sgl_kernel/flash_mla.py index 614c54e2e..144ddc31a 100644 --- a/sgl-kernel/python/sgl_kernel/flash_mla.py +++ b/sgl-kernel/python/sgl_kernel/flash_mla.py @@ -3,7 +3,7 @@ from typing import Optional, Tuple import torch try: - from . import flashmla_ops # triggers TORCH extension registration + from sgl_kernel import flashmla_ops # triggers TORCH extension registration except Exception as _e: _flashmla_import_error = _e else: @@ -35,6 +35,12 @@ def get_mla_metadata( tile_scheduler_metadata: (num_sm_parts, TileSchedulerMetaDataSize), dtype torch.int32. num_splits: (batch_size + 1), dtype torch.int32. """ + if is_fp8_kvcache and topk is None: + return torch.ops.sgl_kernel.get_mla_decoding_metadata_dense_fp8.default( + cache_seqlens, + num_q_tokens_per_head_k, + num_heads_k, + ) return torch.ops.sgl_kernel.get_mla_decoding_metadata.default( cache_seqlens, num_q_tokens_per_head_k, @@ -55,6 +61,8 @@ def flash_mla_with_kvcache( num_splits: torch.Tensor, softmax_scale: Optional[float] = None, causal: bool = False, + descale_q: torch.Tensor | None = None, + descale_k: torch.Tensor | None = None, is_fp8_kvcache: bool = False, indices: Optional[torch.Tensor] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: @@ -69,6 +77,8 @@ def flash_mla_with_kvcache( num_splits: (batch_size + 1), torch.int32, returned by get_mla_metadata. softmax_scale: float. The scale of QK^T before applying softmax. Default to 1 / sqrt(head_dim). causal: bool. Whether to apply causal attention mask. + descale_q: (batch_size), torch.float32. Descaling factors for Q, used for fp8 quantization. + descale_k: (batch_size), torch.float32. Descaling factors for K, used for fp8 quantization. is_fp8_kvcache: bool. Whether the k_cache and v_cache are in fp8 format. For the format of FP8 KV cache, please refer to README.md indices: (batch_size, seq_len_q, topk), torch.int32. If not None, sparse attention will be enabled, and only tokens in the `indices` array will be attended to. Invalid indices should be set to -1 or numbers >= total_seq_len_kv. For details about how to set up `indices`, please refer to README.md. @@ -80,19 +90,38 @@ def flash_mla_with_kvcache( softmax_scale = q.shape[-1] ** (-0.5) if indices is not None: assert causal == False, "causal must be `false` if sparse attention is enabled." - out, softmax_lse = torch.ops.sgl_kernel.fwd_kvcache_mla.default( - q, - k_cache, - head_dim_v, - cache_seqlens, - block_table, - softmax_scale, - causal, - tile_scheduler_metadata, - num_splits, - is_fp8_kvcache, - indices, - ) + assert (descale_q is None) == ( + descale_k is None + ), "descale_q and descale_k should be both None or both not None" + + if indices is None and q.element_size() == 1: + out, softmax_lse = torch.ops.sgl_kernel.fwd_kvcache_mla_fp8.default( + q, + k_cache, + head_dim_v, + cache_seqlens, + block_table, + softmax_scale, + causal, + tile_scheduler_metadata, + num_splits, + descale_q, + descale_k, + ) + else: + out, softmax_lse = torch.ops.sgl_kernel.fwd_kvcache_mla.default( + q, + k_cache, + head_dim_v, + cache_seqlens, + block_table, + softmax_scale, + causal, + tile_scheduler_metadata, + num_splits, + is_fp8_kvcache, + indices, + ) return out, softmax_lse diff --git a/sgl-kernel/tests/test_flashmla.py b/sgl-kernel/tests/test_flashmla.py index b6c049999..0d5a07a46 100644 --- a/sgl-kernel/tests/test_flashmla.py +++ b/sgl-kernel/tests/test_flashmla.py @@ -11,8 +11,6 @@ from sgl_kernel.flash_mla import ( get_mla_metadata, ) -skip_condition = torch.cuda.get_device_capability() < (10, 0) - # ================ prefill usage ================ # S_Q_PREFILL = [1, 62] KV_TOPK_PREFILL = [ @@ -514,5 +512,143 @@ def test_flash_mla_decode( torch.testing.assert_close(lse_ans, lse_ref, atol=1e-6, rtol=8.01 / 65536) +@pytest.mark.parametrize("b", [128]) +@pytest.mark.parametrize("s_q", [1, 2]) +@pytest.mark.parametrize("mean_sk", [4096, 8192, 16384]) +@pytest.mark.parametrize("h_q", [16, 32, 64, 128]) +@pytest.mark.parametrize("h_kv", [1]) +@pytest.mark.parametrize("d", [576]) +@pytest.mark.parametrize("dv", [512]) +@pytest.mark.parametrize("block_size", [64]) +@pytest.mark.parametrize("causal", [True]) +@pytest.mark.parametrize("varlen", [False, True]) +@pytest.mark.parametrize("torch_dtype", [torch.float8_e4m3fn]) +@torch.inference_mode() +def test_flash_mla_fp8( + b, s_q, mean_sk, h_q, h_kv, d, dv, block_size, causal, varlen, torch_dtype +): + device = torch.device("cuda:0") + init_dtype = torch.bfloat16 if torch_dtype == torch.float8_e4m3fn else torch_dtype + torch.set_default_dtype(init_dtype) + torch.set_default_device(device) + torch.cuda.set_device(device) + torch.manual_seed(0) + random.seed(0) + + use_fp8 = torch_dtype == torch.float8_e4m3fn + cache_seqlens = torch.full((b,), mean_sk, dtype=torch.int32) + if varlen: + for i in range(b): + cache_seqlens[i] = max(random.normalvariate(mean_sk, mean_sk / 2), s_q) + total_seqlens = cache_seqlens.sum().item() + max_seqlen = cache_seqlens.max().item() + max_seqlen_pad = triton.cdiv(max_seqlen, 256) * 256 + + q = torch.randn(b, s_q, h_q, d) + block_table = torch.arange( + b * max_seqlen_pad // block_size, dtype=torch.int32 + ).view(b, max_seqlen_pad // block_size) + blocked_k = torch.randn(block_table.numel(), block_size, h_kv, d) + for i in range(b): + blocked_k.view(b, max_seqlen_pad, h_kv, d)[i, cache_seqlens[i].item() :] = ( + float("nan") + ) + blocked_v = blocked_k[..., :dv] + + tile_scheduler_metadata, num_splits = get_mla_metadata( + cache_seqlens, s_q * h_q // h_kv, h_kv, is_fp8_kvcache=use_fp8 + ) + + init_dtype = q.dtype + if use_fp8: + fp8_dtype = torch.float8_e4m3fn + descale_q = torch.ones((1), dtype=torch.float32) + descale_k = torch.ones((1), dtype=torch.float32) + + q = q.to(fp8_dtype) + blocked_k = blocked_k.to(fp8_dtype) + blocked_v = blocked_v.to(fp8_dtype) + else: + descale_q = None + descale_k = None + + def flash_mla(): + return flash_mla_with_kvcache( + q, + blocked_k, + block_table, + cache_seqlens, + dv, + tile_scheduler_metadata, + num_splits, + causal=causal, + descale_q=descale_q, + descale_k=descale_k, + ) + + def scaled_dot_product_attention(query, key, value, is_causal=False): + query = query.float() + key = key.float() + value = value.float() + key = key.repeat_interleave(h_q // h_kv, dim=0) + value = value.repeat_interleave(h_q // h_kv, dim=0) + attn_weight = query @ key.transpose(-2, -1) / math.sqrt(query.size(-1)) + if is_causal: + s_q = query.shape[-2] + s_k = key.shape[-2] + attn_bias = torch.zeros(s_q, s_k, dtype=query.dtype) + temp_mask = torch.ones(s_q, s_k, dtype=torch.bool).tril(diagonal=s_k - s_q) + attn_bias.masked_fill_(temp_mask.logical_not(), float("-inf")) + attn_bias.to(query.dtype) + attn_weight += attn_bias + lse = attn_weight.logsumexp(dim=-1) + attn_weight = torch.softmax(attn_weight, dim=-1, dtype=torch.float32) + return attn_weight @ value, lse + + def ref_mla(): + q_ = (q.to(torch.float) * descale_q).to(init_dtype) if use_fp8 else q + blocked_k_ = ( + (blocked_k.to(torch.float) * descale_k).to(init_dtype) + if use_fp8 + else blocked_k + ) + blocked_v_ = ( + (blocked_v.to(torch.float) * descale_k).to(init_dtype) + if use_fp8 + else blocked_v + ) + out = torch.empty(b, s_q, h_q, dv, dtype=torch.float32) + lse = torch.empty(b, h_q, s_q, dtype=torch.float32) + for i in range(b): + begin = i * max_seqlen_pad + end = begin + cache_seqlens[i] + out_i, lse_i = scaled_dot_product_attention( + q_[i].transpose(0, 1), + blocked_k_.view(-1, h_kv, d)[begin:end].transpose(0, 1), + blocked_v_.view(-1, h_kv, dv)[begin:end].transpose(0, 1), + is_causal=causal, + ) + out[i] = out_i.transpose(0, 1) + lse[i] = lse_i + return out, lse + + def cal_diff( + x: torch.Tensor, y: torch.Tensor, name: str, use_fp8: bool = False + ) -> None: + x, y = x.double(), y.double() + cos_diff = 1 - 2 * (x * y).sum().item() / max( + (x * x + y * y).sum().item(), 1e-12 + ) + if use_fp8: + assert cos_diff < 1e-4 + else: + assert cos_diff < 1e-5 + + out_flash, lse_flash = flash_mla() + out_torch, lse_torch = ref_mla() + cal_diff(out_flash, out_torch, "out", use_fp8) + cal_diff(lse_flash, lse_torch, "lse") + + if __name__ == "__main__": pytest.main([__file__])