[PD & HiSparse] Add DeepSeek V4 support for HiSparse direct Prefill-to-Decode DRAM (#24880)

python/sglang/jit_kernel/hisparse.py

@@ -39,6 +39,39 @@ def _jit_sparse_module(
     )
 
 
+@functools.cache
+def _jit_dsv4_transfer_module(block_size: int) -> Module:
+    template_args = make_cpp_args(block_size)
+    return load_jit(
+        "sparse_cache_dsv4_transfer",
+        block_size,
+        cuda_files=["hisparse.cuh"],
+        cuda_wrappers=[
+            (
+                "transfer_cache_dsv4_mla",
+                f"transfer_cache_dsv4_mla<{template_args}>",
+            )
+        ],
+    )
+
+
+def transfer_cache_dsv4_mla(
+    src_ptrs: torch.Tensor,
+    dst_ptrs: torch.Tensor,
+    src_indices: torch.Tensor,
+    dst_indices: torch.Tensor,
+    block_size: int = 1024,
+) -> None:
+    """Transfer DSv4 C4 tokens between page-padded C4 buffers."""
+    module = _jit_dsv4_transfer_module(block_size)
+    module.transfer_cache_dsv4_mla(
+        src_ptrs,
+        dst_ptrs,
+        src_indices,
+        dst_indices,
+    )
+
+
 def _load_cache_to_device_buffer_mla(
     *,
     is_dsv4_layout: bool,
@@ -156,7 +189,7 @@ def load_cache_to_device_buffer_dsv4_mla(
     block_size: int = 256,
     num_real_reqs: torch.Tensor | None = None,
 ) -> None:
-    """DSv4 hisparse swap-in: page-padded device + linear host (kvcacheio.cuh layout)."""
+    """DSv4 hisparse swap-in: page-padded device + page-padded host C4 layout."""
     _load_cache_to_device_buffer_mla(
         is_dsv4_layout=True,
         top_k_tokens=top_k_tokens,