@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 1aad125815

[CPU] Enable chunked prefill and prefix caching for qwen3.5 (#46202)

Signed-off-by: Li, Tianmu <tianmu.li@intel.com> Co-authored-by: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com> Co-authored-by: Li, Jiang <jiang1.li@intel.com>

Tianmu Li committed Jun 24, 2026 at 20:49 UTC 1aad1258157b9327a1510f3889b6983d6f10004e
7 files changed +388 -32
.buildkite/hardware_tests/cpu.yaml
+4 -2
@@ -53,7 +53,7 @@ steps:
53 - tests/models/language/pooling/
54 commands:
55 - |
56 - bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 40m "
56 + bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 50m "
57 pytest -x -v -s tests/models/language/generation -m cpu_model
58 pytest -x -v -s tests/models/language/pooling -m cpu_model"
59
@@ -68,13 +68,15 @@ steps:
68 - vllm/v1/sample/ops/topk_topp_triton.py
69 - vllm/v1/sample/ops/topk_topp_sampler.py
70 - tests/v1/sample/test_topk_topp_sampler.py
71 + - tests/v1/e2e/test_cpu_linear_attn_chunked_prefix.py
72 commands:
73 - |
74 bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 45m "
75 uv pip install git+https://github.com/triton-lang/triton-cpu.git@270e696d
76 VLLM_USE_V2_MODEL_RUNNER=1 pytest -x -v -s tests/models/language/generation/test_granite.py -m cpu_model
77 # TODO: move to CPU-Kernel Tests once triton-cpu has a pre-built wheel
77 - pytest -x -v -s tests/v1/sample/test_topk_topp_sampler.py::TestTritonTopkTopp"
78 + pytest -x -v -s tests/v1/sample/test_topk_topp_sampler.py::TestTritonTopkTopp
79 + pytest -x -v -s tests/v1/e2e/test_cpu_linear_attn_chunked_prefix.py"
80
81 - label: CPU-Quantization Model Tests
82 depends_on: []
csrc/cpu/sgl-kernels/conv.cpp
+14 -12
@@ -289,19 +289,18 @@ void causal_conv1d_fwd_kernel_impl(
289 }
290 }
291
292 -#define LAUNCH_TINYGEMM_VARLEN_KERNEL(K, NB_SIZE) \
293 - tinygemm_kernel<scalar_t, K, NB_SIZE, has_bias, has_silu>::apply( \
294 - input + batch_offset * dim + mb_start * dim + nb_start, \
295 - weight + nb_start * width, \
296 - out + batch_offset * dim + mb_start * dim + nb_start, \
297 - has_bias ? bias + nb_start : nullptr, \
298 - nullptr, \
299 - false, \
300 - mb_size, \
301 - dim, \
292 +#define LAUNCH_TINYGEMM_VARLEN_KERNEL(K, NB_SIZE) \
293 + tinygemm_kernel<scalar_t, K, NB_SIZE, has_bias, has_silu>::apply( \
294 + input + batch_offset * dim + mb_start * dim + nb_start, \
295 + weight + nb_start * width, \
296 + out + batch_offset * dim + mb_start * dim + nb_start, \
297 + has_bias ? bias + nb_start : nullptr, \
298 + has_conv_states ? conv_states + conv_state_index * conv_state_slot_stride + nb_start : nullptr, \
299 + has_initial_states_value, \
300 + mb_size, \
301 + dim, \
302 mb_start == 0);
303
304 -// TODO: add `has_initial_state` support for varlen kernel
304 template <typename scalar_t>
305 void causal_conv1d_fwd_varlen_kernel_impl(
306 scalar_t* __restrict__ out,
@@ -343,6 +342,9 @@ void causal_conv1d_fwd_varlen_kernel_impl(
342 int64_t nb_start = nb * BLOCK_N;
343 int64_t nb_size = std::min(dim - nb_start, BLOCK_N);
344
345 + const bool has_initial_states_value = has_conv_states ? has_initial_state[bs] : false;
346 + int32_t conv_state_index = has_conv_indices ? conv_indices[bs] : bs;
347 +
348 switch (width << 4 | nb_size >> 4) {
349 case 0x42:
350 LAUNCH_TINYGEMM_VARLEN_KERNEL(4, 32);
@@ -373,7 +375,7 @@ void causal_conv1d_fwd_varlen_kernel_impl(
375 width,
376 dim,
377 seqlen,
376 - /* has_initial_state */ false);
378 + has_initial_state[bs]);
379 }
380 });
381 }
tests/kernels/mamba/cpu/test_cpu_gdn_ops.py
+224
@@ -25,6 +25,8 @@ HEAD_DIMS = [
25 (64, 32),
26 ]
27 CHUNK_SIZE = 64
28 +CONV_DIM = 128
29 +CONV_KERNEL = 4
30 PREFILL_SEQ_LENS = [
31 [1],
32 [1, 2, 3],
@@ -312,3 +314,225 @@ def test_chunk_gated_delta_rule_cpu(
314 atol=1e-2,
315 rtol=1e-2,
316 )
317 +
318 +
319 +# (total_tokens, split) pairs mimicking where chunked prefill breaks a sequence
320 +# across two scheduler steps: chunk-aligned and non-aligned splits.
321 +TWO_CALL_SPLITS = [
322 + (2 * CHUNK_SIZE, CHUNK_SIZE),
323 + (2 * CHUNK_SIZE + 17, CHUNK_SIZE),
324 + (2 * CHUNK_SIZE + 17, CHUNK_SIZE + 9),
325 + (4 * CHUNK_SIZE + 17, 2 * CHUNK_SIZE),
326 + (3 * CHUNK_SIZE, CHUNK_SIZE + 1),
327 +]
328 +
329 +
330 +@pytest.mark.parametrize("total_tokens, split", TWO_CALL_SPLITS)
331 +@pytest.mark.parametrize("num_heads", NUM_HEADS)
332 +@pytest.mark.parametrize("head_dims", HEAD_DIMS)
333 +@torch.inference_mode()
334 +def test_chunk_gated_delta_rule_cpu_two_call_split(
335 + total_tokens: int,
336 + split: int,
337 + num_heads: tuple[int, int],
338 + head_dims: tuple[int, int],
339 +) -> None:
340 + """A prefill split into two calls (the second seeded with the first's
341 + ``final_state`` and a rebased ``cu_seqlens``) must match the single-call
342 + result, mimicking the cross-scheduler-step handoff in
343 + ``cpu_gdn_attention_core``.
344 + """
345 + q, k, v, a, b, A_log, dt_bias = gdn_inputs(
346 + num_tokens=total_tokens,
347 + num_heads=num_heads,
348 + head_dims=head_dims,
349 + )
350 + _, num_v_heads = num_heads
351 + head_dim, v_head_dim = head_dims
352 +
353 + g, beta = ref_gdn_gating(A_log, a, b, dt_bias)
354 + g = g.unsqueeze(0) # [1, T, HV]
355 + beta = beta.unsqueeze(0)
356 +
357 + zero_state = torch.zeros(1, num_v_heads, head_dim, v_head_dim, dtype=torch.float32)
358 +
359 + # Reference: whole sequence in one call, no initial state.
360 + out_full, final_full = ops.chunk_gated_delta_rule_cpu(
361 + query=q,
362 + key=k,
363 + value=v,
364 + g=g,
365 + beta=beta,
366 + initial_state=zero_state,
367 + output_final_state=True,
368 + cu_seqlens=torch.tensor([0, total_tokens], dtype=torch.int32),
369 + head_first=False,
370 + use_qk_l2norm_in_kernel=True,
371 + )
372 +
373 + # Call 1: tokens [0:split], no initial state, capture final state.
374 + out1, state1 = ops.chunk_gated_delta_rule_cpu(
375 + query=q[:, :split],
376 + key=k[:, :split],
377 + value=v[:, :split],
378 + g=g[:, :split],
379 + beta=beta[:, :split],
380 + initial_state=zero_state,
381 + output_final_state=True,
382 + cu_seqlens=torch.tensor([0, split], dtype=torch.int32),
383 + head_first=False,
384 + use_qk_l2norm_in_kernel=True,
385 + )
386 + # Call 2: tokens [split:T] seeded with call 1's final state and a cu_seqlens
387 + # rebased to start at 0, as cpu_gdn_attention_core continues a prefill chunk.
388 + tail = total_tokens - split
389 + out2, state2 = ops.chunk_gated_delta_rule_cpu(
390 + query=q[:, split:],
391 + key=k[:, split:],
392 + value=v[:, split:],
393 + g=g[:, split:],
394 + beta=beta[:, split:],
395 + initial_state=state1.to(torch.float32),
396 + output_final_state=True,
397 + cu_seqlens=torch.tensor([0, tail], dtype=torch.int32),
398 + head_first=False,
399 + use_qk_l2norm_in_kernel=True,
400 + )
401 +
402 + out_split = torch.cat([out1, out2], dim=1)
403 +
404 + # State must be near-exact; output allows a looser bound for the bf16 round-trip.
405 + torch.testing.assert_close(state2, final_full, atol=1e-3, rtol=1e-3)
406 + torch.testing.assert_close(out_split, out_full, atol=2e-2, rtol=2e-2)
407 +
408 +
409 +def _conv_inputs(total_tokens: int):
410 + x = tensor_cache(total_tokens * CONV_DIM, torch.bfloat16).view(
411 + total_tokens, CONV_DIM
412 + )
413 + weight = tensor_cache(CONV_DIM * CONV_KERNEL, torch.bfloat16).view(
414 + CONV_DIM, CONV_KERNEL
415 + )
416 + bias = tensor_cache(CONV_DIM, torch.bfloat16)
417 + return x, weight, bias
418 +
419 +
420 +@pytest.mark.parametrize("total_tokens, split", TWO_CALL_SPLITS)
421 +@torch.inference_mode()
422 +def test_causal_conv1d_torch_two_call_split(total_tokens: int, split: int) -> None:
423 + """Non-AMX conv-state handoff: a two-call split (the second seeded via
424 + ``has_initial_state=True`` from the conv_states the first wrote back) must
425 + match the single-call result.
426 + """
427 + from vllm.model_executor.layers.mamba.ops.cpu.causal_conv1d import (
428 + causal_conv1d_torch,
429 + )
430 +
431 + x, weight, bias = _conv_inputs(total_tokens)
432 + state_len = CONV_KERNEL - 1
433 + # [num_slots, conv_dim, state_len]; slot 0 used here.
434 + conv_states_full = torch.zeros(1, CONV_DIM, state_len, dtype=x.dtype)
435 + conv_states_split = torch.zeros(1, CONV_DIM, state_len, dtype=x.dtype)
436 +
437 + # x is [conv_dim, T] for causal_conv1d_torch.
438 + xt = x.transpose(0, 1).contiguous()
439 +
440 + out_full = causal_conv1d_torch(
441 + x=xt,
442 + weight=weight,
443 + bias=bias,
444 + conv_states=conv_states_full,
445 + query_start_loc=torch.tensor([0, total_tokens], dtype=torch.int32),
446 + cache_indices=torch.tensor([0], dtype=torch.int32),
447 + has_initial_state=torch.tensor([False]),
448 + activation="silu",
449 + )
450 +
451 + out1 = causal_conv1d_torch(
452 + x=xt[:, :split],
453 + weight=weight,
454 + bias=bias,
455 + conv_states=conv_states_split,
456 + query_start_loc=torch.tensor([0, split], dtype=torch.int32),
457 + cache_indices=torch.tensor([0], dtype=torch.int32),
458 + has_initial_state=torch.tensor([False]),
459 + activation="silu",
460 + )
461 + out2 = causal_conv1d_torch(
462 + x=xt[:, split:],
463 + weight=weight,
464 + bias=bias,
465 + conv_states=conv_states_split,
466 + query_start_loc=torch.tensor([0, total_tokens - split], dtype=torch.int32),
467 + cache_indices=torch.tensor([0], dtype=torch.int32),
468 + has_initial_state=torch.tensor([True]),
469 + activation="silu",
470 + )
471 + out_split = torch.cat([out1, out2], dim=1)
472 +
473 + torch.testing.assert_close(out_split, out_full, atol=1e-2, rtol=1e-2)
474 +
475 +
476 +@pytest.mark.skipif(
477 + not torch.cpu._is_amx_tile_supported(),
478 + reason="causal_conv1d_fwd_cpu requires AMX/AVX512",
479 +)
480 +@pytest.mark.parametrize("total_tokens, split", TWO_CALL_SPLITS)
481 +@torch.inference_mode()
482 +def test_causal_conv1d_fwd_cpu_two_call_split(total_tokens: int, split: int) -> None:
483 + """AMX prefill conv op must honor ``has_initial_state`` so a two-call split
484 + matches the single-call result.
485 +
486 + Regression test for ``causal_conv1d_fwd_varlen_kernel_impl`` (``conv.cpp``)
487 + ignoring the carried conv state on continued chunks.
488 + """
489 + state_len = CONV_KERNEL - 1
490 + x, weight, bias = _conv_inputs(total_tokens)
491 +
492 + def amx(x_seg, conv_states, has_init):
493 + seq = x_seg.shape[0]
494 + return ops.causal_conv1d_fwd_cpu(
495 + x=x_seg.transpose(0, 1), # [dim, seq]; stride(-2)==1 (view of [seq,dim])
496 + weight=weight,
497 + bias=bias,
498 + conv_states=conv_states,
499 + query_start_loc=torch.tensor([0, seq], dtype=torch.int32),
500 + cache_indices=torch.tensor([0], dtype=torch.int32),
501 + has_initial_state=torch.tensor([has_init]),
502 + silu_activation=True,
503 + is_vnni=False,
504 + ).contiguous()
505 +
506 + # conv_state layout passed by the AMX branch: [num_slots, dim, state_len].
507 + cs_full = torch.zeros(1, CONV_DIM, state_len, dtype=x.dtype)
508 + out_full = amx(x, cs_full, False)
509 +
510 + cs_split = torch.zeros(1, CONV_DIM, state_len, dtype=x.dtype)
511 + out1 = amx(x[:split], cs_split, False)
512 + out2 = amx(x[split:], cs_split, True)
513 + out_split = torch.cat([out1, out2], dim=1)
514 +
515 + torch.testing.assert_close(out_split, out_full, atol=1e-2, rtol=1e-2)
516 +
517 +
518 +@torch.inference_mode()
519 +def test_batch_memcpy_cpu_fallback() -> None:
520 + """The ctypes batch_memcpy fallback (used when triton-cpu is absent) must
521 + copy each src into its dst, validating the (src_ptrs, dst_ptrs, sizes)
522 + argument order against ctypes.memmove(dst, src, size).
523 + """
524 + from vllm.utils.cpu_triton_utils import batch_memcpy_kernel
525 +
526 + # Varied byte sizes, including a non-power-of-two run.
527 + sizes_bytes = [256, 1024, 17 * 4, 4096]
528 + srcs = [torch.rand(n // 4, dtype=torch.float32) for n in sizes_bytes]
529 + dsts = [torch.zeros_like(s) for s in srcs]
530 +
531 + src_ptrs = torch.tensor([s.data_ptr() for s in srcs], dtype=torch.uint64)
532 + dst_ptrs = torch.tensor([d.data_ptr() for d in dsts], dtype=torch.uint64)
533 + sizes = torch.tensor(sizes_bytes, dtype=torch.int32)
534 +
535 + batch_memcpy_kernel[(len(srcs),)](src_ptrs, dst_ptrs, sizes, BLOCK_SIZE=1024)
536 +
537 + for src, dst in zip(srcs, dsts):
538 + torch.testing.assert_close(dst, src)
tests/v1/e2e/test_cpu_linear_attn_chunked_prefix.py new
+112
@@ -0,0 +1,112 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""CPU chunked-prefill / prefix-caching correctness for linear-attention models."""
4 +
5 +import os
6 +
7 +import pytest
8 +
9 +from tests.models.utils import check_logprobs_close
10 +from vllm import LLM, SamplingParams
11 +from vllm.platforms import current_platform
12 +
13 +if not current_platform.is_cpu():
14 + pytest.skip("skipping CPU-only tests", allow_module_level=True)
15 +
16 +# Bound the KV cache so the run does not scale with host memory; these engines
17 +# only need a few thousand tokens.
18 +os.environ.setdefault("VLLM_CPU_KVCACHE_SPACE", "1")
19 +
20 +MODEL = "Qwen/Qwen3.5-0.8B"
21 +CHUNK_TOKENS = 128 # max_num_batched_tokens for the chunked engine
22 +NUM_LOGPROBS = 5
23 +SP = SamplingParams(max_tokens=32, temperature=0, logprobs=NUM_LOGPROBS)
24 +
25 +
26 +def _long_prompt(repeat: int) -> str:
27 + return "Solve the following arithmetic step by step. " * repeat + "What is 7*8?"
28 +
29 +
30 +# Prompts long enough to span several CHUNK_TOKENS-sized chunks; a single-chunk
31 +# prompt is bit-identical to full prefill regardless of the bug.
32 +PROMPTS = [_long_prompt(r) for r in (40, 60, 80)]
33 +# Spans several full cache blocks; prefix caching only reuses complete blocks.
34 +PREFIX_PROMPT = "You are a helpful assistant. " * 230 + " Now answer: what is 2+2?"
35 +
36 +
37 +def _make_llm(**overrides) -> LLM:
38 + base = dict(
39 + model=MODEL,
40 + dtype="bfloat16",
41 + max_model_len=2048,
42 + enforce_eager=True,
43 + trust_remote_code=True,
44 + )
45 + base.update(overrides)
46 + return LLM(**base)
47 +
48 +
49 +def _tuples(outputs) -> list[tuple[list[int], str, object]]:
50 + """(token_ids, text, sample_logprobs) per request, for check_logprobs_close."""
51 + return [
52 + (list(o.outputs[0].token_ids), o.outputs[0].text, o.outputs[0].logprobs)
53 + for o in outputs
54 + ]
55 +
56 +
57 +@pytest.fixture(scope="module")
58 +def full_prefill_refs():
59 + """Reference (ids, text, logprobs) for PROMPTS and PREFIX_PROMPT, full prefill."""
60 + llm = _make_llm(enable_chunked_prefill=False, enable_prefix_caching=False)
61 + refs = _tuples(llm.generate(PROMPTS, SP))
62 + prefix_ref = _tuples(llm.generate([PREFIX_PROMPT], SP))[0]
63 + del llm
64 + return refs, prefix_ref
65 +
66 +
67 +def test_chunked_prefill_matches_full_prefill(full_prefill_refs):
68 + """Batched multi-chunk prefill must stay close to per-prompt full prefill.
69 +
70 + Prompts are scheduled together so the scheduler interleaves prefill chunks
71 + across requests (the cross-request path where the accuracy gap was strongest).
72 + """
73 + refs, _ = full_prefill_refs
74 + llm = _make_llm(
75 + enable_chunked_prefill=True,
76 + max_num_batched_tokens=CHUNK_TOKENS,
77 + enable_prefix_caching=False,
78 + )
79 + got = _tuples(llm.generate(PROMPTS, SP))
80 + del llm
81 +
82 + check_logprobs_close(
83 + outputs_0_lst=refs,
84 + outputs_1_lst=got,
85 + name_0="full_prefill",
86 + name_1="chunked_prefill",
87 + )
88 +
89 +
90 +def test_prefix_cache_hit_matches_cold_cache(full_prefill_refs):
91 + """A prefix-cache hit must stay close to the cold-cache (reference) output.
92 +
93 + The warm run continues prefill from the restored GDN state; the
94 + num_cached_tokens check guards against a vacuous (no-hit) pass.
95 + """
96 + _, ref = full_prefill_refs
97 + llm = _make_llm(enable_prefix_caching=True)
98 + llm.generate([PREFIX_PROMPT], SP) # prime the cache
99 + warm_out = llm.generate([PREFIX_PROMPT], SP)[0]
100 + warm = _tuples([warm_out])[0]
101 + del llm
102 +
103 + assert warm_out.num_cached_tokens > 0, (
104 + "expected a prefix-cache hit but num_cached_tokens=0; "
105 + "PREFIX_PROMPT may be shorter than one cache block"
106 + )
107 + check_logprobs_close(
108 + outputs_0_lst=[ref],
109 + outputs_1_lst=[warm],
110 + name_0="cold_cache",
111 + name_1="warm_cache",
112 + )
vllm/platforms/cpu.py
-14
@@ -157,20 +157,6 @@ class CpuPlatform(Platform):
157 logger.warning_once("Dual-Batch Overlap is not supported on CPU, disabled.")
158 parallel_config.enable_dbo = False
159
160 - if torch.cpu._is_amx_tile_supported() and (
161 - model_config is not None
162 - and model_config.get_num_layers_by_block_type(
163 - parallel_config, "linear_attention"
164 - )
165 - > 0
166 - ):
167 - cache_config.enable_prefix_caching = False
168 - scheduler_config.enable_chunked_prefill = False
169 - logger.warning_once(
170 - "Disabled unsupported prefix caching and chunked prefill "
171 - "for linear attention on AMX CPU platforms."
172 - )
173 -
160 # Note: workaround for v1 gpu_model_runner
161 from vllm.config import CompilationMode
162
vllm/utils/cpu_triton_utils.py
+10
@@ -5,6 +5,7 @@
5 Contains replacement functions to fallback Triton usages in CPU backend
6 """
7
8 +import ctypes
9 from collections.abc import Callable
10
11 import torch
@@ -336,3 +337,12 @@ rejection_greedy_sample_kernel = _FuncWrapper(_rejection_greedy_sample_kernel_im
337 rejection_random_sample_kernel = _FuncWrapper(_rejection_random_sample_kernel_impl)
338 expand_kernel = _FuncWrapper(_expand_kernel_impl)
339 sample_recovered_tokens_kernel = _FuncWrapper(_sample_recovered_tokens_kernel_impl)
340 +
341 +
342 +def _batch_memcpy_impl(src_ptrs, dst_ptrs, sizes, BLOCK_SIZE=None):
343 + # BLOCK_SIZE is unused; kept for signature parity with the Triton kernel.
344 + for src, dst, size in zip(src_ptrs.tolist(), dst_ptrs.tolist(), sizes.tolist()):
345 + ctypes.memmove(dst, src, size)
346 +
347 +
348 +batch_memcpy_kernel = _FuncWrapper(_batch_memcpy_impl)
vllm/v1/worker/cpu_model_runner.py
+24 -4
@@ -12,7 +12,7 @@ from vllm.logger import init_logger
12 from vllm.model_executor.model_loader import get_model
13 from vllm.tracing import instrument
14 from vllm.v1.core.sched.output import SchedulerOutput
15 -from vllm.v1.kv_cache_interface import KVCacheConfig
15 +from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig
16 from vllm.v1.utils import CpuGpuBuffer
17 from vllm.v1.worker.gpu_model_runner import GPUModelRunner
18
@@ -103,6 +103,10 @@ class CPUModelRunner(GPUModelRunner):
103 cpu_tl.sample_recovered_tokens_kernel
104 )
105
106 + import vllm.v1.worker.mamba_utils
107 +
108 + vllm.v1.worker.mamba_utils.batch_memcpy_kernel = cpu_tl.batch_memcpy_kernel
109 +
110 @instrument(span_name="Loading (CPU)")
111 def load_model(self, load_dummy_weights: bool = False) -> None:
112 if load_dummy_weights:
@@ -153,9 +157,25 @@ class CPUModelRunner(GPUModelRunner):
157 pass
158
159 def _zero_block_ids(self, block_ids: list[int]) -> None:
156 - # CPU attention assigns -INF to logits at invalid positions,
157 - # so stale KV cache data never affects computation.
158 - pass
160 + # Zero full-attention blocks to prevent stale data corruption on partial writes.
161 + # Encoder-only (runner-only) layers are not FullAttentionSpec, so the
162 + # spec filter below already excludes them; no runner-only skip needed.
163 + seen_ptrs: set[int] = set()
164 + for group in self.kv_cache_config.kv_cache_groups:
165 + if not isinstance(group.kv_cache_spec, FullAttentionSpec):
166 + continue
167 + for layer_name in group.layer_names:
168 + ctx = self.compilation_config.static_forward_context.get(layer_name)
169 + if ctx is None:
170 + continue
171 + kv = ctx.kv_cache
172 + if not isinstance(kv, torch.Tensor):
173 + continue
174 + if kv.data_ptr() in seen_ptrs:
175 + continue
176 + seen_ptrs.add(kv.data_ptr())
177 + for block_id in block_ids:
178 + kv[block_id].zero_()
179
180 # =========================================================================
181 # CPU-safe overrides for speculative decoding methods