[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:
Joey
2026-05-07 20:38:33 -07:00
committed by GitHub
co-authored by R0CKSTAR
parent 461bc8af49
commit 15e6572f21
15 changed files with 2513 additions and 8 deletions
+9
View File
@@ -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",