[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
+6
View File
@@ -81,6 +81,7 @@ sources = [
"csrc/common_extension_musa.cc",
"csrc/elementwise/activation.cu",
"csrc/elementwise/concat_mla.cu",
"csrc/elementwise/pos_enc.cu",
"csrc/elementwise/fused_add_rms_norm_kernel.mu",
"csrc/grammar/apply_token_bitmask_inplace_cuda.cu",
"csrc/moe/moe_align_kernel.cu",
@@ -107,6 +108,11 @@ sources = [
"csrc/memory/weak_ref_tensor.cpp",
str(_FLASHINFER_REPO.source_dir / "csrc/norm.cu"),
str(_FLASHINFER_REPO.source_dir / "csrc/renorm.cu"),
# XXX (MUSA): The following files contain MUSA-specific implementations.
"csrc/musa/pos_encoding_contiguous.mu",
"csrc/musa/moe_gemv_swiglu.mu",
"csrc/musa/ternary.mu",
"csrc/musa/top_k_top_p_sampling.mu",
]
cxx_flags = ["force_mcc"]