minor

Signed-off-by: Woosuk Kwon <woosuk@thinkingmachines.ai>

vllm/v1/worker/gpu/sampler.py

@@ -10,8 +10,8 @@ import triton.language as tl
 
 from vllm.config import LogprobsMode
 from vllm.v1.outputs import LogprobsTensors, SamplerOutput
-from vllm.v1.sample.metadata import SamplingMetadata
 from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p
+from vllm.v1.worker.gpu.states import SamplingMetadata
 
 _SAMPLING_EPS = 1e-5
 

vllm/v1/worker/gpu/states.py

@@ -1,13 +1,24 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-from typing import Union
+from dataclasses import dataclass
+from typing import Optional, Union
 
 import numpy as np
 import torch
 
 from vllm.sampling_params import SamplingParams
-from vllm.v1.sample.metadata import SamplingMetadata
+
+
+@dataclass
+class SamplingMetadata:
+
+    temperature: torch.Tensor
+
+    top_p: Optional[torch.Tensor]
+    top_k: Optional[torch.Tensor]
+
+    # None means no logprobs, 0 means sampled token logprobs only
+    max_num_logprobs: Optional[int]
 
 
 class RequestState: