A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 2, 19:30 UTC
[DSv4] Drop _get_compressed_kv_buffer in DeepseekCompressor (#43690)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
Woosuk Kwon committed
May 26, 2026 at 10:11 UTC
193ce8812eb4b43fc48039757072c71ac2eb0a28
1 file changed
+5
-33
vllm/models/deepseek_v4/compressor.py
+5
-33
@@ -173,33 +173,6 @@ class CompressorStateCache(torch.nn.Module, AttentionLayerBase):
173
174
175
class DeepseekCompressor(nn.Module):
176
- _compressed_kv_buffers: ClassVar[dict[tuple[str, int, int], torch.Tensor]] = {}
177
-
178
- @classmethod
179
- def _get_compressed_kv_buffer(
180
- cls,
181
- device: str,
182
- max_num_tokens: int,
183
- head_dim: int,
184
- ) -> torch.Tensor:
185
- if device == "cuda" and torch.accelerator.is_available():
186
- device_key = f"cuda:{torch.accelerator.current_device_index()}"
187
- alloc_device = torch.device(device_key)
188
- else:
189
- device_key = str(device)
190
- alloc_device = torch.device(device)
191
-
192
- key = (device_key, max_num_tokens, head_dim)
193
- buffer = cls._compressed_kv_buffers.get(key)
194
- if buffer is None:
195
- buffer = torch.empty(
196
- (max_num_tokens, head_dim),
197
- dtype=torch.float32,
198
- device=alloc_device,
199
- )
200
- cls._compressed_kv_buffers[key] = buffer
201
- return buffer
202
-
176
def __init__(
177
self,
178
vllm_config: VllmConfig,
@@ -276,11 +249,6 @@ class DeepseekCompressor(nn.Module):
249
self._fused_sparse_kernel = (
250
_fused_kv_compress_norm_rope_insert_sparse_attn_cutedsl
251
)
279
- self._compressed_kv_buffer = self._get_compressed_kv_buffer(
280
- self.device,
281
- vllm_config.scheduler_config.max_num_batched_tokens,
282
- self.head_dim,
283
- )
252
self._quant_block = 64
253
self._token_stride = self.nope_head_dim + self.rope_head_dim * 2
254
self._scale_dim = self.nope_head_dim // 64 + 1 # 7 real + 1 pad
@@ -407,7 +375,11 @@ class DeepseekCompressor(nn.Module):
375
overlap=self.overlap,
376
)
377
else:
410
- compressed_kv = self._compressed_kv_buffer[:num_actual]
378
+ compressed_kv = torch.empty(
379
+ (num_actual, self.head_dim),
380
+ dtype=torch.float32,
381
+ device=state_cache.device,
382
+ )
383
self._compress_kernel(
384
state_cache,
385
token_to_req_indices,