Migrate XTorch operations to Kunlun operations (accelerating iteration) (#177)

Signed-off-by: dongxinyu03 <dongxinyu03@baidu.com>
This commit is contained in:
Xinyu Dong
2026-02-12 18:13:00 +08:00
committed by GitHub
parent 744719587e
commit bf9369f733
15 changed files with 125 additions and 119 deletions

View File

@@ -13,7 +13,7 @@ from typing import Optional
import torch
from vllm.triton_utils import tl, triton
import xtorch_ops
import kunlun_ops
BT_LIST = [8, 16, 32, 64, 128]
@@ -149,5 +149,5 @@ def l2norm_fwd(x: torch.Tensor,
eps: float = 1e-6,
output_dtype: Optional[torch.dtype] = None):
out = torch.empty_like(x)
xtorch_ops.l2norm(x, out, eps)
kunlun_ops.l2norm(x, out, eps)
return out