@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 4ff8c3c8f9

Fork 0

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 6, 23:00 UTC

[BUGFIX][Mamba][Qwen3.5] Zero freed SSM cache blocks on GPU (#35219)

Signed-off-by: Vadim Gimpelson <vadim.gimpelson@gmail.com>

Vadim Gimpelson committed Mar 10, 2026 at 14:32 UTC 4ff8c3c8f9ece010a1d0e376f5cc1b468b95f366

10 files changed +287 -8

vllm/utils/math_utils.py

@@ -30,3 +30,8 @@ def round_up(x: int, y: int) -> int:

30 def round_down(x: int, y: int) -> int:

31 """Round down x to the nearest multiple of y."""

32 return (x // y) * y

33 +

34 +

35 +def largest_power_of_2_divisor(n: int) -> int:

36 + """Return the largest power-of-2 that divides *n* (isolate lowest set bit)."""

37 + return n & (-n)

vllm/v1/attention/backend.py

+20

@@ -86,6 +86,26 @@ class AttentionBackend(ABC):

86 ) -> tuple[int, ...]:

87 raise NotImplementedError

89 + @classmethod

90 + def get_kv_cache_block_dim(

91 + cls,

92 + block_size: int,

93 + num_kv_heads: int,

94 + head_size: int,

95 + cache_dtype_str: str = "auto",

96 + ) -> int:

97 + """Discover which tensor dim is the block index, since different

98 + backends lay out dims differently."""

99 + _S = 1234567

100 + shape = cls.get_kv_cache_shape(

101 + _S,

102 + block_size,

103 + num_kv_heads,

104 + head_size,

105 + cache_dtype_str=cache_dtype_str,

106 + )

107 + return shape.index(_S)

108 +

109 @staticmethod

110 def get_kv_cache_stride_order(

111 include_num_layers_dimension: bool = False,

vllm/v1/core/kv_cache_manager.py

@@ -501,6 +501,13 @@ class KVCacheManager:

501 # Only create new KVCacheBlocks for non-empty blocks

502 return KVCacheBlocks(blocks) if any(blocks) else self.empty_kv_cache_blocks

503

504 + def take_new_block_ids(self) -> list[int]:

505 + """Drain and return new attention block IDs for zeroing."""

506 + ids: list[int] = []

507 + for mgr in self.coordinator.single_type_managers:

508 + ids.extend(mgr.take_new_block_ids())

509 + return ids

510 +

511 def new_step_starts(self) -> None:

512 """Called when a new step is started."""

513 self.coordinator.new_step_starts()

vllm/v1/core/sched/output.py

@@ -233,6 +233,11 @@ class SchedulerOutput:

233 # EC Cache Connector metadata

234 ec_connector_metadata: ECConnectorMetadata | None = None

235

236 + # Block IDs freshly allocated from the pool during this scheduling step.

237 + # The worker zeros the corresponding GPU memory before the blocks are used,

238 + # preventing stale NaN/data from corrupting attention or SSM computation.

239 + new_block_ids_to_zero: list[int] | None = None

240 +

241 @classmethod

242 def make_empty(cls) -> "SchedulerOutput":

243 return cls(

vllm/v1/core/sched/scheduler.py

+10 -8

@@ -48,7 +48,7 @@ from vllm.v1.core.sched.output import (

48 from vllm.v1.core.sched.request_queue import SchedulingPolicy, create_request_queue

49 from vllm.v1.core.sched.utils import check_stop, remove_all

50 from vllm.v1.engine import EngineCoreEventType, EngineCoreOutput, EngineCoreOutputs

51 -from vllm.v1.kv_cache_interface import KVCacheConfig, MambaSpec

51 +from vllm.v1.kv_cache_interface import KVCacheConfig

52 from vllm.v1.metrics.perf import ModelMetrics, PerfStats

53 from vllm.v1.metrics.stats import PrefixCacheStats, SchedulerStats

54 from vllm.v1.outputs import DraftTokenIds, KVConnectorOutput, ModelRunnerOutput

@@ -233,13 +233,8 @@ class Scheduler(SchedulerInterface):

233 self.use_pp = self.parallel_config.pipeline_parallel_size > 1

234 self.use_v2_model_runner = envs.VLLM_USE_V2_MODEL_RUNNER

235

236 - def has_mamba_layers(kv_cache_config: KVCacheConfig) -> bool:

237 - return any(

238 - isinstance(group_spec.kv_cache_spec, MambaSpec)

239 - for group_spec in kv_cache_config.kv_cache_groups

240 - )

241 -

242 - self.has_mamba_layers = has_mamba_layers(kv_cache_config)

236 + self.has_mamba_layers = kv_cache_config.has_mamba_layers

237 + self.needs_kv_cache_zeroing = kv_cache_config.needs_kv_cache_zeroing

238 self.need_mamba_block_aligned_split = (

239 self.has_mamba_layers and self.cache_config.mamba_cache_mode == "align"

240 )

@@ -890,6 +885,12 @@ class Scheduler(SchedulerInterface):

885 self.prev_step_scheduled_req_ids.clear()

886 self.prev_step_scheduled_req_ids.update(num_scheduled_tokens.keys())

887

888 + new_block_ids_to_zero = (

889 + (self.kv_cache_manager.take_new_block_ids() or None)

890 + if self.needs_kv_cache_zeroing

891 + else None

892 + )

893 +

894 scheduler_output = SchedulerOutput(

895 scheduled_new_reqs=new_reqs_data,

896 scheduled_cached_reqs=cached_reqs_data,

@@ -905,6 +906,7 @@ class Scheduler(SchedulerInterface):

906 # the previous and the current steps.

907 finished_req_ids=self.finished_req_ids,

908 free_encoder_mm_hashes=self.encoder_cache_manager.get_freed_mm_hashes(),

909 + new_block_ids_to_zero=new_block_ids_to_zero,

910 )

911

912 # NOTE(Kuntai): this function is designed for multiple purposes:

vllm/v1/core/single_type_kv_cache_manager.py

+11

@@ -55,6 +55,7 @@ class SingleTypeKVCacheManager(ABC):

55 self.kv_cache_spec = kv_cache_spec

56 self.block_pool = block_pool

57 self.enable_caching = enable_caching

58 + self.new_block_ids: list[int] = []

60 # Mapping from request ID to blocks to track the blocks allocated

61 # for each request, so that we can free the blocks when the request

@@ -208,6 +209,8 @@ class SingleTypeKVCacheManager(ABC):

209 cdiv(num_total_computed_tokens, self.block_size) - len(req_blocks)

210 )

211 req_blocks.extend(allocated_blocks)

212 + if type(self.kv_cache_spec) is FullAttentionSpec:

213 + self.new_block_ids.extend(b.block_id for b in allocated_blocks)

214

215 def allocate_new_blocks(

216 self, request_id: str, num_tokens: int, num_tokens_main_model: int

@@ -234,8 +237,16 @@ class SingleTypeKVCacheManager(ABC):

237 else:

238 new_blocks = self.block_pool.get_new_blocks(num_new_blocks)

239 req_blocks.extend(new_blocks)

240 + if type(self.kv_cache_spec) is FullAttentionSpec:

241 + self.new_block_ids.extend(b.block_id for b in new_blocks)

242 return new_blocks

243

244 + def take_new_block_ids(self) -> list[int]:

245 + """Drain and return block IDs allocated since the last call."""

246 + ids = self.new_block_ids

247 + self.new_block_ids = []

248 + return ids

249 +

250 def cache_blocks(self, request: Request, num_tokens: int) -> None:

251 """

252 Cache the blocks for the request.

vllm/v1/kv_cache_interface.py

@@ -489,3 +489,11 @@ class KVCacheConfig:

489 For models with multiple types of attention, there will be multiple groups,

490 see `_get_kv_cache_config_uniform_page_size` for more details.

491 """

492 +

493 + @property

494 + def has_mamba_layers(self) -> bool:

495 + return any(isinstance(g.kv_cache_spec, MambaSpec) for g in self.kv_cache_groups)

496 +

497 + @property

498 + def needs_kv_cache_zeroing(self) -> bool:

499 + return self.has_mamba_layers

vllm/v1/worker/gpu_model_runner.py

+27

@@ -197,6 +197,7 @@ from vllm.v1.worker.workspace import lock_workspace

197

198 from .utils import (

199 AttentionGroup,

200 + KVBlockZeroer,

201 add_kv_sharing_layers_to_kv_cache_groups,

202 bind_kv_cache,

203 prepare_kernel_block_sizes,

@@ -982,6 +983,26 @@ class GPUModelRunner(

983 decode_threshold=self.reorder_batch_threshold,

984 )

985

986 + def _init_kv_zero_meta(self) -> None:

987 + """One-time precomputation for _zero_block_ids.

988 +

989 + Delegates to KVBlockZeroer.init_meta with the runner's state.

990 + Called from gpu_worker.py outside the CuMem pool context.

991 + """

992 + self._kv_block_zeroer = KVBlockZeroer(self.device, self.pin_memory)

993 + self._kv_block_zeroer.init_meta(

994 + attn_groups_iter=self._kv_cache_spec_attn_group_iterator(),

995 + kernel_block_sizes=self._kernel_block_sizes,

996 + cache_dtype=self.cache_config.cache_dtype,

997 + runner_only_attn_layers=self.runner_only_attn_layers,

998 + static_forward_context=(self.compilation_config.static_forward_context),

999 + )

1000 +

1001 + def _zero_block_ids(self, block_ids: list[int]) -> None:

1002 + """Zero the KV cache memory for the given block IDs."""

1003 + if hasattr(self, "_kv_block_zeroer"):

1004 + self._kv_block_zeroer.zero_block_ids(block_ids)

1005 +

1006 # Note: used for model runner override.

1007 def _init_device_properties(self) -> None:

1008 """Initialize attributes from torch.cuda.get_device_properties"""

@@ -1018,6 +1039,11 @@ class GPUModelRunner(

1039 for req_id in scheduler_output.finished_req_ids:

1040 self.input_batch.remove_request(req_id)

1041

1042 + # Zero GPU memory for freshly allocated cache blocks to prevent

1043 + # stale NaN/data from corrupting attention or SSM computation.

1044 + if scheduler_output.new_block_ids_to_zero:

1045 + self._zero_block_ids(scheduler_output.new_block_ids_to_zero)

1046 +

1047 # Free the cached encoder outputs.

1048 for mm_hash in scheduler_output.free_encoder_mm_hashes:

1049 self.encoder_cache.pop(mm_hash, None)

@@ -6476,6 +6502,7 @@ class GPUModelRunner(

6502 kernel_block_sizes = prepare_kernel_block_sizes(

6503 kv_cache_config, self.attn_groups

6504 )

6505 + self._kernel_block_sizes = kernel_block_sizes

6506

6507 # create metadata builders

6508 self.initialize_metadata_builders(kv_cache_config, kernel_block_sizes)

vllm/v1/worker/gpu_worker.py

@@ -556,6 +556,14 @@ class Worker(WorkerBase):

556 else:

557 self.model_runner.initialize_kv_cache(kv_cache_config)

558

559 + # Build KV-zero metadata outside the CuMem pool so the bookkeeping

560 + # GPU tensors (seg_addrs, block-id buffers) use the standard PyTorch

561 + # allocator and are not discarded during sleep/wake cycles.

562 + if kv_cache_config.needs_kv_cache_zeroing and hasattr(

563 + self.model_runner, "_init_kv_zero_meta"

564 + ):

565 + self.model_runner._init_kv_zero_meta()

566 +

567 @instrument(span_name="Warmup (GPU)")

568 def compile_or_warm_up_model(self) -> float:

569 warmup_sizes: list[int] = []

vllm/v1/worker/utils.py

+186

@@ -2,7 +2,10 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 import math

4 from collections import defaultdict

5 +from collections.abc import Iterable

6 from dataclasses import dataclass, field

7 +from itertools import product as iprod

8 +from typing import Any

10 import torch

@@ -12,6 +15,8 @@ from vllm.model_executor.layers.attention import Attention

15 from vllm.model_executor.models.interfaces import MultiModalEmbeddings

16 from vllm.model_executor.models.utils import extract_layer_index

17 from vllm.platforms import current_platform

18 +from vllm.triton_utils import tl, triton

19 +from vllm.utils.math_utils import largest_power_of_2_divisor

20 from vllm.utils.mem_utils import MemorySnapshot, format_gib

21 from vllm.v1.attention.backend import (

22 AttentionBackend,

@@ -21,6 +26,7 @@ from vllm.v1.attention.backend import (

26 from vllm.v1.kv_cache_interface import (

27 AttentionSpec,

28 EncoderOnlyAttentionSpec,

29 + FullAttentionSpec,

30 KVCacheConfig,

31 KVCacheGroupSpec,

32 KVCacheSpec,

@@ -31,6 +37,186 @@ from vllm.v1.kv_cache_interface import (

37 logger = init_logger(__name__)

40 +@triton.jit

41 +def _zero_kv_blocks_kernel(

42 + seg_addrs_ptr,

43 + block_ids_ptr,

44 + n_blocks,

45 + N_SEGS: tl.constexpr,

46 + PAGE_SIZE_EL: tl.constexpr,

47 + BLOCK_SIZE: tl.constexpr,

48 +):

49 + """Zero KV cache blocks across all segments in a single launch.

50 +

51 + Each segment is a contiguous region of one block's data. For backends

52 + where blocks are outermost (block_dim=0) there is one segment per

53 + buffer. For backends where K/V is outermost (block_dim=1) there are

54 + two segments per buffer (one for K, one for V).

55 +

56 + seg_addrs_ptr holds absolute byte addresses (int64) for each segment,

57 + allowing segments to live in different CUDA allocations.

58 +

59 + Programs are mapped as (block_index, seg_index, chunk_index).

60 + """

61 + pid = tl.program_id(0)

62 + chunks = PAGE_SIZE_EL // BLOCK_SIZE

63 + work_per_block = N_SEGS * chunks

64 + block_index = pid // work_per_block

65 + if block_index >= n_blocks:

66 + return

67 + remainder = pid % work_per_block

68 + seg_index = remainder // chunks

69 + chunk_index = remainder % chunks

70 + block_id = tl.load(block_ids_ptr + block_index)

71 + seg_addr = tl.load(seg_addrs_ptr + seg_index)

72 + ptr = tl.cast(seg_addr, tl.pointer_type(tl.int32))

73 + offset = (

74 + block_id.to(tl.int64) * PAGE_SIZE_EL + chunk_index.to(tl.int64) * BLOCK_SIZE

75 + )

76 + cols = tl.arange(0, BLOCK_SIZE).to(tl.int64)

77 + tl.store(ptr + offset + cols, tl.zeros([BLOCK_SIZE], dtype=tl.int32))

78 +

79 +

80 +class KVBlockZeroer:

81 + """Manages efficient zeroing of KV cache blocks via a Triton kernel.

82 +

83 + Call :meth:`init_meta` once after KV caches are allocated to precompute

84 + segment addresses, then call :meth:`zero_block_ids` each step to zero

85 + newly-allocated blocks.

86 + """

87 +

88 + def __init__(self, device: torch.device, pin_memory: bool):

89 + self.device = device

90 + self.pin_memory = pin_memory

91 + self._meta: tuple[torch.Tensor, int, int, int] | None = None

92 + self._id_cap: int = 0

93 + self._ids_pinned: torch.Tensor | None = None

94 + self._ids_gpu: torch.Tensor | None = None

95 +

96 + def init_meta(

97 + self,

98 + attn_groups_iter: Iterable["AttentionGroup"],

99 + kernel_block_sizes: list[int],

100 + cache_dtype: str,

101 + runner_only_attn_layers: set[str],

102 + static_forward_context: dict[str, Any],

103 + ) -> None:

104 + """One-time precomputation for zero_block_ids.

105 +

106 + Builds absolute-address table for the Triton zeroing kernel.

107 + Each entry is the absolute byte address of a segment start on the

108 + GPU, so segments in different CUDA allocations work correctly.

109 +

110 + Block IDs from the scheduler reference logical blocks whose size

111 + may differ from the kernel block size (virtual block splitting).

112 + PAGE_SIZE_EL accounts for this ratio so that

113 + ``block_id * PAGE_SIZE_EL`` lands at the correct offset.

114 +

115 + Only AttentionSpec layers are processed; Mamba layers are skipped.

116 + """

117 + seen_ptrs: set[int] = set()

118 + seg_addrs: list[int] = []

119 + page_size_el: int | None = None

120 +

121 + for group in attn_groups_iter:

122 + spec = group.kv_cache_spec

123 + if type(spec) is not FullAttentionSpec:

124 + continue

125 + if group.kv_cache_group_id >= len(kernel_block_sizes):

126 + continue

127 + kernel_bs = kernel_block_sizes[group.kv_cache_group_id]

128 + ratio = spec.block_size // kernel_bs

129 + block_dim = group.backend.get_kv_cache_block_dim(

130 + kernel_bs,

131 + spec.num_kv_heads,

132 + spec.head_size,

133 + cache_dtype_str=cache_dtype,

134 + )

135 +

136 + for layer_name in group.layer_names:

137 + if layer_name in runner_only_attn_layers:

138 + continue

139 + kv = static_forward_context[layer_name].kv_cache[0]

140 + if isinstance(kv, list):

141 + continue

142 + dp = kv.data_ptr()

143 + if dp in seen_ptrs:

144 + continue

145 + seen_ptrs.add(dp)

146 +

147 + el = kv.element_size()

148 + cur_bytes = kv.stride(block_dim) * el

149 + assert cur_bytes % 4 == 0

150 + kernel_block_el = cur_bytes // 4

151 + cur_page_el = kernel_block_el * ratio

152 + if page_size_el is None:

153 + page_size_el = cur_page_el

154 + else:

155 + assert page_size_el == cur_page_el, (

156 + f"Non-uniform page sizes: {page_size_el} vs {cur_page_el}"

157 + )

158 +

159 + block_stride_bytes = cur_bytes

160 + outer_dims = [

161 + d

162 + for d in range(block_dim)

163 + if kv.stride(d) * el > block_stride_bytes

164 + ]

165 + outer_strides = [kv.stride(d) * el for d in outer_dims]

166 + for outer in iprod(*(range(kv.shape[d]) for d in outer_dims)):

167 + off_bytes = sum(i * s for i, s in zip(outer, outer_strides))

168 + seg_addrs.append(dp + off_bytes)

169 +

170 + if not seg_addrs or page_size_el is None:

171 + self._meta = None

172 + return

173 +

174 + blk_size = min(largest_power_of_2_divisor(page_size_el), 1024)

175 + self._id_cap = 8192

176 + self._ids_pinned = torch.empty(

177 + self._id_cap,

178 + dtype=torch.int64,

179 + pin_memory=self.pin_memory,

180 + )

181 + self._ids_gpu = torch.empty(self._id_cap, dtype=torch.int64, device=self.device)

182 + self._meta = (

183 + torch.tensor(seg_addrs, dtype=torch.int64, device=self.device),

184 + page_size_el,

185 + blk_size,

186 + len(seg_addrs),

187 + )

188 +

189 + def zero_block_ids(self, block_ids: list[int]) -> None:

190 + """Zero the KV cache memory for the given block IDs."""

191 + if not block_ids or self._meta is None:

192 + return

193 + seg_addrs, page_size_el, blk_size, n_segs = self._meta

194 + n_blocks = len(block_ids)

195 + if n_blocks > self._id_cap:

196 + self._id_cap = n_blocks * 2

197 + self._ids_pinned = torch.empty(

198 + self._id_cap,

199 + dtype=torch.int64,

200 + pin_memory=self.pin_memory,

201 + )

202 + self._ids_gpu = torch.empty(

203 + self._id_cap, dtype=torch.int64, device=self.device

204 + )

205 + assert self._ids_pinned is not None and self._ids_gpu is not None

206 + self._ids_pinned[:n_blocks].numpy()[:] = block_ids

207 + idx = self._ids_gpu[:n_blocks]

208 + idx.copy_(self._ids_pinned[:n_blocks], non_blocking=True)

209 + grid = (n_blocks * n_segs * (page_size_el // blk_size),)

210 + _zero_kv_blocks_kernel[grid](

211 + seg_addrs,

212 + idx,

213 + n_blocks,

214 + N_SEGS=n_segs,

215 + PAGE_SIZE_EL=page_size_el,

216 + BLOCK_SIZE=blk_size,

217 + )

218 +

219 +

220 @dataclass

221 class AttentionGroup:

222 backend: type[AttentionBackend]