c59a132f96 · tduncan/Nvidia-Optimization-Stack-VLLM

tests/compile/passes/test_fusion_attn.py

+1 -1

@@ -127,7 +127,7 @@ class AttentionQuantPatternModel(torch.nn.Module):

127 raw_tensor = raw_tensor.view(kv_cache_shape)

128 kv_cache = raw_tensor.permute(*inv_order)

129

130 - self.attn.kv_cache = [kv_cache]

130 + self.attn.kv_cache = kv_cache

131

132 # Build attn metadata

133 self.attn_metadata = self.builder.build(

tests/compile/passes/test_rope_kvcache_fusion.py

+3 -3

@@ -148,7 +148,7 @@ class QKRoPEKVCacheTestModel(torch.nn.Module):

148 raw_tensor = raw_tensor.view(kv_cache_shape)

149 kv_cache = raw_tensor.permute(*inv_order)

150

151 - self.attn.kv_cache = [kv_cache]

151 + self.attn.kv_cache = kv_cache

152

153 # Build attn metadata

154 attn_metadata = self.builder.build(

@@ -295,7 +295,7 @@ def test_rope_kvcache_fusion(

295 }

296 q_unfused, k_unfused, v_unfused, dummy = model(qkv_unfused, pos_unfused)

297 attn_layer = forward_context.no_compile_layers[model.layer_name]

298 - kv_cache_unfused = attn_layer.kv_cache[0]

298 + kv_cache_unfused = attn_layer.kv_cache

299 del dummy

300

301 torch._dynamo.mark_dynamic(qkv, 0)

@@ -309,7 +309,7 @@ def test_rope_kvcache_fusion(

309 }

310 q_fused, k_fused, v_fused, dummy = model_fused(qkv, pos)

311 attn_layer = forward_context.no_compile_layers[model.layer_name]

312 - kv_cache_fused = attn_layer.kv_cache[0]

312 + kv_cache_fused = attn_layer.kv_cache

313 del dummy

314

315 assert fusion_pass.matched_count == 1

tests/v1/e2e/general/test_mamba_prefix_cache.py

+3 -3

@@ -258,8 +258,8 @@ def get_fake_execute_model_fn(original_execute_model_fn: Callable):

258 mamba_kv_cache_dict[

259 num_computed_tokens - num_computed_tokens % BLOCK_SIZE

260 ] = (

261 - kv_cache[0][0][block_id].clone(),

262 - kv_cache[0][1][block_id].clone(),

261 + kv_cache[0][block_id].clone(),

262 + kv_cache[1][block_id].clone(),

263 )

264

265 last_num_computed_tokens = num_computed_tokens

@@ -302,7 +302,7 @@ def get_fake_process_mamba_fn(

302 mamba_layer_name = kv_cache_config.kv_cache_groups[

303 mamba_group_id

304 ].layer_names[0]

305 - mamba_kv_cache = forward_context[mamba_layer_name].kv_cache[0][-1]

305 + mamba_kv_cache = forward_context[mamba_layer_name].kv_cache[-1]

306 mamba_block_table = input_batch.block_table.block_tables[

307 mamba_group_id

308 ].block_table.cpu[0]

tests/v1/worker/test_gpu_model_runner.py

+17 -17

@@ -670,8 +670,8 @@ def test_init_kv_cache_without_kv_sharing(default_vllm_config):

670

671 runner.initialize_kv_cache(kv_cache_config)

672

673 - layer_0_kv = vllm_ctx[layer_0].kv_cache[0]

674 - layer_1_kv = vllm_ctx[layer_1].kv_cache[0]

673 + layer_0_kv = vllm_ctx[layer_0].kv_cache

674 + layer_1_kv = vllm_ctx[layer_1].kv_cache

675 # check layer 1 kv cache does NOT share memory with layer 0

676 assert id(layer_1_kv) != id(layer_0_kv)

677

@@ -740,8 +740,8 @@ def test_init_kv_cache_with_kv_sharing_valid(default_vllm_config):

740 runner.initialize_kv_cache(kv_cache_config)

741 kv_cache_config_after_init = runner.kv_cache_config

742

743 - layer_0_kv = vllm_ctx[layer_0].kv_cache[0]

744 - layer_1_kv = vllm_ctx[layer_1].kv_cache[0]

743 + layer_0_kv = vllm_ctx[layer_0].kv_cache

744 + layer_1_kv = vllm_ctx[layer_1].kv_cache

745 # check layer 1 kv cache shares memory with layer 0

746 assert id(layer_1_kv) == id(layer_0_kv)

747

@@ -864,9 +864,9 @@ def test_hybrid_attention_mamba_tensor_shapes():

864 np.random.shuffle(ind)

865 blocks0, blocks1 = ind[: (num_blocks // 2)], ind[(num_blocks // 2) :]

866

867 - attn_shape = vllm_ctx[layer_0].kv_cache[0].shape

868 - conv_shape = vllm_ctx[layer_2].kv_cache[0][0].shape

869 - ssm_shape = vllm_ctx[layer_2].kv_cache[0][1].shape

867 + attn_shape = vllm_ctx[layer_0].kv_cache.shape

868 + conv_shape = vllm_ctx[layer_2].kv_cache[0].shape

869 + ssm_shape = vllm_ctx[layer_2].kv_cache[1].shape

870

871 # assert we are using FlashInfer

872 assert attn_shape[0] % num_blocks == 0

@@ -905,21 +905,21 @@ def test_hybrid_attention_mamba_tensor_shapes():

905 kernel_blocks_for_attention = kv_blocks_for_attention * block_split_ratio

906

907 for layer in [layer_0, layer_1]:

908 - # attention: kv_cache[0][kernel_block_idx, kv_idx, ...]

908 + # attention: kv_cache[kernel_block_idx, kv_idx, ...]

909 for i, kernel_block in enumerate(kernel_blocks_for_attention):

910 - vllm_ctx[layer].kv_cache[0][kernel_block, :] = attn_blocks_constant[i]

910 + vllm_ctx[layer].kv_cache[kernel_block, :] = attn_blocks_constant[i]

911

912 # fill mamba blocks with constants using kernel block indices

913 for layer in [layer_2, layer_3, layer_4, layer_5]:

914 - # mamba: kv_cache[0][component][kernel_block_idx, ...]

914 + # mamba: kv_cache[component][kernel_block_idx, ...]

915 for i, kv_block in enumerate(kv_blocks_for_mamba):

916 - vllm_ctx[layer].kv_cache[0][0][kv_block, :] = conv_blocks_constant[i]

917 - vllm_ctx[layer].kv_cache[0][1][kv_block, :] = ssm_blocks_constant[i]

916 + vllm_ctx[layer].kv_cache[0][kv_block, :] = conv_blocks_constant[i]

917 + vllm_ctx[layer].kv_cache[1][kv_block, :] = ssm_blocks_constant[i]

918

919 # verify attention and mamba contents are correct

920 for layer in [layer_0, layer_1]:

921 for i, kernel_block in enumerate(kernel_blocks_for_attention):

922 - actual_kv = vllm_ctx[layer].kv_cache[0][kernel_block, :]

922 + actual_kv = vllm_ctx[layer].kv_cache[kernel_block, :]

923 expected = attn_blocks_constant[i]

924

925 # Check K and V separately

@@ -928,8 +928,8 @@ def test_hybrid_attention_mamba_tensor_shapes():

928

929 for layer in [layer_2, layer_3, layer_4, layer_5]:

930 for i, kv_block in enumerate(kv_blocks_for_mamba):

931 - actual_conv = vllm_ctx[layer].kv_cache[0][0][kv_block, :]

932 - actual_ssm = vllm_ctx[layer].kv_cache[0][1][kv_block, :]

931 + actual_conv = vllm_ctx[layer].kv_cache[0][kv_block, :]

932 + actual_ssm = vllm_ctx[layer].kv_cache[1][kv_block, :]

933 expected_conv = conv_blocks_constant[i]

934 expected_ssm = ssm_blocks_constant[i]

935

@@ -938,8 +938,8 @@ def test_hybrid_attention_mamba_tensor_shapes():

938

939 for layer in [layer_2, layer_3, layer_4, layer_5]:

940 for i, kv_block in enumerate(kv_blocks_for_mamba):

941 - actual_conv = vllm_ctx[layer].kv_cache[0][0][kv_block, :]

942 - actual_ssm = vllm_ctx[layer].kv_cache[0][1][kv_block, :]

941 + actual_conv = vllm_ctx[layer].kv_cache[0][kv_block, :]

942 + actual_ssm = vllm_ctx[layer].kv_cache[1][kv_block, :]

943 expected_conv = conv_blocks_constant[i]

944 expected_ssm = ssm_blocks_constant[i]

945 assert torch.equal(actual_conv, expected_conv)

tests/v1/worker/test_utils.py

+10 -10

@@ -23,10 +23,10 @@ def test_bind_kv_cache(default_vllm_config):

23 }

24 runner_kv_caches: list[torch.Tensor] = []

25 bind_kv_cache(kv_cache, ctx, runner_kv_caches)

26 - assert ctx["layers.0.self_attn"].kv_cache[0] is kv_cache["layers.0.self_attn"]

27 - assert ctx["layers.1.self_attn"].kv_cache[0] is kv_cache["layers.1.self_attn"]

28 - assert ctx["layers.2.self_attn"].kv_cache[0] is kv_cache["layers.2.self_attn"]

29 - assert ctx["layers.3.self_attn"].kv_cache[0] is kv_cache["layers.3.self_attn"]

26 + assert ctx["layers.0.self_attn"].kv_cache is kv_cache["layers.0.self_attn"]

27 + assert ctx["layers.1.self_attn"].kv_cache is kv_cache["layers.1.self_attn"]

28 + assert ctx["layers.2.self_attn"].kv_cache is kv_cache["layers.2.self_attn"]

29 + assert ctx["layers.3.self_attn"].kv_cache is kv_cache["layers.3.self_attn"]

30

31 assert runner_kv_caches[0] is kv_cache["layers.0.self_attn"]

32 assert runner_kv_caches[1] is kv_cache["layers.1.self_attn"]

@@ -50,8 +50,8 @@ def test_bind_kv_cache_non_attention(default_vllm_config):

50 runner_kv_caches: list[torch.Tensor] = []

51 bind_kv_cache(kv_cache, ctx, runner_kv_caches)

52

53 - assert ctx["model.layers.20.attn"].kv_cache[0] is kv_cache["model.layers.20.attn"]

54 - assert ctx["model.layers.28.attn"].kv_cache[0] is kv_cache["model.layers.28.attn"]

53 + assert ctx["model.layers.20.attn"].kv_cache is kv_cache["model.layers.20.attn"]

54 + assert ctx["model.layers.28.attn"].kv_cache is kv_cache["model.layers.28.attn"]

55

56 assert runner_kv_caches[0] is kv_cache["model.layers.20.attn"]

57 assert runner_kv_caches[1] is kv_cache["model.layers.28.attn"]

@@ -74,14 +74,14 @@ def test_bind_kv_cache_draft_model(default_vllm_config):

74 runner_kv_caches: list[torch.Tensor] = []

75 bind_kv_cache(kv_cache, ctx, runner_kv_caches)

76

77 - assert ctx["model.layers.0.attn"].kv_cache[0] is kv_cache["model.layers.0.attn"]

78 - assert ctx["model.layers.1.attn"].kv_cache[0] is kv_cache["model.layers.1.attn"]

77 + assert ctx["model.layers.0.attn"].kv_cache is kv_cache["model.layers.0.attn"]

78 + assert ctx["model.layers.1.attn"].kv_cache is kv_cache["model.layers.1.attn"]

79 assert (

80 - ctx["draft_model.layers.0.attn"].kv_cache[0]

80 + ctx["draft_model.layers.0.attn"].kv_cache

81 is kv_cache["draft_model.layers.0.attn"]

82 )

83 assert (

84 - ctx["draft_model.layers.1.attn"].kv_cache[0]

84 + ctx["draft_model.layers.1.attn"].kv_cache

85 is kv_cache["draft_model.layers.1.attn"]

86 )

87

vllm/distributed/kv_transfer/kv_connector/v1/example_connector.py

+2 -4

@@ -181,12 +181,10 @@ class ExampleConnector(KVConnectorBase_V1):

181 # Only process layers that have kv_cache

182 # attribute (attention layers) Skip non-attention

183 # layers like FusedMoE/MLP etc.

184 - kv_cache_attr = getattr(layer, "kv_cache", None)

185 - if kv_cache_attr is None:

184 + kv_cache_layer = getattr(layer, "kv_cache", None)

185 + if kv_cache_layer is None:

186 continue

187

188 - kv_cache_layer = kv_cache_attr[0]

189 -

188 filename = self._generate_filename_debug(

189 layer_name, request.token_ids, request.mm_hashes

190 )

vllm/distributed/kv_transfer/kv_connector/v1/lmcache_integration/vllm_v1_adapter.py

+1 -1

@@ -778,7 +778,7 @@ class LMCacheConnectorV1Impl:

778 continue

779

780 if layer_name not in self.kv_caches:

781 - self.kv_caches[layer_name] = attn_layer.kv_cache[0]

781 + self.kv_caches[layer_name] = attn_layer.kv_cache

782

783 ####################

784 # Worker side APIs

vllm/distributed/kv_transfer/kv_connector/v1/p2p/p2p_nccl_connector.py

+1 -1

@@ -214,7 +214,7 @@ class P2pNcclConnector(KVConnectorBase_V1):

214 if kv_cache is None:

215 continue

216

217 - layer = kv_cache[0]

217 + layer = kv_cache

218

219 kv_cache = self.p2p_nccl_engine.recv_tensor(

220 request.request_id + "#" + layer_name, remote_address

vllm/model_executor/layers/attention/attention.py

+2 -5

@@ -349,10 +349,7 @@ class Attention(nn.Module, AttentionLayerBase):

349 # use a placeholder kv cache tensor during init, which will be replaced

350 # by bind_kv_cache

351 # this variable will not be accessed if use_direct_call is True

352 - self.kv_cache = [

353 - torch.tensor([])

354 - for _ in range(vllm_config.parallel_config.pipeline_parallel_size)

355 - ]

352 + self.kv_cache = torch.tensor([])

353

354 # Initialize KV cache quantization attributes

355 _init_kv_cache_quant(self, quant_config, prefix)

@@ -599,7 +596,7 @@ def get_attention_context(

596 if isinstance(attn_metadata, dict):

597 attn_metadata = attn_metadata[layer_name]

598 attn_layer: Attention | MLAAttention = forward_context.no_compile_layers[layer_name]

602 - kv_cache = attn_layer.kv_cache[0]

599 + kv_cache = attn_layer.kv_cache

600 slot_mapping = forward_context.slot_mapping

601 assert isinstance(slot_mapping, dict), (

602 f"Expected slot_mapping to be a dict, got {type(slot_mapping)}. "

vllm/model_executor/layers/attention/mla_attention.py

+3 -8

@@ -415,12 +415,7 @@ class MLAAttention(nn.Module, AttentionLayerBase):

415 raise ValueError(f"Duplicate layer name: {prefix}")

416 compilation_config.static_forward_context[prefix] = self

417

418 - self.kv_cache = [

419 - torch.tensor([])

420 - for _ in range(

421 - get_current_vllm_config().parallel_config.pipeline_parallel_size

422 - )

423 - ]

418 + self.kv_cache = torch.tensor([])

419

420 self.use_sparse = use_sparse

421

@@ -479,7 +474,7 @@ class MLAAttention(nn.Module, AttentionLayerBase):

474 attn_metadata = forward_context.attn_metadata

475 if isinstance(attn_metadata, dict):

476 attn_metadata = attn_metadata[self.layer_name]

482 - self_kv_cache = self.kv_cache[0]

477 + self_kv_cache = self.kv_cache

478 slot_mapping = forward_context.slot_mapping

479

480 assert isinstance(slot_mapping, dict), (

@@ -939,7 +934,7 @@ def unified_mla_kv_cache_update(

934 return torch.empty(0, device=kv_c_normed.device, dtype=kv_c_normed.dtype)

935

936 attn_layer = forward_context.no_compile_layers[layer_name]

942 - kv_cache = attn_layer.kv_cache[0]

937 + kv_cache = attn_layer.kv_cache

938

939 slot_mapping = forward_context.slot_mapping

940 assert isinstance(slot_mapping, dict), (

vllm/model_executor/layers/attention/static_sink_attention.py

+1 -1

@@ -168,7 +168,7 @@ class StaticSinkAttention(Attention, CustomOp):

168 "sink_key and sink_value have not been prepared"

169 )

170 if not self.sink_populated:

171 - self_kv_cache = self.kv_cache[0]

171 + self_kv_cache = self.kv_cache

172 torch.ops.vllm.maybe_populate_sink(self_kv_cache, self.layer_name)

173

174 return super().forward(query, key, value, output_shape)

vllm/model_executor/layers/kda.py

+1 -1

@@ -306,7 +306,7 @@ class KimiDeltaAttention(nn.Module, MambaBase):

306 non_spec_query_start_loc = attn_metadata.non_spec_query_start_loc

307 non_spec_state_indices_tensor = attn_metadata.non_spec_state_indices_tensor # noqa: E501

308 num_actual_tokens = attn_metadata.num_actual_tokens

309 - constant_caches = self.kv_cache[0]

309 + constant_caches = self.kv_cache

310

311 q_proj_states = q_proj_states[:num_actual_tokens]

312 k_proj_states = k_proj_states[:num_actual_tokens]

vllm/model_executor/layers/mamba/linear_attn.py

+1 -1

@@ -413,7 +413,7 @@ class MiniMaxText01LinearAttention(nn.Module, MambaBase):

413 qkvact = qkvact.view((qkv.shape[0], self.tp_heads, -1))

414 q, k, v = torch.split(qkvact, [self.head_dim] * 3, dim=-1)

415 if attn_metadata is not None:

416 - kv_cache = self.kv_cache[0][0]

416 + kv_cache = self.kv_cache[0]

417 state_indices_tensor = attn_metadata.state_indices_tensor

418 clear_linear_attention_cache_for_new_sequences(

419 kv_cache, state_indices_tensor, attn_metadata

vllm/model_executor/layers/mamba/mamba_mixer.py

+1 -1

@@ -267,7 +267,7 @@ class MambaMixer(MambaBase, PluggableLayer):

267 query_start_loc_p = attn_metadata.query_start_loc_p

268 state_indices_tensor_p = attn_metadata.state_indices_tensor_p

269 state_indices_tensor_d = attn_metadata.state_indices_tensor_d

270 - self_kv_cache = self.kv_cache[0]

270 + self_kv_cache = self.kv_cache

271 conv_state = self_kv_cache[0].transpose(-1, -2)

272 ssm_state = self_kv_cache[1]

273 has_initial_states_p = attn_metadata.has_initial_states_p

vllm/model_executor/layers/mamba/mamba_mixer2.py

+1 -1

@@ -575,7 +575,7 @@ class MambaMixer2(MambaBase, PluggableLayer):

575 assert isinstance(attn_metadata, dict)

576 attn_metadata = attn_metadata[self.prefix]

577 assert isinstance(attn_metadata, Mamba2AttentionMetadata)

578 - self_kv_cache = self.kv_cache[0]

578 + self_kv_cache = self.kv_cache

579 # conv_state = (..., dim, width-1) yet contiguous along 'dim'

580 conv_state = self_kv_cache[0].transpose(-1, -2)

581 ssm_state = self_kv_cache[1]

vllm/model_executor/layers/mamba/short_conv.py

+1 -1

@@ -117,7 +117,7 @@ class ShortConv(MambaBase, CustomOp):

117 assert isinstance(attn_metadata, dict)

118 attn_metadata = attn_metadata[self.prefix]

119 assert isinstance(attn_metadata, ShortConvAttentionMetadata)

120 - self_kv_cache = self.kv_cache[0]

120 + self_kv_cache = self.kv_cache

121 conv_state = self_kv_cache[0].transpose(-1, -2)

122 state_indices_tensor_p = attn_metadata.state_indices_tensor_p

123 state_indices_tensor_d = attn_metadata.state_indices_tensor_d

vllm/model_executor/layers/sparse_attn_indexer.py

+2 -2

@@ -365,7 +365,7 @@ class SparseAttnIndexer(CustomOp):

365 return torch.ops.vllm.sparse_attn_indexer(

366 hidden_states,

367 self.k_cache.prefix,

368 - self.k_cache.kv_cache[0],

368 + self.k_cache.kv_cache,

369 q_fp8,

370 k,

371 weights,

@@ -389,7 +389,7 @@ class SparseAttnIndexer(CustomOp):

389 return torch.ops.vllm.rocm_aiter_sparse_attn_indexer(

390 hidden_states,

391 self.k_cache.prefix,

392 - self.k_cache.kv_cache[0],

392 + self.k_cache.kv_cache,

393 q_fp8,

394 k,

395 weights,

vllm/model_executor/models/bailing_moe_linear.py

+1 -1

@@ -709,7 +709,7 @@ class BailingMoELinearAttention(nn.Module, MambaBase):

709

710 # Get KV cache and state indices

711 if attn_metadata is not None:

712 - kv_cache = self.kv_cache[0][0]

712 + kv_cache = self.kv_cache[0]

713 state_indices_tensor = attn_metadata.state_indices_tensor

714 clear_linear_attention_cache_for_new_sequences(

715 kv_cache, state_indices_tensor, attn_metadata

vllm/model_executor/models/deepseek_v2.py

+1 -1

@@ -586,7 +586,7 @@ class DeepseekV32IndexerCache(torch.nn.Module, AttentionLayerBase):

586 self, head_dim: int, dtype: torch.dtype, prefix: str, cache_config: CacheConfig

587 ):

588 super().__init__()

589 - self.kv_cache = [torch.tensor([])]

589 + self.kv_cache = torch.tensor([])

590 self.head_dim = head_dim

591 self.prefix = prefix

592 self.cache_config = cache_config

vllm/model_executor/models/extract_hidden_states.py

+2 -5

@@ -51,7 +51,7 @@ def unified_kv_cache_update(

51 """

52 forward_context = get_forward_context()

53 attn_layer = forward_context.no_compile_layers[layer_name]

54 - kv_cache = attn_layer.kv_cache[0]

54 + kv_cache = attn_layer.kv_cache

55

56 slot_mapping = forward_context.slot_mapping

57 assert isinstance(slot_mapping, dict), (

@@ -288,10 +288,7 @@ class CacheOnlyAttentionLayer(nn.Module, AttentionLayerBase):

288 )

289

290 # Placeholder KV cache (replaced by bind_kv_cache)

291 - self.kv_cache = [

292 - torch.tensor([])

293 - for _ in range(vllm_config.parallel_config.pipeline_parallel_size)

294 - ]

291 + self.kv_cache = torch.tensor([])

292

293 # Register in compilation context

294 compilation_config = vllm_config.compilation_config

vllm/model_executor/models/olmo_hybrid.py

+1 -1

@@ -428,7 +428,7 @@ class OlmoHybridGatedDeltaNet(nn.Module, MambaBase):

428 non_spec_token_indx = attn_metadata.non_spec_token_indx

429 spec_state_indices_tensor = attn_metadata.spec_state_indices_tensor

430 non_spec_state_indices_tensor = attn_metadata.non_spec_state_indices_tensor

431 - self_kv_cache = self.kv_cache[0]

431 + self_kv_cache = self.kv_cache

432 conv_state = self_kv_cache[0].transpose(-1, -2)

433 ssm_state = self_kv_cache[1]

434 num_actual_tokens = attn_metadata.num_actual_tokens

vllm/model_executor/models/plamo2.py

+1 -1

@@ -262,7 +262,7 @@ class Plamo2MambaMixer(MambaBase, PluggableLayer):

262 assert isinstance(attn_metadata, dict)

263 attn_metadata = attn_metadata[self.prefix]

264 assert isinstance(attn_metadata, Mamba2AttentionMetadata)

265 - self_kv_cache = self.kv_cache[0]

265 + self_kv_cache = self.kv_cache

266 # conv_state = (..., dim, width-1) yet contiguous along 'dim'

267 conv_state = self_kv_cache[0].transpose(-1, -2)

268 ssm_state = self_kv_cache[1]

vllm/model_executor/models/qwen3_next.py

+2 -2

@@ -858,7 +858,7 @@ class Qwen3NextGatedDeltaNet(nn.Module, MambaBase):

858 non_spec_token_indx = attn_metadata.non_spec_token_indx

859 spec_state_indices_tensor = attn_metadata.spec_state_indices_tensor # noqa: E501

860 non_spec_state_indices_tensor = attn_metadata.non_spec_state_indices_tensor # noqa: E501

861 - self_kv_cache = self.kv_cache[0]

861 + self_kv_cache = self.kv_cache

862 conv_state = self_kv_cache[0].transpose(-1, -2)

863 ssm_state = self_kv_cache[1]

864 num_actual_tokens = attn_metadata.num_actual_tokens

@@ -1046,7 +1046,7 @@ class Qwen3NextGatedDeltaNet(nn.Module, MambaBase):

1046 Core attention computation with a packed non-spec decode fast path.

1047 """

1048 non_spec_state_indices_tensor = attn_metadata.non_spec_state_indices_tensor # noqa: E501

1049 - self_kv_cache = self.kv_cache[0]

1049 + self_kv_cache = self.kv_cache

1050 conv_state = self_kv_cache[0].transpose(-1, -2)

1051 ssm_state = self_kv_cache[1]

1052 num_actual_tokens = attn_metadata.num_actual_tokens

vllm/v1/attention/backends/rocm_aiter_fa.py

+3 -7

@@ -481,13 +481,9 @@ class AiterFlashAttentionMetadataBuilder(

481 ):

482 layers = get_layers_from_vllm_config(self.vllm_config, Attention)

483 first_layer_name = [k for k in layers][0]

484 - kv_cache_shape = (

485 - self.vllm_config.compilation_config.static_forward_context[

486 - first_layer_name

487 - ]

488 - .kv_cache[0]

489 - .shape

490 - )

484 + kv_cache_shape = self.vllm_config.compilation_config.static_forward_context[

485 + first_layer_name

486 + ].kv_cache.shape

487 num_blocks = kv_cache_shape[1]

488 self.scale = torch.ones(

489 [num_blocks, self.num_heads_kv, self.block_size],

vllm/v1/worker/gpu_model_runner.py

+4 -1

@@ -5830,7 +5830,10 @@ class GPUModelRunner(

5830

5831 for layer in self.compilation_config.static_forward_context.values():

5832 if hasattr(layer, "kv_cache"):

5833 - layer.kv_cache = []

5833 + kv_cache = layer.kv_cache

5834 + layer.kv_cache = (

5835 + torch.tensor([]) if isinstance(kv_cache, torch.Tensor) else []

5836 + )

5837

5838 gc.collect()

5839 torch.accelerator.empty_cache()

vllm/v1/worker/mamba_utils.py

+1 -1

@@ -119,7 +119,7 @@ def collect_mamba_copy_meta(

119 layer_names = kv_cache_config.kv_cache_groups[mamba_group_id].layer_names

120 for layer_name in layer_names:

121 attention = forward_context[layer_name]

122 - kv_caches: list[torch.Tensor] = attention.kv_cache[0]

122 + kv_caches: list[torch.Tensor] = attention.kv_cache

123 for state, state_copy_func in zip(kv_caches, mamba_state_copy_funcs):

124 copy_spec = state_copy_func(

125 state, block_ids, src_block_idx, accept_token_bias + 1

vllm/v1/worker/utils.py

+3 -4

@@ -136,8 +136,8 @@ class KVBlockZeroer:

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):

139 + kv = static_forward_context[layer_name].kv_cache

140 + if not isinstance(kv, torch.Tensor):

141 continue

142 dp = kv.data_ptr()

143 if dp in seen_ptrs:

@@ -510,8 +510,7 @@ def bind_kv_cache(

510

511 # Bind kv_caches to forward context

512 for layer_name, kv_cache in kv_caches.items():

513 - # NOTE: Keep list wrapper for layers that index kv_cache by engine slot.

514 - forward_context[layer_name].kv_cache = [kv_cache]

513 + forward_context[layer_name].kv_cache = kv_cache

514

515

516 def is_residual_scattered_for_sp(

[V0 Deprecation] Refactor kv cache from list to element (#37487)