sglang 0.4.4.post2__py3-none-any.whl → 0.4.4.post4__py3-none-any.whl

sglang/bench_serving.py

@@ -44,6 +44,12 @@ ASSISTANT_SUFFIX = "Assistant:"
 global args
 
 
+# don't want to import sglang package here
+def _get_bool_env_var(name: str, default: str = "false") -> bool:
+    value = os.getenv(name, default)
+    return value.lower() in ("true", "1")
+
+
 @dataclass
 class RequestFuncInput:
     prompt: str
@@ -965,10 +971,11 @@ async def benchmark(
     request_rate: float,
     max_concurrency: Optional[int],
     disable_tqdm: bool,
-    lora_name: str,
+    lora_names: List[str],
     extra_request_body: Dict[str, Any],
     profile: bool,
     pd_seperated: bool = False,
+    flush_cache: bool = False,
 ):
     if backend in ASYNC_REQUEST_FUNCS:
         request_func = ASYNC_REQUEST_FUNCS[backend]
@@ -986,8 +993,16 @@ async def benchmark(
             return await request_func(request_func_input=request_func_input, pbar=pbar)
 
     # Warmup
-    print("Starting initial single prompt test run...")
+    print(f"Starting warmup with {args.warmup_requests} sequences...")
+
+    # Use the first request for all warmup iterations
     test_prompt, test_prompt_len, test_output_len = input_requests[0]
+    if lora_names != None and len(lora_names) != 0:
+        lora_name = lora_names[0]
+    else:
+        lora_name = None
+
+    # Create the test input once
     test_input = RequestFuncInput(
         model=model_id,
         prompt=test_prompt,
@@ -997,17 +1012,29 @@ async def benchmark(
         lora_name=lora_name,
         extra_request_body=extra_request_body,
     )
-    test_output = await request_func(request_func_input=test_input)
-    if not test_output.success:
+
+    # Run warmup requests
+    warmup_tasks = []
+    for _ in range(args.warmup_requests):
+        warmup_tasks.append(
+            asyncio.create_task(request_func(request_func_input=test_input))
+        )
+
+    warmup_outputs = await asyncio.gather(*warmup_tasks)
+
+    # Check if at least one warmup request succeeded
+    if not any(output.success for output in warmup_outputs):
         raise ValueError(
-            "Initial test run failed - Please make sure benchmark arguments "
-            f"are correctly specified. Error: {test_output.error}"
+            "Warmup failed - Please make sure benchmark arguments "
+            f"are correctly specified. Error: {warmup_outputs[0].error}"
         )
     else:
-        print("Initial test run completed. Starting main benchmark run...")
+        print(
+            f"Warmup completed with {args.warmup_requests} sequences. Starting main benchmark run..."
+        )
 
     # Flush cache
-    if "sglang" in backend:
+    if ("sglang" in backend and _get_bool_env_var("SGLANG_IS_IN_CI")) or flush_cache:
         requests.post(base_url + "/flush_cache", headers=get_auth_headers())
 
     time.sleep(1.0)
@@ -1028,6 +1055,12 @@ async def benchmark(
     tasks: List[asyncio.Task] = []
     async for request in get_request(input_requests, request_rate):
         prompt, prompt_len, output_len = request
+        if lora_names != None and len(lora_names) != 0:
+            idx = random.randint(0, len(lora_names) - 1)
+            lora_name = lora_names[idx]
+        else:
+            lora_name = None
+
         request_func_input = RequestFuncInput(
             model=model_id,
             prompt=prompt,
@@ -1235,6 +1268,10 @@ def run_benchmark(args_: argparse.Namespace):
     if not hasattr(args, "max_concurrency"):
         args.max_concurrency = None
 
+    # Set default value for warmup_requests if not present
+    if not hasattr(args, "warmup_requests"):
+        args.warmup_requests = 1
+
     print(f"benchmark_args={args}")
 
     # Set global environments
@@ -1336,6 +1373,10 @@ def run_benchmark(args_: argparse.Namespace):
     tokenizer = get_tokenizer(tokenizer_id)
     input_requests = get_dataset(args, tokenizer)
 
+    # compatible with SimpleNamespace
+    if not hasattr(args, "flush_cache"):
+        args.flush_cache = False
+
     return asyncio.run(
         benchmark(
             backend=backend,
@@ -1347,10 +1388,11 @@ def run_benchmark(args_: argparse.Namespace):
             request_rate=args.request_rate,
             max_concurrency=args.max_concurrency,
             disable_tqdm=args.disable_tqdm,
-            lora_name=args.lora_name,
+            lora_names=args.lora_name,
             extra_request_body=extra_request_body,
             profile=args.profile,
             pd_seperated=args.pd_seperated,
+            flush_cache=args.flush_cache,
         )
     )
 
@@ -1366,6 +1408,13 @@ def set_ulimit(target_soft_limit=65535):
             print(f"Fail to set RLIMIT_NOFILE: {e}")
 
 
+class LoRAPathAction(argparse.Action):
+    def __call__(self, parser, namespace, values, option_string=None):
+        setattr(namespace, self.dest, [])
+        for lora_name in values:
+            getattr(namespace, self.dest).append(lora_name)
+
+
 if __name__ == "__main__":
     parser = ArgumentParser(description="Benchmark the online serving throughput.")
     parser.add_argument(
@@ -1509,8 +1558,10 @@ if __name__ == "__main__":
     parser.add_argument(
         "--lora-name",
         type=str,
+        nargs="*",
         default=None,
-        help="The name of LoRA adapter",
+        action=LoRAPathAction,
+        help="The names of LoRA adapters. You can provide a list of names in the format {name} {name} {name}...",
     )
     parser.add_argument(
         "--prompt-suffix",
@@ -1523,6 +1574,17 @@ if __name__ == "__main__":
         action="store_true",
         help="Benchmark PD disaggregation server",
     )
+    parser.add_argument(
+        "--flush-cache",
+        action="store_true",
+        help="Flush the cache before running the benchmark",
+    )
+    parser.add_argument(
+        "--warmup-requests",
+        type=int,
+        default=1,
+        help="Number of warmup requests to run before the benchmark",
+    )
 
     group = parser.add_argument_group("generated-shared-prefix dataset arguments")
     group.add_argument(

sglang/srt/_custom_ops.py

@@ -27,17 +27,20 @@ if not is_hpu():
             logger.warning("Failed to import from custom_ar with %r", e)
 
 
-if use_vllm_custom_allreduce and not is_hip():
-    # vLLM custom allreduce
+if not is_hip():
+    if use_vllm_custom_allreduce:
+        custom_op = torch.ops._C_custom_ar
+    else:
+        custom_op = sgl_kernel.allreduce
+
+    # custom allreduce
     def init_custom_ar(
         ipc_tensors: List[torch.Tensor],
         rank_data: torch.Tensor,
         rank: int,
         full_nvlink: bool,
     ) -> int:
-        return torch.ops._C_custom_ar.init_custom_ar(
-            ipc_tensors, rank_data, rank, full_nvlink
-        )
+        return custom_op.init_custom_ar(ipc_tensors, rank_data, rank, full_nvlink)
 
     def all_reduce(
         fa: int,
@@ -46,105 +49,69 @@ if use_vllm_custom_allreduce and not is_hip():
         reg_buffer: int,
         reg_buffer_sz_bytes: int,
     ) -> None:
-        torch.ops._C_custom_ar.all_reduce(fa, inp, out, reg_buffer, reg_buffer_sz_bytes)
+        custom_op.all_reduce(fa, inp, out, reg_buffer, reg_buffer_sz_bytes)
 
     def dispose(fa: int) -> None:
-        torch.ops._C_custom_ar.dispose(fa)
+        custom_op.dispose(fa)
 
     def meta_size() -> int:
-        return torch.ops._C_custom_ar.meta_size()
+        return custom_op.meta_size()
 
     def register_buffer(fa: int, ipc_tensors: List[int]) -> None:
-        return torch.ops._C_custom_ar.register_buffer(fa, ipc_tensors)
+        return custom_op.register_buffer(fa, ipc_tensors)
 
     def get_graph_buffer_ipc_meta(fa: int) -> Tuple[List[int], List[int]]:
-        return torch.ops._C_custom_ar.get_graph_buffer_ipc_meta(fa)
+        return custom_op.get_graph_buffer_ipc_meta(fa)
 
     def register_graph_buffers(
         fa: int, handles: List[List[int]], offsets: List[List[int]]
     ) -> None:
-        torch.ops._C_custom_ar.register_graph_buffers(fa, handles, offsets)
+        custom_op.register_graph_buffers(fa, handles, offsets)
 
 else:
-    if is_hip():
-        # ROCM custom allreduce
-
-        def init_custom_ar(
-            meta: torch.Tensor,
-            rank_data: torch.Tensor,
-            handles: List[str],
-            offsets: List[int],
-            rank: int,
-            full_nvlink: bool,
-        ) -> int:
-            return sgl_kernel.allreduce.init_custom_ar(
-                meta, rank_data, handles, offsets, rank, full_nvlink
-            )
-
-        def all_reduce_reg(fa: int, inp: torch.Tensor, out: torch.Tensor) -> None:
-            sgl_kernel.allreduce.all_reduce_reg(fa, inp, out)
-
-        def all_reduce_unreg(
-            fa: int, inp: torch.Tensor, reg_buffer: torch.Tensor, out: torch.Tensor
-        ) -> None:
-            sgl_kernel.allreduce.all_reduce_unreg(fa, inp, reg_buffer, out)
-
-        def dispose(fa: int) -> None:
-            sgl_kernel.allreduce.dispose(fa)
-
-        def meta_size() -> int:
-            return sgl_kernel.allreduce.meta_size()
-
-        def register_buffer(
-            fa: int, t: torch.Tensor, handles: List[str], offsets: List[int]
-        ) -> None:
-            return sgl_kernel.allreduce.register_buffer(fa, t, handles, offsets)
-
-        def get_graph_buffer_ipc_meta(fa: int) -> Tuple[torch.Tensor, List[int]]:
-            return sgl_kernel.allreduce.get_graph_buffer_ipc_meta(fa)
-
-        def register_graph_buffers(
-            fa: int, handles: List[str], offsets: List[List[int]]
-        ) -> None:
-            sgl_kernel.allreduce.register_graph_buffers(fa, handles, offsets)
-
-        def allocate_meta_buffer(size: int) -> torch.Tensor:
-            return sgl_kernel.allreduce.allocate_meta_buffer(size)
-
-        def get_meta_buffer_ipc_handle(inp: torch.Tensor) -> torch.Tensor:
-            return sgl_kernel.allreduce.get_meta_buffer_ipc_handle(inp)
+    # ROCM custom allreduce
 
-    else:
-        # TRTLLM custom allreduce
-        def init_custom_ar(
-            rank_id: int,
-            world_size: int,
-            rank_data_base: torch.Tensor,
-            buffers: List[int],
-            tmp_result_buffers: List[int],
-            barrier_in: List[int],
-            barrier_out: List[int],
-        ) -> int:
-            return sgl_kernel.init_custom_reduce(
-                rank_id,
-                world_size,
-                rank_data_base,
-                buffers,
-                tmp_result_buffers,
-                barrier_in,
-                barrier_out,
-            )
-
-        def all_reduce(fa: int, inp: torch.Tensor, out: torch.Tensor) -> None:
-            sgl_kernel.custom_reduce(fa, inp, out)
-
-        def dispose(fa: int) -> None:
-            sgl_kernel.custom_dispose(fa)
-
-        def get_graph_buffer_ipc_meta(fa: int) -> Tuple[List[int], List[int]]:
-            return sgl_kernel.get_graph_buffer_ipc_meta(fa)
-
-        def register_graph_buffers(
-            fa: int, handles: List[List[int]], offsets: List[List[int]]
-        ) -> None:
-            sgl_kernel.register_graph_buffers(fa, handles, offsets)
+    def init_custom_ar(
+        meta: torch.Tensor,
+        rank_data: torch.Tensor,
+        handles: List[str],
+        offsets: List[int],
+        rank: int,
+        full_nvlink: bool,
+    ) -> int:
+        return sgl_kernel.allreduce.init_custom_ar(
+            meta, rank_data, handles, offsets, rank, full_nvlink
+        )
+
+    def all_reduce_reg(fa: int, inp: torch.Tensor, out: torch.Tensor) -> None:
+        sgl_kernel.allreduce.all_reduce_reg(fa, inp, out)
+
+    def all_reduce_unreg(
+        fa: int, inp: torch.Tensor, reg_buffer: torch.Tensor, out: torch.Tensor
+    ) -> None:
+        sgl_kernel.allreduce.all_reduce_unreg(fa, inp, reg_buffer, out)
+
+    def dispose(fa: int) -> None:
+        sgl_kernel.allreduce.dispose(fa)
+
+    def meta_size() -> int:
+        return sgl_kernel.allreduce.meta_size()
+
+    def register_buffer(
+        fa: int, t: torch.Tensor, handles: List[str], offsets: List[int]
+    ) -> None:
+        return sgl_kernel.allreduce.register_buffer(fa, t, handles, offsets)
+
+    def get_graph_buffer_ipc_meta(fa: int) -> Tuple[torch.Tensor, List[int]]:
+        return sgl_kernel.allreduce.get_graph_buffer_ipc_meta(fa)
+
+    def register_graph_buffers(
+        fa: int, handles: List[str], offsets: List[List[int]]
+    ) -> None:
+        sgl_kernel.allreduce.register_graph_buffers(fa, handles, offsets)
+
+    def allocate_meta_buffer(size: int) -> torch.Tensor:
+        return sgl_kernel.allreduce.allocate_meta_buffer(size)
+
+    def get_meta_buffer_ipc_handle(inp: torch.Tensor) -> torch.Tensor:
+        return sgl_kernel.allreduce.get_meta_buffer_ipc_handle(inp)

sglang/srt/configs/deepseekvl2.py

@@ -4,7 +4,6 @@ from dataclasses import dataclass
 from typing import Dict, List, Optional, Tuple
 
 import torch
-import torchvision.transforms as T
 from PIL import Image, ImageOps
 from transformers import (
     AutoProcessor,
@@ -76,6 +75,16 @@ class ImageTransform(object):
         self.std = std
         self.normalize = normalize
 
+        # only load torchvision.transforms when needed
+        try:
+            import torchvision.transforms as T
+
+            # FIXME: add version check for gguf
+        except ImportError as err:
+            raise ImportError(
+                "Please install torchvision via `pip install torchvision` to use Deepseek-VL2."
+            ) from err
+
         transform_pipelines = [T.ToTensor()]
 
         if normalize:

sglang/srt/configs/model_config.py

@@ -22,11 +22,7 @@ import torch
 from transformers import PretrainedConfig
 
 from sglang.srt.hf_transformers_utils import get_config, get_context_length
-from sglang.srt.layers.quantization import (
-    BASE_QUANTIZATION_METHODS,
-    QUANTIZATION_METHODS,
-    VLLM_AVAILABLE,
-)
+from sglang.srt.layers.quantization import QUANTIZATION_METHODS
 from sglang.srt.utils import get_bool_env_var, is_hip
 
 logger = logging.getLogger(__name__)
@@ -239,12 +235,7 @@ class ModelConfig:
 
     # adapted from https://github.com/vllm-project/vllm/blob/v0.6.4.post1/vllm/config.py
     def _verify_quantization(self) -> None:
-        # Select supported quantization methods based on vllm availability
-        if VLLM_AVAILABLE:
-            supported_quantization = [*QUANTIZATION_METHODS]
-        else:
-            supported_quantization = [*BASE_QUANTIZATION_METHODS]
-
+        supported_quantization = [*QUANTIZATION_METHODS]
         rocm_supported_quantization = [
             "awq",
             "gptq",
@@ -267,6 +258,7 @@ class ModelConfig:
             "experts_int8",
             "w8a8_int8",
             "w8a8_fp8",
+            "moe_wna16",
         ]
         compatible_quantization_methods = {
             "w8a8_int8": ["compressed-tensors", "compressed_tensors"],
@@ -282,11 +274,7 @@ class ModelConfig:
             quant_method = quant_cfg.get("quant_method", "").lower()
 
             # Detect which checkpoint is it
-            # Only iterate through currently available quantization methods
-            available_methods = (
-                QUANTIZATION_METHODS if VLLM_AVAILABLE else BASE_QUANTIZATION_METHODS
-            )
-            for _, method in available_methods.items():
+            for _, method in QUANTIZATION_METHODS.items():
                 quantization_override = method.override_quantization_method(
                     quant_cfg, self.quantization
                 )
@@ -467,6 +455,7 @@ def is_generation_model(model_architectures: List[str], is_embedding: bool = Fal
         or "InternLM2ForRewardModel" in model_architectures
         or "Qwen2ForRewardModel" in model_architectures
         or "Qwen2ForSequenceClassification" in model_architectures
+        or "CLIPModel" in model_architectures
     ):
         return False
     else:
@@ -488,6 +477,7 @@ multimodal_model_archs = [
     "MllamaForConditionalGeneration",
     "Qwen2VLForConditionalGeneration",
     "Qwen2_5_VLForConditionalGeneration",
+    "CLIPModel",
 ]
 
 

sglang/srt/constrained/base_grammar_backend.py

@@ -169,7 +169,9 @@ class BaseGrammarBackend(ABC):
             self.cache.clear()
 
 
-def create_grammar_backend(server_args: ServerArgs, tokenizer, vocab_size):
+def create_grammar_backend(
+    server_args: ServerArgs, tokenizer, vocab_size: int
+) -> Optional[BaseGrammarBackend]:
     if server_args.grammar_backend == "outlines":
         from sglang.srt.constrained.outlines_backend import OutlinesGrammarBackend
 
@@ -188,6 +190,8 @@ def create_grammar_backend(server_args: ServerArgs, tokenizer, vocab_size):
             tokenizer=tokenizer,
             whitespace_pattern=server_args.constrained_json_whitespace_pattern,
         )
+    elif server_args.grammar_backend == "none":
+        return None
     else:
         raise ValueError(f"Invalid grammar backend: {server_args.grammar_backend}")
 

sglang/srt/custom_op.py

@@ -50,6 +50,7 @@ if _is_cuda:
     def scaled_fp8_quant(
         input: torch.Tensor,
         scale: Optional[torch.Tensor] = None,
+        num_token_padding: Optional[int] = None,
         use_per_token_if_dynamic: bool = False,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         """
@@ -59,6 +60,8 @@ if _is_cuda:
             input (torch.Tensor): Input tensor to be quantized
             scale (Optional[torch.Tensor]): Pre-computed scaling factor for static quantization.
                 If None, scales will be computed dynamically.
+            num_token_padding (Optional[int]): If specified, pad the first dimension
+                of the output to at least this value.
             use_per_token_if_dynamic (bool): When using dynamic scaling (scale=None),
                 determines the quantization granularity:
                 - True: compute scale per token
@@ -75,6 +78,8 @@ if _is_cuda:
         assert input.ndim == 2, f"Expected 2D input tensor, got {input.ndim}D"
         shape = input.shape
         out_dtype = torch.float8_e4m3fnuz if _is_hip else torch.float8_e4m3fn
+        if num_token_padding:
+            shape = (max(num_token_padding, input.shape[0]), shape[1])
         output = torch.empty(shape, device=input.device, dtype=out_dtype)
 
         if scale is None:

sglang/srt/distributed/device_communicators/custom_all_reduce.py

@@ -5,7 +5,7 @@ import logging
 import os
 from contextlib import contextmanager
 from functools import wraps
-from typing import Callable, List, Optional, TypeVar, Union
+from typing import Any, Callable, List, Optional, TypeVar, Union
 
 import torch
 import torch.distributed as dist
@@ -18,7 +18,7 @@ from sglang.srt.distributed.device_communicators.custom_all_reduce_utils import
     gpu_p2p_access_check,
 )
 from sglang.srt.distributed.parallel_state import in_the_same_node_as
-from sglang.srt.utils import cuda_device_count_stateless, is_cuda, is_hip
+from sglang.srt.utils import is_cuda, is_hip
 
 logger = logging.getLogger(__name__)
 
@@ -217,7 +217,7 @@ class CustomAllreduce:
         if cuda_visible_devices:
             device_ids = list(map(int, cuda_visible_devices.split(",")))
         else:
-            device_ids = list(range(cuda_device_count_stateless()))
+            device_ids = list(range(torch.cuda.device_count()))
 
         physical_device_id = device_ids[device.index]
         tensor = torch.tensor([physical_device_id], dtype=torch.int, device="cpu")
@@ -257,7 +257,7 @@ class CustomAllreduce:
         self.world_size = world_size
         self.full_nvlink = full_nvlink
 
-        if ops.use_vllm_custom_allreduce and not _is_hip:
+        if not _is_hip:
             # Buffers memory are owned by this Python class and passed to C++.
             # Meta data composes of two parts: meta data for synchronization and a
             # temporary buffer for storing intermediate allreduce results.
@@ -280,56 +280,24 @@ class CustomAllreduce:
             )
             ops.register_buffer(self._ptr, self.buffer_ptrs)
         else:
-            if _is_hip:
-                # meta data buffers need to be "uncached" for signal on MI200
-                self.meta = ops.allocate_meta_buffer(ops.meta_size() + max_size)
-                self.buffer = torch.empty(
-                    max_size, dtype=torch.uint8, device=self.device
-                )
-                handle = ops.get_meta_buffer_ipc_handle(self.meta)
-                shard_data = (
-                    bytes(handle),  # ipc handle to base ptr
-                    0,  # offset of base ptr
-                )
-                handles, offsets = self._gather_ipc_meta(shard_data)
-                self.rank_data = torch.empty(
-                    8 * 1024 * 1024, dtype=torch.uint8, device=self.device
-                )
-                self._ptr = ops.init_custom_ar(
-                    self.meta, self.rank_data, handles, offsets, rank, self.full_nvlink
-                )
-                self.register_buffer(self.buffer)
-                self.MSCCL = os.getenv("RCCL_MSCCL_ENABLE", "1") == "1"
-            else:
-                # From TensorRT-LLM getMaxRequiredWorkspaceSize
-                self.max_required_workspace_size = [16 * 1024 * 1024, 8 * 1024 * 1024]
-
-                # sizeof(uint32_t) * (MAX_ALL_REDUCE_BLOCKS + 2) * MAX_RANKS_PER_NODE;
-                self.barrier_max_size = 8 * (36 + 2) * 8
-
-                self.buffer_ptrs = self.create_shared_buffer(max_size, group=group)
-                self.tmp_result_buffer_ptrs = self.create_shared_buffer(
-                    max_size, group=group
-                )
-                self.rank_data_base = torch.empty(
-                    8 * 1024 * 1024, dtype=torch.uint8, device=self.device
-                )
-                self.barrier_in_ptrs = self.create_shared_buffer(
-                    self.barrier_max_size, group=group
-                )
-                self.barrier_out_ptrs = self.create_shared_buffer(
-                    self.barrier_max_size, group=group
-                )
+            # meta data buffers need to be "uncached" for signal on MI200
+            self.meta = ops.allocate_meta_buffer(ops.meta_size() + max_size)
+            self.buffer = torch.empty(max_size, dtype=torch.uint8, device=self.device)
+            handle = ops.get_meta_buffer_ipc_handle(self.meta)
+            shard_data = (
+                bytes(handle),  # ipc handle to base ptr
+                0,  # offset of base ptr
+            )
+            handles, offsets = self._gather_ipc_meta(shard_data)
+            self.rank_data = torch.empty(
+                8 * 1024 * 1024, dtype=torch.uint8, device=self.device
+            )
+            self._ptr = ops.init_custom_ar(
+                self.meta, self.rank_data, handles, offsets, rank, self.full_nvlink
+            )
+            self.register_buffer(self.buffer)
+            self.MSCCL = os.getenv("RCCL_MSCCL_ENABLE", "1") == "1"
 
-                self._ptr = ops.init_custom_ar(
-                    rank,
-                    world_size,
-                    self.rank_data_base,
-                    self.buffer_ptrs,
-                    self.tmp_result_buffer_ptrs,
-                    self.barrier_in_ptrs,
-                    self.barrier_out_ptrs,
-                )
         self.disabled = False
 
     @staticmethod
@@ -455,7 +423,7 @@ class CustomAllreduce:
             return False
         # for 4 or more non NVLink-capable GPUs, custom allreduce provides
         # little performance improvement over NCCL.
-        if ops.use_vllm_custom_allreduce and not _is_hip:
+        if not _is_hip:
             if self.world_size == 2 or self.full_nvlink:
                 return inp_size < self.max_size
             return False
@@ -471,18 +439,6 @@ class CustomAllreduce:
                     return inp_size < self.max_size
             return False
 
-        if self.world_size == 2:
-            return (
-                inp_size < self.max_size
-                and inp_size < self.max_required_workspace_size[0]
-            )
-
-        if self.full_nvlink:
-            return (
-                inp_size < self.max_size
-                and inp_size < self.max_required_workspace_size[1]
-            )
-
         return False
 
     # all reduce, assuming inp tensor is IPC registered with register_buffer,
@@ -515,15 +471,12 @@ class CustomAllreduce:
         """
         if out is None:
             out = torch.empty_like(inp)
-        if ops.use_vllm_custom_allreduce:
-            if registered:
-                ops.all_reduce(self._ptr, inp, out, 0, 0)
-            else:
-                ops.all_reduce(
-                    self._ptr, inp, out, self.buffer_ptrs[self.rank], self.max_size
-                )
+        if registered:
+            ops.all_reduce(self._ptr, inp, out, 0, 0)
         else:
-            ops.all_reduce(self._ptr, inp, out)
+            ops.all_reduce(
+                self._ptr, inp, out, self.buffer_ptrs[self.rank], self.max_size
+            )
         return out
 
     def custom_all_reduce(self, input: torch.Tensor) -> Optional[torch.Tensor]:
@@ -554,14 +507,9 @@ class CustomAllreduce:
     def close(self):
         if not self.disabled and self._ptr:
             ops.dispose(self._ptr)
-            if ops.use_vllm_custom_allreduce:
+            if _is_cuda:
                 self.free_shared_buffer(self.meta_ptrs)
                 self.free_shared_buffer(self.buffer_ptrs)
-            elif _is_cuda:
-                self.free_shared_buffer(self.buffer_ptrs)
-                self.free_shared_buffer(self.tmp_result_buffer_ptrs)
-                self.free_shared_buffer(self.barrier_in_ptrs)
-                self.free_shared_buffer(self.barrier_out_ptrs)
             self._ptr = 0
 
     def __del__(self):