fix ut

Author: jikunshang

benchmarks/benchmark_throughput.py

@@ -211,8 +211,7 @@ def main(args: argparse.Namespace):
                                 args.seed, args.n, args.use_beam_search,
                                 args.trust_remote_code, args.dtype,
                                 args.max_model_len, args.enforce_eager,
-                                args.kv_cache_dtype,
-                                args.device)
+                                args.kv_cache_dtype, args.device)
     elif args.backend == "hf":
         assert args.tensor_parallel_size == 1
         elapsed_time = run_hf(requests, args.model, tokenizer, args.n,

benchmarks/kernels/benchmark_paged_attention.py

@@ -64,9 +64,14 @@ def main(
     block_tables = torch.tensor(block_tables, dtype=torch.int, device=device)
 
     # Create the KV cache.
-    key_caches, value_caches = create_kv_caches_with_random(
-        NUM_BLOCKS, block_size, 1, num_kv_heads, head_size, kv_cache_dtype,
-        dtype, device=device)
+    key_caches, value_caches = create_kv_caches_with_random(NUM_BLOCKS,
+                                                            block_size,
+                                                            1,
+                                                            num_kv_heads,
+                                                            head_size,
+                                                            kv_cache_dtype,
+                                                            dtype,
+                                                            device=device)
     key_cache, value_cache = key_caches[0], value_caches[0]
 
     # Prepare for the paged attention kernel.

tests/kernels/test_cache.py

@@ -15,6 +15,7 @@
 NUM_MAPPINGS = [256]  # Arbitrary values for testing
 SEEDS = [0]
 CUDA_DEVICES = [i for i in range(1 if torch.cuda.device_count() == 1 else 2)]
+KV_CACHE_DTYPE = ["auto", "fp8_e5m2"]
 
 
 @pytest.mark.parametrize("num_mappings", NUM_MAPPINGS)

tests/lora/test_worker.py

@@ -5,7 +5,7 @@
 
 from vllm.lora.models import LoRAMapping
 from vllm.lora.request import LoRARequest
-from vllm.config import ModelConfig, ParallelConfig, SchedulerConfig, LoRAConfig
+from vllm.config import ModelConfig, ParallelConfig, SchedulerConfig, LoRAConfig, DeviceConfig
 from vllm.worker.worker import Worker
 
 
@@ -25,6 +25,7 @@ def test_worker_apply_lora(sql_lora_files):
         ),
         parallel_config=ParallelConfig(1, 1, False),
         scheduler_config=SchedulerConfig(32, 32, 32, 256),
+        device_config=DeviceConfig("cuda"),
         local_rank=0,
         rank=0,
         lora_config=LoRAConfig(max_lora_rank=8, max_cpu_loras=32,

vllm/engine/llm_engine.py

@@ -88,6 +88,7 @@ def __init__(
             f"quantization={model_config.quantization}, "
             f"enforce_eager={model_config.enforce_eager}, "
             f"kv_cache_dtype={cache_config.cache_dtype}, "
+            f"device_config={device_config.device}, "
             f"seed={model_config.seed})")
         # TODO(woosuk): Print more configs in debug mode.
 

vllm/worker/model_runner.py

@@ -314,7 +314,9 @@ def _prepare_decode(
                                                   max_len=1,
                                                   pad=_PAD_SLOT_ID,
                                                   dtype=torch.long)
-        context_lens = torch.tensor(context_lens, dtype=torch.int)
+        context_lens = torch.tensor(context_lens,
+                                    dtype=torch.int,
+                                    device=self.device_config.device)
 
         if use_captured_graph:
             # The shape of graph_block_tables is