[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:
@@ -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"]
|
||||
|
||||
Reference in New Issue
Block a user