From f74447bfd0b75c505295c9907bae1921ae7e1269 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Mon, 18 May 2026 23:20:13 +0000 Subject: [PATCH] Proper NVFP4 integration: quantized compressor/indexer + mapper fixes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Weight mapper fixes: - Reorder substr renames: compressor renames first, then .self_attn.compressor. → .attn.mla_attn.compressor., then indexer renames (so indexer keys end up under mla_attn after the compressor rename already fired) - Add compressor param renames: kv_proj→wkv, gate_proj→wgate, kv_norm→norm, position_bias→ape (checkpoint uses NVFP4 naming, model uses internal names) - Add indexer param renames: q_b_proj→wq_b, kv_proj→compressor.wkv, gate_proj→compressor.wgate, kv_norm→k_norm, position_bias→compressor.ape, weights_proj stays (structural: compressor.indexer → indexer.compressor) - Remove broken suffix renames (already fixed in prior commit) Model architecture fixes: - Patch deepseek_compressor.py to pass quant_config (was None, but NVFP4 checkpoint has quantized compressor weights with input_scale/weight_scale) - Patch deepseek_v4_attention.py indexer: weights_proj now uses quant_config (was None, but checkpoint has quantized weights) - Add indexer.compressor.fused_wkv_wgate stacking in load_weights Infrastructure: - Add deepseek_compressor.py to Dockerfile - Force MoE backend to flashinfer_cutedsl (was auto-selecting FLASHINFER_TRTLLM) - Update unit test to 50 cases (compressor + indexer + quantization scales) --- Dockerfile | 1 + docker-compose.yml | 1 + tests/test_nvfp4_mapper.py | 53 ++- vllm/patches/deepseek_v4.py | 26 +- vllm/patches/deepseek_v4_attention.py | 2 +- vllm/patches/layers/deepseek_compressor.py | 436 +++++++++++++++++++++ 6 files changed, 505 insertions(+), 14 deletions(-) create mode 100644 vllm/patches/layers/deepseek_compressor.py diff --git a/Dockerfile b/Dockerfile index 5fef5d56..69742eb6 100644 --- a/Dockerfile +++ b/Dockerfile @@ -37,6 +37,7 @@ ARG VLLM_LOADER_DIR=/usr/local/lib/python3.12/dist-packages/vllm/model_executor/ # Core model patches COPY vllm/patches/deepseek_v4.py ${VLLM_MODELS_DIR}/deepseek_v4.py COPY vllm/patches/deepseek_v4_attention.py ${VLLM_LAYERS_DIR}/deepseek_v4_attention.py +COPY vllm/patches/layers/deepseek_compressor.py ${VLLM_LAYERS_DIR}/deepseek_compressor.py # Config patches (add cutedsl to MoEBackend) ARG VLLM_CONFIG_DIR=/usr/local/lib/python3.12/dist-packages/vllm/config diff --git a/docker-compose.yml b/docker-compose.yml index 4aaee9d8..a11caef5 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -21,6 +21,7 @@ services: - --tool-call-parser=deepseek_v4 - --enable-auto-tool-choice - --reasoning-parser=deepseek_v4 + - --moe-backend=flashinfer_cutedsl - --gpu-memory-utilization=0.9 - --host=0.0.0.0 - --port=8000 diff --git a/tests/test_nvfp4_mapper.py b/tests/test_nvfp4_mapper.py index ef23cee2..aac97200 100644 --- a/tests/test_nvfp4_mapper.py +++ b/tests/test_nvfp4_mapper.py @@ -64,7 +64,21 @@ def _make_deepseek_v4_nvfp4_weights_mapper() -> WeightsMapper: suffix_renames = {} substr_renames = { + # === Compressor (non-indexer) NVFP4 renames === + "compressor.kv_proj.": "compressor.wkv.", + "compressor.gate_proj.": "compressor.wgate.", + "compressor.kv_norm.": "compressor.norm.", + "compressor.position_bias": "compressor.ape", + # === Attention compressor (before indexer renames) === ".self_attn.compressor.": ".attn.mla_attn.compressor.", + # === Indexer params === + "compressor.indexer.q_b_proj.": "indexer.wq_b.", + "compressor.indexer.weights_proj.": "indexer.weights_proj.", + "compressor.indexer.kv_norm.": "indexer.k_norm.", + "compressor.indexer.kv_proj.": "indexer.compressor.wkv.", + "compressor.indexer.gate_proj.": "indexer.compressor.wgate.", + "compressor.indexer.position_bias": "indexer.compressor.ape", + # === Attention projections === ".self_attn.q_a_proj.": ".attn.wq_a.", ".self_attn.kv_proj.": ".attn.wkv.", ".self_attn.q_b_proj.": ".attn.wq_b.", @@ -73,9 +87,11 @@ def _make_deepseek_v4_nvfp4_weights_mapper() -> WeightsMapper: ".self_attn.q_a_norm.": ".attn.q_a_norm.", ".self_attn.kv_norm.": ".attn.kv_norm.", ".self_attn.sinks": ".attn.sinks", + # Shared expert projections ".mlp.shared_experts.gate_proj.": ".ffn.shared_experts.w1.", ".mlp.shared_experts.up_proj.": ".ffn.shared_experts.w3.", ".mlp.shared_experts.down_proj.": ".ffn.shared_experts.down_proj.", + # General renames ".mlp.": ".ffn.", ".self_attn.": ".attn.", } @@ -107,7 +123,6 @@ TEST_CASES = [ ("layers.0.self_attn.q_a_proj.weight", "model.layers.0.attn.wq_a.weight"), ("layers.0.self_attn.q_a_proj.input_scale", "model.layers.0.attn.wq_a.input_scale"), ("layers.0.self_attn.kv_proj.weight", "model.layers.0.attn.wkv.weight"), - ("layers.0.self_attn.kv_proj.input_scale", "model.layers.0.attn.wkv.input_scale"), ("layers.0.self_attn.q_b_proj.weight", "model.layers.0.attn.wq_b.weight"), ("layers.0.self_attn.o_a_proj.weight", "model.layers.0.attn.wo_a.weight"), ("layers.0.self_attn.o_b_proj.weight", "model.layers.0.attn.wo_b.weight"), @@ -116,13 +131,31 @@ TEST_CASES = [ ("layers.0.self_attn.kv_norm.weight", "model.layers.0.attn.kv_norm.weight"), ("layers.0.self_attn.sinks", "model.layers.0.attn.sinks"), - # Attention — compressor (inside self_attn) - ("layers.0.self_attn.compressor.kv_proj.weight", "model.layers.0.attn.mla_attn.compressor.kv_proj.weight"), - ("layers.0.self_attn.compressor.kv_proj.input_scale", "model.layers.0.attn.mla_attn.compressor.kv_proj.input_scale"), - ("layers.0.self_attn.compressor.gate_proj.weight", "model.layers.0.attn.mla_attn.compressor.gate_proj.weight"), - ("layers.0.self_attn.compressor.gate_proj.input_scale", "model.layers.0.attn.mla_attn.compressor.gate_proj.input_scale"), - ("layers.0.self_attn.compressor.kv_norm.weight", "model.layers.0.attn.mla_attn.compressor.kv_norm.weight"), - ("layers.0.self_attn.compressor.position_bias", "model.layers.0.attn.mla_attn.compressor.position_bias"), + # Compressor (non-indexer): kv_proj → wkv, gate_proj → wgate + ("layers.0.self_attn.compressor.kv_proj.weight", "model.layers.0.attn.mla_attn.compressor.wkv.weight"), + ("layers.0.self_attn.compressor.kv_proj.input_scale", "model.layers.0.attn.mla_attn.compressor.wkv.input_scale"), + ("layers.0.self_attn.compressor.kv_proj.weight_scale", "model.layers.0.attn.mla_attn.compressor.wkv.weight_scale"), + ("layers.0.self_attn.compressor.kv_proj.weight_scale_2", "model.layers.0.attn.mla_attn.compressor.wkv.weight_scale_2"), + ("layers.0.self_attn.compressor.gate_proj.weight", "model.layers.0.attn.mla_attn.compressor.wgate.weight"), + ("layers.0.self_attn.compressor.gate_proj.input_scale", "model.layers.0.attn.mla_attn.compressor.wgate.input_scale"), + ("layers.0.self_attn.compressor.gate_proj.weight_scale", "model.layers.0.attn.mla_attn.compressor.wgate.weight_scale"), + ("layers.0.self_attn.compressor.gate_proj.weight_scale_2", "model.layers.0.attn.mla_attn.compressor.wgate.weight_scale_2"), + ("layers.0.self_attn.compressor.kv_norm.weight", "model.layers.0.attn.mla_attn.compressor.norm.weight"), + ("layers.0.self_attn.compressor.position_bias", "model.layers.0.attn.mla_attn.compressor.ape"), + + # Indexer own params + ("layers.10.self_attn.compressor.indexer.q_b_proj.weight", "model.layers.10.attn.mla_attn.indexer.wq_b.weight"), + ("layers.10.self_attn.compressor.indexer.q_b_proj.input_scale", "model.layers.10.attn.mla_attn.indexer.wq_b.input_scale"), + ("layers.10.self_attn.compressor.indexer.weights_proj.weight", "model.layers.10.attn.mla_attn.indexer.weights_proj.weight"), + ("layers.10.self_attn.compressor.indexer.weights_proj.input_scale", "model.layers.10.attn.mla_attn.indexer.weights_proj.input_scale"), + ("layers.10.self_attn.compressor.indexer.kv_norm.weight", "model.layers.10.attn.mla_attn.indexer.k_norm.weight"), + + # Indexer's compressor + ("layers.10.self_attn.compressor.indexer.kv_proj.weight", "model.layers.10.attn.mla_attn.indexer.compressor.wkv.weight"), + ("layers.10.self_attn.compressor.indexer.kv_proj.input_scale", "model.layers.10.attn.mla_attn.indexer.compressor.wkv.input_scale"), + ("layers.10.self_attn.compressor.indexer.gate_proj.weight", "model.layers.10.attn.mla_attn.indexer.compressor.wgate.weight"), + ("layers.10.self_attn.compressor.indexer.gate_proj.input_scale", "model.layers.10.attn.mla_attn.indexer.compressor.wgate.input_scale"), + ("layers.10.self_attn.compressor.indexer.position_bias", "model.layers.10.attn.mla_attn.indexer.compressor.ape"), # MoE gate ("layers.0.mlp.gate.tid2eid", "model.layers.0.ffn.gate.tid2eid"), @@ -135,15 +168,13 @@ TEST_CASES = [ ("layers.0.mlp.experts.0.gate_proj.input_scale", "model.layers.0.ffn.experts.0.w1.input_scale"), ("layers.0.mlp.experts.0.gate_proj.weight_scale", "model.layers.0.ffn.experts.0.w1.weight_scale"), ("layers.0.mlp.experts.0.gate_proj.weight_scale_2", "model.layers.0.ffn.experts.0.w1.weight_scale_2"), - ("layers.0.mlp.experts.255.down_proj.weight", "model.layers.0.ffn.experts.255.w2.weight"), - # Shared experts — gate_proj → w1, up_proj → w3, down_proj stays + # Shared experts ("layers.0.mlp.shared_experts.gate_proj.weight", "model.layers.0.ffn.shared_experts.w1.weight"), ("layers.0.mlp.shared_experts.up_proj.weight", "model.layers.0.ffn.shared_experts.w3.weight"), ("layers.0.mlp.shared_experts.down_proj.weight", "model.layers.0.ffn.shared_experts.down_proj.weight"), ("layers.0.mlp.shared_experts.gate_proj.input_scale", "model.layers.0.ffn.shared_experts.w1.input_scale"), ("layers.0.mlp.shared_experts.down_proj.weight_scale", "model.layers.0.ffn.shared_experts.down_proj.weight_scale"), - ("layers.0.mlp.shared_experts.down_proj.weight_scale_2", "model.layers.0.ffn.shared_experts.down_proj.weight_scale_2"), # Layer norm ("layers.0.post_attention_layernorm.weight", "model.layers.0.post_attention_layernorm.weight"), diff --git a/vllm/patches/deepseek_v4.py b/vllm/patches/deepseek_v4.py index 77f1d227..8b62ff0e 100644 --- a/vllm/patches/deepseek_v4.py +++ b/vllm/patches/deepseek_v4.py @@ -1451,6 +1451,9 @@ class DeepseekV4Model(nn.Module): ("attn.fused_wqa_wkv", "attn.wkv", 1), ("compressor.fused_wkv_wgate", "compressor.wkv", 0), ("compressor.fused_wkv_wgate", "compressor.wgate", 1), + # Indexer's compressor (same stacking pattern) + ("indexer.compressor.fused_wkv_wgate", "indexer.compressor.wkv", 0), + ("indexer.compressor.fused_wkv_wgate", "indexer.compressor.wgate", 1), ] params_dict = dict(self.named_parameters()) loaded_params: set[str] = set() @@ -1654,9 +1657,28 @@ def _make_deepseek_v4_nvfp4_weights_mapper() -> WeightsMapper: # NOTE: specific renames MUST come before general ones (applied in order) substr_renames = { - # Attention compressor (MUST come before .self_attn. → .attn.) + # === Compressor (non-indexer) NVFP4 renames === + # Checkpoint uses kv_proj/gate_proj, model uses wkv/wgate + # (for stacking into fused_wkv_wgate). + "compressor.kv_proj.": "compressor.wkv.", + "compressor.gate_proj.": "compressor.wgate.", + "compressor.kv_norm.": "compressor.norm.", + "compressor.position_bias": "compressor.ape", + # === Attention compressor (MUST come before .self_attn. → .attn. + # and before indexer renames so that .self_attn.compressor.indexer. + # becomes .attn.mla_attn.compressor.indexer.) === ".self_attn.compressor.": ".attn.mla_attn.compressor.", - # Attention projections (specific before .self_attn. → .attn.) + # === Indexer params (now under .attn.mla_attn.compressor.indexer. + # after the compressor rename above) === + # Indexer's own params + "compressor.indexer.q_b_proj.": "indexer.wq_b.", + "compressor.indexer.weights_proj.": "indexer.weights_proj.", + "compressor.indexer.kv_norm.": "indexer.k_norm.", + # Indexer's compressor (compressor.indexer → indexer.compressor) + "compressor.indexer.kv_proj.": "indexer.compressor.wkv.", + "compressor.indexer.gate_proj.": "indexer.compressor.wgate.", + "compressor.indexer.position_bias": "indexer.compressor.ape", + # === Attention projections (specific before .self_attn. → .attn.) === ".self_attn.q_a_proj.": ".attn.wq_a.", ".self_attn.kv_proj.": ".attn.wkv.", ".self_attn.q_b_proj.": ".attn.wq_b.", diff --git a/vllm/patches/deepseek_v4_attention.py b/vllm/patches/deepseek_v4_attention.py index bfe8741b..ca565a85 100644 --- a/vllm/patches/deepseek_v4_attention.py +++ b/vllm/patches/deepseek_v4_attention.py @@ -1106,7 +1106,7 @@ class DeepseekV4Indexer(nn.Module): hidden_size, self.n_head, bias=False, - quant_config=None, + quant_config=quant_config, prefix=f"{prefix}.weights_proj", ) self.k_norm = LayerNorm(self.head_dim, eps=1e-6) diff --git a/vllm/patches/layers/deepseek_compressor.py b/vllm/patches/layers/deepseek_compressor.py new file mode 100644 index 00000000..181ac58d --- /dev/null +++ b/vllm/patches/layers/deepseek_compressor.py @@ -0,0 +1,436 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from dataclasses import dataclass +from typing import Any, ClassVar, cast + +import torch +from torch import nn + +from vllm.config import VllmConfig, get_current_vllm_config +from vllm.forward_context import get_forward_context +from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase +from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.model_executor.layers.linear import ( + MergedColumnParallelLinear, +) +from vllm.platforms import current_platform +from vllm.triton_utils import tl, triton +from vllm.v1.attention.backend import ( + AttentionBackend, + AttentionCGSupport, + AttentionMetadataBuilder, + CommonAttentionMetadata, + MultipleOf, +) +from vllm.v1.attention.ops.deepseek_v4_ops.fused_compress_quant_cache import ( + _fused_kv_compress_norm_rope_insert_indexer_attn, + _fused_kv_compress_norm_rope_insert_indexer_mxfp4_attn, + _fused_kv_compress_norm_rope_insert_sparse_attn, +) +from vllm.v1.attention.ops.deepseek_v4_ops.fused_indexer_q import ( + MXFP4_BLOCK_SIZE, +) +from vllm.v1.kv_cache_interface import ( + KVCacheSpec, + MLAAttentionSpec, + SlidingWindowMLASpec, +) + + +class CompressorBackend(AttentionBackend): + def __init__(self): + super().__init__() + + @staticmethod + def get_name() -> str: + return "CompressorBackend" + + @staticmethod + def get_supported_kernel_block_sizes() -> list[int | MultipleOf]: + return [MultipleOf(1)] + + @classmethod + def get_supported_head_sizes(cls) -> list[int]: + return [512, 1024] + + @staticmethod + def get_builder_cls() -> type["CompressorMetadataBuilder"]: + return CompressorMetadataBuilder + + @staticmethod + def get_kv_cache_shape( + num_blocks: int, + block_size: int, + num_kv_heads: int, + head_size: int, + cache_dtype_str: str = "auto", + ) -> tuple[int, ...]: + assert num_kv_heads == 1 + return (num_blocks, block_size, head_size) + + @staticmethod + def get_kv_cache_stride_order( + include_num_layers_dimension: bool = False, + ) -> tuple[int, ...]: + if include_num_layers_dimension: + return (0, 1, 2, 3) + return (0, 1, 2) + + +@dataclass +class CompressorMetadata: + block_table: torch.Tensor + slot_mapping: torch.Tensor + block_size: int + + token_to_req_indices: torch.Tensor | None = None # [num_tokens] + + +class CompressorMetadataBuilder(AttentionMetadataBuilder): + _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.ALWAYS + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + assert isinstance(self.kv_cache_spec, SlidingWindowMLASpec | MLAAttentionSpec) + mla_spec = cast(SlidingWindowMLASpec | MLAAttentionSpec, self.kv_cache_spec) + self.block_size = mla_spec.block_size + + self.token_to_req_indices = torch.zeros( + self.vllm_config.scheduler_config.max_num_batched_tokens, + dtype=torch.int32, + device=self.device, + ) + + def build( + self, + common_prefix_len: int, + common_attn_metadata: CommonAttentionMetadata, + fast_build: bool = False, + ) -> CompressorMetadata: + query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu + num_reqs = common_attn_metadata.num_reqs + query_lens = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1] + x = torch.repeat_interleave(torch.arange(num_reqs), query_lens).pin_memory() + token_to_req_indices = self.token_to_req_indices[: x.shape[0]] + token_to_req_indices.copy_(x, non_blocking=True) + return CompressorMetadata( + block_table=common_attn_metadata.block_table_tensor.clamp_(min=0), + slot_mapping=common_attn_metadata.slot_mapping, + block_size=self.block_size, + token_to_req_indices=token_to_req_indices, + ) + + +class CompressorStateCache(torch.nn.Module, AttentionLayerBase): + def __init__( + self, + state_dim: int, + dtype: torch.dtype, + compress_ratio: int, + prefix: str, + ): + super().__init__() + self.state_dim = state_dim + self.dtype = dtype + self.prefix = prefix + self.kv_cache = torch.tensor([]) + compilation_config = get_current_vllm_config().compilation_config + if prefix in compilation_config.static_forward_context: + raise ValueError(f"Duplicate layer name: {prefix}") + compilation_config.static_forward_context[prefix] = self + + assert self.dtype == torch.float32 + assert compress_ratio in [4, 128] + coff = 1 + (compress_ratio == 4) + self.sliding_window = coff * compress_ratio + # Block size is constrained by tensor sharing between compressor states + # and KV blocks. Since compressor states share the same physical tensor + # as KV blocks, they must use the same page size. + # The KV block shape [256//4, head_dim] = [64, 584] determines: + # - C4 compressor block shape [4, 2*512*2*4] -> block_size = 4 + # - C128 compressor block shape [8, 512*2*4] -> block_size = 8 + # TODO(yifan): make block size automatically determined and configurable. + if compress_ratio == 4: + self.block_size = 4 + elif compress_ratio == 128: + self.block_size = 8 + else: + raise ValueError(f"Invalid compress ratio: {compress_ratio}") + + def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec: + return SlidingWindowMLASpec( # only has one vector instead of K + V + block_size=self.block_size, + num_kv_heads=1, + head_size=self.state_dim, + dtype=self.dtype, + sliding_window=self.sliding_window, + alignment=576, # NOTE: FlashMLA requires 576B alignment + ) + + def forward(self): ... + + def get_attn_backend(self) -> type[AttentionBackend]: + return CompressorBackend + + +class DeepseekCompressor(nn.Module): + def __init__( + self, + vllm_config: VllmConfig, + compress_ratio: int, + hidden_size: int, + head_dim: int, + rotate: bool = False, + prefix: str = "", + k_cache_prefix="", + use_fp4_cache: bool = False, + ): + super().__init__() + self.compress_ratio = compress_ratio + self.hidden_size = hidden_size + self.head_dim = head_dim + self.rotate = rotate + self.prefix = prefix + self.k_cache_prefix = k_cache_prefix + self.use_fp4_cache = use_fp4_cache + + config = vllm_config.model_config.hf_config + self.rope_head_dim = config.qk_rope_head_dim + self.nope_head_dim = self.head_dim - self.rope_head_dim + self.rms_norm_eps = config.rms_norm_eps + self.device = current_platform.device_type + self.max_num_reqs = vllm_config.scheduler_config.max_num_seqs + self.max_model_len = vllm_config.model_config.max_model_len + + self.overlap = compress_ratio == 4 + self.coff = 1 + self.overlap + + state_dtype = torch.float32 + self.ape = nn.Parameter( + torch.empty( + (compress_ratio, self.coff * self.head_dim), + dtype=state_dtype, + device=self.device, + ), + requires_grad=False, + ) + + quant_config = vllm_config.quant_config + + self.fused_wkv_wgate = MergedColumnParallelLinear( + self.hidden_size, + [self.coff * self.head_dim, self.coff * self.head_dim], + bias=False, + return_bias=False, + quant_config=quant_config, + disable_tp=True, + prefix=f"{prefix}.fused_wkv_wgate", + ) + self.norm = RMSNorm(self.head_dim, self.rms_norm_eps) + + self.state_cache = CompressorStateCache( + state_dim=2 * self.coff * self.head_dim, # kv_state + score_state + dtype=state_dtype, + compress_ratio=compress_ratio, + prefix=f"{prefix}.state_cache", + ) + + # Save reference to static_forward_context for forward-time KV cache lookup. + # get_current_vllm_config() is only available during __init__, not forward. + self._static_forward_context = ( + vllm_config.compilation_config.static_forward_context + ) + + if self.head_dim == 512: + assert not use_fp4_cache, ( + "MXFP4 cache is only supported for indexer (head=128)" + ) + self._fused_kernel = _fused_kv_compress_norm_rope_insert_sparse_attn + self._quant_block = 64 + self._token_stride = self.nope_head_dim + self.rope_head_dim * 2 + self._scale_dim = self.nope_head_dim // 64 + 1 # 7 real + 1 pad + self._num_warps = 4 + elif self.head_dim == 128: + if use_fp4_cache: + self._fused_kernel = ( + _fused_kv_compress_norm_rope_insert_indexer_mxfp4_attn + ) + self._quant_block = MXFP4_BLOCK_SIZE + self._token_stride = self.head_dim // 2 + self._scale_dim = self.head_dim // MXFP4_BLOCK_SIZE + else: + self._fused_kernel = _fused_kv_compress_norm_rope_insert_indexer_attn + self._quant_block = 128 + self._token_stride = self.head_dim + self._scale_dim = 4 # single float32 scale + self._num_warps = 1 + else: + raise ValueError( + f"Unsupported head_dim for fused quant+cache: {self.head_dim}" + ) + + def forward( + self, + # [num_tokens, 2 * self.coff * self.head_dim] + kv_score: torch.Tensor, + # [num_tokens] + positions: torch.Tensor, + rotary_emb, + ) -> None: + # Each of shape [num_tokens, coff * self.head_dim] + # input bf16, output are fp32 + kv, score = kv_score.split( + [self.coff * self.head_dim, self.coff * self.head_dim], dim=-1 + ) + + # Get the metadata and handle dummy profiling run. + attn_metadata = get_forward_context().attn_metadata + if not isinstance(attn_metadata, dict): + return + + state_metadata = cast( + CompressorMetadata, attn_metadata[self.state_cache.prefix] + ) + token_to_req_indices = state_metadata.token_to_req_indices + slot_mapping = state_metadata.slot_mapping + num_actual = slot_mapping.shape[0] + block_table = state_metadata.block_table + block_size = state_metadata.block_size + + # [num_blocks, block_size, kv_dim+score_dim], where kv_dim == score_dim + state_cache = self.state_cache.kv_cache + # kv_state stored in first half, score_state stored in second half + state_width = state_cache.shape[-1] // 2 + pdl_kwargs = {} if current_platform.is_rocm() else {"launch_pdl": False} + + # Store the KV and score (with fused APE addition) in the state. + # NOTE: PDL is disabled — both this kernel and _fused_kernel below + # depend on preceding kernel outputs (kv/score from the cublas GEMM; + # state_cache from this kernel) but neither emits/waits on PDL grid + # dependency primitives, so launch_pdl=True caused a read-after-write + # race and non-deterministic output. + _save_partial_states_kernel[(num_actual,)]( + kv, + kv.stride(0), + score, + score.stride(0), + self.ape, + self.ape.stride(0), + positions, + state_cache, + state_cache.stride(0), + state_cache.stride(1), + slot_mapping, + block_size, + HEAD_SIZE=kv.shape[-1], + TRITON_BLOCK_SIZE=triton.next_power_of_2(kv.shape[-1]), + STATE_WIDTH=state_width, + COMPRESS_RATIO=self.compress_ratio, + **pdl_kwargs, + ) + + # Fused: compress → RMSNorm → RoPE → FP8 quant → KV cache write. + # RoPE requirements (kernel applies forward GPT-J style rotation): + # - is_neox_style=False (interleaved pairs, NOT split-half) + # - cos_sin_cache layout: [max_pos, rope_head_dim] with first half cos, + # second half sin (per-pair, length rope_head_dim // 2 each) + # - applied to LAST rope_head_dim elements of head_dim + # - position used: (positions // compress_ratio) * compress_ratio + cos_sin_cache = rotary_emb.cos_sin_cache + k_cache_metadata = cast(Any, attn_metadata[self.k_cache_prefix]) + kv_cache = self._static_forward_context[self.k_cache_prefix].kv_cache + + self._fused_kernel[(num_actual,)]( + # state cache + state_cache, + state_cache.stride(0), + state_cache.stride(1), + # metadata + token_to_req_indices, + positions, + slot_mapping, + block_table, + block_table.stride(0), + block_size, + # RMSNorm + self.norm.weight, + self.rms_norm_eps, + # RoPE + cos_sin_cache, + cos_sin_cache.stride(0), + # KV cache + kv_cache, + k_cache_metadata.slot_mapping, + kv_cache.shape[1], # paged KV cache block size (tokens per block) + # constexprs + HEAD_SIZE=self.head_dim, + TRITON_BLOCK_SIZE=triton.next_power_of_2(self.head_dim), + STATE_WIDTH=state_width, + COMPRESS_RATIO=self.compress_ratio, + OVERLAP=self.overlap, + ROPE_HEAD_DIM=self.rope_head_dim, + FP8_MAX=448.0, + QUANT_BLOCK=self._quant_block, + TOKEN_STRIDE=self._token_stride, + SCALE_DIM=self._scale_dim, + KV_BLOCK_STRIDE=kv_cache.stride(0), + num_warps=self._num_warps, + **pdl_kwargs, + ) + + +@triton.jit +def _save_partial_states_kernel( + kv_ptr, + kv_stride, + score_ptr, + score_stride, + ape_ptr, + ape_stride, + positions_ptr, + state_cache_ptr, + state_cache_stride0, + state_cache_stride1, + slot_mapping_ptr, + block_size, + HEAD_SIZE: tl.constexpr, + TRITON_BLOCK_SIZE: tl.constexpr, + # state_cache last dim packs [kv_state, score_state], each STATE_WIDTH wide. + STATE_WIDTH: tl.constexpr, + COMPRESS_RATIO: tl.constexpr, +): + token_idx = tl.program_id(0) + slot_id = tl.load(slot_mapping_ptr + token_idx) + + # Skip padded / invalid tokens (slot_id == -1 is the PAD sentinel used + # by vLLM). During CUDA graph replay the batch may contain padding + # tokens whose slot_mapping is -1; writing to kv_state[-1] would be an + # illegal memory access. + if slot_id < 0: + return + + block_idx = slot_id // block_size + pos_in_block = slot_id % block_size + base_ptr = ( + state_cache_ptr + + block_idx * state_cache_stride0 + + pos_in_block * state_cache_stride1 + ) + + block = tl.arange(0, TRITON_BLOCK_SIZE) + mask = block < HEAD_SIZE + + kv = tl.load(kv_ptr + token_idx * kv_stride + block, mask=mask) + tl.store(base_ptr + block, kv, mask=mask) + + # Fused: score += ape[position % compress_ratio] + position = tl.load(positions_ptr + token_idx) + ape_row = position % COMPRESS_RATIO + ape = tl.load(ape_ptr + ape_row * ape_stride + block, mask=mask) + score = tl.load(score_ptr + token_idx * score_stride + block, mask=mask) + tl.store( + base_ptr + STATE_WIDTH + block, + score + ape, + mask=mask, + )