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/__init__.py

@@ -71,6 +71,10 @@ else:
         shuffle_rows,
     )
     from sgl_kernel.grammar import apply_token_bitmask_inplace_cuda
+    from sgl_kernel.infllm_v2 import (
+        infllmv2_attn_stage1,
+        max_pooling_1d_varlen,
+    )
     from sgl_kernel.kvcacheio import (
         transfer_kv_all_layer,
         transfer_kv_all_layer_mla,

sgl-kernel/python/sgl_kernel/infllm_v2/__init__.py

@@ -0,0 +1,7 @@
+from sgl_kernel.infllm_v2.attention import infllmv2_attn_stage1
+from sgl_kernel.infllm_v2.max_pooling import max_pooling_1d_varlen
+
+__all__ = [
+    "infllmv2_attn_stage1",
+    "max_pooling_1d_varlen",
+]

sgl-kernel/python/sgl_kernel/infllm_v2/_loader.py

@@ -0,0 +1,98 @@
+"""Robust loader for the standalone ``infllm_ops`` pybind extension.
+
+The InfLLM-V2 FlashAttention backend is built as its own module ``infllm_ops``
+(installed into the ``sgl_kernel`` package directory). Under editable installs
+the compiled ``.so`` may live in ``site-packages/sgl_kernel`` while the imported
+``sgl_kernel`` package resolves to the source tree, so a plain ``from sgl_kernel
+import infllm_ops`` is not always sufficient. This loader searches the known
+candidate locations and loads the extension by file path.
+"""
+
+import glob
+import importlib.util
+import site
+import sys
+from pathlib import Path
+from typing import List, Optional
+
+_infllm_ops = None
+
+
+def _candidate_dirs() -> List[Path]:
+    dirs: List[Path] = []
+
+    # 1) The directory of the sgl_kernel package as currently imported.
+    try:
+        import sgl_kernel
+
+        dirs.append(Path(sgl_kernel.__file__).parent)
+    except Exception:
+        pass
+
+    # 2) This module's parent package directory (source tree).
+    dirs.append(Path(__file__).resolve().parent.parent)
+
+    # 3) Every ``sgl_kernel`` directory found on the install paths.
+    search_roots: List[str] = []
+    try:
+        search_roots.extend(site.getsitepackages())
+    except Exception:
+        pass
+    try:
+        search_roots.append(site.getusersitepackages())
+    except Exception:
+        pass
+    search_roots.extend(p for p in sys.path if p)
+    for root in search_roots:
+        dirs.append(Path(root) / "sgl_kernel")
+
+    # De-duplicate while preserving order.
+    seen = set()
+    unique: List[Path] = []
+    for d in dirs:
+        key = str(d)
+        if key not in seen:
+            seen.add(key)
+            unique.append(d)
+    return unique
+
+
+def _find_so() -> Optional[Path]:
+    for d in _candidate_dirs():
+        if not d.is_dir():
+            continue
+        matches = sorted(glob.glob(str(d / "infllm_ops*.so")))
+        if matches:
+            return Path(matches[0])
+    return None
+
+
+def load_infllm_ops():
+    """Import and return the ``infllm_ops`` extension module (cached)."""
+    global _infllm_ops
+    if _infllm_ops is not None:
+        return _infllm_ops
+
+    # Fast path: a normal import may already work.
+    try:
+        from sgl_kernel import infllm_ops as _mod  # type: ignore
+
+        _infllm_ops = _mod
+        return _infllm_ops
+    except Exception:
+        pass
+
+    so_path = _find_so()
+    if so_path is None:
+        raise ImportError(
+            "[sgl_kernel] Could not locate the 'infllm_ops' extension (infllm_ops*.so). "
+            "Ensure sgl-kernel was built with the InfLLM-V2 FlashAttention backend."
+        )
+
+    spec = importlib.util.spec_from_file_location("infllm_ops", str(so_path))
+    if spec is None or spec.loader is None:
+        raise ImportError(f"[sgl_kernel] Could not create module spec for {so_path}")
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    _infllm_ops = module
+    return _infllm_ops

sgl-kernel/python/sgl_kernel/infllm_v2/attention.py

@@ -0,0 +1,82 @@
+"""InfLLM-V2 sparse FlashAttention public API.
+
+Ported (drop-in) from ``3rdparty/infllmv2_cuda_impl/infllm_v2/infllmv2_sparse_attention.py``.
+The CUDA backend now lives in the standalone ``infllm_ops`` extension.
+"""
+
+import torch
+from sgl_kernel.infllm_v2._loader import load_infllm_ops
+
+
+def maybe_contiguous(x):
+    return x.contiguous() if x is not None and x.stride(-1) != 1 else x
+
+
+def infllmv2_attn_stage1(
+    q,
+    k,
+    v,
+    cu_seqlens_q,
+    cu_seqlens_k,
+    cu_seqlens_v,
+    max_seqlen_q,
+    max_seqlen_k,
+    dropout_p=0.0,
+    softmax_scale=None,
+    causal=False,
+    window_size=(-1, -1),
+    softcap=0.0,
+    alibi_slopes=None,
+    deterministic=False,
+    return_attn_probs=True,
+    block_table=None,
+):
+    """Neighborhood Sparse Attention (NSA) Stage 1 with varlen support.
+
+    Drop-in replacement for ``infllm_v2.infllmv2_attn_stage1``. Returns the
+    attention-score matrix with the NSA sparsity pattern, shape
+    ``(num_heads_k, total_q, max_seqlen_k)``.
+    """
+    infllm_ops = load_infllm_ops()
+    if softmax_scale is None:
+        softmax_scale = q.shape[-1] ** (-0.5)
+
+    q, k, v = [maybe_contiguous(x) for x in (q, k, v)]
+
+    total_q, nheads, head_dim = q.shape
+    nheads_k = k.shape[1]
+    nheads_per_group = nheads // nheads_k
+
+    q = q.reshape(total_q, nheads_k, nheads_per_group, head_dim)
+    q = (
+        q.transpose(1, 2)
+        .reshape(total_q * nheads_per_group, nheads_k, head_dim)
+        .contiguous()
+    )
+
+    result = infllm_ops.varlen_fwd_stage1(
+        q,
+        k,
+        v,
+        None,
+        cu_seqlens_q,
+        cu_seqlens_k,
+        cu_seqlens_v,
+        None,
+        None,
+        block_table,
+        alibi_slopes,
+        max_seqlen_q,
+        max_seqlen_k,
+        dropout_p,
+        softmax_scale,
+        True,
+        causal,
+        window_size[0],
+        window_size[1],
+        softcap,
+        True,
+        None,
+    )
+
+    return result[0]

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