A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 19:45 UTC
[Bugfix][Kernel][CPU] Fix num_tokens in CPU rotary embedding kernel (#14667)
Signed-off-by: Thien Tran <gau.nernst@yahoo.com.sg>
Thien Tran committed
Mar 14, 2025 at 14:47 UTC
27b50f1fe6e325f73b405263c3ac1fc668531118
1 file changed
+1
-1
csrc/cpu/pos_encoding.cpp
+1
-1
@@ -170,7 +170,7 @@ void rotary_embedding_gptj_impl(
170
void rotary_embedding(torch::Tensor& positions, torch::Tensor& query,
171
torch::Tensor& key, int64_t head_size,
172
torch::Tensor& cos_sin_cache, bool is_neox) {
173
- int num_tokens = query.numel() / query.size(-1);
173
+ int num_tokens = positions.numel();
174
int rot_dim = cos_sin_cache.size(1);
175
int num_heads = query.size(-1) / head_size;
176
int num_kv_heads = key.size(-1) / head_size;