[MUSA][8/N] Port CUDA kernels that are compatible with MUSA (#17946)

Signed-off-by: yafeng.li <yafeng.li@mthreads.com>
Co-authored-by: Alex Nails <alex.nails@radixark.ai>
This commit is contained in:
MARATRIX
2026-04-23 18:04:58 -07:00
committed by GitHub
co-authored by Alex Nails
parent c0166355ae
commit 74c2e5bacd
15 changed files with 1905 additions and 66 deletions
+28 -1
View File
@@ -76,10 +76,37 @@ include_dirs = [
]
sources = [
"csrc/allreduce/custom_all_reduce.cu",
"csrc/attention/merge_attn_states.cu",
"csrc/common_extension_musa.cc",
"csrc/elementwise/activation.cu",
"csrc/elementwise/concat_mla.cu",
"csrc/elementwise/fused_add_rms_norm_kernel.mu",
"csrc/grammar/apply_token_bitmask_inplace_cuda.cu",
"csrc/moe/moe_align_kernel.cu",
"csrc/moe/moe_fused_gate_musa.cu",
"csrc/moe/kimi_k2_moe_fused_gate.cu",
"csrc/moe/moe_sum.cu",
"csrc/moe/moe_sum_reduce.cu",
"csrc/moe/moe_topk_softmax_kernels.cu",
"csrc/quantization/gguf/gguf_kernel.cu",
"csrc/speculative/eagle_utils.cu",
"csrc/speculative/ngram_utils.cu",
"csrc/speculative/packbit.cu",
"csrc/speculative/speculative_sampling.cu",
"csrc/kvcacheio/transfer.cu",
"csrc/gemm/awq_kernel.cu",
"csrc/gemm/bmm_fp8.cu",
"csrc/gemm/dsv3_fused_a_gemm.cu",
"csrc/gemm/dsv3_router_gemm_bf16_out.cu",
"csrc/gemm/dsv3_router_gemm_entry.cu",
"csrc/gemm/dsv3_router_gemm_float_out.cu",
"csrc/gemm/per_token_quant_fp8.cu",
"csrc/gemm/per_token_group_quant_8bit.cu",
"csrc/gemm/per_token_group_quant_8bit_v2.cu",
"csrc/memory/weak_ref_tensor.cpp",
str(_FLASHINFER_REPO.source_dir / "csrc/norm.cu"),
str(_FLASHINFER_REPO.source_dir / "csrc/renorm.cu"),
str(_FLASHINFER_REPO.source_dir / "csrc/sampling.cu"),
]
cxx_flags = ["force_mcc"]