sgl-project--sglang/python/sglang/srt/layers/attention/dsv4/compressor_v2.py

from __future__ import annotations

from typing import TYPE_CHECKING, List, Literal, Optional, TypeAlias, Union, cast

import torch

from sglang.jit_kernel.dsv4 import (
    CompressorDecodePlan,
    CompressorPrefillPlan,
    compress_forward,
    compress_norm_rope_store,
)
from sglang.jit_kernel.utils import is_hip_runtime
from sglang.srt.environ import envs

if TYPE_CHECKING:
    from sglang.srt.layers.attention.deepseek_v4_backend import DSV4Metadata
    from sglang.srt.layers.attention.dsv4.compressor import Compressor
    from sglang.srt.layers.layernorm import RMSNorm
    from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool
    from sglang.srt.model_executor.forward_batch_info import ForwardBatch


CompressMetadata: TypeAlias = Union[CompressorDecodePlan, CompressorPrefillPlan]
# NOTE: alias for backward compatibility
FusedCompressMetadata: TypeAlias = CompressMetadata

_is_hip = is_hip_runtime()

if _is_hip:
    import triton
    import triton.language as tl

    @triton.jit
    def _c128_compress_decode_kernel(
        buf_ptr,
        input_ptr,
        ape_ptr,
        out_ptr,
        plan_ptr,
        buf_stride_slot,
        input_stride_b,
        ape_stride_r,
        out_stride_b,
        bs,
        HEAD_DIM: tl.constexpr,
        BLOCK_D: tl.constexpr,
        COMPRESS_RATIO: tl.constexpr,
    ):
        """Fused C128 decode: write to state buffer + online softmax-pool.

        plan_ptr points to int32 view: [bs, 4] where each row is
        {seq_len, write_loc, read_page_0, read_page_1}.
        """
        bid = tl.program_id(0)
        if bid >= bs:
            return

        # Parse plan
        plan_base = plan_ptr + bid * 4
        seq_len = tl.load(plan_base).to(tl.int32)
        write_loc = tl.load(plan_base + 1).to(tl.int32)
        read_page_0 = tl.load(plan_base + 2).to(tl.int32)

        d = tl.arange(0, BLOCK_D)
        last_dim: tl.constexpr = HEAD_DIM * 2

        # Step 1: Write kv_score_input to state buffer at write_loc
        d_mask_full = d < last_dim
        input_val = tl.load(
            input_ptr + bid * input_stride_b + d, mask=d_mask_full, other=0.0
        )
        tl.store(buf_ptr + write_loc * buf_stride_slot + d, input_val, mask=d_mask_full)

        # Step 2: Check boundary condition
        d_mask_hd = d < HEAD_DIM
        if seq_len % COMPRESS_RATIO != 0:
            tl.store(
                out_ptr + bid * out_stride_b + d,
                tl.zeros([BLOCK_D], tl.float32),
                mask=d_mask_hd,
            )
            return

        # Step 3: Online softmax-pool over 128 slots in the page
        page_base = read_page_0 * COMPRESS_RATIO * buf_stride_slot
        m_prev = tl.full([BLOCK_D], float("-inf"), tl.float32)
        kv_acc = tl.zeros([BLOCK_D], tl.float32)
        w_acc = tl.zeros([BLOCK_D], tl.float32)

        for k in tl.static_range(COMPRESS_RATIO):
            slot_addr = page_base + k * buf_stride_slot
            kv_val = tl.load(buf_ptr + slot_addr + d, mask=d_mask_hd, other=0.0).to(
                tl.float32
            )
            sc_val = tl.load(
                buf_ptr + slot_addr + HEAD_DIM + d, mask=d_mask_hd, other=0.0
            ).to(tl.float32)
            ape_val = tl.load(
                ape_ptr + k * ape_stride_r + d, mask=d_mask_hd, other=0.0
            ).to(tl.float32)
            score_k = sc_val + ape_val

            m_new = tl.maximum(m_prev, score_k)
            exp_old = tl.where(m_prev == float("-inf"), 0.0, tl.exp(m_prev - m_new))
            exp_cur = tl.where(score_k == float("-inf"), 0.0, tl.exp(score_k - m_new))
            kv_acc = kv_acc * exp_old + exp_cur * kv_val
            w_acc = w_acc * exp_old + exp_cur
            m_prev = m_new

        compressed = kv_acc / w_acc
        tl.store(out_ptr + bid * out_stride_b + d, compressed, mask=d_mask_hd)

    @triton.jit
    def _c128_compress_prefill_write_kernel(
        buf_ptr,
        input_ptr,
        plan_w_ptr,
        buf_stride_slot,
        input_stride_b,
        num_w,
        BLOCK_D: tl.constexpr,
        LAST_DIM: tl.constexpr,
    ):
        """Prefill write phase: scatter kv_score_input tokens into state buffer."""
        wid = tl.program_id(0)
        if wid >= num_w:
            return

        # WritePlan: {ragged_id(u32), write_loc(i32)} = 8 bytes = 2 int32s
        plan_base = plan_w_ptr + wid * 2
        ragged_id = (tl.load(plan_base).to(tl.int32)) & 0xFFFF
        write_loc = tl.load(plan_base + 1).to(tl.int32)

        d = tl.arange(0, BLOCK_D)
        d_mask = d < LAST_DIM

        if write_loc >= 0:
            input_val = tl.load(
                input_ptr + ragged_id * input_stride_b + d, mask=d_mask, other=0.0
            )
            tl.store(buf_ptr + write_loc * buf_stride_slot + d, input_val, mask=d_mask)

    @triton.jit
    def _c128_compress_prefill_compress_kernel(
        buf_ptr,
        ape_ptr,
        out_ptr,
        plan_c_ptr,
        buf_stride_slot,
        ape_stride_r,
        out_stride_b,
        num_c,
        HEAD_DIM: tl.constexpr,
        BLOCK_D: tl.constexpr,
        COMPRESS_RATIO: tl.constexpr,
    ):
        """Prefill compress phase: online softmax-pool for each compress plan entry."""
        cid = tl.program_id(0)
        if cid >= num_c:
            return

        # CompressPlan: {seq_len(u32), ragged_id(u16)|buffer_len(u16), read_page_0(i32), read_page_1(i32)}
        plan_base = plan_c_ptr + cid * 4
        read_page_0 = tl.load(plan_base + 2).to(tl.int32)

        d = tl.arange(0, BLOCK_D)
        d_mask_hd = d < HEAD_DIM

        if read_page_0 < 0:
            tl.store(
                out_ptr + cid * out_stride_b + d,
                tl.zeros([BLOCK_D], tl.float32),
                mask=d_mask_hd,
            )
            return

        page_base = read_page_0 * COMPRESS_RATIO * buf_stride_slot
        m_prev = tl.full([BLOCK_D], float("-inf"), tl.float32)
        kv_acc = tl.zeros([BLOCK_D], tl.float32)
        w_acc = tl.zeros([BLOCK_D], tl.float32)

        for k in tl.static_range(COMPRESS_RATIO):
            slot_addr = page_base + k * buf_stride_slot
            kv_val = tl.load(buf_ptr + slot_addr + d, mask=d_mask_hd, other=0.0).to(
                tl.float32
            )
            sc_val = tl.load(
                buf_ptr + slot_addr + HEAD_DIM + d, mask=d_mask_hd, other=0.0
            ).to(tl.float32)
            ape_val = tl.load(
                ape_ptr + k * ape_stride_r + d, mask=d_mask_hd, other=0.0
            ).to(tl.float32)
            score_k = sc_val + ape_val

            m_new = tl.maximum(m_prev, score_k)
            exp_old = tl.where(m_prev == float("-inf"), 0.0, tl.exp(m_prev - m_new))
            exp_cur = tl.where(score_k == float("-inf"), 0.0, tl.exp(score_k - m_new))
            kv_acc = kv_acc * exp_old + exp_cur * kv_val
            w_acc = w_acc * exp_old + exp_cur
            m_prev = m_new

        compressed = kv_acc / w_acc
        tl.store(out_ptr + cid * out_stride_b + d, compressed, mask=d_mask_hd)


def _compress_forward_c128_triton(
    kv_score_buffer: torch.Tensor,
    kv_score_input: torch.Tensor,
    ape: torch.Tensor,
    plan: Union[CompressorDecodePlan, CompressorPrefillPlan],
    head_dim: int,
) -> torch.Tensor:
    """Triton C128 compress_forward for HIP (wave64).

    Fuses write + online-softmax-pool into Triton kernels.
    CUDA graph compatible.
    """
    num_total_slots = kv_score_buffer.shape[0] * kv_score_buffer.shape[1]
    num_pages = kv_score_buffer.shape[0]
    last_dim = kv_score_buffer.shape[-1]
    compress_ratio = 128

    buf_flat = kv_score_buffer.view(-1, last_dim)
    buf_stride_slot = last_dim  # elements per slot

    BLOCK_D = triton.next_power_of_2(last_dim)

    if plan.is_decode:
        # Decode path: single kernel does write + compress
        plan_raw = plan[1].view(torch.int32)  # [bs, 4]
        bs = plan_raw.shape[0]
        out = torch.empty(
            bs, head_dim, dtype=torch.float32, device=kv_score_input.device
        )

        if bs > 0 and num_total_slots > 0:
            grid = (bs,)
            _c128_compress_decode_kernel[grid](
                buf_flat,
                kv_score_input,
                ape,
                out,
                plan_raw,
                buf_stride_slot,
                kv_score_input.stride(0),
                ape.stride(0),
                out.stride(0),
                bs,
                HEAD_DIM=head_dim,
                BLOCK_D=triton.next_power_of_2(head_dim),
                COMPRESS_RATIO=compress_ratio,
                num_warps=8,
            )
        return out
    else:
        # Prefill path: separate write kernel + compress kernel
        plan_c_raw = plan[1].view(torch.int32)  # [num_c, 4]
        plan_w = plan[2]  # [num_w, 8] uint8
        plan_w_raw = plan_w.view(torch.int32)  # [num_w, 2]
        num_c = plan_c_raw.shape[0]
        num_w = plan_w_raw.shape[0]

        out = torch.empty(
            num_c, head_dim, dtype=torch.float32, device=kv_score_input.device
        )

        # Phase 1: Write
        if num_w > 0 and num_total_slots > 0:
            grid_w = (num_w,)
            _c128_compress_prefill_write_kernel[grid_w](
                buf_flat,
                kv_score_input,
                plan_w_raw,
                buf_stride_slot,
                kv_score_input.stride(0),
                num_w,
                BLOCK_D=BLOCK_D,
                LAST_DIM=last_dim,
                num_warps=4,
            )

        # Phase 2: Compress
        if num_c > 0 and num_pages > 0:
            grid_c = (num_c,)
            _c128_compress_prefill_compress_kernel[grid_c](
                buf_flat,
                ape,
                out,
                plan_c_raw,
                buf_stride_slot,
                ape.stride(0),
                out.stride(0),
                num_c,
                HEAD_DIM=head_dim,
                BLOCK_D=triton.next_power_of_2(head_dim),
                COMPRESS_RATIO=compress_ratio,
                num_warps=8,
            )

        return out


def _use_online_compress(compress_ratio: int) -> bool:
    """Online state-pool path is c128-only."""
    return compress_ratio == 128 and envs.SGLANG_OPT_USE_ONLINE_COMPRESS.get()


def _extract_positions_from_plan(
    plan: Union[CompressorDecodePlan, CompressorPrefillPlan],
    compress_ratio: int,
) -> torch.Tensor:
    """Extract RoPE positions from plan tensors (decode or prefill).

    DecodePlan layout: [bs, 16] uint8, first 4 bytes = uint32 seq_len.
    CompressPlan layout: [num_c, 16] uint8, first 4 bytes = uint32 seq_len.
    Position for RoPE = seq_len - compress_ratio.
    """
    plan_tensor = plan[1]  # plan_d or plan_c
    seq_lens = plan_tensor[:, :4].contiguous().view(torch.int32).squeeze(-1)
    positions = seq_lens.to(torch.int32) - compress_ratio
    return positions


def _compress_forward_c128_fallback(
    kv_score_buffer: torch.Tensor,
    kv_score_input: torch.Tensor,
    ape: torch.Tensor,
    plan: Union[CompressorDecodePlan, CompressorPrefillPlan],
    head_dim: int,
) -> torch.Tensor:
    """PyTorch fallback for C128 compress_forward on HIP (wave64).

    Fully vectorized, compatible with CUDA graph capture.
    kv_score_buffer: [num_pages, 128, head_dim * 2]
    ape: [128, head_dim]

    IMPORTANT: This also performs the write to state buffer (like the JIT kernel).
    The JIT kernel does: (1) write kv_score_input to buffer, (2) compress from buffer.
    """
    num_total_slots = kv_score_buffer.shape[0] * kv_score_buffer.shape[1]
    num_pages = kv_score_buffer.shape[0]
    last_dim = kv_score_buffer.shape[-1]

    # Step 1: WRITE kv_score_input to state buffer
    if num_total_slots > 0:
        buf_flat = kv_score_buffer.view(-1, last_dim)
        if plan.is_decode:
            # Decode: plan_d has write_loc per batch item
            plan_raw = plan[1].view(torch.int32)  # [bs, 4]
            write_locs = plan_raw[:, 1].long()
            # Only write valid locations (>= 0 and < buffer size)
            valid_write = (write_locs >= 0) & (write_locs < num_total_slots)
            if valid_write.any():
                buf_flat[write_locs[valid_write]] = kv_score_input[valid_write]
        else:
            # Prefill: plan_w has {ragged_id, write_loc} per write entry
            plan_w = plan[2]  # [num_w, 8] uint8 = WritePlan
            if plan_w.shape[0] > 0:
                plan_w_raw = plan_w.view(torch.int32)  # [num_w, 2]
                ragged_ids = plan_w_raw[:, 0].long() & 0xFFFF
                write_locs = plan_w_raw[:, 1].long()
                valid_write = (write_locs >= 0) & (write_locs < num_total_slots)
                ragged_ids_safe = ragged_ids.clamp(
                    min=0, max=kv_score_input.shape[0] - 1
                )
                if valid_write.any():
                    buf_flat[write_locs[valid_write]] = kv_score_input[
                        ragged_ids_safe[valid_write]
                    ]

    # Step 2: COMPRESS (read from buffer page and do softmax-pool)
    plan_c = plan[1]  # plan_d for decode, plan_c for prefill
    num_tokens = plan_c.shape[0]
    if num_pages == 0 or num_tokens == 0:
        return kv_score_input.new_zeros(num_tokens, head_dim)

    plan_c_raw = plan_c.view(torch.int32)  # [N, 4]
    read_page_0 = plan_c_raw[:, 2].long()
    # Use torch.where instead of clamp to handle -1 (invalid) gracefully
    valid_read = (read_page_0 >= 0) & (read_page_0 < num_pages)
    read_page_0_safe = torch.where(
        valid_read, read_page_0, torch.zeros_like(read_page_0)
    )

    gathered = kv_score_buffer[read_page_0_safe]  # [N, 128, head_dim*2]
    kv = gathered[:, :, :head_dim].float()
    score = gathered[:, :, head_dim:].float() + ape.float().unsqueeze(0)
    weights = score.softmax(dim=1)
    out = (weights * kv).sum(dim=1)

    # For decode: zero out non-boundary tokens (seq_len % 128 != 0)
    # so they don't corrupt kvcache location 0 when stored.
    if plan.is_decode:
        seq_lens = plan_c_raw[:, 0].to(torch.int32)
        is_boundary = (seq_lens % 128 == 0).unsqueeze(-1)  # [N, 1]
        out = torch.where(is_boundary, out, torch.zeros_like(out))

    return out.to(kv_score_input.dtype)


class CompressorBackendMixin:
    def __init__(self):
        super().__init__()
        self.forward_metadata: DSV4Metadata

    def _get_paged_compress_metadata(self, compress_ratio: int) -> CompressMetadata:
        attr_name = f"c{compress_ratio}_compress_metadata"
        return getattr(self.forward_metadata, attr_name)

    def _get_out_loc(self, compress_ratio: int) -> torch.Tensor:
        attr_name = f"c{compress_ratio}_out_loc"
        return getattr(self.forward_metadata.core_metadata, attr_name)

    def _forward_compress_all_in_one(
        self,
        *,
        kv_score_buffer: torch.Tensor,
        kv_score_input: torch.Tensor,
        ape: torch.Tensor,
        head_dim: int,
        norm: RMSNorm,
        freqs_cis_cache: torch.Tensor,
        kv_cache: torch.Tensor,
        is_indexer: bool,
        rotate: bool,
        compress_ratio: int,
        page_size: int,
        out_loc: torch.Tensor,
        use_fp4_indexer: bool = False,
        bf16_store: bool = False,
    ) -> None:
        assert compress_ratio == 4 or compress_ratio == 128
        assert rotate == is_indexer == (head_dim == 128)
        if use_fp4_indexer:
            assert is_indexer
            assert compress_ratio == 4
            assert head_dim == 128

        plan = self._get_paged_compress_metadata(compress_ratio)
        is_online = _use_online_compress(compress_ratio)
        if is_online:
            kv_score_buffer = kv_score_buffer.view(-1, 1, head_dim * 3)
        else:
            coff = 2 if is_overlap_compress(compress_ratio) else 1
            last_dim = 2 * head_dim * coff
            assert kv_score_buffer.shape[-1] == last_dim
            kv_score_buffer = kv_score_buffer.view(-1, compress_ratio, last_dim)

        # Step 1: compress_forward
        kv_compressed = compress_forward(
            kv_score_buffer=kv_score_buffer,
            kv_score_input=kv_score_input,
            ape=ape.view(-1, head_dim),
            plan=plan,
            compress_ratio=compress_ratio,
            head_dim=head_dim,
            is_online=is_online,
        )

        # Step 2: norm + rope + store
        compress_norm_rope_store(
            kv_compressed,
            plan,
            norm_weight=norm.weight,
            norm_eps=norm.variance_epsilon,
            freq_cis=freqs_cis_cache,
            out_loc=out_loc,
            kvcache=kv_cache,
            page_size=page_size,
            use_fp4=use_fp4_indexer,
            bf16_store=bf16_store,
        )

    def forward_unified(
        self,
        x: torch.Tensor,
        forward_batch: ForwardBatch,
        layer_id: int,
        compressor: Compressor,
    ) -> None:
        if forward_batch.forward_mode.is_idle():
            return

        token_to_kv_pool = self.token_to_kv_pool
        token_to_kv_pool = cast("DeepSeekV4TokenToKVPool", token_to_kv_pool)
        kv_score_input = compressor.compute_kv_score(x, forward_batch)

        state_pool = compressor.get_state_pool(self)
        from sglang.srt.layers.attention.dsv4.unified_kv_kernels.env_gate import (
            is_unified_kv_triton,
        )

        if _is_hip and not envs.SGLANG_OPT_USE_JIT_NORM.get():
            self._forward_unified_hip(
                token_to_kv_pool=token_to_kv_pool,
                kv_score_input=kv_score_input,
                state_pool=state_pool,
                compressor=compressor,
                layer_id=layer_id,
            )
        else:
            out_loc = self._get_out_loc(compressor.ratio)
            use_fp4_indexer = (
                compressor.is_in_indexer and self.enable_deepseek_v4_fp4_indexer
            )
            bf16_store = False
            if compressor.is_in_indexer:
                kv_cache = token_to_kv_pool.get_index_k_with_scale_buffer(layer_id)
                page_size = token_to_kv_pool.get_index_k_page_size()
            elif is_unified_kv_triton():
                kv_cache = token_to_kv_pool.get_unified_kv(layer_id)
                page_size = 1
                out_loc = getattr(
                    self.forward_metadata.core_metadata.unified,
                    f"c{compressor.ratio}_out_loc",
                )
                bf16_store = True
            else:
                _, _, compress_kv_pool = token_to_kv_pool.layer_mapping[layer_id]
                assert compress_kv_pool is not None
                kv_cache = token_to_kv_pool.get_extra_key_buffer(layer_id)
                page_size = token_to_kv_pool.get_extra_key_page_size(layer_id)
                if hasattr(compress_kv_pool, "translate_loc_to_hisparse_device"):
                    out_loc = compress_kv_pool._translate_loc_to_hisparse_device(
                        out_loc
                    )
            self._forward_compress_all_in_one(
                kv_score_buffer=state_pool.kv_score_buffer.kv_score,
                kv_score_input=kv_score_input,
                ape=compressor.ape,
                head_dim=compressor.head_dim,
                norm=compressor.norm,
                freqs_cis_cache=compressor.freqs_cis,
                kv_cache=kv_cache.view(dtype=torch.uint8),
                is_indexer=compressor.is_in_indexer,
                rotate=compressor.rotate,
                compress_ratio=compressor.ratio,
                page_size=page_size,
                out_loc=out_loc,
                use_fp4_indexer=use_fp4_indexer,
                bf16_store=bf16_store,
            )
        online_c128_mtp = getattr(self, "online_c128_mtp", None)
        if online_c128_mtp is not None:
            online_c128_mtp.write_prefix_states(
                layer_id=layer_id,
                compressor=compressor,
                kv_score_input=kv_score_input,
                logical_forward_mode=getattr(
                    forward_batch, "_original_forward_mode", None
                )
                or forward_batch.forward_mode,
            )

    def _forward_unified_hip(
        self,
        token_to_kv_pool: DeepSeekV4TokenToKVPool,
        kv_score_input: torch.Tensor,
        state_pool,
        compressor: Compressor,
        layer_id: int,
    ) -> None:
        """HIP-specific forward path using PyTorch/Triton fallbacks."""
        from sglang.srt.layers.attention.dsv4.quant_k_cache import (
            quant_to_nope_fp8_rope_bf16_pack_triton,
        )
        from sglang.srt.layers.attention.nsa.nsa_indexer import rotate_activation
        from sglang.srt.layers.attention.nsa.triton_kernel import act_quant
        from sglang.srt.layers.deepseek_v4_rope import fused_norm_rope_inplace_triton

        compress_ratio = compressor.ratio
        head_dim = compressor.head_dim
        is_indexer = compressor.is_in_indexer

        plan = self._get_paged_compress_metadata(compress_ratio)
        out_loc = self._get_out_loc(compress_ratio)

        # Step 1: compress_forward (always use JIT for both C4 and C128)
        coff = 2 if is_overlap_compress(compress_ratio) else 1
        last_dim = 2 * head_dim * coff
        kv_score_buffer = state_pool.kv_score_buffer.kv_score
        kv_score_buffer = kv_score_buffer.view(-1, compress_ratio, last_dim)

        kv_compressed = compress_forward(
            kv_score_buffer=kv_score_buffer,
            kv_score_input=kv_score_input,
            ape=compressor.ape.view(-1, head_dim),
            plan=plan,
            compress_ratio=compress_ratio,
            head_dim=head_dim,
            is_online=False,
        )

        if kv_compressed.shape[0] == 0:
            return

        # For decode: zero out non-boundary tokens to prevent corrupting kvcache loc 0.
        if plan.is_decode:
            plan_raw = plan[1].view(torch.int32)
            seq_lens_plan = plan_raw[:, 0].to(torch.int32)
            is_boundary = (seq_lens_plan % compress_ratio == 0).unsqueeze(-1)
            kv_compressed = torch.where(
                is_boundary, kv_compressed, torch.zeros_like(kv_compressed)
            )

        # Step 2: norm + rope (Triton fallback for precision parity with V1)
        positions = _extract_positions_from_plan(plan, compress_ratio)
        positions_safe = positions.clamp(min=0)

        fused_norm_rope_inplace_triton(
            kv_compressed,
            compressor.norm.weight,
            compressor.norm.variance_epsilon,
            compressor.freqs_cis,
            positions=positions_safe,
        )

        # Step 3: optional Hadamard rotation for indexer
        if compressor.rotate:
            kv_compressed = rotate_activation(kv_compressed)

        # Step 4: store to kvcache
        # For decode: store ALL tokens. Non-boundary tokens have out_loc=0 (safe).
        # For prefill: plan_c already only contains valid entries.
        if plan.is_decode:
            kv_to_store = kv_compressed
            out_loc_to_store = out_loc
        else:
            kv_to_store = kv_compressed
            plan_raw = plan[1].view(torch.int32)
            ragged_ids = plan_raw[:, 1].to(torch.int32) & 0xFFFF
            out_loc_to_store = out_loc[ragged_ids.long()]

        if kv_to_store.shape[0] == 0:
            return

        if envs.SGLANG_OPT_USE_FUSED_STORE_CACHE.get():
            # fused kernel: BF16 in -> FP8 quant + paged scatter in one launch
            if is_indexer:
                token_to_kv_pool.set_index_k_fused(
                    layer_id=layer_id,
                    loc=out_loc_to_store,
                    cache_k=kv_to_store,
                )
            else:
                token_to_kv_pool.set_extra_key_buffer_fused(
                    layer_id=layer_id,
                    loc=out_loc_to_store,
                    cache_k=kv_to_store,
                )
        else:
            if is_indexer:
                kv_fp8, kv_scale = act_quant(kv_to_store)
                token_to_kv_pool.set_index_k_scale_buffer(
                    layer_id=layer_id,
                    loc=out_loc_to_store,
                    index_k=kv_fp8,
                    index_k_scale=kv_scale,
                )
            else:
                pack = quant_to_nope_fp8_rope_bf16_pack_triton(kv_to_store.bfloat16())
                token_to_kv_pool.set_extra_key_buffer(layer_id, out_loc_to_store, pack)

    # NOTE: alias for backward compatibility
    forward_indexer_compressor = forward_unified
    forward_core_compressor = forward_unified


def is_overlap_compress(compress_ratio: int) -> bool:
    return compress_ratio == 4


def create_paged_compressor_data(
    compress_ratio: Literal[4, 128],
    *,
    is_prefill: bool,
    token_to_kv_pool: DeepSeekV4TokenToKVPool,
    req_to_token: torch.Tensor,
    req_pool_indices: torch.Tensor,
    seq_lens: torch.Tensor,
    extend_lens: Optional[torch.Tensor] = None,
    seq_lens_cpu: Optional[List[int]] = None,
    extend_lens_cpu: Optional[List[int]] = None,
    use_prefill_cuda_graph: bool = False,
    num_q_tokens: Optional[int] = None,
    online_state_slot_offset: int = 0,
) -> CompressMetadata:
    """Build the paged compress metadata (= the plan).

    State-pool slot translation is done inside the C++ planner; the
    Python side just hands the relevant tensors over.
    """
    if _use_online_compress(compress_ratio):
        return _create_online_paged_compressor_data(
            is_prefill=is_prefill,
            token_to_kv_pool=token_to_kv_pool,
            req_to_token=req_to_token,
            req_pool_indices=req_pool_indices,
            seq_lens=seq_lens,
            extend_lens=extend_lens,
            seq_lens_cpu=seq_lens_cpu,
            extend_lens_cpu=extend_lens_cpu,
            use_prefill_cuda_graph=use_prefill_cuda_graph,
            num_q_tokens=num_q_tokens,
            online_state_slot_offset=online_state_slot_offset,
        )

    swa_page_size = token_to_kv_pool.swa_page_size
    ring_size = token_to_kv_pool.get_ring_size(compress_ratio=compress_ratio)
    # NOTE: This is actually a proxy, which encounter some bug with tvm-ffi.
    # As a workaround, we use `.detach()` to get the real tensor.
    full_to_swa = token_to_kv_pool.full_to_swa_index_mapping.detach()
    req_pool_indices_i64 = req_pool_indices.to(torch.int64)

    if is_prefill:
        assert extend_lens is not None
        if seq_lens_cpu is not None:
            assert extend_lens_cpu is not None
            seq_lens_planner = torch.tensor(seq_lens_cpu, dtype=torch.int64)
            extend_lens_planner = torch.tensor(extend_lens_cpu, dtype=torch.int64)
            num_q_tokens = sum(extend_lens_cpu)
        else:
            assert num_q_tokens is not None
            seq_lens_planner = seq_lens.to(torch.int64)
            extend_lens_planner = extend_lens.to(torch.int64)

        return CompressorPrefillPlan.generate(
            compress_ratio=compress_ratio,
            req_pool_indices=req_pool_indices_i64,
            seq_lens=seq_lens_planner,
            extend_lens=extend_lens_planner,
            req_to_token=req_to_token,
            full_to_state=full_to_swa,
            swa_page_size=swa_page_size,
            ring_size=ring_size,
            num_q_tokens=num_q_tokens,
            use_cuda_graph=use_prefill_cuda_graph,
        )
    else:
        return CompressorDecodePlan.generate(
            compress_ratio=compress_ratio,
            req_pool_indices=req_pool_indices_i64,
            req_to_token=req_to_token,
            full_to_state=full_to_swa,
            seq_lens=seq_lens.to(torch.int64),
            swa_page_size=swa_page_size,
            ring_size=ring_size,
        )


def _create_online_paged_compressor_data(
    *,
    is_prefill: bool,
    token_to_kv_pool: DeepSeekV4TokenToKVPool,
    req_to_token: torch.Tensor,
    req_pool_indices: torch.Tensor,
    seq_lens: torch.Tensor,
    extend_lens: Optional[torch.Tensor],
    seq_lens_cpu: Optional[List[int]],
    extend_lens_cpu: Optional[List[int]],
    use_prefill_cuda_graph: bool,
    num_q_tokens: Optional[int],
    online_state_slot_offset: int = 0,
) -> CompressMetadata:
    req_pool_indices = req_pool_indices.to(torch.int64)

    if is_prefill:
        # Sync-on-entry: catch IMA from a prior layer / kernel BEFORE we touch
        # anything in this builder, so blame doesn't land on us spuriously.
        assert extend_lens is not None
        if seq_lens_cpu is not None:
            assert extend_lens_cpu is not None
            seq_lens_planner = torch.tensor(seq_lens_cpu, dtype=torch.int64)
            extend_lens_planner = torch.tensor(extend_lens_cpu, dtype=torch.int64)
            num_q_tokens_planner = sum(extend_lens_cpu)
        else:
            assert num_q_tokens is not None
            seq_lens_planner = seq_lens.to(torch.int64)
            extend_lens_planner = extend_lens.to(torch.int64)
            num_q_tokens_planner = num_q_tokens

        return CompressorPrefillPlan.generate_online(
            seq_lens=seq_lens_planner,
            extend_lens=extend_lens_planner,
            req_pool_indices=req_pool_indices,
            req_to_token=req_to_token,
            num_q_tokens=int(num_q_tokens_planner),
            use_cuda_graph=use_prefill_cuda_graph,
            state_slot_offset=online_state_slot_offset,
        )
    else:
        return CompressorDecodePlan.generate_online(
            seq_lens=seq_lens.to(torch.int64),
            req_pool_indices=req_pool_indices,
            req_to_token=req_to_token,
            state_slot_offset=online_state_slot_offset,
        )