cdc1fa12eb · tduncan/Nvidia-Optimization-Stack-VLLM

docs/source/contributing/model/basic.md

-2

@@ -74,8 +74,6 @@ def forward(

74 self,

75 input_ids: torch.Tensor,

76 positions: torch.Tensor,

77 - kv_caches: List[torch.Tensor],

78 - attn_metadata: AttentionMetadata,

77 ) -> torch.Tensor:

78 ...

79 ```

docs/source/contributing/model/multimodal.md

-2

@@ -16,8 +16,6 @@ Further update the model as follows:

16 self,

17 input_ids: torch.Tensor,

18 positions: torch.Tensor,

19 - kv_caches: List[torch.Tensor],

20 - attn_metadata: AttentionMetadata,

19 + pixel_values: torch.Tensor,

20 ) -> SamplerOutput:

21 ```

tests/kernels/test_encoder_decoder_attn.py

+3 -11

@@ -644,11 +644,7 @@ def _run_encoder_attention_test(

644 # is shaped as [num_tokens, hidden_size] and we can skip the reshape.

645 reshaped_query = packed_qkv.query.view(

646 -1, test_pt.num_heads * test_pt.head_size)

647 - return attn.forward(

648 - reshaped_query, packed_qkv.key, packed_qkv.value,

649 - torch.tensor([],

650 - dtype=torch.float32,

651 - device=packed_qkv.query.device), attn_metadata)

647 + return attn.forward(reshaped_query, packed_qkv.key, packed_qkv.value)

648

649

650 def _run_decoder_self_attention_test(

@@ -682,7 +678,6 @@ def _run_decoder_self_attention_test(

678 & attn_metadata

679 '''

680 attn = test_rsrcs.attn

685 - kv_cache = test_rsrcs.kv_cache

681 packed_qkv = decoder_test_params.packed_qkvo.packed_qkv

682 assert packed_qkv is not None

683 with set_forward_context(attn_metadata, vllm_config):

@@ -695,8 +690,7 @@ def _run_decoder_self_attention_test(

690 # is shaped as [num_tokens, hidden_size] and we can skip the reshape.

691 reshaped_query = packed_qkv.query.view(

692 -1, test_pt.num_heads * test_pt.head_size)

698 - return attn.forward(reshaped_query, packed_qkv.key, packed_qkv.value,

699 - kv_cache, attn_metadata)

693 + return attn.forward(reshaped_query, packed_qkv.key, packed_qkv.value)

694

695

696 def _run_encoder_decoder_cross_attention_test(

@@ -744,7 +738,6 @@ def _run_encoder_decoder_cross_attention_test(

738 assert decoder_test_params.packed_qkvo.packed_qkv is not None

739

740 attn = test_rsrcs.attn

747 - kv_cache = test_rsrcs.kv_cache

741 if cross_test_params is None:

742 key = None

743 value = None

@@ -762,8 +755,7 @@ def _run_encoder_decoder_cross_attention_test(

755 # is shaped as [num_tokens, hidden_size] and we can skip the reshape.

756 reshaped_query = decoder_test_params.packed_qkvo.packed_qkv.query.view(

757 -1, test_pt.num_heads * test_pt.head_size)

765 - return attn.forward(reshaped_query, key, value, kv_cache,

766 - attn_metadata)

758 + return attn.forward(reshaped_query, key, value)

759

760

761 @pytest.fixture(autouse=True)

vllm/attention/layer.py

+7 -12

@@ -7,7 +7,7 @@ import torch.nn as nn

7 import torch.nn.functional as F

8

9 import vllm.envs as envs

10 -from vllm.attention import AttentionMetadata, AttentionType

10 +from vllm.attention import AttentionType

11 from vllm.attention.selector import backend_name_to_enum, get_attn_backend

12 from vllm.config import CacheConfig, get_current_vllm_config

13 from vllm.forward_context import ForwardContext, get_forward_context

@@ -153,15 +153,10 @@ class Attention(nn.Module):

153 query: torch.Tensor,

154 key: torch.Tensor,

155 value: torch.Tensor,

156 - kv_cache: torch.Tensor,

157 - attn_metadata: AttentionMetadata,

156 ) -> torch.Tensor:

159 - # NOTE: please avoid accessing `kv_cache` and `attn_metadata` arguments

160 - # directly, use `self.kv_cache` and

161 - # `get_forward_context().attn_metadata` instead.

157 if self.calculate_kv_scales:

163 - ctx_attn_metadata = get_forward_context().attn_metadata

164 - if ctx_attn_metadata.enable_kv_scales_calculation:

158 + attn_metadata = get_forward_context().attn_metadata

159 + if attn_metadata.enable_kv_scales_calculation:

160 self.calc_kv_scales(key, value)

161 if self.use_output:

162 output = torch.empty_like(query)

@@ -177,14 +172,14 @@ class Attention(nn.Module):

172 value = value.view(-1, self.num_kv_heads, self.head_size)

173 if self.use_direct_call:

174 forward_context: ForwardContext = get_forward_context()

180 - ctx_attn_metadata = forward_context.attn_metadata

175 + attn_metadata = forward_context.attn_metadata

176 self_kv_cache = self.kv_cache[forward_context.virtual_engine]

177 self.impl.forward(self,

178 query,

179 key,

180 value,

181 self_kv_cache,

187 - ctx_attn_metadata,

182 + attn_metadata,

183 output=output)

184 else:

185 torch.ops.vllm.unified_attention_with_output(

@@ -193,10 +188,10 @@ class Attention(nn.Module):

188 else:

189 if self.use_direct_call:

190 forward_context = get_forward_context()

196 - ctx_attn_metadata = forward_context.attn_metadata

191 + attn_metadata = forward_context.attn_metadata

192 self_kv_cache = self.kv_cache[forward_context.virtual_engine]

193 return self.impl.forward(self, query, key, value,

199 - self_kv_cache, ctx_attn_metadata)

194 + self_kv_cache, attn_metadata)

195 else:

196 return torch.ops.vllm.unified_attention(

197 query, key, value, self.layer_name)

vllm/model_executor/layers/mamba/mamba_mixer.py

+3 -2

@@ -7,6 +7,7 @@ from torch.nn.parameter import Parameter

7 from vllm.attention.backends.abstract import AttentionMetadata

8 from vllm.distributed.parallel_state import (

9 get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size)

10 +from vllm.forward_context import get_forward_context

11 from vllm.model_executor.custom_op import CustomOp

12 from vllm.model_executor.layers.layernorm import RMSNorm

13 from vllm.model_executor.layers.linear import (ColumnParallelLinear,

@@ -130,14 +131,14 @@ class MambaMixer(CustomOp):

131 ) if use_rms_norm else None

132

133 def forward_native(self, hidden_states: torch.Tensor,

133 - attn_metadata: AttentionMetadata,

134 conv_state: torch.Tensor, ssm_state: torch.Tensor):

135 pass

136

137 def forward_cuda(self, hidden_states: torch.Tensor,

138 - attn_metadata: AttentionMetadata,

138 mamba_cache_params: MambaCacheParams):

139

140 + attn_metadata: AttentionMetadata = get_forward_context().attn_metadata

141 +

142 # 1. Gated MLP's linear projection

143 projected_states = self.in_proj(hidden_states)[0].transpose(-2, -1)

144 hidden_states, gate = projected_states.chunk(2, dim=-2)

vllm/model_executor/layers/mamba/mamba_mixer2.py

+2 -2

@@ -14,6 +14,7 @@ from vllm.distributed import (divide, get_tensor_model_parallel_rank,

14 get_tensor_model_parallel_world_size,

15 tensor_model_parallel_all_gather,

16 tensor_model_parallel_all_reduce)

17 +from vllm.forward_context import get_forward_context

18 from vllm.model_executor.custom_op import CustomOp

19 from vllm.model_executor.layers.linear import (ColumnParallelLinear,

20 RowParallelLinear)

@@ -376,17 +377,16 @@ class MambaMixer2(CustomOp):

377 eps=rms_norm_eps)

378

379 def forward_native(self, hidden_states: torch.Tensor,

379 - attn_metadata: AttentionMetadata,

380 conv_state: torch.Tensor, ssm_state: torch.Tensor):

381 pass

382

383 def forward_cuda(

384 self,

385 hidden_states: torch.Tensor,

386 - attn_metadata: AttentionMetadata,

386 mamba_cache_params: MambaCacheParams,

387 sequence_idx: Optional[torch.Tensor] = None,

388 ):

389 + attn_metadata: AttentionMetadata = get_forward_context().attn_metadata

390

391 seq_len, _ = hidden_states.shape

392 groups_time_state_size = self.n_groups * self.ssm_state_size

vllm/model_executor/models/adapters.py

+1 -5

@@ -160,7 +160,6 @@ def as_classification_model(cls: _T) -> _T:

160 return cls

161

162 # Lazy import

163 - from vllm.attention import AttentionMetadata

163 from vllm.config import VllmConfig

164 from vllm.model_executor.layers.linear import RowParallelLinear

165 from vllm.model_executor.layers.pooler import PoolingType

@@ -201,13 +200,10 @@ def as_classification_model(cls: _T) -> _T:

200 self,

201 input_ids: torch.Tensor,

202 positions: torch.Tensor,

204 - kv_caches: list[torch.Tensor],

205 - attn_metadata: AttentionMetadata,

203 intermediate_tensors: Optional[IntermediateTensors] = None,

204 inputs_embeds: Optional[torch.Tensor] = None,

205 ) -> torch.Tensor:

209 - hidden_states = super().forward(input_ids, positions, kv_caches,

210 - attn_metadata,

206 + hidden_states = super().forward(input_ids, positions,

207 intermediate_tensors,

208 inputs_embeds)

209 logits, _ = self.score(hidden_states)

vllm/model_executor/models/arctic.py

+5 -19

@@ -5,7 +5,7 @@ from typing import Iterable, List, Optional, Set, Tuple, Union

5 import torch

6 from torch import nn

7

8 -from vllm.attention import Attention, AttentionMetadata

8 +from vllm.attention import Attention

9 from vllm.compilation.decorators import support_torch_compile

10 from vllm.config import CacheConfig, VllmConfig

11 from vllm.distributed import (get_pp_group, get_tensor_model_parallel_rank,

@@ -283,13 +283,11 @@ class ArcticAttention(nn.Module):

283 self,

284 positions: torch.Tensor,

285 hidden_states: torch.Tensor,

286 - kv_cache: torch.Tensor,

287 - attn_metadata: AttentionMetadata,

286 ) -> torch.Tensor:

287 qkv, _ = self.qkv_proj(hidden_states)

288 q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)

289 q, k = self.rotary_emb(positions, q, k)

292 - attn_output = self.attn(q, k, v, kv_cache, attn_metadata)

290 + attn_output = self.attn(q, k, v)

291 output, _ = self.o_proj(attn_output)

292 return output

293

@@ -336,16 +334,12 @@ class ArcticDecoderLayer(nn.Module):

334 self,

335 positions: torch.Tensor,

336 hidden_states: torch.Tensor,

339 - kv_cache: torch.Tensor,

340 - attn_metadata: AttentionMetadata,

337 ) -> torch.Tensor:

338 residual_input = hidden_states

339 hidden_states = self.input_layernorm(hidden_states)

340 hidden_states = self.self_attn(

341 positions=positions,

342 hidden_states=hidden_states,

347 - kv_cache=kv_cache,

348 - attn_metadata=attn_metadata,

343 )

344 hidden_states = residual_input + hidden_states

345

@@ -400,8 +394,6 @@ class ArcticModel(nn.Module):

394 self,

395 input_ids: torch.Tensor,

396 positions: torch.Tensor,

403 - kv_caches: List[torch.Tensor],

404 - attn_metadata: AttentionMetadata,

397 intermediate_tensors: Optional[IntermediateTensors],

398 inputs_embeds: Optional[torch.Tensor] = None,

399 ) -> Union[torch.Tensor, IntermediateTensors]:

@@ -413,11 +405,8 @@ class ArcticModel(nn.Module):

405 else:

406 assert intermediate_tensors is not None

407 hidden_states = intermediate_tensors["hidden_states"]

416 - for i in range(self.start_layer, self.end_layer):

417 - layer = self.layers[i]

418 - hidden_states = layer(positions, hidden_states,

419 - kv_caches[i - self.start_layer],

420 - attn_metadata)

408 + for layer in self.layers[self.start_layer:self.end_layer]:

409 + hidden_states = layer(positions, hidden_states)

410 if not get_pp_group().is_last_rank:

411 return IntermediateTensors({"hidden_states": hidden_states})

412 hidden_states = self.norm(hidden_states)

@@ -458,13 +447,10 @@ class ArcticForCausalLM(nn.Module, SupportsPP, SupportsQuant):

447 self,

448 input_ids: torch.Tensor,

449 positions: torch.Tensor,

461 - kv_caches: List[torch.Tensor],

462 - attn_metadata: AttentionMetadata,

450 intermediate_tensors: Optional[IntermediateTensors] = None,

451 inputs_embeds: Optional[torch.Tensor] = None,

452 ) -> Union[torch.Tensor, IntermediateTensors]:

466 - hidden_states = self.model(input_ids, positions, kv_caches,

467 - attn_metadata, intermediate_tensors,

453 + hidden_states = self.model(input_ids, positions, intermediate_tensors,

454 inputs_embeds)

455 return hidden_states

456

vllm/model_executor/models/aria.py

-5

@@ -9,7 +9,6 @@ from transformers import AriaConfig, AriaTextConfig, BatchFeature

9 from transformers.models.aria.modeling_aria import AriaCrossAttention

10 from transformers.models.aria.processing_aria import AriaProcessor

11

12 -from vllm.attention import AttentionMetadata

12 from vllm.config import CacheConfig, QuantizationConfig, VllmConfig

13 from vllm.distributed import get_tensor_model_parallel_rank

14 from vllm.model_executor.layers.activation import get_act_fn

@@ -626,8 +625,6 @@ class AriaForConditionalGeneration(nn.Module, SupportsMultiModal):

625 self,

626 input_ids: torch.Tensor,

627 positions: torch.Tensor,

629 - kv_caches: List[torch.Tensor],

630 - attn_metadata: AttentionMetadata,

628 intermediate_tensors: Optional[IntermediateTensors] = None,

629 inputs_embeds: Optional[torch.Tensor] = None,

630 **kwargs: object,

@@ -643,8 +640,6 @@ class AriaForConditionalGeneration(nn.Module, SupportsMultiModal):

640 hidden_states = self.language_model(

641 input_ids,

642 positions,

646 - kv_caches,

647 - attn_metadata,

643 intermediate_tensors,

644 inputs_embeds=inputs_embeds,

645 )

vllm/model_executor/models/baichuan.py

+5 -19

@@ -20,13 +20,13 @@

20 # limitations under the License.

21 """Inference-only BaiChuan model compatible with HuggingFace weights."""

22 import math

23 -from typing import Iterable, List, Optional, Set, Tuple, Union

23 +from typing import Iterable, Optional, Set, Tuple, Union

24

25 import torch

26 from torch import nn

27 from transformers import PretrainedConfig

28

29 -from vllm.attention import Attention, AttentionMetadata

29 +from vllm.attention import Attention

30 from vllm.compilation.decorators import support_torch_compile

31 from vllm.config import CacheConfig, VllmConfig

32 from vllm.distributed import (get_pp_group, get_tensor_model_parallel_rank,

@@ -182,14 +182,12 @@ class BaiChuanAttention(nn.Module):

182 self,

183 positions: torch.Tensor,

184 hidden_states: torch.Tensor,

185 - kv_cache: torch.Tensor,

186 - attn_metadata: AttentionMetadata,

185 ) -> torch.Tensor:

186 qkv, _ = self.W_pack(hidden_states)

187 q, k, v = qkv.chunk(chunks=3, dim=-1)

188 if self.postion_embedding != "ALIBI":

189 q, k = self.rotary_emb(positions, q, k)

192 - attn_output = self.attn(q, k, v, kv_cache, attn_metadata)

190 + attn_output = self.attn(q, k, v)

191 output, _ = self.o_proj(attn_output)

192 return output

193

@@ -232,8 +230,6 @@ class BaiChuanDecoderLayer(nn.Module):

230 self,

231 positions: torch.Tensor,

232 hidden_states: torch.Tensor,

235 - kv_cache: torch.Tensor,

236 - attn_metadata: AttentionMetadata,

233 residual: Optional[torch.Tensor],

234 ) -> Tuple[torch.Tensor, torch.Tensor]:

235 # Self Attention

@@ -246,8 +242,6 @@ class BaiChuanDecoderLayer(nn.Module):

242 hidden_states = self.self_attn(

243 positions=positions,

244 hidden_states=hidden_states,

249 - kv_cache=kv_cache,

250 - attn_metadata=attn_metadata,

245 )

246

247 # Fully Connected

@@ -301,8 +295,6 @@ class BaiChuanModel(nn.Module):

295 self,

296 input_ids: torch.Tensor,

297 positions: torch.Tensor,

304 - kv_caches: List[torch.Tensor],

305 - attn_metadata: AttentionMetadata,

298 intermediate_tensors: Optional[IntermediateTensors],

299 inputs_embeds: Optional[torch.Tensor] = None,

300 ) -> Union[torch.Tensor, IntermediateTensors]:

@@ -316,13 +308,10 @@ class BaiChuanModel(nn.Module):

308 assert intermediate_tensors is not None

309 hidden_states = intermediate_tensors["hidden_states"]

310 residual = intermediate_tensors["residual"]

319 - for i in range(self.start_layer, self.end_layer):

320 - layer = self.layers[i]

311 + for layer in self.layers[self.start_layer:self.end_layer]:

312 hidden_states, residual = layer(

313 positions,

314 hidden_states,

324 - kv_caches[i - self.start_layer],

325 - attn_metadata,

315 residual,

316 )

317 if not get_pp_group().is_last_rank:

@@ -379,13 +368,10 @@ class BaiChuanBaseForCausalLM(nn.Module, SupportsLoRA, SupportsPP,

368 self,

369 input_ids: torch.Tensor,

370 positions: torch.Tensor,

382 - kv_caches: List[torch.Tensor],

383 - attn_metadata: AttentionMetadata,

371 intermediate_tensors: Optional[IntermediateTensors] = None,

372 inputs_embeds: Optional[torch.Tensor] = None,

373 ) -> Union[torch.Tensor, IntermediateTensors]:

387 - hidden_states = self.model(input_ids, positions, kv_caches,

388 - attn_metadata, intermediate_tensors,

374 + hidden_states = self.model(input_ids, positions, intermediate_tensors,

375 inputs_embeds)

376 return hidden_states

377

vllm/model_executor/models/bamba.py

+7 -22

@@ -1,17 +1,17 @@

1 # SPDX-License-Identifier: Apache-2.0

2 """Inference-only Bamba model."""

3 # Added by the IBM Team, 2024

4 -from typing import Iterable, List, Optional, Set, Tuple

4 +from typing import Iterable, Optional, Set, Tuple

5

6 import torch

7 from torch import nn

8 from transformers import BambaConfig

9

10 -from vllm.attention.backends.abstract import AttentionMetadata

10 from vllm.attention.layer import Attention

11 from vllm.config import CacheConfig, VllmConfig

12 from vllm.distributed import divide, get_tensor_model_parallel_world_size

13 from vllm.distributed.parallel_state import get_pp_group

14 +from vllm.forward_context import get_forward_context

15 from vllm.model_executor.layers.activation import SiluAndMul

16 from vllm.model_executor.layers.layernorm import RMSNorm

17 from vllm.model_executor.layers.linear import (MergedColumnParallelLinear,

@@ -107,7 +107,6 @@ class BambaMixerDecoderLayer(nn.Module):

107 def forward(

108 self,

109 hidden_states: torch.Tensor,

110 - attn_metadata: AttentionMetadata,

110 residual: Optional[torch.Tensor],

111 mamba_cache_params: MambaCacheParams,

112 sequence_idx: Optional[torch.Tensor] = None,

@@ -120,8 +119,8 @@ class BambaMixerDecoderLayer(nn.Module):

119 hidden_states, residual = self.input_layernorm(

120 hidden_states, residual)

121

123 - hidden_states = self.mamba(hidden_states, attn_metadata,

124 - mamba_cache_params, sequence_idx)

122 + hidden_states = self.mamba(hidden_states, mamba_cache_params,

123 + sequence_idx)

124 # Fully Connected

125 hidden_states, residual = self.pre_ff_layernorm(

126 hidden_states, residual)

@@ -215,15 +214,13 @@ class BambaAttentionDecoderLayer(nn.Module):

214 self,

215 positions: torch.Tensor,

216 hidden_states: torch.Tensor,

218 - kv_cache: torch.Tensor,

219 - attn_metadata: AttentionMetadata,

217 **kwargs,

218 ) -> torch.Tensor:

219 qkv, _ = self.qkv_proj(hidden_states)

220 q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)

221

222 q, k = self.rotary_emb(positions, q, k)

226 - attn_output = self.attn(q, k, v, kv_cache, attn_metadata)

223 + attn_output = self.attn(q, k, v)

224 output, _ = self.o_proj(attn_output)

225 return output

226

@@ -231,8 +228,6 @@ class BambaAttentionDecoderLayer(nn.Module):

228 self,

229 positions: torch.Tensor,

230 hidden_states: torch.Tensor,

234 - kv_cache: torch.Tensor,

235 - attn_metadata: AttentionMetadata,

231 residual: Optional[torch.Tensor],

232 **kwargs,

233 ):

@@ -246,8 +241,6 @@ class BambaAttentionDecoderLayer(nn.Module):

241 hidden_states = self.self_attention(

242 positions=positions,

243 hidden_states=hidden_states,

249 - kv_cache=kv_cache,

250 - attn_metadata=attn_metadata,

244 )

245 # Fully Connected

246 hidden_states, residual = self.pre_ff_layernorm(

@@ -312,8 +305,6 @@ class BambaModel(nn.Module):

305 self,

306 input_ids: torch.Tensor,

307 positions: torch.Tensor,

315 - kv_caches: List[torch.Tensor],

316 - attn_metadata: AttentionMetadata,

308 mamba_cache_params: MambaCacheParams,

309 intermediate_tensors: Optional[IntermediateTensors] = None,

310 inputs_embeds: Optional[torch.Tensor] = None,

@@ -323,6 +314,7 @@ class BambaModel(nn.Module):

314 # proper continuous batching computation including

315 # chunked prefill

316 seq_idx = None

317 + attn_metadata = get_forward_context().attn_metadata

318 if attn_metadata.num_prefills > 0:

319 seq_idx = torch.zeros_like(input_ids, dtype=torch.int32)

320 for i, (srt, end) in enumerate(

@@ -348,9 +340,7 @@ class BambaModel(nn.Module):

340 num_attn = 0

341 for i in range(len(self.layers)):

342 layer = self.layers[i]

351 - kv_cache = None

343 if isinstance(layer, BambaAttentionDecoderLayer):

353 - kv_cache = kv_caches[num_attn]

344 num_attn += 1

345

346 layer_mamba_cache_params = None

@@ -361,8 +351,6 @@ class BambaModel(nn.Module):

351 hidden_states, residual = layer(

352 positions=positions,

353 hidden_states=hidden_states,

364 - kv_cache=kv_cache,

365 - attn_metadata=attn_metadata,

354 residual=residual,

355 mamba_cache_params=layer_mamba_cache_params,

356 sequence_idx=seq_idx,

@@ -440,8 +428,6 @@ class BambaForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP,

428 def forward(self,

429 input_ids: torch.Tensor,

430 positions: torch.Tensor,

443 - kv_caches: List[KVCache],

444 - attn_metadata: AttentionMetadata,

431 intermediate_tensors: Optional[IntermediateTensors] = None,

432 inputs_embeds: Optional[torch.Tensor] = None,

433 **kwargs):

@@ -454,8 +440,7 @@ class BambaForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP,

440 self.vllm_config, self.lm_head.weight.dtype, num_mamba_layers,

441 *self._get_mamba_cache_shape())

442 mamba_cache_params = self.mamba_cache.current_run_tensors(**kwargs)

457 - hidden_states = self.model(input_ids, positions, kv_caches,

458 - attn_metadata, mamba_cache_params,

443 + hidden_states = self.model(input_ids, positions, mamba_cache_params,

444 intermediate_tensors, inputs_embeds)

445

446 return hidden_states

vllm/model_executor/models/bart.py

+20 -73

@@ -19,14 +19,14 @@

19 # limitations under the License.

20 """PyTorch BART model."""

21 import math

22 -from typing import Iterable, List, Optional, Tuple

22 +from typing import Iterable, Optional, Tuple

23

24 import torch

25 from torch import nn

26 from transformers import BartConfig

27 from transformers.utils import logging

28

29 -from vllm.attention import Attention, AttentionMetadata, AttentionType

29 +from vllm.attention import Attention, AttentionType

30 from vllm.config import CacheConfig, LoRAConfig, VllmConfig

31 from vllm.distributed import get_tensor_model_parallel_world_size

32 from vllm.model_executor.layers.activation import get_act_fn

@@ -181,14 +181,13 @@ class BartEncoderAttention(nn.Module):

181 prefix=f"{prefix}.attn",

182 attn_type=AttentionType.ENCODER)

183

184 - def forward(self, hidden_states: torch.Tensor, kv_cache: torch.Tensor,

185 - attn_metadata: AttentionMetadata) -> torch.Tensor:

184 + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:

185 """Input shape: Batch x Time x Channel"""

186

187 qkv, _ = self.qkv_proj(hidden_states)

188 q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)

189

191 - attn_output = self.attn(q, k, v, kv_cache, attn_metadata)

190 + attn_output = self.attn(q, k, v)

191

192 output, _ = self.out_proj(attn_output)

193 return output

@@ -261,14 +260,13 @@ class BartDecoderSelfAttention(nn.Module):

260 prefix=f"{prefix}.attn",

261 attn_type=AttentionType.DECODER)

262

264 - def forward(self, hidden_states: torch.Tensor, kv_cache: torch.Tensor,

265 - attn_metadata: AttentionMetadata) -> torch.Tensor:

263 + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:

264 """Input shape: Batch x Time x Channel"""

265

266 qkv, _ = self.qkv_proj(hidden_states)

267 q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)

268

271 - attn_output = self.attn(q, k, v, kv_cache, attn_metadata)

269 + attn_output = self.attn(q, k, v)

270

271 output, _ = self.out_proj(attn_output)

272 return output

@@ -344,8 +342,6 @@ class BartCrossAttention(nn.Module):

342 def forward(

343 self,

344 decoder_hidden_states: torch.Tensor,

347 - kv_cache: torch.Tensor,

348 - attn_metadata: AttentionMetadata,

345 encoder_hidden_states: Optional[torch.Tensor] = None,

346 ) -> torch.Tensor:

347 """Input shape: Batch x Time x Channel"""

@@ -363,7 +359,7 @@ class BartCrossAttention(nn.Module):

359 _, k, v = qkv_enc.split([self.q_size, self.kv_size, self.kv_size],

360 dim=-1)

361

366 - attn_output = self.attn(q, k, v, kv_cache, attn_metadata)

362 + attn_output = self.attn(q, k, v)

363

364 output, _ = self.out_proj(attn_output)

365 return output

@@ -411,23 +407,16 @@ class BartEncoderLayer(nn.Module):

407

408 self.final_layer_norm = nn.LayerNorm(self.embed_dim)

409

414 - def forward(self, hidden_states: torch.Tensor, kv_cache: torch.Tensor,

415 - attn_metadata: AttentionMetadata) -> torch.Tensor:

410 + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:

411 r"""

412 Args:

413 hidden_states

414 torch.Tensor of *encoder* input embeddings.

420 - kv_cache:

421 - Layer-wise list of KV cache tensors

422 - attn_metadata:

423 - vLLM Attention metadata structure

415 Returns:

416 Encoder layer output torch.Tensor

417 """

418 residual = hidden_states

428 - hidden_states = self.self_attn(hidden_states=hidden_states,

429 - kv_cache=kv_cache,

430 - attn_metadata=attn_metadata)

419 + hidden_states = self.self_attn(hidden_states=hidden_states)

420

421 hidden_states = residual + hidden_states

422 hidden_states = self.self_attn_layer_norm(hidden_states)

@@ -509,18 +498,12 @@ class BartDecoderLayer(nn.Module):

498 def forward(

499 self,

500 decoder_hidden_states: torch.Tensor,

512 - kv_cache: torch.Tensor,

513 - attn_metadata: AttentionMetadata,

501 encoder_hidden_states: Optional[torch.Tensor] = None,

502 ) -> torch.Tensor:

503 r"""

504 Args:

505 decoder_hidden_states

506 torch.Tensor of *decoder* input embeddings.

520 - kv_cache:

521 - KV cache tensor

522 - attn_metadata:

523 - vLLM Attention metadata structure

507 encoder_hidden_states

508 torch.Tensor of *encoder* input embeddings.

509 Returns:

@@ -529,9 +512,7 @@ class BartDecoderLayer(nn.Module):

512 residual = decoder_hidden_states

513

514 # Self Attention

532 - hidden_states = self.self_attn(hidden_states=decoder_hidden_states,

533 - kv_cache=kv_cache,

534 - attn_metadata=attn_metadata)

515 + hidden_states = self.self_attn(hidden_states=decoder_hidden_states)

516

517 hidden_states = residual + hidden_states

518 hidden_states = self.self_attn_layer_norm(hidden_states)

@@ -542,8 +523,6 @@ class BartDecoderLayer(nn.Module):

523

524 hidden_states = self.encoder_attn(

525 decoder_hidden_states=hidden_states,

545 - kv_cache=kv_cache,

546 - attn_metadata=attn_metadata,

526 encoder_hidden_states=encoder_hidden_states,

527 )

528

@@ -609,9 +588,8 @@ class BartEncoder(nn.Module):

588

589 self.layernorm_embedding = nn.LayerNorm(embed_dim)

590

612 - def forward(self, input_ids: torch.Tensor, positions: torch.Tensor,

613 - kv_caches: List[torch.Tensor],

614 - attn_metadata: AttentionMetadata) -> torch.Tensor:

591 + def forward(self, input_ids: torch.Tensor,

592 + positions: torch.Tensor) -> torch.Tensor:

593 r"""

594 Args:

595 input_ids

@@ -620,10 +598,6 @@ class BartEncoder(nn.Module):

598 provide it.

599 positions

600 Positions of *encoder* input sequence tokens.

623 - kv_caches:

624 - Layer-wise list of KV cache tensors

625 - attn_metadata:

626 - vLLM Attention metadata structure

601 Returns:

602 Decoder output torch.Tensor

603 """

@@ -636,12 +610,8 @@ class BartEncoder(nn.Module):

610 hidden_states = inputs_embeds + embed_pos

611 hidden_states = self.layernorm_embedding(hidden_states)

612

639 - for idx, encoder_layer in enumerate(self.layers):

640 - hidden_states = encoder_layer(

641 - hidden_states=hidden_states,

642 - kv_cache=kv_caches[idx],

643 - attn_metadata=attn_metadata,

644 - )

613 + for encoder_layer in self.layers:

614 + hidden_states = encoder_layer(hidden_states=hidden_states)

615

616 return hidden_states

617

@@ -693,9 +663,7 @@ class BartDecoder(nn.Module):

663

664 def forward(self, decoder_input_ids: torch.Tensor,

665 decoder_positions: torch.Tensor,

696 - encoder_hidden_states: Optional[torch.Tensor],

697 - kv_caches: List[torch.Tensor],

698 - attn_metadata: AttentionMetadata) -> torch.Tensor:

666 + encoder_hidden_states: Optional[torch.Tensor]) -> torch.Tensor:

667 r"""

668 Args:

669 decoder_input_ids

@@ -706,10 +674,6 @@ class BartDecoder(nn.Module):

674 Positions of *decoder* input sequence tokens.

675 encoder_hidden_states:

676 Tensor of encoder output embeddings

709 - kv_caches:

710 - Layer-wise list of KV cache tensors

711 - attn_metadata:

712 - vLLM Attention metadata structure

677 Returns:

678 Decoder output torch.Tensor

679 """

@@ -725,11 +689,9 @@ class BartDecoder(nn.Module):

689

690 # decoder layers

691

728 - for idx, decoder_layer in enumerate(self.layers):

692 + for decoder_layer in self.layers:

693 hidden_states = decoder_layer(

694 decoder_hidden_states=hidden_states,

731 - kv_cache=kv_caches[idx],

732 - attn_metadata=attn_metadata,

695 encoder_hidden_states=encoder_hidden_states,

696 )

697

@@ -768,8 +730,7 @@ class BartModel(nn.Module):

730

731 def forward(self, input_ids: torch.Tensor, positions: torch.Tensor,

732 encoder_input_ids: torch.Tensor,

771 - encoder_positions: torch.Tensor, kv_caches: List[torch.Tensor],

772 - attn_metadata: AttentionMetadata) -> torch.Tensor:

733 + encoder_positions: torch.Tensor) -> torch.Tensor:

734 r"""

735 Args:

736 input_ids

@@ -782,10 +743,6 @@ class BartModel(nn.Module):

743 Indices of *encoder* input sequence tokens in the vocabulary.

744 encoder_positions:

745 Positions of *encoder* input sequence tokens.

785 - kv_caches:

786 - Layer-wise list of KV cache tensors

787 - attn_metadata:

788 - vLLM Attention metadata structure

746 Returns:

747 Model output torch.Tensor

748 """

@@ -796,18 +753,14 @@ class BartModel(nn.Module):

753 # Run encoder attention if a non-zero number of encoder tokens

754 # are provided as input

755 encoder_hidden_states = self.encoder(input_ids=encoder_input_ids,

799 - positions=encoder_positions,

800 - kv_caches=kv_caches,

801 - attn_metadata=attn_metadata)

756 + positions=encoder_positions)

757

758 # decoder outputs consists of

759 # (dec_features, past_key_value, dec_hidden, dec_attn)

760 decoder_outputs = self.decoder(

761 decoder_input_ids=input_ids,

762 decoder_positions=positions,

808 - encoder_hidden_states=encoder_hidden_states,

809 - kv_caches=kv_caches,

810 - attn_metadata=attn_metadata)

763 + encoder_hidden_states=encoder_hidden_states)

764

765 return decoder_outputs

766

@@ -845,8 +798,6 @@ class BartForConditionalGeneration(nn.Module):

798 self,

799 input_ids: torch.Tensor,

800 positions: torch.Tensor,

848 - kv_caches: List[torch.Tensor],

849 - attn_metadata: AttentionMetadata,

801 intermediate_tensors: Optional[IntermediateTensors] = None,

802 *,

803 encoder_input_ids: torch.Tensor,

@@ -863,15 +814,11 @@ class BartForConditionalGeneration(nn.Module):

814 torch.Tensor of *encoder* input token ids.

815 encoder_positions

816 torch.Tensor of *encoder* position indices

866 - kv_caches:

867 - Layer-wise list of KV cache tensors

868 - attn_metadata:

869 - vLLM Attention metadata structure

817 Returns:

818 Output torch.Tensor

819 """

820 return self.model(input_ids, positions, encoder_input_ids,

874 - encoder_positions, kv_caches, attn_metadata)

821 + encoder_positions)

822

823 def compute_logits(

824 self,

vllm/model_executor/models/bert.py

+12 -32

@@ -1,15 +1,16 @@

1 # SPDX-License-Identifier: Apache-2.0

2

3 -from typing import Iterable, List, Optional, Set, Tuple

3 +from typing import Iterable, Optional, Set, Tuple

4

5 import torch

6 from torch import nn

7 from transformers import BertConfig

8

9 -from vllm.attention import Attention, AttentionMetadata, AttentionType

9 +from vllm.attention import Attention, AttentionType

10 from vllm.compilation.decorators import support_torch_compile

11 from vllm.config import CacheConfig, PoolerConfig, VllmConfig

12 from vllm.distributed import get_tensor_model_parallel_world_size

13 +from vllm.forward_context import get_forward_context

14 from vllm.model_executor.layers.activation import get_act_fn

15 from vllm.model_executor.layers.linear import (ColumnParallelLinear,

16 QKVParallelLinear,

@@ -113,12 +114,9 @@ class BertEncoder(nn.Module):

114 def forward(

115 self,

116 hidden_states: torch.Tensor,

116 - kv_caches: List[torch.Tensor],

117 - attn_metadata: AttentionMetadata,

117 ) -> torch.Tensor:

119 - for i in range(len(self.layer)):

120 - layer = self.layer[i]

121 - hidden_states = layer(hidden_states, kv_caches[i], attn_metadata)

118 + for layer in self.layer:

119 + hidden_states = layer(hidden_states)

120 return hidden_states

121

122

@@ -152,13 +150,8 @@ class BertLayer(nn.Module):

150 quant_config=quant_config,

151 prefix=f"{prefix}.output")

152

155 - def forward(

156 - self,

157 - hidden_states: torch.Tensor,

158 - kv_cache: Optional[torch.Tensor],

159 - attn_metadata: AttentionMetadata,

160 - ):

161 - attn_output = self.attention(hidden_states, kv_cache, attn_metadata)

153 + def forward(self, hidden_states: torch.Tensor):

154 + attn_output = self.attention(hidden_states)

155 intermediate_output = self.intermediate(attn_output)

156 output = self.output(intermediate_output, attn_output)

157 return output

@@ -191,10 +184,8 @@ class BertAttention(nn.Module):

184 def forward(

185 self,

186 hidden_states: torch.Tensor,

194 - kv_cache: torch.Tensor,

195 - attn_metadata: AttentionMetadata,

187 ) -> torch.Tensor:

197 - self_output = self.self(hidden_states, kv_cache, attn_metadata)

188 + self_output = self.self(hidden_states)

189 return self.output(self_output, hidden_states)

190

191

@@ -246,12 +237,10 @@ class BertSelfAttention(nn.Module):

237 def forward(

238 self,

239 hidden_states: torch.Tensor,

249 - kv_cache: torch.Tensor,

250 - attn_metadata: AttentionMetadata,

240 ) -> torch.Tensor:

241 qkv, _ = self.qkv_proj(hidden_states)

242 q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)

254 - output = self.attn(q, k, v, kv_cache, attn_metadata)

243 + output = self.attn(q, k, v)

244 return output

245

246

@@ -343,8 +332,6 @@ class BertModel(nn.Module):

332 self,

333 input_ids: torch.Tensor,

334 position_ids: torch.Tensor,

346 - kv_caches: List[torch.Tensor],

347 - attn_metadata: AttentionMetadata,

335 intermediate_tensors: Optional[IntermediateTensors] = None,

336 inputs_embeds: Optional[torch.Tensor] = None,

337 token_type_ids: Optional[torch.Tensor] = None,

@@ -352,13 +339,14 @@ class BertModel(nn.Module):

339 if inputs_embeds is not None:

340 hidden_states = inputs_embeds

341 else:

342 + attn_metadata = get_forward_context().attn_metadata

343 assert hasattr(attn_metadata, "seq_lens_tensor")

344 hidden_states = self.embeddings(

345 input_ids=input_ids,

346 seq_lens=attn_metadata.seq_lens_tensor,

347 position_ids=position_ids,

348 token_type_ids=token_type_ids)

361 - return self.encoder(hidden_states, kv_caches, attn_metadata)

349 + return self.encoder(hidden_states)

350

351 def load_weights(self, weights: Iterable[Tuple[str,

352 torch.Tensor]]) -> Set[str]:

@@ -420,17 +408,13 @@ class BertEmbeddingModel(nn.Module):

408 self,

409 input_ids: Optional[torch.Tensor],

410 positions: torch.Tensor,

423 - kv_caches: List[torch.Tensor],

424 - attn_metadata: AttentionMetadata,

411 intermediate_tensors: Optional[IntermediateTensors] = None,

412 inputs_embeds: Optional[torch.Tensor] = None,

413 ) -> torch.Tensor:

414 return self.model(input_ids=input_ids,

415 position_ids=positions,

430 - kv_caches=kv_caches,

416 inputs_embeds=inputs_embeds,

432 - intermediate_tensors=intermediate_tensors,

433 - attn_metadata=attn_metadata)

417 + intermediate_tensors=intermediate_tensors)

418

419 def pooler(

420 self,

@@ -519,16 +503,12 @@ class BertForSequenceClassification(nn.Module, SupportsCrossEncoding):

503 self,

504 input_ids: Optional[torch.Tensor],

505 positions: torch.Tensor,

522 - kv_caches: List[torch.Tensor],

523 - attn_metadata: AttentionMetadata,

506 intermediate_tensors: Optional[IntermediateTensors] = None,

507 inputs_embeds: Optional[torch.Tensor] = None,

508 token_type_ids: Optional[torch.Tensor] = None,

509 ) -> torch.Tensor:

510 return self.bert(input_ids=input_ids,

511 position_ids=positions,

530 - kv_caches=kv_caches,

512 inputs_embeds=inputs_embeds,

513 intermediate_tensors=intermediate_tensors,

533 - attn_metadata=attn_metadata,

514 token_type_ids=token_type_ids)

vllm/model_executor/models/blip2.py

+1 -6

@@ -1,7 +1,7 @@

1 # SPDX-License-Identifier: Apache-2.0

2

3 from functools import cached_property

4 -from typing import (Iterable, List, Literal, Mapping, Optional, Set, Tuple,

4 +from typing import (Iterable, Literal, Mapping, Optional, Set, Tuple,

5 TypedDict, Union)

6

7 import torch

@@ -9,7 +9,6 @@ import torch.nn as nn

9 from transformers import (BatchFeature, Blip2Config, Blip2QFormerConfig,

10 apply_chunking_to_forward)

11

12 -from vllm.attention import AttentionMetadata

12 from vllm.config import CacheConfig, VllmConfig

13 from vllm.model_executor.layers.activation import get_act_fn

14 from vllm.model_executor.layers.quantization import QuantizationConfig

@@ -658,8 +657,6 @@ class Blip2ForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP):

657 self,

658 input_ids: torch.Tensor,

659 positions: torch.Tensor,

661 - kv_caches: List[torch.Tensor],

662 - attn_metadata: AttentionMetadata,

660 intermediate_tensors: Optional[IntermediateTensors] = None,

661 inputs_embeds: Optional[torch.Tensor] = None,

662 **kwargs: object,

@@ -708,8 +705,6 @@ class Blip2ForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP):

705

706 hidden_states = self.language_model.model(input_ids,

707 positions,

711 - kv_caches,

712 - attn_metadata,

708 intermediate_tensors,

709 inputs_embeds=inputs_embeds)

710

vllm/model_executor/models/bloom.py

+7 -24

@@ -18,13 +18,13 @@

18 # limitations under the License.

19 """Inference-only BLOOM model compatible with HuggingFace weights."""

20 import math

21 -from typing import Iterable, List, Optional, Set, Tuple, Union

21 +from typing import Iterable, Optional, Set, Tuple, Union

22

23 import torch

24 from torch import nn

25 from transformers import BloomConfig

26

27 -from vllm.attention import Attention, AttentionMetadata

27 +from vllm.attention import Attention

28 from vllm.compilation.decorators import support_torch_compile

29 from vllm.config import CacheConfig, VllmConfig

30 from vllm.distributed import (get_pp_group, get_tensor_model_parallel_rank,

@@ -126,13 +126,11 @@ class BloomAttention(nn.Module):

126 self,

127 position_ids: torch.Tensor,

128 hidden_states: torch.Tensor,

129 - kv_cache: torch.Tensor,

130 - attn_metadata: AttentionMetadata,

129 ) -> torch.Tensor:

130 del position_ids # Unused.

131 qkv, _ = self.query_key_value(hidden_states)

132 q, k, v = qkv.chunk(chunks=3, dim=-1)

135 - attn_output = self.attn(q, k, v, kv_cache, attn_metadata)

133 + attn_output = self.attn(q, k, v)

134 output, _ = self.dense(attn_output)

135 return output

136

@@ -193,8 +191,6 @@ class BloomBlock(nn.Module):

191 self,

192 position_ids: torch.Tensor,

193 hidden_states: torch.Tensor,

196 - kv_cache: torch.Tensor,

197 - attn_metadata: AttentionMetadata,

194 ) -> torch.Tensor:

195 # Layer norm at the beginning of the transformer layer.

196 layernorm_output = self.input_layernorm(hidden_states)

@@ -209,8 +205,6 @@ class BloomBlock(nn.Module):

205 attention_output = self.self_attention(

206 position_ids=position_ids,

207 hidden_states=layernorm_output,

212 - kv_cache=kv_cache,

213 - attn_metadata=attn_metadata,

208 )

209 attention_output = attention_output + residual

210 layernorm_output = self.post_attention_layernorm(attention_output)

@@ -266,8 +260,6 @@ class BloomModel(nn.Module):

260 self,

261 input_ids: torch.Tensor,

262 position_ids: torch.Tensor,

269 - kv_caches: List[torch.Tensor],

270 - attn_metadata: AttentionMetadata,

263 intermediate_tensors: Optional[IntermediateTensors],

264 inputs_embeds: Optional[torch.Tensor] = None,

265 ) -> Union[torch.Tensor, IntermediateTensors]:

@@ -279,14 +271,8 @@ class BloomModel(nn.Module):

271 else:

272 assert intermediate_tensors is not None

273 hidden_states = intermediate_tensors["hidden_states"]

282 - for i in range(self.start_layer, self.end_layer):

283 - layer = self.h[i]

284 - hidden_states = layer(

285 - position_ids,

286 - hidden_states,

287 - kv_caches[i - self.start_layer],

288 - attn_metadata,

289 - )

274 + for layer in self.h[self.start_layer:self.end_layer]:

275 + hidden_states = layer(position_ids, hidden_states)

276 if not get_pp_group().is_last_rank:

277 return IntermediateTensors({"hidden_states": hidden_states})

278 hidden_states = self.ln_f(hidden_states)

@@ -322,14 +308,11 @@ class BloomForCausalLM(nn.Module, SupportsPP):

308 self,

309 input_ids: torch.Tensor,

310 positions: torch.Tensor,

325 - kv_caches: List[torch.Tensor],

326 - attn_metadata: AttentionMetadata,

311 intermediate_tensors: Optional[IntermediateTensors] = None,

312 inputs_embeds: Optional[torch.Tensor] = None,

313 ) -> Union[torch.Tensor, IntermediateTensors]:

330 - hidden_states = self.transformer(input_ids, positions, kv_caches,

331 - attn_metadata, intermediate_tensors,

332 - inputs_embeds)

314 + hidden_states = self.transformer(input_ids, positions,

315 + intermediate_tensors, inputs_embeds)

316 return hidden_states

317

318 def compute_logits(

vllm/model_executor/models/chameleon.py

+4 -23

@@ -1,7 +1,7 @@

1 # SPDX-License-Identifier: Apache-2.0

2

3 from functools import cached_property

4 -from typing import (Any, Dict, Iterable, List, Literal, Mapping, Optional, Set,

4 +from typing import (Any, Dict, Iterable, Literal, Mapping, Optional, Set,

5 Tuple, TypedDict, Union)

6

7 import torch

@@ -10,7 +10,7 @@ import torch.nn.functional as F

10 from transformers import (BatchFeature, ChameleonConfig, ChameleonProcessor,

11 ChameleonVQVAEConfig)

12

13 -from vllm.attention import Attention, AttentionMetadata

13 +from vllm.attention import Attention

14 from vllm.config import CacheConfig, VllmConfig

15 from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size

16 from vllm.logger import init_logger

@@ -310,15 +310,13 @@ class ChameleonAttention(nn.Module):

310 self,

311 positions: torch.Tensor,

312 hidden_states: torch.Tensor,

313 - kv_cache: torch.Tensor,

314 - attn_metadata: AttentionMetadata,

313 ) -> torch.Tensor:

314 qkv, _ = self.qkv_proj(hidden_states)

315 q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)

316 q, k = self._apply_qk_norm(q, k)

317

318 q, k = self.rotary_emb(positions, q, k)

321 - attn_output = self.attn(q, k, v, kv_cache, attn_metadata)

319 + attn_output = self.attn(q, k, v)

320 output, _ = self.o_proj(attn_output)

321 return output

322

@@ -372,8 +370,6 @@ class ChameleonDecoderLayer(nn.Module):

370 self,

371 positions: torch.Tensor,

372 hidden_states: torch.Tensor,

375 - kv_cache: torch.Tensor,

376 - attn_metadata: AttentionMetadata,

373 residual: Optional[torch.Tensor],

374 ) -> Tuple[torch.Tensor, torch.Tensor]:

375

@@ -386,8 +382,6 @@ class ChameleonDecoderLayer(nn.Module):

382 hidden_states = self.self_attn(

383 positions=positions,

384 hidden_states=hidden_states,

389 - kv_cache=kv_cache,

390 - attn_metadata=attn_metadata,

385 )

386

387 # Fully Connected

@@ -447,8 +441,6 @@ class ChameleonSwinDecoderLayer(nn.Module):

441 self,

442 positions: torch.Tensor,

443 hidden_states: torch.Tensor,

450 - kv_cache: torch.Tensor,

451 - attn_metadata: AttentionMetadata,

444 residual: Optional[torch.Tensor],

445 ) -> Tuple[torch.Tensor, torch.Tensor]:

446

@@ -456,8 +448,6 @@ class ChameleonSwinDecoderLayer(nn.Module):

448 hidden_states = self.self_attn(

449 positions=positions,

450 hidden_states=hidden_states,

459 - kv_cache=kv_cache,

460 - attn_metadata=attn_metadata,

451 )

452

453 hidden_states = self.input_layernorm(hidden_states)

@@ -906,8 +896,6 @@ class ChameleonModel(nn.Module):

896 self,

897 input_ids: Optional[torch.Tensor],

898 positions: torch.Tensor,

909 - kv_caches: List[torch.Tensor],

910 - attn_metadata: AttentionMetadata,

899 intermediate_tensors: Optional[IntermediateTensors],

900 inputs_embeds: Optional[torch.Tensor] = None,

901 ) -> Union[torch.Tensor, IntermediateTensors]:

@@ -921,13 +909,10 @@ class ChameleonModel(nn.Module):

909 assert intermediate_tensors is not None

910 hidden_states = intermediate_tensors["hidden_states"]

911 residual = intermediate_tensors["residual"]

924 - for i in range(self.start_layer, self.end_layer):

925 - layer = self.layers[i]

912 + for layer in self.layers[self.start_layer:self.end_layer]:

913 hidden_states, residual = layer(

914 positions,

915 hidden_states,

929 - kv_caches[i - self.start_layer],

930 - attn_metadata,

916 residual,

917 )

918 if not get_pp_group().is_last_rank:

@@ -1028,8 +1013,6 @@ class ChameleonForConditionalGeneration(nn.Module, SupportsMultiModal,

1013 self,

1014 input_ids: torch.Tensor,

1015 positions: torch.Tensor,

1031 - kv_caches: List[torch.Tensor],

1032 - attn_metadata: AttentionMetadata,

1016 intermediate_tensors: Optional[IntermediateTensors] = None,

1017 inputs_embeds: Optional[torch.Tensor] = None,

1018 **kwargs,

@@ -1048,8 +1031,6 @@ class ChameleonForConditionalGeneration(nn.Module, SupportsMultiModal,

1031

1032 hidden_states = self.model(input_ids,

1033 positions,

1051 - kv_caches,

1052 - attn_metadata,

1034 intermediate_tensors,

1035 inputs_embeds=inputs_embeds)

1036 return hidden_states

vllm/model_executor/models/chatglm.py

+8 -34

@@ -2,13 +2,13 @@

2 # Adapted from

3 # https://github.com/THUDM/ChatGLM2-6B

4 """Inference-only ChatGLM model compatible with THUDM weights."""

5 -from typing import Iterable, List, Optional, Set, Tuple, Union

5 +from typing import Iterable, Optional, Set, Tuple, Union

6

7 import torch

8 from torch import nn

9 from torch.nn import LayerNorm

10

11 -from vllm.attention import Attention, AttentionMetadata

11 +from vllm.attention import Attention

12 from vllm.config import CacheConfig, VllmConfig

13 from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size

14 from vllm.model_executor.layers.activation import SiluAndMul

@@ -108,19 +108,11 @@ class GLMAttention(nn.Module):

108 self,

109 hidden_states: torch.Tensor,

110 position_ids: torch.Tensor,

111 - kv_cache: torch.Tensor,

112 - attn_metadata: AttentionMetadata,

111 ) -> torch.Tensor:

112 qkv, _ = self.query_key_value(hidden_states)

113 q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)

114 q, k = self.rotary_emb(position_ids, q, k)

117 - context_layer = self.attn(

118 - q,

119 - k,

120 - v,

121 - kv_cache,

122 - attn_metadata,

123 - )

115 + context_layer = self.attn(q, k, v)

116 attn_output, _ = self.dense(context_layer)

117 return attn_output

118

@@ -215,8 +207,6 @@ class GLMBlock(nn.Module):

207 self,

208 hidden_states: torch.Tensor,

209 position_ids: torch.Tensor,

218 - kv_cache: torch.Tensor,

219 - attn_metadata: AttentionMetadata,

210 ) -> torch.Tensor:

211 # hidden_states: [num_tokens, h]

212 # Layer norm at the beginning of the transformer layer.

@@ -225,8 +215,6 @@ class GLMBlock(nn.Module):

215 attention_output = self.self_attention(

216 hidden_states=layernorm_output,

217 position_ids=position_ids,

228 - kv_cache=kv_cache,

229 - attn_metadata=attn_metadata,

218 )

219

220 # Residual connection.

@@ -289,17 +277,10 @@ class GLMTransformer(nn.Module):

277 self,

278 hidden_states: torch.Tensor,

279 position_ids: torch.Tensor,

292 - kv_caches: List[torch.Tensor],

293 - attn_metadata: AttentionMetadata,

280 ) -> Union[torch.Tensor, IntermediateTensors]:

295 - for i in range(self.start_layer, self.end_layer):

296 - layer = self.layers[i]

297 - hidden_states = layer(

298 - hidden_states=hidden_states,

299 - position_ids=position_ids,

300 - kv_cache=kv_caches[i - self.start_layer],

301 - attn_metadata=attn_metadata,

302 - )

281 + for layer in self.layers[self.start_layer:self.end_layer]:

282 + hidden_states = layer(hidden_states=hidden_states,

283 + position_ids=position_ids)

284

285 if not get_pp_group().is_last_rank:

286 return IntermediateTensors({"hidden_states": hidden_states})

@@ -350,8 +331,6 @@ class ChatGLMModel(nn.Module):

331 self,

332 input_ids: torch.Tensor,

333 positions: torch.Tensor,

353 - kv_caches: List[torch.Tensor],

354 - attn_metadata: AttentionMetadata,

334 intermediate_tensors: Optional[IntermediateTensors] = None,

335 inputs_embeds: Optional[torch.Tensor] = None,

336 **kwargs: object,

@@ -369,8 +348,6 @@ class ChatGLMModel(nn.Module):

348 hidden_states = self.encoder(

349 hidden_states=hidden_states,

350 position_ids=positions,

372 - kv_caches=kv_caches,

373 - attn_metadata=attn_metadata,

351 )

352

353 return hidden_states

@@ -494,12 +471,9 @@ class ChatGLMForCausalLM(ChatGLMBaseModel, SupportsLoRA, SupportsPP):

471 self,

472 input_ids: torch.Tensor,

473 positions: torch.Tensor,

497 - kv_caches: List[torch.Tensor],

498 - attn_metadata: AttentionMetadata,

474 intermediate_tensors: Optional[IntermediateTensors] = None,

475 inputs_embeds: Optional[torch.Tensor] = None,

476 ) -> Union[torch.Tensor, IntermediateTensors]:

502 - hidden_states = self.transformer(input_ids, positions, kv_caches,

503 - attn_metadata, intermediate_tensors,

504 - inputs_embeds)

477 + hidden_states = self.transformer(input_ids, positions,

478 + intermediate_tensors, inputs_embeds)

479 return hidden_states

vllm/model_executor/models/commandr.py

+5 -19

@@ -21,14 +21,14 @@

21

22 # This file is based on the LLama model definition file in transformers

23 """PyTorch Cohere model."""

24 -from typing import Iterable, List, Optional, Set, Tuple, Union

24 +from typing import Iterable, Optional, Set, Tuple, Union

25

26 import torch

27 import torch.utils.checkpoint

28 from torch import nn

29 from transformers import CohereConfig

30

31 -from vllm.attention import Attention, AttentionMetadata

31 +from vllm.attention import Attention

32 from vllm.compilation.decorators import support_torch_compile

33 from vllm.config import CacheConfig, VllmConfig

34 from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size

@@ -218,8 +218,6 @@ class CohereAttention(nn.Module):

218 self,

219 positions: torch.Tensor,

220 hidden_states: torch.Tensor,

221 - kv_cache: torch.Tensor,

222 - attn_metadata: AttentionMetadata,

221 ) -> torch.Tensor:

222 qkv, _ = self.qkv_proj(hidden_states)

223 q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)

@@ -227,7 +225,7 @@ class CohereAttention(nn.Module):

225 q, k = self._apply_qk_norm(q, k)

226 if self.v1 or self.sliding_window:

227 q, k = self.rotary_emb(positions, q, k)

230 - attn_output = self.attn(q, k, v, kv_cache, attn_metadata)

228 + attn_output = self.attn(q, k, v)

229 output, _ = self.o_proj(attn_output)

230 return output

231

@@ -255,8 +253,6 @@ class CohereDecoderLayer(nn.Module):

253 self,

254 positions: torch.Tensor,

255 hidden_states: torch.Tensor,

258 - kv_cache: torch.Tensor,

259 - attn_metadata: AttentionMetadata,

256 residual: Optional[torch.Tensor],

257 ) -> Tuple[torch.Tensor, torch.Tensor]:

258 # Self Attention

@@ -265,8 +261,6 @@ class CohereDecoderLayer(nn.Module):

261 hidden_states_attention = self.self_attn(

262 positions=positions,

263 hidden_states=hidden_states,

268 - kv_cache=kv_cache,

269 - attn_metadata=attn_metadata,

264 )

265 hidden_states_mlp = self.mlp(hidden_states)

266 # Add everything together

@@ -311,8 +305,6 @@ class CohereModel(nn.Module):

305 self,

306 input_ids: torch.Tensor,

307 positions: torch.Tensor,

314 - kv_caches: List[torch.Tensor],

315 - attn_metadata: AttentionMetadata,

308 intermediate_tensors: Optional[IntermediateTensors],

309 inputs_embeds: Optional[torch.Tensor] = None,

310 ) -> Union[torch.Tensor, IntermediateTensors]:

@@ -326,13 +318,10 @@ class CohereModel(nn.Module):

318 assert intermediate_tensors is not None

319 hidden_states = intermediate_tensors["hidden_states"]

320 residual = intermediate_tensors["residual"]

329 - for i in range(self.start_layer, self.end_layer):

330 - layer = self.layers[i]

321 + for layer in self.layers[self.start_layer:self.end_layer]:

322 hidden_states, residual = layer(

323 positions,

324 hidden_states,

334 - kv_caches[i - self.start_layer],

335 - attn_metadata,

325 residual,

326 )

327 if not get_pp_group().is_last_rank:

@@ -389,13 +378,10 @@ class CohereForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

378 self,

379 input_ids: torch.Tensor,

380 positions: torch.Tensor,

392 - kv_caches: List[torch.Tensor],

393 - attn_metadata: AttentionMetadata,

381 intermediate_tensors: Optional[IntermediateTensors] = None,

382 inputs_embeds: Optional[torch.Tensor] = None,

383 ) -> Union[torch.Tensor, IntermediateTensors]:

397 - hidden_states = self.model(input_ids, positions, kv_caches,

398 - attn_metadata, intermediate_tensors,

384 + hidden_states = self.model(input_ids, positions, intermediate_tensors,

385 inputs_embeds)

386 return hidden_states

387

vllm/model_executor/models/dbrx.py

+7 -28

@@ -1,11 +1,11 @@

1 # SPDX-License-Identifier: Apache-2.0

2

3 -from typing import Iterable, List, Optional, Set, Tuple, Union

3 +from typing import Iterable, Optional, Set, Tuple, Union

4

5 import torch

6 import torch.nn as nn

7

8 -from vllm.attention import Attention, AttentionMetadata

8 +from vllm.attention import Attention

9 from vllm.config import CacheConfig, VllmConfig

10 from vllm.distributed import (get_pp_group, get_tensor_model_parallel_rank,

11 get_tensor_model_parallel_world_size)

@@ -230,15 +230,13 @@ class DbrxAttention(nn.Module):

230 self,

231 position_ids: torch.Tensor,

232 hidden_states: torch.Tensor,

233 - kv_cache: torch.Tensor,

234 - attn_metadata: AttentionMetadata,

233 ) -> torch.Tensor:

234 qkv, _ = self.Wqkv(hidden_states)

235 if self.clip_qkv is not None:

236 qkv.clamp_(min=-self.clip_qkv, max=self.clip_qkv)

237 q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)

238 q, k = self.rotary_emb(position_ids, q, k)

241 - attn_output = self.attn(q, k, v, kv_cache, attn_metadata)

239 + attn_output = self.attn(q, k, v)

240 hidden_states, _ = self.out_proj(attn_output)

241 return hidden_states

242

@@ -265,16 +263,12 @@ class DbrxFusedNormAttention(nn.Module):

263 self,

264 position_ids: torch.Tensor,

265 hidden_states: torch.Tensor,

268 - kv_cache: torch.Tensor,

269 - attn_metadata: AttentionMetadata,

266 ) -> torch.Tensor:

267 residual = hidden_states

268 hidden_states = self.norm_1(hidden_states)

269 x = self.attn(

270 position_ids=position_ids,

271 hidden_states=hidden_states,

276 - kv_cache=kv_cache,

277 - attn_metadata=attn_metadata,

272 )

273 hidden_states = residual + x

274 residual = hidden_states

@@ -303,14 +297,10 @@ class DbrxBlock(nn.Module):

297 self,

298 position_ids: torch.Tensor,

299 hidden_states: torch.Tensor,

306 - kv_cache: torch.Tensor,

307 - attn_metadata: AttentionMetadata,

300 ) -> torch.Tensor:

301 hidden_states, residual = self.norm_attn_norm(

302 position_ids=position_ids,

303 hidden_states=hidden_states,

312 - kv_cache=kv_cache,

313 - attn_metadata=attn_metadata,

304 )

305 hidden_states = self.ffn(hidden_states)

306 hidden_states = hidden_states + residual

@@ -353,8 +343,6 @@ class DbrxModel(nn.Module):

343 self,

344 input_ids: torch.Tensor,

345 position_ids: torch.Tensor,

356 - kv_caches: List[torch.Tensor],

357 - attn_metadata: AttentionMetadata,

346 intermediate_tensors: Optional[IntermediateTensors],

347 inputs_embeds: Optional[torch.Tensor] = None,

348 ) -> Union[torch.Tensor, IntermediateTensors]:

@@ -366,14 +354,8 @@ class DbrxModel(nn.Module):

354 else:

355 assert intermediate_tensors

356 hidden_states = intermediate_tensors["hidden_states"]

369 - for i in range(self.start_layer, self.end_layer):

370 - block = self.blocks[i]

371 - hidden_states = block(

372 - position_ids,

373 - hidden_states,

374 - kv_caches[i - self.start_layer],

375 - attn_metadata,

376 - )

357 + for block in self.blocks[self.start_layer:self.end_layer]:

358 + hidden_states = block(position_ids, hidden_states)

359 if not get_pp_group().is_last_rank:

360 return IntermediateTensors({"hidden_states": hidden_states})

361 hidden_states = self.norm_f(hidden_states)

@@ -415,14 +397,11 @@ class DbrxForCausalLM(nn.Module, SupportsPP):

397 self,

398 input_ids: torch.Tensor,

399 positions: torch.Tensor,

418 - kv_caches: List[torch.Tensor],

419 - attn_metadata: AttentionMetadata,

400 intermediate_tensors: Optional[IntermediateTensors] = None,

401 inputs_embeds: Optional[torch.Tensor] = None,

402 ) -> Union[torch.Tensor, IntermediateTensors]:

423 - hidden_states = self.transformer(input_ids, positions, kv_caches,

424 - attn_metadata, intermediate_tensors,

425 - inputs_embeds)

403 + hidden_states = self.transformer(input_ids, positions,

404 + intermediate_tensors, inputs_embeds)

405 return hidden_states

406

407 def compute_logits(

vllm/model_executor/models/deepseek.py

+6 -20

@@ -22,13 +22,13 @@

22 # See the License for the specific language governing permissions and

23 # limitations under the License.

24 """Inference-only Deepseek model."""

25 -from typing import Any, Dict, Iterable, List, Optional, Set, Tuple, Union

25 +from typing import Any, Dict, Iterable, Optional, Set, Tuple, Union

26

27 import torch

28 from torch import nn

29 from transformers import PretrainedConfig

30

31 -from vllm.attention import Attention, AttentionMetadata

31 +from vllm.attention import Attention

32 from vllm.config import CacheConfig, VllmConfig

33 from vllm.distributed import (get_pp_group, get_tensor_model_parallel_rank,

34 get_tensor_model_parallel_world_size,

@@ -248,13 +248,11 @@ class DeepseekAttention(nn.Module):

248 self,

249 positions: torch.Tensor,

250 hidden_states: torch.Tensor,

251 - kv_cache: torch.Tensor,

252 - attn_metadata: AttentionMetadata,

251 ) -> torch.Tensor:

252 qkv, _ = self.qkv_proj(hidden_states)

253 q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)

254 q, k = self.rotary_emb(positions, q, k)

257 - attn_output = self.attn(q, k, v, kv_cache, attn_metadata)

255 + attn_output = self.attn(q, k, v)

256 output, _ = self.o_proj(attn_output)

257 return output

258

@@ -309,8 +307,6 @@ class DeepseekDecoderLayer(nn.Module):

307 self,

308 positions: torch.Tensor,

309 hidden_states: torch.Tensor,

312 - kv_cache: torch.Tensor,

313 - attn_metadata: AttentionMetadata,

310 residual: Optional[torch.Tensor],

311 ) -> torch.Tensor:

312 # Self Attention

@@ -323,8 +319,6 @@ class DeepseekDecoderLayer(nn.Module):

319 hidden_states = self.self_attn(

320 positions=positions,

321 hidden_states=hidden_states,

326 - kv_cache=kv_cache,

327 - attn_metadata=attn_metadata,

322 )

323

324 # Fully Connected

@@ -370,8 +364,6 @@ class DeepseekModel(nn.Module):

364 self,

365 input_ids: torch.Tensor,

366 positions: torch.Tensor,

373 - kv_caches: List[torch.Tensor],

374 - attn_metadata: AttentionMetadata,

367 intermediate_tensors: Optional[IntermediateTensors],

368 inputs_embeds: Optional[torch.Tensor] = None,

369 ) -> Union[torch.Tensor, IntermediateTensors]:

@@ -384,11 +376,8 @@ class DeepseekModel(nn.Module):

376 else:

377 hidden_states = intermediate_tensors["hidden_states"]

378 residual = intermediate_tensors["residual"]

387 - for i in range(self.start_layer, self.end_layer):

388 - layer = self.layers[i]

389 - hidden_states, residual = layer(positions, hidden_states,

390 - kv_caches[i - self.start_layer],

391 - attn_metadata, residual)

379 + for layer in self.layers[self.start_layer:self.end_layer]:

380 + hidden_states, residual = layer(positions, hidden_states, residual)

381 if not get_pp_group().is_last_rank:

382 return IntermediateTensors({

383 "hidden_states": hidden_states,

@@ -425,13 +414,10 @@ class DeepseekForCausalLM(nn.Module, SupportsPP):

414 self,

415 input_ids: torch.Tensor,

416 positions: torch.Tensor,

428 - kv_caches: List[torch.Tensor],

429 - attn_metadata: AttentionMetadata,

417 intermediate_tensors: Optional[IntermediateTensors] = None,

418 inputs_embeds: Optional[torch.Tensor] = None,

419 ) -> Union[torch.Tensor, IntermediateTensors]:

433 - hidden_states = self.model(input_ids, positions, kv_caches,

434 - attn_metadata, intermediate_tensors,

420 + hidden_states = self.model(input_ids, positions, intermediate_tensors,

421 inputs_embeds)

422 return hidden_states

423

vllm/model_executor/models/deepseek_mtp.py

+4 -15

@@ -1,11 +1,10 @@

1 # SPDX-License-Identifier: Apache-2.0

2 -from typing import Iterable, List, Optional, Set, Tuple

2 +from typing import Iterable, Optional, Set, Tuple

3

4 import torch

5 import torch.nn as nn

6 from transformers import PretrainedConfig

7

8 -from vllm.attention.backends.abstract import AttentionMetadata

8 from vllm.config import CacheConfig, ModelConfig, VllmConfig

9 from vllm.model_executor.layers.fused_moe import FusedMoE

10 from vllm.model_executor.layers.layernorm import RMSNorm

@@ -69,8 +68,6 @@ class DeepSeekMultiTokenPredictorLayer(nn.Module):

68 self,

69 input_ids: torch.Tensor,

70 positions: torch.Tensor,

72 - kv_cache: torch.Tensor,

73 - attn_metadata: AttentionMetadata,

71 previous_hidden_states: torch.Tensor,

72 inputs_embeds: Optional[torch.Tensor] = None,

73 spec_step_index: int = 0,

@@ -88,8 +85,6 @@ class DeepSeekMultiTokenPredictorLayer(nn.Module):

85

86 hidden_states, residual = self.mtp_block(positions=positions,

87 hidden_states=hidden_states,

91 - kv_cache=kv_cache,

92 - attn_metadata=attn_metadata,

88 residual=None)

89 hidden_states = residual + hidden_states

90 return self.shared_head(hidden_states)

@@ -122,8 +117,6 @@ class DeepSeekMultiTokenPredictor(nn.Module):

117 self,

118 input_ids: torch.Tensor,

119 positions: torch.Tensor,

125 - kv_caches: List[torch.Tensor],

126 - attn_metadata: AttentionMetadata,

120 previous_hidden_states: torch.Tensor,

121 inputs_embeds: Optional[torch.Tensor] = None,

122 spec_step_idx: int = 0,

@@ -131,8 +124,6 @@ class DeepSeekMultiTokenPredictor(nn.Module):

124 return self.layers[str(self.mtp_start_layer_idx + spec_step_idx)](

125 input_ids,

126 positions,

134 - kv_caches[spec_step_idx],

135 - attn_metadata,

127 previous_hidden_states,

128 inputs_embeds,

129 spec_step_idx,

@@ -165,16 +156,14 @@ class DeepSeekMTP(nn.Module):

156 self,

157 input_ids: torch.Tensor,

158 positions: torch.Tensor,

168 - kv_caches: List[torch.Tensor],

169 - attn_metadata: AttentionMetadata,

159 previous_hidden_states: torch.Tensor,

160 intermediate_tensors: Optional[IntermediateTensors] = None,

161 inputs_embeds: Optional[torch.Tensor] = None,

162 spec_step_idx: int = 0,

163 ) -> torch.Tensor:

175 - hidden_states = self.model(input_ids, positions, kv_caches,

176 - attn_metadata, previous_hidden_states,

177 - inputs_embeds, spec_step_idx)

164 + hidden_states = self.model(input_ids, positions,

165 + previous_hidden_states, inputs_embeds,

166 + spec_step_idx)

167 return hidden_states

168

169 def compute_logits(

vllm/model_executor/models/deepseek_v2.py

+7 -24

@@ -22,13 +22,13 @@

22 # See the License for the specific language governing permissions and

23 # limitations under the License.

24 """Inference-only DeepseekV2/DeepseekV3 model."""

25 -from typing import Any, Dict, Iterable, List, Optional, Set, Tuple, Union

25 +from typing import Any, Dict, Iterable, Optional, Set, Tuple, Union

26

27 import torch

28 from torch import nn

29 from transformers import PretrainedConfig

30

31 -from vllm.attention import Attention, AttentionMetadata

31 +from vllm.attention import Attention

32 from vllm.compilation.decorators import support_torch_compile

33 from vllm.config import CacheConfig, ModelConfig, VllmConfig

34 from vllm.distributed import (get_pp_group,

@@ -279,8 +279,6 @@ class DeepseekV2Attention(nn.Module):

279 self,

280 positions: torch.Tensor,

281 hidden_states: torch.Tensor,

282 - kv_cache: torch.Tensor,

283 - attn_metadata: AttentionMetadata,

282 ) -> torch.Tensor:

283 if self.q_lora_rank is not None:

284 q = self.q_a_proj(hidden_states)[0]

@@ -313,7 +311,7 @@ class DeepseekV2Attention(nn.Module):

311 v = torch.nn.functional.pad(

312 v, [0, self.qk_head_dim - self.v_head_dim],

313 value=0).view(-1, self.num_local_heads * self.qk_head_dim)

316 - attn_output = self.attn(q, k, v, kv_cache, attn_metadata)

314 + attn_output = self.attn(q, k, v)

315 attn_output = attn_output.view(

316 -1, self.num_local_heads,

317 self.qk_head_dim)[..., :self.v_head_dim].reshape(

@@ -451,8 +449,6 @@ class DeepseekV2MLAAttention(nn.Module):

449 self,

450 positions: torch.Tensor,

451 hidden_states: torch.Tensor,

454 - kv_cache: torch.Tensor,

455 - attn_metadata: AttentionMetadata,

452 ) -> torch.Tensor:

453 if self.q_lora_rank is not None:

454 ckq = self.q_a_proj(hidden_states)[0]

@@ -462,8 +458,7 @@ class DeepseekV2MLAAttention(nn.Module):

458 kv_c, k_pe = self.kv_a_proj_with_mqa(hidden_states)[0].split(

459 [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)

460 kv_c_normed = self.kv_a_layernorm(kv_c.contiguous())

465 - return self.mla_attn(hidden_states_or_q_c, kv_c_normed, k_pe, kv_cache,

466 - attn_metadata)

461 + return self.mla_attn(hidden_states_or_q_c, kv_c_normed, k_pe)

462

463

464 class DeepseekV2DecoderLayer(nn.Module):

@@ -532,8 +527,6 @@ class DeepseekV2DecoderLayer(nn.Module):

527 self,

528 positions: torch.Tensor,

529 hidden_states: torch.Tensor,

535 - kv_cache: torch.Tensor,

536 - attn_metadata: AttentionMetadata,

530 residual: Optional[torch.Tensor],

531 ) -> torch.Tensor:

532 # Self Attention

@@ -546,8 +539,6 @@ class DeepseekV2DecoderLayer(nn.Module):

539 hidden_states = self.self_attn(

540 positions=positions,

541 hidden_states=hidden_states,

549 - kv_cache=kv_cache,

550 - attn_metadata=attn_metadata,

542 )

543

544 # Fully Connected

@@ -608,8 +599,6 @@ class DeepseekV2Model(nn.Module):

599 self,

600 input_ids: torch.Tensor,

601 positions: torch.Tensor,

611 - kv_caches: List[torch.Tensor],

612 - attn_metadata: AttentionMetadata,

602 intermediate_tensors: Optional[IntermediateTensors],

603 inputs_embeds: Optional[torch.Tensor] = None,

604 ) -> Union[torch.Tensor, IntermediateTensors]:

@@ -624,11 +613,8 @@ class DeepseekV2Model(nn.Module):

613 hidden_states = intermediate_tensors["hidden_states"]

614 residual = intermediate_tensors["residual"]

615

627 - for i in range(self.start_layer, self.end_layer):

628 - layer = self.layers[i]

629 - hidden_states, residual = layer(positions, hidden_states,

630 - kv_caches[i - self.start_layer],

631 - attn_metadata, residual)

616 + for layer in self.layers[self.start_layer:self.end_layer]:

617 + hidden_states, residual = layer(positions, hidden_states)

618

619 if not get_pp_group().is_last_rank:

620 return IntermediateTensors({

@@ -665,13 +651,10 @@ class DeepseekV2ForCausalLM(nn.Module, SupportsPP):

651 self,

652 input_ids: torch.Tensor,

653 positions: torch.Tensor,

668 - kv_caches: List[torch.Tensor],

669 - attn_metadata: AttentionMetadata,

654 intermediate_tensors: Optional[IntermediateTensors] = None,

655 inputs_embeds: Optional[torch.Tensor] = None,

656 ) -> Union[torch.Tensor, IntermediateTensors]:

673 - hidden_states = self.model(input_ids, positions, kv_caches,

674 - attn_metadata, intermediate_tensors,

657 + hidden_states = self.model(input_ids, positions, intermediate_tensors,

658 inputs_embeds)

659 return hidden_states

660

vllm/model_executor/models/deepseek_vl2.py

-5

@@ -13,7 +13,6 @@ import torch.nn.functional as F

13 from einops import rearrange, repeat

14 from transformers import BatchFeature

15

16 -from vllm.attention import AttentionMetadata

16 from vllm.config import VllmConfig

17 from vllm.logger import init_logger

18 from vllm.model_executor import SamplingMetadata

@@ -595,8 +594,6 @@ class DeepseekVLV2ForCausalLM(nn.Module, SupportsMultiModal, SupportsPP):

594 def forward(self,

595 input_ids: torch.Tensor,

596 positions: torch.Tensor,

598 - kv_caches: List[torch.Tensor],

599 - attn_metadata: AttentionMetadata,

597 intermediate_tensors: Optional[IntermediateTensors] = None,

598 inputs_embeds: Optional[torch.Tensor] = None,

599 **kwargs: object):

@@ -614,8 +611,6 @@ class DeepseekVLV2ForCausalLM(nn.Module, SupportsMultiModal, SupportsPP):

611

612 hidden_states = self.language_model(input_ids,

613 positions,

617 - kv_caches,

618 - attn_metadata,

614 intermediate_tensors,

615 inputs_embeds=inputs_embeds)

616

vllm/model_executor/models/eagle.py

+1 -6

@@ -1,11 +1,10 @@

1 # SPDX-License-Identifier: Apache-2.0

2

3 -from typing import Iterable, List, Optional, Tuple

3 +from typing import Iterable, Optional, Tuple

4

5 import torch

6 import torch.nn as nn

7

8 -from vllm.attention.backends.abstract import AttentionMetadata

8 from vllm.config import VllmConfig

9 from vllm.logger import init_logger

10 from vllm.model_executor.layers.logits_processor import LogitsProcessor

@@ -121,8 +120,6 @@ class EAGLE(nn.Module):

120 self,

121 input_ids: torch.Tensor,

122 positions: torch.Tensor,

124 - kv_caches: List[torch.Tensor],

125 - attn_metadata: AttentionMetadata,

123 previous_hidden_states: torch.Tensor,

124 intermediate_tensors: Optional[IntermediateTensors] = None,

125 inputs_embeds: Optional[torch.Tensor] = None,

@@ -140,8 +137,6 @@ class EAGLE(nn.Module):

137 input_ids=None,

138 inputs_embeds=inputs_embeds,

139 positions=positions,

143 - kv_caches=kv_caches,

144 - attn_metadata=attn_metadata,

140 intermediate_tensors=intermediate_tensors,

141 )

142 return hidden_states

vllm/model_executor/models/exaone.py

+6 -24

@@ -24,12 +24,12 @@

24 # limitations under the License.

25 """Inference-only Exaone model compatible with HuggingFace weights."""

26

27 -from typing import Any, Dict, Iterable, List, Optional, Set, Tuple, Union

27 +from typing import Any, Dict, Iterable, Optional, Set, Tuple, Union

28

29 import torch

30 from torch import nn

31

32 -from vllm.attention import Attention, AttentionMetadata

32 +from vllm.attention import Attention

33 from vllm.compilation.decorators import support_torch_compile

34 from vllm.config import CacheConfig, VllmConfig

35 from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size

@@ -179,13 +179,11 @@ class ExaoneAttention(nn.Module):

179 self,

180 positions: torch.Tensor,

181 hidden_states: torch.Tensor,

182 - kv_cache: torch.Tensor,

183 - attn_metadata: AttentionMetadata,

182 ) -> torch.Tensor:

183 qkv, _ = self.qkv_proj(hidden_states)

184 q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)

185 q, k = self.rotary_emb(positions, q, k)

188 - attn_output = self.attn(q, k, v, kv_cache, attn_metadata)

186 + attn_output = self.attn(q, k, v)

187 output, _ = self.out_proj(attn_output)

188 return output

189

@@ -225,14 +223,10 @@ class ExaoneBlockAttention(nn.Module):

223 self,

224 positions: torch.Tensor,

225 hidden_states: torch.Tensor,

228 - kv_cache: torch.Tensor,

229 - attn_metadata: AttentionMetadata,

226 ) -> torch.Tensor:

227 return self.attention(

228 positions=positions,

229 hidden_states=hidden_states,

234 - kv_cache=kv_cache,

235 - attn_metadata=attn_metadata,

230 )

231

232

@@ -288,8 +282,6 @@ class ExaoneDecoderLayer(nn.Module):

282 self,

283 positions: torch.Tensor,

284 hidden_states: torch.Tensor,

291 - kv_cache: torch.Tensor,

292 - attn_metadata: AttentionMetadata,

285 residual: Optional[torch.Tensor],

286 ) -> Tuple[torch.Tensor, torch.Tensor]:

287 # Self Attention

@@ -301,8 +293,6 @@ class ExaoneDecoderLayer(nn.Module):

293 hidden_states = self.attn(

294 positions=positions,

295 hidden_states=hidden_states,

304 - kv_cache=kv_cache,

305 - attn_metadata=attn_metadata,

296 )

297

298 # Fully Connected

@@ -365,8 +355,6 @@ class ExaoneModel(nn.Module):

355 self,

356 input_ids: Optional[torch.Tensor],

357 positions: torch.Tensor,

368 - kv_caches: List[torch.Tensor],

369 - attn_metadata: AttentionMetadata,

358 intermediate_tensors: Optional[IntermediateTensors],

359 inputs_embeds: Optional[torch.Tensor] = None,

360 ) -> Union[torch.Tensor, IntermediateTensors]:

@@ -381,13 +369,10 @@ class ExaoneModel(nn.Module):

369 hidden_states = intermediate_tensors["hidden_states"]

370 residual = intermediate_tensors["residual"]

371

384 - for i in range(self.start_layer, self.end_layer):

385 - layer = self.h[i]

372 + for layer in self.h[self.start_layer:self.end_layer]:

373 hidden_states, residual = layer(

374 positions,

375 hidden_states,

389 - kv_caches[i - self.start_layer],

390 - attn_metadata,

376 residual,

377 )

378

@@ -471,14 +456,11 @@ class ExaoneForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

456 self,

457 input_ids: torch.Tensor,

458 positions: torch.Tensor,

474 - kv_caches: List[torch.Tensor],

475 - attn_metadata: AttentionMetadata,

459 intermediate_tensors: Optional[IntermediateTensors] = None,

460 inputs_embeds: Optional[torch.Tensor] = None,

461 ) -> Union[torch.Tensor, IntermediateTensors]:

479 - model_output = self.transformer(input_ids, positions, kv_caches,

480 - attn_metadata, intermediate_tensors,

481 - inputs_embeds)

462 + model_output = self.transformer(input_ids, positions,

463 + intermediate_tensors, inputs_embeds)

464 return model_output

465

466 def compute_logits(

vllm/model_executor/models/falcon.py

+7 -24

@@ -20,14 +20,14 @@

20 """PyTorch Falcon model."""

21

22 import math

23 -from typing import Iterable, List, Optional, Set, Tuple, Union

23 +from typing import Iterable, Optional, Set, Tuple, Union

24

25 import torch

26 from torch import nn

27 from torch.nn import LayerNorm

28 from transformers import FalconConfig as HF_FalconConfig

29

30 -from vllm.attention import Attention, AttentionMetadata

30 +from vllm.attention import Attention

31 from vllm.compilation.decorators import support_torch_compile

32 from vllm.config import CacheConfig, VllmConfig

33 from vllm.distributed import (get_pp_group, get_tensor_model_parallel_rank,

@@ -190,8 +190,6 @@ class FalconAttention(nn.Module):

190 self,

191 positions: torch.Tensor,

192 hidden_states: torch.Tensor,

193 - kv_cache: torch.Tensor,

194 - attn_metadata: AttentionMetadata,

193 ) -> torch.Tensor:

194 qkv, bias = self.query_key_value(hidden_states)

195 if bias is not None:

@@ -199,7 +197,7 @@ class FalconAttention(nn.Module):

197 q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)

198 if self.use_rotary:

199 q, k = self.rotary_emb(positions, q, k)

202 - attn_output = self.attn(q, k, v, kv_cache, attn_metadata)

200 + attn_output = self.attn(q, k, v)

201 attn_output, bias = self.dense(attn_output)

202 return attn_output, bias

203

@@ -291,8 +289,6 @@ class FalconDecoderLayer(nn.Module):

289 self,

290 positions: torch.Tensor,

291 hidden_states: torch.Tensor,

294 - kv_cache: torch.Tensor,

295 - attn_metadata: AttentionMetadata,

292 ) -> torch.Tensor:

293 residual = hidden_states

294

@@ -306,8 +302,6 @@ class FalconDecoderLayer(nn.Module):

302 attention_output, attention_bias = self.self_attention(

303 positions=positions,

304 hidden_states=attention_layernorm_out,

309 - kv_cache=kv_cache,

310 - attn_metadata=attn_metadata,

305 )

306 if self.reduce_row_parallel_results and attention_bias is not None:

307 attention_output += attention_bias

@@ -384,8 +378,6 @@ class FalconModel(nn.Module):

378 self,

379 input_ids: torch.Tensor,

380 positions: torch.Tensor,

387 - kv_caches: List[torch.Tensor],

388 - attn_metadata: AttentionMetadata,

381 intermediate_tensors: Optional[IntermediateTensors],

382 inputs_embeds: Optional[torch.Tensor] = None,

383 ) -> Union[torch.Tensor, IntermediateTensors]:

@@ -396,14 +388,8 @@ class FalconModel(nn.Module):

388 hidden_states = self.get_input_embeddings(input_ids)

389 else:

390 hidden_states = intermediate_tensors["hidden_states"]

399 - for i in range(self.start_layer, self.end_layer):

400 - layer = self.h[i]

401 - hidden_states = layer(

402 - positions,

403 - hidden_states,

404 - kv_caches[i - self.start_layer],

405 - attn_metadata,

406 - )

391 + for layer in self.h[self.start_layer:self.end_layer]:

392 + hidden_states = layer(positions, hidden_states)

393 if not get_pp_group().is_last_rank:

394 return IntermediateTensors({"hidden_states": hidden_states})

395 hidden_states = self.ln_f(hidden_states)

@@ -450,14 +436,11 @@ class FalconForCausalLM(nn.Module, SupportsPP):

436 self,

437 input_ids: torch.LongTensor,

438 positions: torch.Tensor,

453 - kv_caches: List[torch.Tensor],

454 - attn_metadata: AttentionMetadata,

439 intermediate_tensors: Optional[IntermediateTensors] = None,

440 inputs_embeds: Optional[torch.Tensor] = None,

441 ) -> torch.Tensor:

458 - hidden_states = self.transformer(input_ids, positions, kv_caches,

459 - attn_metadata, intermediate_tensors,

460 - inputs_embeds)

442 + hidden_states = self.transformer(input_ids, positions,

443 + intermediate_tensors, inputs_embeds)

444 return hidden_states

445

446 def compute_logits(

vllm/model_executor/models/florence2.py

+6 -28

@@ -1,12 +1,11 @@

1 # SPDX-License-Identifier: Apache-2.0

2

3 import math

4 -from typing import Iterable, List, Optional, Set, Tuple

4 +from typing import Iterable, Optional, Set, Tuple

5

6 import torch

7 import torch.nn as nn

8

9 -from vllm.attention import AttentionMetadata

9 from vllm.config import VllmConfig

10 from vllm.model_executor.layers.logits_processor import LogitsProcessor

11 from vllm.model_executor.layers.sampler import SamplerOutput, get_sampler

@@ -50,8 +49,7 @@ class Florence2LanguageModel(nn.Module):

49

50 def forward(self, input_ids: torch.Tensor, positions: torch.Tensor,

51 encoder_input_ids: torch.Tensor,

53 - encoder_positions: torch.Tensor, kv_caches: List[torch.Tensor],

54 - attn_metadata: AttentionMetadata) -> torch.Tensor:

52 + encoder_positions: torch.Tensor) -> torch.Tensor:

53 r"""

54 Args:

55 input_ids

@@ -64,10 +62,6 @@ class Florence2LanguageModel(nn.Module):

62 Indices of *encoder* input sequence tokens in the vocabulary.

63 encoder_positions:

64 Positions of *encoder* input sequence tokens.

67 - kv_caches:

68 - Layer-wise list of KV cache tensors

69 - attn_metadata:

70 - vLLM Attention metadata structure

65 Returns:

66 Model output torch.Tensor

67 """

@@ -78,18 +72,14 @@ class Florence2LanguageModel(nn.Module):

72 # Run encoder attention if a non-zero number of encoder tokens

73 # are provided as input

74 encoder_hidden_states = self.encoder(input_ids=encoder_input_ids,

81 - positions=encoder_positions,

82 - kv_caches=kv_caches,

83 - attn_metadata=attn_metadata)

75 + positions=encoder_positions)

76

77 # decoder outputs consists of

78 # (dec_features, past_key_value, dec_hidden, dec_attn)

79 decoder_outputs = self.decoder(

80 decoder_input_ids=input_ids,

81 decoder_positions=positions,

90 - encoder_hidden_states=encoder_hidden_states,

91 - kv_caches=kv_caches,

92 - attn_metadata=attn_metadata)

82 + encoder_hidden_states=encoder_hidden_states)

83

84 return decoder_outputs

85

@@ -122,8 +112,6 @@ class Florence2LanguageForConditionalGeneration(nn.Module):

112 positions: torch.Tensor,

113 encoder_input_ids: torch.Tensor,

114 encoder_positions: torch.Tensor,

125 - kv_caches: List[torch.Tensor],

126 - attn_metadata: AttentionMetadata,

115 **kwargs,

116 ) -> torch.Tensor:

117 r"""

@@ -136,15 +124,11 @@ class Florence2LanguageForConditionalGeneration(nn.Module):

124 torch.Tensor of *encoder* input token ids.

125 encoder_positions

126 torch.Tensor of *encoder* position indices

139 - kv_caches:

140 - Layer-wise list of KV cache tensors

141 - attn_metadata:

142 - vLLM Attention metadata structure

127 Returns:

128 Output torch.Tensor

129 """

130 return self.model(input_ids, positions, encoder_input_ids,

147 - encoder_positions, kv_caches, attn_metadata)

131 + encoder_positions)

132

133 def compute_logits(

134 self,

@@ -213,8 +197,6 @@ class Florence2ForConditionalGeneration(nn.Module):

197 self,

198 input_ids: torch.Tensor,

199 positions: torch.Tensor,

216 - kv_caches: List[torch.Tensor],

217 - attn_metadata: AttentionMetadata,

200 intermediate_tensors: Optional[IntermediateTensors] = None,

201 *,

202 encoder_input_ids: torch.Tensor,

@@ -231,15 +213,11 @@ class Florence2ForConditionalGeneration(nn.Module):

213 torch.Tensor of *encoder* input token ids.

214 encoder_positions

215 torch.Tensor of *encoder* position indices

234 - kv_caches:

235 - Layer-wise list of KV cache tensors

236 - attn_metadata:

237 - vLLM Attention metadata structure

216 Returns:

217 Output torch.Tensor

218 """

219 return self.language_model(input_ids, positions, encoder_input_ids,

242 - encoder_positions, kv_caches, attn_metadata)

220 + encoder_positions)

221

222 def compute_logits(

223 self,

vllm/model_executor/models/fuyu.py

-5

@@ -25,7 +25,6 @@ import torch.nn as nn

25 from transformers import (BatchFeature, FuyuConfig, FuyuImageProcessor,

26 FuyuProcessor)

27

28 -from vllm.attention import AttentionMetadata

28 from vllm.config import VllmConfig

29 from vllm.model_executor.layers.linear import ColumnParallelLinear

30 from vllm.model_executor.layers.sampler import SamplerOutput

@@ -351,8 +350,6 @@ class FuyuForCausalLM(nn.Module, SupportsMultiModal, SupportsPP):

350 self,

351 input_ids: torch.Tensor,

352 positions: torch.Tensor,

354 - kv_caches: List[torch.Tensor],

355 - attn_metadata: AttentionMetadata,

353 intermediate_tensors: Optional[IntermediateTensors] = None,

354 inputs_embeds: Optional[torch.Tensor] = None,

355 **kwargs: object,

@@ -371,8 +368,6 @@ class FuyuForCausalLM(nn.Module, SupportsMultiModal, SupportsPP):

368 hidden_states = self.language_model(

369 input_ids=input_ids,

370 positions=positions,

374 - kv_caches=kv_caches,

375 - attn_metadata=attn_metadata,

371 intermediate_tensors=intermediate_tensors,

372 inputs_embeds=inputs_embeds,

373 )

vllm/model_executor/models/gemma.py

+5 -19

@@ -16,13 +16,13 @@

16 # limitations under the License.

17 """Inference-only Gemma model compatible with HuggingFace weights."""

18 from functools import cache

19 -from typing import Iterable, List, Optional, Set, Tuple, Union

19 +from typing import Iterable, Optional, Set, Tuple, Union

20

21 import torch

22 from torch import nn

23 from transformers import GemmaConfig

24

25 -from vllm.attention import Attention, AttentionMetadata

25 +from vllm.attention import Attention

26 from vllm.compilation.decorators import support_torch_compile

27 from vllm.config import CacheConfig, VllmConfig

28 from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size

@@ -183,13 +183,11 @@ class GemmaAttention(nn.Module):

183 self,

184 positions: torch.Tensor,

185 hidden_states: torch.Tensor,

186 - kv_cache: torch.Tensor,

187 - attn_metadata: AttentionMetadata,

186 ) -> torch.Tensor:

187 qkv, _ = self.qkv_proj(hidden_states)

188 q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)

189 q, k = self.rotary_emb(positions, q, k)

192 - attn_output = self.attn(q, k, v, kv_cache, attn_metadata)

190 + attn_output = self.attn(q, k, v)

191 output, _ = self.o_proj(attn_output)

192 return output

193

@@ -233,8 +231,6 @@ class GemmaDecoderLayer(nn.Module):

231 self,

232 positions: torch.Tensor,

233 hidden_states: torch.Tensor,

236 - kv_cache: torch.Tensor,

237 - attn_metadata: AttentionMetadata,

234 residual: Optional[torch.Tensor],

235 ) -> Tuple[torch.Tensor, torch.Tensor]:

236 # Self Attention

@@ -247,8 +243,6 @@ class GemmaDecoderLayer(nn.Module):

243 hidden_states = self.self_attn(

244 positions=positions,

245 hidden_states=hidden_states,

250 - kv_cache=kv_cache,

251 - attn_metadata=attn_metadata,

246 )

247

248 # Fully Connected

@@ -298,8 +292,6 @@ class GemmaModel(nn.Module):

292 self,

293 input_ids: torch.Tensor,

294 positions: torch.Tensor,

301 - kv_caches: List[torch.Tensor],

302 - attn_metadata: AttentionMetadata,

295 intermediate_tensors: Optional[IntermediateTensors],

296 inputs_embeds: Optional[torch.Tensor] = None,

297 ) -> Union[torch.Tensor, IntermediateTensors]:

@@ -313,13 +305,10 @@ class GemmaModel(nn.Module):

305 else:

306 hidden_states = intermediate_tensors["hidden_states"]

307 residual = intermediate_tensors["residual"]

316 - for i in range(self.start_layer, self.end_layer):

317 - layer = self.layers[i]

308 + for layer in self.layers[self.start_layer:self.end_layer]:

309 hidden_states, residual = layer(

310 positions,

311 hidden_states,

321 - kv_caches[i - self.start_layer],

322 - attn_metadata,

312 residual,

313 )

314 if not get_pp_group().is_last_rank:

@@ -370,13 +359,10 @@ class GemmaForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

359 self,

360 input_ids: torch.Tensor,

361 positions: torch.Tensor,

373 - kv_caches: List[torch.Tensor],

374 - attn_metadata: AttentionMetadata,

362 intermediate_tensors: Optional[IntermediateTensors] = None,

363 inputs_embeds: Optional[torch.Tensor] = None,

364 ) -> Union[torch.Tensor, IntermediateTensors]:

378 - hidden_states = self.model(input_ids, positions, kv_caches,

379 - attn_metadata, intermediate_tensors,

365 + hidden_states = self.model(input_ids, positions, intermediate_tensors,

366 inputs_embeds)

367 return hidden_states

368

vllm/model_executor/models/gemma2.py

+5 -19

@@ -15,13 +15,13 @@

15 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

16 # See the License for the specific language governing permissions and

17 # limitations under the License.

18 -from typing import Iterable, List, Optional, Set, Tuple, Union

18 +from typing import Iterable, Optional, Set, Tuple, Union

19

20 import torch

21 from torch import nn

22 from transformers import Gemma2Config

23

24 -from vllm.attention import Attention, AttentionMetadata

24 +from vllm.attention import Attention

25 from vllm.compilation.decorators import support_torch_compile

26 from vllm.config import CacheConfig, VllmConfig

27 from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size

@@ -164,13 +164,11 @@ class Gemma2Attention(nn.Module):

164 self,

165 positions: torch.Tensor,

166 hidden_states: torch.Tensor,

167 - kv_cache: torch.Tensor,

168 - attn_metadata: AttentionMetadata,

167 ) -> torch.Tensor:

168 qkv, _ = self.qkv_proj(hidden_states)

169 q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)

170 q, k = self.rotary_emb(positions, q, k)

173 - attn_output = self.attn(q, k, v, kv_cache, attn_metadata)

171 + attn_output = self.attn(q, k, v)

172 output, _ = self.o_proj(attn_output)

173 return output

174

@@ -220,8 +218,6 @@ class Gemma2DecoderLayer(nn.Module):

218 self,

219 positions: torch.Tensor,

220 hidden_states: torch.Tensor,

223 - kv_cache: torch.Tensor,

224 - attn_metadata: AttentionMetadata,

221 residual: Optional[torch.Tensor],

222 ) -> Tuple[torch.Tensor, torch.Tensor]:

223 if residual is None:

@@ -233,8 +229,6 @@ class Gemma2DecoderLayer(nn.Module):

229 hidden_states = self.self_attn(

230 positions=positions,

231 hidden_states=hidden_states,

236 - kv_cache=kv_cache,

237 - attn_metadata=attn_metadata,

232 )

233 hidden_states = self.post_attention_layernorm(hidden_states)

234

@@ -284,8 +278,6 @@ class Gemma2Model(nn.Module):

278 self,

279 input_ids: Optional[torch.Tensor],

280 positions: torch.Tensor,

287 - kv_caches: List[torch.Tensor],

288 - attn_metadata: AttentionMetadata,

281 intermediate_tensors: Optional[IntermediateTensors],

282 inputs_embeds: Optional[torch.Tensor] = None,

283 ) -> Union[torch.Tensor, IntermediateTensors]:

@@ -300,13 +292,10 @@ class Gemma2Model(nn.Module):

292 assert intermediate_tensors is not None

293 hidden_states = intermediate_tensors["hidden_states"]

294 residual = intermediate_tensors["residual"]

303 - for i in range(self.start_layer, self.end_layer):

304 - layer = self.layers[i]

295 + for layer in self.layers[self.start_layer:self.end_layer]:

296 hidden_states, residual = layer(

297 positions,

298 hidden_states,

308 - kv_caches[i - self.start_layer],

309 - attn_metadata,

299 residual,

300 )

301 if not get_pp_group().is_last_rank:

@@ -415,13 +404,10 @@ class Gemma2ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

404 self,

405 input_ids: torch.Tensor,

406 positions: torch.Tensor,

418 - kv_caches: List[torch.Tensor],

419 - attn_metadata: AttentionMetadata,

407 intermediate_tensors: Optional[IntermediateTensors] = None,

408 inputs_embeds: Optional[torch.Tensor] = None,

409 ) -> Union[torch.Tensor, IntermediateTensors]:

423 - hidden_states = self.model(input_ids, positions, kv_caches,

424 - attn_metadata, intermediate_tensors,

410 + hidden_states = self.model(input_ids, positions, intermediate_tensors,

411 inputs_embeds)

412 return hidden_states

413

vllm/model_executor/models/glm4v.py

+3 -7

@@ -4,7 +4,7 @@

4 # https://github.com/THUDM/CogAgent

5 """Inference-only CogAgent model compatible with THUDM weights."""

6 from argparse import Namespace

7 -from typing import List, Literal, Mapping, Optional, TypedDict, Union

7 +from typing import Literal, Mapping, Optional, TypedDict, Union

8

9 import torch

10 from torch import nn

@@ -15,7 +15,6 @@ from transformers import PreTrainedTokenizer, TensorType

15 from transformers.image_utils import ImageInput

16 from transformers.tokenization_utils_base import TextInput

17

18 -from vllm.attention import AttentionMetadata

18 from vllm.attention.layer import MultiHeadAttention

19 from vllm.config import VllmConfig

20 from vllm.distributed import get_tensor_model_parallel_world_size

@@ -628,8 +627,6 @@ class GLM4VForCausalLM(ChatGLMBaseModel, SupportsLoRA, SupportsPP,

627 self,

628 input_ids: torch.Tensor,

629 positions: torch.Tensor,

631 - kv_caches: List[torch.Tensor],

632 - attn_metadata: AttentionMetadata,

630 intermediate_tensors: Optional[IntermediateTensors] = None,

631 inputs_embeds: Optional[torch.Tensor] = None,

632 **kwargs: object,

@@ -645,8 +642,7 @@ class GLM4VForCausalLM(ChatGLMBaseModel, SupportsLoRA, SupportsPP,

642 vision_embeddings)

643 input_ids = None

644

648 - hidden_states = self.transformer(input_ids, positions, kv_caches,

649 - attn_metadata, intermediate_tensors,

650 - inputs_embeds)

645 + hidden_states = self.transformer(input_ids, positions,

646 + intermediate_tensors, inputs_embeds)

647

648 return hidden_states

vllm/model_executor/models/gpt2.py

+8 -24

@@ -18,13 +18,13 @@

18 # See the License for the specific language governing permissions and

19 # limitations under the License.

20 """Inference-only GPT-2 model compatible with HuggingFace weights."""

21 -from typing import Iterable, List, Optional, Set, Tuple, Union

21 +from typing import Iterable, Optional, Set, Tuple, Union

22

23 import torch

24 from torch import nn

25 from transformers import GPT2Config

26

27 -from vllm.attention import Attention, AttentionMetadata

27 +from vllm.attention import Attention

28 from vllm.compilation.decorators import support_torch_compile

29 from vllm.config import CacheConfig, VllmConfig

30 from vllm.distributed.parallel_state import (

@@ -92,12 +92,10 @@ class GPT2Attention(nn.Module):

92 def forward(

93 self,

94 hidden_states: torch.Tensor,

95 - kv_cache: torch.Tensor,

96 - attn_metadata: AttentionMetadata,

95 ) -> torch.Tensor:

96 qkv, _ = self.c_attn(hidden_states)

97 q, k, v = qkv.chunk(chunks=3, dim=-1)

100 - attn_output = self.attn(q, k, v, kv_cache, attn_metadata)

98 + attn_output = self.attn(q, k, v)

99 attn_output, _ = self.c_proj(attn_output)

100 return attn_output

101

@@ -164,16 +162,10 @@ class GPT2Block(nn.Module):

162 def forward(

163 self,

164 hidden_states: torch.Tensor,

167 - kv_cache: torch.Tensor,

168 - attn_metadata: AttentionMetadata,

165 ) -> torch.Tensor:

166 residual = hidden_states

167 hidden_states = self.ln_1(hidden_states)

172 - attn_output = self.attn(

173 - hidden_states=hidden_states,

174 - kv_cache=kv_cache,

175 - attn_metadata=attn_metadata,

176 - )

168 + attn_output = self.attn(hidden_states=hidden_states)

169 # residual connection

170 hidden_states = attn_output + residual

171

@@ -222,8 +214,6 @@ class GPT2Model(nn.Module):

214 self,

215 input_ids: torch.Tensor,

216 position_ids: torch.Tensor,

225 - kv_caches: List[torch.Tensor],

226 - attn_metadata: AttentionMetadata,

217 intermediate_tensors: Optional[IntermediateTensors],

218 inputs_embeds: Optional[torch.Tensor],

219 ) -> Union[torch.Tensor, IntermediateTensors]:

@@ -236,11 +226,8 @@ class GPT2Model(nn.Module):

226 assert intermediate_tensors is not None

227 hidden_states = intermediate_tensors["hidden_states"]

228

239 - for i in range(self.start_layer, self.end_layer):

240 - layer = self.h[i]

241 - hidden_states = layer(hidden_states,

242 - kv_caches[i - self.start_layer],

243 - attn_metadata)

229 + for layer in self.h[self.start_layer:self.end_layer]:

230 + hidden_states = layer(hidden_states)

231

232 if not get_pp_group().is_last_rank:

233 return IntermediateTensors({"hidden_states": hidden_states})

@@ -279,14 +266,11 @@ class GPT2LMHeadModel(nn.Module, SupportsPP):

266 self,

267 input_ids: torch.Tensor,

268 positions: torch.Tensor,

282 - kv_caches: List[torch.Tensor],

283 - attn_metadata: AttentionMetadata,

269 intermediate_tensors: Optional[IntermediateTensors] = None,

270 inputs_embeds: Optional[torch.Tensor] = None,

271 ) -> Union[torch.Tensor, IntermediateTensors]:

287 - hidden_states = self.transformer(input_ids, positions, kv_caches,

288 - attn_metadata, intermediate_tensors,

289 - inputs_embeds)

272 + hidden_states = self.transformer(input_ids, positions,

273 + intermediate_tensors, inputs_embeds)

274 return hidden_states

275

276 def compute_logits(

vllm/model_executor/models/gpt_bigcode.py

+8 -24

@@ -19,13 +19,13 @@

19 # See the License for the specific language governing permissions and

20 # limitations under the License.

21 """Inference-only GPTBigCode model compatible with HuggingFace weights."""

22 -from typing import Iterable, List, Optional, Set, Tuple, Union

22 +from typing import Iterable, Optional, Set, Tuple, Union

23

24 import torch

25 from torch import nn

26 from transformers import GPTBigCodeConfig

27

28 -from vllm.attention import Attention, AttentionMetadata

28 +from vllm.attention import Attention

29 from vllm.compilation.decorators import support_torch_compile

30 from vllm.config import CacheConfig, VllmConfig

31 from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size

@@ -101,8 +101,6 @@ class GPTBigCodeAttention(nn.Module):

101 def forward(

102 self,

103 hidden_states: torch.Tensor,

104 - kv_cache: torch.Tensor,

105 - attn_metadata: AttentionMetadata,

104 ) -> torch.Tensor:

105 qkv, _ = self.c_attn(hidden_states)

106 q, k, v = qkv.split(

@@ -112,7 +110,7 @@ class GPTBigCodeAttention(nn.Module):

110 ],

111 dim=-1,

112 )

115 - attn_output = self.attn(q, k, v, kv_cache, attn_metadata)

113 + attn_output = self.attn(q, k, v)

114 attn_output, _ = self.c_proj(attn_output)

115 return attn_output

116

@@ -173,16 +171,10 @@ class GPTBigCodeBlock(nn.Module):

171 def forward(

172 self,

173 hidden_states: torch.Tensor,

176 - kv_cache: torch.Tensor,

177 - attn_metadata: AttentionMetadata,

174 ) -> torch.Tensor:

175 residual = hidden_states

176 hidden_states = self.ln_1(hidden_states)

181 - attn_output = self.attn(

182 - hidden_states=hidden_states,

183 - kv_cache=kv_cache,

184 - attn_metadata=attn_metadata,

185 - )

177 + attn_output = self.attn(hidden_states=hidden_states, )

178 # residual connection

179 hidden_states = attn_output + residual

180

@@ -234,8 +226,6 @@ class GPTBigCodeModel(nn.Module):

226 self,

227 input_ids: torch.Tensor,

228 position_ids: torch.Tensor,

237 - kv_caches: List[torch.Tensor],

238 - attn_metadata: AttentionMetadata,

229 intermediate_tensors: Optional[IntermediateTensors],

230 inputs_embeds: Optional[torch.Tensor] = None,

231 ) -> Union[torch.Tensor, IntermediateTensors]:

@@ -246,11 +236,8 @@ class GPTBigCodeModel(nn.Module):

236 else:

237 hidden_states = intermediate_tensors["hidden_states"]

238

249 - for i in range(self.start_layer, self.end_layer):

250 - layer = self.h[i]

251 - hidden_states = layer(hidden_states,

252 - kv_caches[i - self.start_layer],

253 - attn_metadata)

239 + for layer in self.h[self.start_layer:self.end_layer]:

240 + hidden_states = layer(hidden_states)

241

242 if not get_pp_group().is_last_rank:

243 return IntermediateTensors({"hidden_states": hidden_states})

@@ -302,14 +289,11 @@ class GPTBigCodeForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

289 self,

290 input_ids: torch.Tensor,

291 positions: torch.Tensor,

305 - kv_caches: List[torch.Tensor],

306 - attn_metadata: AttentionMetadata,

292 intermediate_tensors: Optional[IntermediateTensors] = None,

293 inputs_embeds: Optional[torch.Tensor] = None,

294 ) -> Union[torch.Tensor, IntermediateTensors]:

310 - hidden_states = self.transformer(input_ids, positions, kv_caches,

311 - attn_metadata, intermediate_tensors,

312 - inputs_embeds)

295 + hidden_states = self.transformer(input_ids, positions,

296 + intermediate_tensors, inputs_embeds)

297 return hidden_states

298

299 def compute_logits(

vllm/model_executor/models/gpt_j.py

+7 -24

@@ -17,13 +17,13 @@

17 # See the License for the specific language governing permissions and

18 # limitations under the License.

19 """Inference-only GPT-J model compatible with HuggingFace weights."""

20 -from typing import Iterable, List, Optional, Set, Tuple, Union

20 +from typing import Iterable, Optional, Set, Tuple, Union

21

22 import torch

23 from torch import nn

24 from transformers import GPTJConfig

25

26 -from vllm.attention import Attention, AttentionMetadata

26 +from vllm.attention import Attention

27 from vllm.compilation.decorators import support_torch_compile

28 from vllm.config import CacheConfig, VllmConfig

29 from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size

@@ -104,13 +104,11 @@ class GPTJAttention(nn.Module):

104 self,

105 position_ids: torch.Tensor,

106 hidden_states: torch.Tensor,

107 - kv_cache: torch.Tensor,

108 - attn_metadata: AttentionMetadata,

107 ) -> torch.Tensor:

108 qkv, _ = self.qkv_proj(hidden_states)

109 q, k, v = qkv.chunk(chunks=3, dim=-1)

110 q, k = self.rotary_emb(position_ids, q, k)

113 - attn_output = self.attn(q, k, v, kv_cache, attn_metadata)

111 + attn_output = self.attn(q, k, v)

112 attn_output, _ = self.out_proj(attn_output)

113 return attn_output

114

@@ -167,16 +165,12 @@ class GPTJBlock(nn.Module):

165 self,

166 position_ids: torch.Tensor,

167 hidden_states: torch.Tensor,

170 - kv_cache: torch.Tensor,

171 - attn_metadata: AttentionMetadata,

168 ) -> torch.Tensor:

169 residual = hidden_states

170 hidden_states = self.ln_1(hidden_states)

171 attn_output = self.attn(

172 position_ids=position_ids,

173 hidden_states=hidden_states,

178 - kv_cache=kv_cache,

179 - attn_metadata=attn_metadata,

174 )

175 mlp_output = self.mlp(hidden_states)

176 hidden_states = attn_output + mlp_output + residual

@@ -217,8 +211,6 @@ class GPTJModel(nn.Module):

211 self,

212 input_ids: torch.Tensor,

213 position_ids: torch.Tensor,

220 - kv_caches: List[torch.Tensor],

221 - attn_metadata: AttentionMetadata,

214 intermediate_tensors: Optional[IntermediateTensors],

215 inputs_embeds: Optional[torch.Tensor] = None,

216 ) -> Union[torch.Tensor, IntermediateTensors]:

@@ -229,14 +221,8 @@ class GPTJModel(nn.Module):

221 hidden_states = self.get_input_embeddings(input_ids)

222 else:

223 hidden_states = intermediate_tensors["hidden_states"]

232 - for i in range(self.start_layer, self.end_layer):

233 - layer = self.h[i]

234 - hidden_states = layer(

235 - position_ids,

236 - hidden_states,

237 - kv_caches[i - self.start_layer],

238 - attn_metadata,

239 - )

224 + for layer in self.h[self.start_layer:self.end_layer]:

225 + hidden_states = layer(position_ids, hidden_states)

226 if not get_pp_group().is_last_rank:

227 return IntermediateTensors({"hidden_states": hidden_states})

228 hidden_states = self.ln_f(hidden_states)

@@ -273,14 +259,11 @@ class GPTJForCausalLM(nn.Module, SupportsPP):

259 self,

260 input_ids: torch.Tensor,

261 positions: torch.Tensor,

276 - kv_caches: List[torch.Tensor],

277 - attn_metadata: AttentionMetadata,

262 intermediate_tensors: Optional[IntermediateTensors] = None,

263 inputs_embeds: Optional[torch.Tensor] = None,

264 ) -> Union[torch.Tensor, IntermediateTensors]:

281 - hidden_states = self.transformer(input_ids, positions, kv_caches,

282 - attn_metadata, intermediate_tensors,

283 - inputs_embeds)

265 + hidden_states = self.transformer(input_ids, positions,

266 + intermediate_tensors, inputs_embeds)

267 return hidden_states

268

269 def compute_logits(

vllm/model_executor/models/gpt_neox.py

+7 -24

@@ -17,13 +17,13 @@

17 # See the License for the specific language governing permissions and

18 # limitations under the License.

19 """Inference-only GPT-NeoX model compatible with HuggingFace weights."""

20 -from typing import Iterable, List, Optional, Set, Tuple, Union

20 +from typing import Iterable, Optional, Set, Tuple, Union

21

22 import torch

23 from torch import nn

24 from transformers import GPTNeoXConfig

25

26 -from vllm.attention import Attention, AttentionMetadata

26 +from vllm.attention import Attention

27 from vllm.compilation.decorators import support_torch_compile

28 from vllm.config import CacheConfig, VllmConfig

29 from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size

@@ -104,13 +104,11 @@ class GPTNeoXAttention(nn.Module):

104 self,

105 position_ids: torch.Tensor,

106 hidden_states: torch.Tensor,

107 - kv_cache: torch.Tensor,

108 - attn_metadata: AttentionMetadata,

107 ) -> torch.Tensor:

108 qkv, _ = self.query_key_value(hidden_states)

109 q, k, v = qkv.chunk(chunks=3, dim=-1)

110 q, k = self.rotary_emb(position_ids, q, k)

113 - attn_output = self.attn(q, k, v, kv_cache, attn_metadata)

111 + attn_output = self.attn(q, k, v)

112 output, _ = self.dense(attn_output)

113 return output

114

@@ -167,15 +165,11 @@ class GPTNeoXLayer(nn.Module):

165 self,

166 position_ids: torch.Tensor,

167 hidden_states: torch.Tensor,

170 - kv_cache: torch.Tensor,

171 - attn_metadata: AttentionMetadata,

168 ) -> torch.Tensor:

169 attn_input = self.input_layernorm(hidden_states)

170 attn_output = self.attention(

171 position_ids=position_ids,

172 hidden_states=attn_input,

177 - kv_cache=kv_cache,

178 - attn_metadata=attn_metadata,

173 )

174

175 if self.use_parallel_residual:

@@ -230,8 +224,6 @@ class GPTNeoXModel(nn.Module):

224 self,

225 input_ids: torch.Tensor,

226 position_ids: torch.Tensor,

233 - kv_caches: List[torch.Tensor],

234 - attn_metadata: AttentionMetadata,

227 intermediate_tensors: Optional[IntermediateTensors],

228 inputs_embeds: Optional[torch.Tensor] = None,

229 ) -> Union[torch.Tensor, IntermediateTensors]:

@@ -242,14 +234,8 @@ class GPTNeoXModel(nn.Module):

234 hidden_states = self.get_input_embeddings(input_ids)

235 else:

236 hidden_states = intermediate_tensors["hidden_states"]

245 - for i in range(self.start_layer, self.end_layer):

246 - layer = self.layers[i]

247 - hidden_states = layer(

248 - position_ids,

249 - hidden_states,

250 - kv_caches[i - self.start_layer],

251 - attn_metadata,

252 - )

237 + for layer in self.layers[self.start_layer:self.end_layer]:

238 + hidden_states = layer(position_ids, hidden_states)

239 if not get_pp_group().is_last_rank:

240 return IntermediateTensors({"hidden_states": hidden_states})

241 hidden_states = self.final_layer_norm(hidden_states)

@@ -285,14 +271,11 @@ class GPTNeoXForCausalLM(nn.Module, SupportsPP):

271 self,

272 input_ids: torch.Tensor,

273 positions: torch.Tensor,

288 - kv_caches: List[torch.Tensor],

289 - attn_metadata: AttentionMetadata,

274 intermediate_tensors: Optional[IntermediateTensors] = None,

275 inputs_embeds: Optional[torch.Tensor] = None,

276 ) -> Union[torch.Tensor, IntermediateTensors]:

293 - hidden_states = self.gpt_neox(input_ids, positions, kv_caches,

294 - attn_metadata, intermediate_tensors,

295 - inputs_embeds)

277 + hidden_states = self.gpt_neox(input_ids, positions,

278 + intermediate_tensors, inputs_embeds)

279 return hidden_states

280

281 def compute_logits(

vllm/model_executor/models/granite.py

+6 -23

@@ -22,13 +22,13 @@

22 # See the License for the specific language governing permissions and

23 # limitations under the License.

24 """Inference-only IBM Granite model compatible with HuggingFace weights."""

25 -from typing import Any, Dict, Iterable, List, Optional, Set, Tuple, Union

25 +from typing import Any, Dict, Iterable, Optional, Set, Tuple, Union

26

27 import torch

28 from torch import nn

29 from transformers import GraniteConfig

30

31 -from vllm.attention import Attention, AttentionMetadata

31 +from vllm.attention import Attention

32 from vllm.compilation.decorators import support_torch_compile

33 from vllm.config import CacheConfig, VllmConfig

34 from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size

@@ -166,13 +166,11 @@ class GraniteAttention(nn.Module):

166 self,

167 positions: torch.Tensor,

168 hidden_states: torch.Tensor,

169 - kv_cache: torch.Tensor,

170 - attn_metadata: AttentionMetadata,

169 ) -> torch.Tensor:

170 qkv, _ = self.qkv_proj(hidden_states)

171 q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)

172 q, k = self.rotary_emb(positions, q, k)

175 - attn_output = self.attn(q, k, v, kv_cache, attn_metadata)

173 + attn_output = self.attn(q, k, v)

174 output, _ = self.o_proj(attn_output)

175 return output

176

@@ -233,8 +231,6 @@ class GraniteDecoderLayer(nn.Module):

231 self,

232 positions: torch.Tensor,

233 hidden_states: torch.Tensor,

236 - kv_cache: torch.Tensor,

237 - attn_metadata: AttentionMetadata,

234 ) -> Tuple[torch.Tensor, torch.Tensor]:

235 # Self Attention

236 residual = hidden_states

@@ -242,8 +238,6 @@ class GraniteDecoderLayer(nn.Module):

238 hidden_states = self.self_attn(

239 positions=positions,

240 hidden_states=hidden_states,

245 - kv_cache=kv_cache,

246 - attn_metadata=attn_metadata,

241 )

242 hidden_states = residual + hidden_states * self.residual_multiplier

243 # Fully Connected

@@ -300,8 +294,6 @@ class GraniteModel(nn.Module):

294 self,

295 input_ids: Optional[torch.Tensor],

296 positions: torch.Tensor,

303 - kv_caches: List[torch.Tensor],

304 - attn_metadata: AttentionMetadata,

297 intermediate_tensors: Optional[IntermediateTensors],

298 inputs_embeds: Optional[torch.Tensor] = None,

299 ) -> Union[torch.Tensor, IntermediateTensors]:

@@ -318,14 +310,8 @@ class GraniteModel(nn.Module):

310 hidden_states = intermediate_tensors["hidden_states"]

311 residual = intermediate_tensors["residual"]

312

321 - for i in range(self.start_layer, self.end_layer):

322 - layer = self.layers[i]

323 - hidden_states = layer(

324 - positions,

325 - hidden_states,

326 - kv_caches[i - self.start_layer],

327 - attn_metadata,

328 - )

313 + for layer in self.layers[self.start_layer:self.end_layer]:

314 + hidden_states = layer(positions, hidden_states)

315

316 if not get_pp_group().is_last_rank:

317 return IntermediateTensors({

@@ -405,13 +391,10 @@ class GraniteForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

391 self,

392 input_ids: torch.Tensor,

393 positions: torch.Tensor,

408 - kv_caches: List[torch.Tensor],

409 - attn_metadata: AttentionMetadata,

394 intermediate_tensors: Optional[IntermediateTensors] = None,

395 inputs_embeds: Optional[torch.Tensor] = None,

396 ) -> Union[torch.Tensor, IntermediateTensors]:

413 - model_output = self.model(input_ids, positions, kv_caches,

414 - attn_metadata, intermediate_tensors,

397 + model_output = self.model(input_ids, positions, intermediate_tensors,

398 inputs_embeds)

399 return model_output

400

vllm/model_executor/models/granitemoe.py

+6 -20

@@ -22,13 +22,13 @@

22 # See the License for the specific language governing permissions and

23 # limitations under the License.

24 """Inference-only GraniteMoe model."""

25 -from typing import Iterable, List, Optional, Set, Tuple

25 +from typing import Iterable, Optional, Set, Tuple

26

27 import torch

28 from torch import nn

29 from transformers.models.granitemoe import GraniteMoeConfig

30

31 -from vllm.attention import Attention, AttentionMetadata

31 +from vllm.attention import Attention

32 from vllm.compilation.decorators import support_torch_compile

33 from vllm.config import CacheConfig, VllmConfig

34 from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size

@@ -173,13 +173,11 @@ class GraniteMoeAttention(nn.Module):

173 self,

174 positions: torch.Tensor,

175 hidden_states: torch.Tensor,

176 - kv_cache: torch.Tensor,

177 - attn_metadata: AttentionMetadata,

176 ) -> torch.Tensor:

177 qkv, _ = self.qkv_proj(hidden_states)

178 q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)

179 q, k = self.rotary_emb(positions, q, k)

182 - attn_output = self.attn(q, k, v, kv_cache, attn_metadata)

180 + attn_output = self.attn(q, k, v)

181 output, _ = self.o_proj(attn_output)

182 return output

183

@@ -226,8 +224,6 @@ class GraniteMoeDecoderLayer(nn.Module):

224 self,

225 positions: torch.Tensor,

226 hidden_states: torch.Tensor,

229 - kv_cache: torch.Tensor,

230 - attn_metadata: AttentionMetadata,

227 ) -> torch.Tensor:

228 # Self Attention

229 residual = hidden_states

@@ -235,8 +231,6 @@ class GraniteMoeDecoderLayer(nn.Module):

231 hidden_states = self.self_attn(

232 positions=positions,

233 hidden_states=hidden_states,

238 - kv_cache=kv_cache,

239 - attn_metadata=attn_metadata,

234 )

235 hidden_states = residual + hidden_states * self.residual_multiplier

236 residual = hidden_states

@@ -287,8 +281,6 @@ class GraniteMoeModel(nn.Module):

281 self,

282 input_ids: torch.Tensor,

283 positions: torch.Tensor,

290 - kv_caches: List[torch.Tensor],

291 - attn_metadata: AttentionMetadata,

284 intermediate_tensors: Optional[IntermediateTensors],

285 inputs_embeds: Optional[torch.Tensor] = None,

286 ) -> torch.Tensor:

@@ -303,11 +295,8 @@ class GraniteMoeModel(nn.Module):

295 assert intermediate_tensors is not None

296 hidden_states = intermediate_tensors["hidden_states"]

297 residual = intermediate_tensors["residual"]

306 - for i in range(self.start_layer, self.end_layer):

307 - layer = self.layers[i]

308 - hidden_states = layer(positions, hidden_states,

309 - kv_caches[i - self.start_layer],

310 - attn_metadata)

298 + for layer in self.layers[self.start_layer:self.end_layer]:

299 + hidden_states = layer(positions, hidden_states)

300 if not get_pp_group().is_last_rank:

301 return IntermediateTensors({

302 "hidden_states": hidden_states,

@@ -377,13 +366,10 @@ class GraniteMoeForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

366 self,

367 input_ids: torch.Tensor,

368 positions: torch.Tensor,

380 - kv_caches: List[torch.Tensor],

381 - attn_metadata: AttentionMetadata,

369 intermediate_tensors: Optional[IntermediateTensors] = None,

370 inputs_embeds: Optional[torch.Tensor] = None,

371 ) -> torch.Tensor:

385 - hidden_states = self.model(input_ids, positions, kv_caches,

386 - attn_metadata, intermediate_tensors,

372 + hidden_states = self.model(input_ids, positions, intermediate_tensors,

373 inputs_embeds)

374 return hidden_states

375

vllm/model_executor/models/gritlm.py

+3 -6

@@ -1,15 +1,15 @@

1 # SPDX-License-Identifier: Apache-2.0

2

3 from array import array

4 -from typing import List, Optional, Union

4 +from typing import Optional, Union

5

6 import torch

7 import torch.nn as nn

8 from xformers.ops.fmha.attn_bias import BlockDiagonalMask

9

10 -from vllm.attention import AttentionMetadata

10 from vllm.attention.backends.xformers import XFormersImpl

11 from vllm.config import ModelConfig, VllmConfig

12 +from vllm.forward_context import get_forward_context

13 from vllm.logger import init_logger

14 from vllm.model_executor.layers.pooler import PoolerHead

15 from vllm.model_executor.models.llama import LlamaForCausalLM

@@ -217,13 +217,12 @@ class GritLM(LlamaForCausalLM):

217 self,

218 input_ids: torch.Tensor,

219 positions: torch.Tensor,

220 - kv_caches: List[torch.Tensor],

221 - attn_metadata: AttentionMetadata,

220 **kwargs,

221 ) -> Union[torch.Tensor, IntermediateTensors]:

222

223 # Change attention to non-causal for pooling tasks.

224 if self.runner_type == "pooling":

225 + attn_metadata = get_forward_context().attn_metadata

226 assert attn_metadata.prefill_metadata.attn_bias is None

227 attn_metadata.prefill_metadata.attn_bias = [

228 BlockDiagonalMask.from_seqlens(attn_metadata.seq_lens)

@@ -232,8 +231,6 @@ class GritLM(LlamaForCausalLM):

231 return super().forward(

232 input_ids=input_ids,

233 positions=positions,

235 - kv_caches=kv_caches,

236 - attn_metadata=attn_metadata,

234 **kwargs,

235 )

236

vllm/model_executor/models/idefics3.py

-9

@@ -25,7 +25,6 @@ from torch import nn

25 from transformers import (BatchFeature, Idefics3Config, Idefics3ImageProcessor,

26 Idefics3Processor)

27

28 -from vllm.attention import AttentionMetadata

28 from vllm.config import VllmConfig

29 from vllm.logger import init_logger

30 from vllm.model_executor.layers.linear import ReplicatedLinear

@@ -563,8 +562,6 @@ class Idefics3Model(nn.Module):

562 self,

563 input_ids: torch.Tensor,

564 positions: torch.Tensor,

566 - kv_caches: List[torch.Tensor],

567 - attn_metadata: AttentionMetadata,

565 intermediate_tensors: Optional[IntermediateTensors] = None,

566 inputs_embeds: Optional[torch.Tensor] = None,

567 ) -> Union[torch.Tensor, IntermediateTensors]:

@@ -572,8 +569,6 @@ class Idefics3Model(nn.Module):

569 hidden_states = self.text_model(

570 input_ids,

571 positions,

575 - kv_caches,

576 - attn_metadata,

572 intermediate_tensors,

573 inputs_embeds=inputs_embeds,

574 )

@@ -645,8 +640,6 @@ class Idefics3ForConditionalGeneration(nn.Module, SupportsMultiModal,

640 self,

641 input_ids: torch.Tensor,

642 positions: torch.Tensor,

648 - kv_caches: List[torch.Tensor],

649 - attn_metadata: AttentionMetadata,

643 intermediate_tensors: Optional[IntermediateTensors] = None,

644 inputs_embeds: Optional[torch.Tensor] = None,

645 **kwargs: object,

@@ -664,8 +657,6 @@ class Idefics3ForConditionalGeneration(nn.Module, SupportsMultiModal,

657

658 hidden_states = self.model.text_model(input_ids,

659 positions,

667 - kv_caches,

668 - attn_metadata,

660 intermediate_tensors,

661 inputs_embeds=inputs_embeds)

662

vllm/model_executor/models/interfaces_base.py

+3 -6

@@ -1,7 +1,7 @@

1 # SPDX-License-Identifier: Apache-2.0

2

3 -from typing import (TYPE_CHECKING, List, Optional, Protocol, Type, Union,

4 - overload, runtime_checkable)

3 +from typing import (TYPE_CHECKING, Optional, Protocol, Type, Union, overload,

4 + runtime_checkable)

5

6 import torch

7 import torch.nn as nn

@@ -11,7 +11,6 @@ from vllm.logger import init_logger

11 from vllm.utils import supports_kw

12

13 if TYPE_CHECKING:

14 - from vllm.attention import AttentionMetadata

14 from vllm.config import VllmConfig

15 from vllm.model_executor.layers.pooler import PoolerOutput

16 from vllm.model_executor.layers.sampler import SamplerOutput

@@ -46,8 +45,6 @@ class VllmModel(Protocol[T_co]):

45 self,

46 input_ids: torch.Tensor,

47 positions: torch.Tensor,

49 - kv_caches: List[torch.Tensor],

50 - attn_metadata: "AttentionMetadata",

48 ) -> T_co:

49 ...

50

@@ -62,7 +59,7 @@ def _check_vllm_model_forward(model: Union[Type[object], object]) -> bool:

59 if not callable(model_forward):

60 return False

61

65 - vllm_kws = ("input_ids", "positions", "kv_caches", "attn_metadata")

62 + vllm_kws = ("input_ids", "positions")

63 missing_kws = tuple(kw for kw in vllm_kws

64 if not supports_kw(model_forward, kw))

65

vllm/model_executor/models/internlm2.py

+7 -28

@@ -1,13 +1,13 @@

1 # SPDX-License-Identifier: Apache-2.0

2

3 from functools import partial

4 -from typing import Any, Dict, Iterable, List, Optional, Set, Tuple, Type, Union

4 +from typing import Any, Dict, Iterable, Optional, Set, Tuple, Type, Union

5

6 import torch

7 from torch import nn

8 from transformers import PretrainedConfig

9

10 -from vllm.attention import Attention, AttentionMetadata

10 +from vllm.attention import Attention

11 from vllm.compilation.decorators import support_torch_compile

12 from vllm.config import CacheConfig, VllmConfig

13 from vllm.distributed import (get_pp_group, get_tensor_model_parallel_rank,

@@ -175,13 +175,11 @@ class InternLM2Attention(nn.Module):

175 self,

176 positions: torch.Tensor,

177 hidden_states: torch.Tensor,

178 - kv_cache: torch.Tensor,

179 - attn_metadata: AttentionMetadata,

178 ) -> torch.Tensor:

179 qkv, _ = self.wqkv(hidden_states)

180 q, k, v = self.split_qkv(qkv)

181 q, k = self.rotary_emb(positions, q, k)

184 - attn_output = self.attn(q, k, v, kv_cache, attn_metadata)

182 + attn_output = self.attn(q, k, v)

183 output, _ = self.wo(attn_output)

184 return output

185

@@ -227,8 +225,6 @@ class InternLMDecoderLayer(nn.Module):

225 self,

226 positions: torch.Tensor,

227 hidden_states: torch.Tensor,

230 - kv_cache: torch.Tensor,

231 - attn_metadata: AttentionMetadata,

228 residual: Optional[torch.Tensor],

229 ) -> Tuple[torch.Tensor, torch.Tensor]:

230 # Self Attention

@@ -241,8 +237,6 @@ class InternLMDecoderLayer(nn.Module):

237 hidden_states = self.attention(

238 positions=positions,

239 hidden_states=hidden_states,

244 - kv_cache=kv_cache,

245 - attn_metadata=attn_metadata,

240 )

241

242 # Fully Connected

@@ -290,8 +284,6 @@ class InternLM2Model(nn.Module):

284 self,

285 input_ids: torch.Tensor,

286 positions: torch.Tensor,

293 - kv_caches: List[torch.Tensor],

294 - attn_metadata: AttentionMetadata,

287 intermediate_tensors: Optional[IntermediateTensors] = None,

288 inputs_embeds: Optional[torch.Tensor] = None,

289 ) -> Union[torch.Tensor, IntermediateTensors]:

@@ -305,15 +297,8 @@ class InternLM2Model(nn.Module):

297 assert intermediate_tensors is not None

298 hidden_states = intermediate_tensors["hidden_states"]

299 residual = intermediate_tensors["residual"]

308 - for i in range(self.start_layer, self.end_layer):

309 - layer = self.layers[i]

310 - hidden_states, residual = layer(

311 - positions,

312 - hidden_states,

313 - kv_caches[i - self.start_layer],

314 - attn_metadata,

315 - residual,

316 - )

300 + for layer in self.layers[self.start_layer:self.end_layer]:

301 + hidden_states, residual = layer(positions, hidden_states, residual)

302 if not get_pp_group().is_last_rank:

303 return IntermediateTensors({

304 "hidden_states": hidden_states,

@@ -363,13 +348,10 @@ class InternLM2ForCausalLM(nn.Module, SupportsPP, SupportsLoRA):

348 self,

349 input_ids: torch.Tensor,

350 positions: torch.Tensor,

366 - kv_caches: List[torch.Tensor],

367 - attn_metadata: AttentionMetadata,

351 intermediate_tensors: Optional[IntermediateTensors],

352 inputs_embeds: Optional[torch.Tensor] = None,

353 ) -> torch.Tensor:

371 - hidden_states = self.model(input_ids, positions, kv_caches,

372 - attn_metadata, intermediate_tensors,

354 + hidden_states = self.model(input_ids, positions, intermediate_tensors,

355 inputs_embeds)

356 return hidden_states

357

@@ -466,13 +448,10 @@ class InternLM2ForRewardModel(InternLM2ForCausalLM):

448 self,

449 input_ids: torch.Tensor,

450 positions: torch.Tensor,

469 - kv_caches: List[torch.Tensor],

470 - attn_metadata: AttentionMetadata,

451 intermediate_tensors: Optional[IntermediateTensors] = None,

452 inputs_embeds: Optional[torch.Tensor] = None,

453 ) -> Union[torch.Tensor, IntermediateTensors]:

474 - hidden_states = self.model(input_ids, positions, kv_caches,

475 - attn_metadata, intermediate_tensors,

454 + hidden_states = self.model(input_ids, positions, intermediate_tensors,

455 inputs_embeds)

456 logits, _ = self.v_head(hidden_states)

457 return logits

vllm/model_executor/models/internlm2_ve.py

+2 -12

@@ -1,12 +1,11 @@

1 # SPDX-License-Identifier: Apache-2.0

2

3 -from typing import List, Optional, Tuple, Union

3 +from typing import Optional, Tuple, Union

4

5 import torch

6 from torch import nn

7 from transformers import PretrainedConfig

8

9 -from vllm.attention import AttentionMetadata

9 from vllm.config import CacheConfig, VllmConfig

10 from vllm.distributed import get_pp_group

11 from vllm.model_executor.layers.layernorm import RMSNorm

@@ -65,8 +64,6 @@ class InternLM2VEDecoderLayer(nn.Module):

64 self,

65 positions: torch.Tensor,

66 hidden_states: torch.Tensor,

68 - kv_cache: torch.Tensor,

69 - attn_metadata: AttentionMetadata,

67 residual: Optional[torch.Tensor],

68 visual_token_mask: Optional[torch.Tensor] = None,

69 ) -> Tuple[torch.Tensor, torch.Tensor]:

@@ -80,8 +77,6 @@ class InternLM2VEDecoderLayer(nn.Module):

77 hidden_states = self.attention(

78 positions=positions,

79 hidden_states=hidden_states,

83 - kv_cache=kv_cache,

84 - attn_metadata=attn_metadata,

80 )

81

82 # Fully Connected

@@ -113,8 +108,6 @@ class InternLM2VEModel(InternLM2Model):

108 self,

109 input_ids: torch.Tensor,

110 positions: torch.Tensor,

116 - kv_caches: List[torch.Tensor],

117 - attn_metadata: AttentionMetadata,

111 intermediate_tensors: Optional[IntermediateTensors] = None,

112 inputs_embeds: Optional[torch.Tensor] = None,

113 visual_token_mask: Optional[torch.Tensor] = None,

@@ -129,13 +122,10 @@ class InternLM2VEModel(InternLM2Model):

122 assert intermediate_tensors is not None

123 hidden_states = intermediate_tensors["hidden_states"]

124 residual = intermediate_tensors["residual"]

132 - for i in range(self.start_layer, self.end_layer):

133 - layer = self.layers[i]

125 + for layer in self.layers[self.start_layer:self.end_layer]:

126 hidden_states, residual = layer(

127 positions,

128 hidden_states,

137 - kv_caches[i - self.start_layer],

138 - attn_metadata,

129 residual,

130 visual_token_mask=visual_token_mask,

131 )

vllm/model_executor/models/internvl.py

-5

@@ -17,7 +17,6 @@ import torchvision.transforms as T

17 from PIL import Image

18 from transformers import BatchFeature, PretrainedConfig, TensorType

19

20 -from vllm.attention import AttentionMetadata

20 from vllm.config import VllmConfig

21 from vllm.model_executor.layers.quantization import QuantizationConfig

22 from vllm.model_executor.layers.quantization.awq import AWQConfig

@@ -929,8 +928,6 @@ class InternVLChatModel(nn.Module, SupportsMultiModal, SupportsPP):

928 self,

929 input_ids: torch.Tensor,

930 positions: torch.Tensor,

932 - kv_caches: List[torch.Tensor],

933 - attn_metadata: AttentionMetadata,

931 intermediate_tensors: Optional[IntermediateTensors] = None,

932 inputs_embeds: Optional[torch.Tensor] = None,

933 **kwargs: object,

@@ -951,8 +948,6 @@ class InternVLChatModel(nn.Module, SupportsMultiModal, SupportsPP):

948 forward_kwargs = {

949 "input_ids": input_ids,

950 "positions": positions,

954 - "kv_caches": kv_caches,

955 - "attn_metadata": attn_metadata,

951 "intermediate_tensors": intermediate_tensors,

952 "inputs_embeds": inputs_embeds,

953 }

vllm/model_executor/models/jais.py

+8 -24

@@ -21,12 +21,12 @@

21 """Inference-only Jais model compatible with HuggingFace weights."""

22

23 import math

24 -from typing import Iterable, List, Optional, Set, Tuple, Union

24 +from typing import Iterable, Optional, Set, Tuple, Union

25

26 import torch

27 from torch import nn

28

29 -from vllm.attention import Attention, AttentionMetadata

29 +from vllm.attention import Attention

30 from vllm.compilation.decorators import support_torch_compile

31 from vllm.config import CacheConfig, VllmConfig

32 from vllm.distributed import (get_pp_group, get_tensor_model_parallel_rank,

@@ -123,12 +123,10 @@ class JAISAttention(nn.Module):

123 def forward(

124 self,

125 hidden_states: torch.Tensor,

126 - kv_cache: torch.Tensor,

127 - attn_metadata: AttentionMetadata,

126 ) -> torch.Tensor:

127 qkv, _ = self.c_attn(hidden_states)

128 q, k, v = qkv.chunk(chunks=3, dim=-1)

131 - attn_output = self.attn(q, k, v, kv_cache, attn_metadata)

129 + attn_output = self.attn(q, k, v)

130 attn_output, _ = self.c_proj(attn_output)

131 return attn_output

132

@@ -200,16 +198,10 @@ class JAISBlock(nn.Module):

198 def forward(

199 self,

200 hidden_states: torch.Tensor,

203 - kv_cache: torch.Tensor,

204 - attn_metadata: AttentionMetadata,

201 ) -> torch.Tensor:

202 residual = hidden_states

203 hidden_states = self.ln_1(hidden_states)

208 - attn_output = self.attn(

209 - hidden_states=hidden_states,

210 - kv_cache=kv_cache,

211 - attn_metadata=attn_metadata,

212 - )

204 + attn_output = self.attn(hidden_states=hidden_states, )

205 # residual connection

206 hidden_states = attn_output + residual

207

@@ -266,8 +258,6 @@ class JAISModel(nn.Module):

258 self,

259 input_ids: torch.Tensor,

260 position_ids: torch.Tensor,

269 - kv_caches: List[torch.Tensor],

270 - attn_metadata: AttentionMetadata,

261 intermediate_tensors: Optional[IntermediateTensors] = None,

262 inputs_embeds: Optional[torch.Tensor] = None,

263 ) -> Union[IntermediateTensors, torch.Tensor]:

@@ -285,11 +275,8 @@ class JAISModel(nn.Module):

275 assert intermediate_tensors is not None

276 hidden_states = intermediate_tensors["hidden_states"]

277

288 - for i in range(self.start_layer, self.end_layer):

289 - layer = self.h[i]

290 - hidden_states = layer(hidden_states,

291 - kv_caches[i - self.start_layer],

292 - attn_metadata)

278 + for layer in self.h[self.start_layer:self.end_layer]:

279 + hidden_states = layer(hidden_states)

280

281 if not get_pp_group().is_last_rank:

282 return IntermediateTensors({"hidden_states": hidden_states})

@@ -332,14 +319,11 @@ class JAISLMHeadModel(nn.Module, SupportsPP):

319 self,

320 input_ids: torch.Tensor,

321 positions: torch.Tensor,

335 - kv_caches: List[torch.Tensor],

336 - attn_metadata: AttentionMetadata,

322 intermediate_tensors: Optional[IntermediateTensors] = None,

323 inputs_embeds: Optional[torch.Tensor] = None,

324 ) -> Union[IntermediateTensors, torch.Tensor]:

340 - hidden_states = self.transformer(input_ids, positions, kv_caches,

341 - attn_metadata, intermediate_tensors,

342 - inputs_embeds)

325 + hidden_states = self.transformer(input_ids, positions,

326 + intermediate_tensors, inputs_embeds)

327 return hidden_states

328

329 def compute_logits(

vllm/model_executor/models/jamba.py

+5 -24

@@ -1,12 +1,11 @@

1 # SPDX-License-Identifier: Apache-2.0

2 """Inference-only Jamba model."""

3 -from typing import Iterable, List, Optional, Set, Tuple

3 +from typing import Iterable, Optional, Set, Tuple

4

5 import torch

6 from torch import nn

7 from transformers import JambaConfig

8

9 -from vllm.attention.backends.abstract import AttentionMetadata

9 from vllm.attention.layer import Attention

10 from vllm.config import CacheConfig, VllmConfig

11 from vllm.distributed import get_tensor_model_parallel_world_size

@@ -138,7 +137,6 @@ class JambaMambaDecoderLayer(nn.Module):

137 def forward(

138 self,

139 hidden_states: torch.Tensor,

141 - attn_metadata: AttentionMetadata,

140 residual: Optional[torch.Tensor],

141 mamba_cache_params: MambaCacheParams,

142 **kwargs,

@@ -150,8 +148,7 @@ class JambaMambaDecoderLayer(nn.Module):

148 hidden_states, residual = self.input_layernorm(

149 hidden_states, residual)

150

153 - hidden_states = self.mamba(hidden_states, attn_metadata,

154 - mamba_cache_params)

151 + hidden_states = self.mamba(hidden_states, mamba_cache_params)

152 # Fully Connected

153 hidden_states, residual = self.pre_ff_layernorm(

154 hidden_states, residual)

@@ -223,13 +220,11 @@ class JambaAttentionDecoderLayer(nn.Module):

220 self,

221 positions: torch.Tensor,

222 hidden_states: torch.Tensor,

226 - kv_cache: torch.Tensor,

227 - attn_metadata: AttentionMetadata,

223 **kwargs,

224 ) -> torch.Tensor:

225 qkv, _ = self.qkv_proj(hidden_states)

226 q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)

232 - attn_output = self.attn(q, k, v, kv_cache, attn_metadata)

227 + attn_output = self.attn(q, k, v)

228 output, _ = self.o_proj(attn_output)

229 return output

230

@@ -237,8 +232,6 @@ class JambaAttentionDecoderLayer(nn.Module):

232 self,

233 positions: torch.Tensor,

234 hidden_states: torch.Tensor,

240 - kv_cache: torch.Tensor,

241 - attn_metadata: AttentionMetadata,

235 residual: Optional[torch.Tensor],

236 **kwargs,

237 ):

@@ -252,8 +245,6 @@ class JambaAttentionDecoderLayer(nn.Module):

245 hidden_states = self.self_attention(

246 positions=positions,

247 hidden_states=hidden_states,

255 - kv_cache=kv_cache,

256 - attn_metadata=attn_metadata,

248 )

249 # Fully Connected

250 hidden_states, residual = self.pre_ff_layernorm(

@@ -320,8 +311,6 @@ class JambaModel(nn.Module):

311 self,

312 input_ids: torch.Tensor,

313 positions: torch.Tensor,

323 - kv_caches: List[torch.Tensor],

324 - attn_metadata: AttentionMetadata,

314 mamba_cache_params: MambaCacheParams,

315 intermediate_tensors: Optional[IntermediateTensors] = None,

316 inputs_embeds: Optional[torch.Tensor] = None,

@@ -339,12 +328,9 @@ class JambaModel(nn.Module):

328

329 kv_cache_index = 0

330 mamba_cache_index = 0

342 - for i in range(self.start_layer, self.end_layer):

343 - layer = self.layers[i]

344 - kv_cache = None

331 + for layer in self.layers[self.start_layer:self.end_layer]:

332 layer_mamba_cache_params = None

333 if isinstance(layer, JambaAttentionDecoderLayer):

347 - kv_cache = kv_caches[kv_cache_index]

334 kv_cache_index += 1

335 if isinstance(layer, JambaMambaDecoderLayer):

336 current_state_layer = mamba_cache_index

@@ -355,8 +341,6 @@ class JambaModel(nn.Module):

341 hidden_states, residual = layer(

342 positions=positions,

343 hidden_states=hidden_states,

358 - kv_cache=kv_cache,

359 - attn_metadata=attn_metadata,

344 residual=residual,

345 mamba_cache_params=layer_mamba_cache_params)

346 if not get_pp_group().is_last_rank:

@@ -429,8 +413,6 @@ class JambaForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP,

413 def forward(self,

414 input_ids: torch.Tensor,

415 positions: torch.Tensor,

432 - kv_caches: List[KVCache],

433 - attn_metadata: AttentionMetadata,

416 intermediate_tensors: Optional[IntermediateTensors] = None,

417 inputs_embeds: Optional[torch.Tensor] = None,

418 **kwargs):

@@ -443,8 +425,7 @@ class JambaForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP,

425

426 mamba_cache_params = self.mamba_cache.current_run_tensors(**kwargs)

427

446 - hidden_states = self.model(input_ids, positions, kv_caches,

447 - attn_metadata, mamba_cache_params,

428 + hidden_states = self.model(input_ids, positions, mamba_cache_params,

429 intermediate_tensors, inputs_embeds)

430 return hidden_states

431

vllm/model_executor/models/llama.py

+7 -21

@@ -22,13 +22,13 @@

22 # See the License for the specific language governing permissions and

23 # limitations under the License.

24 """Inference-only LLaMA model compatible with HuggingFace weights."""

25 -from typing import Any, Dict, Iterable, List, Optional, Set, Tuple, Type, Union

25 +from typing import Any, Dict, Iterable, Optional, Set, Tuple, Type, Union

26

27 import torch

28 from torch import nn

29 from transformers import LlamaConfig

30

31 -from vllm.attention import Attention, AttentionMetadata

31 +from vllm.attention import Attention

32 from vllm.compilation.decorators import support_torch_compile

33 from vllm.config import CacheConfig, VllmConfig

34 from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size

@@ -197,13 +197,11 @@ class LlamaAttention(nn.Module):

197 self,

198 positions: torch.Tensor,

199 hidden_states: torch.Tensor,

200 - kv_cache: torch.Tensor,

201 - attn_metadata: AttentionMetadata,

200 ) -> torch.Tensor:

201 qkv, _ = self.qkv_proj(hidden_states)

202 q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)

203 q, k = self.rotary_emb(positions, q, k)

206 - attn_output = self.attn(q, k, v, kv_cache, attn_metadata)

204 + attn_output = self.attn(q, k, v)

205 output, _ = self.o_proj(attn_output)

206 return output

207

@@ -268,8 +266,6 @@ class LlamaDecoderLayer(nn.Module):

266 self,

267 positions: torch.Tensor,

268 hidden_states: torch.Tensor,

271 - kv_cache: torch.Tensor,

272 - attn_metadata: AttentionMetadata,

269 residual: Optional[torch.Tensor],

270 ) -> Tuple[torch.Tensor, torch.Tensor]:

271 # Self Attention

@@ -280,9 +276,7 @@ class LlamaDecoderLayer(nn.Module):

276 hidden_states, residual = self.input_layernorm(

277 hidden_states, residual)

278 hidden_states = self.self_attn(positions=positions,

283 - hidden_states=hidden_states,

284 - kv_cache=kv_cache,

285 - attn_metadata=attn_metadata)

279 + hidden_states=hidden_states)

280

281 # Fully Connected

282 hidden_states, residual = self.post_attention_layernorm(

@@ -347,8 +341,6 @@ class LlamaModel(nn.Module):

341 self,

342 input_ids: Optional[torch.Tensor],

343 positions: torch.Tensor,

350 - kv_caches: List[torch.Tensor],

351 - attn_metadata: AttentionMetadata,

344 intermediate_tensors: Optional[IntermediateTensors],

345 inputs_embeds: Optional[torch.Tensor] = None,

346 ) -> Union[torch.Tensor, IntermediateTensors]:

@@ -363,11 +355,8 @@ class LlamaModel(nn.Module):

355 hidden_states = intermediate_tensors["hidden_states"]

356 residual = intermediate_tensors["residual"]

357

366 - for i in range(self.start_layer, self.end_layer):

367 - layer = self.layers[i]

368 - hidden_states, residual = layer(positions, hidden_states,

369 - kv_caches[i - self.start_layer],

370 - attn_metadata, residual)

358 + for layer in self.layers[self.start_layer:self.end_layer]:

359 + hidden_states, residual = layer(positions, hidden_states, residual)

360

361 if not get_pp_group().is_last_rank:

362 return IntermediateTensors({

@@ -535,13 +524,10 @@ class LlamaForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

524 self,

525 input_ids: torch.Tensor,

526 positions: torch.Tensor,

538 - kv_caches: List[torch.Tensor],

539 - attn_metadata: AttentionMetadata,

527 intermediate_tensors: Optional[IntermediateTensors] = None,

528 inputs_embeds: Optional[torch.Tensor] = None,

529 ) -> Union[torch.Tensor, IntermediateTensors]:

543 - model_output = self.model(input_ids, positions, kv_caches,

544 - attn_metadata, intermediate_tensors,

530 + model_output = self.model(input_ids, positions, intermediate_tensors,

531 inputs_embeds)

532 return model_output

533

vllm/model_executor/models/llava.py

-5

@@ -15,7 +15,6 @@ from transformers import __version__ as TRANSFORMERS_VERSION

15 from transformers.models.llava import LlavaProcessor

16 from transformers.models.pixtral import PixtralProcessor

17

18 -from vllm.attention import AttentionMetadata

18 from vllm.config import VllmConfig

19 from vllm.inputs import InputProcessingContext

20 from vllm.model_executor.layers.activation import get_act_fn

@@ -658,8 +657,6 @@ class LlavaForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP):

657 self,

658 input_ids: torch.Tensor,

659 positions: torch.Tensor,

661 - kv_caches: List[torch.Tensor],

662 - attn_metadata: AttentionMetadata,

660 intermediate_tensors: Optional[IntermediateTensors] = None,

661 inputs_embeds: Optional[torch.Tensor] = None,

662 **kwargs: object,

@@ -712,8 +709,6 @@ class LlavaForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP):

709

710 hidden_states = self.language_model.model(input_ids,

711 positions,

715 - kv_caches,

716 - attn_metadata,

712 intermediate_tensors,

713 inputs_embeds=inputs_embeds)

714

vllm/model_executor/models/llava_next.py

-5

@@ -12,7 +12,6 @@ from transformers.models.llava_next.modeling_llava_next import (

12 get_anyres_image_grid_shape, unpad_image)

13 from typing_extensions import NotRequired

14

15 -from vllm.attention import AttentionMetadata

15 from vllm.config import VllmConfig

16 from vllm.model_executor.layers.sampler import SamplerOutput, get_sampler

17 from vllm.model_executor.sampling_metadata import SamplingMetadata

@@ -508,8 +507,6 @@ class LlavaNextForConditionalGeneration(nn.Module, SupportsMultiModal,

507 self,

508 input_ids: torch.Tensor,

509 positions: torch.Tensor,

511 - kv_caches: List[torch.Tensor],

512 - attn_metadata: AttentionMetadata,

510 intermediate_tensors: Optional[IntermediateTensors] = None,

511 inputs_embeds: Optional[torch.Tensor] = None,

512 **kwargs: object,

@@ -571,8 +568,6 @@ class LlavaNextForConditionalGeneration(nn.Module, SupportsMultiModal,

568

569 hidden_states = self.language_model.model(input_ids,

570 positions,

574 - kv_caches,

575 - attn_metadata,

571 intermediate_tensors,

572 inputs_embeds=inputs_embeds)

573 return hidden_states

vllm/model_executor/models/llava_next_video.py

-5

@@ -10,7 +10,6 @@ import torch.nn as nn

10 from transformers import (BatchFeature, LlavaNextVideoConfig,

11 LlavaNextVideoProcessor)

12

13 -from vllm.attention import AttentionMetadata

13 from vllm.config import VllmConfig

14 from vllm.model_executor.layers.activation import get_act_fn

15 from vllm.model_executor.layers.sampler import SamplerOutput, get_sampler

@@ -443,8 +442,6 @@ class LlavaNextVideoForConditionalGeneration(nn.Module, SupportsMultiModal,

442 self,

443 input_ids: torch.Tensor,

444 positions: torch.Tensor,

446 - kv_caches: List[torch.Tensor],

447 - attn_metadata: AttentionMetadata,

445 intermediate_tensors: Optional[IntermediateTensors] = None,

446 inputs_embeds: Optional[torch.Tensor] = None,

447 **kwargs: object,

@@ -468,8 +465,6 @@ class LlavaNextVideoForConditionalGeneration(nn.Module, SupportsMultiModal,

465

466 hidden_states = self.language_model.model(input_ids,

467 positions,

471 - kv_caches,

472 - attn_metadata,

468 intermediate_tensors,

469 inputs_embeds=inputs_embeds)

470

vllm/model_executor/models/llava_onevision.py

-5

@@ -13,7 +13,6 @@ from transformers.models.llava_onevision.modeling_llava_onevision import (

13 get_anyres_image_grid_shape, unpad_image)

14 from typing_extensions import NotRequired

15

16 -from vllm.attention import AttentionMetadata

16 from vllm.config import VllmConfig

17 from vllm.model_executor.layers.activation import get_act_fn

18 from vllm.model_executor.layers.sampler import SamplerOutput, get_sampler

@@ -922,8 +921,6 @@ class LlavaOnevisionForConditionalGeneration(nn.Module, SupportsMultiModal,

921 self,

922 input_ids: torch.Tensor,

923 positions: torch.Tensor,

925 - kv_caches: List[torch.Tensor],

926 - attn_metadata: AttentionMetadata,

924 intermediate_tensors: Optional[IntermediateTensors] = None,

925 inputs_embeds: Optional[torch.Tensor] = None,

926 **kwargs: object,

@@ -955,8 +952,6 @@ class LlavaOnevisionForConditionalGeneration(nn.Module, SupportsMultiModal,

952

953 hidden_states = self.language_model.model(input_ids,

954 positions,

958 - kv_caches,

959 - attn_metadata,

955 intermediate_tensors,

956 inputs_embeds=inputs_embeds)

957

vllm/model_executor/models/mamba.py

+4 -12

@@ -1,12 +1,11 @@

1 # SPDX-License-Identifier: Apache-2.0

2 """PyTorch MAMBA model."""

3 -from typing import Iterable, List, Optional, Set, Tuple

3 +from typing import Iterable, Optional, Set, Tuple

4

5 import torch

6 from torch import nn

7 from transformers import MambaConfig

8

9 -from vllm.attention.backends.abstract import AttentionMetadata

9 from vllm.config import CacheConfig, VllmConfig

10 from vllm.distributed import get_tensor_model_parallel_world_size

11 from vllm.distributed.parallel_state import get_pp_group

@@ -64,7 +63,6 @@ class MambaDecoderLayer(nn.Module):

63 def forward(

64 self,

65 hidden_states: torch.Tensor,

67 - attn_metadata: AttentionMetadata,

66 residual: Optional[torch.Tensor],

67 mamba_cache_params: MambaCacheParams,

68 **kwargs,

@@ -75,8 +73,7 @@ class MambaDecoderLayer(nn.Module):

73 else:

74 hidden_states, residual = self.norm(hidden_states, residual)

75

78 - hidden_states = self.mixer(hidden_states, attn_metadata,

79 - mamba_cache_params)

76 + hidden_states = self.mixer(hidden_states, mamba_cache_params)

77 return hidden_states, residual

78

79

@@ -125,7 +122,6 @@ class MambaModel(nn.Module):

122 self,

123 input_ids: torch.Tensor,

124 positions: torch.Tensor,

128 - attn_metadata: AttentionMetadata,

125 mamba_cache_params: MambaCacheParams,

126 intermediate_tensors: Optional[IntermediateTensors] = None,

127 inputs_embeds: Optional[torch.Tensor] = None,

@@ -146,7 +142,6 @@ class MambaModel(nn.Module):

142 hidden_states, residual = layer(

143 positions=positions,

144 hidden_states=hidden_states,

149 - attn_metadata=attn_metadata,

145 residual=residual,

146 mamba_cache_params=mamba_cache_params.at_layer_idx(

147 i - self.start_layer))

@@ -208,8 +203,6 @@ class MambaForCausalLM(nn.Module, HasInnerState, IsAttentionFree, SupportsPP):

203 def forward(self,

204 input_ids: torch.Tensor,

205 positions: torch.Tensor,

211 - kv_caches: List[KVCache],

212 - attn_metadata: AttentionMetadata,

206 intermediate_tensors: Optional[IntermediateTensors] = None,

207 inputs_embeds: Optional[torch.Tensor] = None,

208 **kwargs):

@@ -222,9 +215,8 @@ class MambaForCausalLM(nn.Module, HasInnerState, IsAttentionFree, SupportsPP):

215

216 mamba_cache_params = self.mamba_cache.current_run_tensors(**kwargs)

217

225 - hidden_states = self.backbone(input_ids, positions, attn_metadata,

226 - mamba_cache_params, intermediate_tensors,

227 - inputs_embeds)

218 + hidden_states = self.backbone(input_ids, positions, mamba_cache_params,

219 + intermediate_tensors, inputs_embeds)

220

221 return hidden_states

222

vllm/model_executor/models/mamba2.py

+7 -11

@@ -1,6 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 """PyTorch MAMBA2 model."""

3 -from typing import Iterable, List, Optional, Set, Tuple

3 +from typing import Iterable, Optional, Set, Tuple

4

5 import torch

6 from torch import nn

@@ -10,6 +10,7 @@ from vllm.attention.backends.abstract import AttentionMetadata

10 from vllm.config import VllmConfig

11 from vllm.distributed import divide, get_tensor_model_parallel_world_size

12 from vllm.distributed.parallel_state import get_pp_group

13 +from vllm.forward_context import get_forward_context

14 from vllm.model_executor.layers.layernorm import RMSNorm

15 from vllm.model_executor.layers.logits_processor import LogitsProcessor

16 from vllm.model_executor.layers.mamba.mamba_mixer2 import (

@@ -63,7 +64,6 @@ class Mamba2DecoderLayer(nn.Module):

64 def forward(

65 self,

66 hidden_states: torch.Tensor,

66 - attn_metadata: AttentionMetadata,

67 residual: Optional[torch.Tensor],

68 mamba_cache_params: MambaCacheParams,

69 sequence_idx: Optional[torch.Tensor],

@@ -75,8 +75,8 @@ class Mamba2DecoderLayer(nn.Module):

75 else:

76 hidden_states, residual = self.norm(hidden_states, residual)

77

78 - hidden_states = self.mixer(hidden_states, attn_metadata,

79 - mamba_cache_params, sequence_idx)

78 + hidden_states = self.mixer(hidden_states, mamba_cache_params,

79 + sequence_idx)

80 return hidden_states, residual

81

82

@@ -122,7 +122,6 @@ class Mamba2Model(nn.Module):

122 self,

123 input_ids: torch.Tensor,

124 positions: torch.Tensor,

125 - attn_metadata: AttentionMetadata,

125 mamba_cache_params: MambaCacheParams,

126 intermediate_tensors: Optional[IntermediateTensors] = None,

127 inputs_embeds: Optional[torch.Tensor] = None,

@@ -142,6 +141,7 @@ class Mamba2Model(nn.Module):

141 # proper continuous batching computation including

142 # chunked prefill

143 seq_idx = None

144 + attn_metadata: AttentionMetadata = get_forward_context().attn_metadata

145 if attn_metadata.num_prefills > 0:

146 seq_idx = torch.zeros_like(input_ids, dtype=torch.int32)

147 for i, (srt, end) in enumerate(

@@ -158,7 +158,6 @@ class Mamba2Model(nn.Module):

158 hidden_states, residual = layer(

159 positions=positions,

160 hidden_states=hidden_states,

161 - attn_metadata=attn_metadata,

161 residual=residual,

162 mamba_cache_params=mamba_cache_params.at_layer_idx(

163 i - self.start_layer),

@@ -224,8 +223,6 @@ class Mamba2ForCausalLM(nn.Module, HasInnerState, IsAttentionFree):

223 def forward(self,

224 input_ids: torch.Tensor,

225 positions: torch.Tensor,

227 - kv_caches: List[KVCache],

228 - attn_metadata: AttentionMetadata,

226 intermediate_tensors: Optional[IntermediateTensors] = None,

227 inputs_embeds: Optional[torch.Tensor] = None,

228 **kwargs):

@@ -238,9 +235,8 @@ class Mamba2ForCausalLM(nn.Module, HasInnerState, IsAttentionFree):

235

236 mamba_cache_params = self.mamba_cache.current_run_tensors(**kwargs)

237

241 - hidden_states = self.backbone(input_ids, positions, attn_metadata,

242 - mamba_cache_params, intermediate_tensors,

243 - inputs_embeds)

238 + hidden_states = self.backbone(input_ids, positions, mamba_cache_params,

239 + intermediate_tensors, inputs_embeds)

240

241 return hidden_states

242

vllm/model_executor/models/minicpm.py

+5 -19

@@ -23,13 +23,13 @@

23 # limitations under the License.

24 """Inference-only MiniCPM model compatible with HuggingFace weights."""

25 import math

26 -from typing import Any, Dict, Iterable, List, Optional, Set, Tuple, Union

26 +from typing import Any, Dict, Iterable, Optional, Set, Tuple, Union

27

28 import torch

29 from torch import nn

30 from transformers import PretrainedConfig

31

32 -from vllm.attention import Attention, AttentionMetadata

32 +from vllm.attention import Attention

33 from vllm.compilation.decorators import support_torch_compile

34 from vllm.config import CacheConfig, VllmConfig

35 from vllm.distributed import (get_pp_group, get_tensor_model_parallel_rank,

@@ -257,8 +257,6 @@ class MiniCPMAttention(nn.Module):

257 self,

258 positions: torch.Tensor,

259 hidden_states: torch.Tensor,

260 - kv_cache: torch.Tensor,

261 - attn_metadata: AttentionMetadata,

260 ) -> torch.Tensor:

261 qkv, _ = self.qkv_proj(hidden_states)

262 q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)

@@ -266,7 +264,7 @@ class MiniCPMAttention(nn.Module):

264 q, k = q.float(), k.float()

265 q, k = self.rotary_emb(positions, q, k)

266 q, k = q.to(orig_dtype), k.to(orig_dtype)

269 - attn_output = self.attn(q, k, v, kv_cache, attn_metadata)

267 + attn_output = self.attn(q, k, v)

268 output, _ = self.o_proj(attn_output)

269 return output

270

@@ -331,8 +329,6 @@ class MiniCPMDecoderLayer(nn.Module):

329 self,

330 positions: torch.Tensor,

331 hidden_states: torch.Tensor,

334 - kv_cache: torch.Tensor,

335 - attn_metadata: AttentionMetadata,

332 residual: Optional[torch.Tensor],

333 ) -> Tuple[torch.Tensor, torch.Tensor]:

334 # Self Attention

@@ -341,8 +337,6 @@ class MiniCPMDecoderLayer(nn.Module):

337 hidden_states = self.self_attn(

338 positions=positions,

339 hidden_states=hidden_states,

344 - kv_cache=kv_cache,

345 - attn_metadata=attn_metadata,

340 )

341 hidden_states = residual + hidden_states * \

342 (self.config.scale_depth / math.sqrt(self.config.num_hidden_layers))

@@ -409,8 +403,6 @@ class MiniCPMModel(nn.Module):

403 self,

404 input_ids: torch.Tensor,

405 positions: torch.Tensor,

412 - kv_caches: List[torch.Tensor],

413 - attn_metadata: AttentionMetadata,

406 intermediate_tensors: Optional[IntermediateTensors] = None,

407 inputs_embeds: Optional[torch.Tensor] = None,

408 ) -> Union[torch.Tensor, IntermediateTensors]:

@@ -424,13 +416,10 @@ class MiniCPMModel(nn.Module):

416 hidden_states = intermediate_tensors["hidden_states"]

417 residual = intermediate_tensors["residual"]

418

427 - for i in range(self.start_layer, self.end_layer):

428 - layer = self.layers[i]

419 + for layer in self.layers[self.start_layer:self.end_layer]:

420 hidden_states, residual = layer(

421 positions,

422 hidden_states,

432 - kv_caches[i - self.start_layer],

433 - attn_metadata,

423 residual,

424 )

425 if not get_pp_group().is_last_rank:

@@ -579,13 +568,10 @@ class MiniCPMForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

568 self,

569 input_ids: torch.Tensor,

570 positions: torch.Tensor,

582 - kv_caches: List[torch.Tensor],

583 - attn_metadata: AttentionMetadata,

571 intermediate_tensors: Optional[IntermediateTensors] = None,

572 inputs_embeds: Optional[torch.Tensor] = None,

573 ) -> Union[torch.Tensor, IntermediateTensors]:

587 - hidden_states = self.model(input_ids, positions, kv_caches,

588 - attn_metadata, intermediate_tensors,

574 + hidden_states = self.model(input_ids, positions, intermediate_tensors,

575 inputs_embeds)

576 return hidden_states

577

vllm/model_executor/models/minicpm3.py

+2 -4

@@ -29,7 +29,7 @@ import torch

29 from torch import nn

30 from transformers import PretrainedConfig

31

32 -from vllm.attention import Attention, AttentionMetadata

32 +from vllm.attention import Attention

33 from vllm.config import CacheConfig, VllmConfig

34 from vllm.distributed import get_tensor_model_parallel_world_size

35 from vllm.model_executor.layers.layernorm import RMSNorm

@@ -129,8 +129,6 @@ class MiniCPM3Attention(nn.Module):

129 self,

130 positions: torch.Tensor,

131 hidden_states: torch.Tensor,

132 - kv_cache: torch.Tensor,

133 - attn_metadata: AttentionMetadata,

132 ) -> torch.Tensor:

133 q, _ = self.q_a_proj(hidden_states)

134 q = self.q_a_layernorm(q)

@@ -170,7 +168,7 @@ class MiniCPM3Attention(nn.Module):

168 v, [0, self.qk_head_dim - self.v_head_dim],

169 value=0).view(-1, self.num_local_heads * self.qk_head_dim)

170

173 - attn_output = self.attn(q, k, v, kv_cache, attn_metadata)

171 + attn_output = self.attn(q, k, v)

172 attn_output = attn_output.view(

173 -1, self.num_local_heads,

174 self.qk_head_dim)[..., :self.v_head_dim].reshape(

vllm/model_executor/models/minicpmo.py

-5

@@ -33,7 +33,6 @@ from transformers.modeling_outputs import BaseModelOutputWithPast

33 from transformers.models.whisper.modeling_whisper import (

34 ACT2FN, WHISPER_ATTENTION_CLASSES, WhisperConfig, WhisperEncoder)

35

36 -from vllm.attention import AttentionMetadata

36 from vllm.config import VllmConfig

37 from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalKwargs

38 from vllm.multimodal.inputs import MultiModalFieldConfig

@@ -792,8 +791,6 @@ class MiniCPMO(MiniCPMV2_6):

791 self,

792 input_ids: torch.Tensor,

793 positions: torch.Tensor,

795 - kv_caches: List[torch.Tensor],

796 - attn_metadata: AttentionMetadata,

794 intermediate_tensors: Optional[IntermediateTensors] = None,

795 **kwargs: Any,

796 ) -> torch.Tensor:

@@ -818,8 +815,6 @@ class MiniCPMO(MiniCPMV2_6):

815 output = self.llm.model(

816 input_ids=input_ids,

817 positions=positions,

821 - kv_caches=kv_caches,

822 - attn_metadata=attn_metadata,

818 intermediate_tensors=intermediate_tensors,

819 inputs_embeds=vlm_embeddings,

820 )

vllm/model_executor/models/minicpmv.py

-5

@@ -37,7 +37,6 @@ from torch import nn

37 from transformers import BatchFeature, PretrainedConfig

38 from typing_extensions import TypeVar

39

40 -from vllm.attention import AttentionMetadata

40 from vllm.config import VllmConfig

41 from vllm.model_executor.layers.quantization import QuantizationConfig

42 from vllm.model_executor.layers.resampler import (BaseResampler, Resampler2,

@@ -1030,8 +1029,6 @@ class MiniCPMVBaseModel(nn.Module, SupportsMultiModal, SupportsPP):

1029 self,

1030 input_ids: torch.Tensor,

1031 positions: torch.Tensor,

1033 - kv_caches: List[torch.Tensor],

1034 - attn_metadata: AttentionMetadata,

1032 intermediate_tensors: Optional[IntermediateTensors] = None,

1033 **kwargs: Any,

1034 ) -> torch.Tensor:

@@ -1051,8 +1048,6 @@ class MiniCPMVBaseModel(nn.Module, SupportsMultiModal, SupportsPP):

1048 output = self.llm.model(

1049 input_ids=input_ids,

1050 positions=positions,

1054 - kv_caches=kv_caches,

1055 - attn_metadata=attn_metadata,

1051 intermediate_tensors=intermediate_tensors,

1052 inputs_embeds=vlm_embeddings,

1053 )

vllm/model_executor/models/mixtral.py

+6 -20

@@ -22,13 +22,13 @@

22 # See the License for the specific language governing permissions and

23 # limitations under the License.

24 """Inference-only Mixtral model."""

25 -from typing import Iterable, List, Optional, Set, Tuple, Union

25 +from typing import Iterable, Optional, Set, Tuple, Union

26

27 import torch

28 from torch import nn

29 from transformers import MixtralConfig

30

31 -from vllm.attention import Attention, AttentionMetadata

31 +from vllm.attention import Attention

32 from vllm.compilation.decorators import support_torch_compile

33 from vllm.config import CacheConfig, VllmConfig

34 from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size

@@ -175,13 +175,11 @@ class MixtralAttention(nn.Module):

175 self,

176 positions: torch.Tensor,

177 hidden_states: torch.Tensor,

178 - kv_cache: torch.Tensor,

179 - attn_metadata: AttentionMetadata,

178 ) -> torch.Tensor:

179 qkv, _ = self.qkv_proj(hidden_states)

180 q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)

181 q, k = self.rotary_emb(positions, q, k)

184 - attn_output = self.attn(q, k, v, kv_cache, attn_metadata)

182 + attn_output = self.attn(q, k, v)

183 output, _ = self.o_proj(attn_output)

184 return output

185

@@ -224,8 +222,6 @@ class MixtralDecoderLayer(nn.Module):

222 self,

223 positions: torch.Tensor,

224 hidden_states: torch.Tensor,

227 - kv_cache: torch.Tensor,

228 - attn_metadata: AttentionMetadata,

225 residual: Optional[torch.Tensor],

226 ) -> torch.Tensor:

227 # Self Attention

@@ -238,8 +234,6 @@ class MixtralDecoderLayer(nn.Module):

234 hidden_states = self.self_attn(

235 positions=positions,

236 hidden_states=hidden_states,

241 - kv_cache=kv_cache,

242 - attn_metadata=attn_metadata,

237 )

238

239 # Fully Connected

@@ -291,8 +285,6 @@ class MixtralModel(nn.Module):

285 self,

286 input_ids: torch.Tensor,

287 positions: torch.Tensor,

294 - kv_caches: List[torch.Tensor],

295 - attn_metadata: AttentionMetadata,

288 intermediate_tensors: Optional[IntermediateTensors],

289 inputs_embeds: Optional[torch.Tensor] = None,

290 ) -> Union[torch.Tensor, IntermediateTensors]:

@@ -306,11 +298,8 @@ class MixtralModel(nn.Module):

298 assert intermediate_tensors is not None

299 hidden_states = intermediate_tensors["hidden_states"]

300 residual = intermediate_tensors["residual"]

309 - for i in range(self.start_layer, self.end_layer):

310 - layer = self.layers[i]

311 - hidden_states, residual = layer(positions, hidden_states,

312 - kv_caches[i - self.start_layer],

313 - attn_metadata, residual)

301 + for layer in self.layers[self.start_layer:self.end_layer]:

302 + hidden_states, residual = layer(positions, hidden_states, residual)

303 if not get_pp_group().is_last_rank:

304 return IntermediateTensors({

305 "hidden_states": hidden_states,

@@ -377,13 +366,10 @@ class MixtralForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

366 self,

367 input_ids: torch.Tensor,

368 positions: torch.Tensor,

380 - kv_caches: List[torch.Tensor],

381 - attn_metadata: AttentionMetadata,

369 intermediate_tensors: Optional[IntermediateTensors] = None,

370 inputs_embeds: Optional[torch.Tensor] = None,

371 ) -> Union[torch.Tensor, IntermediateTensors]:

385 - hidden_states = self.model(input_ids, positions, kv_caches,

386 - attn_metadata, intermediate_tensors,

372 + hidden_states = self.model(input_ids, positions, intermediate_tensors,

373 inputs_embeds)

374 return hidden_states

375

vllm/model_executor/models/mixtral_quant.py

+6 -20

@@ -22,7 +22,7 @@

22 # See the License for the specific language governing permissions and

23 # limitations under the License.

24 """Inference-only Mixtral model."""

25 -from typing import Iterable, List, Optional, Set, Tuple, Union

25 +from typing import Iterable, Optional, Set, Tuple, Union

26

27 import numpy as np

28 import torch

@@ -30,7 +30,7 @@ import torch.nn.functional as F

30 from torch import nn

31 from transformers import MixtralConfig

32

33 -from vllm.attention import Attention, AttentionMetadata

33 +from vllm.attention import Attention

34 from vllm.config import CacheConfig, VllmConfig

35 from vllm.distributed import (get_pp_group, get_tensor_model_parallel_rank,

36 get_tensor_model_parallel_world_size,

@@ -229,13 +229,11 @@ class MixtralAttention(nn.Module):

229 self,

230 positions: torch.Tensor,

231 hidden_states: torch.Tensor,

232 - kv_cache: torch.Tensor,

233 - attn_metadata: AttentionMetadata,

232 ) -> torch.Tensor:

233 qkv, _ = self.qkv_proj(hidden_states)

234 q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)

235 q, k = self.rotary_emb(positions, q, k)

238 - attn_output = self.attn(q, k, v, kv_cache, attn_metadata)

236 + attn_output = self.attn(q, k, v)

237 output, _ = self.o_proj(attn_output)

238 return output

239

@@ -274,8 +272,6 @@ class MixtralDecoderLayer(nn.Module):

272 self,

273 positions: torch.Tensor,

274 hidden_states: torch.Tensor,

277 - kv_cache: torch.Tensor,

278 - attn_metadata: AttentionMetadata,

275 residual: Optional[torch.Tensor],

276 ) -> torch.Tensor:

277 # Self Attention

@@ -288,8 +284,6 @@ class MixtralDecoderLayer(nn.Module):

284 hidden_states = self.self_attn(

285 positions=positions,

286 hidden_states=hidden_states,

291 - kv_cache=kv_cache,

292 - attn_metadata=attn_metadata,

287 )

288

289 # Fully Connected

@@ -333,8 +327,6 @@ class MixtralModel(nn.Module):

327 self,

328 input_ids: torch.Tensor,

329 positions: torch.Tensor,

336 - kv_caches: List[torch.Tensor],

337 - attn_metadata: AttentionMetadata,

330 intermediate_tensors: Optional[IntermediateTensors],

331 inputs_embeds: Optional[torch.Tensor] = None,

332 ) -> Union[torch.Tensor, IntermediateTensors]:

@@ -348,11 +340,8 @@ class MixtralModel(nn.Module):

340 assert intermediate_tensors is not None

341 hidden_states = intermediate_tensors["hidden_states"]

342 residual = intermediate_tensors["residual"]

351 - for i in range(self.start_layer, self.end_layer):

352 - layer = self.layers[i]

353 - hidden_states, residual = layer(positions, hidden_states,

354 - kv_caches[i - self.start_layer],

355 - attn_metadata, residual)

343 + for layer in self.layers[self.start_layer:self.end_layer]:

344 + hidden_states, residual = layer(positions, hidden_states, residual)

345 if not get_pp_group().is_last_rank:

346 return IntermediateTensors({

347 "hidden_states": hidden_states,

@@ -390,13 +379,10 @@ class MixtralForCausalLM(nn.Module, SupportsPP):

379 self,

380 input_ids: torch.Tensor,

381 positions: torch.Tensor,

393 - kv_caches: List[torch.Tensor],

394 - attn_metadata: AttentionMetadata,

382 intermediate_tensors: Optional[IntermediateTensors] = None,

383 inputs_embeds: Optional[torch.Tensor] = None,

384 ) -> Union[torch.Tensor, IntermediateTensors]:

398 - hidden_states = self.model(input_ids, positions, kv_caches,

399 - attn_metadata, intermediate_tensors,

385 + hidden_states = self.model(input_ids, positions, intermediate_tensors,

386 inputs_embeds)

387 return hidden_states

388

vllm/model_executor/models/mllama.py

+15 -35

@@ -38,7 +38,8 @@ from vllm.attention import Attention, AttentionMetadata, AttentionType

38 from vllm.attention.ops.paged_attn import PagedAttention

39 from vllm.attention.selector import _Backend

40 from vllm.config import VllmConfig

41 -from vllm.distributed import get_tensor_model_parallel_world_size

41 +from vllm.distributed import get_pp_group, get_tp_group

42 +from vllm.forward_context import get_forward_context

43 from vllm.logger import init_logger

44 from vllm.model_executor.layers.layernorm import RMSNorm

45 from vllm.model_executor.layers.linear import (ColumnParallelLinear,

@@ -416,11 +417,11 @@ class MllamaVisionSdpaAttention(nn.Module):

417 prefix: str = ""):

418 super().__init__()

419

419 - model_parallel_size = get_tensor_model_parallel_world_size()

420 + tensor_parallel_size = get_tp_group().world_size

421 self.embed_dim = config.hidden_size

422 self.num_heads = config.attention_heads

423 self.head_dim = config.hidden_size // config.attention_heads

423 - self.num_local_heads = self.num_heads // model_parallel_size

424 + self.num_local_heads = self.num_heads // tensor_parallel_size

425 self.q_size = self.num_local_heads * self.head_dim

426 self.kv_size = self.num_local_heads * self.head_dim

427

@@ -771,12 +772,13 @@ class MllamaTextCrossAttention(nn.Module):

772 ):

773 super().__init__()

774 self.config = config

774 - self.model_parallel_size = get_tensor_model_parallel_world_size()

775 + self.pipeline_parallel_rank = get_pp_group().rank_in_group

776 + self.tensor_parallel_size = get_tp_group().world_size

777 self.num_heads = self.config.num_attention_heads

776 - self.num_local_heads = self.num_heads // self.model_parallel_size

778 + self.num_local_heads = self.num_heads // self.tensor_parallel_size

779 self.num_key_value_heads = self.config.num_key_value_heads

780 self.num_local_key_value_heads = \

779 - self.num_key_value_heads // self.model_parallel_size

781 + self.num_key_value_heads // self.tensor_parallel_size

782 self.dropout = config.dropout

783 self.hidden_size = config.hidden_size

784 self.head_dim = config.hidden_size // self.num_heads

@@ -824,8 +826,6 @@ class MllamaTextCrossAttention(nn.Module):

826 attention_mask: Optional[torch.Tensor],

827 kv_range_for_decode: Optional[List[Tuple[int, int]]],

828 cross_attention_states: Optional[torch.Tensor],

827 - kv_cache: torch.Tensor,

828 - attn_metadata: AttentionMetadata,

829 ) -> torch.Tensor:

830 qkv_dec, _ = self.qkv_proj(hidden_states)

831 q, _, _ = qkv_dec.split(

@@ -846,14 +846,11 @@ class MllamaTextCrossAttention(nn.Module):

846 q = self.q_norm(q)

847

848 if attention_mask is not None:

849 - output = self._attention_with_mask(q, k, v, kv_cache,

850 - attention_mask,

851 - kv_range_for_decode,

852 - attn_metadata)

849 + output = self._attention_with_mask(q, k, v, attention_mask,

850 + kv_range_for_decode)

851 else:

852 output = self.attn(

855 - q.view(-1, self.num_local_heads * self.head_dim), k, v,

856 - kv_cache, attn_metadata)

853 + q.view(-1, self.num_local_heads * self.head_dim), k, v)

854 out, _ = self.o_proj(output)

855 return out

856

@@ -862,11 +859,11 @@ class MllamaTextCrossAttention(nn.Module):

859 q: torch.Tensor,

860 k: torch.Tensor,

861 v: torch.Tensor,

865 - kv_cache: torch.Tensor,

862 attention_mask: torch.Tensor,

863 kv_range_for_decode: List[Tuple[int, int]],

868 - attn_metadata: AttentionMetadata,

864 ) -> torch.Tensor:

865 + kv_cache = self.attn.kv_cache[self.pipeline_parallel_rank]

866 + attn_metadata: AttentionMetadata = get_forward_context().attn_metadata

867 # Skip writing kv-cache for the initial profiling run.

868 if len(kv_cache.shape) > 1:

869 i = torch.ones(1, dtype=torch.float32)

@@ -978,8 +975,6 @@ class MllamaCrossAttentionDecoderLayer(torch.nn.Module):

975 cross_attention_mask: torch.Tensor,

976 kv_range_for_decode: Optional[List[Tuple[int, int]]],

977 full_text_row_masked_out_mask: torch.Tensor,

981 - kv_cache: List[torch.Tensor],

982 - attn_metadata: AttentionMetadata,

978 ) -> torch.Tensor:

979 residual = hidden_states

980 hidden_states = self.input_layernorm(hidden_states)

@@ -989,8 +984,6 @@ class MllamaCrossAttentionDecoderLayer(torch.nn.Module):

984 attention_mask=cross_attention_mask,

985 kv_range_for_decode=kv_range_for_decode,

986 cross_attention_states=cross_attention_states,

992 - kv_cache=kv_cache,

993 - attn_metadata=attn_metadata,

987 )

988 hidden_states = full_text_row_masked_out_mask * hidden_states

989 hidden_states = residual + self.cross_attn_attn_gate.tanh(

@@ -1054,14 +1047,12 @@ class MllamaTextModel(nn.Module):

1047 kv_range_for_decode: Optional[List[Tuple[int, int]]],

1048 full_text_row_masked_out_mask: Optional[Tuple[torch.Tensor,

1049 torch.Tensor]],

1057 - kv_caches: List[torch.Tensor],

1058 - attn_metadata: AttentionMetadata,

1050 skip_cross_attention: bool,

1051 ) -> torch.Tensor:

1052 inputs_embeds = self.embed_tokens(input_ids)

1053 hidden_states = inputs_embeds

1054

1064 - for idx, decoder_layer in enumerate(self.layers):

1055 + for decoder_layer in self.layers:

1056 if isinstance(decoder_layer, MllamaCrossAttentionDecoderLayer):

1057 if not skip_cross_attention:

1058 hidden_states = decoder_layer(

@@ -1071,15 +1062,11 @@ class MllamaTextModel(nn.Module):

1062 kv_range_for_decode=kv_range_for_decode,

1063 full_text_row_masked_out_mask=

1064 full_text_row_masked_out_mask,

1074 - kv_cache=kv_caches[idx],

1075 - attn_metadata=attn_metadata,

1065 )

1066 elif isinstance(decoder_layer, LlamaDecoderLayer):

1067 hidden_states, residual = decoder_layer(

1068 positions=positions,

1069 hidden_states=hidden_states,

1081 - kv_cache=kv_caches[idx],

1082 - attn_metadata=attn_metadata,

1070 residual=None,

1071 )

1072 hidden_states = hidden_states + residual

@@ -1124,8 +1111,6 @@ class MllamaForCausalLM(nn.Module):

1111 kv_range_for_decode: Optional[List[Tuple[int, int]]],

1112 full_text_row_masked_out_mask: Optional[Tuple[torch.Tensor,

1113 torch.Tensor]],

1127 - kv_caches: List[torch.Tensor],

1128 - attn_metadata: AttentionMetadata,

1114 skip_cross_attention: bool,

1115 ) -> torch.Tensor:

1116 hidden_states = self.model(

@@ -1135,8 +1120,6 @@ class MllamaForCausalLM(nn.Module):

1120 cross_attention_mask=cross_attention_mask,

1121 kv_range_for_decode=kv_range_for_decode,

1122 full_text_row_masked_out_mask=full_text_row_masked_out_mask,

1138 - kv_caches=kv_caches,

1139 - attn_metadata=attn_metadata,

1123 skip_cross_attention=skip_cross_attention,

1124 )

1125 return hidden_states

@@ -1353,10 +1336,9 @@ class MllamaForConditionalGeneration(nn.Module, SupportsMultiModal):

1336 self,

1337 input_ids: torch.Tensor,

1338 positions: torch.Tensor,

1356 - kv_caches: List[torch.Tensor],

1357 - attn_metadata: AttentionMetadata,

1339 **kwargs: object,

1340 ) -> Union[Tuple, CausalLMOutputWithPast]:

1341 + attn_metadata = get_forward_context().attn_metadata

1342 if attn_metadata.num_prefill_tokens > 0 and \

1343 attn_metadata.num_decode_tokens > 0:

1344 raise ValueError("Chunk prefill not supported")

@@ -1410,8 +1392,6 @@ class MllamaForConditionalGeneration(nn.Module, SupportsMultiModal):

1392 cross_attention_mask=cross_attention_mask,

1393 kv_range_for_decode=kv_range_for_decode,

1394 full_text_row_masked_out_mask=full_text_row_masked_out_mask,

1413 - kv_caches=kv_caches,

1414 - attn_metadata=attn_metadata,

1395 skip_cross_attention=skip_cross_attention,

1396 )

1397

vllm/model_executor/models/molmo.py

+3 -22

@@ -16,7 +16,7 @@ from transformers import (BatchFeature, PretrainedConfig, ProcessorMixin,

16 from transformers.image_utils import ImageInput

17 from transformers.tokenization_utils_base import TextInput

18

19 -from vllm.attention import Attention, AttentionMetadata

19 +from vllm.attention import Attention

20 from vllm.attention.layer import MultiHeadAttention

21 from vllm.compilation.decorators import support_torch_compile

22 from vllm.config import CacheConfig, VllmConfig

@@ -460,15 +460,13 @@ class MolmoAttention(nn.Module):

460 self,

461 positions: torch.Tensor,

462 hidden_states: torch.Tensor,

463 - kv_cache: torch.Tensor,

464 - attn_metadata: AttentionMetadata,

463 ) -> torch.Tensor:

464 qkv, _ = self.qkv_proj(hidden_states)

465 q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)

466 if self.q_norm is not None and self.k_norm is not None:

467 q, k = self._apply_qk_norm(q, k)

468 q, k = self.rotary_emb(positions, q, k)

471 - attn_output = self.attn(q, k, v, kv_cache, attn_metadata)

469 + attn_output = self.attn(q, k, v)

470 output, _ = self.o_proj(attn_output)

471 return output

472

@@ -580,8 +578,6 @@ class MolmoDecoderLayer(nn.Module):

578 self,

579 positions: torch.Tensor,

580 hidden_states: torch.Tensor,

583 - kv_cache: torch.Tensor,

584 - attn_metadata: AttentionMetadata,

581 residual: Optional[torch.Tensor],

582 ) -> Tuple[torch.Tensor, Optional[Tuple[torch.Tensor, torch.Tensor]]]:

583 # Self Attention

@@ -594,8 +590,6 @@ class MolmoDecoderLayer(nn.Module):

590 hidden_states = self.self_attn(

591 positions=positions,

592 hidden_states=hidden_states,

597 - kv_cache=kv_cache,

598 - attn_metadata=attn_metadata,

593 )

594

595 hidden_states, residual = self.post_attention_layernorm(

@@ -610,8 +604,6 @@ class MolmoDecoderNormAfterLayer(MolmoDecoderLayer):

604 self,

605 positions: torch.Tensor,

606 hidden_states: torch.Tensor,

613 - kv_cache: torch.Tensor,

614 - attn_metadata: AttentionMetadata,

607 residual: Optional[torch.Tensor],

608 ) -> Tuple[torch.Tensor, Optional[Tuple[torch.Tensor, torch.Tensor]]]:

609 # Self Attention

@@ -619,8 +611,6 @@ class MolmoDecoderNormAfterLayer(MolmoDecoderLayer):

611 hidden_states = self.self_attn(

612 positions=positions,

613 hidden_states=hidden_states,

622 - kv_cache=kv_cache,

623 - attn_metadata=attn_metadata,

614 )

615

616 hidden_states = self.input_layernorm(hidden_states)

@@ -841,8 +831,6 @@ class MolmoModel(nn.Module, SupportsQuant):

831 self,

832 input_ids: torch.Tensor,

833 positions: torch.Tensor,

844 - kv_caches: List[torch.Tensor],

845 - attn_metadata: AttentionMetadata,

834 intermediate_tensors: Optional[IntermediateTensors] = None,

835 inputs_embeds: Optional[torch.Tensor] = None,

836 ) -> torch.Tensor:

@@ -858,13 +846,10 @@ class MolmoModel(nn.Module, SupportsQuant):

846 residual = intermediate_tensors["residual"]

847

848 # Apply blocks one-by-one.

861 - for i in range(self.start_layer, self.end_layer):

862 - layer = self.layers[i]

849 + for layer in self.layers[self.start_layer:self.end_layer]:

850 hidden_states, residual = layer(

851 positions,

852 hidden_states,

866 - kv_caches[i - self.start_layer],

867 - attn_metadata,

853 residual,

854 )

855 if not get_pp_group().is_last_rank:

@@ -1643,8 +1628,6 @@ class MolmoForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA,

1628 self,

1629 input_ids: torch.LongTensor,

1630 positions: torch.LongTensor,

1646 - kv_caches: List[torch.Tensor],

1647 - attn_metadata: AttentionMetadata,

1631 intermediate_tensors: Optional[IntermediateTensors] = None,

1632 inputs_embeds: Optional[torch.Tensor] = None,

1633 **kwargs: object,

@@ -1663,8 +1646,6 @@ class MolmoForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA,

1646

1647 hidden_states = self.model(input_ids,

1648 positions,

1666 - kv_caches,

1667 - attn_metadata,

1649 intermediate_tensors,

1650 inputs_embeds=inputs_embeds)

1651

vllm/model_executor/models/mpt.py

+7 -24

@@ -2,12 +2,12 @@

2

3 # Adapted from https://huggingface.co/mosaicml/mpt-7b/tree/main

4 import math

5 -from typing import Iterable, List, Optional, Set, Tuple, Union

5 +from typing import Iterable, Optional, Set, Tuple, Union

6

7 import torch

8 import torch.nn as nn

9

10 -from vllm.attention import Attention, AttentionMetadata

10 +from vllm.attention import Attention

11 from vllm.compilation.decorators import support_torch_compile

12 from vllm.config import CacheConfig, VllmConfig

13 from vllm.distributed import (get_pp_group, get_tensor_model_parallel_rank,

@@ -125,8 +125,6 @@ class MPTAttention(nn.Module):

125 self,

126 position_ids: torch.Tensor,

127 hidden_states: torch.Tensor,

128 - kv_cache: torch.Tensor,

129 - attn_metadata: AttentionMetadata,

128 ) -> torch.Tensor:

129 del position_ids # unused.

130 qkv, _ = self.Wqkv(hidden_states)

@@ -136,7 +134,7 @@ class MPTAttention(nn.Module):

134 if self.qk_ln:

135 q = self.q_ln(q)

136 k = self.k_ln(k)

139 - attn_output = self.attn(q, k, v, kv_cache, attn_metadata)

137 + attn_output = self.attn(q, k, v)

138 output, _ = self.out_proj(attn_output)

139 return output

140

@@ -196,15 +194,11 @@ class MPTBlock(nn.Module):

194 self,

195 position_ids: torch.Tensor,

196 hidden_states: torch.Tensor,

199 - kv_cache: torch.Tensor,

200 - attn_metadata: AttentionMetadata,

197 ) -> torch.Tensor:

198 x = self.norm_1(hidden_states)

199 x = self.attn(

200 position_ids=position_ids,

201 hidden_states=x,

206 - kv_cache=kv_cache,

207 - attn_metadata=attn_metadata,

202 )

203 hidden_states = hidden_states + x

204 x = self.norm_2(hidden_states)

@@ -253,8 +247,6 @@ class MPTModel(nn.Module):

247 self,

248 input_ids: torch.Tensor,

249 position_ids: torch.Tensor,

256 - kv_caches: List[torch.Tensor],

257 - attn_metadata: AttentionMetadata,

250 intermediate_tensors: Optional[IntermediateTensors],

251 inputs_embeds: Optional[torch.Tensor] = None,

252 ) -> Union[torch.Tensor, IntermediateTensors]:

@@ -267,14 +259,8 @@ class MPTModel(nn.Module):

259 assert intermediate_tensors is not None

260 hidden_states = intermediate_tensors["hidden_states"]

261

270 - for i in range(self.start_layer, self.end_layer):

271 - block = self.blocks[i]

272 - hidden_states = block(

273 - position_ids,

274 - hidden_states,

275 - kv_caches[i - self.start_layer],

276 - attn_metadata,

277 - )

262 + for block in self.blocks[self.start_layer:self.end_layer]:

263 + hidden_states = block(position_ids, hidden_states)

264 if not get_pp_group().is_last_rank:

265 return IntermediateTensors({"hidden_states": hidden_states})

266 hidden_states = self.norm_f(hidden_states)

@@ -306,14 +292,11 @@ class MPTForCausalLM(nn.Module, SupportsPP):

292 self,

293 input_ids: torch.Tensor,

294 positions: torch.Tensor,

309 - kv_caches: List[torch.Tensor],

310 - attn_metadata: AttentionMetadata,

295 intermediate_tensors: Optional[IntermediateTensors] = None,

296 inputs_embeds: Optional[torch.Tensor] = None,

297 ) -> Union[torch.Tensor, IntermediateTensors]:

314 - hidden_states = self.transformer(input_ids, positions, kv_caches,

315 - attn_metadata, intermediate_tensors,

316 - inputs_embeds)

298 + hidden_states = self.transformer(input_ids, positions,

299 + intermediate_tensors, inputs_embeds)

300 return hidden_states

301

302 def compute_logits(

vllm/model_executor/models/nemotron.py

+5 -23

@@ -27,7 +27,7 @@ from typing import Any, Dict, Iterable, List, Optional, Set, Tuple, Union

27 import torch

28 from torch import nn

29

30 -from vllm.attention import Attention, AttentionMetadata

30 +from vllm.attention import Attention

31 from vllm.compilation.decorators import support_torch_compile

32 from vllm.config import CacheConfig, VllmConfig

33 from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size

@@ -204,13 +204,11 @@ class NemotronAttention(nn.Module):

204 self,

205 positions: torch.Tensor,

206 hidden_states: torch.Tensor,

207 - kv_cache: torch.Tensor,

208 - attn_metadata: AttentionMetadata,

207 ) -> torch.Tensor:

208 qkv, _ = self.qkv_proj(hidden_states)

209 q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)

210 q, k = self.rotary_emb(positions, q, k)

213 - attn_output = self.attn(q, k, v, kv_cache, attn_metadata)

211 + attn_output = self.attn(q, k, v)

212 output, _ = self.o_proj(attn_output)

213 return output

214

@@ -269,8 +267,6 @@ class NemotronDecoderLayer(nn.Module):

267 self,

268 positions: torch.Tensor,

269 hidden_states: torch.Tensor,

272 - kv_cache: torch.Tensor,

273 - attn_metadata: AttentionMetadata,

270 residual: Optional[torch.Tensor],

271 ) -> Tuple[torch.Tensor, torch.Tensor]:

272 # Self Attention

@@ -283,8 +279,6 @@ class NemotronDecoderLayer(nn.Module):

279 hidden_states = self.self_attn(

280 positions=positions,

281 hidden_states=hidden_states,

286 - kv_cache=kv_cache,

287 - attn_metadata=attn_metadata,

282 )

283

284 # Fully Connected

@@ -343,8 +337,6 @@ class NemotronModel(nn.Module):

337 self,

338 input_ids: Optional[torch.Tensor],

339 positions: torch.Tensor,

346 - kv_caches: List[torch.Tensor],

347 - attn_metadata: AttentionMetadata,

340 intermediate_tensors: Optional[IntermediateTensors],

341 inputs_embeds: Optional[torch.Tensor] = None,

342 ) -> Union[torch.Tensor, IntermediateTensors]:

@@ -359,15 +351,8 @@ class NemotronModel(nn.Module):

351 hidden_states = intermediate_tensors["hidden_states"]

352 residual = intermediate_tensors["residual"]

353

362 - for i in range(self.start_layer, self.end_layer):

363 - layer = self.layers[i]

364 - hidden_states, residual = layer(

365 - positions,

366 - hidden_states,

367 - kv_caches[i - self.start_layer],

368 - attn_metadata,

369 - residual,

370 - )

354 + for layer in self.layers[self.start_layer:self.end_layer]:

355 + hidden_states, residual = layer(positions, hidden_states, residual)

356

357 if not get_pp_group().is_last_rank:

358 return IntermediateTensors({

@@ -444,13 +429,10 @@ class NemotronForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

429 self,

430 input_ids: torch.Tensor,

431 positions: torch.Tensor,

447 - kv_caches: List[torch.Tensor],

448 - attn_metadata: AttentionMetadata,

432 intermediate_tensors: Optional[IntermediateTensors] = None,

433 inputs_embeds: Optional[torch.Tensor] = None,

434 ) -> Union[torch.Tensor, IntermediateTensors]:

452 - model_output = self.model(input_ids, positions, kv_caches,

453 - attn_metadata, intermediate_tensors,

435 + model_output = self.model(input_ids, positions, intermediate_tensors,

436 inputs_embeds)

437 return model_output

438

vllm/model_executor/models/olmo.py

+6 -22

@@ -22,13 +22,13 @@

22 # See the License for the specific language governing permissions and

23 # limitations under the License.

24 """Inference-only OLMo model compatible with HuggingFace weights."""

25 -from typing import Iterable, List, Optional, Set, Tuple, Union

25 +from typing import Iterable, Optional, Set, Tuple, Union

26

27 import torch

28 from torch import nn

29 from transformers import OlmoConfig

30

31 -from vllm.attention import Attention, AttentionMetadata

31 +from vllm.attention import Attention

32 from vllm.compilation.decorators import support_torch_compile

33 from vllm.config import CacheConfig, VllmConfig

34 from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size

@@ -119,15 +119,13 @@ class OlmoAttention(nn.Module):

119 self,

120 positions: torch.Tensor,

121 hidden_states: torch.Tensor,

122 - kv_cache: torch.Tensor,

123 - attn_metadata: AttentionMetadata,

122 ) -> torch.Tensor:

123 qkv, _ = self.qkv_proj(hidden_states)

124 if self.clip_qkv is not None:

125 qkv.clamp_(min=-self.clip_qkv, max=self.clip_qkv)

126 q, k, v = qkv.chunk(chunks=3, dim=-1)

127 q, k = self.rotary_emb(positions, q, k)

130 - attn_output = self.attn(q, k, v, kv_cache, attn_metadata)

128 + attn_output = self.attn(q, k, v)

129 output, _ = self.o_proj(attn_output)

130 return output

131

@@ -212,14 +210,11 @@ class OlmoDecoderLayer(nn.Module):

210 self,

211 positions: torch.Tensor,

212 hidden_states: torch.Tensor,

215 - kv_cache: torch.Tensor,

216 - attn_metadata: AttentionMetadata,

213 ) -> Tuple[torch.Tensor, Optional[Tuple[torch.Tensor, torch.Tensor]]]:

214 # Attention block.

215 residual = hidden_states

216 hidden_states = self.input_layernorm(hidden_states)

221 - hidden_states = self.self_attn(positions, hidden_states, kv_cache,

222 - attn_metadata)

217 + hidden_states = self.self_attn(positions, hidden_states)

218 hidden_states = hidden_states + residual

219

220 # MLP block.

@@ -263,8 +258,6 @@ class OlmoModel(nn.Module):

258 self,

259 input_ids: torch.Tensor,

260 positions: torch.Tensor,

266 - kv_caches: List[torch.Tensor],

267 - attn_metadata: AttentionMetadata,

261 intermediate_tensors: Optional[IntermediateTensors],

262 inputs_embeds: Optional[torch.Tensor] = None,

263 ) -> Union[torch.Tensor, IntermediateTensors]:

@@ -281,14 +274,9 @@ class OlmoModel(nn.Module):

274 hidden_states = intermediate_tensors["hidden_states"]

275

276 # Apply blocks one-by-one.

284 - for i in range(self.start_layer, self.end_layer):

277 + for layer in self.layers[self.start_layer:self.end_layer]:

278 # shape: (batch_size, seq_len, d_model)

286 - hidden_states = self.layers[i](

287 - positions,

288 - hidden_states,

289 - kv_caches[i - self.start_layer],

290 - attn_metadata,

291 - )

279 + hidden_states = layer(positions, hidden_states)

280

281 if not get_pp_group().is_last_rank:

282 return IntermediateTensors({"hidden_states": hidden_states})

@@ -332,16 +320,12 @@ class OlmoForCausalLM(nn.Module, SupportsPP):

320 self,

321 input_ids: torch.Tensor,

322 positions: torch.Tensor,

335 - kv_caches: List[torch.Tensor],

336 - attn_metadata: AttentionMetadata,

323 intermediate_tensors: Optional[IntermediateTensors] = None,

324 inputs_embeds: Optional[torch.Tensor] = None,

325 ) -> Union[torch.Tensor, IntermediateTensors]:

326 hidden_states = self.model(

327 input_ids=input_ids,

328 positions=positions,

343 - kv_caches=kv_caches,

344 - attn_metadata=attn_metadata,

329 intermediate_tensors=intermediate_tensors,

330 inputs_embeds=inputs_embeds,

331 )

vllm/model_executor/models/olmo2.py

+6 -22

@@ -24,12 +24,12 @@

24 """Inference-only OLMo2 model compatible with HuggingFace weights."""

25

26 from functools import partial

27 -from typing import Iterable, List, Optional, Tuple, Union

27 +from typing import Iterable, Optional, Tuple, Union

28

29 import torch

30 from torch import nn

31

32 -from vllm.attention import Attention, AttentionMetadata

32 +from vllm.attention import Attention

33 from vllm.config import VllmConfig

34 from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size

35 from vllm.distributed.communication_op import tensor_model_parallel_all_gather

@@ -153,14 +153,12 @@ class Olmo2Attention(nn.Module):

153 self,

154 positions: torch.Tensor,

155 hidden_states: torch.Tensor,

156 - kv_cache: torch.Tensor,

157 - attn_metadata: AttentionMetadata,

156 ) -> torch.Tensor:

157 qkv, _ = self.qkv_proj(hidden_states)

158 q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)

159 q, k = self._apply_qk_norm(q, k)

160 q, k = self.rotary_emb(positions, q, k)

163 - attn_output = self.attn(q, k, v, kv_cache, attn_metadata)

161 + attn_output = self.attn(q, k, v)

162 output, _ = self.o_proj(attn_output)

163 return output

164

@@ -239,13 +237,10 @@ class Olmo2DecoderLayer(nn.Module):

237 self,

238 positions: torch.Tensor,

239 hidden_states: torch.Tensor,

242 - kv_cache: torch.Tensor,

243 - attn_metadata: AttentionMetadata,

240 ) -> torch.Tensor:

241 # Attention block.

242 residual = hidden_states

247 - hidden_states = self.self_attn(positions, hidden_states, kv_cache,

248 - attn_metadata)

243 + hidden_states = self.self_attn(positions, hidden_states)

244 hidden_states = self.post_attention_layernorm(hidden_states)

245 hidden_states = hidden_states + residual

246

@@ -287,8 +282,6 @@ class Olmo2Model(nn.Module):

282 self,

283 input_ids: torch.Tensor,

284 positions: torch.Tensor,

290 - kv_caches: List[torch.Tensor],

291 - attn_metadata: AttentionMetadata,

285 intermediate_tensors: Optional[IntermediateTensors],

286 ) -> Union[torch.Tensor, IntermediateTensors]:

287 """

@@ -307,14 +300,9 @@ class Olmo2Model(nn.Module):

300 assert isinstance(hidden_states, torch.Tensor)

301

302 # Apply blocks one-by-one.

310 - for i in range(self.start_layer, self.end_layer):

303 + for layer in self.layers[self.start_layer:self.end_layer]:

304 # shape: (batch_size, seq_len, d_model)

312 - hidden_states = self.layers[i](

313 - positions,

314 - hidden_states,

315 - kv_caches[i - self.start_layer],

316 - attn_metadata,

317 - )

305 + hidden_states = layer(positions, hidden_states)

306

307 if not get_pp_group().is_last_rank:

308 return IntermediateTensors({"hidden_states": hidden_states})

@@ -357,15 +345,11 @@ class Olmo2ForCausalLM(nn.Module, SupportsPP):

345 self,

346 input_ids: torch.Tensor,

347 positions: torch.Tensor,

360 - kv_caches: List[torch.Tensor],

361 - attn_metadata: AttentionMetadata,

348 intermediate_tensors: Optional[IntermediateTensors] = None,

349 ) -> Union[torch.Tensor, IntermediateTensors]:

350 hidden_states = self.model(

351 input_ids=input_ids,

352 positions=positions,

367 - kv_caches=kv_caches,

368 - attn_metadata=attn_metadata,

353 intermediate_tensors=intermediate_tensors,

354 )

355 return hidden_states

vllm/model_executor/models/olmoe.py

+5 -19

@@ -12,13 +12,13 @@

12 # See the License for the specific language governing permissions and

13 # limitations under the License.

14 """Inference-only OLMoE model compatible with HuggingFace weights."""

15 -from typing import Any, Dict, Iterable, List, Optional, Set, Tuple, Union

15 +from typing import Any, Dict, Iterable, Optional, Set, Tuple, Union

16

17 import torch

18 from torch import nn

19 from transformers import PretrainedConfig

20

21 -from vllm.attention import Attention, AttentionMetadata

21 +from vllm.attention import Attention

22 from vllm.compilation.decorators import support_torch_compile

23 from vllm.config import CacheConfig, VllmConfig

24 from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size

@@ -168,14 +168,12 @@ class OlmoeAttention(nn.Module):

168 self,

169 positions: torch.Tensor,

170 hidden_states: torch.Tensor,

171 - kv_cache: torch.Tensor,

172 - attn_metadata: AttentionMetadata,

171 ) -> torch.Tensor:

172 qkv, _ = self.qkv_proj(hidden_states)

173 q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)

174 q, k = self.q_norm(q.contiguous()), self.k_norm(k.contiguous())

175 q, k = self.rotary_emb(positions, q, k)

178 - attn_output = self.attn(q, k, v, kv_cache, attn_metadata)

176 + attn_output = self.attn(q, k, v)

177 output, _ = self.o_proj(attn_output)

178 return output

179

@@ -222,8 +220,6 @@ class OlmoeDecoderLayer(nn.Module):

220 self,

221 positions: torch.Tensor,

222 hidden_states: torch.Tensor,

225 - kv_cache: torch.Tensor,

226 - attn_metadata: AttentionMetadata,

223 residual: Optional[torch.Tensor],

224 ) -> torch.Tensor:

225 # Self Attention

@@ -237,8 +233,6 @@ class OlmoeDecoderLayer(nn.Module):

233 hidden_states = self.self_attn(

234 positions=positions,

235 hidden_states=hidden_states,

240 - kv_cache=kv_cache,

241 - attn_metadata=attn_metadata,

236 )

237

238 # Fully Connected

@@ -283,8 +277,6 @@ class OlmoeModel(nn.Module):

277 self,

278 input_ids: torch.Tensor,

279 positions: torch.Tensor,

286 - kv_caches: List[torch.Tensor],

287 - attn_metadata: AttentionMetadata,

280 intermediate_tensors: Optional[IntermediateTensors],

281 inputs_embeds: Optional[torch.Tensor] = None,

282 ) -> Union[torch.Tensor, IntermediateTensors]:

@@ -299,13 +291,10 @@ class OlmoeModel(nn.Module):

291 hidden_states = intermediate_tensors["hidden_states"]

292 residual = intermediate_tensors["residual"]

293

302 - for i in range(self.start_layer, self.end_layer):

303 - layer = self.layers[i]

294 + for layer in self.layers[self.start_layer:self.end_layer]:

295 hidden_states, residual = layer(

296 positions,

297 hidden_states,

307 - kv_caches[i - self.start_layer],

308 - attn_metadata,

298 residual,

299 )

300

@@ -347,13 +336,10 @@ class OlmoeForCausalLM(nn.Module, SupportsPP):

336 self,

337 input_ids: torch.Tensor,

338 positions: torch.Tensor,

350 - kv_caches: List[torch.Tensor],

351 - attn_metadata: AttentionMetadata,

339 intermediate_tensors: Optional[IntermediateTensors] = None,

340 inputs_embeds: Optional[torch.Tensor] = None,

341 ) -> Union[torch.Tensor, IntermediateTensors]:

355 - hidden_states = self.model(input_ids, positions, kv_caches,

356 - attn_metadata, intermediate_tensors,

342 + hidden_states = self.model(input_ids, positions, intermediate_tensors,

343 inputs_embeds)

344 return hidden_states

345

vllm/model_executor/models/opt.py

+7 -25

@@ -18,13 +18,13 @@

18 # See the License for the specific language governing permissions and

19 # limitations under the License.

20 """Inference-only OPT model compatible with HuggingFace weights."""

21 -from typing import Iterable, List, Optional, Set, Tuple, Union

21 +from typing import Iterable, Optional, Set, Tuple, Union

22

23 import torch

24 from torch import nn

25 from transformers import OPTConfig

26

27 -from vllm.attention import Attention, AttentionMetadata

27 +from vllm.attention import Attention

28 from vllm.compilation.decorators import support_torch_compile

29 from vllm.config import CacheConfig, VllmConfig

30 from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size

@@ -107,12 +107,10 @@ class OPTAttention(nn.Module):

107 def forward(

108 self,

109 hidden_states: torch.Tensor,

110 - kv_cache: torch.Tensor,

111 - attn_metadata: AttentionMetadata,

110 ) -> torch.Tensor:

111 qkv, _ = self.qkv_proj(hidden_states)

112 q, k, v = qkv.chunk(chunks=3, dim=-1)

115 - attn_output = self.attn(q, k, v, kv_cache, attn_metadata)

113 + attn_output = self.attn(q, k, v)

114 output, _ = self.out_proj(attn_output)

115 return output

116

@@ -164,17 +162,13 @@ class OPTDecoderLayer(nn.Module):

162 def forward(

163 self,

164 hidden_states: torch.Tensor,

167 - kv_cache: torch.Tensor,

168 - attn_metadata: AttentionMetadata,

165 ) -> torch.Tensor:

166 # Self Attention

167 residual = hidden_states

168 # 125m, 1.7B, ..., 175B applies layer norm BEFORE attention

169 if self.do_layer_norm_before:

170 hidden_states = self.self_attn_layer_norm(hidden_states)

175 - hidden_states = self.self_attn(hidden_states=hidden_states,

176 - kv_cache=kv_cache,

177 - attn_metadata=attn_metadata)

171 + hidden_states = self.self_attn(hidden_states=hidden_states)

172 hidden_states = residual + hidden_states

173 # 350m applies layer norm AFTER attention

174 if not self.do_layer_norm_before:

@@ -261,8 +255,6 @@ class OPTDecoder(nn.Module):

255 self,

256 input_ids: torch.Tensor,

257 positions: torch.Tensor,

264 - kv_caches: List[torch.Tensor],

265 - attn_metadata: AttentionMetadata,

258 intermediate_tensors: Optional[IntermediateTensors],

259 inputs_embeds: Optional[torch.Tensor] = None,

260 ) -> Union[torch.Tensor, IntermediateTensors]:

@@ -277,11 +269,8 @@ class OPTDecoder(nn.Module):

269 assert intermediate_tensors is not None

270 hidden_states = intermediate_tensors["hidden_states"]

271

280 - for i in range(self.start_layer, self.end_layer):

281 - layer = self.layers[i]

282 - hidden_states = layer(hidden_states,

283 - kv_caches[i - self.start_layer],

284 - attn_metadata)

272 + for layer in self.layers[self.start_layer:self.end_layer]:

273 + hidden_states = layer(hidden_states)

274

275 if not get_pp_group().is_last_rank:

276 return IntermediateTensors({"hidden_states": hidden_states})

@@ -317,15 +306,11 @@ class OPTModel(nn.Module):

306 self,

307 input_ids: torch.Tensor,

308 positions: torch.Tensor,

320 - kv_caches: List[torch.Tensor],

321 - attn_metadata: AttentionMetadata,

309 intermediate_tensors: Optional[IntermediateTensors],

310 inputs_embeds: Optional[torch.Tensor] = None,

311 ) -> Union[torch.Tensor, IntermediateTensors]:

312 return self.decoder(input_ids,

313 positions,

327 - kv_caches,

328 - attn_metadata,

314 intermediate_tensors,

315 inputs_embeds=inputs_embeds)

316

@@ -362,13 +347,10 @@ class OPTForCausalLM(nn.Module, SupportsPP):

347 self,

348 input_ids: torch.Tensor,

349 positions: torch.Tensor,

365 - kv_caches: List[torch.Tensor],

366 - attn_metadata: AttentionMetadata,

350 intermediate_tensors: Optional[IntermediateTensors] = None,

351 inputs_embeds: Optional[torch.Tensor] = None,

352 ) -> Union[torch.Tensor, IntermediateTensors]:

370 - hidden_states = self.model(input_ids, positions, kv_caches,

371 - attn_metadata, intermediate_tensors,

353 + hidden_states = self.model(input_ids, positions, intermediate_tensors,

354 inputs_embeds)

355 return hidden_states

356

vllm/model_executor/models/orion.py

+6 -23

@@ -5,13 +5,13 @@

5 # Copyright (c) OrionStar Inc.

6 # LICENSE: https://huggingface.co/OrionStarAI/Orion-14B-Base/blob/main/LICENSE

7 """Inference-only Orion-14B model compatible with HuggingFace weights."""

8 -from typing import Any, Dict, Iterable, List, Optional, Set, Tuple, Union

8 +from typing import Any, Dict, Iterable, Optional, Set, Tuple, Union

9

10 import torch

11 from torch import nn

12 from transformers import PretrainedConfig

13

14 -from vllm.attention import Attention, AttentionMetadata

14 +from vllm.attention import Attention

15 from vllm.compilation.decorators import support_torch_compile

16 from vllm.config import CacheConfig, VllmConfig

17 from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size

@@ -136,13 +136,11 @@ class OrionAttention(nn.Module):

136 self,

137 positions: torch.Tensor,

138 hidden_states: torch.Tensor,

139 - kv_cache: torch.Tensor,

140 - attn_metadata: AttentionMetadata,

139 ) -> torch.Tensor:

140 qkv, _ = self.qkv_proj(hidden_states)

141 q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)

142 q, k = self.rotary_emb(positions, q, k)

145 - attn_output = self.attn(q, k, v, kv_cache, attn_metadata)

143 + attn_output = self.attn(q, k, v)

144 output, _ = self.o_proj(attn_output)

145 return output

146

@@ -189,8 +187,6 @@ class OrionDecoderLayer(nn.Module):

187 self,

188 positions: torch.Tensor,

189 hidden_states: torch.Tensor,

192 - kv_cache: torch.Tensor,

193 - attn_metadata: AttentionMetadata,

190 ) -> Tuple[torch.Tensor, torch.Tensor]:

191 # Self Attention

192 residual = hidden_states

@@ -198,8 +194,6 @@ class OrionDecoderLayer(nn.Module):

194 hidden_states = self.self_attn(

195 positions=positions,

196 hidden_states=hidden_states,

201 - kv_cache=kv_cache,

202 - attn_metadata=attn_metadata,

197 )

198

199 hidden_states = residual + hidden_states

@@ -247,8 +241,6 @@ class OrionModel(nn.Module):

241 self,

242 input_ids: torch.Tensor,

243 positions: torch.Tensor,

250 - kv_caches: List[torch.Tensor],

251 - attn_metadata: AttentionMetadata,

244 intermediate_tensors: Optional[IntermediateTensors],

245 inputs_embeds: Optional[torch.Tensor] = None,

246 ) -> Union[torch.Tensor, IntermediateTensors]:

@@ -260,14 +252,8 @@ class OrionModel(nn.Module):

252 else:

253 assert intermediate_tensors is not None

254 hidden_states = intermediate_tensors["hidden_states"]

263 - for i in range(self.start_layer, self.end_layer):

264 - layer = self.layers[i]

265 - hidden_states = layer(

266 - positions,

267 - hidden_states,

268 - kv_caches[i - self.start_layer],

269 - attn_metadata,

270 - )

255 + for layer in self.layers[self.start_layer:self.end_layer]:

256 + hidden_states = layer(positions, hidden_states)

257 if not get_pp_group().is_last_rank:

258 return IntermediateTensors({

259 "hidden_states": hidden_states,

@@ -303,13 +289,10 @@ class OrionForCausalLM(nn.Module, SupportsPP):

289 self,

290 input_ids: torch.Tensor,

291 positions: torch.Tensor,

306 - kv_caches: List[torch.Tensor],

307 - attn_metadata: AttentionMetadata,

292 intermediate_tensors: Optional[IntermediateTensors] = None,

293 inputs_embeds: Optional[torch.Tensor] = None,

294 ) -> Union[torch.Tensor, IntermediateTensors]:

311 - hidden_states = self.model(input_ids, positions, kv_caches,

312 - attn_metadata, intermediate_tensors,

295 + hidden_states = self.model(input_ids, positions, intermediate_tensors,

296 inputs_embeds)

297 return hidden_states

298

vllm/model_executor/models/paligemma.py

+1 -6

@@ -1,13 +1,12 @@

1 # SPDX-License-Identifier: Apache-2.0

2

3 -from typing import (Iterable, List, Literal, Mapping, Optional, Set, Tuple,

3 +from typing import (Iterable, Literal, Mapping, Optional, Set, Tuple,

4 TypedDict, Union)

5

6 import torch

7 from torch import nn

8 from transformers import PaliGemmaConfig

9

10 -from vllm.attention import AttentionMetadata

10 from vllm.config import VllmConfig

11 from vllm.inputs import (INPUT_REGISTRY, DecoderOnlyInputs, DummyData,

12 InputContext, token_inputs)

@@ -288,8 +287,6 @@ class PaliGemmaForConditionalGeneration(nn.Module, SupportsMultiModal,

287 def forward(self,

288 input_ids: torch.Tensor,

289 positions: torch.Tensor,

291 - kv_caches: List[torch.Tensor],

292 - attn_metadata: AttentionMetadata,

290 intermediate_tensors: Optional[IntermediateTensors] = None,

291 inputs_embeds: Optional[torch.Tensor] = None,

292 **kwargs: object) -> Union[SamplerOutput, IntermediateTensors]:

@@ -306,8 +303,6 @@ class PaliGemmaForConditionalGeneration(nn.Module, SupportsMultiModal,

303

304 hidden_states = self.language_model.model(input_ids,

305 positions,

309 - kv_caches,

310 - attn_metadata,

306 intermediate_tensors,

307 inputs_embeds=inputs_embeds)

308

vllm/model_executor/models/persimmon.py

+5 -22

@@ -21,13 +21,13 @@

21 # See the License for the specific language governing permissions and

22 # limitations under the License.

23 """Inference-only persimmon model compatible with HuggingFace weights."""

24 -from typing import Iterable, List, Optional, Set, Tuple, Union

24 +from typing import Iterable, Optional, Set, Tuple, Union

25

26 import torch

27 from torch import nn

28 from transformers import PersimmonConfig

29

30 -from vllm.attention import Attention, AttentionMetadata

30 +from vllm.attention import Attention

31 from vllm.compilation.decorators import support_torch_compile

32 from vllm.config import CacheConfig, VllmConfig

33 from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size

@@ -142,8 +142,6 @@ class PersimmonAttention(nn.Module):

142 self,

143 position_ids: torch.Tensor,

144 hidden_states: torch.Tensor,

145 - kv_cache: torch.Tensor,

146 - attn_metadata: AttentionMetadata,

145 ) -> torch.Tensor:

146 # [seq_length, 3 x hidden_size]

147 qkv, _ = self.query_key_value(hidden_states)

@@ -161,7 +159,7 @@ class PersimmonAttention(nn.Module):

159 k = self._merge_heads(k)

160

161 q, k = self.rotary_emb(position_ids, q, k)

164 - attn_output = self.attn(q, k, v, kv_cache, attn_metadata)

162 + attn_output = self.attn(q, k, v)

163 output, _ = self.dense(attn_output)

164 return output

165

@@ -189,8 +187,6 @@ class PersimmonDecoderLayer(nn.Module):

187 self,

188 position_ids: torch.Tensor,

189 hidden_states: torch.Tensor,

192 - kv_cache: torch.Tensor,

193 - attn_metadata: AttentionMetadata,

190 ) -> torch.Tensor:

191 residual = hidden_states

192

@@ -200,8 +196,6 @@ class PersimmonDecoderLayer(nn.Module):

196 hidden_states = self.self_attn(

197 position_ids=position_ids,

198 hidden_states=hidden_states,

203 - kv_cache=kv_cache,

204 - attn_metadata=attn_metadata,

199 )

200 hidden_states = residual + hidden_states

201

@@ -248,8 +242,6 @@ class PersimmonModel(nn.Module):

242 self,

243 input_ids: torch.Tensor,

244 positions: torch.Tensor,

251 - kv_caches: List[torch.Tensor],

252 - attn_metadata: AttentionMetadata,

245 intermediate_tensors: Optional[IntermediateTensors],

246 inputs_embeds: Optional[torch.Tensor] = None,

247 ) -> Union[torch.Tensor, IntermediateTensors]:

@@ -261,13 +253,8 @@ class PersimmonModel(nn.Module):

253 else:

254 assert intermediate_tensors is not None

255 hidden_states = intermediate_tensors["hidden_states"]

264 - for i in range(self.start_layer, self.end_layer):

265 - hidden_states = self.layers[i](

266 - positions,

267 - hidden_states,

268 - kv_caches[i - self.start_layer],

269 - attn_metadata,

270 - )

256 + for layer in self.layers[self.start_layer:self.end_layer]:

257 + hidden_states = layer(positions, hidden_states)

258 if not get_pp_group().is_last_rank:

259 return IntermediateTensors({"hidden_states": hidden_states})

260 hidden_states = self.final_layernorm(hidden_states)

@@ -298,16 +285,12 @@ class PersimmonForCausalLM(nn.Module, SupportsPP):

285 self,

286 input_ids: torch.Tensor,

287 positions: torch.Tensor,

301 - kv_caches: List[torch.Tensor],

302 - attn_metadata: AttentionMetadata,

288 intermediate_tensors: Optional[IntermediateTensors] = None,

289 inputs_embeds: Optional[torch.Tensor] = None,

290 ):

291 hidden_states = self.model(

292 input_ids=input_ids,

293 positions=positions,

309 - kv_caches=kv_caches,

310 - attn_metadata=attn_metadata,

294 intermediate_tensors=intermediate_tensors,

295 inputs_embeds=inputs_embeds,

296 )

vllm/model_executor/models/phi.py

+6 -23

@@ -36,13 +36,13 @@

36 # OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE

37 # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

38 """Inference-only Phi-1.5 model compatible with HuggingFace weights."""

39 -from typing import Iterable, List, Optional, Set, Tuple, Union

39 +from typing import Iterable, Optional, Set, Tuple, Union

40

41 import torch

42 from torch import nn

43 from transformers import PhiConfig

44

45 -from vllm.attention import Attention, AttentionMetadata

45 +from vllm.attention import Attention

46 from vllm.compilation.decorators import support_torch_compile

47 from vllm.config import CacheConfig, VllmConfig

48 from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size

@@ -126,13 +126,11 @@ class PhiAttention(nn.Module):

126 self,

127 position_ids: torch.Tensor,

128 hidden_states: torch.Tensor,

129 - kv_cache: torch.Tensor,

130 - attn_metadata: AttentionMetadata,

129 ) -> torch.Tensor:

130 qkv, _ = self.qkv_proj(hidden_states)

131 q, k, v = qkv.chunk(chunks=3, dim=-1)

132 q, k = self.rotary_emb(position_ids, q, k)

135 - attn_output = self.attn(q, k, v, kv_cache, attn_metadata)

133 + attn_output = self.attn(q, k, v)

134 output, _ = self.dense(attn_output)

135 return output

136

@@ -186,16 +184,12 @@ class PhiLayer(nn.Module):

184 self,

185 position_ids: torch.Tensor,

186 hidden_states: torch.Tensor,

189 - kv_cache: torch.Tensor,

190 - attn_metadata: AttentionMetadata,

187 ) -> torch.Tensor:

188 residual = hidden_states

189 hidden_states = self.input_layernorm(hidden_states)

190 attn_outputs = self.self_attn(

191 position_ids=position_ids,

192 hidden_states=hidden_states,

197 - kv_cache=kv_cache,

198 - attn_metadata=attn_metadata,

193 )

194 feed_forward_hidden_states = self.mlp(hidden_states)

195 hidden_states = attn_outputs + feed_forward_hidden_states + residual

@@ -234,8 +228,6 @@ class PhiModel(nn.Module):

228 self,

229 input_ids: torch.Tensor,

230 positions: torch.Tensor,

237 - kv_caches: List[torch.Tensor],

238 - attn_metadata: AttentionMetadata,

231 intermediate_tensors: Optional[IntermediateTensors],

232 inputs_embeds: Optional[torch.Tensor] = None,

233 ) -> Union[torch.Tensor, IntermediateTensors]:

@@ -247,14 +239,8 @@ class PhiModel(nn.Module):

239 else:

240 assert intermediate_tensors is not None

241 hidden_states = intermediate_tensors["hidden_states"]

250 - for i in range(self.start_layer, self.end_layer):

251 - layer = self.layers[i]

252 - hidden_states = layer(

253 - positions,

254 - hidden_states,

255 - kv_caches[i - self.start_layer],

256 - attn_metadata,

257 - )

242 + for layer in self.layers[self.start_layer:self.end_layer]:

243 + hidden_states = layer(positions, hidden_states)

244

245 if not get_pp_group().is_last_rank:

246 return IntermediateTensors({"hidden_states": hidden_states})

@@ -304,13 +290,10 @@ class PhiForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

290 self,

291 input_ids: torch.Tensor,

292 positions: torch.Tensor,

307 - kv_caches: List[torch.Tensor],

308 - attn_metadata: AttentionMetadata,

293 intermediate_tensors: Optional[IntermediateTensors] = None,

294 inputs_embeds: Optional[torch.Tensor] = None,

295 ) -> Union[torch.Tensor, IntermediateTensors]:

312 - hidden_states = self.model(input_ids, positions, kv_caches,

313 - attn_metadata, intermediate_tensors,

296 + hidden_states = self.model(input_ids, positions, intermediate_tensors,

297 inputs_embeds)

298

299 return hidden_states

vllm/model_executor/models/phi3_small.py

+5 -23

@@ -1,13 +1,13 @@

1 # SPDX-License-Identifier: Apache-2.0

2

3 import math

4 -from typing import Iterable, List, Optional, Set, Tuple, Union

4 +from typing import Iterable, Optional, Set, Tuple, Union

5

6 import torch

7 from torch import nn

8 from transformers.configuration_utils import PretrainedConfig

9

10 -from vllm.attention import Attention, AttentionMetadata

10 +from vllm.attention import Attention

11 from vllm.config import CacheConfig, VllmConfig

12 from vllm.distributed import (get_pp_group, get_tensor_model_parallel_rank,

13 get_tensor_model_parallel_world_size)

@@ -231,8 +231,6 @@ class Phi3SmallSelfAttention(nn.Module):

231 self,

232 positions: torch.Tensor,

233 hidden_states: torch.Tensor,

234 - kv_cache: torch.Tensor,

235 - attn_metadata: AttentionMetadata,

234 ) -> Tuple[torch.Tensor, Optional[torch.Tensor],

235 Optional[Tuple[torch.Tensor]]]:

236 qkv, _ = self.query_key_value(hidden_states)

@@ -248,7 +246,7 @@ class Phi3SmallSelfAttention(nn.Module):

246 v = v.reshape(-1, self.head_dim * self.num_kv_heads_per_partion)

247

248 q, k = self.rotary_emb(positions, q, k)

251 - attn_output = self.attn(q, k, v, kv_cache, attn_metadata=attn_metadata)

249 + attn_output = self.attn(q, k, v)

250 output, _ = self.dense(attn_output)

251

252 return output

@@ -282,8 +280,6 @@ class Phi3SmallDecoderLayer(nn.Module):

280 self,

281 positions: torch.Tensor,

282 hidden_states: torch.Tensor,

285 - kv_cache: torch.Tensor,

286 - attn_metadata: AttentionMetadata,

283 ) -> torch.Tensor:

284 residual = hidden_states

285 hidden_states = self.input_layernorm(hidden_states)

@@ -291,8 +287,6 @@ class Phi3SmallDecoderLayer(nn.Module):

287 hidden_states = self.self_attn(

288 positions=positions,

289 hidden_states=hidden_states,

294 - kv_cache=kv_cache,

295 - attn_metadata=attn_metadata,

290 )

291 hidden_states = residual + hidden_states

292

@@ -338,8 +332,6 @@ class Phi3SmallModel(nn.Module):

332 self,

333 input_ids: torch.LongTensor,

334 positions: Optional[torch.LongTensor],

341 - kv_caches: List[torch.Tensor],

342 - attn_metadata: AttentionMetadata,

335 intermediate_tensors: Optional[IntermediateTensors],

336 inputs_embeds: Optional[torch.Tensor],

337 ) -> Union[torch.Tensor, IntermediateTensors]:

@@ -354,14 +346,8 @@ class Phi3SmallModel(nn.Module):

346 else:

347 assert intermediate_tensors

348 hidden_states = intermediate_tensors["hidden_states"]

357 - for i in range(self.start_layer, self.end_layer):

358 - layer = self.layers[i]

359 - hidden_states = layer(

360 - positions,

361 - hidden_states,

362 - kv_caches[i - self.start_layer],

363 - attn_metadata,

364 - )

349 + for layer in self.layers[self.start_layer:self.end_layer]:

350 + hidden_states = layer(positions, hidden_states)

351 if not get_pp_group().is_last_rank:

352 return IntermediateTensors({"hidden_states": hidden_states})

353 hidden_states = self.final_layernorm(hidden_states)

@@ -438,16 +424,12 @@ class Phi3SmallForCausalLM(nn.Module, SupportsPP):

424 self,

425 input_ids: torch.LongTensor,

426 positions: Optional[torch.LongTensor],

441 - kv_caches: List[torch.Tensor],

442 - attn_metadata: AttentionMetadata,

427 intermediate_tensors: Optional[IntermediateTensors] = None,

428 inputs_embeds: Optional[torch.Tensor] = None,

429 ) -> Union[torch.Tensor, IntermediateTensors]:

430 output_hidden_states = self.model(

431 input_ids=input_ids,

432 positions=positions,

449 - kv_caches=kv_caches,

450 - attn_metadata=attn_metadata,

433 intermediate_tensors=intermediate_tensors,

434 inputs_embeds=inputs_embeds,

435 )

vllm/model_executor/models/phi3v.py

-5

@@ -23,7 +23,6 @@ import torch.nn as nn

23 from transformers import (BatchFeature, CLIPVisionConfig, PretrainedConfig,

24 ProcessorMixin)

25

26 -from vllm.attention import AttentionMetadata

26 from vllm.config import VllmConfig

27 from vllm.logger import init_logger

28 from vllm.model_executor.layers.quantization import QuantizationConfig

@@ -672,8 +671,6 @@ class Phi3VForCausalLM(nn.Module, SupportsMultiModal, SupportsPP,

671 def forward(self,

672 input_ids: torch.Tensor,

673 positions: torch.Tensor,

675 - kv_caches: List[torch.Tensor],

676 - attn_metadata: AttentionMetadata,

674 intermediate_tensors: Optional[IntermediateTensors] = None,

675 inputs_embeds: Optional[torch.Tensor] = None,

676 **kwargs: object):

@@ -691,8 +688,6 @@ class Phi3VForCausalLM(nn.Module, SupportsMultiModal, SupportsPP,

688

689 hidden_states = self.language_model.model(input_ids,

690 positions,

694 - kv_caches,

695 - attn_metadata,

691 intermediate_tensors,

692 inputs_embeds=inputs_embeds)

693

vllm/model_executor/models/phimoe.py

+5 -19

@@ -22,13 +22,13 @@

22 # See the License for the specific language governing permissions and

23 # limitations under the License.

24 """Inference-only PhiMoE model."""

25 -from typing import Iterable, List, Optional, Set, Tuple, Union

25 +from typing import Iterable, Optional, Set, Tuple, Union

26

27 import torch

28 from torch import nn

29 from transformers.configuration_utils import PretrainedConfig

30

31 -from vllm.attention import Attention, AttentionMetadata

31 +from vllm.attention import Attention

32 from vllm.compilation.decorators import support_torch_compile

33 from vllm.config import CacheConfig, VllmConfig

34 from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size

@@ -357,13 +357,11 @@ class PhiMoEAttention(nn.Module):

357 self,

358 positions: torch.Tensor,

359 hidden_states: torch.Tensor,

360 - kv_cache: torch.Tensor,

361 - attn_metadata: AttentionMetadata,

360 ) -> torch.Tensor:

361 qkv, _ = self.qkv_proj(hidden_states)

362 q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)

363 q, k = self.rotary_emb(positions, q, k)

366 - attn_output = self.attn(q, k, v, kv_cache, attn_metadata)

364 + attn_output = self.attn(q, k, v)

365 output, _ = self.o_proj(attn_output)

366 return output

367

@@ -410,8 +408,6 @@ class PhiMoEDecoderLayer(nn.Module):

408 self,

409 positions: torch.Tensor,

410 hidden_states: torch.Tensor,

413 - kv_cache: torch.Tensor,

414 - attn_metadata: AttentionMetadata,

411 residual: Optional[torch.Tensor],

412 ) -> torch.Tensor:

413 residual = hidden_states

@@ -422,8 +418,6 @@ class PhiMoEDecoderLayer(nn.Module):

418 hidden_states = self.self_attn(

419 positions=positions,

420 hidden_states=hidden_states,

425 - kv_cache=kv_cache,

426 - attn_metadata=attn_metadata,

421 )

422 hidden_states = hidden_states + residual

423

@@ -478,8 +472,6 @@ class PhiMoEModel(nn.Module):

472 self,

473 input_ids: torch.Tensor,

474 positions: torch.Tensor,

481 - kv_caches: List[torch.Tensor],

482 - attn_metadata: AttentionMetadata,

475 intermediate_tensors: Optional[IntermediateTensors],

476 inputs_embeds: Optional[torch.Tensor] = None,

477 ) -> Union[torch.Tensor, IntermediateTensors]:

@@ -494,13 +486,10 @@ class PhiMoEModel(nn.Module):

486 hidden_states = intermediate_tensors["hidden_states"]

487 residual = intermediate_tensors["residual"]

488

497 - for i in range(self.start_layer, self.end_layer):

498 - layer = self.layers[i]

489 + for layer in self.layers[self.start_layer:self.end_layer]:

490 hidden_states, residual = layer(

491 positions,

492 hidden_states,

502 - kv_caches[i - self.start_layer],

503 - attn_metadata,

493 residual,

494 )

495

@@ -571,13 +560,10 @@ class PhiMoEForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

560 self,

561 input_ids: torch.Tensor,

562 positions: torch.Tensor,

574 - kv_caches: List[torch.Tensor],

575 - attn_metadata: AttentionMetadata,

563 intermediate_tensors: Optional[IntermediateTensors] = None,

564 inputs_embeds: Optional[torch.Tensor] = None,

565 ) -> Union[torch.Tensor, IntermediateTensors]:

579 - hidden_states = self.model(input_ids, positions, kv_caches,

580 - attn_metadata, intermediate_tensors,

566 + hidden_states = self.model(input_ids, positions, intermediate_tensors,

567 inputs_embeds)

568 return hidden_states

569

vllm/model_executor/models/pixtral.py

-5

@@ -16,7 +16,6 @@ from transformers.models.pixtral.image_processing_pixtral import (

16 from transformers.models.pixtral.modeling_pixtral import (

17 PixtralRotaryEmbedding, apply_rotary_pos_emb, position_ids_in_meshgrid)

18

19 -from vllm.attention import AttentionMetadata

19 from vllm.config import VllmConfig

20 from vllm.distributed import divide, get_tensor_model_parallel_world_size

21 from vllm.inputs import (INPUT_REGISTRY, DecoderOnlyInputs, DummyData,

@@ -270,8 +269,6 @@ class PixtralForConditionalGeneration(nn.Module, SupportsMultiModal,

269 self,

270 input_ids: torch.Tensor,

271 positions: torch.Tensor,

273 - kv_caches: List[torch.Tensor],

274 - attn_metadata: AttentionMetadata,

272 intermediate_tensors: Optional[IntermediateTensors] = None,

273 inputs_embeds: Optional[torch.Tensor] = None,

274 **kwargs: object,

@@ -291,8 +288,6 @@ class PixtralForConditionalGeneration(nn.Module, SupportsMultiModal,

288

289 hidden_states = self.language_model.model(input_ids,

290 positions,

294 - kv_caches,

295 - attn_metadata,

291 intermediate_tensors,

292 inputs_embeds=inputs_embeds)

293

vllm/model_executor/models/prithvi_geospatial_mae.py

+1 -4

@@ -15,13 +15,12 @@

15 # See the License for the specific language governing permissions and

16 # limitations under the License.

17 """Inference-only IBM/NASA Prithvi Geospatial model."""

18 -from typing import Iterable, List, Mapping, Optional, Set, Tuple, Union

18 +from typing import Iterable, Mapping, Optional, Set, Tuple, Union

19

20 import torch

21 import torch.nn as nn

22 from transformers import BatchFeature

23

24 -from vllm.attention import AttentionMetadata

24 from vllm.config import VllmConfig

25 from vllm.model_executor.model_loader.weight_utils import default_weight_loader

26 from vllm.model_executor.models.interfaces import (IsAttentionFree,

@@ -181,8 +180,6 @@ class PrithviGeoSpatialMAE(nn.Module, IsAttentionFree, SupportsMultiModal):

180 self,

181 input_ids: Optional[torch.Tensor],

182 positions: torch.Tensor,

184 - kv_caches: List[torch.Tensor],

185 - attn_metadata: AttentionMetadata,

183 intermediate_tensors: Optional[IntermediateTensors] = None,

184 inputs_embeds: Optional[torch.Tensor] = None,

185 **kwargs: object,

vllm/model_executor/models/qwen.py

+6 -20

@@ -6,13 +6,13 @@

6 # LICENSE: https://huggingface.co/Qwen/Qwen-7B/blob/main/LICENSE

7 """Inference-only QWen model compatible with HuggingFace weights."""

8

9 -from typing import Any, Dict, Iterable, List, Optional, Set, Tuple, Union

9 +from typing import Any, Dict, Iterable, Optional, Set, Tuple, Union

10

11 import torch

12 from torch import nn

13 from transformers import PretrainedConfig

14

15 -from vllm.attention import Attention, AttentionMetadata

15 +from vllm.attention import Attention

16 from vllm.compilation.decorators import support_torch_compile

17 from vllm.config import CacheConfig, VllmConfig

18 from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size

@@ -124,13 +124,11 @@ class QWenAttention(nn.Module):

124 self,

125 positions: torch.Tensor,

126 hidden_states: torch.Tensor,

127 - kv_cache: torch.Tensor,

128 - attn_metadata: AttentionMetadata,

127 ) -> torch.Tensor:

128 qkv, _ = self.c_attn(hidden_states)

129 q, k, v = qkv.chunk(chunks=3, dim=-1)

130 q, k = self.rotary_emb(positions, q, k)

133 - attn_output = self.attn(q, k, v, kv_cache, attn_metadata)

131 + attn_output = self.attn(q, k, v)

132 output, _ = self.c_proj(attn_output)

133 return output

134

@@ -168,8 +166,6 @@ class QWenBlock(nn.Module):

166 self,

167 positions: torch.Tensor,

168 hidden_states: torch.Tensor,

171 - kv_cache: torch.Tensor,

172 - attn_metadata: AttentionMetadata,

169 residual: Optional[torch.Tensor],

170 ) -> Tuple[torch.Tensor, torch.Tensor]:

171 # Self Attention

@@ -181,8 +177,6 @@ class QWenBlock(nn.Module):

177 hidden_states = self.attn(

178 positions=positions,

179 hidden_states=hidden_states,

184 - kv_cache=kv_cache,

185 - attn_metadata=attn_metadata,

180 )

181

182 # Fully Connected

@@ -225,8 +219,6 @@ class QWenModel(nn.Module):

219 self,

220 input_ids: torch.Tensor,

221 positions: torch.Tensor,

228 - kv_caches: List[torch.Tensor],

229 - attn_metadata: AttentionMetadata,

222 intermediate_tensors: Optional[IntermediateTensors],

223 inputs_embeds: Optional[torch.Tensor] = None,

224 ) -> Union[torch.Tensor, IntermediateTensors]:

@@ -241,13 +233,10 @@ class QWenModel(nn.Module):

233 hidden_states = intermediate_tensors["hidden_states"]

234 residual = intermediate_tensors["residual"]

235

244 - for i in range(self.start_layer, self.end_layer):

245 - layer = self.h[i]

236 + for layer in self.h[self.start_layer:self.end_layer]:

237 hidden_states, residual = layer(

238 positions,

239 hidden_states,

249 - kv_caches[i - self.start_layer],

250 - attn_metadata,

240 residual,

241 )

242 if not get_pp_group().is_last_rank:

@@ -373,12 +362,9 @@ class QWenLMHeadModel(QWenBaseModel, SupportsPP, SupportsLoRA):

362 self,

363 input_ids: torch.Tensor,

364 positions: torch.Tensor,

376 - kv_caches: List[torch.Tensor],

377 - attn_metadata: AttentionMetadata,

365 intermediate_tensors: Optional[IntermediateTensors] = None,

366 inputs_embeds: Optional[torch.Tensor] = None,

367 ) -> Union[torch.Tensor, IntermediateTensors]:

381 - hidden_states = self.transformer(input_ids, positions, kv_caches,

382 - attn_metadata, intermediate_tensors,

383 - inputs_embeds)

368 + hidden_states = self.transformer(input_ids, positions,

369 + intermediate_tensors, inputs_embeds)

370 return hidden_states

vllm/model_executor/models/qwen2.py

+6 -23

@@ -23,13 +23,13 @@

23 # See the License for the specific language governing permissions and

24 # limitations under the License.

25 """Inference-only Qwen2 model compatible with HuggingFace weights."""

26 -from typing import Iterable, List, Optional, Set, Tuple, Union

26 +from typing import Iterable, Optional, Set, Tuple, Union

27

28 import torch

29 from torch import nn

30 from transformers import Qwen2Config

31

32 -from vllm.attention import Attention, AttentionMetadata, AttentionType

32 +from vllm.attention import Attention, AttentionType

33 from vllm.compilation.decorators import support_torch_compile

34 from vllm.config import CacheConfig, VllmConfig

35 from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size

@@ -170,13 +170,11 @@ class Qwen2Attention(nn.Module):

170 self,

171 positions: torch.Tensor,

172 hidden_states: torch.Tensor,

173 - kv_cache: torch.Tensor,

174 - attn_metadata: AttentionMetadata,

173 ) -> torch.Tensor:

174 qkv, _ = self.qkv_proj(hidden_states)

175 q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)

176 q, k = self.rotary_emb(positions, q, k)

179 - attn_output = self.attn(q, k, v, kv_cache, attn_metadata)

177 + attn_output = self.attn(q, k, v)

178 output, _ = self.o_proj(attn_output)

179 return output

180

@@ -233,8 +231,6 @@ class Qwen2DecoderLayer(nn.Module):

231 self,

232 positions: torch.Tensor,

233 hidden_states: torch.Tensor,

236 - kv_cache: torch.Tensor,

237 - attn_metadata: AttentionMetadata,

234 residual: Optional[torch.Tensor],

235 ) -> Tuple[torch.Tensor, torch.Tensor]:

236 # Self Attention

@@ -247,8 +243,6 @@ class Qwen2DecoderLayer(nn.Module):

243 hidden_states = self.self_attn(

244 positions=positions,

245 hidden_states=hidden_states,

250 - kv_cache=kv_cache,

251 - attn_metadata=attn_metadata,

246 )

247

248 # Fully Connected

@@ -328,8 +322,6 @@ class Qwen2Model(nn.Module):

322 self,

323 input_ids: torch.Tensor,

324 positions: torch.Tensor,

331 - kv_caches: List[torch.Tensor],

332 - attn_metadata: AttentionMetadata,

325 intermediate_tensors: Optional[IntermediateTensors] = None,

326 inputs_embeds: Optional[torch.Tensor] = None,

327 ) -> Union[torch.Tensor, IntermediateTensors]:

@@ -343,13 +335,10 @@ class Qwen2Model(nn.Module):

335 assert intermediate_tensors is not None

336 hidden_states = intermediate_tensors["hidden_states"]

337 residual = intermediate_tensors["residual"]

346 - for i in range(self.start_layer, self.end_layer):

347 - layer = self.layers[i]

338 + for layer in self.layers[self.start_layer:self.end_layer]:

339 hidden_states, residual = layer(

340 positions,

341 hidden_states,

351 - kv_caches[i - self.start_layer],

352 - attn_metadata,

342 residual,

343 )

344 if not get_pp_group().is_last_rank:

@@ -468,13 +457,10 @@ class Qwen2ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

457 self,

458 input_ids: torch.Tensor,

459 positions: torch.Tensor,

471 - kv_caches: List[torch.Tensor],

472 - attn_metadata: AttentionMetadata,

460 intermediate_tensors: Optional[IntermediateTensors] = None,

461 inputs_embeds: Optional[torch.Tensor] = None,

462 ) -> Union[torch.Tensor, IntermediateTensors]:

476 - hidden_states = self.model(input_ids, positions, kv_caches,

477 - attn_metadata, intermediate_tensors,

463 + hidden_states = self.model(input_ids, positions, intermediate_tensors,

464 inputs_embeds)

465 return hidden_states

466

@@ -553,12 +539,9 @@ class Qwen2EmbeddingModel(nn.Module, SupportsLoRA, SupportsPP):

539 self,

540 input_ids: torch.Tensor,

541 positions: torch.Tensor,

556 - kv_caches: List[torch.Tensor],

557 - attn_metadata: AttentionMetadata,

542 intermediate_tensors: Optional[IntermediateTensors] = None,

543 ) -> torch.Tensor:

560 - return self.model(input_ids, positions, kv_caches, attn_metadata,

561 - intermediate_tensors)

544 + return self.model(input_ids, positions, intermediate_tensors)

545

546 def pooler(

547 self,

vllm/model_executor/models/qwen2_5_vl.py

-5

@@ -37,7 +37,6 @@ from transformers.models.qwen2_5_vl import Qwen2_5_VLProcessor

37 from transformers.models.qwen2_5_vl.configuration_qwen2_5_vl import (

38 Qwen2_5_VLConfig, Qwen2_5_VLVisionConfig)

39

40 -from vllm.attention import AttentionMetadata

40 from vllm.config import VllmConfig

41 from vllm.distributed import parallel_state, tensor_model_parallel_all_gather

42 from vllm.distributed import utils as dist_utils

@@ -992,8 +991,6 @@ class Qwen2_5_VLForConditionalGeneration(nn.Module, SupportsMultiModal,

991 self,

992 input_ids: torch.Tensor,

993 positions: torch.Tensor,

995 - kv_caches: List[torch.Tensor],

996 - attn_metadata: AttentionMetadata,

994 intermediate_tensors: Optional[IntermediateTensors] = None,

995 inputs_embeds: Optional[torch.Tensor] = None,

996 **kwargs: object,

@@ -1047,8 +1044,6 @@ class Qwen2_5_VLForConditionalGeneration(nn.Module, SupportsMultiModal,

1044 hidden_states = self.language_model.model(

1045 input_ids=input_ids,

1046 positions=positions,

1050 - kv_caches=kv_caches,

1051 - attn_metadata=attn_metadata,

1047 intermediate_tensors=intermediate_tensors,

1048 inputs_embeds=inputs_embeds,

1049 )

vllm/model_executor/models/qwen2_audio.py

+2 -7

@@ -22,8 +22,8 @@

22 # limitations under the License.

23 """Inference-only Qwen2-Audio model compatible with HuggingFace weights."""

24 from functools import cached_property

25 -from typing import (Any, Iterable, List, Mapping, Optional, Set, Tuple,

26 - TypedDict, Union)

25 +from typing import (Any, Iterable, Mapping, Optional, Set, Tuple, TypedDict,

26 + Union)

27

28 import torch

29 import torch.nn as nn

@@ -33,7 +33,6 @@ from transformers.models.qwen2_audio import (Qwen2AudioConfig,

33 Qwen2AudioProcessor)

34 from transformers.models.whisper import WhisperFeatureExtractor

35

36 -from vllm.attention import AttentionMetadata

36 from vllm.config import VllmConfig

37 from vllm.model_executor.layers.sampler import SamplerOutput, get_sampler

38 from vllm.model_executor.sampling_metadata import SamplingMetadata

@@ -380,8 +379,6 @@ class Qwen2AudioForConditionalGeneration(nn.Module, SupportsMultiModal,

379 self,

380 input_ids: torch.Tensor,

381 positions: torch.Tensor,

383 - kv_caches: List[torch.Tensor],

384 - attn_metadata: AttentionMetadata,

382 intermediate_tensors: Optional[IntermediateTensors] = None,

383 inputs_embeds: Optional[torch.Tensor] = None,

384 **kwargs: object,

@@ -400,8 +397,6 @@ class Qwen2AudioForConditionalGeneration(nn.Module, SupportsMultiModal,

397

398 hidden_states = self.language_model.model(input_ids,

399 positions,

403 - kv_caches,

404 - attn_metadata,

400 intermediate_tensors,

401 inputs_embeds=inputs_embeds)

402 return hidden_states

vllm/model_executor/models/qwen2_moe.py

+6 -20

@@ -23,14 +23,14 @@

23 # See the License for the specific language governing permissions and

24 # limitations under the License.

25 """Inference-only Qwen2MoE model compatible with HuggingFace weights."""

26 -from typing import Any, Dict, Iterable, List, Optional, Set, Tuple, Union

26 +from typing import Any, Dict, Iterable, Optional, Set, Tuple, Union

27

28 import torch

29 import torch.nn.functional as F

30 from torch import nn

31 from transformers import PretrainedConfig

32

33 -from vllm.attention import Attention, AttentionMetadata

33 +from vllm.attention import Attention

34 from vllm.compilation.decorators import support_torch_compile

35 from vllm.config import CacheConfig, VllmConfig

36 from vllm.distributed import (get_pp_group,

@@ -232,13 +232,11 @@ class Qwen2MoeAttention(nn.Module):

232 self,

233 positions: torch.Tensor,

234 hidden_states: torch.Tensor,

235 - kv_cache: torch.Tensor,

236 - attn_metadata: AttentionMetadata,

235 ) -> torch.Tensor:

236 qkv, _ = self.qkv_proj(hidden_states)

237 q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)

238 q, k = self.rotary_emb(positions, q, k)

241 - attn_output = self.attn(q, k, v, kv_cache, attn_metadata)

239 + attn_output = self.attn(q, k, v)

240 output, _ = self.o_proj(attn_output)

241 return output

242

@@ -296,8 +294,6 @@ class Qwen2MoeDecoderLayer(nn.Module):

294 self,

295 positions: torch.Tensor,

296 hidden_states: torch.Tensor,

299 - kv_cache: torch.Tensor,

300 - attn_metadata: AttentionMetadata,

297 residual: Optional[torch.Tensor],

298 ) -> torch.Tensor:

299 # Self Attention

@@ -310,8 +306,6 @@ class Qwen2MoeDecoderLayer(nn.Module):

306 hidden_states = self.self_attn(

307 positions=positions,

308 hidden_states=hidden_states,

313 - kv_cache=kv_cache,

314 - attn_metadata=attn_metadata,

309 )

310

311 # Fully Connected

@@ -358,8 +352,6 @@ class Qwen2MoeModel(nn.Module):

352 self,

353 input_ids: torch.Tensor,

354 positions: torch.Tensor,

361 - kv_caches: List[torch.Tensor],

362 - attn_metadata: AttentionMetadata,

355 intermediate_tensors: Optional[IntermediateTensors] = None,

356 inputs_embeds: Optional[torch.Tensor] = None,

357 ) -> Union[torch.Tensor, IntermediateTensors]:

@@ -373,11 +365,8 @@ class Qwen2MoeModel(nn.Module):

365 assert intermediate_tensors is not None

366 hidden_states = intermediate_tensors["hidden_states"]

367 residual = intermediate_tensors["residual"]

376 - for i in range(self.start_layer, self.end_layer):

377 - layer = self.layers[i]

378 - hidden_states, residual = layer(positions, hidden_states,

379 - kv_caches[i - self.start_layer],

380 - attn_metadata, residual)

368 + for layer in self.layers[self.start_layer:self.end_layer]:

369 + hidden_states, residual = layer(positions, hidden_states, residual)

370 if not get_pp_group().is_last_rank:

371 return IntermediateTensors({

372 "hidden_states": hidden_states,

@@ -416,13 +405,10 @@ class Qwen2MoeForCausalLM(nn.Module, SupportsPP):

405 self,

406 input_ids: torch.Tensor,

407 positions: torch.Tensor,

419 - kv_caches: List[torch.Tensor],

420 - attn_metadata: AttentionMetadata,

408 intermediate_tensors: Optional[IntermediateTensors] = None,

409 inputs_embeds: Optional[torch.Tensor] = None,

410 ) -> Union[torch.Tensor, IntermediateTensors]:

424 - hidden_states = self.model(input_ids, positions, kv_caches,

425 - attn_metadata, intermediate_tensors,

411 + hidden_states = self.model(input_ids, positions, intermediate_tensors,

412 inputs_embeds)

413 return hidden_states

414

vllm/model_executor/models/qwen2_rm.py

+2 -6

@@ -5,12 +5,11 @@

7 """Inference-only Qwen2-RM model compatible with HuggingFace weights."""

8 -from typing import Iterable, List, Optional, Set, Tuple, Union

8 +from typing import Iterable, Optional, Set, Tuple, Union

9

10 import torch

11 from torch import nn

12

13 -from vllm.attention import AttentionMetadata

13 from vllm.config import VllmConfig

14 from vllm.model_executor.layers.linear import (ColumnParallelLinear,

15 RowParallelLinear)

@@ -80,13 +79,10 @@ class Qwen2RewardBaseModel(nn.Module, SupportsLoRA, SupportsPP):

79 self,

80 input_ids: torch.Tensor,

81 positions: torch.Tensor,

83 - kv_caches: List[torch.Tensor],

84 - attn_metadata: AttentionMetadata,

82 intermediate_tensors: Optional[IntermediateTensors] = None,

83 inputs_embeds: Optional[torch.Tensor] = None,

84 ) -> Union[torch.Tensor, IntermediateTensors]:

88 - hidden_states = self.model(input_ids, positions, kv_caches,

89 - attn_metadata, intermediate_tensors,

85 + hidden_states = self.model(input_ids, positions, intermediate_tensors,

86 inputs_embeds)

87 logits, _ = self.score(hidden_states)

88 return logits

vllm/model_executor/models/qwen2_vl.py

+2 -7

@@ -24,8 +24,8 @@

24 # limitations under the License.

25 """Inference-only Qwen2-VL model compatible with HuggingFace weights."""

26 from functools import cached_property, partial

27 -from typing import (Any, Callable, Iterable, List, Literal, Mapping, Optional,

28 - Set, Tuple, Type, TypedDict, Union)

27 +from typing import (Any, Callable, Iterable, Literal, Mapping, Optional, Set,

28 + Tuple, Type, TypedDict, Union)

29

30 import torch

31 import torch.nn as nn

@@ -38,7 +38,6 @@ from transformers.models.qwen2_vl.configuration_qwen2_vl import (

38 Qwen2VLConfig, Qwen2VLVisionConfig)

39 from transformers.models.qwen2_vl.image_processing_qwen2_vl import smart_resize

40

41 -from vllm.attention import AttentionMetadata

41 from vllm.config import VllmConfig

42 from vllm.distributed import parallel_state, tensor_model_parallel_all_gather

43 from vllm.distributed import utils as dist_utils

@@ -1302,8 +1301,6 @@ class Qwen2VLForConditionalGeneration(nn.Module, SupportsMultiModal,

1301 self,

1302 input_ids: torch.Tensor,

1303 positions: torch.Tensor,

1305 - kv_caches: List[torch.Tensor],

1306 - attn_metadata: AttentionMetadata,

1304 intermediate_tensors: Optional[IntermediateTensors] = None,

1305 inputs_embeds: Optional[torch.Tensor] = None,

1306 **kwargs: object,

@@ -1354,8 +1351,6 @@ class Qwen2VLForConditionalGeneration(nn.Module, SupportsMultiModal,

1351 hidden_states = self.language_model.model(

1352 input_ids=input_ids,

1353 positions=positions,

1357 - kv_caches=kv_caches,

1358 - attn_metadata=attn_metadata,

1354 intermediate_tensors=intermediate_tensors,

1355 inputs_embeds=inputs_embeds,

1356 )

vllm/model_executor/models/qwen_vl.py

+2 -6

@@ -22,7 +22,6 @@ from transformers import (BatchFeature, PretrainedConfig, PreTrainedTokenizer,

22 from transformers.image_utils import ImageInput

23 from transformers.tokenization_utils_base import TextInput

24

25 -from vllm.attention import AttentionMetadata

25 from vllm.config import VllmConfig

26 from vllm.model_executor.layers.activation import get_act_fn

27 from vllm.model_executor.layers.linear import (ColumnParallelLinear,

@@ -766,8 +765,6 @@ class QwenVLForConditionalGeneration(QWenBaseModel, SupportsPP, SupportsLoRA,

765 self,

766 input_ids: torch.Tensor,

767 positions: torch.Tensor,

769 - kv_caches: List[torch.Tensor],

770 - attn_metadata: AttentionMetadata,

768 intermediate_tensors: Optional[IntermediateTensors] = None,

769 inputs_embeds: Optional[torch.Tensor] = None,

770 **kwargs: object,

@@ -783,7 +780,6 @@ class QwenVLForConditionalGeneration(QWenBaseModel, SupportsPP, SupportsLoRA,

780 vision_embeddings)

781 input_ids = None

782

786 - hidden_states = self.transformer(input_ids, positions, kv_caches,

787 - attn_metadata, intermediate_tensors,

788 - inputs_embeds)

783 + hidden_states = self.transformer(input_ids, positions,

784 + intermediate_tensors, inputs_embeds)

785 return hidden_states

vllm/model_executor/models/roberta.py

+1 -6

@@ -1,13 +1,12 @@

1 # SPDX-License-Identifier: Apache-2.0

2

3 import itertools

4 -from typing import Iterable, List, Optional, Tuple

4 +from typing import Iterable, Optional, Tuple

5

6 import torch

7 from torch import nn

8 from transformers import RobertaConfig

9

10 -from vllm.attention import AttentionMetadata

10 from vllm.config import VllmConfig

11 from vllm.model_executor.layers.pooler import CrossEncodingPooler

12 from vllm.model_executor.layers.vocab_parallel_embedding import (

@@ -243,16 +242,12 @@ class RobertaForSequenceClassification(nn.Module, SupportsCrossEncoding):

242 self,

243 input_ids: Optional[torch.Tensor],

244 positions: torch.Tensor,

246 - kv_caches: List[torch.Tensor],

247 - attn_metadata: AttentionMetadata,

245 intermediate_tensors: Optional[IntermediateTensors] = None,

246 inputs_embeds: Optional[torch.Tensor] = None,

247 token_type_ids: Optional[torch.Tensor] = None,

248 ) -> torch.Tensor:

249 return self.roberta(input_ids=input_ids,

250 position_ids=positions,

254 - kv_caches=kv_caches,

251 inputs_embeds=inputs_embeds,

252 intermediate_tensors=intermediate_tensors,

257 - attn_metadata=attn_metadata,

253 token_type_ids=token_type_ids)

vllm/model_executor/models/solar.py

+4 -17

@@ -23,13 +23,13 @@

23 # limitations under the License.

24 """Inference-only Solar model compatible with HuggingFace weights."""

25

26 -from typing import Any, Dict, Iterable, List, Optional, Set, Tuple, Union

26 +from typing import Any, Dict, Iterable, Optional, Set, Tuple, Union

27

28 import torch

29 from torch import nn

30 from transformers import PretrainedConfig

31

32 -from vllm.attention import Attention, AttentionMetadata

32 +from vllm.attention import Attention

33 from vllm.compilation.decorators import support_torch_compile

34 from vllm.config import CacheConfig, VllmConfig

35 from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size

@@ -172,13 +172,11 @@ class SolarAttention(nn.Module):

172 self,

173 positions: torch.Tensor,

174 hidden_states: torch.Tensor,

175 - kv_cache: torch.Tensor,

176 - attn_metadata: AttentionMetadata,

175 ) -> torch.Tensor:

176 qkv, _ = self.qkv_proj(hidden_states)

177 q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)

178 q, k = self.rotary_emb(positions, q, k)

181 - attn_output = self.attn(q, k, v, kv_cache, attn_metadata)

179 + attn_output = self.attn(q, k, v)

180 output, _ = self.o_proj(attn_output)

181 return output

182

@@ -238,8 +236,6 @@ class SolarDecoderLayer(nn.Module):

236 self,

237 positions: torch.Tensor,

238 hidden_states: torch.Tensor,

241 - kv_cache: torch.Tensor,

242 - attn_metadata: AttentionMetadata,

239 residual: Optional[torch.Tensor],

240 ) -> Tuple[torch.Tensor, torch.Tensor]:

241 # Self Attention

@@ -252,8 +248,6 @@ class SolarDecoderLayer(nn.Module):

248 hidden_states = self.self_attn(

249 positions=positions,

250 hidden_states=hidden_states,

255 - kv_cache=kv_cache,

256 - attn_metadata=attn_metadata,

251 )

252

253 # Fully Connected

@@ -315,8 +309,6 @@ class SolarModel(nn.Module):

309 self,

310 input_ids: Optional[torch.Tensor],

311 positions: torch.Tensor,

318 - kv_caches: List[torch.Tensor],

319 - attn_metadata: AttentionMetadata,

312 intermediate_tensors: Optional[IntermediateTensors],

313 inputs_embeds: Optional[torch.Tensor] = None,

314 ) -> Union[torch.Tensor, IntermediateTensors]:

@@ -357,8 +349,6 @@ class SolarModel(nn.Module):

349 hidden_states, residual = layer(

350 positions,

351 hidden_states,

360 - kv_caches[i - self.start_layer],

361 - attn_metadata,

352 residual,

353 )

354

@@ -438,13 +428,10 @@ class SolarForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

428 self,

429 input_ids: torch.Tensor,

430 positions: torch.Tensor,

441 - kv_caches: List[torch.Tensor],

442 - attn_metadata: AttentionMetadata,

431 intermediate_tensors: Optional[IntermediateTensors] = None,

432 inputs_embeds: Optional[torch.Tensor] = None,

433 ) -> Union[torch.Tensor, IntermediateTensors]:

446 - model_output = self.model(input_ids, positions, kv_caches,

447 - attn_metadata, intermediate_tensors,

434 + model_output = self.model(input_ids, positions, intermediate_tensors,

435 inputs_embeds)

436 return model_output

437

vllm/model_executor/models/stablelm.py

+6 -23

@@ -20,13 +20,13 @@

20 # https://huggingface.co/stabilityai/stablelm-3b-4e1t/blob/main/config.json

21 """Inference-only StabeLM (https://github.com/Stability-AI/StableLM)

22 model compatible with HuggingFace weights."""

23 -from typing import Iterable, List, Optional, Set, Tuple, Union

23 +from typing import Iterable, Optional, Set, Tuple, Union

24

25 import torch

26 from torch import nn

27 from transformers import StableLmConfig

28

29 -from vllm.attention import Attention, AttentionMetadata

29 +from vllm.attention import Attention

30 from vllm.config import CacheConfig, VllmConfig

31 from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size

32 from vllm.model_executor.layers.activation import SiluAndMul

@@ -147,13 +147,11 @@ class StablelmAttention(nn.Module):

147 self,

148 positions: torch.Tensor,

149 hidden_states: torch.Tensor,

150 - kv_cache: torch.Tensor,

151 - attn_metadata: AttentionMetadata,

150 ) -> torch.Tensor:

151 qkv, _ = self.qkv_proj(hidden_states)

152 q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)

153 q, k = self.rotary_emb(positions, q, k)

156 - attn_output = self.attn(q, k, v, kv_cache, attn_metadata)

154 + attn_output = self.attn(q, k, v)

155 output, _ = self.o_proj(attn_output)

156 return output

157

@@ -183,8 +181,6 @@ class StablelmDecoderLayer(nn.Module):

181 self,

182 positions: torch.Tensor,

183 hidden_states: torch.Tensor,

186 - kv_cache: torch.Tensor,

187 - attn_metadata: AttentionMetadata,

184 ) -> Tuple[torch.Tensor, torch.Tensor]:

185 # Self Attention

186 residual = hidden_states

@@ -192,8 +188,6 @@ class StablelmDecoderLayer(nn.Module):

188 hidden_states = self.self_attn(

189 positions=positions,

190 hidden_states=hidden_states,

195 - kv_cache=kv_cache,

196 - attn_metadata=attn_metadata,

191 )

192 hidden_states = residual + hidden_states

193

@@ -241,8 +235,6 @@ class StableLMEpochModel(nn.Module):

235 self,

236 input_ids: torch.Tensor,

237 positions: torch.Tensor,

244 - kv_caches: List[torch.Tensor],

245 - attn_metadata: AttentionMetadata,

238 intermediate_tensors: Optional[IntermediateTensors],

239 inputs_embeds: Optional[torch.Tensor] = None,

240 ) -> Union[torch.Tensor, IntermediateTensors]:

@@ -254,14 +246,8 @@ class StableLMEpochModel(nn.Module):

246 else:

247 assert intermediate_tensors is not None

248 hidden_states = intermediate_tensors["hidden_states"]

257 - for i in range(self.start_layer, self.end_layer):

258 - layer = self.layers[i]

259 - hidden_states, residual = layer(

260 - positions,

261 - hidden_states,

262 - kv_caches[i - self.start_layer],

263 - attn_metadata,

264 - )

249 + for layer in self.layers[self.start_layer:self.end_layer]:

250 + hidden_states, residual = layer(positions, hidden_states)

251 if not get_pp_group().is_last_rank:

252 return IntermediateTensors({"hidden_states": hidden_states})

253 hidden_states = self.norm(hidden_states)

@@ -296,13 +282,10 @@ class StablelmForCausalLM(nn.Module, SupportsPP):

282 self,

283 input_ids: torch.Tensor,

284 positions: torch.Tensor,

299 - kv_caches: List[torch.Tensor],

300 - attn_metadata: AttentionMetadata,

285 intermediate_tensors: Optional[IntermediateTensors] = None,

286 inputs_embeds: Optional[torch.Tensor] = None,

287 ) -> Union[torch.Tensor, IntermediateTensors]:

304 - hidden_states = self.model(input_ids, positions, kv_caches,

305 - attn_metadata, intermediate_tensors,

288 + hidden_states = self.model(input_ids, positions, intermediate_tensors,

289 inputs_embeds)

290 return hidden_states

291

vllm/model_executor/models/starcoder2.py

+6 -20

@@ -19,13 +19,13 @@

19 # See the License for the specific language governing permissions and

20 # limitations under the License.

21 """ PyTorch Starcoder2 model."""

22 -from typing import Iterable, List, Optional, Set, Tuple, Union

22 +from typing import Iterable, Optional, Set, Tuple, Union

23

24 import torch

25 from torch import nn

26 from transformers import Starcoder2Config

27

28 -from vllm.attention import Attention, AttentionMetadata

28 +from vllm.attention import Attention

29 from vllm.compilation.decorators import support_torch_compile

30 from vllm.config import CacheConfig, VllmConfig

31 from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size

@@ -118,13 +118,11 @@ class Starcoder2Attention(nn.Module):

118 self,

119 positions: torch.Tensor,

120 hidden_states: torch.Tensor,

121 - kv_cache: torch.Tensor,

122 - attn_metadata: AttentionMetadata,

121 ) -> torch.Tensor:

122 qkv, _ = self.qkv_proj(hidden_states)

123 q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)

124 q, k = self.rotary_emb(positions, q, k)

127 - attn_output = self.attn(q, k, v, kv_cache, attn_metadata)

125 + attn_output = self.attn(q, k, v)

126 output, _ = self.o_proj(attn_output)

127 return output

128

@@ -184,8 +182,6 @@ class Starcoder2DecoderLayer(nn.Module):

182 self,

183 positions: torch.Tensor,

184 hidden_states: torch.Tensor,

187 - kv_cache: torch.Tensor,

188 - attn_metadata: AttentionMetadata,

185 ) -> torch.Tensor:

186 # Self Attention

187 residual = hidden_states

@@ -193,8 +189,6 @@ class Starcoder2DecoderLayer(nn.Module):

189 hidden_states = self.self_attn(

190 positions=positions,

191 hidden_states=hidden_states,

196 - kv_cache=kv_cache,

197 - attn_metadata=attn_metadata,

192 )

193 hidden_states = residual + hidden_states

194

@@ -246,8 +240,6 @@ class Starcoder2Model(nn.Module):

240 self,

241 input_ids: torch.Tensor,

242 positions: torch.Tensor,

249 - kv_caches: List[torch.Tensor],

250 - attn_metadata: AttentionMetadata,

243 intermediate_tensors: Optional[IntermediateTensors],

244 inputs_embeds: Optional[torch.Tensor] = None,

245 ) -> Union[torch.Tensor, IntermediateTensors]:

@@ -259,11 +251,8 @@ class Starcoder2Model(nn.Module):

251 else:

252 assert intermediate_tensors is not None

253 hidden_states = intermediate_tensors["hidden_states"]

262 - for i in range(self.start_layer, self.end_layer):

263 - layer = self.layers[i]

264 - hidden_states = layer(positions, hidden_states,

265 - kv_caches[i - self.start_layer],

266 - attn_metadata)

254 + for layer in self.layers[self.start_layer:self.end_layer]:

255 + hidden_states = layer(positions, hidden_states)

256 if not get_pp_group().is_last_rank:

257 return IntermediateTensors({"hidden_states": hidden_states})

258 hidden_states = self.norm(hidden_states)

@@ -306,13 +295,10 @@ class Starcoder2ForCausalLM(nn.Module, SupportsPP):

295 self,

296 input_ids: torch.Tensor,

297 positions: torch.Tensor,

309 - kv_caches: List[torch.Tensor],

310 - attn_metadata: AttentionMetadata,

298 intermediate_tensors: Optional[IntermediateTensors] = None,

299 inputs_embeds: Optional[torch.Tensor] = None,

300 ) -> Union[torch.Tensor, IntermediateTensors]:

314 - hidden_states = self.model(input_ids, positions, kv_caches,

315 - attn_metadata, intermediate_tensors,

301 + hidden_states = self.model(input_ids, positions, intermediate_tensors,

302 inputs_embeds)

303 return hidden_states

304

vllm/model_executor/models/transformers.py

+2 -11

@@ -22,7 +22,7 @@ from torch import nn

22 from transformers import AutoModel, PreTrainedModel

23 from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS

24

25 -from vllm.attention import Attention, AttentionMetadata

25 +from vllm.attention import Attention

26 from vllm.config import VllmConfig

27 from vllm.distributed import get_tensor_model_parallel_world_size

28 from vllm.distributed.utils import divide

@@ -59,7 +59,6 @@ def vllm_flash_attention_forward(

59 # Transformers kwargs

60 scaling: Optional[float] = None,

61 # vLLM kwargs

62 - attn_metadata: Optional[AttentionMetadata] = None,

62 attention_instances: Optional[list[Attention]] = None,

63 **kwargs):

64 self_attn = attention_instances[module.layer_idx]

@@ -68,12 +67,7 @@ def vllm_flash_attention_forward(

67 hidden = query.shape[-2]

68 query, key, value = (x.transpose(1, 2) for x in (query, key, value))

69 query, key, value = (x.reshape(hidden, -1) for x in (query, key, value))

71 - return self_attn.forward(

72 - query,

73 - key,

74 - value,

75 - kv_cache=None, # argument not used

76 - attn_metadata=attn_metadata), None

70 + return self_attn.forward(query, key, value), None

71

72

73 ALL_ATTENTION_FUNCTIONS["vllm"] = vllm_flash_attention_forward

@@ -251,8 +245,6 @@ class TransformersModel(nn.Module, SupportsQuant):

245 self,

246 input_ids: torch.Tensor,

247 positions: torch.Tensor,

254 - kv_caches: list[torch.Tensor], # argument not used

255 - attn_metadata: AttentionMetadata,

248 intermediate_tensors: Optional[IntermediateTensors] = None,

249 inputs_embeds: Optional[torch.Tensor] = None,

250 ) -> Union[torch.Tensor, IntermediateTensors]:

@@ -260,7 +252,6 @@ class TransformersModel(nn.Module, SupportsQuant):

252 input_ids[None, ...],

253 use_cache=False,

254 position_ids=positions[None, ...],

263 - attn_metadata=attn_metadata,

255 intermediate_tensors=intermediate_tensors,

256 attention_instances=self.attention_instances,

257 return_dict=False)[0][0, ...] # we remove batch dimension for now

vllm/model_executor/models/ultravox.py

+5 -12

@@ -4,8 +4,8 @@

4 """PyTorch Ultravox model."""

5 import math

6 from functools import cached_property

7 -from typing import (Any, Iterable, List, Literal, Mapping, Optional, Set,

8 - Tuple, TypedDict, Union)

7 +from typing import (Any, Iterable, Literal, Mapping, Optional, Set, Tuple,

8 + TypedDict, Union)

9

10 import torch

11 import torch.utils.checkpoint

@@ -16,8 +16,8 @@ from transformers.models.whisper import WhisperFeatureExtractor

16 from transformers.models.whisper.modeling_whisper import WhisperEncoder

17

18 from vllm import envs

19 -from vllm.attention import AttentionMetadata

19 from vllm.config import VllmConfig

20 +from vllm.forward_context import get_forward_context

21 from vllm.model_executor.layers.activation import MulAndSilu, get_act_fn

22 from vllm.model_executor.layers.layernorm import RMSNorm

23 from vllm.model_executor.layers.sampler import SamplerOutput, get_sampler

@@ -495,13 +495,13 @@ class UltravoxModel(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA):

495 self,

496 input_ids: torch.Tensor,

497 multimodal_embeddings: Optional[NestedTensors] = None,

498 - attn_metadata: Optional[AttentionMetadata] = None,

498 ) -> torch.Tensor:

499 inputs_embeds = self.language_model.get_input_embeddings(input_ids)

500 if multimodal_embeddings is not None:

501

502 # TODO(ywang96): remove this block after v0 is deprecated.

503 if not envs.VLLM_USE_V1:

504 + attn_metadata = get_forward_context().attn_metadata

505 merge_multimodal_embeddings_from_map(

506 inputs_embeds, multimodal_embeddings,

507 attn_metadata.multi_modal_placeholder_index_maps["audio"])

@@ -514,8 +514,6 @@ class UltravoxModel(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA):

514 def forward(self,

515 input_ids: torch.Tensor,

516 positions: torch.Tensor,

517 - kv_caches: List[torch.Tensor],

518 - attn_metadata: AttentionMetadata,

517 intermediate_tensors: Optional[torch.Tensor] = None,

518 inputs_embeds: Optional[torch.Tensor] = None,

519 **kwargs) -> Union[torch.Tensor, IntermediateTensors]:

@@ -540,17 +538,12 @@ class UltravoxModel(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA):

538 elif inputs_embeds is None:

539 multimodal_embeddings = self.get_multimodal_embeddings(**kwargs)

540

543 - # TODO(ywang96): remove attn_metadata from get_input_embeddings

544 - # after v0 is deprecated

541 inputs_embeds = self.get_input_embeddings(input_ids,

546 - multimodal_embeddings,

547 - attn_metadata)

542 + multimodal_embeddings)

543 input_ids = None

544

545 hidden_states = self.language_model.model(input_ids,

546 positions,

552 - kv_caches,

553 - attn_metadata,

547 intermediate_tensors,

548 inputs_embeds=inputs_embeds)

549 return hidden_states

vllm/model_executor/models/whisper.py

+13 -76

@@ -10,7 +10,7 @@ from transformers import (BatchFeature, WhisperConfig, WhisperFeatureExtractor,

10 WhisperProcessor)

11 from transformers.models.whisper.modeling_whisper import sinusoids

12

13 -from vllm.attention import Attention, AttentionMetadata, AttentionType

13 +from vllm.attention import Attention, AttentionType

14 from vllm.config import CacheConfig, VllmConfig

15 from vllm.distributed import get_tensor_model_parallel_world_size

16 from vllm.logger import init_logger

@@ -134,13 +134,11 @@ class WhisperAttention(nn.Module):

134 def forward(

135 self,

136 hidden_states: torch.Tensor,

137 - kv_cache: torch.Tensor,

138 - attn_metadata: AttentionMetadata,

137 ):

138 qkv, _ = self.qkv_proj(hidden_states)

139 q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)

140

143 - attn_output = self.attn(q, k, v, kv_cache, attn_metadata)

141 + attn_output = self.attn(q, k, v)

142

143 output, _ = self.out_proj(attn_output)

144

@@ -196,8 +194,6 @@ class WhisperCrossAttention(WhisperAttention):

194 self,

195 hidden_states: torch.Tensor,

196 encoder_hidden_states: Optional[torch.Tensor],

199 - kv_cache: torch.Tensor,

200 - attn_metadata: AttentionMetadata,

197 ):

198 q, _ = self.q_proj(hidden_states)

199

@@ -209,13 +205,7 @@ class WhisperCrossAttention(WhisperAttention):

205 else:

206 k = v = None

207

212 - attn_output = self.attn(

213 - q,

214 - k,

215 - v,

216 - kv_cache,

217 - attn_metadata,

218 - )

208 + attn_output = self.attn(q, k, v)

209

210 output, _ = self.out_proj(attn_output)

211

@@ -285,16 +275,10 @@ class WhisperEncoderLayer(nn.Module):

275 def forward(

276 self,

277 hidden_states: torch.Tensor,

288 - kv_cache: torch.Tensor,

289 - attn_metadata: AttentionMetadata,

278 ):

279 residual = hidden_states

280 hidden_states = self.self_attn_layer_norm(hidden_states)

293 - hidden_states = self.self_attn(

294 - hidden_states=hidden_states,

295 - kv_cache=kv_cache,

296 - attn_metadata=attn_metadata,

297 - )

281 + hidden_states = self.self_attn(hidden_states=hidden_states)

282 hidden_states = residual + hidden_states

283 residual = hidden_states

284 hidden_states = self.final_layer_norm(hidden_states)

@@ -348,14 +332,10 @@ class WhisperDecoderLayer(nn.Module):

332 self,

333 hidden_states: torch.Tensor,

334 encoder_hidden_states: Optional[torch.Tensor],

351 - kv_cache: torch.Tensor,

352 - attn_metadata: AttentionMetadata,

335 ):

336 residual = hidden_states

337 hidden_states = self.self_attn_layer_norm(hidden_states)

356 - hidden_states = self.self_attn(hidden_states=hidden_states,

357 - kv_cache=kv_cache,

358 - attn_metadata=attn_metadata)

338 + hidden_states = self.self_attn(hidden_states=hidden_states)

339 hidden_states = residual + hidden_states

340

341 residual = hidden_states

@@ -363,8 +343,6 @@ class WhisperDecoderLayer(nn.Module):

343 hidden_states = self.encoder_attn(

344 hidden_states=hidden_states,

345 encoder_hidden_states=encoder_hidden_states,

366 - kv_cache=kv_cache,

367 - attn_metadata=attn_metadata,

346 )

347 hidden_states = residual + hidden_states

348

@@ -411,12 +389,7 @@ class WhisperEncoder(nn.Module):

389 self.embed_positions.weight.copy_(

390 sinusoids(*self.embed_positions.weight.shape))

391

414 - def forward(

415 - self,

416 - input_features: Union[torch.Tensor, List[torch.Tensor]],

417 - kv_caches: List[torch.Tensor],

418 - attn_metadata: AttentionMetadata,

419 - ):

392 + def forward(self, input_features: Union[torch.Tensor, List[torch.Tensor]]):

393 hidden_states = []

394 for features in input_features:

395 embeds = nn.functional.gelu(self.conv1(features))

@@ -426,12 +399,8 @@ class WhisperEncoder(nn.Module):

399 hidden_states.append(embeds)

400 hidden_states = torch.cat(hidden_states)

401

429 - for idx, encoder_layer in enumerate(self.layers):

430 - hidden_states = encoder_layer(

431 - hidden_states,

432 - kv_cache=kv_caches[idx],

433 - attn_metadata=attn_metadata,

434 - )

402 + for encoder_layer in self.layers:

403 + hidden_states = encoder_layer(hidden_states)

404

405 hidden_states = self.layer_norm(hidden_states)

406 return hidden_states

@@ -466,19 +435,15 @@ class WhisperDecoder(nn.Module):

435 input_ids,

436 positions: torch.Tensor,

437 encoder_hidden_states: Optional[torch.Tensor],

469 - kv_caches: List[torch.Tensor],

470 - attn_metadata: AttentionMetadata,

438 ):

439 inputs_embeds = self.get_input_embeddings(input_ids)

440 positions = self.embed_positions(positions)

441 hidden_states = inputs_embeds + positions

442

476 - for idx, decoder_layer in enumerate(self.layers):

443 + for decoder_layer in self.layers:

444 hidden_states = decoder_layer(

445 hidden_states,

446 encoder_hidden_states=encoder_hidden_states,

480 - kv_cache=kv_caches[idx],

481 - attn_metadata=attn_metadata,

447 )

448

449 hidden_states = self.layer_norm(hidden_states)

@@ -505,36 +470,22 @@ class WhisperModel(nn.Module):

470 input_features: Optional[Union[torch.Tensor, List[torch.Tensor]]],

471 input_ids: Optional[torch.Tensor],

472 positions: torch.Tensor,

508 - kv_caches: List[torch.Tensor],

509 - attn_metadata: AttentionMetadata,

473 ) -> torch.Tensor:

511 - encoder_outputs = self.get_encoder_outputs(

512 - input_features,

513 - kv_caches=kv_caches,

514 - attn_metadata=attn_metadata,

515 - )

474 + encoder_outputs = self.get_encoder_outputs(input_features)

475 decoder_outputs = self.decoder(

476 input_ids=input_ids,

477 positions=positions,

478 encoder_hidden_states=encoder_outputs,

520 - kv_caches=kv_caches,

521 - attn_metadata=attn_metadata,

479 )

480 return decoder_outputs

481

482 def get_encoder_outputs(

483 self,

484 input_features: Optional[Union[torch.Tensor, List[torch.Tensor]]],

528 - kv_caches: List[torch.Tensor],

529 - attn_metadata: AttentionMetadata,

485 ) -> Optional[torch.Tensor]:

486 if input_features is None:

487 return None

533 - return self.encoder(

534 - input_features,

535 - kv_caches=kv_caches,

536 - attn_metadata=attn_metadata,

537 - )

488 + return self.encoder(input_features)

489

490 def load_weights(self, weights: Iterable[Tuple[str,

491 torch.Tensor]]) -> Set[str]:

@@ -733,8 +684,6 @@ class WhisperForConditionalGeneration(nn.Module, SupportsTranscription,

684 self,

685 input_ids: torch.Tensor,

686 positions: torch.Tensor,

736 - kv_caches: List[torch.Tensor],

737 - attn_metadata: AttentionMetadata,

687 **kwargs,

688 ) -> torch.Tensor:

689 audio_input = self._parse_and_validate_audio_input(**kwargs)

@@ -742,31 +691,19 @@ class WhisperForConditionalGeneration(nn.Module, SupportsTranscription,

691 input_features=audio_input["input_features"],

692 input_ids=input_ids,

693 positions=positions,

745 - kv_caches=kv_caches,

746 - attn_metadata=attn_metadata,

694 )

695 return decoder_outputs

696

750 - def get_multimodal_embeddings(

751 - self,

752 - kv_caches: List[torch.Tensor],

753 - attn_metadata: AttentionMetadata,

754 - **kwargs,

755 - ) -> Optional[NestedTensors]:

697 + def get_multimodal_embeddings(self, **kwargs) -> Optional[NestedTensors]:

698 # TODO: This method does not obey the interface for SupportsMultiModal.

699 # Refactor this once encoder/decoder support is implemented in V1.

700 audio_input = self._parse_and_validate_audio_input(**kwargs)

759 - return self.model.get_encoder_outputs(

760 - audio_input["input_features"],

761 - kv_caches=kv_caches,

762 - attn_metadata=attn_metadata,

763 - )

701 + return self.model.get_encoder_outputs(audio_input["input_features"])

702

703 def get_input_embeddings(

704 self,

705 input_ids: torch.Tensor,

706 multimodal_embeddings: Optional[NestedTensors] = None,

769 - attn_metadata: Optional[AttentionMetadata] = None,

707 ) -> torch.Tensor:

708 # TODO: This method just returns the decoder sequence embeddings since

709 # Whisper does not have encoder text tokens. Refactor this once

vllm/spec_decode/draft_model_runner.py

-2

@@ -288,8 +288,6 @@ class TP1DraftModelRunner(ModelRunnerWrapperBase):

288 hidden_states = model_executable(

289 input_ids=model_input.input_tokens,

290 positions=model_input.input_positions,

291 - kv_caches=kv_caches,

292 - attn_metadata=model_input.attn_metadata,

291 intermediate_tensors=intermediate_tensors,

292 **MultiModalKwargs.as_kwargs(multi_modal_kwargs,

293 device=self.device),

vllm/v1/worker/gpu_model_runner.py

+1 -21

@@ -939,8 +939,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):

939 hidden_states = self.model(

940 input_ids=input_ids,

941 positions=positions,

942 - kv_caches=self.kv_caches,

943 - attn_metadata=None,

942 intermediate_tensors=intermediate_tensors,

943 inputs_embeds=inputs_embeds,

944 )

@@ -1137,11 +1135,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):

1135 def _dummy_run(

1136 self,

1137 num_tokens: int,

1140 - kv_caches: Optional[List[torch.Tensor]] = None,

1138 ) -> torch.Tensor:

1139 model = self.model

1143 - if kv_caches is None:

1144 - kv_caches = self.kv_caches

1140 if self.is_multimodal_model:

1141 input_ids = None

1142 inputs_embeds = self.inputs_embeds[:num_tokens]

@@ -1172,26 +1167,12 @@ class GPUModelRunner(LoRAModelRunnerMixin):

1167 hidden_states = model(

1168 input_ids=input_ids,

1169 positions=positions,

1175 - kv_caches=kv_caches,

1176 - attn_metadata=None,

1170 intermediate_tensors=intermediate_tensors,

1171 inputs_embeds=inputs_embeds,

1172 )

1173 return hidden_states

1174

1175 def profile_run(self) -> None:

1183 - # use an empty tensor instead of `None`` to force Dynamo to pass

1184 - # it by reference, rather by specializing on the value `None`.

1185 - # the `dtype` argument does not matter, and we use `float32` as

1186 - # a placeholder (it has wide hardware support).

1187 - # it is important to create tensors inside the loop, rather than

1188 - # multiplying the list, to avoid Dynamo from treating them as

1189 - # tensor aliasing.

1190 - dummy_kv_caches = [

1191 - torch.tensor((), dtype=torch.float32, device=self.device)

1192 - for _ in range(self.num_attn_layers)

1193 - ]

1194 -

1176 # Profile with multimodal encoder & encoder cache.

1177 # TODO: handle encoder-decoder models once we support them.

1178 if (self.is_multimodal_model and self.max_num_encoder_input_tokens > 0

@@ -1302,8 +1283,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

1283 with self.maybe_profile_with_lora(self.lora_config,

1284 num_scheduled_tokens):

1285 # Trigger compilation for general shape.

1305 - hidden_states = self._dummy_run(self.max_num_tokens,

1306 - dummy_kv_caches)

1286 + hidden_states = self._dummy_run(self.max_num_tokens)

1287 if get_pp_group().is_last_rank:

1288 hidden_states = hidden_states[logit_indices]

1289 logits = self.model.compute_logits(hidden_states, None)

vllm/v1/worker/tpu_model_runner.py

+5 -14

@@ -13,11 +13,10 @@ import torch.nn as nn

13 import torch_xla.core.xla_model as xm

14 import torch_xla.runtime as xr

15

16 -from vllm.attention import AttentionMetadata

16 from vllm.attention.backends.abstract import AttentionType

17 from vllm.attention.layer import Attention

18 from vllm.config import VllmConfig

20 -from vllm.forward_context import set_forward_context

19 +from vllm.forward_context import get_forward_context, set_forward_context

20 from vllm.logger import init_logger

21 from vllm.model_executor.model_loader import get_model

22 from vllm.sampling_params import SamplingType

@@ -623,7 +622,6 @@ class TPUModelRunner:

622 assert self.model is not None

623 selected_token_ids = self.model(prompt_data.input_tokens,

624 prompt_data.input_positions,

626 - prompt_data.attn_metadata,

625 self.kv_caches)

626

627 # In parallel to TPU execution, prepare the next iteration

@@ -662,7 +660,6 @@ class TPUModelRunner:

660 assert self.model is not None

661 selected_token_ids = self.model(decode_data.input_tokens,

662 decode_data.input_positions,

665 - decode_data.attn_metadata,

663 self.kv_caches)

664

665 # Transfer sampled tokens from TPU to CPU

@@ -839,7 +836,7 @@ class TPUModelRunner:

836

837 with set_forward_context(attn_metadata, self.vllm_config, 0):

838 assert self.model is not None

842 - self.model(token_ids, position_ids, attn_metadata, kv_caches)

839 + self.model(token_ids, position_ids, kv_caches)

840

841 def capture_model(self) -> None:

842 """Compile the model."""

@@ -963,7 +960,6 @@ class ModelWrapperV1(nn.Module):

960 self,

961 token_ids: torch.Tensor,

962 position_ids: torch.Tensor,

966 - attn_metadata: AttentionMetadata,

963 kv_caches: List[Tuple[torch.Tensor, torch.Tensor]],

964 ) -> torch.Tensor:

965 """Executes the forward pass of the model and samples the next token.

@@ -971,7 +967,6 @@ class ModelWrapperV1(nn.Module):

967 Args:

968 token_ids: The input token IDs of shape [batch_size, seq_len].

969 position_ids: The input position IDs of shape [batch_size, seq_len].

974 - attn_metadata: The Pallas attention metadata.

970 input_lens: The actual input lengths of shape [batch_size].

971 t: The sampling temperature of shape [batch_size].

972 p: The top-p probability of shape [batch_size].

@@ -980,7 +975,8 @@ class ModelWrapperV1(nn.Module):

975 memory profiling at initialization.

976 """

977 # Skip this in memory profiling at initialization.

983 - if attn_metadata is not None and kv_caches[0][0].numel() > 0:

978 + if kv_caches[0][0].numel() > 0:

979 + attn_metadata = get_forward_context().attn_metadata

980 # index_copy_(slot_mapping) only works when the inserted dimension

981 # is 0. However, the KV cache in the Pallas backend has the shape

982 # [num_kv_heads, num_blocks, block_size, head_size]. To make it

@@ -1001,12 +997,7 @@ class ModelWrapperV1(nn.Module):

997 attn_metadata.slot_mapping = slot_mapping

998

999 assert self.model is not None

1004 - hidden_states = self.model(

1005 - token_ids,

1006 - position_ids,

1007 - kv_caches,

1008 - attn_metadata,

1009 - )

1000 + hidden_states = self.model(token_ids, position_ids)

1001

1002 hidden_states = hidden_states.flatten(0, 1)

1003 logits = self.model.compute_logits(hidden_states, None)

vllm/worker/cpu_enc_dec_model_runner.py

-4

@@ -297,10 +297,6 @@ class CPUEncoderDecoderModelRunner(

297 model_input.encoder_input_tokens,

298 "encoder_positions":

299 model_input.encoder_input_positions,

300 - "kv_caches":

301 - kv_caches,

302 - "attn_metadata":

303 - model_input.attn_metadata,

300 **MultiModalKwargs.as_kwargs(model_input.multi_modal_kwargs or {},

301 device=self.device),

302 "intermediate_tensors":

vllm/worker/cpu_model_runner.py

-2

@@ -654,8 +654,6 @@ class CPUModelRunner(CPUModelRunnerBase[ModelInputForCPUWithSamplingMetadata]):

654 hidden_states = model_executable(

655 input_ids=model_input.input_tokens,

656 positions=model_input.input_positions,

657 - kv_caches=kv_caches,

658 - attn_metadata=model_input.attn_metadata,

657 intermediate_tensors=intermediate_tensors,

658 **execute_model_kwargs,

659 **multimodal_kwargs,

vllm/worker/cpu_pooling_model_runner.py

-14

@@ -41,16 +41,6 @@ class CPUPoolingModelRunner(

41 raise ValueError(

42 "CPU worker does not support multi-step execution.")

43

44 - num_layers = self.model_config.get_num_layers(self.parallel_config)

45 - # use an empty tensor instead of `None`` to force Dynamo to pass

46 - # it by reference, rather by specializing on the value ``None``.

47 - # the `dtype` argument does not matter, and we use `float32` as

48 - # a placeholder (it has wide hardware support).

49 - kv_caches = [

50 - torch.tensor([], dtype=torch.float32, device=self.device)

51 - for _ in range(num_layers)

52 - ]

53 -

44 model_executable = self.model

45 cross_enc_kwargs = {}

46 if model_input.token_type_ids is not None:

@@ -60,10 +50,6 @@ class CPUPoolingModelRunner(

50 model_input.input_tokens,

51 "positions":

52 model_input.input_positions,

63 - "kv_caches":

64 - kv_caches,

65 - "attn_metadata":

66 - model_input.attn_metadata,

53 **MultiModalKwargs.as_kwargs(model_input.multi_modal_kwargs or {},

54 device=self.device),

55 **cross_enc_kwargs,

vllm/worker/enc_dec_model_runner.py

+1 -13

@@ -184,8 +184,6 @@ class EncoderDecoderModelRunner(GPUModelRunnerBase[EncoderDecoderModelInput]):

184 positions=model_input.input_positions,

185 encoder_input_ids=model_input.encoder_input_tokens,

186 encoder_positions=model_input.encoder_input_positions,

187 - kv_caches=kv_caches,

188 - attn_metadata=model_input.attn_metadata,

187 intermediate_tensors=intermediate_tensors,

188 **MultiModalKwargs.as_kwargs(multi_modal_kwargs,

189 device=self.device),

@@ -324,21 +322,11 @@ class EncoderDecoderModelRunner(GPUModelRunnerBase[EncoderDecoderModelInput]):

322 or encoder_dummy_data.multi_modal_placeholders)

323 seqs.append(seq)

324

327 - # Run the model with the dummy inputs.

328 - num_layers = self.model_config.get_num_layers(self.parallel_config)

329 - # use an empty tensor instead of `None`` to force Dynamo to pass

330 - # it by reference, rather by specializing on the value ``None``.

331 - # the `dtype` argument does not matter, and we use `float32` as

332 - # a placeholder (it has wide hardware support).

333 - kv_caches = [

334 - torch.tensor([], dtype=torch.float32, device=self.device)

335 - for _ in range(num_layers)

336 - ]

325 finished_requests_ids = [seq.request_id for seq in seqs]

326 model_input = self.prepare_model_input(

327 seqs, finished_requests_ids=finished_requests_ids)

328 intermediate_tensors = None

341 - self.execute_model(model_input, kv_caches, intermediate_tensors)

329 + self.execute_model(model_input, None, intermediate_tensors)

330 torch.cuda.synchronize()

331 return

332

vllm/worker/hpu_model_runner.py

+15 -21

@@ -384,11 +384,12 @@ class HpuModelAdapter:

384 if 'virtual_engine' in kwargs:

385 virtual_engine = kwargs.pop('virtual_engine')

386 input_ids = kwargs['input_ids']

387 - kwargs['attn_metadata'] = self._update_metadata(

388 - kwargs['attn_metadata'], input_ids.size(0), input_ids.size(1),

389 - input_ids.device, self.dtype)

387 + attn_metadata = self._update_metadata(kwargs.pop('attn_metadata'),

388 + input_ids.size(0),

389 + input_ids.size(1),

390 + input_ids.device, self.dtype)

391 LoraMask.setLoraMask(kwargs.pop('lora_mask'))

391 - with set_forward_context(kwargs['attn_metadata'], self.vllm_config,

392 + with set_forward_context(attn_metadata, self.vllm_config,

393 virtual_engine):

394 hidden_states = self.model(*args, **kwargs)

395 hidden_states = hidden_states.view(-1, hidden_states.shape[-1])

@@ -1346,15 +1347,13 @@ class HPUModelRunnerBase(ModelRunnerBase[TModelInputForHPU]):

1347 max_seq_len = self.bucketing_global_state.prompt_seq_bucket_cfg[-1]

1348 max_batch_size = min(self.max_num_batched_tokens // max_seq_len,

1349 self.scheduler_config.max_num_seqs)

1349 - self.warmup_scenario(max_batch_size, max_seq_len, True, kv_caches,

1350 - False, True)

1350 + self.warmup_scenario(max_batch_size, max_seq_len, True, False, True)

1351 return

1352

1353 def warmup_scenario(self,

1354 batch_size,

1355 seq_len,

1356 is_prompt,

1357 - kv_caches,

1357 is_pt_profiler_run=False,

1358 is_lora_profile_run=False) -> None:

1359 use_graphs = self._use_graphs(batch_size, seq_len, is_prompt)

@@ -1418,7 +1417,7 @@ class HPUModelRunnerBase(ModelRunnerBase[TModelInputForHPU]):

1417 profiler.start()

1418 for _ in range(times):

1419 inputs = self.prepare_model_input(seqs)

1421 - self.execute_model(inputs, kv_caches, warmup_mode=True)

1420 + self.execute_model(inputs, None, warmup_mode=True)

1421 torch.hpu.synchronize()

1422 if profiler:

1423 profiler.step()

@@ -1470,17 +1469,16 @@ class HPUModelRunnerBase(ModelRunnerBase[TModelInputForHPU]):

1469 f"free_mem:{free_mem}")

1470 logger.info(msg)

1471

1473 - def warmup_all_buckets(self, buckets, is_prompt, kv_caches):

1472 + def warmup_all_buckets(self, buckets, is_prompt):

1473 for i, (batch_size, seq_len) in enumerate(reversed(buckets)):

1474 self.log_warmup('Prompt' if is_prompt else 'Decode', i,

1475 len(buckets), batch_size, seq_len)

1477 - self.warmup_scenario(batch_size, seq_len, is_prompt, kv_caches)

1476 + self.warmup_scenario(batch_size, seq_len, is_prompt)

1477

1478 def warmup_graphs(self,

1479 strategy,

1480 buckets,

1481 is_prompt,

1483 - kv_caches,

1482 available_mem,

1483 starting_mem=0,

1484 total_batch_seq=0.001):

@@ -1512,7 +1510,7 @@ class HPUModelRunnerBase(ModelRunnerBase[TModelInputForHPU]):

1510 self.graphed_buckets.add(graphed_bucket)

1511 self.log_warmup(phase, idx, num_candidates, batch_size, seq_len)

1512 with HabanaMemoryProfiler() as mem_prof:

1515 - self.warmup_scenario(batch_size, seq_len, is_prompt, kv_caches)

1513 + self.warmup_scenario(batch_size, seq_len, is_prompt)

1514 used_mem = align_workers(mem_prof.consumed_device_memory,

1515 torch.distributed.ReduceOp.MAX)

1516 available_mem -= used_mem

@@ -1542,8 +1540,7 @@ class HPUModelRunnerBase(ModelRunnerBase[TModelInputForHPU]):

1540 graphs = graph == 't'

1541 if graphs:

1542 self.graphed_buckets.add((int(bs), int(seq_len), is_prompt))

1545 - self.warmup_scenario(int(bs), int(seq_len), is_prompt, kv_caches,

1546 - True)

1543 + self.warmup_scenario(int(bs), int(seq_len), is_prompt, True)

1544 raise AssertionError("Finished profiling")

1545 if self.skip_warmup:

1546 logger.info("Skipping warmup...")

@@ -1608,9 +1605,9 @@ class HPUModelRunnerBase(ModelRunnerBase[TModelInputForHPU]):

1605 with compile_only_mode_context(

1606 ) if can_use_compile_only_mode else contextlib.nullcontext():

1607 self.warmup_all_buckets(self.bucketing_global_state.prompt_buckets,

1611 - True, kv_caches)

1608 + True)

1609 self.warmup_all_buckets(self.bucketing_global_state.decode_buckets,

1613 - False, kv_caches)

1610 + False)

1611

1612 if not self.enforce_eager and htorch.utils.internal.is_lazy():

1613 assert self.mem_margin is not None, \

@@ -1641,11 +1638,11 @@ class HPUModelRunnerBase(ModelRunnerBase[TModelInputForHPU]):

1638 mem_post_prompt, prompt_batch_seq, prompt_captured_all = \

1639 self.warmup_graphs(

1640 prompt_strategy, self.bucketing_global_state.prompt_buckets,

1644 - True, kv_caches, prompt_available_memory)

1641 + True, prompt_available_memory)

1642 mem_post_decode, decode_batch_seq, decode_captured_all = \

1643 self.warmup_graphs(

1644 decode_strategy, self.bucketing_global_state.decode_buckets,

1648 - False, kv_caches, decode_available_memory)

1645 + False, decode_available_memory)

1646

1647 # Not all prompt buckets were captured, but all decode buckets

1648 # were captured and we have some free graph-allocated space

@@ -1656,7 +1653,6 @@ class HPUModelRunnerBase(ModelRunnerBase[TModelInputForHPU]):

1653 self.warmup_graphs(

1654 prompt_strategy,

1655 self.bucketing_global_state.prompt_buckets, True,

1659 - kv_caches,

1656 graph_free_mem - mem_post_prompt - mem_post_decode,

1657 mem_post_prompt, prompt_batch_seq))

1658

@@ -1669,7 +1665,6 @@ class HPUModelRunnerBase(ModelRunnerBase[TModelInputForHPU]):

1665 mem_post_decode, _, _ = self.warmup_graphs(

1666 decode_strategy,

1667 self.bucketing_global_state.decode_buckets, False,

1672 - kv_caches,

1668 graph_free_mem - mem_post_prompt - mem_post_decode,

1669 mem_post_decode, decode_batch_seq)

1670

@@ -1982,7 +1977,6 @@ class HPUModelRunner(HPUModelRunnerBase[ModelInputForHPUWithSamplingMetadata]):

1977 execute_model_kwargs = {

1978 "input_ids": input_tokens,

1979 "positions": input_positions,

1985 - "kv_caches": kv_caches,

1980 "attn_metadata": self.trim_attn_metadata(attn_metadata),

1981 "intermediate_tensors": intermediate_tensors,

1982 "lora_mask": lora_mask,

vllm/worker/model_runner.py

+2 -11

@@ -26,7 +26,7 @@ from vllm.core.scheduler import SchedulerOutputs

26 from vllm.distributed import get_kv_transfer_group, get_pp_group

27 from vllm.distributed.parallel_state import (get_tensor_model_parallel_rank,

28 graph_capture)

29 -from vllm.forward_context import set_forward_context

29 +from vllm.forward_context import get_forward_context, set_forward_context

30 from vllm.inputs import INPUT_REGISTRY, InputRegistry

31 from vllm.logger import init_logger

32 from vllm.lora.layers import LoRAMapping

@@ -1727,8 +1727,6 @@ class ModelRunner(GPUModelRunnerBase[ModelInputForGPUWithSamplingMetadata]):

1727 hidden_or_intermediate_states = model_executable(

1728 input_ids=model_input.input_tokens,

1729 positions=model_input.input_positions,

1730 - kv_caches=kv_caches,

1731 - attn_metadata=model_input.attn_metadata,

1730 intermediate_tensors=intermediate_tensors,

1731 **MultiModalKwargs.as_kwargs(multi_modal_kwargs,

1732 device=self.device),

@@ -1913,8 +1911,6 @@ class CUDAGraphRunner(nn.Module):

1911 self.model(

1912 input_ids=input_ids,

1913 positions=positions,

1916 - kv_caches=kv_caches,

1917 - attn_metadata=attn_metadata,

1914 intermediate_tensors=intermediate_inputs,

1915 **kwargs,

1916 )

@@ -1927,8 +1923,6 @@ class CUDAGraphRunner(nn.Module):

1923 output_hidden_or_intermediate_states = self.model(

1924 input_ids=input_ids,

1925 positions=positions,

1930 - kv_caches=kv_caches,

1931 - attn_metadata=attn_metadata,

1926 intermediate_tensors=intermediate_inputs,

1927 **kwargs,

1928 )

@@ -1976,13 +1970,10 @@ class CUDAGraphRunner(nn.Module):

1970 self,

1971 input_ids: torch.Tensor,

1972 positions: torch.Tensor,

1979 - kv_caches: List[torch.Tensor],

1980 - attn_metadata: AttentionMetadata,

1973 intermediate_tensors: Optional[IntermediateTensors],

1974 **kwargs,

1975 ) -> torch.Tensor:

1984 - # KV caches are fixed tensors, so we don't need to copy them.

1985 - del kv_caches

1976 + attn_metadata: AttentionMetadata = get_forward_context().attn_metadata

1977

1978 # Copy the input tensors to the input buffers.

1979 self.input_buffers["input_ids"].copy_(input_ids, non_blocking=True)

vllm/worker/multi_step_model_runner.py

+2 -2

@@ -476,7 +476,7 @@ class MultiStepModelRunner(GPUModelRunnerBase[StatefulModelInput]):

476 # path for warm up runs

477 if not model_input.is_multi_step:

478 return self._base_model_runner.execute_model(

479 - frozen_model_input, kv_caches, intermediate_tensors, num_steps)

479 + frozen_model_input, None, intermediate_tensors, num_steps)

480

481 # make sure we skip the sampler on the lask rank and only pythonize

482 # if CPU is ahead.

@@ -538,7 +538,7 @@ class MultiStepModelRunner(GPUModelRunnerBase[StatefulModelInput]):

538

539 # Execute the model

540 output = self._base_model_runner.execute_model(frozen_model_input,

541 - kv_caches,

541 + None,

542 intermediate_tensors,

543 num_steps=1)

544

vllm/worker/openvino_model_runner.py

-4

@@ -346,10 +346,6 @@ class OpenVINOModelRunner(ModelRunnerBase):

346 input_tokens,

347 "positions":

348 input_positions,

349 - "kv_caches":

350 - kv_caches,

351 - "attn_metadata":

352 - attn_metadata,

349 **MultiModalKwargs.as_kwargs(multi_modal_kwargs or {},

350 device=self.device),

351 }

vllm/worker/pooling_model_runner.py

-12

@@ -91,16 +91,6 @@ class PoolingModelRunner(

91 else:

92 model_executable = self.model

93

94 - num_layers = self.model_config.get_num_layers(self.parallel_config)

95 - # use an empty tensor instead of `None`` to force Dynamo to pass

96 - # it by reference, rather by specializing on the value ``None``.

97 - # the `dtype` argument does not matter, and we use `float32` as

98 - # a placeholder (it has wide hardware support).

99 - kv_caches = [

100 - torch.tensor([], dtype=torch.float32, device=self.device)

101 - for _ in range(num_layers)

102 - ]

103 -

94 multi_modal_kwargs = model_input.multi_modal_kwargs or {}

95 seqlen_agnostic_kwargs = {

96 "finished_requests_ids": model_input.finished_requests_ids,

@@ -121,8 +111,6 @@ class PoolingModelRunner(

111 hidden_or_intermediate_states = model_executable(

112 input_ids=model_input.input_tokens,

113 positions=model_input.input_positions,

124 - kv_caches=kv_caches,

125 - attn_metadata=model_input.attn_metadata,

114 intermediate_tensors=intermediate_tensors,

115 **MultiModalKwargs.as_kwargs(multi_modal_kwargs,

116 device=self.device),

vllm/worker/tpu_model_runner.py

+9 -15

@@ -15,7 +15,7 @@ import torch_xla.runtime as xr

15

16 from vllm.attention import AttentionMetadata, get_attn_backend

17 from vllm.config import VllmConfig

18 -from vllm.forward_context import set_forward_context

18 +from vllm.forward_context import get_forward_context, set_forward_context

19 from vllm.logger import init_logger

20 from vllm.model_executor.layers.sampler import SamplerOutput

21 from vllm.model_executor.model_loader import get_model

@@ -275,8 +275,8 @@ class TPUModelRunner(ModelRunnerBase[ModelInputForTPU]):

275 torch._dynamo.mark_dynamic(p, 0)

276 # Dummy run.

277 with set_forward_context(attn_metadata, self.vllm_config, 0):

278 - self.model(token_ids, position_ids, attn_metadata, input_lens, t,

279 - p, num_samples, kv_caches)

278 + self.model(token_ids, position_ids, input_lens, t, p, num_samples,

279 + kv_caches)

280

281 def warmup_model(

282 self,

@@ -679,8 +679,8 @@ class TPUModelRunner(ModelRunnerBase[ModelInputForTPU]):

679 self.vllm_config,

680 model_input.virtual_engine):

681 output_token_ids = self.model(token_ids, position_ids,

682 - attn_metadata, input_lens, t,

683 - p, model_input.num_samples,

682 + input_lens, t, p,

683 + model_input.num_samples,

684 kv_caches)

685 next_token_ids.append(output_token_ids[0])

686 start_idx = end_idx

@@ -730,8 +730,8 @@ class TPUModelRunner(ModelRunnerBase[ModelInputForTPU]):

730 self.vllm_config,

731 model_input.virtual_engine):

732 output_token_ids = self.model(token_ids, position_ids,

733 - attn_metadata, input_lens, t,

734 - p, model_input.num_samples,

733 + input_lens, t, p,

734 + model_input.num_samples,

735 kv_caches)

736 self.cached_step_outputs.append(output_token_ids)

737

@@ -777,7 +777,6 @@ class ModelWrapper(nn.Module):

777 self,

778 token_ids: torch.Tensor,

779 position_ids: torch.Tensor,

780 - attn_metadata: AttentionMetadata,

780 input_lens: torch.Tensor,

781 t: torch.Tensor,

782 p: torch.Tensor,

@@ -789,7 +788,6 @@ class ModelWrapper(nn.Module):

788 Args:

789 token_ids: The input token IDs of shape [batch_size, seq_len].

790 position_ids: The input position IDs of shape [batch_size, seq_len].

792 - attn_metadata: The Pallas attention metadata.

791 input_lens: The actual input lengths of shape [batch_size].

792 t: The sampling temperature of shape [batch_size].

793 p: The top-p probability of shape [batch_size].

@@ -802,6 +800,7 @@ class ModelWrapper(nn.Module):

800 start_indicies = torch.arange(

801 batch_size, dtype=torch.int32, device=input_lens.device) * seq_len

802 logits_indices = start_indicies + input_lens - 1

803 + attn_metadata = get_forward_context().attn_metadata

804

805 # FIXME(woosuk): This is a temporary hack to avoid using the existing

806 # sampler and sampling metadata.

@@ -833,12 +832,7 @@ class ModelWrapper(nn.Module):

832 slot_mapping = slot_mapping.flatten()

833 attn_metadata.slot_mapping = slot_mapping

834

836 - hidden_states = self.model(

837 - token_ids,

838 - position_ids,

839 - kv_caches,

840 - attn_metadata,

841 - )

835 + hidden_states = self.model(token_ids, position_ids)

836 hidden_states = hidden_states.flatten(0, 1)

837 logits = self.model.compute_logits(hidden_states, sampling_metadata)

838

vllm/worker/xpu_model_runner.py

+1 -12

@@ -484,15 +484,6 @@ class XPUModelRunner(ModelRunnerBase[ModelInputForXPUWithSamplingMetadata]):

484 multi_modal_placeholders=dummy_data.multi_modal_placeholders)

485 seqs.append(seq)

486

487 - # Run the model with the dummy inputs.

488 - num_layers = self.model_config.get_num_layers(self.parallel_config)

489 - # use an empty tensor instead of `None`` to force Dynamo to pass

490 - # it by reference, rather by specializing on the value ``None``.

491 - # the `dtype` argument does not matter, and we use `float32` as

492 - # a placeholder (it has wide hardware support).

493 - kv_caches = [

494 - torch.tensor([], dtype=torch.float32, device=self.device)

495 - ] * num_layers

487 finished_requests_ids = [seq.request_id for seq in seqs]

488 model_input = self.prepare_model_input(

489 seqs, finished_requests_ids=finished_requests_ids)

@@ -502,7 +493,7 @@ class XPUModelRunner(ModelRunnerBase[ModelInputForXPUWithSamplingMetadata]):

493 batch_size=batch_size,

494 dtype=self.model_config.dtype,

495 device=self.device)

505 - self.execute_model(model_input, kv_caches, intermediate_tensors)

496 + self.execute_model(model_input, None, intermediate_tensors)

497 torch.xpu.synchronize()

498 return

499

@@ -581,8 +572,6 @@ class XPUModelRunner(ModelRunnerBase[ModelInputForXPUWithSamplingMetadata]):

572 hidden_or_intermediate_states = model_executable(

573 input_ids=model_input.input_tokens,

574 positions=model_input.input_positions,

584 - kv_caches=kv_caches,

585 - attn_metadata=model_input.attn_metadata,

575 intermediate_tensors=intermediate_tensors,

576 **MultiModalKwargs.as_kwargs(model_input.multi_modal_kwargs

577 or {},

Remove unused kwargs from model definitions (#13555)