[MUSA][18/N] Add MUSA-optimized kernel implementations for hot ops (#23255)
Signed-off-by: Joey-gvwal <joey_gvwal@yeah.net> Co-authored-by: R0CKSTAR <yeahdongcn@gmail.com>
This commit is contained in:
@@ -111,6 +111,15 @@ from sgl_kernel.version import __version__
|
||||
if torch.version.hip is not None:
|
||||
from sgl_kernel.elementwise import gelu_quick
|
||||
|
||||
if hasattr(torch.version, "musa") and torch.version.musa is not None:
|
||||
from sgl_kernel.musa import (
|
||||
musa_batched_rotary_embedding_contiguous,
|
||||
musa_fused_gemv,
|
||||
musa_fused_moe_gemv,
|
||||
musa_fused_mul_add,
|
||||
musa_rotary_embedding_contiguous,
|
||||
)
|
||||
|
||||
|
||||
_DEBUG_EXPORT_NAMES = [
|
||||
"apply_shuffle_mul_sum",
|
||||
|
||||
Reference in New Issue
Block a user