biondizzle/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

Add contributing guideline and mypy config (#122)

Browse Source
This commit is contained in:
Woosuk Kwon
2023-05-23 17:58:51 -07:00
committed by GitHub
parent 3f942acfe1
commit a283ec2eec
16 changed files with 128 additions and 44 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
cacheflow/model_executor/layers/attention.py
View File

@@ -61,7 +61,7 @@ class GPTCacheFlowAttention(nn.Module):
key: torch.Tensor, # [num_prompt_tokens, num_heads, head_size]
value: torch.Tensor, # [num_prompt_tokens, num_heads, head_size]
attn_bias: xops.AttentionBias,
) -> None:
) -> torch.Tensor:
# TODO(woosuk): The unsqueeze op may incur some CPU overhead. Optimize.
out = xops.memory_efficient_attention_forward(
query.unsqueeze(0),
@@ -197,7 +197,7 @@ class GPTNeoXCacheFlowAttention(GPTCacheFlowAttention):
def forward(
self,
positions: torch.LongTensor, # [num_tokens]
positions: torch.Tensor, # [num_tokens]
query: torch.Tensor, # [num_tokens, num_heads * head_size]
key: torch.Tensor, # [num_tokens, num_heads * head_size]
value: torch.Tensor, # [num_tokens, num_heads * head_size]

2
cacheflow/model_executor/layers/sampler.py
View File

@@ -347,7 +347,7 @@ def _sample_from_generation_tokens(
# Greedy sampling.
assert len(seq_ids) == 1
next_token_id = torch.argmax(probs, dim=-1)
next_token_ids = [next_token_id.item()]
next_token_ids = [int(next_token_id.item())]
parent_seq_ids = seq_ids
else:
# Random sampling.

4
cacheflow/model_executor/model_loader.py
View File

@@ -1,4 +1,6 @@
"""Utilities for selecting and loading models."""
from typing import Type
import torch
import torch.nn as nn
from transformers import PretrainedConfig
@@ -17,7 +19,7 @@ _MODEL_REGISTRY = {
}
def _get_model_architecture(config: PretrainedConfig) -> nn.Module:
def _get_model_architecture(config: PretrainedConfig) -> Type[nn.Module]:
architectures = getattr(config, "architectures", [])
for arch in architectures:
if arch in _MODEL_REGISTRY:

8
cacheflow/model_executor/models/gpt2.py
View File

@@ -168,8 +168,8 @@ class GPT2Model(nn.Module):
def forward(
self,
input_ids: torch.LongTensor,
position_ids: torch.LongTensor,
input_ids: torch.Tensor,
position_ids: torch.Tensor,
kv_caches: List[KVCache],
input_metadata: InputMetadata,
cache_events: Optional[List[torch.cuda.Event]],
@@ -204,8 +204,8 @@ class GPT2LMHeadModel(nn.Module):
def forward(
self,
input_ids: torch.LongTensor,
positions: torch.LongTensor,
input_ids: torch.Tensor,
positions: torch.Tensor,
kv_caches: List[KVCache],
input_metadata: InputMetadata,
cache_events: Optional[List[torch.cuda.Event]],

12
cacheflow/model_executor/models/gpt_neox.py
View File

@@ -67,7 +67,7 @@ class GPTNeoXAttention(nn.Module):
def forward(
self,
position_ids: torch.LongTensor,
position_ids: torch.Tensor,
hidden_states: torch.Tensor,
kv_cache: KVCache,
input_metadata: InputMetadata,
@@ -118,7 +118,7 @@ class GPTNeoXLayer(nn.Module):
def forward(
self,
position_ids: torch.LongTensor,
position_ids: torch.Tensor,
hidden_states: torch.Tensor,
kv_cache: KVCache,
input_metadata: InputMetadata,
@@ -162,8 +162,8 @@ class GPTNeoXModel(nn.Module):
def forward(
self,
input_ids: torch.LongTensor,
position_ids: torch.LongTensor,
input_ids: torch.Tensor,
position_ids: torch.Tensor,
kv_caches: List[KVCache],
input_metadata: InputMetadata,
cache_events: Optional[List[torch.cuda.Event]],
@@ -199,8 +199,8 @@ class GPTNeoXForCausalLM(nn.Module):
def forward(
self,
input_ids: torch.LongTensor,
positions: torch.LongTensor,
input_ids: torch.Tensor,
positions: torch.Tensor,
kv_caches: List[KVCache],
input_metadata: InputMetadata,
cache_events: Optional[List[torch.cuda.Event]],

12
cacheflow/model_executor/models/llama.py
View File

@@ -109,7 +109,7 @@ class LlamaAttention(nn.Module):
def forward(
self,
positions: torch.LongTensor,
positions: torch.Tensor,
hidden_states: torch.Tensor,
kv_cache: KVCache,
input_metadata: InputMetadata,
@@ -143,7 +143,7 @@ class LlamaDecoderLayer(nn.Module):
def forward(
self,
positions: torch.LongTensor,
positions: torch.Tensor,
hidden_states: torch.Tensor,
kv_cache: KVCache,
input_metadata: InputMetadata,
@@ -184,8 +184,8 @@ class LlamaModel(nn.Module):
def forward(
self,
input_ids: torch.LongTensor,
positions: torch.LongTensor,
input_ids: torch.Tensor,
positions: torch.Tensor,
kv_caches: List[KVCache],
input_metadata: InputMetadata,
cache_events: Optional[List[torch.cuda.Event]],
@@ -222,8 +222,8 @@ class LlamaForCausalLM(nn.Module):
def forward(
self,
input_ids: torch.LongTensor,
positions: torch.LongTensor,
input_ids: torch.Tensor,
positions: torch.Tensor,
kv_caches: List[KVCache],
input_metadata: InputMetadata,
cache_events: Optional[List[torch.cuda.Event]],

14
cacheflow/model_executor/models/opt.py
View File

@@ -47,7 +47,7 @@ class OPTLearnedPositionalEmbedding(nn.Embedding):
self.offset = 2
super().__init__(num_embeddings + self.offset, embedding_dim)
def forward(self, positions: torch.LongTensor):
def forward(self, positions: torch.Tensor):
return super().forward(positions + self.offset)
@@ -199,8 +199,8 @@ class OPTDecoder(nn.Module):
def forward(
self,
input_ids: torch.LongTensor,
positions: torch.LongTensor,
input_ids: torch.Tensor,
positions: torch.Tensor,
kv_caches: List[KVCache],
input_metadata: InputMetadata,
cache_events: Optional[List[torch.cuda.Event]],
@@ -235,8 +235,8 @@ class OPTModel(nn.Module):
def forward(
self,
input_ids: torch.LongTensor,
positions: torch.LongTensor,
input_ids: torch.Tensor,
positions: torch.Tensor,
kv_caches: List[KVCache],
input_metadata: InputMetadata,
cache_events: Optional[List[torch.cuda.Event]],
@@ -258,8 +258,8 @@ class OPTForCausalLM(nn.Module):
def forward(
self,
input_ids: torch.LongTensor,
positions: torch.LongTensor,
input_ids: torch.Tensor,
positions: torch.Tensor,
kv_caches: List[KVCache],
input_metadata: InputMetadata,
cache_events: Optional[List[torch.cuda.Event]],