from __future__ import annotations from typing import TYPE_CHECKING, List, Literal, NamedTuple, Optional, Union import torch import torch.nn as nn from sglang.jit_kernel.dsv4 import linear_bf16_fp32, triton_create_paged_compress_data from sglang.jit_kernel.dsv4.compress_old import ( CompressorDecodePlan, CompressorPrefillPlan, compress_forward, compress_fused_norm_rope_inplace, ) from sglang.srt.configs.deepseek_v4 import DeepSeekV4Config from sglang.srt.environ import envs from sglang.srt.layers.attention.dsa.triton_kernel import act_quant from sglang.srt.layers.attention.dsa.utils import dsa_use_prefill_cp from sglang.srt.layers.attention.dsv4.quant_k_cache import ( quant_to_nope_fp8_rope_bf16_pack_triton, ) from sglang.srt.layers.layernorm import RMSNorm from sglang.srt.layers.linear import ReplicatedLinear from sglang.srt.layers.utils.cp_utils import cp_all_gather_rerange_output from sglang.srt.layers.utils.multi_platform import MultiPlatformOp from sglang.srt.mem_cache.deepseek_v4_compress_state import ( CompressStatePool, ) from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool from sglang.srt.model_executor.forward_context import get_attn_backend from sglang.srt.models.deepseek_v2 import _is_hip from sglang.srt.runtime_context import get_parallel from sglang.srt.utils import add_prefix, is_npu, set_weight_attrs _is_npu = is_npu() if TYPE_CHECKING: from sglang.srt.layers.attention.base_attn_backend import AttentionBackend from sglang.srt.layers.attention.deepseek_v4_backend import DeepseekV4AttnBackend from sglang.srt.layers.rotary_embedding import RotaryEmbedding from sglang.srt.model_executor.forward_batch_info import ForwardBatch class FusedCompressMetadata(NamedTuple): write_loc: torch.Tensor extra_data: Optional[torch.Tensor] plan: Union[CompressorDecodePlan, CompressorPrefillPlan] def copy_(self, other: FusedCompressMetadata) -> None: from .metadata import maybe_copy_inplace self.write_loc.copy_(other.write_loc) maybe_copy_inplace(self.extra_data, src=other.extra_data) self.plan.copy_(other.plan) class CompressorBackendMixin: def get_paged_compress_metadata(self, compress_ratio: int) -> FusedCompressMetadata: attr_name = f"c{compress_ratio}_compress_metadata" metadata = getattr(self.forward_metadata, attr_name) assert isinstance(metadata, FusedCompressMetadata) return metadata def forward_compress( self, *, kv_score_buffer: torch.Tensor, kv_score_input: torch.Tensor, ape: torch.Tensor, head_dim: int, norm: RMSNorm, freqs_cis_cache: torch.Tensor, rotate: bool, forward_batch: ForwardBatch, compress_ratio: int, is_paged: bool = False, ) -> torch.Tensor: from sglang.srt.layers.attention.dsa.dsa_indexer import rotate_activation assert compress_ratio in ( 4, 128, ), f"DSV4 supports CSA(4x) and HCA(128x) only, got {compress_ratio=}" if is_paged: metadata = self.get_paged_compress_metadata(compress_ratio) coff = 2 if is_overlap_compress(compress_ratio) else 1 if compress_ratio == 128 and envs.SGLANG_OPT_USE_ONLINE_COMPRESS.get(): kv_score_buffer = kv_score_buffer.view(-1, 1, head_dim * 3) else: 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) else: plan = make_compressor_plan(compress_ratio, forward_batch) metadata = (forward_batch.req_pool_indices.to(torch.int32), None, plan) indices, extra_data, plan = metadata if _is_hip: if not is_paged: raise NotImplementedError("HIP fused compressor expects paged metadata") from sglang.srt.layers.attention.dsv4.fused_compress_triton import ( hip_compress_forward, hip_compress_fused_norm_rope_hadamard_inplace, hip_compress_fused_norm_rope_inplace, ) kv_compressed = hip_compress_forward( kv_score_buffer=kv_score_buffer, kv_score_input=kv_score_input, ape=ape, indices=indices, plan=plan, compress_ratio=compress_ratio, head_dim=head_dim, extra_data=extra_data, ) norm_eps = ( norm.variance_epsilon if hasattr(norm, "variance_epsilon") else norm.eps ) if rotate: hip_compress_fused_norm_rope_hadamard_inplace( kv_compressed, norm.weight, norm_eps, freqs_cis_cache, plan, head_dim, ) else: hip_compress_fused_norm_rope_inplace( kv_compressed, norm.weight, norm_eps, freqs_cis_cache, plan, ) return kv_compressed kv_compressed = compress_forward( kv_score_buffer=kv_score_buffer, kv_score_input=kv_score_input, ape=ape, indices=indices, plan=plan, compress_ratio=compress_ratio, head_dim=head_dim, extra_data=extra_data, ) compress_fused_norm_rope_inplace( kv_compressed, norm.weight, norm.variance_epsilon, freqs_cis_cache, plan, ) return rotate_activation(kv_compressed) if rotate else kv_compressed def forward_core_compressor( 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 if TYPE_CHECKING: assert isinstance(token_to_kv_pool, DeepSeekV4TokenToKVPool) new_compressed_kv = compressor(x, forward_batch, attn_backend=self) core_metadata = self.forward_metadata.core_metadata out_loc = ( core_metadata.c4_out_loc if compressor.ratio == 4 else core_metadata.c128_out_loc ) if out_loc.shape[0] > new_compressed_kv.shape[0]: out_loc = out_loc[: new_compressed_kv.shape[0]] if envs.SGLANG_OPT_USE_FUSED_STORE_CACHE.get(): token_to_kv_pool.set_extra_key_buffer_fused( layer_id=layer_id, loc=out_loc, cache_k=new_compressed_kv, ) else: pack = quant_to_nope_fp8_rope_bf16_pack_triton(new_compressed_kv.bfloat16()) token_to_kv_pool.set_extra_key_buffer(layer_id, out_loc, pack) def forward_indexer_compressor( self, x: torch.Tensor, forward_batch: ForwardBatch, layer_id: int, compressor: Compressor, ) -> None: assert is_overlap_compress(compressor.ratio) token_to_kv_pool = self.token_to_kv_pool if TYPE_CHECKING: assert isinstance(token_to_kv_pool, DeepSeekV4TokenToKVPool) new_compressed_kv = compressor(x, forward_batch, attn_backend=self) out_loc = self.forward_metadata.core_metadata.c4_out_loc if out_loc.shape[0] > new_compressed_kv.shape[0]: out_loc = out_loc[: new_compressed_kv.shape[0]] if self.enable_deepseek_v4_fp4_indexer: token_to_kv_pool.set_index_k_fp4( layer_id=layer_id, loc=out_loc, cache_k=new_compressed_kv, ) elif envs.SGLANG_OPT_USE_FUSED_STORE_CACHE.get(): token_to_kv_pool.set_index_k_fused( layer_id=layer_id, loc=out_loc, cache_k=new_compressed_kv, ) else: new_compressed_kv_fp8, new_compressed_kv_scale = act_quant( new_compressed_kv ) token_to_kv_pool.set_index_k_scale_buffer( layer_id=layer_id, loc=out_loc, index_k=new_compressed_kv_fp8, index_k_scale=new_compressed_kv_scale, ) def is_overlap_compress(compress_ratio: int) -> bool: return compress_ratio == 4 def make_compressor_plan( compress_ratio: Literal[4, 128], forward_batch: ForwardBatch, ) -> Union[CompressorDecodePlan, CompressorPrefillPlan]: if forward_batch.forward_mode.is_decode(): seq_lens_32 = forward_batch.seq_lens.to(torch.int32) return CompressorDecodePlan(compress_ratio, seq_lens_32) if forward_batch.forward_mode.is_prefill(): assert not forward_batch.forward_mode.is_target_verify() extend_lens_list = forward_batch.extend_seq_lens_cpu seq_lens_cpu = forward_batch.seq_lens_cpu assert extend_lens_list is not None and seq_lens_cpu is not None return CompressorPrefillPlan.generate( compress_ratio=compress_ratio, num_q_tokens=sum(extend_lens_list), seq_lens=seq_lens_cpu, extend_lens=torch.tensor(extend_lens_list), device=forward_batch.seq_lens.device, ) elif forward_batch.forward_mode.is_target_verify(): raise NotImplementedError("target verify mode to be implemented") else: raise NotImplementedError(f"unsupported mode {forward_batch.forward_mode=}") 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, ) -> FusedCompressMetadata: swa_page_size = token_to_kv_pool.swa_page_size ring_size = token_to_kv_pool.get_ring_size(compress_ratio=compress_ratio) # assert ring_size % compress_ratio == 0 def clip_down(positions: torch.Tensor) -> torch.Tensor: return positions // compress_ratio * compress_ratio def get_raw_loc(positions: torch.Tensor) -> torch.Tensor: positions = positions.masked_fill(positions < 0, 0) if compress_ratio == 128: state_loc = req_pool_indices * ring_size + positions % ring_size else: loc = req_to_token[req_pool_indices, positions] swa_loc = token_to_kv_pool.translate_loc_from_full_to_swa(loc) swa_pages = swa_loc // swa_page_size state_loc = swa_pages * ring_size + swa_loc % ring_size return (state_loc // compress_ratio).to(torch.int32) is_overlap = is_overlap_compress(compress_ratio) if is_prefill: assert extend_lens is not None write_loc, extra_data = triton_create_paged_compress_data( compress_ratio=compress_ratio, is_overlap=is_overlap, swa_page_size=swa_page_size, ring_size=ring_size, req_pool_indices=req_pool_indices, seq_lens=seq_lens, extend_seq_lens=extend_lens, req_to_token=req_to_token, full_to_swa_index_mapping=token_to_kv_pool.full_to_swa_index_mapping, ) plan_kwargs: dict if seq_lens_cpu is None: assert num_q_tokens is not None plan_kwargs = dict( num_q_tokens=num_q_tokens, seq_lens=seq_lens, extend_lens=extend_lens, ) else: assert extend_lens_cpu is not None plan_kwargs = dict( num_q_tokens=sum(extend_lens_cpu), seq_lens=torch.tensor(seq_lens_cpu), extend_lens=torch.tensor(extend_lens_cpu), ) plan = CompressorPrefillPlan.generate( compress_ratio=compress_ratio, device=seq_lens.device, use_cuda_graph=use_prefill_cuda_graph, **plan_kwargs, ) else: write_positions = clip_down(seq_lens - 1) write_loc = get_raw_loc(write_positions) if is_overlap: write_overlap_loc = get_raw_loc(write_positions - compress_ratio) extra_data = write_overlap_loc.view(-1, 1) elif _is_hip: extra_data = get_raw_loc(write_positions - compress_ratio) else: extra_data = None plan = CompressorDecodePlan(compress_ratio, seq_lens.to(torch.int32)) return FusedCompressMetadata(write_loc=write_loc, extra_data=extra_data, plan=plan) class Compressor(MultiPlatformOp): def __init__( self, config: DeepSeekV4Config, layer_id: int, is_in_indexer: bool, freqs_cis: torch.Tensor, compress_ratio: Literal[0, 4, 128], head_dim: int, rotate: bool = False, prefix: str = "", rotary_emb: Optional[RotaryEmbedding] = None, ) -> None: super().__init__() self.layer_id = layer_id self.is_in_indexer = is_in_indexer self.dim = config.hidden_size self.head_dim = head_dim self.rope_head_dim = getattr(config, "qk_rope_head_dim", 64) assert compress_ratio != 0, "compress_ratio should not be 0" self.ratio = compress_ratio self.overlap = self.ratio == 4 self.rotate = rotate self.coff = coff = 1 + self.overlap self.ape = nn.Parameter( torch.empty(self.ratio, coff * self.head_dim, dtype=torch.float32) ) set_weight_attrs(self.ape, {"weight_loader": self.load_ape_weight}) wkv_gate_dtype = torch.bfloat16 self.wkv_gate = ReplicatedLinear( self.dim, 2 * coff * self.head_dim, bias=False, quant_config=None, prefix=add_prefix("wkv_gate", prefix), params_dtype=wkv_gate_dtype, ) self.norm = RMSNorm( self.head_dim, eps=config.rms_norm_eps, weight_dtype=torch.float32 ) self.rotary_emb = rotary_emb self.freqs_cis = freqs_cis self.ape_converted = False def _apply_ape_hotfix(self): self.ape_converted = True if _is_npu: return if self.overlap: ape = torch.chunk(self.ape.data, 2, dim=-1) ape = torch.cat([ape[0], ape[1]], dim=0) self.ape.data.copy_(ape.view(self.ratio, -1)) def apply_ape_hotfix(self): assert not self.ape_converted self._apply_ape_hotfix() def load_ape_weight(self, param: torch.Tensor, loaded_weight: torch.Tensor) -> None: assert param is self.ape assert loaded_weight.shape == param.shape param.data.copy_(loaded_weight) self._apply_ape_hotfix() def get_state_pool(self, attn_backend: AttentionBackend) -> CompressStatePool: token_to_kv_pool = attn_backend.token_to_kv_pool assert isinstance(token_to_kv_pool, DeepSeekV4TokenToKVPool) if self.is_in_indexer: ret = token_to_kv_pool.get_indexer_compress_states(self.layer_id) else: ret = token_to_kv_pool.get_attention_compress_states(self.layer_id) assert isinstance(ret, CompressStatePool) return ret def compute_kv_score(self, x: torch.Tensor, forward_batch: ForwardBatch): kv_score = linear_bf16_fp32(x, self.wkv_gate.weight) # CUDA path: delegate to backend if dsa_use_prefill_cp(forward_batch): kv_score = cp_all_gather_rerange_output( kv_score, get_parallel().attn_cp_size, forward_batch, torch.cuda.current_stream(), ) return kv_score def forward_native( self, x: torch.Tensor, forward_batch: ForwardBatch, attn_backend: Optional[AttentionBackend] = None, ) -> torch.Tensor: if forward_batch.forward_mode.is_idle(): assert x.shape[0] == 0 return x.new_empty(0, self.head_dim) kv_score = self.compute_kv_score(x, forward_batch) if TYPE_CHECKING: assert isinstance(attn_backend, DeepseekV4AttnBackend) kv_score_buffer = self.get_state_pool(attn_backend).kv_score_buffer.kv_score return attn_backend.forward_compress( kv_score_buffer=kv_score_buffer, kv_score_input=kv_score, ape=self.ape.view(-1, self.head_dim), head_dim=self.head_dim, norm=self.norm, freqs_cis_cache=self.freqs_cis, rotate=self.rotate, compress_ratio=self.ratio, forward_batch=forward_batch, is_paged=True, ) def forward_npu( self, x: torch.Tensor, forward_batch: ForwardBatch, attn_backend: Optional[AttentionBackend] = None, ) -> torch.Tensor: if forward_batch.forward_mode.is_idle(): assert x.shape[0] == 0 return x.new_empty(0, self.head_dim) if dsa_use_prefill_cp(forward_batch): x = cp_all_gather_rerange_output( x, get_parallel().attn_cp_size, forward_batch, torch.cuda.current_stream(), ) return get_attn_backend().forward_compress(self, x, forward_batch)