fix

Signed-off-by: Woosuk Kwon <[email protected]>

Author: WoosukKwon

vllm/v1/worker/gpu/cudagraph_utils.py

@@ -78,7 +78,7 @@ def capture_graph(
         kv_cache_config: KVCacheConfig,
     ) -> None:
         num_reqs = min(num_tokens, self.max_num_reqs)
-        input_ids = input_buffers.input_ids.gpu[:num_tokens]
+        input_ids = input_buffers.input_ids[:num_tokens]
         positions = input_buffers.positions[:num_tokens]
         attn_metadata = prepare_inputs_to_capture(
             num_reqs,

vllm/v1/worker/gpu/input_batch.py

@@ -29,7 +29,7 @@ def __init__(
         self.pin_memory = pin_memory
 
         self.idx_mapping = self._make_buffer(max_num_reqs, dtype=torch.int32)
-        self.input_ids = self._make_buffer(max_num_tokens, dtype=torch.int32)
+        self.input_ids = torch.zeros(max_num_tokens, dtype=torch.int32, device=device)
         self.positions = torch.zeros(max_num_tokens, dtype=torch.int64, device=device)
         self.query_start_loc = self._make_buffer(max_num_reqs + 1, dtype=torch.int32)
         self.seq_lens = torch.zeros(max_num_reqs, dtype=torch.int32, device=device)
@@ -116,7 +116,7 @@ def make_dummy(
         input_buffers.seq_lens[num_reqs:] = 0
         seq_lens = input_buffers.seq_lens[:num_reqs]
 
-        input_ids = input_buffers.input_ids.copy_to_gpu(num_tokens)
+        input_ids = input_buffers.input_ids[:num_tokens]
         positions = input_buffers.positions[:num_tokens]
         # attn_metadata = defaultdict(lambda: None)
         logits_indices = query_start_loc[1:] - 1

vllm/v1/worker/gpu/model_runner.py

@@ -410,9 +410,6 @@ def update_states(self, scheduler_output: SchedulerOutput) -> None:
                 cu_num_new_blocks[i].append(x + len(block_ids))
                 new_block_ids[i].extend(block_ids)
             overwrite.append(True)
-        # Update the GPU tensors for request states.
-        if scheduler_output.scheduled_new_reqs:
-            self.req_states.prefill_len.copy_to_gpu()
 
         # Add new blocks for the existing requests.
         cached_reqs = scheduler_output.scheduled_cached_reqs
@@ -507,7 +504,7 @@ def prepare_inputs(
 
         # Get prefill tokens.
         prepare_prefill_inputs(
-            self.input_buffers.input_ids.gpu,
+            self.input_buffers.input_ids,
             self.req_states.next_prefill_tokens,
             idx_mapping,
             query_start_loc_gpu,
@@ -529,7 +526,7 @@ def prepare_inputs(
         # Some input token ids are directly read from the last sampled tokens
         # and draft tokens. Also, get the logits indices to sample tokens from.
         logits_indices = combine_sampled_and_draft_tokens(
-            self.input_buffers.input_ids.gpu,
+            self.input_buffers.input_ids,
             idx_mapping,
             self.req_states.last_sampled_tokens,
             query_start_loc_gpu,
@@ -570,7 +567,7 @@ def prepare_inputs(
             kv_cache_config=self.kv_cache_config,
         )
 
-        input_ids = self.input_buffers.input_ids.gpu[:num_tokens_after_padding]
+        input_ids = self.input_buffers.input_ids[:num_tokens_after_padding]
         positions = self.input_buffers.positions[:num_tokens_after_padding]
         return InputBatch(
             req_ids=req_ids,

vllm/v1/worker/gpu/spec_decode/eagle.py

@@ -121,7 +121,7 @@ def run_model(
             num_tokens_across_dp=num_tokens_across_dp,
         ):
             ret_hidden_states = self.model(
-                input_ids=self.input_buffers.input_ids.gpu[:num_tokens],
+                input_ids=self.input_buffers.input_ids[:num_tokens],
                 positions=self.input_buffers.positions[:num_tokens],
                 hidden_states=self.hidden_states[:num_tokens],
             )
@@ -139,7 +139,7 @@ def generate_draft(
         num_tokens_across_dp: torch.Tensor | None,
     ) -> None:
         pos = self.input_buffers.positions[:num_reqs]
-        query_start_loc = self.input_buffers.query_start_loc.gpu[: num_reqs + 1]
+        query_start_loc = self.input_buffers.query_start_loc[: num_reqs + 1]
         for step in range(1, self.num_speculative_steps):
             # Run the eagle model.
             last_hidden_states, hidden_states = self.run_model(
@@ -379,7 +379,7 @@ def prepare_eagle_inputs(
     )
     _prepare_eagle_inputs_kernel[(num_reqs,)](
         last_token_indices,
-        input_buffers.input_ids.gpu,
+        input_buffers.input_ids,
         input_buffers.positions,
         input_batch.input_ids,
         input_batch.positions,
@@ -482,7 +482,7 @@ def prepare_eagle_decode(
         last_token_indices,
         target_seq_lens,
         num_rejected,
-        input_buffers.input_ids.gpu,
+        input_buffers.input_ids,
         input_buffers.positions,
         input_hidden_states,
         input_hidden_states.stride(0),
@@ -550,7 +550,7 @@ def update_eagle_inputs(
 ):
     num_reqs, hidden_size = output_hidden_states.shape
     _update_eagle_inputs_kernel[(num_reqs,)](
-        input_buffers.input_ids.gpu,
+        input_buffers.input_ids,
         input_buffers.positions,
         hidden_states,
         hidden_states.stride(0),

vllm/v1/worker/gpu/states.py

@@ -117,8 +117,7 @@ def __init__(
         self.prefill_token_ids = UvaBuffer(
             self.max_num_reqs, self.max_model_len, dtype=torch.int32
         )
-        self.prefill_len = self._make_buffer(self.max_num_reqs, dtype=torch.int32)
-
+        self.prefill_len = UvaBuffer(self.max_num_reqs, dtype=torch.int32)
         # Number of computed tokens.
         self.num_computed_prefill_tokens = np.zeros(self.max_num_reqs, dtype=np.int32)
         self.num_computed_tokens = torch.zeros(