moves types around

Signed-off-by: Sage Moore <sage@neuralmagic.com>

Author: SageMoore

vllm/compilation/ubatch_utils.py

@@ -0,0 +1,8 @@
+from dataclasses import dataclass
+from typing import TypeAlias
+
+@dataclass
+class UbatchSlice:
+    request_slice: slice
+    token_slice: slice
+UBatchSlices: TypeAlias = list[UbatchSlice]

vllm/compilation/ubatch_wrapper.py

@@ -19,7 +19,6 @@
 
 logger = init_logger(__name__)
 
-
 @dataclasses.dataclass
 class UbatchMetadata:
     context: UBatchContext

vllm/forward_context.py

@@ -13,13 +13,7 @@
 import vllm.envs as envs
 from vllm.config import CUDAGraphMode, ParallelConfig, VllmConfig
 from vllm.logger import init_logger
-from typing import TypeAlias
-
-@dataclass
-class UbatchSlice:
-    request_slice: slice
-    token_slice: slice
-UBatchSlices: TypeAlias = list[UbatchSlice]
+from vllm.compilation.ubatch_utils import (UbatchSlice, UBatchSlices)
 
 if TYPE_CHECKING:
     from vllm.attention.backends.abstract import AttentionMetadata

vllm/v1/attention/backends/utils.py

@@ -29,6 +29,8 @@
 from vllm.logger import init_logger
 from vllm.v1.kv_cache_interface import AttentionSpec
 
+from vllm.compilation.ubatch_utils import UbatchSlice
+
 logger = init_logger(__name__)
 _KV_CACHE_LAYOUT_OVERRIDE = None
 
@@ -73,12 +75,6 @@ class CommonAttentionMetadata:
     num_logits_indices: Optional[int] = None
 
 
-@dataclass
-class UbatchSlice:
-    request_slice: slice
-    token_slice: slice
-
-
 def slice_query_start_locs(
     query_start_loc: torch.Tensor,
     request_slice: slice,

vllm/v1/worker/gpu_model_runner.py

@@ -24,6 +24,7 @@
 from vllm.compilation.cuda_graph import CUDAGraphWrapper
 from vllm.compilation.monitor import set_cudagraph_capturing_enabled
 from vllm.compilation.ubatch_wrapper import UBatchWrapper
+from vllm.compilation.ubatch_utils import (UbatchSlice, UBatchSlices)
 from vllm.config import (CompilationLevel, CUDAGraphMode, VllmConfig,
                          get_layers_from_vllm_config, update_config)
 from vllm.distributed.eplb.eplb_state import EplbState
@@ -109,8 +110,6 @@
 PerLayerAttnMetadata: TypeAlias = Union[list[AttnMetadataDict],
                                         AttnMetadataDict]
 
-UBatchSlices: TypeAlias = list[UbatchSlice]
-
 
 class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):