diff --git a/python/sglang/srt/_custom_ops.py b/python/sglang/srt/_custom_ops.py index de47707c1..3353aa2ea 100644 --- a/python/sglang/srt/_custom_ops.py +++ b/python/sglang/srt/_custom_ops.py @@ -4,32 +4,20 @@ from typing import List, Optional, Tuple import torch -from sglang.srt.utils import get_bool_env_var, is_hip, is_hpu, is_npu +from sglang.srt.utils import is_hip, is_hpu, is_npu logger = logging.getLogger(__name__) -use_vllm_custom_allreduce = get_bool_env_var( - "USE_VLLM_CUSTOM_ALLREDUCE", default="false" -) + if not is_hpu(): - # ROCm does not use vllm custom allreduce - if use_vllm_custom_allreduce and not is_hip(): - try: - import vllm._C # noqa: F401 - except ImportError as e: - logger.warning("Failed to import from vllm._C with %r", e) - else: - try: - import sgl_kernel - except ImportError as e: - logger.warning("Failed to import from custom_ar with %r", e) + try: + import sgl_kernel + except ImportError as e: + logger.warning("Failed to import from custom_ar with %r", e) if not is_hip() and not is_npu(): - if use_vllm_custom_allreduce: - custom_op = torch.ops._C_custom_ar - else: - custom_op = sgl_kernel.allreduce + custom_op = sgl_kernel.allreduce # custom allreduce def init_custom_ar( diff --git a/python/sglang/srt/configs/falcon_h1.py b/python/sglang/srt/configs/falcon_h1.py index b8869b4ff..1f524b892 100644 --- a/python/sglang/srt/configs/falcon_h1.py +++ b/python/sglang/srt/configs/falcon_h1.py @@ -19,7 +19,6 @@ from transformers.configuration_utils import PretrainedConfig from transformers.utils import logging from sglang.srt.configs.mamba_utils import Mamba2CacheParams, Mamba2StateShape -from sglang.srt.layers.dp_attention import get_tensor_model_parallel_world_size logger = logging.get_logger(__name__) @@ -297,8 +296,10 @@ class FalconH1Config(PretrainedConfig): @property def mamba2_cache_params(self): + from sglang.srt.layers.dp_attention import get_attention_tp_size + shape = Mamba2StateShape.create( - tp_world_size=get_tensor_model_parallel_world_size(), + tp_world_size=get_attention_tp_size(), intermediate_size=self.mamba_intermediate, n_groups=self.mamba_n_groups, num_heads=self.mamba_n_heads, diff --git a/python/sglang/srt/configs/nemotron_h.py b/python/sglang/srt/configs/nemotron_h.py index 9e156f6a7..b73b146fe 100644 --- a/python/sglang/srt/configs/nemotron_h.py +++ b/python/sglang/srt/configs/nemotron_h.py @@ -20,7 +20,6 @@ from transformers.configuration_utils import PretrainedConfig from transformers.utils import logging from sglang.srt.configs.mamba_utils import Mamba2CacheParams, Mamba2StateShape -from sglang.srt.layers.dp_attention import get_attention_tp_size logger = logging.get_logger(__name__) @@ -273,6 +272,8 @@ class NemotronHConfig(PretrainedConfig): @property def mamba2_cache_params(self) -> Mamba2CacheParams: + from sglang.srt.layers.dp_attention import get_attention_tp_size + shape = Mamba2StateShape.create( tp_world_size=get_attention_tp_size(), intermediate_size=self.mamba_num_heads * self.mamba_head_dim, diff --git a/python/sglang/srt/configs/qwen3_next.py b/python/sglang/srt/configs/qwen3_next.py index 630227a2c..cd1b6f1ea 100644 --- a/python/sglang/srt/configs/qwen3_next.py +++ b/python/sglang/srt/configs/qwen3_next.py @@ -21,7 +21,6 @@ from transformers.modeling_rope_utils import rope_config_validation from transformers.utils import logging from sglang.srt.configs.mamba_utils import Mamba2CacheParams, Mamba2StateShape -from sglang.srt.layers.dp_attention import get_attention_tp_size logger = logging.get_logger(__name__) @@ -277,6 +276,8 @@ class Qwen3NextConfig(PretrainedConfig): @property def mamba2_cache_params(self) -> Mamba2CacheParams: + from sglang.srt.layers.dp_attention import get_attention_tp_size + shape = Mamba2StateShape.create( tp_world_size=get_attention_tp_size(), intermediate_size=self.linear_value_head_dim * self.linear_num_value_heads, diff --git a/python/sglang/srt/distributed/device_communicators/custom_all_reduce.py b/python/sglang/srt/distributed/device_communicators/custom_all_reduce.py index 634bd4aa4..452341a5a 100644 --- a/python/sglang/srt/distributed/device_communicators/custom_all_reduce.py +++ b/python/sglang/srt/distributed/device_communicators/custom_all_reduce.py @@ -21,24 +21,19 @@ from sglang.srt.distributed.parallel_state import in_the_same_node_as from sglang.srt.environ import envs from sglang.srt.utils import is_cuda, is_hip, log_info_on_rank0 -logger = logging.getLogger(__name__) +try: + # Use custom allreduce from sgl kernel (ROCM and TRT-LLM) + import sgl_kernel # noqa: F401 + + custom_ar = True +except ImportError: + # For CPUs + custom_ar = False + _is_cuda = is_cuda() _is_hip = is_hip() - -try: - if ops.use_vllm_custom_allreduce and not _is_hip: - # Use vLLM custom allreduce - ops.meta_size() - else: - # Use custom allreduce from sgl kernel (ROCM and TRT-LLM) - import sgl_kernel # noqa: F401 - custom_ar = True -except Exception: - # For CPUs - custom_ar = False - logger = logging.getLogger(__name__) diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index df17263d1..e53204651 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -229,7 +229,6 @@ class Envs: SGLANG_SKIP_SGL_KERNEL_VERSION_CHECK = EnvBool(False) # vLLM dependencies (TODO: they have been deprecated, we can remove them safely) - USE_VLLM_CUSTOM_ALLREDUCE = EnvBool(False) USE_VLLM_CUTLASS_W8A8_FP8_KERNEL = EnvBool(False) USE_TRITON_W8A8_FP8_KERNEL = EnvBool(False) diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py index 375b9b56d..c0f2fb26c 100644 --- a/python/sglang/srt/utils/common.py +++ b/python/sglang/srt/utils/common.py @@ -303,6 +303,7 @@ def xpu_has_xmx_support(): return False +@lru_cache(maxsize=1) def is_flashinfer_available(): """ Check whether flashinfer is available. diff --git a/sgl-kernel/README.md b/sgl-kernel/README.md index 5084bdda8..dc5925a55 100644 --- a/sgl-kernel/README.md +++ b/sgl-kernel/README.md @@ -52,8 +52,8 @@ make build ```cpp // We need def with schema here for torch.compile m.def( - "bmm_fp8(Tensor A, Tensor B, Tensor! D, Tensor A_scale, Tensor B_scale, Tensor workspace_buffer, int " - "cublas_handle, int cuda_stream) -> ()"); + "bmm_fp8(Tensor A, Tensor B, Tensor! D, Tensor A_scale, Tensor B_scale, Tensor workspace_buffer, " + "int cublas_handle) -> ()"); m.impl("bmm_fp8", torch::kCUDA, &bmm_fp8); ``` diff --git a/sgl-kernel/csrc/common_extension.cc b/sgl-kernel/csrc/common_extension.cc index a36e9f495..44f99424d 100644 --- a/sgl-kernel/csrc/common_extension.cc +++ b/sgl-kernel/csrc/common_extension.cc @@ -90,13 +90,13 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) { m.def( "apply_rope_pos_ids_cos_sin_cache(Tensor q, Tensor k, Tensor! q_rope, Tensor! k_rope, Tensor cos_sin_cache, " - "Tensor pos_ids, bool interleave, bool enable_pdl, int cuda_stream, " + "Tensor pos_ids, bool interleave, bool enable_pdl, " "Tensor? v, Tensor!? k_buffer, Tensor!? v_buffer, Tensor? kv_cache_loc) -> ()"); m.impl("apply_rope_pos_ids_cos_sin_cache", torch::kCUDA, &apply_rope_pos_ids_cos_sin_cache); m.def( - "downcast_fp8(Tensor k, Tensor v, Tensor k_out, Tensor v_out, Tensor k_scale, Tensor v_scale, Tensor loc, int " - "mult, int offset, int cuda_stream) -> ()"); + "downcast_fp8(Tensor k, Tensor v, Tensor k_out, Tensor v_out, Tensor k_scale, Tensor v_scale, Tensor loc, " + "int mult, int offset) -> ()"); m.impl("downcast_fp8", torch::kCUDA, &downcast_fp8); m.def("copy_to_gpu_no_ce(Tensor input, Tensor! output) -> ()"); @@ -303,13 +303,13 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) { "Tensor candidates, Tensor retrive_index, Tensor retrive_next_token, Tensor retrive_next_sibling, " "Tensor uniform_samples, Tensor uniform_samples_for_final_sampling, Tensor target_probs, Tensor draft_probs, " "float threshold_single, float threshold_acc, " - "bool deterministic, int cuda_stream) -> ()"); + "bool deterministic) -> ()"); m.impl("tree_speculative_sampling_target_only", torch::kCUDA, &tree_speculative_sampling_target_only); m.def( "verify_tree_greedy(Tensor! predicts, Tensor! accept_index, Tensor! accept_token_num, " "Tensor candidates, Tensor retrive_index, Tensor retrive_next_token, Tensor retrive_next_sibling, " - "Tensor target_predict, int cuda_stream) -> ()"); + "Tensor target_predict) -> ()"); m.impl("verify_tree_greedy", torch::kCUDA, &verify_tree_greedy); m.def( @@ -403,8 +403,8 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) { * From FlashInfer */ m.def( - "bmm_fp8(Tensor A, Tensor B, Tensor! D, Tensor A_scale, Tensor B_scale, Tensor workspace_buffer, int " - "cublas_handle, int cuda_stream) -> ()", + "bmm_fp8(Tensor A, Tensor B, Tensor! D, Tensor A_scale, Tensor B_scale, Tensor workspace_buffer, " + "int cublas_handle) -> ()", {at::Tag::needs_fixed_stride_order}); m.impl("bmm_fp8", torch::kCUDA, &bmm_fp8); diff --git a/sgl-kernel/csrc/common_extension_rocm.cc b/sgl-kernel/csrc/common_extension_rocm.cc index 6ca8091de..54c520641 100644 --- a/sgl-kernel/csrc/common_extension_rocm.cc +++ b/sgl-kernel/csrc/common_extension_rocm.cc @@ -106,7 +106,7 @@ TORCH_LIBRARY_EXPAND(sgl_kernel, m) { m.def( "verify_tree_greedy(Tensor! predicts, Tensor! accept_index, Tensor! accept_token_num, " "Tensor candidates, Tensor retrive_index, Tensor retrive_next_token, Tensor retrive_next_sibling, " - "Tensor target_predict, int cuda_stream) -> ()"); + "Tensor target_predict) -> ()"); m.impl("verify_tree_greedy", torch::kCUDA, &verify_tree_greedy); m.def( diff --git a/sgl-kernel/csrc/elementwise/cast.cu b/sgl-kernel/csrc/elementwise/cast.cu index a1ff8703f..3ce8135de 100644 --- a/sgl-kernel/csrc/elementwise/cast.cu +++ b/sgl-kernel/csrc/elementwise/cast.cu @@ -150,14 +150,13 @@ void downcast_fp8( at::Tensor& v_scale, at::Tensor& loc, int64_t mult, - int64_t offset, - int64_t cuda_stream) { + int64_t offset) { CHECK_INPUT(k); CHECK_INPUT(v); CHECK_INPUT(k_out); CHECK_INPUT(v_out); - cudaStream_t stream = reinterpret_cast(cuda_stream); + cudaStream_t stream = at::cuda::getCurrentCUDAStream(); switch (k.scalar_type()) { case at::ScalarType::BFloat16: downcast_fp8_impl<__nv_bfloat16>(k, v, k_out, v_out, k_scale, v_scale, loc, mult, offset, stream); diff --git a/sgl-kernel/csrc/elementwise/rope.cu b/sgl-kernel/csrc/elementwise/rope.cu index 041558f61..23bc87660 100644 --- a/sgl-kernel/csrc/elementwise/rope.cu +++ b/sgl-kernel/csrc/elementwise/rope.cu @@ -28,7 +28,6 @@ void apply_rope_pos_ids_cos_sin_cache( at::Tensor pos_ids, bool interleave, bool enable_pdl, - int64_t cuda_stream, const std::optional& v, const std::optional& k_buffer, const std::optional& v_buffer, @@ -88,7 +87,7 @@ void apply_rope_pos_ids_cos_sin_cache( size_t k_rope_stride_n = k_rope.stride(0); size_t k_rope_stride_h = k_rope.stride(1); - cudaStream_t stream = reinterpret_cast(cuda_stream); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FP16(q.scalar_type(), c_type, [&] { // TODO temporarily only use `BatchQKApplyRotaryPosIdsCosSinCacheEnhanced` when save_kv_cache // to avoid changing original code path; but this branch is feature-complete and should switch to this later diff --git a/sgl-kernel/csrc/gemm/bmm_fp8.cu b/sgl-kernel/csrc/gemm/bmm_fp8.cu index 4a82b4b27..cef85a7de 100644 --- a/sgl-kernel/csrc/gemm/bmm_fp8.cu +++ b/sgl-kernel/csrc/gemm/bmm_fp8.cu @@ -27,8 +27,7 @@ void bmm_fp8( at::Tensor A_scale, at::Tensor B_scale, at::Tensor workspace_buffer, - int64_t cublas_handle, - int64_t cuda_stream) { + int64_t cublas_handle) { TORCH_CHECK(A.is_cuda(), "A must be a CUDA tensor"); TORCH_CHECK(B.is_cuda(), "B must be a CUDA tensor"); TORCH_CHECK(D.is_cuda(), "D must be a CUDA tensor"); @@ -51,7 +50,7 @@ void bmm_fp8( auto n = B.size(2); auto lt_handle = reinterpret_cast(cublas_handle); - auto stream = reinterpret_cast(cuda_stream); + auto stream = at::cuda::getCurrentCUDAStream(); auto status = flashinfer::bmm_fp8::bmm_fp8_internal_cublaslt( workspace_buffer.data_ptr(), diff --git a/sgl-kernel/csrc/speculative/eagle_utils.cu b/sgl-kernel/csrc/speculative/eagle_utils.cu index 7bf5db274..e8e306325 100644 --- a/sgl-kernel/csrc/speculative/eagle_utils.cu +++ b/sgl-kernel/csrc/speculative/eagle_utils.cu @@ -328,8 +328,7 @@ void verify_tree_greedy( at::Tensor retrive_index, at::Tensor retrive_next_token, at::Tensor retrive_next_sibling, - at::Tensor target_predict, - int64_t cuda_stream = 0) { + at::Tensor target_predict) { CHECK_INPUT(candidates); CHECK_INPUT(retrive_index); CHECK_INPUT(retrive_next_token); @@ -389,7 +388,7 @@ void verify_tree_greedy( throw std::runtime_error("Expected 'target_predict' to be of type long (torch.int64)."); } - cudaStream_t stream = reinterpret_cast(cuda_stream); + cudaStream_t stream = at::cuda::getCurrentCUDAStream(); dim3 grid(batch_size); dim3 block(1); diff --git a/sgl-kernel/csrc/speculative/speculative_sampling.cu b/sgl-kernel/csrc/speculative/speculative_sampling.cu index ca545e99e..7a1400f8e 100644 --- a/sgl-kernel/csrc/speculative/speculative_sampling.cu +++ b/sgl-kernel/csrc/speculative/speculative_sampling.cu @@ -42,8 +42,7 @@ void tree_speculative_sampling_target_only( at::Tensor draft_probs, double threshold_single, double threshold_acc, - bool deterministic = true, - int64_t cuda_stream = 0) { + bool deterministic = true) { CHECK_INPUT(candidates); CHECK_INPUT(retrive_index); CHECK_INPUT(retrive_next_token); @@ -124,7 +123,7 @@ void tree_speculative_sampling_target_only( CHECK_GE(threshold_acc, 0); CHECK_GE(1, threshold_acc); - cudaStream_t stream = reinterpret_cast(cuda_stream); + cudaStream_t stream = at::cuda::getCurrentCUDAStream(); cudaError_t status = sampling::TreeSpeculativeSamplingTargetOnly( static_cast(predicts.data_ptr()), static_cast(accept_index.data_ptr()), diff --git a/sgl-kernel/include/sgl_kernel_ops.h b/sgl-kernel/include/sgl_kernel_ops.h index a022b80dc..688910d02 100644 --- a/sgl-kernel/include/sgl_kernel_ops.h +++ b/sgl-kernel/include/sgl_kernel_ops.h @@ -152,7 +152,6 @@ void apply_rope_pos_ids_cos_sin_cache( at::Tensor pos_ids, bool interleave, bool enable_pdl, - int64_t cuda_stream, const std::optional& v, const std::optional& k_buffer, const std::optional& v_buffer, @@ -167,8 +166,7 @@ void downcast_fp8( at::Tensor& v_scale, at::Tensor& loc, int64_t mult, - int64_t offset, - int64_t cuda_stream); + int64_t offset); void copy_to_gpu_no_ce(const at::Tensor& input, at::Tensor& output); void concat_mla_k(torch::Tensor k, torch::Tensor k_nope, torch::Tensor k_rope); @@ -253,8 +251,7 @@ void bmm_fp8( at::Tensor A_scale, at::Tensor B_scale, at::Tensor workspace_buffer, - int64_t cublas_handle, - int64_t cuda_stream); + int64_t cublas_handle); void dsv3_router_gemm(torch::Tensor& output, const torch::Tensor& mat_a, const torch::Tensor& mat_b); void dsv3_fused_a_gemm(torch::Tensor& output, torch::Tensor const& mat_a, torch::Tensor const& mat_b); @@ -471,8 +468,7 @@ void tree_speculative_sampling_target_only( at::Tensor draft_probs, double threshold_single = 1, double threshold_acc = 1, - bool deterministic = true, - int64_t cuda_stream = 0); + bool deterministic = true); void verify_tree_greedy( at::Tensor predicts, // mutable @@ -482,8 +478,7 @@ void verify_tree_greedy( at::Tensor retrive_index, at::Tensor retrive_next_token, at::Tensor retrive_next_sibling, - at::Tensor target_predict, - int64_t cuda_stream = 0); + at::Tensor target_predict); void reconstruct_indices_from_tree_mask( at::Tensor tree_mask, diff --git a/sgl-kernel/python/sgl_kernel/elementwise.py b/sgl-kernel/python/sgl_kernel/elementwise.py index 13bb11be3..3b402366e 100644 --- a/sgl-kernel/python/sgl_kernel/elementwise.py +++ b/sgl-kernel/python/sgl_kernel/elementwise.py @@ -2,7 +2,7 @@ from dataclasses import dataclass from typing import List, Optional import torch -from sgl_kernel.utils import get_cuda_stream, is_arch_support_pdl +from sgl_kernel.utils import is_arch_support_pdl # These implementations extensively draw from and build upon the FlashInfer project https://github.com/flashinfer-ai/flashinfer @@ -263,6 +263,10 @@ class FusedSetKVBufferArg: cache_loc: torch.Tensor +def _view_3d(x, head_size): + return x.view(x.shape[0], -1, head_size) + + def apply_rope_with_cos_sin_cache_inplace( positions: torch.Tensor, query: torch.Tensor, @@ -317,31 +321,27 @@ def apply_rope_with_cos_sin_cache_inplace( assert a.v_scale is None, "v_scale is not yet supported" assert a.cache_loc.dtype == torch.int64, f"{a.cache_loc.dtype=}" - def _view_3d(x): - return x.view(x.shape[0], -1, head_size) - torch.ops.sgl_kernel.apply_rope_pos_ids_cos_sin_cache.default( - _view_3d(query), - _view_3d(key), - _view_3d(query), - _view_3d(key), + _view_3d(query, head_size), + _view_3d(key, head_size), + _view_3d(query, head_size), + _view_3d(key, head_size), cos_sin_cache, positions.long(), (not is_neox), enable_pdl, - get_cuda_stream(), ( - _view_3d(fused_set_kv_buffer_arg.value) + _view_3d(fused_set_kv_buffer_arg.value, head_size) if fused_set_kv_buffer_arg is not None else None ), ( - _view_3d(fused_set_kv_buffer_arg.k_buffer) + _view_3d(fused_set_kv_buffer_arg.k_buffer, head_size) if fused_set_kv_buffer_arg is not None else None ), ( - _view_3d(fused_set_kv_buffer_arg.v_buffer) + _view_3d(fused_set_kv_buffer_arg.v_buffer, head_size) if fused_set_kv_buffer_arg is not None else None ), @@ -365,7 +365,7 @@ def downcast_fp8( offset: int = 0, ) -> None: torch.ops.sgl_kernel.downcast_fp8( - k, v, k_out, v_out, k_scale, v_scale, loc, mult, offset, get_cuda_stream() + k, v, k_out, v_out, k_scale, v_scale, loc, mult, offset ) diff --git a/sgl-kernel/python/sgl_kernel/gemm.py b/sgl-kernel/python/sgl_kernel/gemm.py index 3af4751fc..4e23ebdd9 100644 --- a/sgl-kernel/python/sgl_kernel/gemm.py +++ b/sgl-kernel/python/sgl_kernel/gemm.py @@ -2,7 +2,7 @@ from typing import Optional, Tuple import torch from sgl_kernel.scalar_type import ScalarType -from sgl_kernel.utils import _get_cache_buf, get_cuda_stream +from sgl_kernel.utils import _get_cache_buf def awq_dequantize( @@ -60,7 +60,6 @@ def _bmm_fp8_internal( B_scale, workspace_buffer, cublas_handle, - get_cuda_stream(), ) diff --git a/sgl-kernel/python/sgl_kernel/speculative.py b/sgl-kernel/python/sgl_kernel/speculative.py index ee4bf7fb0..17e424972 100644 --- a/sgl-kernel/python/sgl_kernel/speculative.py +++ b/sgl-kernel/python/sgl_kernel/speculative.py @@ -1,5 +1,4 @@ import torch -from sgl_kernel.utils import get_cuda_stream def tree_speculative_sampling_target_only( @@ -33,7 +32,6 @@ def tree_speculative_sampling_target_only( threshold_single, threshold_acc, deterministic, - get_cuda_stream(), ) @@ -56,7 +54,6 @@ def verify_tree_greedy( retrive_next_token, retrive_next_sibling, target_predict, - get_cuda_stream(), ) diff --git a/sgl-kernel/python/sgl_kernel/utils.py b/sgl-kernel/python/sgl_kernel/utils.py index f2fa0b617..d03476eff 100644 --- a/sgl-kernel/python/sgl_kernel/utils.py +++ b/sgl-kernel/python/sgl_kernel/utils.py @@ -18,11 +18,6 @@ from typing import Dict, Tuple import torch - -def get_cuda_stream() -> int: - return torch.cuda.current_stream().cuda_stream - - _cache_buf: Dict[Tuple[str, torch.device], torch.Tensor] = {}