chore: import upstream snapshot with attribution

python/sglang/jit_kernel/dsv4/__init__.py

@@ -0,0 +1,63 @@
+from .attn import (
+    fused_store_cache,
+    get_paged_mqa_logits_metadata,
+    triton_create_paged_compress_data,
+)
+from .c128_cleanup import clear_unaccepted_c128_draft_states
+from .compress import (
+    CompressorDecodePlan,
+    CompressorPrefillPlan,
+    compress_forward,
+    compress_norm_rope_store,
+)
+from .compress_old import fused_norm_rope_inplace
+from .elementwise import (
+    fused_k_norm_rope_flashmla,
+    fused_q_indexer_rope_first_quant,
+    fused_q_indexer_rope_hadamard_fp4_quant,
+    fused_q_indexer_rope_hadamard_quant,
+    fused_q_norm_rope,
+    fused_rope_inplace,
+)
+from .fp8_wo_a import sglang_per_token_group_quant_fp8_dsv4_wo_a
+from .gemm import linear_bf16_fp32
+from .moe import (
+    hash_topk,
+    mask_topk_ids,
+    mega_moe_pre_dispatch,
+    silu_and_mul_clamp,
+    silu_and_mul_contig_post_quant,
+    silu_and_mul_masked_post_quant,
+)
+from .topk import plan_topk_v2, topk_transform_512, topk_transform_512_v2
+from .utils import make_name
+
+__all__ = [
+    "CompressorDecodePlan",
+    "CompressorPrefillPlan",
+    "compress_forward",
+    "compress_norm_rope_store",
+    "clear_unaccepted_c128_draft_states",
+    "fused_norm_rope_inplace",
+    "fused_store_cache",
+    "fused_rope_inplace",
+    "fused_q_norm_rope",
+    "fused_q_indexer_rope_first_quant",
+    "fused_q_indexer_rope_hadamard_fp4_quant",
+    "fused_q_indexer_rope_hadamard_quant",
+    "fused_k_norm_rope_flashmla",
+    "sglang_per_token_group_quant_fp8_dsv4_wo_a",
+    "make_name",
+    "linear_bf16_fp32",
+    "get_paged_mqa_logits_metadata",
+    "triton_create_paged_compress_data",
+    "topk_transform_512",
+    "topk_transform_512_v2",
+    "plan_topk_v2",
+    "hash_topk",
+    "mega_moe_pre_dispatch",
+    "mask_topk_ids",
+    "silu_and_mul_clamp",
+    "silu_and_mul_masked_post_quant",
+    "silu_and_mul_contig_post_quant",
+]

python/sglang/jit_kernel/dsv4/attn.py

@@ -0,0 +1,204 @@
+from typing import Literal, Tuple
+
+import torch
+import triton
+import triton.language as tl
+
+from sglang.jit_kernel.utils import (
+    cache_once,
+    is_arch_support_pdl,
+    is_hip_runtime,
+    load_jit,
+    make_cpp_args,
+)
+
+from .utils import make_name
+
+
+@cache_once
+def _jit_metadata_module():
+    return load_jit(
+        make_name("metadata"),
+        cuda_files=["deepseek_v4/paged_mqa_metadata.cuh"],
+        cuda_wrappers=[("run", "IndexerMetadataKernel::run")],
+    )
+
+
+@cache_once
+def _jit_fused_store_module(
+    name: Literal["flashmla", "indexer"],
+    input_dtype: torch.dtype,
+    index_dtype: torch.dtype,
+    page_size: int,
+):
+    args = make_cpp_args(input_dtype, index_dtype, page_size, is_arch_support_pdl())
+    cname = "FlashMLA" if name == "flashmla" else "Indexer"
+    kernel_class = f"FusedStoreCache{cname}Kernel<{args}>"
+    return load_jit(
+        make_name("store_" + name),
+        *args,
+        cuda_files=["deepseek_v4/store.cuh"],
+        cuda_wrappers=[("run", f"{kernel_class}::run")],
+    )
+
+
+def get_paged_mqa_logits_metadata(seq_lens: torch.Tensor, page_size: int, num_sm: int):
+    assert page_size == 64
+    seq_lens = seq_lens.view(-1).to(torch.int32)
+    metadata = seq_lens.new_empty(num_sm + 1, 2)
+    module = _jit_metadata_module()
+    module.run(seq_lens, metadata)
+    return metadata
+
+
+def fused_store_cache(
+    input: torch.Tensor,
+    cache: torch.Tensor,
+    indices: torch.Tensor,
+    *,
+    page_size: int,
+    type: Literal["flashmla", "indexer"],
+) -> None:
+    if is_hip_runtime():
+        from sglang.jit_kernel.triton_store_cache import triton_fused_store_cache
+
+        triton_fused_store_cache(input, cache, indices, page_size=page_size, type=type)
+    else:
+        module = _jit_fused_store_module(
+            name=type,
+            input_dtype=input.dtype,
+            index_dtype=indices.dtype,
+            page_size=page_size,
+        )
+        module.run(input, cache, indices)
+
+
+@triton.jit
+def create_paged_compress_data_kernel(
+    req_pool_indices_ptr,
+    seq_lens_ptr,
+    extend_seq_lens_ptr,
+    req_to_token_ptr,
+    full_to_swa_index_mapping_ptr,
+    out_0_ptr,
+    out_1_ptr,
+    batch_size,
+    stride_req_to_token_0,
+    stride_req_to_token_1: tl.constexpr,
+    stride_out_1_0,
+    stride_out_1_1: tl.constexpr,
+    compress_ratio: tl.constexpr,
+    is_overlap: tl.constexpr,
+    swa_page_size: tl.constexpr,
+    ring_size: tl.constexpr,
+    BLOCK: tl.constexpr,
+) -> None:
+    pid = tl.program_id(0)
+    offs = pid * BLOCK + tl.arange(0, BLOCK)
+    mask = offs < batch_size
+
+    rid = tl.load(req_pool_indices_ptr + offs, mask=mask, other=0).to(tl.int32)
+    seq_len = tl.load(seq_lens_ptr + offs, mask=mask, other=0).to(tl.int32)
+    extend_len = tl.load(extend_seq_lens_ptr + offs, mask=mask, other=0).to(tl.int32)
+    prefix_len = seq_len - extend_len
+
+    cr = compress_ratio
+    write_pos = ((seq_len - 1) // cr) * cr
+    load_pos = ((prefix_len - 1) // cr) * cr
+    write_overlap_pos = write_pos - cr
+    load_overlap_pos = load_pos - cr
+    v0 = tl.zeros([BLOCK], tl.int32)
+    v1 = tl.zeros([BLOCK], tl.int32)
+    v2 = tl.zeros([BLOCK], tl.int32)
+    v3 = tl.zeros([BLOCK], tl.int32)
+
+    for i in tl.static_range(4):
+        if i == 0:
+            pos = load_pos
+        elif i == 1:
+            pos = write_pos
+        elif i == 2:
+            pos = load_overlap_pos
+        else:
+            pos = write_overlap_pos
+        pos = tl.maximum(pos, 0)
+        if compress_ratio == 128:
+            state_loc = rid * ring_size + (pos % ring_size)
+        else:
+            loc = tl.load(
+                req_to_token_ptr
+                + rid.to(tl.int64) * stride_req_to_token_0
+                + pos.to(tl.int64) * stride_req_to_token_1,
+                mask=mask,
+                other=0,
+            ).to(tl.int32)
+            swa_loc = tl.load(
+                full_to_swa_index_mapping_ptr + loc, mask=mask, other=0
+            ).to(tl.int32)
+            swa_page = swa_loc // swa_page_size
+            state_loc = swa_page * ring_size + (swa_loc % ring_size)
+        state_loc = state_loc // cr
+        if i == 0:
+            v0 = state_loc
+        elif i == 1:
+            v1 = state_loc
+        elif i == 2:
+            v2 = state_loc
+        else:
+            v3 = state_loc
+
+    tl.store(out_0_ptr + offs, v1, mask=mask)
+
+    if is_overlap:
+        base = out_1_ptr + offs * stride_out_1_0
+        tl.store(base + 0 * stride_out_1_1, v2, mask=mask)
+        tl.store(base + 1 * stride_out_1_1, v0, mask=mask)
+        tl.store(base + 2 * stride_out_1_1, v3, mask=mask)
+        tl.store(base + 3 * stride_out_1_1, write_pos.to(tl.int32), mask=mask)
+    else:
+        base = out_1_ptr + offs * stride_out_1_0
+        tl.store(base + 0 * stride_out_1_1, v0, mask=mask)
+
+
+def triton_create_paged_compress_data(
+    *,
+    compress_ratio: int,
+    is_overlap: bool,
+    swa_page_size: int,
+    ring_size: int,
+    req_pool_indices: torch.Tensor,
+    seq_lens: torch.Tensor,
+    extend_seq_lens: torch.Tensor,
+    req_to_token: torch.Tensor,
+    full_to_swa_index_mapping: torch.Tensor,
+    block: int = 128,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    batch_size = req_pool_indices.shape[0]
+    out_dim = 4 if is_overlap else 1
+    device_args: dict = dict(device=req_pool_indices.device, dtype=torch.int32)
+    out_0 = torch.empty((batch_size,), **device_args)
+    out_1 = torch.empty((batch_size, out_dim), **device_args)
+    grid = (triton.cdiv(batch_size, block),)
+    create_paged_compress_data_kernel[grid](
+        req_pool_indices,
+        seq_lens,
+        extend_seq_lens,
+        req_to_token,
+        full_to_swa_index_mapping,
+        out_0,
+        out_1,
+        batch_size=batch_size,
+        stride_req_to_token_0=req_to_token.stride(0),
+        stride_req_to_token_1=req_to_token.stride(1),  # type: ignore
+        stride_out_1_0=out_1.stride(0),
+        stride_out_1_1=out_1.stride(1),  # type: ignore
+        compress_ratio=compress_ratio,  # type: ignore
+        is_overlap=1 if is_overlap else 0,  # type: ignore
+        swa_page_size=swa_page_size,  # type: ignore
+        ring_size=ring_size,  # type: ignore
+        BLOCK=block,  # type: ignore
+    )
+
+    if not is_overlap:
+        out_1.squeeze_(1)
+    return out_0, out_1

python/sglang/jit_kernel/dsv4/c128_cleanup.py

@@ -0,0 +1,58 @@
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _clear_unaccepted_c128_draft_states_kernel(
+    state,
+    req_pool_indices,
+    seq_lens,
+    accept_lens,
+    ring_size: tl.constexpr,
+    half: tl.constexpr,
+    num_draft_tokens: tl.constexpr,
+    BLOCK_D: tl.constexpr,
+):
+    bid = tl.program_id(0)
+    draft_offset = tl.program_id(1)
+    block_id = tl.program_id(2)
+
+    accept_len = tl.load(accept_lens + bid)
+    if draft_offset < accept_len:
+        return
+
+    req_pool_idx = tl.load(req_pool_indices + bid).to(tl.int64)
+    seq_len = tl.load(seq_lens + bid).to(tl.int64)
+    slot = (seq_len + draft_offset) % ring_size
+    row = req_pool_idx * ring_size + slot
+
+    offsets = block_id * BLOCK_D + tl.arange(0, BLOCK_D)
+    mask = offsets < half
+    row_base = row * (half * 2)
+    tl.store(state + row_base + offsets, 0.0, mask=mask)
+    tl.store(state + row_base + half + offsets, float("-inf"), mask=mask)
+
+
+def clear_unaccepted_c128_draft_states(
+    state: torch.Tensor,
+    req_pool_indices: torch.Tensor,
+    seq_lens: torch.Tensor,
+    accept_lens: torch.Tensor,
+    *,
+    ring_size: int,
+    num_draft_tokens: int,
+) -> None:
+    half = state.shape[-1] // 2
+    _clear_unaccepted_c128_draft_states_kernel[
+        (req_pool_indices.numel(), num_draft_tokens, triton.cdiv(half, 256))
+    ](
+        state,
+        req_pool_indices,
+        seq_lens,
+        accept_lens,
+        ring_size,
+        half,
+        num_draft_tokens,
+        BLOCK_D=256,
+    )

python/sglang/jit_kernel/dsv4/compress.py

@@ -0,0 +1,371 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Literal, NamedTuple, Optional, Union
+
+import torch
+
+from sglang.jit_kernel.utils import (
+    cache_once,
+    is_arch_support_pdl,
+    load_jit,
+    make_cpp_args,
+)
+
+from .utils import make_name
+
+if TYPE_CHECKING:
+    from tvm_ffi.module import Module
+
+
+@cache_once
+def _jit_compress_norm_rope_module(
+    dtype: torch.dtype,
+    head_dim: int,
+    rope_dim: int,
+    page_size: int,
+    bf16_store: bool = False,
+) -> Module:
+    args = make_cpp_args(
+        dtype, head_dim, rope_dim, page_size, is_arch_support_pdl(), bf16_store
+    )
+    cuda_wrappers = [("forward", f"FusedNormRopeKernel<{args}>::forward")]
+    if head_dim == 128:
+        cuda_wrappers.append(
+            ("forward_fp4", f"FusedNormRopeKernel<{args}>::forward_fp4")
+        )
+    return load_jit(
+        make_name(f"fused_norm_rope_v2"),
+        *args,
+        cuda_files=[f"deepseek_v4/fused_norm_rope_v2.cuh"],
+        cuda_wrappers=cuda_wrappers,
+    )
+
+
+@cache_once
+def _jit_compress_module(
+    head_dim: int,
+    dtype_buffer: torch.dtype,
+    dtype_in: torch.dtype,
+    dtype_out: torch.dtype,
+    ratio: Literal[4, 128],
+) -> Module:
+    args = make_cpp_args(
+        head_dim, dtype_buffer, dtype_in, dtype_out, is_arch_support_pdl()
+    )
+    kernel_class = f"FlashCompress{ratio}Kernel<{args}>"
+    return load_jit(
+        make_name(f"compress_{ratio}_v2"),
+        *args,
+        cuda_files=[f"deepseek_v4/c{ratio}_v2.cuh"],
+        cuda_wrappers=[
+            ("decode", f"{kernel_class}::run_decode"),
+            ("prefill", f"{kernel_class}::run_prefill"),
+        ],
+        extra_cuda_cflags=["-use_fast_math"],
+    )
+
+
+@cache_once
+def _jit_compress_128_online_module(head_dim: int) -> Module:
+    assert head_dim == 512
+    args = make_cpp_args(head_dim, is_arch_support_pdl())
+    kernel_class = f"FlashCompress128OnlineKernel<{args}>"
+    return load_jit(
+        make_name(f"compress_128_online_v2"),
+        *args,
+        cuda_files=["deepseek_v4/c128_online_v2.cuh"],
+        cuda_wrappers=[
+            ("decode", f"{kernel_class}::run_decode"),
+            ("prefill", f"{kernel_class}::run_prefill"),
+            ("plan_decode", "plan_compress_128_online_decode"),
+            ("plan_prefill", "plan_compress_128_online_prefill"),
+        ],
+        extra_cuda_cflags=["-use_fast_math"],
+    )
+
+
+@cache_once
+def _jit_compress_plan_module() -> Module:
+    return load_jit(
+        make_name(f"compress_plan"),
+        cuda_files=[f"deepseek_v4/c_plan.cuh"],
+        cuda_wrappers=[
+            ("plan_prefill", "plan_compress_prefill"),
+            ("plan_decode", "plan_compress_decode"),
+            ("plan_prefill_legacy", "plan_compress_prefill_legacy"),
+            ("plan_decode_legacy", "plan_compress_decode_legacy"),
+        ],
+    )
+
+
+# ----------------------------------------------------------------------------
+# Plan tensor sizes (must match the C++ structs in compress.cuh).
+# ----------------------------------------------------------------------------
+_PREFILL_PLAN_BYTES = 24
+
+
+# ----------------------------------------------------------------------------
+# Plan dataclasses. The element at index 1 is the consumer for
+# `compress_fused_norm_rope_inplace` (which reads ragged_id / seq_len from a
+# 16-byte plan tensor --- both DecodePlan and CompressPlan satisfy that layout).
+# ----------------------------------------------------------------------------
+
+
+class CompressorDecodePlan(NamedTuple):
+    compress_ratio: int
+    plan_d: torch.Tensor  # [batch_size, 16] uint8 --- DecodePlan
+
+    def copy_(self, other) -> None:
+        assert isinstance(other, CompressorDecodePlan)
+        assert self.compress_ratio == other.compress_ratio
+        self.plan_d.copy_(other.plan_d)
+
+    @staticmethod
+    def generate(
+        compress_ratio: Literal[4, 128],
+        req_pool_indices: torch.Tensor,
+        req_to_token: torch.Tensor,
+        full_to_state: torch.Tensor,
+        seq_lens: torch.Tensor,
+        swa_page_size: int,
+        ring_size: int,
+    ) -> CompressorDecodePlan:
+        module = _jit_compress_plan_module()
+        plan_d = module.plan_decode(
+            req_pool_indices,
+            req_to_token,
+            full_to_state,
+            seq_lens,
+            int(compress_ratio),
+            int(swa_page_size),
+            int(ring_size),
+        )
+        return CompressorDecodePlan(compress_ratio, torch.from_dlpack(plan_d))
+
+    @staticmethod
+    def generate_legacy(
+        compress_ratio: Literal[4, 128],
+        req_pool_indices: torch.Tensor,
+        seq_lens: torch.Tensor,
+    ) -> CompressorDecodePlan:
+        module = _jit_compress_plan_module()
+        plan_d = module.plan_decode_legacy(req_pool_indices, seq_lens, compress_ratio)
+        return CompressorDecodePlan(compress_ratio, torch.from_dlpack(plan_d))
+
+    @staticmethod
+    def generate_online(
+        seq_lens: torch.Tensor,
+        req_pool_indices: torch.Tensor,
+        req_to_token: torch.Tensor,
+        state_slot_offset: int = 0,
+    ) -> CompressorDecodePlan:
+        batch_size = int(seq_lens.shape[0])
+        module = _jit_compress_128_online_module(512)
+        plan_d = torch.empty(
+            (batch_size, 16),
+            dtype=torch.uint8,
+            device=req_pool_indices.device,
+        )
+        module.plan_decode(
+            seq_lens,
+            req_pool_indices,
+            req_to_token,
+            plan_d,
+            int(state_slot_offset),
+        )
+        return CompressorDecodePlan(128, plan_d)
+
+    @property
+    def is_decode(self) -> bool:
+        return True
+
+
+class CompressorPrefillPlan(NamedTuple):
+    compress_ratio: int
+    plan_c: torch.Tensor  # [num_q_tokens, 16] uint8 --- CompressPlan
+    plan_w: torch.Tensor  # [num_q_tokens,  8] uint8 --- WritePlan
+    pin_buffer: Optional[torch.Tensor] = None  # keep alive
+
+    def copy_(self, other) -> None:
+        assert isinstance(other, CompressorPrefillPlan)
+        assert self.compress_ratio == other.compress_ratio
+        self.plan_c.copy_(other.plan_c)
+        self.plan_w.copy_(other.plan_w)
+
+    @staticmethod
+    def generate(
+        compress_ratio: Literal[4, 128],
+        req_pool_indices: torch.Tensor,
+        seq_lens: torch.Tensor,
+        extend_lens: torch.Tensor,
+        req_to_token: torch.Tensor,
+        full_to_state: torch.Tensor,
+        swa_page_size: int,
+        ring_size: int,
+        num_q_tokens: int,
+        use_cuda_graph: bool = False,
+    ) -> CompressorPrefillPlan:
+        is_gpu_input = seq_lens.device.type == "cuda"
+        pin_buffer = torch.empty(
+            0 if is_gpu_input else num_q_tokens * _PREFILL_PLAN_BYTES,
+            dtype=torch.uint8,
+            pin_memory=not is_gpu_input,
+        )
+        module = _jit_compress_plan_module()
+        plan_c, plan_w = module.plan_prefill(
+            req_pool_indices,
+            req_to_token,
+            full_to_state,
+            seq_lens,
+            extend_lens,
+            pin_buffer,
+            int(num_q_tokens),
+            int(compress_ratio),
+            int(swa_page_size),
+            int(ring_size),
+            bool(use_cuda_graph),
+        )
+        return CompressorPrefillPlan(
+            compress_ratio,
+            torch.from_dlpack(plan_c),
+            torch.from_dlpack(plan_w),
+            pin_buffer,
+        )
+
+    @staticmethod
+    def generate_legacy(
+        compress_ratio: Literal[4, 128],
+        req_pool_indices: torch.Tensor,
+        seq_lens: torch.Tensor,
+        extend_lens: torch.Tensor,
+        num_q_tokens: int,
+        device: torch.device,
+        use_cuda_graph: bool = False,
+    ) -> CompressorPrefillPlan:
+        pin_buffer = torch.empty(
+            num_q_tokens * _PREFILL_PLAN_BYTES,
+            dtype=torch.uint8,
+            pin_memory=True,
+        )
+        module = _jit_compress_plan_module()
+        plan_c, plan_w = module.plan_prefill_legacy(
+            req_pool_indices,
+            seq_lens,
+            extend_lens,
+            pin_buffer,
+            int(num_q_tokens),
+            int(compress_ratio),
+            bool(use_cuda_graph),
+        )
+        return CompressorPrefillPlan(
+            compress_ratio,
+            torch.from_dlpack(plan_c),
+            torch.from_dlpack(plan_w),
+            pin_buffer,
+        )
+
+    @staticmethod
+    def generate_online(
+        seq_lens: torch.Tensor,
+        extend_lens: torch.Tensor,
+        req_pool_indices: torch.Tensor,
+        req_to_token: torch.Tensor,
+        num_q_tokens: int,
+        use_cuda_graph: bool = False,
+        state_slot_offset: int = 0,
+    ) -> CompressorPrefillPlan:
+        seq_lens_cpu = seq_lens.detach().to(torch.int64).cpu()
+        extend_lens_cpu = extend_lens.detach().to(torch.int64).cpu()
+        rid_i64 = req_pool_indices.to(torch.int64)
+        r2t_i32 = req_to_token.to(torch.int32)
+        pin_buffer = torch.empty(
+            (2, num_q_tokens, 16), dtype=torch.uint8, pin_memory=True
+        )
+        plan_c_pin, plan_w_pin = pin_buffer[0], pin_buffer[1]
+        device = req_pool_indices.device
+        plan_c_dev = torch.empty((num_q_tokens, 16), dtype=torch.uint8, device=device)
+        plan_w_dev = torch.empty((num_q_tokens, 16), dtype=torch.uint8, device=device)
+        module = _jit_compress_128_online_module(512)  # NOTE: only support dim=512
+        num_c, num_w = module.plan_prefill(
+            seq_lens_cpu,
+            extend_lens_cpu,
+            rid_i64,
+            r2t_i32,
+            plan_c_pin,
+            plan_w_pin,
+            plan_c_dev,
+            plan_w_dev,
+            int(state_slot_offset),
+            bool(use_cuda_graph),
+        )
+        return CompressorPrefillPlan(
+            128,
+            plan_c_dev[: int(num_c)],
+            plan_w_dev[: int(num_w)],
+            pin_buffer,
+        )
+
+    @property
+    def is_decode(self) -> bool:
+        return False
+
+
+def compress_forward(
+    kv_score_buffer: torch.Tensor,
+    kv_score_input: torch.Tensor,
+    ape: torch.Tensor,
+    plan: Union[CompressorDecodePlan, CompressorPrefillPlan],
+    *,
+    head_dim: int,
+    compress_ratio: Literal[4, 128],
+    out: Optional[torch.Tensor] = None,
+    is_online: bool = False,
+) -> torch.Tensor:
+    if out is None:
+        num_q_tokens = plan[1].shape[0]  # NOTE: decode = bs, prefill = dynamic
+        out = kv_score_input.new_empty((num_q_tokens, head_dim))
+    assert plan.compress_ratio == compress_ratio
+    if is_online:
+        assert compress_ratio == 128 and head_dim == 512
+        module = _jit_compress_128_online_module(512)
+    else:
+        dtype_in, dtype_out = kv_score_input.dtype, out.dtype
+        module = _jit_compress_module(
+            head_dim, kv_score_buffer.dtype, dtype_in, dtype_out, compress_ratio
+        )
+    fn = module.decode if plan.is_decode else module.prefill
+    fn(kv_score_buffer, kv_score_input, out, ape, *plan[1:3])
+    return out
+
+
+def compress_norm_rope_store(
+    kv: torch.Tensor,
+    plan: Union[CompressorDecodePlan, CompressorPrefillPlan],
+    *,
+    norm_weight: torch.Tensor,
+    norm_eps: float,
+    freq_cis: torch.Tensor,
+    out_loc: torch.Tensor,
+    kvcache: torch.Tensor,
+    page_size: int,
+    use_fp4: bool = False,
+    bf16_store: bool = False,
+) -> None:
+    if use_fp4:
+        assert kv.shape[-1] == 128
+    freq_cis = torch.view_as_real(freq_cis).flatten(-2)
+    module = _jit_compress_norm_rope_module(
+        kv.dtype, kv.shape[-1], freq_cis.shape[-1], page_size, bf16_store
+    )
+    fn = module.forward_fp4 if use_fp4 else module.forward
+    fn(
+        kv,
+        plan[1],
+        norm_weight,
+        norm_eps,
+        freq_cis,
+        out_loc,
+        kvcache,
+        plan.is_decode,
+        plan.compress_ratio,
+    )

python/sglang/jit_kernel/dsv4/compress_old.py

@@ -0,0 +1,308 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Literal, NamedTuple, Optional, Union
+
+import torch
+
+from sglang.jit_kernel.utils import (
+    cache_once,
+    is_arch_support_pdl,
+    load_jit,
+    make_cpp_args,
+)
+from sglang.srt.environ import envs
+
+from .utils import make_name
+
+if TYPE_CHECKING:
+    from tvm_ffi.module import Module
+
+
+@cache_once
+def _jit_common_module() -> Module:
+    return load_jit(
+        make_name("common"),
+        cuda_files=["deepseek_v4/common.cuh"],
+        cuda_wrappers=[("plan_compress_prefill", "plan_compress_prefill")],
+    )
+
+
+@cache_once
+def _jit_compress_128_online_plan_module() -> Module:
+    """Host-side plan generator for online compress 128 (no template args)."""
+    return load_jit(
+        make_name("compress_128_online_plan"),
+        cuda_files=["deepseek_v4/c128_online.cuh"],
+        cuda_wrappers=[
+            ("plan_compress_online_prefill", "plan_compress_online_prefill"),
+        ],
+    )
+
+
+@cache_once
+def _jit_compress_128_online_module(head_dim: int) -> Module:
+    """Online compress 128 kernel: ring_size=1, per-index (max, sum, kv) state."""
+    args = make_cpp_args(head_dim, is_arch_support_pdl())
+    kernel_class = f"FlashCompress128OnlineKernel<{args}>"
+    return load_jit(
+        make_name("compress_128_online"),
+        *args,
+        cuda_files=["deepseek_v4/c128_online.cuh"],
+        cuda_wrappers=[
+            ("decode", f"{kernel_class}::run_decode"),
+            ("prefill", f"{kernel_class}::run_prefill"),
+        ],
+        extra_cuda_cflags=["-use_fast_math"],
+    )
+
+
+@cache_once
+def _jit_norm_rope_module(
+    dtype: torch.dtype,
+    head_dim: int,
+    rope_dim: int,
+) -> Module:
+    args = make_cpp_args(dtype, head_dim, rope_dim, is_arch_support_pdl())
+    return load_jit(
+        make_name("fused_norm_rope"),
+        *args,
+        cuda_files=["deepseek_v4/fused_norm_rope.cuh"],
+        cuda_wrappers=[
+            ("forward", f"FusedNormRopeKernel<{args}>::forward"),
+        ],
+    )
+
+
+@cache_once
+def _jit_compress_module(
+    head_dim: int,
+    dtype_in: torch.dtype,
+    dtype_out: torch.dtype,
+    ratio: Literal[4, 128],
+) -> Module:
+    args = make_cpp_args(head_dim, dtype_in, dtype_out, is_arch_support_pdl())
+    kernel_class = f"FlashCompress{ratio}Kernel<{args}>"
+    return load_jit(
+        make_name(f"compress_{ratio}"),
+        *args,
+        cuda_files=[f"deepseek_v4/c{ratio}.cuh"],
+        cuda_wrappers=[
+            ("decode", f"{kernel_class}::run_decode"),
+            ("prefill", f"{kernel_class}::run_prefill"),
+        ],
+        extra_cuda_cflags=["-use_fast_math"],
+    )
+
+
+class CompressorPrefillPlan(NamedTuple):
+    compress_ratio: int
+    compress_plan: torch.Tensor
+    write_plan: torch.Tensor
+
+    def copy_(self, other: CompressorPrefillPlan) -> None:
+        assert self.compress_ratio == other.compress_ratio
+        self.compress_plan.copy_(other.compress_plan)
+        self.write_plan.copy_(other.write_plan)
+
+    @staticmethod
+    def generate(
+        compress_ratio: Literal[4, 128],
+        num_q_tokens: int,
+        seq_lens: torch.Tensor,
+        extend_lens: torch.Tensor,
+        device: torch.device,
+        use_cuda_graph: bool = False,
+    ) -> CompressorPrefillPlan:
+        from sglang.srt.environ import envs
+
+        # Online c128 keeps the same NamedTuple shape (compress_plan, write_plan)
+        # so call sites that splat `*plan[1:]` continue to work, but the C++
+        # plan struct semantics differ (last-token coords + window_len).
+        if compress_ratio == 128 and envs.SGLANG_OPT_USE_ONLINE_COMPRESS.get():
+            return CompressorPrefillPlan._generate_online(
+                num_q_tokens=num_q_tokens,
+                seq_lens=seq_lens,
+                extend_lens=extend_lens,
+                device=device,
+                use_cuda_graph=use_cuda_graph,
+            )
+        assert seq_lens.device == extend_lens.device
+        seq_lens = seq_lens.to(torch.int64)
+        extend_lens = extend_lens.to(torch.int64)
+        plan_tensor = torch.empty(
+            (2, num_q_tokens, 16),
+            dtype=torch.uint8,
+            device=seq_lens.device,
+            pin_memory=seq_lens.is_cpu,
+        )
+        module = _jit_common_module()
+        is_overlap = compress_ratio == 4
+        plan_lens = module.plan_compress_prefill(
+            extend_lens,
+            seq_lens,
+            plan_tensor[0],
+            plan_tensor[1],
+            compress_ratio,
+            is_overlap,
+            use_cuda_graph,
+        )
+        return CompressorPrefillPlan(
+            compress_ratio,
+            plan_tensor[0, : plan_lens[0]].to(device, non_blocking=True),
+            plan_tensor[1, : plan_lens[1]].to(device, non_blocking=True),
+        )
+
+    @staticmethod
+    def _generate_online(
+        num_q_tokens: int,
+        seq_lens: torch.Tensor,
+        extend_lens: torch.Tensor,
+        device: torch.device,
+        use_cuda_graph: bool,
+    ) -> CompressorPrefillPlan:
+        # Online plan host-side path: only CPU/cuda-host implemented today.
+        # Move inputs to CPU pinned memory then bounce the result to device.
+        seq_lens_cpu = seq_lens.detach().to(torch.int64).cpu()
+        extend_lens_cpu = extend_lens.detach().to(torch.int64).cpu()
+        plan_tensor = torch.empty(
+            (2, num_q_tokens, 16),
+            dtype=torch.uint8,
+            device="cpu",
+            pin_memory=True,
+        )
+        module = _jit_compress_128_online_plan_module()
+        plan_lens = module.plan_compress_online_prefill(
+            extend_lens_cpu,
+            seq_lens_cpu,
+            plan_tensor[0],
+            plan_tensor[1],
+            use_cuda_graph,
+        )
+        return CompressorPrefillPlan(
+            128,
+            plan_tensor[0, : plan_lens[0]].to(device, non_blocking=True),
+            plan_tensor[1, : plan_lens[1]].to(device, non_blocking=True),
+        )
+
+    @property
+    def is_decode(self) -> bool:
+        return False
+
+
+class CompressorDecodePlan(NamedTuple):
+    compress_ratio: int
+    seq_lens: torch.Tensor
+
+    def copy_(self, other: CompressorDecodePlan) -> None:
+        assert self.compress_ratio == other.compress_ratio
+        self.seq_lens.copy_(other.seq_lens)
+
+    @property
+    def is_decode(self) -> bool:
+        return True
+
+
+def compress_plan(
+    compress_ratio: Literal[4, 128],
+    num_q_tokens: int,
+    seq_lens: torch.Tensor,
+    extend_lens: Optional[torch.Tensor],
+    device: torch.device,
+) -> Union[CompressorDecodePlan, CompressorPrefillPlan]:
+    if extend_lens is not None:
+        return CompressorPrefillPlan.generate(
+            compress_ratio,
+            num_q_tokens,
+            seq_lens,
+            extend_lens,
+            device,
+        )
+    else:
+        assert num_q_tokens == len(seq_lens)
+        seq_lens = seq_lens.to(device, non_blocking=True)
+        return CompressorDecodePlan(compress_ratio, seq_lens)
+
+
+def compress_forward(
+    kv_score_buffer: torch.Tensor,
+    kv_score_input: torch.Tensor,
+    ape: torch.Tensor,
+    indices: torch.Tensor,
+    plan: Union[CompressorDecodePlan, CompressorPrefillPlan, None] = None,
+    extra_data: Optional[torch.Tensor] = None,
+    *,
+    head_dim: int,
+    compress_ratio: Literal[4, 128],
+    out: Optional[torch.Tensor] = None,
+    seq_lens: Optional[torch.Tensor] = None,
+    extend_lens: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    assert head_dim % 128 == 0
+    num_q_tokens = kv_score_input.shape[0]
+    if out is None:
+        out = kv_score_input.new_empty((num_q_tokens, head_dim))
+    if plan is None:
+        assert seq_lens is not None
+        plan = compress_plan(
+            compress_ratio,
+            num_q_tokens,
+            seq_lens,
+            extend_lens,
+            kv_score_input.device,
+        )
+    assert plan.compress_ratio == compress_ratio, "Mismatched compress ratio in plan!"
+    # Online c128: separate JIT module, fp32 state, no compile-time dtypes.
+    if compress_ratio == 128 and envs.SGLANG_OPT_USE_ONLINE_COMPRESS.get():
+        online_module = _jit_compress_128_online_module(head_dim=head_dim)
+        F = online_module.decode if plan.is_decode else online_module.prefill
+        F(kv_score_buffer, kv_score_input, out, ape, indices, *plan[1:], extra_data)
+        return out
+    module = _jit_compress_module(
+        head_dim,
+        kv_score_input.dtype,
+        out.dtype,
+        compress_ratio,
+    )
+    F = module.decode if plan.is_decode else module.prefill
+    F(kv_score_buffer, kv_score_input, out, ape, indices, *plan[1:], extra_data)
+    return out
+
+
+def compress_fused_norm_rope_inplace(
+    kv: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float,
+    freq_cis: torch.Tensor,
+    plan: Union[CompressorDecodePlan, CompressorPrefillPlan],
+) -> None:
+    freq_cis = torch.view_as_real(freq_cis).flatten(-2)
+    module = _jit_norm_rope_module(kv.dtype, kv.shape[-1], freq_cis.shape[-1])
+    module.forward(
+        kv,
+        weight,
+        plan[1],
+        freq_cis,
+        int(plan.is_decode),
+        eps,
+        plan.compress_ratio,
+    )
+
+
+def fused_norm_rope_inplace(
+    kv: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float,
+    freq_cis: torch.Tensor,
+    positions: torch.Tensor,
+) -> None:
+    freq_cis = torch.view_as_real(freq_cis).flatten(-2)
+    module = _jit_norm_rope_module(kv.dtype, kv.shape[-1], freq_cis.shape[-1])
+    module.forward(
+        kv,
+        weight,
+        positions,
+        freq_cis,
+        2,
+        eps,
+        0,
+    )

python/sglang/jit_kernel/dsv4/elementwise.py

@@ -0,0 +1,262 @@
+from typing import Optional, Tuple
+
+import torch
+
+from sglang.jit_kernel.utils import (
+    cache_once,
+    is_arch_support_pdl,
+    load_jit,
+    make_cpp_args,
+)
+from sglang.srt.utils import is_hip, is_xpu
+
+from .utils import make_name
+
+_is_hip = is_hip()
+_is_xpu = is_xpu()
+
+
+@cache_once
+def _jit_fused_rope_module():
+    args = make_cpp_args(is_arch_support_pdl())
+    return load_jit(
+        make_name("fused_rope"),
+        *args,
+        cuda_files=["deepseek_v4/rope.cuh"],
+        cuda_wrappers=[("forward", f"FusedQKRopeKernel<{args}>::forward")],
+    )
+
+
+@cache_once
+def _jit_main_q_norm_rope_module(
+    dtype: torch.dtype,
+    head_dim: int,
+    rope_dim: int,
+):
+    """Main MLA path Q kernel: rmsnorm-self + RoPE, warp per (token, head)."""
+    args = make_cpp_args(dtype, head_dim, rope_dim, is_arch_support_pdl())
+    return load_jit(
+        make_name("main_q_norm_rope"),
+        *args,
+        cuda_files=["deepseek_v4/main_norm_rope.cuh"],
+        cuda_wrappers=[
+            ("forward", f"FusedQNormRopeKernel<{args}>::forward"),
+        ],
+    )
+
+
+@cache_once
+def _jit_main_k_norm_rope_flashmla_module(
+    dtype: torch.dtype,
+    head_dim: int,
+    rope_dim: int,
+    page_size: int,
+):
+    """Main MLA path K kernel: rmsnorm + RoPE + write to FlashMLA paged cache."""
+    args = make_cpp_args(dtype, head_dim, rope_dim, page_size, is_arch_support_pdl())
+    return load_jit(
+        make_name("main_k_norm_rope_flashmla"),
+        *args,
+        cuda_files=["deepseek_v4/main_norm_rope.cuh"],
+        cuda_wrappers=[
+            ("forward", f"FusedKNormRopeFlashMLAKernel<{args}>::forward"),
+        ],
+    )
+
+
+@cache_once
+def _jit_main_q_indexer_rope_hadamard_quant_module(dtype: torch.dtype):
+    """C4 indexer Q kernel: RoPE + 128-pt Hadamard + fp8 act-quant"""
+    args = make_cpp_args(dtype, is_arch_support_pdl())
+    return load_jit(
+        make_name("main_q_indexer_rope_hadamard_quant"),
+        *args,
+        cuda_files=["deepseek_v4/main_norm_rope.cuh"],
+        cuda_wrappers=[
+            ("forward", f"FusedQIndexerRopeHadamardQuantKernel<{args}>::forward"),
+        ],
+    )
+
+
+# V3.2 lays q out as [rope | nope] (V4 is [nope | rope]) -> kRopeFirst=true, and
+# drops the Hadamard rotation (kHadamard=false).
+@cache_once
+def _jit_main_q_indexer_rope_first_quant_module(dtype: torch.dtype):
+    args = make_cpp_args(dtype, is_arch_support_pdl(), True, False)
+    return load_jit(
+        make_name("main_q_indexer_rope_first_quant"),
+        *args,
+        cuda_files=["deepseek_v4/main_norm_rope.cuh"],
+        cuda_wrappers=[
+            ("forward", f"FusedQIndexerRopeHadamardQuantKernel<{args}>::forward"),
+        ],
+    )
+
+
+@cache_once
+def _jit_main_q_indexer_rope_hadamard_fp4_quant_module(dtype: torch.dtype):
+    args = make_cpp_args(dtype, is_arch_support_pdl())
+    return load_jit(
+        make_name("main_q_indexer_rope_hadamard_fp4_quant"),
+        *args,
+        cuda_files=["deepseek_v4/main_norm_rope.cuh"],
+        cuda_wrappers=[
+            ("forward", f"FusedQIndexerRopeHadamardFp4QuantKernel<{args}>::forward"),
+        ],
+    )
+
+
+def fused_rope_inplace(
+    q: torch.Tensor,
+    k: Optional[torch.Tensor],
+    freqs_cis: torch.Tensor,
+    positions: torch.Tensor,
+    inverse: bool = False,
+) -> None:
+    """Apply rotary embeddings to both Q and K in a single fused CUDA kernel.
+
+    Args:
+        q: [batch_size, num_q_heads, rope_dim] bfloat16
+        k: [batch_size, num_k_heads, rope_dim] bfloat16 or None
+        freqs_cis: [max_seq_len, rope_dim // 2] complex64 (full table)
+        positions: [batch_size] int32 or int64, indices into freqs_cis
+        inverse: if True, apply inverse rotation (conjugate freqs)
+    """
+    if _is_hip or _is_xpu:
+        from sglang.srt.layers.deepseek_v4_rope import apply_rotary_emb_triton
+
+        apply_rotary_emb_triton(q, freqs_cis, positions=positions, inverse=inverse)
+        if k is not None:
+            apply_rotary_emb_triton(k, freqs_cis, positions=positions, inverse=inverse)
+        return
+
+    freqs_real = torch.view_as_real(freqs_cis).flatten(-2).contiguous()
+    module = _jit_fused_rope_module()
+    module.forward(q, k, freqs_real, positions, inverse)
+
+
+def fused_q_norm_rope(
+    q_input: torch.Tensor,
+    q_output: torch.Tensor,
+    eps: float,
+    freqs_cis: torch.Tensor,
+    positions: torch.Tensor,
+) -> None:
+    freqs_real = torch.view_as_real(freqs_cis).flatten(-2)
+    head_dim = q_input.shape[-1]
+    rope_dim = freqs_real.shape[-1]
+    module = _jit_main_q_norm_rope_module(q_input.dtype, head_dim, rope_dim)
+    module.forward(q_input, q_output, freqs_real, positions, eps)
+
+
+def fused_q_indexer_rope_hadamard_quant(
+    q_input: torch.Tensor,
+    weight: torch.Tensor,
+    weight_scale: float,
+    freqs_cis: torch.Tensor,
+    positions: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    freqs_real = torch.view_as_real(freqs_cis).flatten(-2)
+    q_fp8 = torch.empty(q_input.shape, dtype=torch.float8_e4m3fn, device=q_input.device)
+    weights_out = torch.empty(
+        (*q_input.shape[:-1], 1), dtype=torch.float32, device=q_input.device
+    )
+    if _is_hip:
+        torch.ops.sgl_kernel.dsv4_fused_q_indexer_rope_hadamard_quant(
+            q_input,
+            q_fp8,
+            weight,
+            weights_out,
+            float(weight_scale),
+            freqs_real,
+            positions,
+        )
+    else:
+        module = _jit_main_q_indexer_rope_hadamard_quant_module(q_input.dtype)
+        module.forward(
+            q_input,
+            q_fp8,
+            weight,
+            weights_out,
+            float(weight_scale),
+            freqs_real,
+            positions,
+        )
+    return q_fp8, weights_out
+
+
+def fused_q_indexer_rope_first_quant(
+    q_input: torch.Tensor,
+    weight: torch.Tensor,
+    weight_scale: float,
+    cos_sin_cache: torch.Tensor,
+    positions: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """DeepSeek-V3.2 only. Indexer Q: RoPE on the leading dims + fp8 act-quant. CUDA only."""
+    q_fp8 = torch.empty(q_input.shape, dtype=torch.float8_e4m3fn, device=q_input.device)
+    weights_out = torch.empty(
+        (*q_input.shape[:-1], 1), dtype=torch.float32, device=q_input.device
+    )
+    module = _jit_main_q_indexer_rope_first_quant_module(q_input.dtype)
+    module.forward(
+        q_input,
+        q_fp8,
+        weight,
+        weights_out,
+        float(weight_scale),
+        cos_sin_cache,
+        positions,
+    )
+    return q_fp8, weights_out
+
+
+def fused_q_indexer_rope_hadamard_fp4_quant(
+    q_input: torch.Tensor,
+    weight: torch.Tensor,
+    weight_scale: float,
+    freqs_cis: torch.Tensor,
+    positions: torch.Tensor,
+) -> Tuple[Tuple[torch.Tensor, torch.Tensor], torch.Tensor]:
+    if _is_hip:
+        raise RuntimeError("DeepSeek V4 FP4 indexer requires the CUDA fused Q path.")
+    freqs_real = torch.view_as_real(freqs_cis).flatten(-2)
+    q_fp4 = torch.empty(
+        (*q_input.shape[:-1], q_input.shape[-1] // 2),
+        dtype=torch.int8,
+        device=q_input.device,
+    )
+    q_sf = torch.empty(q_input.shape[:-1], dtype=torch.int32, device=q_input.device)
+    weights_out = torch.empty(
+        (*q_input.shape[:-1], 1), dtype=torch.float32, device=q_input.device
+    )
+    module = _jit_main_q_indexer_rope_hadamard_fp4_quant_module(q_input.dtype)
+    module.forward(
+        q_input,
+        q_fp4,
+        q_sf,
+        weight,
+        weights_out,
+        float(weight_scale),
+        freqs_real,
+        positions,
+    )
+    return (q_fp4, q_sf), weights_out
+
+
+def fused_k_norm_rope_flashmla(
+    kv: torch.Tensor,
+    kv_weight: torch.Tensor,
+    eps: float,
+    freqs_cis: torch.Tensor,
+    positions: torch.Tensor,
+    out_loc: torch.Tensor,
+    kvcache: torch.Tensor,
+    page_size: int,
+) -> None:
+    freqs_real = torch.view_as_real(freqs_cis).flatten(-2)
+    head_dim = kv.shape[-1]
+    rope_dim = freqs_real.shape[-1]
+    module = _jit_main_k_norm_rope_flashmla_module(
+        kv.dtype, head_dim, rope_dim, page_size
+    )
+    module.forward(kv, kv_weight, freqs_real, positions, out_loc, kvcache, eps)

python/sglang/jit_kernel/dsv4/fp8_wo_a.py

@@ -0,0 +1,93 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Tuple
+
+import torch
+
+from sglang.jit_kernel.utils import (
+    cache_once,
+    is_arch_support_pdl,
+    load_jit,
+    make_cpp_args,
+)
+from sglang.kernel_api_logging import debug_kernel_api
+from sglang.srt.utils.custom_op import register_custom_op
+
+from .utils import make_name
+
+if TYPE_CHECKING:
+    from tvm_ffi.module import Module
+
+_GROUP_SIZE = 128
+
+
+@cache_once
+def _jit_module(in_dtype: torch.dtype, use_pdl: bool) -> Module:
+    args = make_cpp_args(in_dtype, use_pdl)
+    return load_jit(
+        make_name("fp8_wo_a_group_major_quant_ue8m0"),
+        *args,
+        cuda_files=["deepseek_v4/fp8_wo_a_group_major_quant.cuh"],
+        cuda_wrappers=[
+            (
+                "fp8_wo_a_group_major_quant_ue8m0",
+                f"FP8WoAGroupMajorQuantUE8M0Kernel<{args}>::run",
+            )
+        ],
+        # Match the AOT/JIT v2 quant path's fast-math build so FP8 rounding stays
+        # bit-identical for the DSV4 wo_a replacement.
+        extra_cuda_cflags=["--use_fast_math"],
+    )
+
+
+@register_custom_op(
+    op_name="fp8_wo_a_group_major_quant_ue8m0",
+    mutates_args=["output_q", "output_s"],
+)
+def _fp8_wo_a_group_major_quant_ue8m0_custom_op(
+    input: torch.Tensor,
+    output_q: torch.Tensor,
+    output_s: torch.Tensor,
+) -> None:
+    """Opaque custom-op boundary for the DeepSeek-V4 wo_a quant JIT kernel."""
+    assert input.dtype in (torch.bfloat16, torch.float16)
+
+    module = _jit_module(input.dtype, is_arch_support_pdl())
+    module.fp8_wo_a_group_major_quant_ue8m0(input, output_q, output_s)
+
+
+@debug_kernel_api
+def fp8_wo_a_group_major_quant_ue8m0(
+    input: torch.Tensor,
+    output_q: torch.Tensor,
+    output_s: torch.Tensor,
+) -> None:
+    _fp8_wo_a_group_major_quant_ue8m0_custom_op(input, output_q, output_s)
+
+
+def sglang_per_token_group_quant_fp8_dsv4_wo_a(
+    x: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Quantize DSV4 wo_a activations for DeepGEMM fp8_einsum.
+
+    The input is a [T, G, D] bf16/fp16 tensor whose hidden dimension is
+    contiguous. The output codes are contiguous [T, G, D] fp8 values. The scale
+    tensor is returned as logical [T, G, D/128] fp32 UE8M0 values backed by
+    contiguous [G, T, D/128] storage, so each group/head [T, S] panel is
+    contiguous for the DeepGEMM recipe=(1, 1, 128) consumer. Group size is fixed
+    to 128 and the absmax floor is fixed to 1e-10.
+    """
+    num_tokens, num_groups, hidden = x.shape
+    hidden_groups = hidden // _GROUP_SIZE
+    x_q = torch.empty(x.shape, device=x.device, dtype=torch.float8_e4m3fn)
+    x_s_storage = torch.empty(
+        (num_groups, num_tokens, hidden_groups),
+        device=x.device,
+        dtype=torch.float32,
+    )
+
+    if x.numel() > 0:
+        fp8_wo_a_group_major_quant_ue8m0(x, x_q, x_s_storage)
+
+    # DeepGEMM fp8_einsum consumes each group/head [T, S] scale panel contiguously.
+    return x_q, x_s_storage.transpose(0, 1)

python/sglang/jit_kernel/dsv4/gemm.py

@@ -0,0 +1,24 @@
+import torch
+
+from sglang.srt.environ import envs
+from sglang.srt.layers import deep_gemm_wrapper
+from sglang.srt.utils import get_bool_env_var, is_hip
+
+_is_hip = is_hip()
+_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip
+
+if _use_aiter:
+    from aiter.tuned_gemm import tgemm
+
+_linear_bf16_fp32_algo = envs.SGLANG_OPT_BF16_FP32_GEMM_ALGO.get()
+
+
+def linear_bf16_fp32(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+    if _use_aiter:
+        return tgemm.mm(x, y, otype=x.dtype).float()
+    elif _linear_bf16_fp32_algo == "deep_gemm":
+        z = torch.empty(x.size(0), y.size(0), dtype=torch.float32, device=x.device)
+        deep_gemm_wrapper.gemm_nt_bf16bf16f32(x, y, z)
+        return z
+    else:
+        return torch.mm(x, y.t(), out_dtype=torch.float32)

python/sglang/jit_kernel/dsv4/moe.py

@@ -0,0 +1,229 @@
+from typing import Optional, Tuple
+
+import torch
+
+from sglang.jit_kernel.utils import (
+    cache_once,
+    is_arch_support_pdl,
+    is_hip_runtime,
+    load_jit,
+    make_cpp_args,
+)
+
+from .utils import make_name
+
+
+@cache_once
+def _jit_mask_topk_module():
+    return load_jit(
+        make_name("mask_topk"),
+        cuda_files=["deepseek_v4/hash_topk.cuh"],
+        cuda_wrappers=[("run", "MaskKernel::run")],
+    )
+
+
+@cache_once
+def _jit_hash_topk_module():
+    args = make_cpp_args("act_sqrt_softplus", is_arch_support_pdl())
+    return load_jit(
+        make_name("hash_topk"),
+        *args,
+        cuda_files=["deepseek_v4/hash_topk.cuh"],
+        cuda_wrappers=[("hash_topk", f"HashTopKKernel<{args}>::run")],
+    )
+
+
+@cache_once
+def _jit_mega_moe_pre_dispatch_module(quant_group_size: int):
+    args = make_cpp_args(quant_group_size, is_arch_support_pdl())
+    return load_jit(
+        make_name("mega_moe_pre_dispatch"),
+        *args,
+        cuda_files=["deepseek_v4/mega_moe_pre_dispatch.cuh"],
+        cuda_wrappers=[("run", f"MegaMoEPreDispatchKernel<{args}>::run")],
+    )
+
+
+@cache_once
+def _jit_silu_mul_quant_varlen_module(
+    quant_group_size: int,
+    scale_ue8m0: bool,
+    swizzle: bool,
+    apply_swiglu_limit: bool,
+):
+    args = make_cpp_args(
+        quant_group_size,
+        scale_ue8m0,
+        swizzle,
+        is_arch_support_pdl(),
+        apply_swiglu_limit,
+    )
+    return load_jit(
+        make_name("silu_mul_quant_varlen"),
+        *args,
+        cuda_files=["deepseek_v4/silu_and_mul_masked_post_quant.cuh"],
+        cuda_wrappers=[("run", f"SiluAndMulMaskedPostQuantKernel<{args}>::run")],
+        extra_cuda_cflags=["-use_fast_math"],
+    )
+
+
+@cache_once
+def _jit_silu_mul_quant_contig_module(
+    quant_group_size: int,
+    scale_ue8m0: bool,
+    swizzle: bool,
+    apply_swiglu_limit: bool,
+):
+    args = make_cpp_args(
+        quant_group_size,
+        scale_ue8m0,
+        swizzle,
+        is_arch_support_pdl(),
+        apply_swiglu_limit,
+    )
+    return load_jit(
+        make_name("silu_mul_quant_contig"),
+        *args,
+        cuda_files=["deepseek_v4/silu_and_mul_masked_post_quant.cuh"],
+        cuda_wrappers=[("run", f"SiluAndMulContigPostQuantKernel<{args}>::run")],
+        extra_cuda_cflags=["-use_fast_math"],
+    )
+
+
+@cache_once
+def _jit_silu_and_mul_clamp_module(dtype: torch.dtype):
+    args = make_cpp_args(dtype, is_arch_support_pdl())
+    return load_jit(
+        make_name("silu_and_mul_clamp"),
+        *args,
+        cuda_files=["deepseek_v4/silu_and_mul_masked_post_quant.cuh"],
+        cuda_wrappers=[("run", f"SiluAndMulClampKernel<{args}>::run")],
+        extra_cuda_cflags=["-use_fast_math"],
+    )
+
+
+def mask_topk_ids(topk_ids: torch.Tensor, num_token_non_padded: torch.Tensor):
+    return _jit_mask_topk_module().run(topk_ids, num_token_non_padded)
+
+
+def hash_topk(
+    router_logits: torch.Tensor,
+    input_ids: torch.Tensor,
+    tid2eid: torch.Tensor,
+    num_fused_shared_experts: int = 0,
+    routed_scaling_factor: float = 1.0,
+    scoring_func: str = "sqrtsoftplus",
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    assert scoring_func == "sqrtsoftplus"
+    if is_hip_runtime():
+        from sglang.jit_kernel.triton.hash_topk import hash_topk_triton
+
+        return hash_topk_triton(
+            router_logits,
+            input_ids,
+            tid2eid,
+            num_fused_shared_experts,
+            routed_scaling_factor,
+            scoring_func,
+        )
+    else:
+        num_tokens = router_logits.size(0)
+        topk_routed = tid2eid.size(1)
+        topk_fused = topk_routed + num_fused_shared_experts
+        topk_ids = torch.empty(
+            (num_tokens, topk_fused), dtype=torch.int32, device=router_logits.device
+        )
+        topk_weights = torch.empty(
+            (num_tokens, topk_fused), dtype=torch.float32, device=router_logits.device
+        )
+        module = _jit_hash_topk_module()
+        module.hash_topk(
+            router_logits,
+            input_ids,
+            tid2eid,
+            topk_weights,
+            topk_ids,
+            routed_scaling_factor,
+        )
+        return topk_weights, topk_ids
+
+
+def mega_moe_pre_dispatch(
+    x: torch.Tensor,
+    topk_idx: torch.Tensor,
+    topk_weights: torch.Tensor,
+    buf_x: torch.Tensor,
+    buf_x_sf: torch.Tensor,
+    buf_topk_idx: torch.Tensor,
+    buf_topk_weights: torch.Tensor,
+    quant_group_size: int = 32,
+) -> None:
+    module = _jit_mega_moe_pre_dispatch_module(quant_group_size)
+    module.run(
+        x,
+        topk_idx,
+        topk_weights,
+        buf_x,
+        buf_x_sf,
+        buf_topk_idx,
+        buf_topk_weights,
+    )
+
+
+def silu_and_mul_clamp(
+    input: torch.Tensor,
+    output: torch.Tensor,
+    swiglu_limit: float,
+) -> None:
+    module = _jit_silu_and_mul_clamp_module(input.dtype)
+    module.run(input, output, float(swiglu_limit))
+
+
+def silu_and_mul_masked_post_quant(
+    input: torch.Tensor,
+    output: torch.Tensor,
+    output_scale: torch.Tensor,
+    quant_group_size: int,
+    masked_m: torch.Tensor,
+    scale_ue8m0: bool = False,
+    topk: int = 8,
+    transposed: bool = False,
+    swiglu_limit: Optional[float] = None,
+    swizzle: bool = False,
+) -> None:
+    apply_swiglu_limit = swiglu_limit is not None
+    module = _jit_silu_mul_quant_varlen_module(
+        quant_group_size, scale_ue8m0, swizzle, apply_swiglu_limit
+    )
+    module.run(
+        input,
+        output,
+        output_scale,
+        masked_m,
+        topk,
+        transposed,
+        float(swiglu_limit) if apply_swiglu_limit else 0.0,
+    )
+
+
+def silu_and_mul_contig_post_quant(
+    input: torch.Tensor,
+    output: torch.Tensor,
+    output_scale: torch.Tensor,
+    quant_group_size: int,
+    scale_ue8m0: bool = False,
+    transposed: bool = False,
+    swiglu_limit: Optional[float] = None,
+    swizzle: bool = False,
+) -> None:
+    apply_swiglu_limit = swiglu_limit is not None
+    module = _jit_silu_mul_quant_contig_module(
+        quant_group_size, scale_ue8m0, swizzle, apply_swiglu_limit
+    )
+    module.run(
+        input,
+        output,
+        output_scale,
+        transposed,
+        float(swiglu_limit) if apply_swiglu_limit else 0.0,
+    )

python/sglang/jit_kernel/dsv4/online_c128_mtp.py

@@ -0,0 +1,262 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, List, Optional
+
+import torch
+
+from sglang.jit_kernel.dsv4.utils import make_name
+from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args
+from sglang.srt.environ import envs
+
+if TYPE_CHECKING:
+    from tvm_ffi.module import Module
+
+
+@cache_once
+def _jit_online_c128_mtp_module(
+    head_dim: int, seq_dtype: torch.dtype, req_dtype: torch.dtype
+) -> Module:
+    args = make_cpp_args(head_dim, seq_dtype, req_dtype)
+    return load_jit(
+        make_name(f"online_c128_mtp_{head_dim}"),
+        *args,
+        cuda_files=["deepseek_v4/online_c128_mtp.cuh"],
+        cuda_wrappers=[
+            ("write_prefix_states", f"OnlineC128MTPWritePrefixKernel<{args}>::run"),
+            ("mark_pending", f"OnlineC128MTPMarkPendingKernel<{args}>::run"),
+            ("commit_pending", f"OnlineC128MTPCommitPendingKernel<{args}>::run"),
+        ],
+        extra_cuda_cflags=["-use_fast_math"],
+    )
+
+
+@dataclass
+class _OnlineC128LayerRuntime:
+    head_dim: int
+    main_state: torch.Tensor
+    state_slot_offset: int
+
+
+@dataclass
+class _OnlineC128VerifyContext:
+    req_pool_indices: torch.Tensor
+    seq_lens: torch.Tensor
+
+
+class OnlineC128MTPController:
+    def __init__(self, backend: Any):
+        self.backend = backend
+        self._verify_ctx: Optional[_OnlineC128VerifyContext] = None
+        self._layer_runtimes: Optional[List[_OnlineC128LayerRuntime]] = None
+
+    def enabled(self) -> bool:
+        return (
+            envs.SGLANG_OPT_USE_ONLINE_COMPRESS.get()
+            and envs.SGLANG_EXPERIMENTAL_ONLINE_C128_MTP.get()
+            and self.backend.mtp_enabled
+        )
+
+    def state_slot_offset(self) -> int:
+        if not self.enabled():
+            return 0
+        return self.backend.token_to_kv_pool.get_online_c128_mtp_state_slot_offset()
+
+    def begin_verify(
+        self,
+        req_pool_indices: torch.Tensor,
+        seq_lens: torch.Tensor,
+    ) -> None:
+        if not self.enabled():
+            self.clear()
+            return
+
+        self._verify_ctx = _OnlineC128VerifyContext(
+            req_pool_indices=req_pool_indices.detach(),
+            seq_lens=seq_lens.detach(),
+        )
+        head_dim = self._head_dim()
+        if head_dim is None or self._num_verify_tokens() == 0:
+            return
+        token_to_kv_pool = self.backend.token_to_kv_pool
+        _jit_online_c128_mtp_module(
+            head_dim, seq_lens.dtype, req_pool_indices.dtype
+        ).mark_pending(
+            seq_lens,
+            req_pool_indices,
+            token_to_kv_pool.get_online_c128_mtp_pending_seq_lens(),
+            min(seq_lens.shape[0], req_pool_indices.shape[0]),
+            token_to_kv_pool.get_online_c128_state_num_req_slots(),
+        )
+
+    def clear(self) -> None:
+        self._verify_ctx = None
+
+    def prepare_forward(
+        self,
+        logical_forward_mode,
+        req_pool_indices: torch.Tensor,
+        seq_lens: torch.Tensor,
+        *,
+        verify_bs: Optional[int] = None,
+    ) -> int:
+        if not self.enabled():
+            self.clear()
+            return 0
+        if logical_forward_mode is None or logical_forward_mode.is_idle():
+            self.clear()
+            return 0
+
+        active_req_pool_indices = req_pool_indices
+        active_seq_lens = seq_lens
+        if logical_forward_mode.is_target_verify():
+            if verify_bs is None:
+                verify_bs = req_pool_indices.shape[0]
+            active_req_pool_indices = req_pool_indices[:verify_bs]
+            active_seq_lens = seq_lens[:verify_bs]
+            if verify_bs == 0:
+                self.clear()
+                return 0
+
+        self.commit_pending(
+            req_pool_indices=active_req_pool_indices,
+            seq_lens=active_seq_lens,
+        )
+        if not logical_forward_mode.is_target_verify():
+            return 0
+
+        self.begin_verify(
+            req_pool_indices=active_req_pool_indices,
+            seq_lens=active_seq_lens,
+        )
+        return self.state_slot_offset()
+
+    def write_prefix_states(
+        self,
+        layer_id: int,
+        compressor: Any,
+        kv_score_input: torch.Tensor,
+        logical_forward_mode,
+    ) -> None:
+        if (
+            not self.enabled()
+            or logical_forward_mode is None
+            or not logical_forward_mode.is_target_verify()
+            or compressor.is_in_indexer
+            or compressor.ratio != 128
+            or kv_score_input.numel() == 0
+        ):
+            return
+
+        ctx = self._active_ctx()
+        num_verify_tokens = self._num_verify_tokens()
+        if ctx is None or num_verify_tokens == 0:
+            return
+
+        token_to_kv_pool = self.backend.token_to_kv_pool
+        head_dim = compressor.head_dim
+        state_pool = token_to_kv_pool.get_attention_compress_states(layer_id)
+        total_bs = kv_score_input.numel() // (num_verify_tokens * head_dim * 2)
+        layer_bs = min(ctx.seq_lens.shape[0], ctx.req_pool_indices.shape[0], total_bs)
+        if layer_bs <= 0:
+            return
+
+        _jit_online_c128_mtp_module(
+            head_dim, ctx.seq_lens.dtype, ctx.req_pool_indices.dtype
+        ).write_prefix_states(
+            kv_score_input,
+            ctx.seq_lens,
+            ctx.req_pool_indices,
+            self.backend.req_to_token,
+            compressor.ape.reshape(128, head_dim),
+            state_pool.kv_score_buffer.kv_score,
+            layer_bs,
+            num_verify_tokens,
+            state_pool.online_mtp_state_slot_offset,
+        )
+
+    def commit_pending(
+        self,
+        req_pool_indices: torch.Tensor,
+        seq_lens: torch.Tensor,
+    ) -> None:
+        if self._verify_ctx is None:
+            return
+        if not self.enabled():
+            self.clear()
+            return
+        if req_pool_indices.numel() == 0 or seq_lens.numel() == 0:
+            return
+
+        num_verify_tokens = self._num_verify_tokens()
+        if num_verify_tokens == 0:
+            self.clear()
+            return
+
+        backend = self.backend
+        token_to_kv_pool = backend.token_to_kv_pool
+        pending_seq_lens = token_to_kv_pool.get_online_c128_mtp_pending_seq_lens()
+        cur_bs = min(seq_lens.shape[0], req_pool_indices.shape[0])
+
+        for runtime in self._iter_layer_runtimes():
+            _jit_online_c128_mtp_module(
+                runtime.head_dim, seq_lens.dtype, req_pool_indices.dtype
+            ).commit_pending(
+                seq_lens,
+                req_pool_indices,
+                backend.req_to_token,
+                pending_seq_lens,
+                runtime.main_state,
+                cur_bs,
+                num_verify_tokens,
+                runtime.state_slot_offset,
+                token_to_kv_pool.get_online_c128_state_num_req_slots(),
+            )
+
+        self.clear()
+
+    def _num_verify_tokens(self) -> int:
+        if not self.enabled():
+            return 0
+        num_verify_tokens = int(self.backend.speculative_num_draft_tokens)
+        max_draft_tokens = (
+            self.backend.token_to_kv_pool.get_online_c128_mtp_max_draft_tokens()
+        )
+        return num_verify_tokens if 0 < num_verify_tokens <= max_draft_tokens else 0
+
+    def _active_ctx(self) -> Optional[_OnlineC128VerifyContext]:
+        ctx = self._verify_ctx
+        if (
+            ctx is None
+            or ctx.seq_lens.numel() == 0
+            or ctx.req_pool_indices.numel() == 0
+        ):
+            return None
+        return ctx
+
+    def _head_dim(self) -> Optional[int]:
+        for runtime in self._iter_layer_runtimes():
+            return runtime.head_dim
+        return None
+
+    def _iter_layer_runtimes(self):
+        if self._layer_runtimes is None:
+            runtimes = []
+            token_to_kv_pool = self.backend.token_to_kv_pool
+            for layer in self.backend.model_runner.model.model.layers:
+                attn = getattr(layer, "self_attn", None)
+                compressor = getattr(attn, "compressor", None)
+                if compressor is None or compressor.ratio != 128:
+                    continue
+                state_pool = token_to_kv_pool.get_attention_compress_states(
+                    compressor.layer_id
+                )
+                runtimes.append(
+                    _OnlineC128LayerRuntime(
+                        head_dim=compressor.head_dim,
+                        main_state=state_pool.kv_score_buffer.kv_score,
+                        state_slot_offset=state_pool.online_mtp_state_slot_offset,
+                    )
+                )
+            self._layer_runtimes = runtimes
+        return iter(self._layer_runtimes)

python/sglang/jit_kernel/dsv4/topk.py

@@ -0,0 +1,115 @@
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+
+from sglang.jit_kernel.utils import (
+    cache_once,
+    is_arch_support_pdl,
+    is_hip_runtime,
+    load_jit,
+    make_cpp_args,
+)
+
+from .utils import make_name
+
+
+@cache_once
+def _jit_topk_v1_module(topk: int):
+    args = make_cpp_args(is_arch_support_pdl())
+    assert topk in (512, 1024), "Only support topk=512 or 1024"
+    return load_jit(
+        make_name(f"topk_v1_{topk}"),
+        *args,
+        cuda_files=["deepseek_v4/topk_v1.cuh"],
+        cuda_wrappers=[("topk_transform", f"TopKKernel<{args}>::transform")],
+        extra_cuda_cflags=[f"-DSGL_TOPK={topk}"],
+    )
+
+
+@cache_once
+def _jit_topk_v2_module():
+    # v2 is universal: topk (<= 2048) is a runtime argument, not a compile-time
+    # constant, so a single module serves every k.
+    return load_jit(
+        make_name("topk_v2"),
+        cuda_files=["deepseek_v4/topk_v2.cuh"],
+        cuda_wrappers=[
+            ("topk_transform", "TopKKernel::transform"),
+            ("topk_plan", "TopKKernel::plan"),
+        ],
+    )
+
+
+def topk_transform_512(
+    scores: torch.Tensor,
+    seq_lens: torch.Tensor,
+    page_tables: torch.Tensor,
+    out_page_indices: torch.Tensor,
+    page_size: int,
+    out_raw_indices: Optional[torch.Tensor] = None,
+) -> None:
+    if is_hip_runtime():
+        torch.ops.sgl_kernel.deepseek_v4_topk_transform_512(
+            scores, seq_lens, page_tables, out_page_indices, page_size, out_raw_indices
+        )
+    else:
+        module = _jit_topk_v1_module(out_page_indices.shape[1])
+        module.topk_transform(
+            scores, seq_lens, page_tables, out_page_indices, page_size, out_raw_indices
+        )
+
+
+# metadata is (batch+1, 2) int32: row 0 = {cluster_threshold, num_cluster_items};
+# rows 1..N = {batch_id, seq_len} of items routed to the persistent cluster pool.
+_PLAN_METADATA_INTS_PER_BATCH = 2
+
+
+def plan_topk_v2(seq_lens: torch.Tensor, static_threshold: int = 0) -> torch.Tensor:
+    """Preprocess the per-batch routing plan for :func:`topk_transform_512_v2`.
+
+    IMPORTANT: every entry of ``seq_lens`` must be NON-NEGATIVE. The device
+    kernel reads the int32 buffer as ``uint32_t``, so a negative length (e.g.
+    -4 from a DP-padded / idle-companion row) reinterprets as ~4e9, poisons
+    the plan, and drives the transform kernel into an illegal memory access.
+    Producers of padded rows must clamp their lengths to 0 (0 selects the
+    trivial all-(-1) output path, which is safe).
+    """
+    module = _jit_topk_v2_module()
+    bs = seq_lens.shape[0]
+    metadata = seq_lens.new_empty(bs + 1, _PLAN_METADATA_INTS_PER_BATCH)
+    module.topk_plan(seq_lens, metadata, static_threshold)
+    return metadata
+
+
+def topk_transform_512_v2(
+    scores: torch.Tensor,
+    seq_lens: torch.Tensor,
+    page_tables: torch.Tensor,
+    out_page_indices: torch.Tensor,
+    page_size: int,
+    metadata: torch.Tensor,
+    out_raw_indices: Optional[torch.Tensor] = None,
+) -> None:
+    """Fused top-k + page-table transform (DeepSeek-V4 top-k v2 kernel).
+
+    IMPORTANT: every entry of ``seq_lens`` must be NON-NEGATIVE, and
+    ``metadata`` must come from :func:`plan_topk_v2` over the same ``seq_lens``
+    values. The kernel reads lengths as ``uint32_t``: a negative entry
+    reinterprets as a ~4e9-token sequence, sending the row down the cluster
+    path over garbage scores and crashing with an illegal memory access
+    (GLM 5.2 MTP DP-idle companion rows hit exactly this). A length of 0 is
+    the valid way to express "no tokens": the row takes the trivial path and
+    the output is all -1.
+    """
+    module = _jit_topk_v2_module()
+    module.topk_transform(
+        scores,
+        seq_lens,
+        page_tables,
+        out_page_indices,
+        page_size,
+        metadata,
+        out_raw_indices,
+    )

python/sglang/jit_kernel/dsv4/utils.py

@@ -0,0 +1,2 @@
+def make_name(name: str) -> str:
+    return f"dpsk_v4_{name}"