fix llama PA

predibase · May 22, 2024 · a20bc03 · a20bc03
1 parent 6d41fbf
commit a20bc03
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/server/lorax_server/models/custom_modeling/flash_llama_modeling.py b/server/lorax_server/models/custom_modeling/flash_llama_modeling.py
@@ -331,13 +331,14 @@ def forward(
  query,
  kv_cache[0],
  kv_cache[1],
- self.num_key_value_heads,
+ self.kv_head_mapping,
  self.softmax_scale,
  block_tables,
  input_lengths,
  max_s,
  )
 
+
  return self.o_proj(attn_output.view(-1, self.num_heads * self.head_size), adapter_data)