minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

fix: support FA4 backend for GLM4.7-flash (#33436)

Browse Source
This commit is contained in:
Yuzhen Zhou
2026-08-09 20:05:38 +08:00
committed by GitHub
parent 71043b9dbb
commit 7120f3ee13
2 changed files with 260 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/ops/attention/flash_attention_v4.py
+65
View File
@@ -4,6 +4,7 @@ import os
from typing import Callable, Optional, Tuple, Union
import torch
import torch.nn.functional as F
from sglang.kernel_api_logging import debug_kernel_api
@@ -31,6 +32,60 @@ def _maybe_contiguous(x: Optional[torch.Tensor]) -> Optional[torch.Tensor]:
return x.contiguous() if x is not None and x.stride(-1) != 1 else x
def _pad_mla_q_heads(q, qv, v, pack_gqa):
if qv is None or pack_gqa is False:
return q, qv, None
num_heads = qv.shape[-2]
num_kv_heads = v.shape[-2]
qhead_per_kvhead = num_heads // num_kv_heads
if 128 % qhead_per_kvhead == 0 or qhead_per_kvhead % 128 == 0:
return q, qv, None
qhead_per_kvhead_padded = 1 << (qhead_per_kvhead - 1).bit_length()
def pad(x):
if x is None:
return None
prefix = x.shape[:-2]
x = x.reshape(*prefix, num_kv_heads, qhead_per_kvhead, x.shape[-1])
x = F.pad(x, (0, 0, 0, qhead_per_kvhead_padded - qhead_per_kvhead))
return x.reshape(*prefix, num_kv_heads * qhead_per_kvhead_padded, x.shape[-1])
# Pad each KV group to a valid ratio so MLA stays on the packed kernel.
return (
pad(q),
pad(qv),
(
num_kv_heads,
qhead_per_kvhead,
qhead_per_kvhead_padded,
),
)
def _unpad_mla_result(result, head_padding):
if head_padding is None:
return result
num_kv_heads, qhead_per_kvhead, qhead_per_kvhead_padded = head_padding
out, lse = result
prefix = out.shape[:-2]
out = out.reshape(*prefix, num_kv_heads, qhead_per_kvhead_padded, out.shape[-1])[
..., :qhead_per_kvhead, :
]
out = out.reshape(
*prefix, num_kv_heads * qhead_per_kvhead, out.shape[-1]
).contiguous()
if lse is not None:
prefix = lse.shape[:-1]
lse = lse.reshape(*prefix, num_kv_heads, qhead_per_kvhead_padded)[
..., :qhead_per_kvhead
]
lse = lse.reshape(*prefix, num_kv_heads * qhead_per_kvhead).contiguous()
return out, lse
@debug_kernel_api
def flash_attn_varlen_func(
q: torch.Tensor,
@@ -81,6 +136,15 @@ def flash_attn_varlen_func(
) from _flash_attn_import_error
q, k, v, qv = [_maybe_contiguous(t) for t in (q, k, v, qv)]
if qv is None and q.shape[-1] == 256 and k.shape[-1] == 256 and v.shape[-1] == 256:
# The vendored hd256 kernel assumes dense Q/K/V strides.
# TODO: Remove this workaround after the FA4 in current environment includes
# https://github.com/Dao-AILab/flash-attention/pull/2670 (flash-attn-4 >= 4.0.0b20).
q, k, v = [t.contiguous() for t in (q, k, v)]
q, qv, mla_head_padding = _pad_mla_q_heads(q, qv, v, pack_gqa)
if qv is not None and num_splits < 1:
# FA4 MLA does not implement split-KV; auto mode must use one split.
num_splits = 1
cu_seqlens_q, cu_seqlens_k = [
_maybe_contiguous(t) for t in (cu_seqlens_q, cu_seqlens_k)
]
@@ -145,6 +209,7 @@ def flash_attn_varlen_func(
**descale_kwargs,
**rel_bias_kwargs,
)
result = _unpad_mla_result(result, mla_head_padding)
if return_softmax_lse:
return result
test/registered/kernels/ops/attention/test_flash_attention_4.py
+195 -18
View File
@@ -1515,29 +1515,49 @@ def _generate_block_kvcache(
not is_sm100_or_sm110_supported(),
reason="flash_attn.cute implements qv on SM100/SM110 only (not SM120).",
)
@pytest.mark.parametrize("mha_type", ["mqa", "gqa"])
@pytest.mark.parametrize(
"seqlen_q,seqlen_k",
"seqlen_q,seqlen_k,nheads,nheads_k,num_splits",
[
(1, 128), # plain decode
(4, 1024), # speculative decode (multiple q rows per request)
(64, 800), # chunked extend
(16, 20000), # long context
(*sequence_shape, *head_config)
for sequence_shape, head_config in itertools.product(
[
(1, 128), # plain decode
(4, 1024), # speculative decode (multiple q rows per request)
(64, 800), # chunked extend
(16, 20000), # long context
],
[
(8, 1, 1), # DeepSeek-style MQA, tile-compatible head ratio
(8, 4, 1), # GQA
(20, 1, 0), # GLM-4.7-Flash TP1, padded head ratio
],
)
]
+ [
pytest.param(1, 128, 10, 1, 0, id="glm-tp2"),
pytest.param(1, 128, 5, 1, 0, id="glm-tp4"),
],
)
def test_flash_attn_varlen_qv_deepseek_absorbed(seqlen_q, seqlen_k, mha_type):
"""DeepSeek absorbed-MLA FA4 shape: rope q/k head_dim 64, latent v/qv
head_dim 512, varlen q over a paged KV cache, num_splits=1. Mirrors the
production calls in flashattention_backend.py, where extend
(flash_attn_varlen_func) and decode (flash_attn_with_kvcache) share this
qv-threaded path.
def test_flash_attn_varlen_qv_deepseek_absorbed(
seqlen_q, seqlen_k, nheads, nheads_k, num_splits
):
"""Absorbed-MLA FA4 shape: rope q/k head_dim 64, latent v/qv head_dim 512,
varlen q over a paged KV cache. Mirrors the production calls in
flashattention_backend.py, where extend (flash_attn_varlen_func) and decode
(flash_attn_with_kvcache) share this qv-threaded path.
The (20, 1, 0) case is the GLM-4.7-Flash TP1 shape: a 20:1 head ratio is
incompatible with the MLA kernel's 128-row cluster tile, so the SGLang
wrapper must pad each KV group to 32 q heads and crop the output, and must
coerce num_splits=0 (the production non-deterministic decode default) to 1
because FA4 MLA has no split-KV. Without the wrapper fix this case dies on
"split kv not supported with qv" (num_splits=0) or the cluster_tile_m ratio
assert in flash_fwd_mla_sm100.py (num_splits=1).
"""
device = "cuda"
dtype = torch.bfloat16
torch.random.manual_seed(seqlen_q + seqlen_k)
batch_size = 5
nheads = 8
nheads_k = 1 if mha_type == "mqa" else 4
d, dv = 64, 512
page_size = 128
@@ -1556,7 +1576,7 @@ def test_flash_attn_varlen_qv_deepseek_absorbed(seqlen_q, seqlen_k, mha_type):
torch.arange(batch_size + 1, dtype=torch.int32, device=device) * seqlen_q
)
out_unpad = flash_attn_varlen_func(
out_unpad, lse_unpad = flash_attn_varlen_func(
rearrange(q, "b s h d -> (b s) h d"),
k_cache_paged,
v_cache_paged,
@@ -1566,14 +1586,17 @@ def test_flash_attn_varlen_qv_deepseek_absorbed(seqlen_q, seqlen_k, mha_type):
seqused_k=cache_seqlens,
page_table=page_table,
causal=True,
num_splits=1,
num_splits=num_splits,
return_softmax_lse=True,
ver=4,
)
assert out_unpad.shape == (batch_size * seqlen_q, nheads, dv)
assert lse_unpad.shape == (batch_size * seqlen_q, nheads)
out = rearrange(out_unpad, "(b s) h d -> b s h d", b=batch_size)
# Decode enters through the flash_attn_with_kvcache wrapper; it must
# thread qv/num_splits down to the same varlen kernel call bit-for-bit.
out_kvcache = flash_attn_with_kvcache(
out_kvcache, lse_kvcache = flash_attn_with_kvcache(
q=rearrange(q, "b s h d -> (b s) h d"),
k_cache=k_cache_paged,
v_cache=v_cache_paged,
@@ -1583,16 +1606,37 @@ def test_flash_attn_varlen_qv_deepseek_absorbed(seqlen_q, seqlen_k, mha_type):
cu_seqlens_q=cu_seqlens_q,
max_seqlen_q=seqlen_q,
causal=True,
num_splits=1,
num_splits=num_splits,
return_softmax_lse=True,
ver=4,
)
assert torch.equal(out_kvcache, out_unpad)
assert torch.equal(lse_kvcache, lse_unpad)
key_padding_mask = rearrange(
torch.arange(seqlen_k, device=device), "s -> 1 s"
) < rearrange(cache_seqlens, "b -> b 1")
k_rep = repeat(k_cache, "b s h d -> b s (h g) d", g=nheads // nheads_k)
v_rep = repeat(v_cache, "b s h d -> b s (h g) d", g=nheads // nheads_k)
scores = torch.einsum(
"bthd,bshd->bhts", q.float() / math.sqrt(d + dv), k_rep.float()
)
scores += torch.einsum(
"bthd,bshd->bhts", qv.float() / math.sqrt(d + dv), v_rep.float()
)
scores.masked_fill_(rearrange(~key_padding_mask, "b s -> b 1 1 s"), float("-inf"))
scores.masked_fill_(
construct_local_mask(
seqlen_q,
seqlen_k,
(None, 0),
key_padding_mask=key_padding_mask,
device=device,
),
float("-inf"),
)
lse_ref = rearrange(torch.logsumexp(scores, dim=-1), "b h s -> (b s) h")
torch.testing.assert_close(lse_unpad, lse_ref, rtol=1e-4, atol=1e-4)
out_ref, _ = attention_ref(
q, k_rep, v_rep, None, key_padding_mask, causal=True, qv=qv
)
@@ -1618,5 +1662,138 @@ def test_flash_attn_varlen_qv_deepseek_absorbed(seqlen_q, seqlen_k, mha_type):
).abs().mean().item()
@pytest.mark.skipif(
not is_sm100_or_sm110_supported(),
reason="flash_attn.cute implements qv on SM100/SM110 only (not SM120).",
)
def test_flash_attn_qv_paged_decode_cuda_graph():
device = "cuda"
dtype = torch.bfloat16
batch_size, seqlen_q, seqlen_k = 2, 1, 128
nheads, nheads_k, d, dv = 20, 1, 64, 512
torch.random.manual_seed(0)
q = torch.randn(batch_size, nheads, d, device=device, dtype=dtype)
qv = torch.randn(batch_size, nheads, dv, device=device, dtype=dtype)
_, _, page_table, k_cache_paged, v_cache_paged, _ = _generate_block_kvcache(
seqlen_k,
128,
batch_size,
nheads_k,
d,
dv,
device,
dtype,
dtype,
)
cache_seqlens = torch.full(
(batch_size,), seqlen_k, dtype=torch.int32, device=device
)
cu_seqlens_q = torch.arange(batch_size + 1, dtype=torch.int32, device=device)
def run(q_input, qv_input):
return flash_attn_with_kvcache(
q=q_input,
k_cache=k_cache_paged,
v_cache=v_cache_paged,
qv=qv_input,
page_table=page_table,
cache_seqlens=cache_seqlens,
cu_seqlens_q=cu_seqlens_q,
max_seqlen_q=seqlen_q,
causal=True,
num_splits=0,
return_softmax_lse=True,
ver=4,
)
warmup_stream = torch.cuda.Stream()
warmup_stream.wait_stream(torch.cuda.current_stream())
with torch.cuda.stream(warmup_stream):
run(q, qv)
torch.cuda.current_stream().wait_stream(warmup_stream)
graph = torch.cuda.CUDAGraph()
with torch.cuda.graph(graph):
graph_out, graph_lse = run(q, qv)
q_replay = torch.randn_like(q)
qv_replay = torch.randn_like(qv)
q.copy_(q_replay)
qv.copy_(qv_replay)
graph.replay()
torch.cuda.synchronize()
replay_out, replay_lse = graph_out.clone(), graph_lse.clone()
eager_out, eager_lse = run(q_replay, qv_replay)
torch.testing.assert_close(replay_out, eager_out, rtol=0, atol=0)
torch.testing.assert_close(replay_lse, eager_lse, rtol=0, atol=0)
@pytest.mark.skipif(
not is_sm100_or_sm110_supported(),
reason="The dedicated hd256 kernel runs on SM100/SM110 only.",
)
@pytest.mark.parametrize("strided_input", ["q", "k", "v"])
def test_flash_attn_hd256_noncontiguous_inputs(strided_input):
device = "cuda"
dtype = torch.bfloat16
batch_size, seqlen_q, seqlen_k, nheads, d = 2, 4, 128, 4, 256
torch.random.manual_seed(0)
inputs = {}
for name, seqlen in (("q", seqlen_q), ("k", seqlen_k), ("v", seqlen_k)):
fused = torch.randn(
batch_size * seqlen,
nheads * 2,
d,
device=device,
dtype=dtype,
)
view = fused[:, ::2, :]
inputs[name] = view if name == strided_input else view.contiguous()
strided = inputs[strided_input]
assert not strided.is_contiguous()
assert strided.stride(-1) == 1
cu_seqlens_q = (
torch.arange(batch_size + 1, dtype=torch.int32, device=device) * seqlen_q
)
cu_seqlens_k = (
torch.arange(batch_size + 1, dtype=torch.int32, device=device) * seqlen_k
)
dense_inputs = {name: tensor.contiguous() for name, tensor in inputs.items()}
out, lse = flash_attn_varlen_func(
inputs["q"],
inputs["k"],
inputs["v"],
cu_seqlens_q=cu_seqlens_q,
cu_seqlens_k=cu_seqlens_k,
max_seqlen_q=seqlen_q,
max_seqlen_k=seqlen_k,
num_splits=1,
pack_gqa=False,
return_softmax_lse=True,
ver=4,
)
out_dense, lse_dense = flash_attn_varlen_func(
dense_inputs["q"],
dense_inputs["k"],
dense_inputs["v"],
cu_seqlens_q=cu_seqlens_q,
cu_seqlens_k=cu_seqlens_k,
max_seqlen_q=seqlen_q,
max_seqlen_k=seqlen_k,
num_splits=1,
pack_gqa=False,
return_softmax_lse=True,
ver=4,
)
assert torch.equal(out, out_dense)
assert torch.equal(lse, lse_dense)
if __name__ == "__main__":
sys.exit(pytest.main([__file__, "-v", "-s"]))