Deepseek V4 (#23882)

Co-authored-by: Baizhou Zhang <sobereddiezhang@gmail.com>
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Co-authored-by: fzyzcjy <ch271828n@outlook.com>
Co-authored-by: ispobock <ispobaoke@gmail.com>
Co-authored-by: Zhiqiang Xie <xiezhq@stanford.edu>
Co-authored-by: yueming-yuan <yym022502@gmail.com>
Co-authored-by: DarkSharpness <2040703891@qq.com>
Co-authored-by: Yuhao Yang <47235274+yhyang201@users.noreply.github.com>
Co-authored-by: yhyang201 <yhyang201@users.noreply.github.com>
Co-authored-by: yhyang201 <yhyang201@gmail.com>
Co-authored-by: Qiaolin Yu <90088090+qiaolin-yu@users.noreply.github.com>
Co-authored-by: Ethan (Yusheng) Su <11704492+yushengsu-thu@users.noreply.github.com>
Co-authored-by: Mingyi <27337995+wisclmy0611@users.noreply.github.com>
Co-authored-by: Cheng Wan <54331508+ch-wan@users.noreply.github.com>
Co-authored-by: Yihao Wang <42559837+againstentropy@users.noreply.github.com>

python/sglang/jit_kernel/hisparse.py

@@ -18,10 +18,13 @@ def _jit_sparse_module(
     num_top_k: int,
     hot_buffer_size: int,
     is_mla: bool = False,
+    is_dsv4_layout: bool = False,
 ) -> Module:
-    template_args = make_cpp_args(block_size, num_top_k, hot_buffer_size, is_mla)
+    template_args = make_cpp_args(
+        block_size, num_top_k, hot_buffer_size, is_mla, is_dsv4_layout
+    )
     cache_args = make_cpp_args(
-        item_size_bytes, block_size, num_top_k, hot_buffer_size, is_mla
+        item_size_bytes, block_size, num_top_k, hot_buffer_size, is_mla, is_dsv4_layout
     )
     return load_jit(
         "sparse_cache",
@@ -36,7 +39,9 @@ def _jit_sparse_module(
     )
 
 
-def load_cache_to_device_buffer_mla(
+def _load_cache_to_device_buffer_mla(
+    *,
+    is_dsv4_layout: bool,
     top_k_tokens: torch.Tensor,
     device_buffer_tokens: torch.Tensor,
     host_cache_locs: torch.Tensor,
@@ -50,16 +55,21 @@ def load_cache_to_device_buffer_mla(
     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,
+    page_size: int,
+    block_size: int,
+    num_real_reqs: torch.Tensor | 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
+        item_size_bytes,
+        block_size,
+        num_top_k,
+        hot_buffer_size,
+        is_mla=True,
+        is_dsv4_layout=is_dsv4_layout,
     )
 
     empty = torch.empty(0)
@@ -86,3 +96,83 @@ def load_cache_to_device_buffer_mla(
         page_size,
         item_size_bytes,
     )
+
+
+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:
+    """Generic MLA hisparse swap-in: device + host both linear (stride=item_size_bytes)."""
+    _load_cache_to_device_buffer_mla(
+        is_dsv4_layout=False,
+        top_k_tokens=top_k_tokens,
+        device_buffer_tokens=device_buffer_tokens,
+        host_cache_locs=host_cache_locs,
+        device_buffer_locs=device_buffer_locs,
+        host_cache=host_cache,
+        device_buffer=device_buffer,
+        top_k_device_locs=top_k_device_locs,
+        req_pool_indices=req_pool_indices,
+        seq_lens=seq_lens,
+        lru_slots=lru_slots,
+        item_size_bytes=item_size_bytes,
+        num_top_k=num_top_k,
+        hot_buffer_size=hot_buffer_size,
+        page_size=page_size,
+        block_size=block_size,
+        num_real_reqs=num_real_reqs,
+    )
+
+
+def load_cache_to_device_buffer_dsv4_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:
+    """DSv4 hisparse swap-in: page-padded device + linear host (kvcacheio.cuh layout)."""
+    _load_cache_to_device_buffer_mla(
+        is_dsv4_layout=True,
+        top_k_tokens=top_k_tokens,
+        device_buffer_tokens=device_buffer_tokens,
+        host_cache_locs=host_cache_locs,
+        device_buffer_locs=device_buffer_locs,
+        host_cache=host_cache,
+        device_buffer=device_buffer,
+        top_k_device_locs=top_k_device_locs,
+        req_pool_indices=req_pool_indices,
+        seq_lens=seq_lens,
+        lru_slots=lru_slots,
+        item_size_bytes=item_size_bytes,
+        num_top_k=num_top_k,
+        hot_buffer_size=hot_buffer_size,
+        page_size=page_size,
+        block_size=block_size,
+        num_real_reqs=num_real_reqs,
+    )