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

python/sglang/kernels/ops/attention/flash_attention_v4.py

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

@@ -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
+        (*sequence_shape, *head_config)
-        (4, 1024),  # speculative decode (multiple q rows per request)
+        for sequence_shape, head_config in itertools.product(
-        (64, 800),  # chunked extend
+            [
-        (16, 20000),  # long context
+                (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):
+def test_flash_attn_varlen_qv_deepseek_absorbed(
-    """DeepSeek absorbed-MLA FA4 shape: rope q/k head_dim 64, latent v/qv
+    seqlen_q, seqlen_k, nheads, nheads_k, num_splits
-    head_dim 512, varlen q over a paged KV cache, num_splits=1. Mirrors the
+):
-    production calls in flashattention_backend.py, where extend
+    """Absorbed-MLA FA4 shape: rope q/k head_dim 64, latent v/qv head_dim 512,
-    (flash_attn_varlen_func) and decode (flash_attn_with_kvcache) share this
+    varlen q over a paged KV cache. Mirrors the production calls in
-    qv-threaded path.
+    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"]))