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/python/sgl_kernel/infllm_v2/max_pooling.py

@@ -0,0 +1,61 @@
+import torch
+
+
+def max_pooling_1d_varlen(
+    input: torch.Tensor,  # num_heads x total_q x max_k
+    cu_seqlens_q: torch.Tensor,  # batch_size + 1
+    cu_seqlens_k: torch.Tensor,  # batch_size + 1
+    cache_lens: torch.Tensor,  # batch_size
+    max_seqlen_q: int,
+    max_context_len: int,
+    local_blocks: int,
+    init_blocks: int,
+    block_size: int = 64,
+    stride: int = 16,
+    total_q: int = -1,
+) -> torch.Tensor:
+    """Variable-length 1D max pooling over packed sequences.
+
+    Drop-in replacement for ``infllm_v2.max_pooling_1d_varlen``.
+    """
+    assert input.dtype in (torch.float16, torch.bfloat16)
+    assert cu_seqlens_q.dtype == torch.int32
+    assert cu_seqlens_k.dtype == torch.int32
+    assert cache_lens.dtype == torch.int32
+    assert input.dim() == 3, f"Expected 3D input, got {input.dim()}D"
+
+    input = input.contiguous()
+    cu_seqlens_q = cu_seqlens_q.contiguous()
+    cu_seqlens_k = cu_seqlens_k.contiguous()
+    cache_lens = cache_lens.contiguous()
+
+    max_seqlen_k = max_context_len // stride
+    out_len = (max_context_len + block_size - 1) // block_size
+
+    stride = block_size // stride
+    kernel_size = stride + 1
+    padding = 1
+
+    num_heads = input.shape[0]
+    total_q = input.shape[1]
+
+    output = torch.zeros(
+        num_heads, total_q, out_len, device=input.device, dtype=input.dtype
+    )
+    torch.ops.sgl_kernel.infllm_v2_max_pooling_1d_varlen.default(
+        input,
+        output,
+        cu_seqlens_q,
+        cu_seqlens_k,
+        cache_lens,
+        max_seqlen_q,
+        max_seqlen_k,
+        kernel_size,
+        stride,
+        padding,
+        block_size,
+        local_blocks,
+        init_blocks,
+        total_q,
+    )
+    return output