chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[sgl-kernel] support custom fp8 flashmla kernel (#13087)

Browse Source
This commit is contained in:
Fan Yin
2025-11-14 04:45:21 +08:00
committed by GitHub
parent 8779100768
commit 2966367a31
5 changed files with 214 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
sgl-kernel/cmake/flashmla.cmake
View File

@@ -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 $<$<COMPILE_LANGUAGE:CUDA>:${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
)

9
sgl-kernel/csrc/flashmla_extension.cc
View File

@@ -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)

18
sgl-kernel/include/sgl_kernel_ops.h
View File

@@ -935,3 +935,21 @@ void FMHACutlassSM100FwdRun(
std::vector<at::Tensor>
sparse_prefill_fwd(const at::Tensor& q, const at::Tensor& kv, const at::Tensor& indices, double sm_scale, int64_t d_v);
std::vector<at::Tensor> 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<at::Tensor>& descale_q, // None or batch_size
const std::optional<at::Tensor>& descale_k // None or batch_size
);
std::vector<at::Tensor> get_mla_decoding_metadata_dense_fp8(
at::Tensor& seqlens_k, const int64_t num_heads_per_head_k, const int64_t num_heads_k);

57
sgl-kernel/python/sgl_kernel/flash_mla.py
View File

@@ -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

140
sgl-kernel/tests/test_flashmla.py
View File

@@ -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__])