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