diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index f5255dc85..4e1088894 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -396,6 +396,7 @@ class Envs: SGLANG_DG_CACHE_DIR = EnvStr(os.path.expanduser("~/.cache/deep_gemm")) SGLANG_DG_USE_NVRTC = EnvBool(False) SGLANG_USE_DEEPGEMM_BMM = EnvBool(False) + SGLANG_DEEPGEMM_SANITY_CHECK = EnvBool(False) # DeepSeek MHA Optimization SGLANG_CHUNKED_PREFIX_CACHE_THRESHOLD = EnvInt(8192) diff --git a/python/sglang/srt/layers/activation.py b/python/sglang/srt/layers/activation.py index a642fea95..ea3abef7d 100644 --- a/python/sglang/srt/layers/activation.py +++ b/python/sglang/srt/layers/activation.py @@ -36,6 +36,7 @@ from sglang.srt.utils import ( is_cpu, is_cuda, is_hip, + is_musa, is_npu, is_xpu, set_weight_attrs, @@ -43,6 +44,7 @@ from sglang.srt.utils import ( from sglang.utils import resolve_obj_by_qualname _is_cuda = is_cuda() +_is_musa = is_musa() _is_npu = is_npu() _is_cpu_amx_available = cpu_has_amx_support() _is_cpu = is_cpu() @@ -53,6 +55,8 @@ if _is_cuda or _is_xpu: from sgl_kernel import gelu_and_mul, gelu_tanh_and_mul, silu_and_mul elif _is_hip: from sgl_kernel import gelu_and_mul, gelu_quick, gelu_tanh_and_mul, silu_and_mul +elif _is_musa: + from sgl_kernel import silu_and_mul if is_npu(): import torch_npu @@ -95,6 +99,15 @@ class SiluAndMul(MultiPlatformOp): silu_and_mul(x, out) return out + def forward_musa(self, x: torch.Tensor) -> torch.Tensor: + if not get_global_server_args().disable_piecewise_cuda_graph: + return self.forward_native(x) + + if not hasattr(self, "_musa_swish_glu"): + # XXX (MUSA): nn.SwishGLU seems to have better performance than silu_and_mul on MUSA, we can switch to it for now. We can consider implementing a silu_and_mul kernel for MUSA in the future if needed. + self._musa_swish_glu = nn.SwishGLU() + return self._musa_swish_glu(x) + class GeluAndMul(MultiPlatformOp): def __init__(self, approximate="tanh"): diff --git a/python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py b/python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py index 6aef866f5..a30bf5e14 100644 --- a/python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py +++ b/python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py @@ -1,6 +1,6 @@ import logging import os -from contextlib import contextmanager +from contextlib import contextmanager, nullcontext from enum import IntEnum, auto from typing import Dict, List, Tuple @@ -14,10 +14,12 @@ from sglang.srt.distributed.device_communicators.pynccl_allocator import ( from sglang.srt.environ import envs from sglang.srt.layers.deep_gemm_wrapper.configurer import ENABLE_JIT_DEEPGEMM from sglang.srt.server_args import ServerArgs -from sglang.srt.utils import ceil_div, get_available_gpu_memory +from sglang.srt.utils import ceil_div, get_available_gpu_memory, is_musa logger = logging.getLogger(__name__) +_is_musa = is_musa() + if ENABLE_JIT_DEEPGEMM: import deep_gemm @@ -332,9 +334,18 @@ class _BF16F32WarmupExecutor(_BaseWarmupExecutor): deep_gemm.bf16_gemm_nt(self.lhs[:m], self.rhs, self.out[:m]) -@contextmanager def deep_gemm_execution_hook( m: int, n: int, k: int, num_groups: int, kernel_type: DeepGemmKernelType +): + if _is_musa: + return nullcontext() + + return _deep_gemm_execution_hook(m, n, k, num_groups, kernel_type) + + +@contextmanager +def _deep_gemm_execution_hook( + m: int, n: int, k: int, num_groups: int, kernel_type: DeepGemmKernelType ): if m > 0: _maybe_compile_deep_gemm_one_type_all(kernel_type, n, k, num_groups) diff --git a/python/sglang/srt/layers/deep_gemm_wrapper/configurer.py b/python/sglang/srt/layers/deep_gemm_wrapper/configurer.py index 34494f599..6ae5ee431 100644 --- a/python/sglang/srt/layers/deep_gemm_wrapper/configurer.py +++ b/python/sglang/srt/layers/deep_gemm_wrapper/configurer.py @@ -1,14 +1,22 @@ import logging from sglang.srt.environ import envs -from sglang.srt.utils import get_device_sm, is_blackwell_supported +from sglang.srt.utils import ( + get_device_sm, + is_blackwell_supported, + is_cuda, + is_musa, +) logger = logging.getLogger(__name__) +_is_cuda = is_cuda() +_is_musa = is_musa() + def _compute_enable_deep_gemm(): sm_version = get_device_sm() - if sm_version < 90: + if (_is_cuda and sm_version < 90) or (_is_musa and sm_version < 31): return False try: @@ -23,3 +31,4 @@ ENABLE_JIT_DEEPGEMM = _compute_enable_deep_gemm() DEEPGEMM_BLACKWELL = ENABLE_JIT_DEEPGEMM and is_blackwell_supported() DEEPGEMM_SCALE_UE8M0 = DEEPGEMM_BLACKWELL +DEEPGEMM_NEED_TMA_ALIGNED_SCALES = not (DEEPGEMM_SCALE_UE8M0 or _is_musa) diff --git a/python/sglang/srt/layers/deep_gemm_wrapper/entrypoint.py b/python/sglang/srt/layers/deep_gemm_wrapper/entrypoint.py index 1087d7784..db4e23a94 100644 --- a/python/sglang/srt/layers/deep_gemm_wrapper/entrypoint.py +++ b/python/sglang/srt/layers/deep_gemm_wrapper/entrypoint.py @@ -4,14 +4,15 @@ from typing import Any, Optional, Tuple import torch +from sglang.srt.environ import envs from sglang.srt.layers.deep_gemm_wrapper import compile_utils from sglang.srt.layers.deep_gemm_wrapper.configurer import ( # noqa: F401 DEEPGEMM_BLACKWELL, + DEEPGEMM_NEED_TMA_ALIGNED_SCALES, DEEPGEMM_SCALE_UE8M0, ENABLE_JIT_DEEPGEMM, ) from sglang.srt.server_args import ServerArgs -from sglang.srt.utils import get_bool_env_var logger = logging.getLogger(__name__) @@ -19,7 +20,7 @@ if ENABLE_JIT_DEEPGEMM: import deep_gemm from deep_gemm.utils.layout import get_mn_major_tma_aligned_tensor # noqa: F401 -_SANITY_CHECK = get_bool_env_var("SGLANG_DEEPGEMM_SANITY_CHECK") +_SANITY_CHECK = envs.SGLANG_DEEPGEMM_SANITY_CHECK.get() # TODO maybe rename these functions diff --git a/python/sglang/srt/layers/layernorm.py b/python/sglang/srt/layers/layernorm.py index 30dc41451..8c4f92969 100644 --- a/python/sglang/srt/layers/layernorm.py +++ b/python/sglang/srt/layers/layernorm.py @@ -34,6 +34,7 @@ from sglang.srt.utils import ( is_cuda, is_flashinfer_available, is_hip, + is_musa, is_npu, is_xpu, ) @@ -41,6 +42,7 @@ from sglang.srt.utils import ( _is_cuda = is_cuda() _is_flashinfer_available = is_flashinfer_available() _is_hip = is_hip() +_is_musa = is_musa() _is_npu = is_npu() _use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip _is_cpu_amx_available = cpu_has_amx_support() @@ -48,7 +50,7 @@ _is_cpu = is_cpu() _is_xpu = is_xpu() _flashinfer_layernorm_available = False -if _is_cuda or _is_xpu: +if _is_cuda or _is_xpu or _is_musa: if _is_flashinfer_available: try: from flashinfer.norm import layernorm @@ -323,6 +325,29 @@ class RMSNorm(MultiPlatformOp): rms_norm(out, x, self.weight.data, self.variance_epsilon) return out + def forward_musa( + self, + x: torch.Tensor, + residual: Optional[torch.Tensor] = None, + post_residual_addition: Optional[torch.Tensor] = None, + ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: + if not get_global_server_args().disable_piecewise_cuda_graph: + return self.forward_native(x, residual, post_residual_addition) + + if not x.is_contiguous(): + x = x.contiguous() + + if residual is not None: + if post_residual_addition is not None: + residual = residual + post_residual_addition + fused_add_rmsnorm(x, residual, self.weight.data, self.variance_epsilon) + return x, residual + + out = nn.functional.rms_norm( + x, (self.hidden_size,), self.weight.data, self.variance_epsilon + ) + return out + def forward_native( self, x: torch.Tensor, diff --git a/python/sglang/srt/layers/moe/ep_moe/kernels.py b/python/sglang/srt/layers/moe/ep_moe/kernels.py index 4cd4b4f81..40de48e72 100644 --- a/python/sglang/srt/layers/moe/ep_moe/kernels.py +++ b/python/sglang/srt/layers/moe/ep_moe/kernels.py @@ -3,12 +3,14 @@ import logging import torch import triton -from sglang.srt.utils import ceil_div, is_cuda +from sglang.srt.utils import ceil_div, is_cuda, is_musa logger = logging.getLogger(__name__) _is_cuda = is_cuda() -if _is_cuda: +_is_musa = is_musa() + +if _is_cuda or _is_musa: from sglang.srt.layers.quantization.fp8_kernel import ( sglang_per_token_group_quant_fp8 as per_token_group_quant_fp8, ) @@ -665,6 +667,8 @@ def _fwd_kernel_ep_scatter_2( HIDDEN_SIZE_PAD: tl.constexpr, SCALE_HIDDEN_SIZE: tl.constexpr, SCALE_HIDDEN_SIZE_PAD: tl.constexpr, + # Platform-specific semaphore for atomic_add performance tuning + ATOMIC_ADD_SEM: tl.constexpr, ): start_token_id = tl.program_id(0) grid_num = tl.num_programs(0) @@ -689,7 +693,9 @@ def _fwd_kernel_ep_scatter_2( topk_index = topk_idx_int32.to(tl.int64) expert_id = tl.load(recv_topk + token_id * recv_topk_stride0 + topk_index) if expert_id >= 0: - dest_token_index_int32 = tl.atomic_add(expert_start_loc + expert_id, 1) + dest_token_index_int32 = tl.atomic_add( + expert_start_loc + expert_id, 1, sem=ATOMIC_ADD_SEM + ) dest_token_index = dest_token_index_int32.to(tl.int64) tl.store( @@ -783,6 +789,8 @@ def ep_scatter( HIDDEN_SIZE_PAD=triton.next_power_of_2(hidden_size), SCALE_HIDDEN_SIZE=scale_hidden_size, SCALE_HIDDEN_SIZE_PAD=triton.next_power_of_2(scale_hidden_size), + # XXX (MUSA): Atomic add with "relaxed" semaphore on musa backend for better performance + ATOMIC_ADD_SEM=None if not _is_musa else "relaxed", ) return diff --git a/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py b/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py index 6578fece6..f4b4585cb 100644 --- a/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py +++ b/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py @@ -22,6 +22,7 @@ from sglang.srt.utils import ( get_bool_env_var, is_cuda, is_hip, + is_musa, is_npu, ) from sglang.srt.utils.offloader import get_offloader @@ -42,6 +43,7 @@ _is_hip = is_hip() _is_npu = is_npu() _is_cuda = is_cuda() _use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip +_is_musa = is_musa() if not (_is_npu or _is_hip) and _is_cuda: from sgl_kernel import silu_and_mul @@ -166,8 +168,9 @@ class DeepGemmRunnerCore(MoeRunnerCore): device=hidden_states_device, dtype=torch.bfloat16, ) - if not deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0: + if deep_gemm_wrapper.DEEPGEMM_NEED_TMA_ALIGNED_SCALES: hidden_states_scale = tma_align_input_scale(hidden_states_scale) + deep_gemm_wrapper.grouped_gemm_nt_f8f8bf16_contig( (hidden_states, hidden_states_scale), w13_weight_fp8, @@ -203,7 +206,7 @@ class DeepGemmRunnerCore(MoeRunnerCore): device=hidden_states_device, dtype=torch.bfloat16, ) - if not deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0: + if deep_gemm_wrapper.DEEPGEMM_NEED_TMA_ALIGNED_SCALES: down_input_scale = tma_align_input_scale(down_input_scale) deep_gemm_wrapper.grouped_gemm_nt_f8f8bf16_contig( @@ -251,7 +254,7 @@ class DeepGemmRunnerCore(MoeRunnerCore): hidden_states_scale = _cast_to_e8m0_with_rounding_up( hidden_states_scale ) - else: + elif deep_gemm_wrapper.DEEPGEMM_NEED_TMA_ALIGNED_SCALES: hidden_states_scale = deep_gemm_wrapper.get_mn_major_tma_aligned_tensor( hidden_states_scale ) @@ -317,7 +320,7 @@ class DeepGemmRunnerCore(MoeRunnerCore): # GroupGemm-1 n = w2_weight.shape[1] - if not deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0: + if deep_gemm_wrapper.DEEPGEMM_NEED_TMA_ALIGNED_SCALES: down_input_scale = deep_gemm_wrapper.get_mn_major_tma_aligned_tensor( down_input_scale ) diff --git a/python/sglang/srt/layers/moe/moe_runner/triton_utils/fused_moe.py b/python/sglang/srt/layers/moe/moe_runner/triton_utils/fused_moe.py index a16f42fac..b93add78f 100644 --- a/python/sglang/srt/layers/moe/moe_runner/triton_utils/fused_moe.py +++ b/python/sglang/srt/layers/moe/moe_runner/triton_utils/fused_moe.py @@ -20,6 +20,7 @@ from sglang.srt.utils import ( is_cpu, is_cuda, is_hip, + is_musa, is_xpu, use_intel_xpu_backend, ) @@ -44,6 +45,7 @@ _is_cpu = is_cpu() _use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip _is_xpu = is_xpu() _use_sgl_xpu = use_intel_xpu_backend() +_is_musa = is_musa() if _is_cuda: @@ -62,6 +64,10 @@ elif _is_hip: # because the code uses moe_sum_reduce_triton as fallback (line 619) elif _is_xpu: from sgl_kernel import moe_sum_reduce, silu_and_mul +elif _is_musa: + from sgl_kernel import moe_sum_reduce + + _silu_and_mul_musa = torch.nn.SwishGLU() # Try to import vllm_ops for non-CUDA/HIP/XPU platforms _has_vllm_ops = False @@ -534,6 +540,8 @@ def _fused_moe_kernel_sequence( down_moe_use_tma, activation, ) + elif _is_musa: + intermediate_cache2 = _silu_and_mul_musa(intermediate_cache1.view(-1, N)) else: if _has_vllm_ops: vllm_ops.silu_and_mul( @@ -647,7 +655,7 @@ def _fused_moe_kernel_sequence( if no_combine: pass - elif _is_cuda: + elif _is_cuda or _is_musa: if use_fused_moe_sum_all_reduce: if routed_scaling_factor != 1.0: assert out_slice is not None diff --git a/python/sglang/srt/layers/moe/moe_runner/triton_utils/moe_align_block_size.py b/python/sglang/srt/layers/moe/moe_runner/triton_utils/moe_align_block_size.py index 12d1d5b1d..5be9b5136 100644 --- a/python/sglang/srt/layers/moe/moe_runner/triton_utils/moe_align_block_size.py +++ b/python/sglang/srt/layers/moe/moe_runner/triton_utils/moe_align_block_size.py @@ -5,13 +5,14 @@ from typing import Tuple import torch import triton -from sglang.srt.utils import is_cuda, is_hip, is_xpu +from sglang.srt.utils import is_cuda, is_hip, is_musa, is_xpu _is_cuda = is_cuda() _is_hip = is_hip() _is_xpu = is_xpu() +_is_musa = is_musa() -if _is_cuda or _is_hip or _is_xpu: +if _is_cuda or _is_hip or _is_xpu or _is_musa: from sgl_kernel import moe_align_block_size as sgl_moe_align_block_size diff --git a/python/sglang/srt/layers/moe/topk.py b/python/sglang/srt/layers/moe/topk.py index 77660ed55..f67880552 100644 --- a/python/sglang/srt/layers/moe/topk.py +++ b/python/sglang/srt/layers/moe/topk.py @@ -62,6 +62,7 @@ from sglang.srt.utils import ( is_cpu, is_cuda, is_hip, + is_musa, is_npu, is_xpu, ) @@ -80,8 +81,9 @@ _is_xpu = is_xpu() _is_npu = is_npu() _is_xpu = is_xpu() _use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip +_is_musa = is_musa() -if _is_cuda: +if _is_cuda or _is_musa: from sgl_kernel import moe_fused_gate try: @@ -124,7 +126,7 @@ if _is_cuda: except ImportError as e: pass -if _is_cuda or _is_hip or _is_xpu: +if _is_cuda or _is_hip or _is_xpu or _is_musa: from sgl_kernel import topk_softmax try: @@ -851,7 +853,7 @@ def biased_grouped_topk_gpu( return topk_weights, topk_ids elif ( - _is_cuda + (_is_cuda or _is_musa) # moe_fused_gate kernel ensures that num_experts/num_expert_group does not exceed MAX_VPT=32 now. And when kernel can handle MAX_VPT > 32, we can remove this assertion. and experts_per_group <= 32 and is_power_of_two(num_experts) @@ -1077,7 +1079,6 @@ def select_experts( num_token_non_padded: Optional[torch.Tensor] = None, expert_location_dispatch_info: Optional[ExpertLocationDispatchInfo] = None, ) -> StandardTopKOutput: - top_k = topk_config.top_k use_grouped_topk = topk_config.use_grouped_topk topk_group = topk_config.topk_group @@ -1094,12 +1095,13 @@ def select_experts( scoring_func = topk_config.scoring_func - router_logits, correction_bias = ( - expert_location_dispatch.transform_select_experts_inputs( - router_logits=router_logits, - correction_bias=correction_bias, - info=expert_location_dispatch_info, - ) + ( + router_logits, + correction_bias, + ) = expert_location_dispatch.transform_select_experts_inputs( + router_logits=router_logits, + correction_bias=correction_bias, + info=expert_location_dispatch_info, ) # DeepSeek V2/V3/R1 series models use grouped_top_k diff --git a/python/sglang/srt/layers/quantization/fp8.py b/python/sglang/srt/layers/quantization/fp8.py index cafc7a6a6..b1a1df37a 100644 --- a/python/sglang/srt/layers/quantization/fp8.py +++ b/python/sglang/srt/layers/quantization/fp8.py @@ -85,6 +85,7 @@ from sglang.srt.utils import ( is_cpu, is_cuda, is_hip, + is_musa, is_npu, is_sm90_supported, is_sm100_supported, @@ -103,6 +104,7 @@ if TYPE_CHECKING: _is_hip = is_hip() _is_cuda = is_cuda() +_is_musa = is_musa() _is_npu = is_npu() _is_cpu_amx_available = cpu_has_amx_support() _is_cpu = is_cpu() @@ -185,6 +187,9 @@ class Fp8Config(QuantizationConfig): return [torch.bfloat16, torch.half] def get_min_capability(self) -> int: + if _is_musa: + return 31 + return 100 if self.use_mxfp8 else 80 @classmethod diff --git a/python/sglang/srt/layers/quantization/fp8_kernel.py b/python/sglang/srt/layers/quantization/fp8_kernel.py index bf7b6ba17..bc83b166a 100644 --- a/python/sglang/srt/layers/quantization/fp8_kernel.py +++ b/python/sglang/srt/layers/quantization/fp8_kernel.py @@ -37,6 +37,7 @@ from sglang.srt.utils import ( is_cpu, is_cuda, is_hip, + is_musa, is_sm100_supported, is_sm120_supported, log_info_on_rank0, @@ -47,11 +48,12 @@ from sglang.srt.utils.patch_torch import register_fake_if_exists _is_hip = is_hip() _is_cuda = is_cuda() _is_cpu = is_cpu() +_is_musa = is_musa() _is_sm100_supported = is_sm100_supported() _is_sm120_supported = is_sm120_supported() _use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip -if _is_cuda: +if _is_cuda or _is_musa: from sgl_kernel import sgl_per_token_quant_fp8 from sglang.jit_kernel.per_tensor_quant_fp8 import ( @@ -506,7 +508,7 @@ def sglang_per_token_group_quant_fp8( # Enable v2 kernel by default on supported group sizes _V2_KERNEL_SUPPORTED_GROUP_SIZES = [16, 32, 64, 128] if enable_v2 is None: - enable_v2 = group_size in _V2_KERNEL_SUPPORTED_GROUP_SIZES + enable_v2 = group_size in _V2_KERNEL_SUPPORTED_GROUP_SIZES or _is_musa if x.shape[0] > 0: # Temporary @@ -1111,6 +1113,11 @@ def w8a8_block_fp8_matmul_deepgemm( # Deepgemm only supports output tensor type as bfloat16 assert C.dtype == torch.bfloat16 and deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM + if _is_musa: + # XXX (MUSA): `deep_gemm_fp8_fp8_bf16_nt` on MUSA requires contiguous tensors + As = As.contiguous() + Bs = Bs.contiguous() + deep_gemm_fp8_fp8_bf16_nt(A, As, B, Bs, C) return C diff --git a/python/sglang/srt/layers/quantization/fp8_utils.py b/python/sglang/srt/layers/quantization/fp8_utils.py index 8f3c84bf8..8d075b688 100755 --- a/python/sglang/srt/layers/quantization/fp8_utils.py +++ b/python/sglang/srt/layers/quantization/fp8_utils.py @@ -39,6 +39,7 @@ from sglang.srt.utils import ( is_flashinfer_available, is_gfx95_supported, is_hip, + is_musa, is_sm90_supported, is_sm100_supported, is_sm120_supported, @@ -54,6 +55,7 @@ _is_fp8_fnuz = is_fp8_fnuz() _is_sm100_supported = is_sm100_supported() _is_sm120_supported = is_sm120_supported() _is_gfx95_supported = is_gfx95_supported() +_is_musa = is_musa() _use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip _use_aiter_gfx95 = _use_aiter and _is_gfx95_supported diff --git a/python/sglang/srt/layers/quantization/unquant.py b/python/sglang/srt/layers/quantization/unquant.py index 2dcdd49b7..1ade4ed9e 100644 --- a/python/sglang/srt/layers/quantization/unquant.py +++ b/python/sglang/srt/layers/quantization/unquant.py @@ -695,4 +695,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, MultiPlatformOp): def forward_tpu(self, *args, **kwargs) -> CombineInput: raise NotImplementedError("The TPU backend currently does not support MoE.") + def forward_musa(self, *args, **kwargs) -> CombineInput: + return self.forward_cuda(*args, **kwargs) + forward_native = forward_cpu diff --git a/python/sglang/srt/layers/sampler.py b/python/sglang/srt/layers/sampler.py index e947a48cb..caa468b9f 100644 --- a/python/sglang/srt/layers/sampler.py +++ b/python/sglang/srt/layers/sampler.py @@ -16,7 +16,12 @@ from sglang.srt.layers.utils.logprob import get_token_ids_logprobs, get_top_logp from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo from sglang.srt.sampling.sampling_params import TOP_K_ALL from sglang.srt.server_args import get_global_server_args -from sglang.srt.utils.common import crash_on_warnings, get_bool_env_var, is_cuda, is_npu +from sglang.srt.utils.common import ( + crash_on_warnings, + get_bool_env_var, + is_cuda, + is_npu, +) if is_cuda(): from flashinfer.sampling import ( @@ -27,6 +32,7 @@ if is_cuda(): top_k_renorm_prob, top_p_renorm_prob, ) + if is_npu(): import torch_npu diff --git a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py index 5ff13ebbb..d710c9018 100644 --- a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py +++ b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py @@ -13,6 +13,7 @@ from sglang.srt.model_executor.forward_batch_info import ForwardBatch from sglang.srt.models.deepseek_common.utils import ( _is_cuda, _is_hip, + _is_musa, _is_npu, _use_aiter_gfx95, ) @@ -491,7 +492,7 @@ class DeepseekMHAForwardMixin: # Temporary for DeepSeek V3/R1 only, but can generalize if needed k_shape = (k_nope.shape[0], self.num_local_heads, self.qk_head_dim) if ( - _is_cuda + (_is_cuda or _is_musa) and (self.num_local_heads == 128) and (self.qk_nope_head_dim == 128) and (self.qk_rope_head_dim == 64) diff --git a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py index 646591909..26fef866b 100644 --- a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py +++ b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py @@ -21,6 +21,7 @@ from sglang.srt.models.deepseek_common.utils import ( _is_cuda, _is_gfx95_supported, _is_hip, + _is_musa, _use_aiter, _use_aiter_gfx95, ) @@ -553,6 +554,11 @@ class DeepseekMLAForwardMixin: torch.bfloat16, ) attn_bmm_output = attn_bmm_output.transpose(0, 1).flatten(1, 2) + elif _is_musa: + attn_bmm_output = torch.bmm( + attn_output.to(torch.bfloat16).transpose(0, 1), self.w_vc + ) + attn_bmm_output = attn_bmm_output.transpose(0, 1).flatten(1, 2) else: if is_in_piecewise_cuda_graph(): # torch dynamo requires out= op was called where output tensor was non-contiguous diff --git a/python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py b/python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py index 77760007a..8112b321e 100644 --- a/python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py +++ b/python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py @@ -49,6 +49,7 @@ from sglang.srt.models.deepseek_common.utils import ( _is_cuda, _is_fp8_fnuz, _is_hip, + _is_musa, _is_npu, _is_xpu, _use_aiter_gfx95, @@ -498,7 +499,7 @@ class DeepseekV2WeightLoaderMixin: ) if ( - (_is_cuda or _is_xpu) + (_is_cuda or _is_musa or _is_xpu) and weight_block_size[0] == 128 and weight_block_size[1] == 128 ): @@ -585,6 +586,14 @@ class DeepseekV2WeightLoaderMixin: ) if _is_hip: self_attn.w_scale *= 2.0 + # XXX (MUSA): Remove this after adding FP8 support in bmm kernel on MUSA + if _is_musa and w.dtype == torch.float8_e4m3fn: + self_attn.w_kc = ( + self_attn.w_kc.to(torch.bfloat16) * self_attn.w_scale + ) + self_attn.w_vc = ( + self_attn.w_vc.to(torch.bfloat16) * self_attn.w_scale + ) else: num_tiles_k = self_attn.qk_nope_head_dim // weight_block_size[1] num_tiles_n = self_attn.v_head_dim // weight_block_size[0] diff --git a/python/sglang/srt/models/deepseek_common/utils.py b/python/sglang/srt/models/deepseek_common/utils.py index a5579d528..fab9594ef 100644 --- a/python/sglang/srt/models/deepseek_common/utils.py +++ b/python/sglang/srt/models/deepseek_common/utils.py @@ -29,6 +29,7 @@ from sglang.srt.utils import ( is_cuda, is_gfx95_supported, is_hip, + is_musa, is_npu, is_nvidia_cublas_version_ge_12_9, is_xpu, @@ -37,6 +38,7 @@ from sglang.srt.utils import ( _is_hip = is_hip() _is_cuda = is_cuda() _is_npu = is_npu() +_is_musa = is_musa() _is_fp8_fnuz = is_fp8_fnuz() _use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip _is_cpu_amx_available = cpu_has_amx_support() diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index 4d49cd03b..6b23c2c88 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -141,6 +141,7 @@ from sglang.srt.models.deepseek_common.utils import ( _is_cuda, _is_gfx95_supported, _is_hip, + _is_musa, _is_npu, _is_xpu, _use_aiter, @@ -182,6 +183,8 @@ elif _is_npu: forward_mla_core_npu, forward_mla_prepare_npu, ) +elif _is_musa: + from sgl_kernel import dsv3_fused_a_gemm, dsv3_router_gemm else: pass @@ -640,7 +643,9 @@ class DeepseekV2MoE(nn.Module): expert_location_dispatch_info=dispatch_info, ) final_hidden_states = self.experts(hidden_states, topk_output) - if not _is_cuda or isinstance(self.experts.quant_method, KTEPWrapperMethod): + if not (_is_cuda or _is_musa) or isinstance( + self.experts.quant_method, KTEPWrapperMethod + ): final_hidden_states *= self.routed_scaling_factor current_stream.wait_stream(self.alt_stream) @@ -725,6 +730,7 @@ class DeepseekV2MoE(nn.Module): ) if ( not _is_cuda + and not _is_musa and not _is_xpu and not _use_aiter or isinstance(self.experts.quant_method, KTEPWrapperMethod) @@ -1910,7 +1916,7 @@ class DeepseekV2Model(nn.Module): self.alt_stream = ( torch.cuda.Stream() - if _is_cuda or envs.SGLANG_NPU_USE_MULTI_STREAM.get() + if _is_cuda or _is_musa or envs.SGLANG_NPU_USE_MULTI_STREAM.get() else None ) @@ -2249,12 +2255,15 @@ class DeepseekV2ForCausalLM(nn.Module, DeepseekV2WeightLoaderMixin): or self.config.n_shared_experts != 1 ): disable_reason = "Config does not support fused shared expert(s)." - elif (not _is_cuda or torch.cuda.get_device_capability("cuda") < (8, 0)) and ( - not _is_hip or torch.cuda.get_device_capability("cuda") < (9, 4) + elif ( + (not _is_cuda or torch.cuda.get_device_capability("cuda") < (8, 0)) + and (not _is_hip or torch.cuda.get_device_capability("cuda") < (9, 4)) + and (not _is_musa or torch.musa.get_device_capability("musa") < (3, 1)) ): disable_reason = ( "Only Deepseek V3/R1 on NV-platform with capability >= 80 " "or AMD-platform with capability >= gfx942(MI30x) can use shared experts fusion optimization." + "or MT-platform with capability >= 31 can use shared experts fusion optimization." ) elif get_moe_expert_parallel_world_size() > 1 and ( not _is_hip or torch.cuda.get_device_capability("cuda") < (9, 4) diff --git a/python/sglang/srt/observability/scheduler_metrics_mixin.py b/python/sglang/srt/observability/scheduler_metrics_mixin.py index 0c9da952e..4b75fb877 100644 --- a/python/sglang/srt/observability/scheduler_metrics_mixin.py +++ b/python/sglang/srt/observability/scheduler_metrics_mixin.py @@ -376,6 +376,7 @@ class SchedulerMetricsMixin: { "cpu": "cpu graph", "npu": "npu graph", + "musa": "musa graph", }, ) @@ -561,6 +562,7 @@ class SchedulerMetricsMixin: { "cpu": "cpu graph", "npu": "npu graph", + "musa": "musa graph", }, ) msg += ( diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 56b9565db..094ca7e83 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -4506,7 +4506,7 @@ class ServerArgs: "--device", type=str, default=ServerArgs.device, - help="The device to use ('cuda', 'xpu', 'hpu', 'npu', 'cpu'). Defaults to auto-detection if not specified.", + help="The device to use ('cuda', 'xpu', 'hpu', 'npu', 'cpu', 'musa'). Defaults to auto-detection if not specified.", ) parser.add_argument( "--tensor-parallel-size", diff --git a/python/sglang/srt/speculative/eagle_utils.py b/python/sglang/srt/speculative/eagle_utils.py index 8ce0a27b0..8b6f85ce8 100644 --- a/python/sglang/srt/speculative/eagle_utils.py +++ b/python/sglang/srt/speculative/eagle_utils.py @@ -4,13 +4,14 @@ from typing import List, Optional import torch -from sglang.srt.utils import is_cuda, is_hip, is_npu +from sglang.srt.utils import is_cuda, is_hip, is_musa, is_npu _is_cuda = is_cuda() _is_hip = is_hip() _is_npu = is_npu() +_is_musa = is_musa() -if _is_cuda or _is_hip: +if _is_cuda or _is_hip or _is_musa: from sgl_kernel import ( build_tree_kernel_efficient as sgl_build_tree_kernel_efficient, ) @@ -169,7 +170,7 @@ def verify_tree_greedy_func( target_predict: torch.Tensor, topk: int = -1, ): - if _is_cuda or _is_hip: + if _is_cuda or _is_hip or _is_musa: from sgl_kernel import verify_tree_greedy verify_tree_greedy( diff --git a/python/sglang/srt/speculative/eagle_worker.py b/python/sglang/srt/speculative/eagle_worker.py index cf99854fb..384b90621 100644 --- a/python/sglang/srt/speculative/eagle_worker.py +++ b/python/sglang/srt/speculative/eagle_worker.py @@ -71,12 +71,14 @@ from sglang.srt.utils import ( empty_context, get_available_gpu_memory, is_cuda, + is_musa, is_npu, next_power_of_2, ) from sglang.srt.utils.patch_torch import monkey_patch_torch_reductions _is_npu = is_npu() +_is_musa = is_musa() if is_cuda(): from sgl_kernel import segment_packbits # noqa: F401 @@ -1214,7 +1216,7 @@ class EAGLEWorker(TpModelWorker): return success, message -@torch.compile(dynamic=True, disable=_is_npu) +@torch.compile(dynamic=True, disable=(_is_npu or _is_musa)) def get_last_loc_large_page_size_top_k_1( req_to_token: torch.Tensor, req_pool_indices: torch.Tensor, diff --git a/python/sglang/srt/speculative/spec_utils.py b/python/sglang/srt/speculative/spec_utils.py index 05421c202..3199ed4ce 100644 --- a/python/sglang/srt/speculative/spec_utils.py +++ b/python/sglang/srt/speculative/spec_utils.py @@ -177,7 +177,8 @@ def assign_draft_cache_locs( mask = copy_offset < copy_len data = tl.load(out_cache_ptr + copy_offset, mask=mask) tl.store(token_pool + kv_start + copy_offset, data, mask=mask) - if page_size != 1 and topk != 1 and duplicate_cache_len > 0: + # XXX (MUSA): Triton issue: chained boolean operators (A or B or C) are not supported. + if (page_size != 1 and topk != 1) and duplicate_cache_len > 0: # Part 2: Copy indices into source_cache_loc and target_cache_loc # Expected output: src:[8,9,10,8,9,10...] tgt:[16,17,18,24,25,26...] prefix_len = tl.load(seq_lens + pid) diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py index d5f4cec31..8c7cb4ae2 100644 --- a/python/sglang/srt/utils/common.py +++ b/python/sglang/srt/utils/common.py @@ -1530,7 +1530,7 @@ def get_amdgpu_memory_capacity(): def get_device_sm(): - if torch.cuda.is_available(): + if torch.cuda.is_available() or is_musa(): major, minor = torch.cuda.get_device_capability() return major * 10 + minor return 0 @@ -2074,6 +2074,8 @@ def direct_register_custom_op( my_lib.impl(op_name, op_func, "PrivateUse1") elif is_xpu(): my_lib.impl(op_name, op_func, "XPU") + elif is_musa(): + my_lib.impl(op_name, op_func, "MUSA") else: my_lib.impl(op_name, op_func, "CUDA") if fake_impl is not None: