github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-09 04:50:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

[Hotfix] Fix bugs in testing continuous batching (#5270)

Browse Source 
* fix bug

* fix bugs

* fix bugs

* fix bugs and add padding

* add funcs and fix bugs

* fix typos

* fix bugs

* add func
This commit is contained in:
Jianghai
2024-01-18 16:31:14 +08:00
committed by GitHub
parent 5ae9099f92
commit 9e2342bde2
6 changed files with 86 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
colossalai/inference/modeling/layers/attention.py
View File

@@ -69,7 +69,7 @@ def convert_kvcache(cache, lengths, block_tables, pad_id=0):
)
padding = seq_len - _cache.size(0)
if padding > 0:
_cache = F.pad(_cache, (0, 0, 0, 0, 0, 1), value=pad_id)
_cache = F.pad(_cache, (0, 0, 0, 0, 0, padding), value=pad_id)
padded_cache.append(_cache)
return torch.stack(padded_cache, dim=0)

3
colossalai/inference/modeling/models/llama.py
View File

@@ -173,7 +173,10 @@ def llama_attn_forward(
key_states = self.k_proj(hidden_states).view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
value_states = self.v_proj(hidden_states).view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
kv_seq_len = max(sequence_lengths).item()
cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
query_states = query_states.transpose(1, 2)