96b2cdf6d8
### 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>
31 lines
1.0 KiB
Python
31 lines
1.0 KiB
Python
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
|