EngineX/xc-llm-kunlun
3
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

[Feature] Support glmx (#194)

Browse Source 
Signed-off-by: Li Wei <liwei.109@outlook.com>
Co-authored-by: tangshiwen <tangshiwen@baidu.com>
Co-authored-by: Xinyu Dong <dongxinyu03@baidu.com>
This commit is contained in:
Li Wei
2026-02-12 15:40:42 +08:00
committed by GitHub
parent 070bfa4a73
commit 744719587e
9 changed files with 159 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
vllm_kunlun/ops/_kunlun_ops.py
View File

@@ -195,10 +195,6 @@ class KunlunOps:
query_x = query.contiguous()
key_x = key.contiguous()
num_tokens = query_x.shape[0]
num_heads = query_x.shape[1] // head_size
num_kv_heads = key_x.shape[1] // head_size
torch.ops._C.rotary_embedding(
positions,
query_x,
@@ -207,9 +203,6 @@ class KunlunOps:
cos_sin_cache,
is_neox_style)
query_x = query_x.view(num_tokens, num_heads * head_size)
key_x = key_x.view(num_tokens, num_kv_heads * head_size)
return query_x, key_x
# Rotary embedding