[Triton] Centralize Ascend extension op dispatch in triton_utils (#6937)
### What this PR does / why we need it? This pull request refactors the dispatch mechanism for the **triton-ascend-specific operators** `insert_slice`, `extract_slice`, and `get_element` to ensure compatibility with both CANN 8.5 and 9.0. A unified helper function, `_resolve_triton_ascend_op`, has been introduced in `vllm_ascend/ops/triton/triton_utils.py`. This function dynamically resolves these operators by first attempting to import them from the `triton.language.extra.cann.extension` module, which is present in newer CANN versions. If that fails, it falls back to the standard `triton.language` module. This approach centralizes operator dispatch logic, allowing individual Triton kernels to use these functions without being aware of the underlying Triton/CANN version. All call sites have been updated to use these new unified functions. ### Does this PR introduce _any_ user-facing change? No. This is an internal refactoring of operator implementations and does not introduce any user-facing changes. ### How was this patch tested? CI is expected to pass with existing tests. **Testing Context:** - vLLM version: v0.16.0 - vLLM main: `15d76f74e2fdb12a95ea00f0ca283acf6219a2b7` Signed-off-by: linfeng-yuan <1102311262@qq.com>
This commit is contained in:
@@ -1,10 +1,46 @@
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from vllm.triton_utils import HAS_TRITON, triton
|
||||
from vllm.triton_utils import HAS_TRITON, tl, triton
|
||||
|
||||
_NUM_AICORE = -1
|
||||
_NUM_VECTORCORE = -1
|
||||
_extension_module = None
|
||||
|
||||
if HAS_TRITON:
|
||||
try:
|
||||
import triton.language.extra.cann.extension as _extension_module # type: ignore
|
||||
except ImportError:
|
||||
_extension_module = None
|
||||
|
||||
|
||||
def _resolve_triton_ascend_op(op_name: str):
|
||||
if not HAS_TRITON:
|
||||
raise RuntimeError(f"Triton op '{op_name}' cannot be resolved because HAS_TRITON is False")
|
||||
|
||||
if _extension_module is not None:
|
||||
extension_op = getattr(_extension_module, op_name, None)
|
||||
if extension_op is not None:
|
||||
return extension_op
|
||||
|
||||
tl_op = getattr(tl, op_name, None)
|
||||
if tl_op is not None:
|
||||
return tl_op
|
||||
|
||||
raise RuntimeError(
|
||||
f"Failed to resolve Triton op '{op_name}': "
|
||||
"neither triton.language.extra.cann.extension nor triton.language provides it."
|
||||
)
|
||||
|
||||
|
||||
if HAS_TRITON:
|
||||
insert_slice = _resolve_triton_ascend_op("insert_slice")
|
||||
extract_slice = _resolve_triton_ascend_op("extract_slice")
|
||||
get_element = _resolve_triton_ascend_op("get_element")
|
||||
else:
|
||||
insert_slice = None
|
||||
extract_slice = None
|
||||
get_element = None
|
||||
|
||||
|
||||
def init_device_properties_triton():
|
||||
|
||||
Reference in New Issue
Block a user