sglang 0.5.4__py3-none-any.whl → 0.5.4.post2__py3-none-any.whl

sglang/srt/configs/model_config.py

@@ -205,6 +205,14 @@ class ModelConfig:
             self.hf_config, "image_token_id", None
         ) or getattr(self.hf_config, "image_token_index", None)
 
+        # matryoshka embeddings
+        self.matryoshka_dimensions = getattr(
+            self.hf_config, "matryoshka_dimensions", None
+        )
+        self.is_matryoshka = self.matryoshka_dimensions or getattr(
+            self.hf_config, "is_matryoshka", False
+        )
+
     @staticmethod
     def from_server_args(
         server_args: ServerArgs,
@@ -358,6 +366,13 @@ class ModelConfig:
             self.qk_rope_head_dim = self.hf_text_config.qk_rope_head_dim
             self.v_head_dim = self.hf_text_config.v_head_dim
             self.qk_nope_head_dim = self.hf_text_config.qk_nope_head_dim
+        elif "KimiLinearForCausalLM" in self.hf_config.architectures:
+            self.head_dim = 72
+            self.attention_arch = AttentionArch.MLA
+            self.kv_lora_rank = self.hf_config.kv_lora_rank
+            self.qk_rope_head_dim = self.hf_config.qk_rope_head_dim
+            self.v_head_dim = self.hf_config.v_head_dim
+            self.qk_nope_head_dim = self.hf_config.qk_nope_head_dim
         else:
             if (
                 "MistralModel" in self.hf_config.architectures
@@ -535,7 +550,7 @@ class ModelConfig:
                 quant_cfg = self._parse_modelopt_quant_config(quant_config_dict)
         return quant_cfg
 
-    def _parse_modelopt_quant_config(self, quant_config_dict: dict) -> dict:
+    def _parse_modelopt_quant_config(self, quant_config_dict: dict) -> Optional[dict]:
         """Parse ModelOpt quantization config and return the appropriate quant_method."""
         json_quant_configs = quant_config_dict["quantization"]
         quant_algo = json_quant_configs.get("quant_algo", None)
@@ -547,8 +562,7 @@ class ModelConfig:
         elif quant_algo and "FP8" in quant_algo:
             return {"quant_method": "modelopt_fp8"}
         else:
-            # Default to FP8 for backward compatibility
-            return {"quant_method": "modelopt_fp8"}
+            return None
 
     def _is_already_quantized(self) -> bool:
         """Check if the model is already quantized based on config files."""
@@ -583,14 +597,20 @@ class ModelConfig:
             return
 
         # Check if ModelOpt quantization is specified
-        modelopt_quantization_specified = self.quantization in [
+        _MODELOPT_QUANTIZATION_METHODS = [
             "modelopt",
             "modelopt_fp8",
             "modelopt_fp4",
         ]
+        modelopt_quantization_specified = (
+            self.quantization in _MODELOPT_QUANTIZATION_METHODS
+        )
 
         if not modelopt_quantization_specified:
-            raise ValueError("quantize_and_serve requires ModelOpt quantization")
+            raise ValueError(
+                "quantize_and_serve requires ModelOpt quantization (set with --quantization "
+                f"{{{', '.join(sorted(_MODELOPT_QUANTIZATION_METHODS))}}})"
+            )
 
         # quantize_and_serve is disabled due to compatibility issues
         raise NotImplementedError(
@@ -614,6 +634,7 @@ class ModelConfig:
             "petit_nvfp4",
             "quark",
             "mxfp4",
+            "auto-round",
         ]
         optimized_quantization_methods = [
             "fp8",
@@ -635,6 +656,7 @@ class ModelConfig:
             "petit_nvfp4",
         ]
         compatible_quantization_methods = {
+            "modelopt_fp8": ["modelopt"],
             "modelopt_fp4": ["modelopt"],
             "petit_nvfp4": ["modelopt"],
             "w8a8_int8": ["compressed-tensors", "compressed_tensors"],
@@ -806,7 +828,7 @@ def _get_and_verify_dtype(
 ) -> torch.dtype:
     # NOTE: getattr(config, "torch_dtype", torch.float32) is not correct
     # because config.torch_dtype can be None.
-    config_dtype = getattr(config, "torch_dtype", None)
+    config_dtype = getattr(config, "dtype", None)
     if isinstance(config_dtype, str):
         config_dtype = _STR_DTYPE_TO_TORCH_DTYPE.get(config_dtype, None)
     if config_dtype is None:
@@ -915,12 +937,13 @@ multimodal_model_archs = [
     "InternVLChatModel",
     "InternS1ForConditionalGeneration",
     "Phi4MMForCausalLM",
-    "VILAForConditionalGeneration",
     "Step3VLForConditionalGeneration",
     "POINTSV15ChatModel",
     "DotsVLMForCausalLM",
     "DotsOCRForCausalLM",
     "Sarashina2VisionForCausalLM",
+    "NVILAForConditionalGeneration",
+    "NVILALiteForConditionalGeneration",
     "DeepseekOCRForCausalLM",
 ]
 

sglang/srt/constants.py

@@ -1,3 +1,10 @@
 # GPU Memory Types
 GPU_MEMORY_TYPE_KV_CACHE = "kv_cache"
 GPU_MEMORY_TYPE_WEIGHTS = "weights"
+GPU_MEMORY_TYPE_CUDA_GRAPH = "cuda_graph"
+
+GPU_MEMORY_ALL_TYPES = [
+    GPU_MEMORY_TYPE_KV_CACHE,
+    GPU_MEMORY_TYPE_WEIGHTS,
+    GPU_MEMORY_TYPE_CUDA_GRAPH,
+]

sglang/srt/debug_utils/tensor_dump_forward_hook.py

@@ -0,0 +1,149 @@
+"""
+This file provides a function `register_forward_hook_for_model` that registers a forward hook on every operator of the model.
+After registration, during model inference, all tensors generated throughout the forward pass will be recorded.
+
+Usage:
+Specify the output directory for dumping tensors using the argument `--debug-tensor-dump-output-folder`.
+A separate directory will be created for each GPU rank, named in the format `f"TP{tp_rank}_PP{pp_rank}_Rank{rank}_pid{pid}"`.
+Each complete forward pass of the model generates a `.pt` file named `f"Pass{pass_num}.pt"`, which can be loaded using `torch.load`.
+The file contains a series of key-value pairs, where the keys correspond to operator names in the model
+(similar to those in model.safetensors.index.json), and the values are the outputs produced by the respective operators.
+"""
+
+import logging
+import os
+from pathlib import Path
+
+import torch
+
+from sglang.srt.layers.logits_processor import LogitsProcessorOutput
+from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors
+
+logger = logging.getLogger(__name__)
+
+
+class TensorDumper:
+    def __init__(
+        self, dump_dir: str, dump_layers: int, tp_size: int, tp_rank: int, pp_rank: int
+    ):
+        self._dump_layers = dump_layers
+        self._forward_pass_id = 0
+        self._pid = os.getpid()
+        self._current_tensors = {}
+        self._base_dir = Path(dump_dir)
+        rank = tp_size * pp_rank + tp_rank
+        self._process_dir = (
+            self._base_dir / f"TP{tp_rank}_PP{pp_rank}_Rank{rank}_pid{self._pid}"
+        )
+        self._process_dir.mkdir(parents=True, exist_ok=True)
+
+    def get_dump_dir(self):
+        return str(self._process_dir)
+
+    def add_tensor(self, name, tensor_item):
+        if isinstance(tensor_item, (tuple, list)):
+            tensors = [t.cpu() for t in tensor_item if t is not None]
+            if len(tensors) == 1:
+                self._current_tensors[name] = tensors[0]
+            else:
+                self._current_tensors[name] = tensors
+        elif isinstance(tensor_item, torch.Tensor):
+            self._current_tensors[name] = tensor_item.cpu()
+        elif isinstance(tensor_item, LogitsProcessorOutput):
+            self._current_tensors[name] = tensor_item.next_token_logits.cpu()
+        elif isinstance(tensor_item, ForwardBatch):
+            self._current_tensors[name + ".forward_batch_info.input_ids"] = (
+                tensor_item.input_ids.cpu()
+            )
+            self._current_tensors[name + ".forward_batch_info.seq_lens"] = (
+                tensor_item.seq_lens.cpu()
+            )
+            self._current_tensors[name + ".forward_batch_info.positions"] = (
+                tensor_item.positions.cpu()
+            )
+        elif isinstance(tensor_item, PPProxyTensors):
+            for tensor_name in tensor_item.tensors.keys():
+                self._current_tensors[name + ".pp_proxy_tensors." + tensor_name] = (
+                    tensor_item.tensors[tensor_name].cpu()
+                )
+        else:
+            logger.warning(f"Unsupported type: {type(tensor_item)}: {tensor_item}")
+
+    def dump_current_tensors(self):
+        if len(self._current_tensors) == 0:
+            return
+        tensor_file_for_pass = self._process_dir / f"Pass{self._forward_pass_id:05d}.pt"
+        logger.info(
+            f"Dump {self._forward_pass_id:05d}th pass to {tensor_file_for_pass}"
+        )
+        torch.save(self._current_tensors, str(tensor_file_for_pass))
+        self._current_tensors = {}
+        self._forward_pass_id += 1
+
+    def _add_hook_recursive(
+        self, model, prefix, top_level_module_name, layers_module_name
+    ):
+        model_top_level_module_matched = False
+        layers_prefix = top_level_module_name + "." + layers_module_name
+        for name, module in model._modules.items():
+            top_level_model = False
+            if len(prefix) == 0:
+                cur_name = name
+                if cur_name == top_level_module_name:
+                    model_top_level_module_matched = True
+                    top_level_model = True
+            else:
+                cur_name = prefix + "." + name
+            if self._dump_layers > 0 and name.isdigit() and prefix == layers_prefix:
+                # If we only need n layers, skip the reset layers.
+                # Most models' layout is like model.layers.0.
+                cur_layer = int(name)
+                if cur_layer >= self._dump_layers:
+                    continue
+            if module is not None:
+                _, sub_count = self._add_hook_recursive(
+                    module, cur_name, top_level_module_name, layers_module_name
+                )
+                if sub_count == 0 or top_level_model:
+                    # Avoid duplicated output hooks, e.g. self_attn may contain:
+                    # self_attn.qkv_proj, self_attn.attn & self_attn.o_proj.
+                    # Therefore, we do not need to add output hooks for self_attn,
+                    # since the output of self_attn should be the same to self_attn.o_proj.
+                    module.register_forward_hook(
+                        self._dump_hook(cur_name, top_level_model)
+                    )
+        return model_top_level_module_matched, len(model._modules.items())
+
+    def _dump_hook(self, tensor_name, do_dump):
+        def inner_dump_hook(module, input, output):
+            if do_dump:
+                # This is the top-level model, so we will record the input for it.
+                for item in input:
+                    if isinstance(item, ForwardBatch):
+                        self.add_tensor(tensor_name, item)
+                self.dump_current_tensors()
+            if output is not None:
+                self.add_tensor(tensor_name, output)
+
+        return inner_dump_hook
+
+
+def register_forward_hook_for_model(
+    model, dump_dir: str, dump_layers: int, tp_size: int, tp_rank: int, pp_rank: int
+):
+    tensor_dumper = TensorDumper(dump_dir, dump_layers, tp_size, tp_rank, pp_rank)
+    # Most models have the layerout like:
+    # XxxxForCausalLM
+    #     (model): XxxxModel
+    #         (layers): ModuleList
+    # If the model is not constructed with this layout,
+    # environment variable can be used to specify the module names.
+    top_level_module_name = os.getenv("TENSOR_DUMP_TOP_LEVEL_MODULE_NAME", "model")
+    layers_module_name = os.getenv("TENSOR_DUMP_LAYERS_MODULE_NAME", "layers")
+    model_top_level_module_matched, _ = tensor_dumper._add_hook_recursive(
+        model, "", top_level_module_name, layers_module_name
+    )
+    assert (
+        model_top_level_module_matched
+    ), f"model should have a module named {top_level_module_name}"
+    return tensor_dumper

sglang/srt/disaggregation/decode.py

@@ -58,6 +58,11 @@ from sglang.srt.mem_cache.memory_pool import (
     ReqToTokenPool,
     SWAKVPool,
 )
+from sglang.srt.tracing.trace import (
+    trace_event_batch,
+    trace_slice_batch,
+    trace_slice_end,
+)
 from sglang.srt.utils import get_int_env_var, require_mlp_sync
 from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter
 
@@ -313,6 +318,7 @@ class DecodePreallocQueue:
             )
 
             req.add_latency(RequestStage.DECODE_PREPARE)
+            trace_slice_end(RequestStage.DECODE_PREPARE, req.rid, auto_next_anon=True)
             self.queue.append(
                 DecodeRequest(req=req, kv_receiver=kv_receiver, waiting_for_input=False)
             )
@@ -521,13 +527,15 @@ class DecodePreallocQueue:
             decode_req.kv_receiver.init(
                 page_indices, decode_req.metadata_buffer_index, state_indices
             )
-            decode_req.req.add_latency(RequestStage.DECODE_BOOTSTRAP)
             preallocated_reqs.append(decode_req)
             indices_to_remove.add(i)
             decode_req.req.time_stats.decode_transfer_queue_entry_time = (
                 time.perf_counter()
             )
             decode_req.req.add_latency(RequestStage.DECODE_BOOTSTRAP)
+            trace_slice_end(
+                RequestStage.DECODE_BOOTSTRAP, decode_req.req.rid, auto_next_anon=True
+            )
 
         self.queue = [
             entry for i, entry in enumerate(self.queue) if i not in indices_to_remove
@@ -765,8 +773,12 @@ class DecodeTransferQueue:
                 indices_to_remove.add(i)
                 decode_req.req.time_stats.wait_queue_entry_time = time.perf_counter()
 
-                # special handling for sampling_params.max_new_tokens == 1
-                if decode_req.req.sampling_params.max_new_tokens == 1:
+                # special handling for corner cases
+                should_finish = (
+                    decode_req.req.sampling_params.max_new_tokens == 1
+                    or output_id in decode_req.req.eos_token_ids
+                )
+                if should_finish:
                     # finish immediately
                     decode_req.req.time_stats.forward_entry_time = (
                         decode_req.req.time_stats.completion_time
@@ -776,8 +788,19 @@ class DecodeTransferQueue:
                         [decode_req.req], decode_req.req.return_logprob
                     )
                     self.tree_cache.cache_finished_req(decode_req.req)
+                    trace_slice_end(
+                        RequestStage.DECODE_QUICK_FINISH,
+                        decode_req.req.rid,
+                        thread_finish_flag=True,
+                    )
                 else:
                     transferred_reqs.append(decode_req.req)
+                    trace_slice_end(
+                        RequestStage.DECODE_TRANSFERRED,
+                        decode_req.req.rid,
+                        auto_next_anon=True,
+                    )
+
             elif poll in [
                 KVPoll.Bootstrapping,
                 KVPoll.WaitingForInput,
@@ -823,6 +846,7 @@ class SchedulerDisaggregationDecodeMixin:
                     self.stream_output(
                         batch.reqs, any(req.return_logprob for req in batch.reqs)
                     )
+                    trace_slice_batch(RequestStage.DECODE_FAKE_OUTPUT, batch.reqs)
                     if prepare_mlp_sync_flag:
                         self._prepare_idle_batch_and_run(None)
                 else:
@@ -872,6 +896,7 @@ class SchedulerDisaggregationDecodeMixin:
                     self.stream_output(
                         batch.reqs, any(req.return_logprob for req in batch.reqs)
                     )
+                    trace_slice_batch(RequestStage.DECODE_FAKE_OUTPUT, batch.reqs)
                     if prepare_mlp_sync_flag:
                         batch_, batch_result = self._prepare_idle_batch_and_run(
                             None, delay_process=True
@@ -954,6 +979,9 @@ class SchedulerDisaggregationDecodeMixin:
                 self.running_batch = self.update_running_batch(self.running_batch)
                 ret = self.running_batch if not self.running_batch.is_empty() else None
 
+        if ret:
+            attrs = {"bid": hex(id(ret)), "batch_size": ret.batch_size()}
+            trace_event_batch("schedule", ret.reqs, attrs=attrs)
         return ret
 
     def get_new_prebuilt_batch(self: Scheduler) -> Optional[ScheduleBatch]:
@@ -1009,6 +1037,9 @@ class SchedulerDisaggregationDecodeMixin:
         return new_batch
 
     def process_decode_queue(self: Scheduler):
+        if self.server_args.disaggregation_decode_enable_offload_kvcache:
+            self.decode_offload_manager.check_offload_progress()
+
         # try to resume retracted requests if there are enough space for another `num_reserved_decode_tokens` decode steps
         resumed_reqs = self.disagg_decode_prealloc_queue.resume_retracted_reqs()
         self.waiting_queue.extend(resumed_reqs)
@@ -1031,6 +1062,3 @@ class SchedulerDisaggregationDecodeMixin:
                 self.disagg_decode_transfer_queue.pop_transferred()
             )  # the requests which kv has arrived
             self.waiting_queue.extend(alloc_reqs)
-
-        if self.server_args.disaggregation_decode_enable_offload_kvcache:
-            self.decode_offload_manager.check_offload_progress()

sglang/srt/disaggregation/nixl/conn.py

@@ -231,8 +231,8 @@ class NixlKVManager(CommonKVManager):
             ]
             for k in keys_to_remove:
                 del self.connection_pool[k]
-            if failed_bootstrap_addr in self.prefill_tp_size_table:
-                del self.prefill_tp_size_table[failed_bootstrap_addr]
+            if failed_bootstrap_addr in self.prefill_attn_tp_size_table:
+                del self.prefill_attn_tp_size_table[failed_bootstrap_addr]
             if failed_bootstrap_addr in self.prefill_dp_size_table:
                 del self.prefill_dp_size_table[failed_bootstrap_addr]
             if failed_bootstrap_addr in self.prefill_pp_size_table:

sglang/srt/disaggregation/prefill.py

@@ -53,6 +53,7 @@ from sglang.srt.mem_cache.memory_pool import (
     NSATokenToKVPool,
     SWAKVPool,
 )
+from sglang.srt.tracing.trace import trace_event_batch, trace_slice, trace_slice_end
 from sglang.srt.utils import broadcast_pyobj, point_to_point_pyobj, require_mlp_sync
 
 if TYPE_CHECKING:
@@ -198,6 +199,7 @@ class PrefillBootstrapQueue:
         self._process_req(req)
         req.add_latency(RequestStage.PREFILL_PREPARE)
         self.queue.append(req)
+        trace_slice_end(RequestStage.PREFILL_PREPARE, req.rid, auto_next_anon=True)
 
     def extend(self, reqs: List[Req], num_kv_heads: int) -> None:
         for req in reqs:
@@ -289,6 +291,10 @@ class PrefillBootstrapQueue:
             req.time_stats.wait_queue_entry_time = time.perf_counter()
             req.add_latency(RequestStage.PREFILL_BOOTSTRAP)
 
+            trace_slice_end(
+                RequestStage.PREFILL_BOOTSTRAP, req.rid, auto_next_anon=True
+            )
+
         self.queue = [
             entry for i, entry in enumerate(self.queue) if i not in indices_to_remove
         ]
@@ -316,6 +322,9 @@ class SchedulerDisaggregationPrefillMixin:
             )
             self.process_prefill_chunk()
             batch = self.get_new_batch_prefill()
+            if batch:
+                attrs = {"bid": hex(id(batch)), "batch_size": batch.batch_size()}
+                trace_event_batch("schedule", batch.reqs, attrs=attrs)
 
             if require_mlp_sync(self.server_args):
                 batch = self.prepare_mlp_sync_batch(batch)
@@ -348,6 +357,9 @@ class SchedulerDisaggregationPrefillMixin:
             )
             self.process_prefill_chunk()
             batch = self.get_new_batch_prefill()
+            if batch:
+                attrs = {"bid": hex(id(batch)), "batch_size": batch.batch_size()}
+                trace_event_batch("schedule", batch.reqs, attrs=attrs)
 
             if require_mlp_sync(self.server_args):
                 batch = self.prepare_mlp_sync_batch(batch)
@@ -423,6 +435,7 @@ class SchedulerDisaggregationPrefillMixin:
                 req.output_ids.append(next_token_id)
                 self.tree_cache.cache_unfinished_req(req)  # update the tree and lock
                 req.add_latency(RequestStage.PREFILL_FORWARD)
+                trace_slice(RequestStage.PREFILL_FORWARD, req.rid, auto_next_anon=True)
                 self.disagg_prefill_inflight_queue.append(req)
                 if self.spec_algorithm.is_eagle() and batch.spec_info is not None:
                     req.output_topk_p = batch.spec_info.topk_p[i]
@@ -487,6 +500,9 @@ class SchedulerDisaggregationPrefillMixin:
 
                 if self.enable_overlap:
                     self.send_kv_chunk(req, last_chunk=False, end_idx=req.tmp_end_idx)
+                trace_slice(
+                    RequestStage.PREFILL_CHUNKED_FORWARD, req.rid, auto_next_anon=True
+                )
 
         self.maybe_send_health_check_signal()
 
@@ -558,6 +574,9 @@ class SchedulerDisaggregationPrefillMixin:
             req.add_latency(RequestStage.PREFILL_TRANSFER_KV_CACHE)
             self.req_to_metadata_buffer_idx_allocator.free(req.metadata_buffer_index)
             req.metadata_buffer_index = -1
+            trace_slice(
+                RequestStage.PREFILL_TRANSFER_KV_CACHE, req.rid, thread_finish_flag=True
+            )
 
         self.disagg_prefill_inflight_queue = undone_reqs
 
@@ -569,7 +588,7 @@ class SchedulerDisaggregationPrefillMixin:
         """
         polls = poll_and_all_reduce(
             [req.disagg_kv_sender for req in self.disagg_prefill_inflight_queue],
-            self.tp_worker.get_tp_group().cpu_group,
+            self.tp_worker.get_attention_tp_cpu_group(),
         )
 
         transferred_rids: List[str] = []
@@ -703,8 +722,11 @@ class SchedulerDisaggregationPrefillMixin:
         else:
             data = None
 
-        if self.tp_size != 1:
+        if self.attn_tp_size != 1:
             data = broadcast_pyobj(
-                data, self.tp_group.rank, self.tp_cpu_group, src=self.tp_group.ranks[0]
+                data,
+                self.attn_tp_group.rank,
+                self.attn_tp_cpu_group,
+                src=self.attn_tp_group.ranks[0],
             )
         return data

sglang/srt/distributed/device_communicators/custom_all_reduce.py

@@ -18,6 +18,7 @@ from sglang.srt.distributed.device_communicators.custom_all_reduce_utils import
     is_weak_contiguous,
 )
 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__)
@@ -210,6 +211,7 @@ class CustomAllreduce:
             self.register_buffer(self.buffer)
 
         self.disabled = False
+        self.tms_cudagraph = envs.SGLANG_MEMORY_SAVER_CUDA_GRAPH.get()
 
     @staticmethod
     def create_shared_buffer(
@@ -394,7 +396,7 @@ class CustomAllreduce:
                 if _is_hip:
                     return self.all_reduce_reg(input)
                 else:
-                    return self.all_reduce(input, registered=True)
+                    return self.all_reduce(input, registered=not self.tms_cudagraph)
             else:
                 # If warm up, mimic the allocation pattern since custom
                 # allreduce is out-of-place.

sglang/srt/distributed/parallel_state.py

@@ -68,7 +68,7 @@ REDUCE_OP_SUM = int(torch.distributed.ReduceOp.SUM)
 
 @dataclass
 class GraphCaptureContext:
-    stream: torch.cuda.Stream if not _is_npu else torch.npu.Stream
+    stream: torch.get_device_module().Stream
 
 
 @dataclass
@@ -340,17 +340,10 @@ class GroupCoordinator:
         self.qr_comm: Optional[QuickAllReduce] = None
         if use_custom_allreduce and self.world_size > 1:
             # Initialize a custom fast all-reduce implementation.
-            if torch_compile is not None and torch_compile:
-                # For piecewise CUDA graph, the requirement for custom allreduce is larger to
-                # avoid illegal cuda memory access.
-                ca_max_size = 256 * 1024 * 1024
-            else:
-                ca_max_size = 8 * 1024 * 1024
             try:
                 self.ca_comm = CustomAllreduce(
                     group=self.cpu_group,
                     device=self.device,
-                    max_size=ca_max_size,
                 )
             except Exception as e:
                 logger.warning(
@@ -505,7 +498,7 @@ class GroupCoordinator:
                 maybe_pynccl_context = nullcontext()
             else:
                 maybe_pynccl_context = pynccl_comm.change_state(
-                    enable=True, stream=torch.cuda.current_stream()
+                    enable=True, stream=torch.get_device_module().current_stream()
                 )
 
             pymscclpp_comm = self.pymscclpp_comm
@@ -562,7 +555,7 @@ class GroupCoordinator:
             and input_.symmetric_memory
         ):
             with self.pynccl_comm.change_state(
-                enable=True, stream=torch.cuda.current_stream()
+                enable=True, stream=torch.get_device_module().current_stream()
             ):
                 self.pynccl_comm.all_reduce(input_)
                 return input_
@@ -662,7 +655,9 @@ class GroupCoordinator:
         world_size = self.world_size
         pynccl_comm = self.pynccl_comm
 
-        with pynccl_comm.change_state(enable=True, stream=torch.cuda.current_stream()):
+        with pynccl_comm.change_state(
+            enable=True, stream=torch.get_device_module().current_stream()
+        ):
             assert (
                 pynccl_comm is not None and not pynccl_comm.disabled
             ), "pynccl is required for reduce_scatterv"
@@ -786,7 +781,9 @@ class GroupCoordinator:
         world_size = self.world_size
         pynccl_comm = self.pynccl_comm
 
-        with pynccl_comm.change_state(enable=True, stream=torch.cuda.current_stream()):
+        with pynccl_comm.change_state(
+            enable=True, stream=torch.get_device_module().current_stream()
+        ):
             assert (
                 pynccl_comm is not None and not pynccl_comm.disabled
             ), "pynccl is required for all_gatherv"