sglang 0.1.20__tar.gz → 0.1.21__tar.gz

{sglang-0.1.20/sglang.egg-info → sglang-0.1.21}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: sglang
-Version: 0.1.20
+Version: 0.1.21
 Summary: A structured generation langauge for LLMs.
 License:                                  Apache License
                                    Version 2.0, January 2004
@@ -623,6 +623,14 @@ python -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --port
 python -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --port 30000 --mem-fraction-static 0.7
 ```
 - See [hyperparameter_tuning.md](docs/hyperparameter_tuning.md) on tuning hyperparameters for better performance.
+- Add `--nnodes 2` to run tensor parallelism on multiple nodes. If you have two nodes with two GPUs on each node and want to run TP=4, let `sgl-dev-1` be the hostname of the first node and `50000` be an available port.
+```
+# Node 0
+python -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --tp 4 --nccl-init sgl-dev-1:50000 --nnodes 2 --node-rank 0
+
+# Node 1
+python -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --tp 4 --nccl-init sgl-dev-1:50000 --nnodes 2 --node-rank 1
+```
 
 ### Supported Models
 - Llama

{sglang-0.1.20 → sglang-0.1.21}/README.md

@@ -377,6 +377,14 @@ python -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --port
 python -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --port 30000 --mem-fraction-static 0.7
 ```
 - See [hyperparameter_tuning.md](docs/hyperparameter_tuning.md) on tuning hyperparameters for better performance.
+- Add `--nnodes 2` to run tensor parallelism on multiple nodes. If you have two nodes with two GPUs on each node and want to run TP=4, let `sgl-dev-1` be the hostname of the first node and `50000` be an available port.
+```
+# Node 0
+python -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --tp 4 --nccl-init sgl-dev-1:50000 --nnodes 2 --node-rank 0
+
+# Node 1
+python -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --tp 4 --nccl-init sgl-dev-1:50000 --nnodes 2 --node-rank 1
+```
 
 ### Supported Models
 - Llama

{sglang-0.1.20 → sglang-0.1.21}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "sglang"
-version = "0.1.20"
+version = "0.1.21"
 description = "A structured generation langauge for LLMs."
 readme = "README.md"
 requires-python = ">=3.8"

{sglang-0.1.20 → sglang-0.1.21}/sglang/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "0.1.20"
+__version__ = "0.1.21"
 
 # SGL API Components
 from sglang.api import (

{sglang-0.1.20 → sglang-0.1.21}/sglang/backend/runtime_endpoint.py

@@ -12,7 +12,6 @@ from sglang.utils import http_request
 
 
 class RuntimeEndpoint(BaseBackend):
-
     def __init__(
         self,
         base_url: str,
@@ -38,7 +37,8 @@ class RuntimeEndpoint(BaseBackend):
         self.model_info = res.json()
 
         self.chat_template = get_chat_template_by_model_path(
-            self.model_info["model_path"])
+            self.model_info["model_path"]
+        )
 
     def get_model_name(self):
         return self.model_info["model_path"]
@@ -124,7 +124,12 @@ class RuntimeEndpoint(BaseBackend):
         else:
             raise RuntimeError(f"Invalid dtype: {sampling_params.dtype}")
 
-        for item in ["return_logprob", "logprob_start_len", "top_logprobs_num", "return_text_in_logprobs"]:
+        for item in [
+            "return_logprob",
+            "logprob_start_len",
+            "top_logprobs_num",
+            "return_text_in_logprobs",
+        ]:
             value = getattr(sampling_params, item, None)
             if value is not None:
                 data[item] = value
@@ -171,7 +176,12 @@ class RuntimeEndpoint(BaseBackend):
         else:
             raise RuntimeError(f"Invalid dtype: {sampling_params.dtype}")
 
-        for item in ["return_logprob", "logprob_start_len", "top_logprobs_num", "return_text_in_logprobs"]:
+        for item in [
+            "return_logprob",
+            "logprob_start_len",
+            "top_logprobs_num",
+            "return_text_in_logprobs",
+        ]:
             value = getattr(sampling_params, item, None)
             if value is not None:
                 data[item] = value

{sglang-0.1.20 → sglang-0.1.21}/sglang/bench_latency.py

@@ -32,7 +32,6 @@ import logging
 import multiprocessing
 import time
 
-
 import numpy as np
 import torch
 import torch.distributed as dist

{sglang-0.1.20 → sglang-0.1.21}/sglang/global_config.py

@@ -25,7 +25,8 @@ class GlobalConfig:
         # This can improve the speed for large batch sizes during prefill.
         self.layer_sync_threshold = 8192
 
-        # Runtime constants: Flashinfer
+        # Runtime constants: others
+        self.num_continue_decode_steps = 10
         self.flashinfer_workspace_size = 192 * 1024 * 1024
 
         # Output tokenization configs
@@ -44,4 +45,5 @@ class GlobalConfig:
         # adjust_cache: Adjust the position embedding of KV cache.
         self.concate_and_append_mode = "no_adjust"
 
+
 global_config = GlobalConfig()

{sglang-0.1.20 → sglang-0.1.21}/sglang/lang/chat_template.py

@@ -84,7 +84,7 @@ register_chat_template(
             "system": ("SYSTEM:", "\n"),
             "user": ("USER:", "\n"),
             "assistant": ("ASSISTANT:", "\n"),
-        }
+        },
     )
 )
 
@@ -177,7 +177,7 @@ register_chat_template(
             "assistant": ("", "<|im_end|>\n"),
         },
         style=ChatTemplateStyle.PLAIN,
-        stop_str=("<|im_end|>",)
+        stop_str=("<|im_end|>",),
     )
 )
 

{sglang-0.1.20 → sglang-0.1.21}/sglang/lang/ir.py

@@ -24,9 +24,9 @@ class SglSamplingParams:
     presence_penalty: float = 0.0
     ignore_eos: bool = False
     return_logprob: Optional[bool] = None
-    logprob_start_len: Optional[int] = None,
-    top_logprobs_num: Optional[int] = None,
-    return_text_in_logprobs: Optional[bool] = None,
+    logprob_start_len: Optional[int] = (None,)
+    top_logprobs_num: Optional[int] = (None,)
+    return_text_in_logprobs: Optional[bool] = (None,)
 
     # for constrained generation, not included in to_xxx_kwargs
     dtype: Optional[str] = None

{sglang-0.1.20 → sglang-0.1.21}/sglang/srt/managers/controller/cuda_graph_runner.py

@@ -8,7 +8,10 @@ from vllm.distributed.parallel_state import graph_capture
 from sglang.global_config import global_config
 from sglang.srt.layers.logits_processor import LogitProcessorOutput
 from sglang.srt.managers.controller.infer_batch import (
-    Batch, ForwardMode, InputMetadata, init_flashinfer_args
+    Batch,
+    ForwardMode,
+    InputMetadata,
+    init_flashinfer_args,
 )
 
 
@@ -24,18 +27,28 @@ class CudaGraphRunner:
         # Common inputs
         self.max_bs = max_batch_size_to_capture
         self.input_ids = torch.zeros((self.max_bs,), dtype=torch.int32, device="cuda")
-        self.req_pool_indices = torch.zeros((self.max_bs,), dtype=torch.int32, device="cuda")
+        self.req_pool_indices = torch.zeros(
+            (self.max_bs,), dtype=torch.int32, device="cuda"
+        )
         self.seq_lens = torch.ones((self.max_bs,), dtype=torch.int32, device="cuda")
-        self.position_ids_offsets = torch.zeros((self.max_bs,), dtype=torch.int32, device="cuda")
-        self.out_cache_loc = torch.zeros((self.max_bs,), dtype=torch.int32, device="cuda")
+        self.position_ids_offsets = torch.zeros(
+            (self.max_bs,), dtype=torch.int32, device="cuda"
+        )
+        self.out_cache_loc = torch.zeros(
+            (self.max_bs,), dtype=torch.int32, device="cuda"
+        )
 
         # FlashInfer inputs
-        self.flashinfer_workspace_buffer = self.model_runner.flashinfer_workspace_buffers[0]
+        self.flashinfer_workspace_buffer = (
+            self.model_runner.flashinfer_workspace_buffers[0]
+        )
         self.flashinfer_kv_indptr = torch.zeros(
             (self.max_bs + 1,), dtype=torch.int32, device="cuda"
         )
         self.flashinfer_kv_indices = torch.zeros(
-            (self.max_bs * model_runner.model_config.context_len,), dtype=torch.int32, device="cuda"
+            (self.max_bs * model_runner.model_config.context_len,),
+            dtype=torch.int32,
+            device="cuda",
         )
         self.flashinfer_kv_last_page_len = torch.ones(
             (self.max_bs,), dtype=torch.int32, device="cuda"
@@ -49,7 +62,12 @@ class CudaGraphRunner:
         with graph_capture() as graph_capture_context:
             self.stream = graph_capture_context.stream
             for bs in batch_size_list:
-                graph, input_buffers, output_buffers, flashinfer_handler = self.capture_one_batch_size(bs)
+                (
+                    graph,
+                    input_buffers,
+                    output_buffers,
+                    flashinfer_handler,
+                ) = self.capture_one_batch_size(bs)
                 self.graphs[bs] = graph
                 self.input_buffers[bs] = input_buffers
                 self.output_buffers[bs] = output_buffers
@@ -71,17 +89,19 @@ class CudaGraphRunner:
 
         # FlashInfer inputs
         if not _grouped_size_compiled_for_decode_kernels(
-            self.model_runner.model_config.num_attention_heads // self.model_runner.tp_size,
+            self.model_runner.model_config.num_attention_heads
+            // self.model_runner.tp_size,
             self.model_runner.model_config.get_num_kv_heads(self.model_runner.tp_size),
         ):
             use_tensor_cores = True
         else:
             use_tensor_cores = False
         flashinfer_decode_wrapper = BatchDecodeWithPagedKVCacheWrapper(
-            self.flashinfer_workspace_buffer, "NHD",
+            self.flashinfer_workspace_buffer,
+            "NHD",
             use_cuda_graph=True,
             use_tensor_cores=use_tensor_cores,
-            paged_kv_indptr_buffer=self.flashinfer_kv_indptr[:bs+1],
+            paged_kv_indptr_buffer=self.flashinfer_kv_indptr[: bs + 1],
             paged_kv_indices_buffer=self.flashinfer_kv_indices,
             paged_kv_last_page_len_buffer=self.flashinfer_kv_last_page_len[:bs],
         )
@@ -163,10 +183,14 @@ class CudaGraphRunner:
         else:
             output = LogitProcessorOutput(
                 next_token_logits=output.next_token_logits[:raw_bs],
-                next_token_logprobs=output.next_token_logprobs[:raw_bs] if output.next_token_logprobs is not None else None,
+                next_token_logprobs=output.next_token_logprobs[:raw_bs]
+                if output.next_token_logprobs is not None
+                else None,
                 normalized_prompt_logprobs=None,
                 prefill_token_logprobs=None,
                 prefill_top_logprobs=None,
-                decode_top_logprobs=output.decode_top_logprobs[:raw_bs] if output.decode_top_logprobs is not None else None,
+                decode_top_logprobs=output.decode_top_logprobs[:raw_bs]
+                if output.decode_top_logprobs is not None
+                else None,
             )
         return output

{sglang-0.1.20 → sglang-0.1.21}/sglang/srt/managers/controller/infer_batch.py

@@ -174,9 +174,6 @@ class Req:
 
         return False, ""
 
-    def max_new_tokens(self):
-        return self.sampling_params.max_new_tokens
-
     def check_finished(self):
         if self.finished():
             return
@@ -352,7 +349,7 @@ class Batch:
         extend_num_tokens = seq_lens.sum() - prefix_lens.sum()
         out_cache_loc = self.token_to_kv_pool.alloc(extend_num_tokens)
         if out_cache_loc is None:
-            self.tree_cache.evict(extend_num_tokens, self.token_to_kv_pool.dec_refs)
+            self.tree_cache.evict(extend_num_tokens, self.token_to_kv_pool.free)
             out_cache_loc = self.token_to_kv_pool.alloc(extend_num_tokens)
 
             if out_cache_loc is None:
@@ -422,7 +419,7 @@ class Batch:
         if self.token_to_kv_pool.available_size() >= bs:
             return True
 
-        self.tree_cache.evict(bs, self.token_to_kv_pool.dec_refs)
+        self.tree_cache.evict(bs, self.token_to_kv_pool.free)
 
         if self.token_to_kv_pool.available_size() >= bs:
             return True
@@ -453,7 +450,7 @@ class Batch:
             token_indices = self.req_to_token_pool.req_to_token[
                 req_pool_indices_cpu[idx]
             ][last_uncached_pos : seq_lens_cpu[idx]]
-            self.token_to_kv_pool.dec_refs(token_indices)
+            self.token_to_kv_pool.free(token_indices)
 
             # release the last node
             self.tree_cache.dec_lock_ref(req.last_node)
@@ -596,8 +593,7 @@ class Batch:
             "logit_bias",
         ]:
             self_val = getattr(self, item, None)
-            # logit_bias can be None
-            if self_val is not None:
+            if self_val is not None:  # logit_bias can be None
                 setattr(self, item, self_val[new_indices])
 
     def merge(self, other: "Batch"):
@@ -668,7 +664,9 @@ class Batch:
             sampled_index = torch.multinomial(probs_sort, num_samples=1)
         except RuntimeError as e:
             warnings.warn(f"Ignore errors in sampling: {e}")
-            sampled_index = torch.ones(probs_sort.shape[:-1] + (1,), dtype=torch.int64, device=probs.device)
+            sampled_index = torch.ones(
+                probs_sort.shape[:-1] + (1,), dtype=torch.int64, device=probs.device
+            )
         batch_next_token_ids = torch.gather(probs_idx, dim=1, index=sampled_index).view(
             -1
         )
@@ -749,8 +747,14 @@ class InputMetadata:
         skip_flashinfer_init=False,
     ):
         if not skip_flashinfer_init and not model_runner.server_args.disable_flashinfer:
-            init_flashinfer_args(forward_mode, model_runner, req_pool_indices, seq_lens, prefix_lens,
-                                 model_runner.flashinfer_decode_wrapper)
+            init_flashinfer_args(
+                forward_mode,
+                model_runner,
+                req_pool_indices,
+                seq_lens,
+                prefix_lens,
+                model_runner.flashinfer_decode_wrapper,
+            )
 
         batch_size = len(req_pool_indices)
 
@@ -807,16 +811,24 @@ class InputMetadata:
         )
 
         if model_runner.server_args.disable_flashinfer:
-            (ret.triton_max_seq_len,
-             ret.triton_max_extend_len,
-             ret.triton_start_loc,
-             ret.triton_prefix_lens) = init_triton_args(forward_mode, seq_lens, prefix_lens)
+            (
+                ret.triton_max_seq_len,
+                ret.triton_max_extend_len,
+                ret.triton_start_loc,
+                ret.triton_prefix_lens,
+            ) = init_triton_args(forward_mode, seq_lens, prefix_lens)
 
         return ret
 
 
-def init_flashinfer_args(forward_mode, model_runner, req_pool_indices, seq_lens, prefix_lens,
-                         flashinfer_decode_wrapper):
+def init_flashinfer_args(
+    forward_mode,
+    model_runner,
+    req_pool_indices,
+    seq_lens,
+    prefix_lens,
+    flashinfer_decode_wrapper,
+):
     num_qo_heads = model_runner.model_config.num_attention_heads // model_runner.tp_size
     num_kv_heads = model_runner.model_config.get_num_kv_heads(model_runner.tp_size)
     head_dim = model_runner.model_config.head_dim
@@ -827,9 +839,7 @@ def init_flashinfer_args(forward_mode, model_runner, req_pool_indices, seq_lens,
     else:
         paged_kernel_lens = prefix_lens
 
-    kv_indptr = torch.zeros(
-        (batch_size + 1,), dtype=torch.int32, device="cuda"
-    )
+    kv_indptr = torch.zeros((batch_size + 1,), dtype=torch.int32, device="cuda")
     kv_indptr[1:] = torch.cumsum(paged_kernel_lens, dim=0)
     req_pool_indices_cpu = req_pool_indices.cpu().numpy()
     paged_kernel_lens_cpu = paged_kernel_lens.cpu().numpy()
@@ -842,9 +852,7 @@ def init_flashinfer_args(forward_mode, model_runner, req_pool_indices, seq_lens,
         ],
         dim=0,
     ).contiguous()
-    kv_last_page_len = torch.ones(
-        (batch_size,), dtype=torch.int32, device="cuda"
-    )
+    kv_last_page_len = torch.ones((batch_size,), dtype=torch.int32, device="cuda")
 
     if forward_mode == ForwardMode.DECODE:
         flashinfer_decode_wrapper.end_forward()
@@ -859,9 +867,7 @@ def init_flashinfer_args(forward_mode, model_runner, req_pool_indices, seq_lens,
         )
     else:
         # extend part
-        qo_indptr = torch.zeros(
-            (batch_size + 1,), dtype=torch.int32, device="cuda"
-        )
+        qo_indptr = torch.zeros((batch_size + 1,), dtype=torch.int32, device="cuda")
         qo_indptr[1:] = torch.cumsum(seq_lens - prefix_lens, dim=0)
 
         model_runner.flashinfer_prefill_wrapper_ragged.end_forward()

{sglang-0.1.20 → sglang-0.1.21}/sglang/srt/managers/controller/manager_multi.py

@@ -42,6 +42,8 @@ class LoadBalanceMethod(Enum):
 
 
 class Controller:
+    """A controller that manages multiple data parallel workers."""
+
     def __init__(
         self,
         load_balance_method: str,
@@ -183,9 +185,11 @@ def start_controller_process(
     except Exception:
         pipe_writer.send(get_exception_traceback())
         raise
-
     pipe_writer.send("init ok")
-    loop = asyncio.get_event_loop()
+
+    loop = asyncio.new_event_loop()
+    loop.set_default_executor(ThreadPoolExecutor(max_workers=256))
+
     asyncio.set_event_loop(loop)
     loop.create_task(controller.loop_for_recv_requests())
     loop.run_until_complete(controller.loop_for_forward())

sglang-0.1.21/sglang/srt/managers/controller/manager_single.py

@@ -0,0 +1,177 @@
+"""A controller that manages a group of tensor parallel workers."""
+
+import multiprocessing
+import logging
+import os
+import pickle
+
+import torch
+import torch.distributed as dist
+import zmq
+import zmq.asyncio
+
+from sglang.srt.managers.controller.tp_worker import ModelTpServer
+from sglang.srt.server_args import PortArgs, ServerArgs, ModelPortArgs
+from sglang.srt.utils import kill_parent_process
+from sglang.utils import get_exception_traceback
+
+
+logger = logging.getLogger("srt.controller")
+
+
+def run_tp_server(
+    gpu_id: int,
+    tp_rank: int,
+    server_args: ServerArgs,
+    model_port_args: ModelPortArgs,
+    model_overide_args: dict,
+):
+    """Run a tp server."""
+    try:
+        model_server = ModelTpServer(
+            gpu_id,
+            tp_rank,
+            server_args,
+            model_port_args,
+            model_overide_args,
+        )
+        tp_cpu_group = model_server.model_runner.tp_group.cpu_group
+
+        while True:
+            recv_reqs = broadcast_recv_input(None, tp_rank, tp_cpu_group)
+            model_server.exposed_step(recv_reqs)
+    except Exception:
+        logger.error("Exception in run_tp_server:\n" + get_exception_traceback())
+        raise
+
+
+def launch_tp_servers(gpu_ids, tp_rank_range, server_args,
+                      model_port_args, model_overide_args):
+    """Launch multiple tp servers."""
+    procs = []
+    for i in tp_rank_range:
+        proc = multiprocessing.Process(target=run_tp_server, args=(
+            gpu_ids[i], i, server_args, model_port_args, model_overide_args
+        ))
+        proc.start()
+        procs.append(proc)
+
+    return procs
+
+
+def broadcast_recv_input(data, rank, dist_group):
+    """Broadcast inputs from rank=0 to all other ranks with torch.dist backend."""
+
+    if rank == 0:
+        if len(data) == 0:
+            tensor_size = torch.tensor([0], dtype=torch.long)
+            dist.broadcast(tensor_size, src=0, group=dist_group)
+        else:
+            serialized_data = pickle.dumps(data)
+            size = len(serialized_data)
+            tensor_data = torch.ByteTensor(list(serialized_data))
+            tensor_size = torch.tensor([size], dtype=torch.long)
+
+            dist.broadcast(tensor_size, src=0, group=dist_group)
+            dist.broadcast(tensor_data, src=0, group=dist_group)
+    else:
+        tensor_size = torch.tensor([0], dtype=torch.long)
+        dist.broadcast(tensor_size, src=0, group=dist_group)
+        size = tensor_size.item()
+
+        if size == 0:
+            return []
+
+        tensor_data = torch.empty(size, dtype=torch.uint8)
+        dist.broadcast(tensor_data, src=0, group=dist_group)
+
+        serialized_data = bytes(tensor_data.tolist())
+        data = pickle.loads(serialized_data)
+        return data
+
+
+class ControllerSingle:
+    """A controller that manages a group of tensor parallel workers."""
+
+    def __init__(self, server_args: ServerArgs, port_args: PortArgs, model_overide_args: dict):
+        # Parse args
+        self.server_args = server_args
+
+        # Init communication
+        context = zmq.Context(2)
+        self.recv_from_tokenizer = context.socket(zmq.PULL)
+        self.recv_from_tokenizer.bind(f"tcp://127.0.0.1:{port_args.router_port}")
+
+        self.send_to_detokenizer = context.socket(zmq.PUSH)
+        self.send_to_detokenizer.connect(
+            f"tcp://127.0.0.1:{port_args.detokenizer_port}"
+        )
+
+        # Init model server
+        tp_size_local = server_args.tp_size // server_args.nnodes
+        gpu_ids = [i for _ in range(server_args.nnodes) for i in range(tp_size_local)]
+
+        # Launch other tp ranks
+        if tp_size_local > 1:
+            tp_rank_range = range(1, tp_size_local)
+            self.tp_procs = launch_tp_servers(
+                gpu_ids, tp_rank_range, server_args,
+                port_args.model_port_args[0], model_overide_args)
+
+        # Launch tp rank 0
+        self.tp_server = ModelTpServer(
+            gpu_ids[0],
+            0,
+            server_args,
+            port_args.model_port_args[0],
+            model_overide_args,
+        )
+        self.tp_cpu_group = self.tp_server.model_runner.tp_group.cpu_group
+
+    def loop_for_forward(self):
+        while True:
+            recv_reqs = self.recv_requests()
+
+            if self.server_args.tp_size > 1:
+                broadcast_recv_input(recv_reqs, 0, self.tp_cpu_group)
+
+            out_pyobjs = self.tp_server.exposed_step(recv_reqs)
+
+            for obj in out_pyobjs:
+                self.send_to_detokenizer.send_pyobj(obj)
+
+    def recv_requests(self):
+        recv_reqs = []
+        while True:
+            try:
+                recv_req = self.recv_from_tokenizer.recv_pyobj(zmq.NOBLOCK)
+                recv_reqs.append(recv_req)
+            except zmq.ZMQError:
+                break
+        return recv_reqs
+
+
+def start_controller_process(
+    server_args: ServerArgs, port_args: PortArgs, pipe_writer, model_overide_args: dict
+):
+    logging.basicConfig(
+        level=getattr(logging, server_args.log_level.upper()),
+        format="%(message)s",
+    )
+
+    try:
+        controller = ControllerSingle(server_args, port_args, model_overide_args)
+    except Exception:
+        pipe_writer.send(get_exception_traceback())
+        raise
+
+    pipe_writer.send("init ok")
+
+    try:
+        controller.loop_for_forward()
+    except Exception:
+        logger.error("Exception in ControllerSingle:\n" + get_exception_traceback())
+    finally:
+        for t in controller.tp_procs:
+            os.kill(t.pid, 9)
+        kill_parent_process()

{sglang-0.1.20 → sglang-0.1.21}/sglang/srt/managers/controller/model_runner.py

@@ -11,12 +11,17 @@ import torch
 import torch.nn as nn
 from vllm.config import DeviceConfig, LoadConfig
 from vllm.config import ModelConfig as VllmModelConfig
-from vllm.distributed import init_distributed_environment, initialize_model_parallel
+from vllm.distributed import init_distributed_environment, initialize_model_parallel, get_tp_group
 from vllm.model_executor.model_loader import get_model
 from vllm.model_executor.models import ModelRegistry
 
 from sglang.global_config import global_config
-from sglang.srt.managers.controller.infer_batch import Batch, ForwardMode, InputMetadata, global_server_args_dict
+from sglang.srt.managers.controller.infer_batch import (
+    Batch,
+    ForwardMode,
+    InputMetadata,
+    global_server_args_dict,
+)
 from sglang.srt.memory_pool import ReqToTokenPool, TokenToKVPool
 from sglang.srt.server_args import ServerArgs
 from sglang.srt.utils import (
@@ -70,6 +75,7 @@ class ModelRunner:
             distributed_init_method=nccl_init_method,
         )
         initialize_model_parallel(tensor_model_parallel_size=self.tp_size)
+        self.tp_group = get_tp_group()
         total_gpu_memory = get_available_gpu_memory(
             self.gpu_id, distributed=self.tp_size > 1
         )
@@ -83,7 +89,9 @@ class ModelRunner:
 
         # Set some global args
         global_server_args_dict["disable_flashinfer"] = server_args.disable_flashinfer
-        global_server_args_dict["attention_reduce_in_fp32"] = server_args.attention_reduce_in_fp32
+        global_server_args_dict[
+            "attention_reduce_in_fp32"
+        ] = server_args.attention_reduce_in_fp32
 
         # Load the model and create memory pool
         self.load_model()
@@ -217,7 +225,9 @@ class ModelRunner:
             self.flashinfer_workspace_buffers[1], "NHD"
         )
         self.flashinfer_decode_wrapper = BatchDecodeWithPagedKVCacheWrapper(
-            self.flashinfer_workspace_buffers[0], "NHD", use_tensor_cores=use_tensor_cores
+            self.flashinfer_workspace_buffers[0],
+            "NHD",
+            use_tensor_cores=use_tensor_cores,
         )
 
     def init_cuda_graphs(self):
@@ -229,7 +239,9 @@ class ModelRunner:
 
         logger.info(f"[gpu_id={self.gpu_id}] Capture cuda graph begin.")
         batch_size_list = [1, 2, 4] + [i * 8 for i in range(1, 16)]
-        self.cuda_graph_runner = CudaGraphRunner(self, max_batch_size_to_capture=max(batch_size_list))
+        self.cuda_graph_runner = CudaGraphRunner(
+            self, max_batch_size_to_capture=max(batch_size_list)
+        )
         self.cuda_graph_runner.capture(batch_size_list)
 
     @torch.inference_mode()