diff --git a/python/sglang/kernels/ops/attention/flash_attention_v4.py b/python/sglang/kernels/ops/attention/flash_attention_v4.py index eaf935363..613c6daf0 100644 --- a/python/sglang/kernels/ops/attention/flash_attention_v4.py +++ b/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 diff --git a/test/registered/kernels/ops/attention/test_flash_attention_4.py b/test/registered/kernels/ops/attention/test_flash_attention_4.py index 5a16481f5..69132bba3 100644 --- a/test/registered/kernels/ops/attention/test_flash_attention_4.py +++ b/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 - (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"]))