fix pre-commit

Signed-off-by: Zhuohan Li <zhuohan123@gmail.com>

vllm/benchmarks/datasets.py

@@ -75,9 +75,10 @@ class SampleRequest:
     Represents a single inference request for benchmarking.
     """
 
-    prompt: str | list[str]
+    prompt: str | list[str] | None
     prompt_len: int
     expected_output_len: int
+    prompt_token_ids: list[int] | None = None
     multi_modal_data: MultiModalDataDict | dict | list[dict] | None = None
     lora_request: LoRARequest | None = None
     request_id: str | None = None
@@ -385,7 +386,7 @@ def gen_prompt_decode_to_target_len(
     max_retry: int = 10,
     add_special_tokens: bool = False,
     rng: np.random.Generator | None = None,
-) -> tuple[str, list[int]]:
+) -> tuple[str, list[int], int]:
     """
     Ensure decoded-then-encoded prompt length matches the target token length.
 
@@ -438,9 +439,10 @@ def gen_prompt_decode_to_target_len(
 # -----------------------------------------------------------------------------
 
 
-class RandomDataset(BenchmarkDataset):
+class RandomTokenIDDataset(BenchmarkDataset):
     """
-    Synthetic text-only dataset for serving/throughput benchmarks.
+    Synthetic token-id-only dataset for serving/throughput benchmarks.
+    No need to use a tokenizer with this dataset.
 
     Strategy:
     - Sample input/output token lengths per request from integer-uniform ranges
@@ -448,7 +450,6 @@ class RandomDataset(BenchmarkDataset):
     - Prepend a fixed random prefix of length prefix_len.
     - Generate the remaining tokens as a reproducible sequence:
       (offset + index + arange(input_len)) % vocab_size.
-    - Decode then re-encode/truncate to ensure prompt token counts match.
     - Uses numpy.default_rng seeded with random_seed for reproducible sampling.
     """
 
@@ -467,14 +468,155 @@ class RandomDataset(BenchmarkDataset):
 
     def sample(
         self,
-        tokenizer: PreTrainedTokenizerBase,
         num_requests: int,
         request_id_prefix: str = "",
-        no_oversample: bool = False,
         prefix_len: int = DEFAULT_PREFIX_LEN,
         range_ratio: float = DEFAULT_RANGE_RATIO,
         input_len: int = DEFAULT_INPUT_LEN,
         output_len: int = DEFAULT_OUTPUT_LEN,
+        vocab_size: int = 1,
+        **kwargs,
+    ) -> list[SampleRequest]:
+        # validate total input tokens (prefix + sampled) is at least 1.
+        min_sampled_input = math.floor(input_len * (1.0 - float(range_ratio)))
+        min_total_input = int(prefix_len) + min_sampled_input
+        if min_total_input < 1:
+            raise ValueError(
+                f"--random-input-len is too small: with --random-range-ratio "
+                f"{range_ratio}, the minimum possible total input tokens (prefix "
+                f"+ sampled) is {min_total_input}. Increase --random-input-len and/or "
+                "--random-prefix-len, or decrease --random-range-ratio so that "
+                "prefix_len + floor(random_input_len * (1 - range_ratio)) >= 1."
+            )
+
+        input_lens, output_lens, offsets = self.get_sampling_params(
+            num_requests, range_ratio, input_len, output_len, vocab_size
+        )
+
+        # Generate prefix once
+        prefix_token_ids = self.get_prefix(vocab_size, prefix_len)
+
+        requests = []
+        for i in range(num_requests):
+            prompt_token_ids, total_input_len = self.generate_token_id_sequence(
+                prefix_token_ids=prefix_token_ids,
+                prefix_len=prefix_len,
+                vocab_size=vocab_size,
+                input_len=int(input_lens[i]),
+                offset=int(offsets[i]),
+                index=i,
+            )
+            requests.append(
+                SampleRequest(
+                    prompt=None,
+                    prompt_token_ids=prompt_token_ids,
+                    prompt_len=total_input_len,
+                    expected_output_len=int(output_lens[i]),
+                    request_id=request_id_prefix + str(i),
+                )
+            )
+
+        return requests
+
+    def get_prefix(self, vocab_size: int, prefix_len: int) -> list[int]:
+        """
+        Get the prefix for the dataset.
+        """
+        return (
+            self._rng.integers(0, vocab_size, size=prefix_len).tolist()
+            if prefix_len > 0
+            else []
+        )
+
+    def get_sampling_params(
+        self,
+        num_requests: int,
+        range_ratio: float,
+        input_len: int,
+        output_len: int,
+        vocab_size: int,
+    ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+        """
+        Get the sampling parameters for the dataset.
+        """
+        # Enforce range_ratio < 1
+        if not (0.0 <= range_ratio < 1.0):
+            raise ValueError("range_ratio must be in [0, 1).")
+        # Bounds use floor for low and ceil for high
+        input_low = math.floor(input_len * (1 - range_ratio))
+        input_high = math.ceil(input_len * (1 + range_ratio))
+        output_low = math.floor(output_len * (1 - range_ratio))
+        output_high = math.ceil(output_len * (1 + range_ratio))
+        # Ensure the lower bound for output length is at least 1 to
+        # prevent sampling 0 tokens.
+        output_low = max(output_low, 1)
+        output_high = max(output_high, 1)
+
+        if input_low > input_high:
+            raise ValueError(
+                f"Invalid input sampling interval: low={input_low} > high={input_high}"
+            )
+        if output_low > output_high:
+            raise ValueError(
+                "Invalid output sampling interval: "
+                f"low={output_low} > high={output_high}"
+            )
+
+        logger.info(
+            "Sampling input_len from [%s, %s] and output_len from [%s, %s]",
+            input_low,
+            input_high,
+            output_low,
+            output_high,
+        )
+
+        input_lens = self._rng.integers(input_low, input_high + 1, size=num_requests)
+        output_lens = self._rng.integers(output_low, output_high + 1, size=num_requests)
+        offsets = self._rng.integers(0, vocab_size, size=num_requests)
+        return input_lens, output_lens, offsets
+
+    def generate_token_id_sequence(
+        self,
+        *,
+        prefix_token_ids: list[int],
+        prefix_len: int,
+        vocab_size: int,
+        input_len: int,
+        offset: int,
+        index: int,
+    ) -> tuple[list[int], int]:
+        """
+        Returns (token_sequence, total_input_len).
+        """
+        # Build the inner sequence by sampling sequentially from the vocab
+        inner_seq = ((offset + index + np.arange(input_len)) % vocab_size).tolist()
+        token_sequence = prefix_token_ids + inner_seq
+        total_input_len = prefix_len + int(input_len)
+        return token_sequence, total_input_len
+
+
+# -----------------------------------------------------------------------------
+# Random Dataset Implementation (Synthetic Data)
+# -----------------------------------------------------------------------------
+
+
+class RandomDataset(RandomTokenIDDataset):
+    """
+    Synthetic text-only dataset for serving/throughput benchmarks.
+    Additionally to RandomTokenIDDataset, we perform a decode then re-encode/truncate
+    to ensure prompt token counts match.
+    """
+
+    def sample(
+        self,
+        tokenizer: PreTrainedTokenizerBase,
+        num_requests: int,
+        request_id_prefix: str = "",
+        no_oversample: bool = False,
+        prefix_len: int = RandomTokenIDDataset.DEFAULT_PREFIX_LEN,
+        range_ratio: float = RandomTokenIDDataset.DEFAULT_RANGE_RATIO,
+        input_len: int = RandomTokenIDDataset.DEFAULT_INPUT_LEN,
+        output_len: int = RandomTokenIDDataset.DEFAULT_OUTPUT_LEN,
         batchsize: int = 1,
         **kwargs,
     ) -> list[SampleRequest]:
@@ -490,17 +632,17 @@ class RandomDataset(BenchmarkDataset):
                 "the minimum possible total input tokens (prefix + sampled) is "
                 f"{min_total_input}. Increase --random-input-len and/or "
                 "--random-prefix-len, or decrease --random-range-ratio so that "
-                "prefix_len + floor(max(0, random_input_len - num_special)) "
-                "* (1 - range_ratio) >= 1."
+                "prefix_len + floor(max(0, random_input_len - num_special) "
+                "* (1 - range_ratio)) >= 1."
             )
 
+        vocab_size = tokenizer.vocab_size
         input_lens, output_lens, offsets = self.get_sampling_params(
-            num_requests, range_ratio, input_len, output_len, tokenizer
+            num_requests, range_ratio, input_len, output_len, vocab_size
         )
 
         # Generate prefix once
-        prefix_token_ids = self.get_prefix(tokenizer, prefix_len)
-        vocab_size = tokenizer.vocab_size
+        prefix_token_ids = self.get_prefix(vocab_size, prefix_len)
 
         requests = []
         token_mismatch_total = 0
@@ -552,67 +694,6 @@ class RandomDataset(BenchmarkDataset):
 
         return requests
 
-    def get_prefix(
-        self, tokenizer: PreTrainedTokenizerBase, prefix_len: int
-    ) -> list[int]:
-        """
-        Get the prefix for the dataset.
-        """
-        return (
-            self._rng.integers(0, tokenizer.vocab_size, size=prefix_len).tolist()
-            if prefix_len > 0
-            else []
-        )
-
-    def get_sampling_params(
-        self,
-        num_requests: int,
-        range_ratio: float,
-        input_len: int,
-        output_len: int,
-        tokenizer: PreTrainedTokenizerBase,
-    ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
-        """
-        Get the sampling parameters for the dataset.
-        """
-        # Enforce range_ratio < 1
-        if not (0.0 <= range_ratio < 1.0):
-            raise ValueError("range_ratio must be in [0, 1).")
-        num_special_tokens = int(tokenizer.num_special_tokens_to_add())
-        real_input_len = max(0, int(input_len) - num_special_tokens)
-        # Bounds use floor for low and ceil for high
-        input_low = math.floor(real_input_len * (1 - range_ratio))
-        input_high = math.ceil(real_input_len * (1 + range_ratio))
-        output_low = math.floor(output_len * (1 - range_ratio))
-        output_high = math.ceil(output_len * (1 + range_ratio))
-        # Ensure the lower bound for output length is at least 1 to
-        # prevent sampling 0 tokens.
-        output_low = max(output_low, 1)
-        output_high = max(output_high, 1)
-
-        if input_low > input_high:
-            raise ValueError(
-                f"Invalid input sampling interval: low={input_low} > high={input_high}"
-            )
-        if output_low > output_high:
-            raise ValueError(
-                "Invalid output sampling interval: "
-                f"low={output_low} > high={output_high}"
-            )
-
-        logger.info(
-            "Sampling input_len from [%s, %s] and output_len from [%s, %s]",
-            input_low,
-            input_high,
-            output_low,
-            output_high,
-        )
-
-        input_lens = self._rng.integers(input_low, input_high + 1, size=num_requests)
-        output_lens = self._rng.integers(output_low, output_high + 1, size=num_requests)
-        offsets = self._rng.integers(0, tokenizer.vocab_size, size=num_requests)
-        return input_lens, output_lens, offsets
-
     def generate_token_sequence(
         self,
         *,
@@ -656,7 +737,7 @@ class RandomDataset(BenchmarkDataset):
 
 
 # -----------------------------------------------------------------------------
-# Random Dataset Implementation (Synthetic Data)
+# Random Dataset Implementation (Reranking)
 # -----------------------------------------------------------------------------
 
 
@@ -684,9 +765,10 @@ class RandomDatasetForReranking(RandomDataset):
         n_sep_tokens = int(is_reranker)
 
         query_len_param = (input_len // 2) - n_sep_tokens if is_reranker else input_len
+        vocab_size = tokenizer.vocab_size
 
         query_lens, _, query_offsets = self.get_sampling_params(
-            1, range_ratio, query_len_param, 0, tokenizer
+            1, range_ratio, query_len_param, 0, vocab_size
         )
 
         query_len = int(query_lens[0])
@@ -700,9 +782,8 @@ class RandomDatasetForReranking(RandomDataset):
             doc_len_param = input_len - query_len - n_sep_tokens
 
         doc_lens, _, doc_offsets = self.get_sampling_params(
-            num_requests, range_ratio, doc_len_param, 0, tokenizer
+            num_requests, range_ratio, doc_len_param, 0, vocab_size
         )
-        vocab_size = tokenizer.vocab_size
 
         query_prompt, query_input_len, token_mismatch_total = (
             self.generate_token_sequence(
@@ -1054,9 +1135,11 @@ class RandomMultiModalDataset(RandomDataset):
                 "Video sampling not implemented; set its probability to 0."
             )
 
+        vocab_size = tokenizer.vocab_size
+
         # Get the sampling parameters for the dataset
         input_lens, output_lens, offsets = self.get_sampling_params(
-            num_requests, range_ratio, input_len, output_len, tokenizer
+            num_requests, range_ratio, input_len, output_len, vocab_size
         )
 
         (
@@ -1072,8 +1155,8 @@ class RandomMultiModalDataset(RandomDataset):
         )
 
         # Generate prefix once
-        prefix_token_ids = self.get_prefix(tokenizer, prefix_len)
-        vocab_size = tokenizer.vocab_size
+        prefix_token_ids = self.get_prefix(vocab_size, prefix_len)
+
         # Add synthetic multimodal items to each request
         mm_requests = []
         token_mismatch_total = 0

vllm/benchmarks/throughput.py

@@ -23,6 +23,7 @@ from vllm.benchmarks.datasets import (
     InstructCoderDataset,
     PrefixRepetitionRandomDataset,
     RandomDataset,
+    RandomTokenIDDataset,
     SampleRequest,
     ShareGPTDataset,
     SonnetDataset,
@@ -340,6 +341,10 @@ def get_requests(args, tokenizer):
         "output_len": args.output_len,
     }
 
+    if args.dataset_name == "random_token_id":
+        sample_kwargs["range_ratio"] = args.random_range_ratio
+        sample_kwargs["prefix_len"] = args.prefix_len
+        dataset_cls = RandomTokenIDDataset
     if args.dataset_path is None or args.dataset_name == "random":
         sample_kwargs["range_ratio"] = args.random_range_ratio
         sample_kwargs["prefix_len"] = args.prefix_len
@@ -691,9 +696,14 @@ def main(args: argparse.Namespace):
         args.seed = 0
     random.seed(args.seed)
     # Sample the requests.
-    tokenizer = AutoTokenizer.from_pretrained(
-        args.tokenizer, trust_remote_code=args.trust_remote_code
+    tokenizer = (
+        AutoTokenizer.from_pretrained(
+            args.tokenizer, trust_remote_code=args.trust_remote_code
+        )
+        if args.skip_tokenizer_init
+        else None
     )
+
     requests = get_requests(args, tokenizer)
     is_multi_modal = any(request.multi_modal_data is not None for request in requests)
     request_outputs: list[RequestOutput] | None = None

vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py

@@ -991,14 +991,11 @@ class NixlConnectorWorker:
         # Enable different block lengths for different layers when MLA is used.
         self.block_len_per_layer = list[int]()
         self.slot_size_per_layer = list[int]()  # HD bytes in kv terms
-        self.device_id = self.tp_rank
         for layer_name, cache_or_caches in xfer_buffers.items():
             cache_list = cache_or_caches if split_k_and_v else [cache_or_caches]
 
             for cache in cache_list:
                 base_addr = cache.data_ptr()
-                if not self.use_host_buffer and current_platform.is_cuda_alike():
-                    self.device_id = cache.device.index
                 if base_addr in seen_base_addresses:
                     continue
 
@@ -1026,7 +1023,7 @@ class NixlConnectorWorker:
                         "All kv cache tensors must have the same size"
                     )
                 caches_data.append(
-                    (base_addr, curr_tensor_size_bytes, self.device_id, "")
+                    (base_addr, curr_tensor_size_bytes, self.tp_rank, "")
                 )
 
         logger.debug(
@@ -1073,7 +1070,7 @@ class NixlConnectorWorker:
                 block_offset = block_id * self.block_len_per_layer[i]
                 addr = base_addr + block_offset
                 # (addr, len, device id)
-                blocks_data.append((addr, kv_block_len, self.device_id))
+                blocks_data.append((addr, kv_block_len, self.tp_rank))
 
             if self._use_flashinfer:
                 # Separate and interleave K/V regions to maintain the same
@@ -1084,13 +1081,12 @@ class NixlConnectorWorker:
                     addr = base_addr + block_offset
                     # Register addresses for V cache (K registered first).
                     v_addr = addr + kv_block_len
-                    blocks_data.append((v_addr, kv_block_len, self.device_id))
+                    blocks_data.append((v_addr, kv_block_len, self.tp_rank))
         logger.debug(
-            "Created %s blocks for src engine %s and rank %s on device id %s",
+            "Created %s blocks for src engine %s and rank %s",
             len(blocks_data),
             self.engine_id,
             self.tp_rank,
-            self.device_id,
         )
 
         descs = self.nixl_wrapper.get_xfer_descs(blocks_data, self.nixl_memory_type)

vllm/v1/engine/utils.py

@@ -134,12 +134,9 @@ class CoreEngineProcManager:
         data_parallel = vllm_config.parallel_config.data_parallel_size > 1
         try:
             for proc, local_dp_rank in zip(self.processes, local_dp_ranks):
-                # Adjust device control in DP for non-CUDA platforms
-                # For CUDA platforms, setting same device id for different DP
-                # processes affects NCCL init performance.
                 with (
                     set_device_control_env_var(vllm_config, local_dp_rank)
-                    if (data_parallel and not current_platform.is_cuda_alike())
+                    if (data_parallel)
                     else contextlib.nullcontext()
                 ):
                     proc.start()

vllm/v1/worker/dp_utils.py

@@ -8,6 +8,7 @@ import torch.distributed as dist
 from vllm.config import ParallelConfig
 from vllm.distributed.parallel_state import get_dp_group
 from vllm.logger import init_logger
+from vllm.platforms import current_platform
 from vllm.v1.worker.ubatch_utils import (
     UBatchSlices,
     check_ubatch_thresholds,
@@ -19,8 +20,7 @@ logger = init_logger(__name__)
 
 
 def _get_device_and_group(parallel_config: ParallelConfig):
-    # Use the actual device assigned to the DP group, not just the device type
-    device = get_dp_group().device
+    device = current_platform.device_type
     group = get_dp_group().device_group
 
     # Transfering this tensor from GPU to CPU will introduce a GPU sync

vllm/v1/worker/gpu_worker.py

@@ -172,27 +172,6 @@ class Worker(WorkerBase):
         if self.device_config.device.type == "cuda":
             # This env var set by Ray causes exceptions with graph building.
             os.environ.pop("NCCL_ASYNC_ERROR_HANDLING", None)
-            if (
-                self.parallel_config.data_parallel_size > 1
-                and self.parallel_config.data_parallel_size_local > 0
-                and self.parallel_config.data_parallel_backend != "ray"
-            ):
-                # Use local DP rank if available, otherwise use global DP rank.
-                dp_local_rank = self.parallel_config.data_parallel_rank_local
-                if dp_local_rank is None:
-                    dp_local_rank = self.parallel_config.data_parallel_rank
-
-                tp_pp_world_size = (
-                    self.parallel_config.pipeline_parallel_size
-                    * self.parallel_config.tensor_parallel_size
-                )
-
-                # DP_LOCAL_RANK * TP_PP_WORLD_SIZE + TP_LOCAL_RANK
-                self.local_rank += dp_local_rank * tp_pp_world_size
-                assert self.local_rank <= torch.cuda.device_count(), (
-                    f"DP adjusted local rank {self.local_rank} is out of bounds. "
-                )
-
             self.device = torch.device(f"cuda:{self.local_rank}")
             current_platform.set_device(self.device)