[Ops][Triton] Add a triton kernel supporting partial rope. (#4413)
### What this PR does / why we need it? This PR adds a triton rope kernel witch supports scenarios of `rope_dim != head_dim`. This can save the split op before rope and the concat op after rope. Profiling shows improvement. ### Does this PR introduce _any_ user-facing change? None ### How was this patch tested? I will add related ut after ci integrated with triton. - vLLM version: v0.11.2 - vLLM main: https://github.com/vllm-project/vllm/commit/v0.11.2 --------- Signed-off-by: whx-sjtu <2952154980@qq.com>
This commit is contained in:
30
vllm_ascend/ops/triton/triton_utils.py
Normal file
30
vllm_ascend/ops/triton/triton_utils.py
Normal file
@@ -0,0 +1,30 @@
|
||||
from typing import Any, Dict
|
||||
|
||||
import torch
|
||||
from vllm.triton_utils import HAS_TRITON, triton
|
||||
|
||||
_NUM_AICORE = -1
|
||||
_NUM_VECTORCORE = -1
|
||||
|
||||
|
||||
def init_device_properties_triton():
|
||||
global _NUM_AICORE, _NUM_VECTORCORE
|
||||
if _NUM_AICORE == -1 and HAS_TRITON:
|
||||
device_properties: Dict[str, Any] = (
|
||||
triton.runtime.driver.active.utils.get_device_properties(
|
||||
torch.npu.current_device()))
|
||||
_NUM_AICORE = device_properties.get("num_aicore", -1)
|
||||
_NUM_VECTORCORE = device_properties.get("num_vectorcore", -1)
|
||||
assert _NUM_AICORE > 0 and _NUM_VECTORCORE > 0, "Failed to detect device properties."
|
||||
|
||||
|
||||
def get_aicore_num():
|
||||
global _NUM_AICORE
|
||||
assert _NUM_AICORE > 0, "Device properties not initialized. Please call init_device_properties_triton() first."
|
||||
return _NUM_AICORE
|
||||
|
||||
|
||||
def get_vectorcore_num():
|
||||
global _NUM_VECTORCORE
|
||||
assert _NUM_VECTORCORE > 0, "Device properties not initialized. Please call init_device_properties_triton() first."
|
||||
return _NUM_VECTORCORE
|
||||
Reference in New Issue
Block a user