feat(sgl-kernel): add InfLLM v2 attention kernels (#29383)

Co-authored-by: Size Wang <paulgeorge13hhhhh@gmail.com>
Co-authored-by: lijiayi <lijiayi@modelbest.cn>
Co-authored-by: suhmily10 <suhmily@gmail.com>
Co-authored-by: Xiaoyue Xu <xiaoyue.xu.me@gmail.com>
Co-authored-by: hansjohn <74091612+hansjohn@users.noreply.github.com>
Co-authored-by: zhangyan <1762895426@qq.com>

sgl-kernel/tests/test_infllm_v2_attention.py

@@ -0,0 +1,61 @@
+"""Equivalence tests for the migrated InfLLM-V2 FlashAttention API.
+
+These compare the ``sgl_kernel.infllm_v2`` implementations against the original
+``infllm_v2`` package (3rdparty/infllmv2_cuda_impl). Both call the same CUDA
+kernels, so outputs are expected to match closely. The whole module is skipped
+if the reference ``infllm_v2`` package is not importable.
+"""
+
+import pytest
+import torch
+
+sgl = pytest.importorskip("sgl_kernel.infllm_v2")
+ref = pytest.importorskip("infllm_v2")
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available(), reason="CUDA is required for InfLLM-V2 kernels"
+)
+
+
+def _assert_close(a, b, name):
+    a = a.float()
+    b = b.float()
+    assert a.shape == b.shape, f"{name}: shape mismatch {a.shape} vs {b.shape}"
+    max_diff = (a - b).abs().max().item()
+    assert torch.allclose(a, b, atol=1e-2, rtol=1e-2), f"{name}: max diff {max_diff}"
+
+
+@pytest.mark.parametrize("head_dim", [64, 128])
+@pytest.mark.parametrize("causal", [False, True])
+@pytest.mark.parametrize("seqlen_q,seqlen_k", [(256, 16), (64, 17)])
+def test_stage1_matches_reference(head_dim, causal, seqlen_q, seqlen_k):
+    torch.manual_seed(0)
+    n_heads, n_kv_heads = 32, 2
+    dtype = torch.bfloat16
+
+    q = torch.randn(n_heads, seqlen_q, head_dim, dtype=dtype, device="cuda")
+    k = torch.randn(n_kv_heads, seqlen_k, head_dim, dtype=dtype, device="cuda")
+
+    cu_seqlens_q = torch.tensor([0, seqlen_q], dtype=torch.int32, device="cuda")
+    cu_seqlens_k = torch.tensor([0, seqlen_k], dtype=torch.int32, device="cuda")
+
+    q = q.transpose(0, 1).contiguous()
+    k = k.transpose(0, 1).contiguous()
+
+    common = dict(
+        cu_seqlens_q=cu_seqlens_q,
+        cu_seqlens_k=cu_seqlens_k,
+        cu_seqlens_v=cu_seqlens_k,
+        max_seqlen_q=seqlen_q,
+        max_seqlen_k=seqlen_k,
+        causal=causal,
+    )
+    out_ref = ref.infllmv2_attn_stage1(q, k, k, **common)
+    out_sgl = sgl.infllmv2_attn_stage1(q, k, k, **common)
+    _assert_close(out_sgl, out_ref, "stage1")
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main([__file__, "-v"]))

sgl-kernel/tests/test_infllm_v2_max_pooling.py

@@ -0,0 +1,109 @@
+import pytest
+import torch
+from sgl_kernel import max_pooling_1d_varlen
+
+
+def _ref_varlen(
+    score: torch.Tensor,  # [num_heads, total_q, max_k]
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    cache_lens: torch.Tensor,
+    max_context_len: int,
+    local_blocks: int,
+    init_blocks: int,
+    block_size: int,
+    kernel_stride: int,
+) -> torch.Tensor:
+    """Pure-torch reference mirroring the CUDA kernel exactly (fp32 math)."""
+    num_heads, total_q, _ = score.shape
+    out_len = (max_context_len + block_size - 1) // block_size
+    stride = block_size // kernel_stride
+    kernel_size = stride + 1
+    padding = 1
+
+    cu_q = cu_seqlens_q.tolist()
+    cu_k = cu_seqlens_k.tolist()
+    cache = cache_lens.tolist()
+    batch_size = len(cache)
+
+    out = torch.zeros(num_heads, total_q, out_len, dtype=torch.float32)
+    s = score.float().cpu()
+    for q in range(total_q):
+        b = 0
+        for bb in range(batch_size):
+            if cu_q[bb] <= q < cu_q[bb + 1]:
+                b = bb
+                break
+        bidq_local = q - cu_q[b]
+        seqlen_k = cu_k[b + 1] - cu_k[b]
+        off_bq = (bidq_local + cache[b]) // block_size
+        for h in range(num_heads):
+            for k in range(out_len):
+                if (k < init_blocks) or (off_bq >= k and off_bq <= k + local_blocks):
+                    out[h, q, k] = float("inf")
+                else:
+                    start = max(k * stride - padding, 0)
+                    end = min(start + kernel_size, seqlen_k)
+                    if end > start:
+                        out[h, q, k] = s[h, q, start:end].max()
+                    else:
+                        out[h, q, k] = float("-inf")
+    return out
+
+
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
+@pytest.mark.parametrize("num_heads", [1, 4])
+@pytest.mark.parametrize("seq_lens", [[37], [16, 48], [8, 8, 24]])
+def test_max_pooling_varlen_matches_reference(dtype, num_heads, seq_lens):
+    torch.manual_seed(0)
+    block_size = 64
+    kernel_stride = 16
+    local_blocks = 1
+    init_blocks = 1
+    max_context_len = 512
+
+    total_q = sum(seq_lens)
+    max_k = max_context_len // kernel_stride
+    cu = [0]
+    for n in seq_lens:
+        cu.append(cu[-1] + n)
+    cu_seqlens_q = torch.tensor(cu, dtype=torch.int32, device="cuda")
+    cu_seqlens_k = torch.tensor(cu, dtype=torch.int32, device="cuda")
+    cache_lens = torch.zeros(len(seq_lens), dtype=torch.int32, device="cuda")
+
+    score = torch.randn(num_heads, total_q, max_k, dtype=dtype, device="cuda")
+
+    out = max_pooling_1d_varlen(
+        score,
+        cu_seqlens_q,
+        cu_seqlens_k,
+        cache_lens,
+        max_seqlen_q=max(seq_lens),
+        max_context_len=max_context_len,
+        local_blocks=local_blocks,
+        init_blocks=init_blocks,
+        block_size=block_size,
+        stride=kernel_stride,
+        total_q=total_q,
+    )
+    ref = _ref_varlen(
+        score,
+        cu_seqlens_q,
+        cu_seqlens_k,
+        cache_lens,
+        max_context_len,
+        local_blocks,
+        init_blocks,
+        block_size,
+        kernel_stride,
+    ).to(out.device)
+
+    assert torch.equal(torch.isinf(out) & (out > 0), torch.isinf(ref) & (ref > 0))
+    finite = torch.isfinite(ref)
+    torch.testing.assert_close(out[finite].float(), ref[finite], rtol=1e-2, atol=1e-2)
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main([__file__, "-v", "-s"]))