Skip to content

vllm.models.deepseek_v4.cpu.cpu_compressor

CPU DeepSeek-V4 compressor subclass and dispatch: partial-state-cache write and the compress -> RMSNorm -> RoPE -> FP8 quant -> KV cache store step, calling torch.ops._C.save_partial_states_cpu/ compress_norm_rope_store_cpu/compress_norm_rope_store_indexer_cpu (csrc/cpu/sgl-kernels/compressor.cpp) in place of triton-cpu.

Covers both the head_dim=512 main-attention compressor path and the head_dim=128 indexer compressor path (fp8 only -- MXFP4, use_fp4_cache, never occurs on CPU: the indexer's MXFP4 cache requires a Blackwell GPU).

Classes:

DeepseekV4CPUCompressor

Bases: DeepseekCompressor

CPU compressor: same state/weights as the shared base, but forward's cache-write and compress->RMSNorm->RoPE->quant->store steps always dispatch straight to the ported CPU kernels above instead of through the shared method's platform-dispatch chain. This class only ever runs on CPU, and use_fp4_cache is always False here (the indexer's MXFP4 cache requires a Blackwell GPU), so the shared method's MXFP4/two-stage/cutedsl/triton branches never apply.

Methods:

Source code in vllm/models/deepseek_v4/cpu/cpu_compressor.py 
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
class DeepseekV4CPUCompressor(DeepseekCompressor):
    """CPU compressor: same state/weights as the shared base, but
    ``forward``'s cache-write and compress->RMSNorm->RoPE->quant->store
    steps always dispatch straight to the ported CPU kernels above instead
    of through the shared method's platform-dispatch chain. This class
    only ever runs on CPU, and ``use_fp4_cache`` is always ``False`` here
    (the indexer's MXFP4 cache requires a Blackwell GPU), so the shared
    method's MXFP4/two-stage/cutedsl/triton branches never apply.
    """

    _norm_weight_fp32: torch.Tensor

    def cache_norm_weight_fp32(self) -> None:
        """Cache a fp32 contiguous copy of ``norm.weight`` once, right after
        weight loading (see
        ``DeepseekV4CPUAttention.process_weights_after_loading``) -- the CPU
        kernel requires fp32/contiguous but the checkpoint loads this weight
        in bf16, and it never changes after loading."""
        self._norm_weight_fp32 = self.norm.weight.to(torch.float32).contiguous()

    def forward(
        self,
        kv_score: torch.Tensor,
        positions: torch.Tensor,
        rotary_emb,
    ) -> None:
        kv, score = kv_score.split(
            [self.coff * self.head_dim, self.coff * self.head_dim], dim=-1
        )

        attn_metadata = get_forward_context().attn_metadata
        if not isinstance(attn_metadata, dict):
            return

        state_metadata = attn_metadata[self.state_cache.prefix]
        assert isinstance(state_metadata, CompressorMetadata)
        token_to_req_indices = state_metadata.token_to_req_indices
        assert token_to_req_indices is not None
        slot_mapping = state_metadata.slot_mapping
        num_actual = slot_mapping.shape[0]
        block_table = state_metadata.block_table
        block_size = state_metadata.block_size

        state_cache = self.state_cache.kv_cache

        save_partial_states_cpu(
            kv, score, self.ape, positions, state_cache, slot_mapping
        )

        cos_sin_cache = rotary_emb.cos_sin_cache
        k_cache_metadata = cast(Any, attn_metadata[self.k_cache_prefix])
        k_cache_layer = self._static_forward_context[self.k_cache_prefix]
        kv_cache = k_cache_layer.kv_cache

        assert self.head_dim in (512, 128), (
            f"unsupported compressor head_dim: {self.head_dim}"
        )
        window = (1 + self.overlap) * self.compress_ratio
        win_offsets = torch.arange(window, device=positions.device)
        win_local = positions.unsqueeze(-1) - window + 1 + win_offsets
        req_idx = token_to_req_indices[:num_actual]
        gather_slots = map_local_to_global_slots_cpu(
            win_local, req_idx, block_table, block_size
        )

        kv_cache_2d = kv_cache.view(kv_cache.shape[0], -1)
        kv_cache_block_size = kv_cache.shape[1]

        store_fn = (
            compress_norm_rope_store_cpu
            if self.head_dim == 512
            else compress_norm_rope_store_indexer_cpu
        )
        store_fn(
            state_cache,
            gather_slots,
            positions,
            k_cache_metadata.slot_mapping,
            self._norm_weight_fp32,
            self.rms_norm_eps,
            cos_sin_cache,
            kv_cache_2d,
            kv_cache_block_size,
            self.compress_ratio,
        )

cache_norm_weight_fp32()

Cache a fp32 contiguous copy of norm.weight once, right after weight loading (see DeepseekV4CPUAttention.process_weights_after_loading) -- the CPU kernel requires fp32/contiguous but the checkpoint loads this weight in bf16, and it never changes after loading.

Source code in vllm/models/deepseek_v4/cpu/cpu_compressor.py 
40
41
42
43
44
45
46
def cache_norm_weight_fp32(self) -> None:
    """Cache a fp32 contiguous copy of ``norm.weight`` once, right after
    weight loading (see
    ``DeepseekV4CPUAttention.process_weights_after_loading``) -- the CPU
    kernel requires fp32/contiguous but the checkpoint loads this weight
    in bf16, and it never changes after loading."""
    self._norm_weight_fp32 = self.norm.weight.to(torch.float32).contiguous()