HiSparse for Sparse Attention (#20343)

python/sglang/jit_kernel/hisparse.py

@@ -0,0 +1,88 @@
+from __future__ import annotations
+
+import functools
+from typing import TYPE_CHECKING
+
+import torch
+
+from sglang.jit_kernel.utils import load_jit, make_cpp_args
+
+if TYPE_CHECKING:
+    from tvm_ffi.module import Module
+
+
+@functools.cache
+def _jit_sparse_module(
+    item_size_bytes: int,
+    block_size: int,
+    num_top_k: int,
+    hot_buffer_size: int,
+    is_mla: bool = False,
+) -> Module:
+    template_args = make_cpp_args(block_size, num_top_k, hot_buffer_size, is_mla)
+    cache_args = make_cpp_args(
+        item_size_bytes, block_size, num_top_k, hot_buffer_size, is_mla
+    )
+    return load_jit(
+        "sparse_cache",
+        *cache_args,
+        cuda_files=["hisparse.cuh"],
+        cuda_wrappers=[
+            (
+                "load_cache_to_device_buffer",
+                f"load_cache_to_device_buffer<{template_args}>",
+            )
+        ],
+    )
+
+
+def load_cache_to_device_buffer_mla(
+    top_k_tokens: torch.Tensor,
+    device_buffer_tokens: torch.Tensor,
+    host_cache_locs: torch.Tensor,
+    device_buffer_locs: torch.Tensor,
+    host_cache: torch.Tensor,
+    device_buffer: torch.Tensor,
+    top_k_device_locs: torch.Tensor,
+    req_pool_indices: torch.Tensor,
+    seq_lens: torch.Tensor,
+    lru_slots: torch.Tensor,
+    item_size_bytes: int,
+    num_top_k: int,
+    hot_buffer_size: int,
+    page_size: int = 1,
+    block_size: int = 256,
+    num_real_reqs: torch.Tensor | None = None,
+) -> None:
+    assert (
+        hot_buffer_size >= num_top_k
+    ), f"hot_buffer_size ({hot_buffer_size}) must be >= num_top_k ({num_top_k})"
+
+    module = _jit_sparse_module(
+        item_size_bytes, block_size, num_top_k, hot_buffer_size, is_mla=True
+    )
+
+    empty = torch.empty(0)
+
+    if num_real_reqs is None:
+        num_real_reqs = torch.tensor(
+            [top_k_tokens.size(0)], dtype=torch.int32, device=top_k_tokens.device
+        )
+
+    module.load_cache_to_device_buffer(
+        top_k_tokens,
+        device_buffer_tokens,
+        host_cache_locs,
+        device_buffer_locs,
+        host_cache,
+        empty,
+        device_buffer,
+        empty,
+        top_k_device_locs,
+        req_pool_indices,
+        seq_lens,
+        lru_slots,
+        num_real_reqs,
+        page_size,
+        item_size_bytes,
+    )