[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
+2
View File
@@ -127,6 +127,8 @@ int64_t cutlass_mla_get_workspace_size(
void rmsnorm(at::Tensor& output, at::Tensor& input, at::Tensor& weight, double eps, bool enable_pdl);
void sgl_fused_add_rmsnorm(
torch::Tensor input, torch::Tensor residual, torch::Tensor weight, double eps, bool enable_pdl);
void musa_fused_add_rms_norm(
torch::Tensor& input, torch::Tensor& residual, torch::Tensor& weight, double epsilon, bool enable_pdl);
void gemma_rmsnorm(at::Tensor& output, at::Tensor& input, at::Tensor& weight, double eps, bool enable_pdl);
void gemma_fused_add_rmsnorm(at::Tensor& input, at::Tensor& residual, at::Tensor& weight, double eps, bool enable_pdl);
void silu_and_mul(at::Tensor& out, at::Tensor& input);