zero out attn output

fixes and zero out input to fi kernel

.buildkite/scripts/hardware_ci/run-amd-test.sh

@@ -333,15 +333,15 @@ apply_rocm_test_overrides() {
   # --- Entrypoint ignores ---
   if [[ $cmds == *" entrypoints/openai "* ]]; then
     cmds=${cmds//" entrypoints/openai "/" entrypoints/openai \
-    --ignore=entrypoints/openai/test_audio.py \
-    --ignore=entrypoints/openai/test_shutdown.py \
+    --ignore=entrypoints/openai/chat_completion/test_audio.py \
+    --ignore=entrypoints/openai/completion/test_shutdown.py \
     --ignore=entrypoints/openai/test_completion.py \
     --ignore=entrypoints/openai/test_models.py \
     --ignore=entrypoints/openai/test_lora_adapters.py \
     --ignore=entrypoints/openai/test_return_tokens_as_ids.py \
-    --ignore=entrypoints/openai/test_root_path.py \
+    --ignore=entrypoints/openai/chat_completion/test_root_path.py \
     --ignore=entrypoints/openai/test_tokenization.py \
-    --ignore=entrypoints/openai/test_prompt_validation.py "}
+    --ignore=entrypoints/openai/completion/test_prompt_validation.py "}
   fi
 
   if [[ $cmds == *" entrypoints/llm "* ]]; then

.buildkite/test-amd.yaml

@@ -162,7 +162,7 @@ steps:
   - tests/entrypoints/test_chat_utils
   commands:
   - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s entrypoints/openai --ignore=entrypoints/openai/test_chat_with_tool_reasoning.py --ignore=entrypoints/openai/test_oot_registration.py --ignore=entrypoints/openai/test_tensorizer_entrypoint.py --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses
+  - pytest -v -s entrypoints/openai --ignore=entrypoints/openai/chat_completion/test_chat_with_tool_reasoning.py --ignore=entrypoints/openai/chat_completion/test_oot_registration.py --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses
   - pytest -v -s entrypoints/test_chat_utils.py
 
 - label: Entrypoints Integration Test (API Server 2)
@@ -674,12 +674,12 @@ steps:
   - vllm/config/model.py
   - vllm/model_executor
   - tests/model_executor
-  - tests/entrypoints/openai/test_tensorizer_entrypoint.py
+  - tests/entrypoints/openai/completion/test_tensorizer_entrypoint.py
   commands:
     - apt-get update && apt-get install -y curl libsodium23
     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
     - pytest -v -s model_executor
-    - pytest -v -s entrypoints/openai/test_tensorizer_entrypoint.py
+    - pytest -v -s entrypoints/openai/completion/test_tensorizer_entrypoint.py
 
 - label: Benchmarks # 11min
   timeout_in_minutes: 20
@@ -1143,7 +1143,7 @@ steps:
   - pytest -v -s plugins_tests/test_scheduler_plugins.py
   - pip install -e ./plugins/vllm_add_dummy_model
   - pytest -v -s distributed/test_distributed_oot.py
-  - pytest -v -s entrypoints/openai/test_oot_registration.py
+  - pytest -v -s entrypoints/openai/chat_completion/test_oot_registration.py
   - pytest -v -s models/test_oot_registration.py
   - pytest -v -s plugins/lora_resolvers
 
@@ -1502,7 +1502,7 @@ steps:
   - tests/entrypoints/test_chat_utils
   commands:
   - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s entrypoints/openai --ignore=entrypoints/openai/test_chat_with_tool_reasoning.py --ignore=entrypoints/openai/test_oot_registration.py --ignore=entrypoints/openai/test_tensorizer_entrypoint.py --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses
+  - pytest -v -s entrypoints/openai --ignore=entrypoints/openai/chat_completion/test_chat_with_tool_reasoning.py --ignore=entrypoints/openai/chat_completion/test_oot_registration.py --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses
   - pytest -v -s entrypoints/test_chat_utils.py
 
 - label: Entrypoints Integration Test (API Server 2)
@@ -2133,12 +2133,12 @@ steps:
   - vllm/config/model.py
   - vllm/model_executor
   - tests/model_executor
-  - tests/entrypoints/openai/test_tensorizer_entrypoint.py
+  - tests/entrypoints/openai/completion/test_tensorizer_entrypoint.py
   commands:
     - apt-get update && apt-get install -y curl libsodium23
     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
     - pytest -v -s model_executor
-    - pytest -v -s entrypoints/openai/test_tensorizer_entrypoint.py
+    - pytest -v -s entrypoints/openai/completion/test_tensorizer_entrypoint.py
 
 - label: Benchmarks # 11min
   timeout_in_minutes: 20
@@ -2735,7 +2735,7 @@ steps:
   - pytest -v -s plugins_tests/test_scheduler_plugins.py
   - pip install -e ./plugins/vllm_add_dummy_model
   - pytest -v -s distributed/test_distributed_oot.py
-  - pytest -v -s entrypoints/openai/test_oot_registration.py # it needs a clean process
+  - pytest -v -s entrypoints/openai/chat_completion/test_oot_registration.py # it needs a clean process
   - pytest -v -s models/test_oot_registration.py # it needs a clean process
   - pytest -v -s plugins/lora_resolvers # unit tests for in-tree lora resolver plugins
 
@@ -3257,7 +3257,7 @@ steps:
   - tests/entrypoints/test_chat_utils
   commands:
   - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s entrypoints/openai --ignore=entrypoints/openai/test_chat_with_tool_reasoning.py --ignore=entrypoints/openai/test_oot_registration.py --ignore=entrypoints/openai/test_tensorizer_entrypoint.py --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses
+  - pytest -v -s entrypoints/openai --ignore=entrypoints/openai/chat_completion/test_chat_with_tool_reasoning.py --ignore=entrypoints/openai/chat_completion/test_oot_registration.py --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses
   - pytest -v -s entrypoints/test_chat_utils.py
 
 - label: Entrypoints Integration Test (API Server 2)
@@ -3872,12 +3872,12 @@ steps:
   - vllm/config/model.py
   - vllm/model_executor
   - tests/model_executor
-  - tests/entrypoints/openai/test_tensorizer_entrypoint.py
+  - tests/entrypoints/openai/completion/test_tensorizer_entrypoint.py
   commands:
     - apt-get update && apt-get install -y curl libsodium23
     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
     - pytest -v -s model_executor
-    - pytest -v -s entrypoints/openai/test_tensorizer_entrypoint.py
+    - pytest -v -s entrypoints/openai/completion/test_tensorizer_entrypoint.py
 
 - label: Benchmarks # 11min
   timeout_in_minutes: 20
@@ -4508,7 +4508,7 @@ steps:
   - pytest -v -s plugins_tests/test_scheduler_plugins.py
   - pip install -e ./plugins/vllm_add_dummy_model
   - pytest -v -s distributed/test_distributed_oot.py
-  - pytest -v -s entrypoints/openai/test_oot_registration.py # it needs a clean process
+  - pytest -v -s entrypoints/openai/chat_completion/test_oot_registration.py # it needs a clean process
   - pytest -v -s models/test_oot_registration.py # it needs a clean process
   - pytest -v -s plugins/lora_resolvers # unit tests for in-tree lora resolver plugins
 

.buildkite/test_areas/distributed.yaml

@@ -15,8 +15,29 @@ steps:
   - pytest -v -s distributed/test_shm_buffer.py
   - pytest -v -s distributed/test_shm_storage.py
 
-- label: Distributed (2 GPUs)
-  timeout_in_minutes: 60
+- label: Distributed DP Tests (2 GPUs)
+  timeout_in_minutes: 20
+  working_dir: "/vllm-workspace/tests"
+  num_devices: 2
+  source_file_dependencies:
+  - vllm/distributed/
+  - vllm/engine/
+  - vllm/executor/
+  - vllm/worker/worker_base.py
+  - vllm/v1/engine/
+  - vllm/v1/worker/
+  - tests/v1/distributed
+  - tests/v1/entrypoints/openai/test_multi_api_servers.py
+  commands:
+  # https://github.com/NVIDIA/nccl/issues/1838
+  - export NCCL_CUMEM_HOST_ENABLE=0
+  - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_async_llm_dp.py
+  - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_eagle_dp.py
+  - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_external_lb_dp.py
+  - DP_SIZE=2 pytest -v -s v1/entrypoints/openai/test_multi_api_servers.py
+
+- label: Distributed Compile + RPC Tests (2 GPUs)
+  timeout_in_minutes: 20
   working_dir: "/vllm-workspace/tests"
   num_devices: 2
   source_file_dependencies:
@@ -29,22 +50,31 @@ steps:
   - vllm/v1/worker/
   - tests/compile/fullgraph/test_basic_correctness.py
   - tests/compile/test_wrapper.py
-  - tests/distributed/
   - tests/entrypoints/llm/test_collective_rpc.py
-  - tests/v1/distributed
-  - tests/v1/entrypoints/openai/test_multi_api_servers.py
+  commands:
+  # https://github.com/NVIDIA/nccl/issues/1838
+  - export NCCL_CUMEM_HOST_ENABLE=0
+  - pytest -v -s entrypoints/llm/test_collective_rpc.py
+  - pytest -v -s ./compile/fullgraph/test_basic_correctness.py
+  - pytest -v -s ./compile/test_wrapper.py
+
+- label: Distributed Torchrun + Shutdown Tests (2 GPUs)
+  timeout_in_minutes: 20
+  working_dir: "/vllm-workspace/tests"
+  num_devices: 2
+  source_file_dependencies:
+  - vllm/distributed/
+  - vllm/engine/
+  - vllm/executor/
+  - vllm/worker/worker_base.py
+  - vllm/v1/engine/
+  - vllm/v1/worker/
+  - tests/distributed/
   - tests/v1/shutdown
   - tests/v1/worker/test_worker_memory_snapshot.py
   commands:
   # https://github.com/NVIDIA/nccl/issues/1838
   - export NCCL_CUMEM_HOST_ENABLE=0
-  - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_async_llm_dp.py
-  - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_eagle_dp.py
-  - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_external_lb_dp.py
-  - DP_SIZE=2 pytest -v -s v1/entrypoints/openai/test_multi_api_servers.py
-  - pytest -v -s entrypoints/llm/test_collective_rpc.py
-  - pytest -v -s ./compile/fullgraph/test_basic_correctness.py
-  - pytest -v -s ./compile/test_wrapper.py
   - VLLM_TEST_SAME_HOST=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed'
   - VLLM_TEST_SAME_HOST=1 VLLM_TEST_WITH_DEFAULT_DEVICE_SET=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed'
   - CUDA_VISIBLE_DEVICES=0,1 pytest -v -s v1/shutdown

.buildkite/test_areas/entrypoints.yaml

@@ -34,7 +34,7 @@ steps:
   - tests/entrypoints/test_chat_utils
   commands:
   - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s entrypoints/openai --ignore=entrypoints/openai/chat_completion/test_chat_with_tool_reasoning.py --ignore=entrypoints/openai/test_oot_registration.py --ignore=entrypoints/openai/test_tensorizer_entrypoint.py --ignore=entrypoints/openai/correctness/  --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses
+  - pytest -v -s entrypoints/openai --ignore=entrypoints/openai/chat_completion/test_chat_with_tool_reasoning.py --ignore=entrypoints/openai/chat_completion/test_oot_registration.py --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py --ignore=entrypoints/openai/correctness/  --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses
   - pytest -v -s entrypoints/test_chat_utils.py
   mirror:
     amd:

.buildkite/test_areas/kernels.yaml

@@ -35,7 +35,7 @@ steps:
   parallelism: 2
 
 - label: Kernels MoE Test %N
-  timeout_in_minutes: 60
+  timeout_in_minutes: 25
   source_file_dependencies:
   - csrc/quantization/cutlass_w8a8/moe/
   - csrc/moe/
@@ -47,7 +47,7 @@ steps:
   commands:
     - pytest -v -s kernels/moe --ignore=kernels/moe/test_modular_oai_triton_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
     - pytest -v -s kernels/moe/test_modular_oai_triton_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
-  parallelism: 2
+  parallelism: 5
 
 - label: Kernels Mamba Test
   timeout_in_minutes: 45

.buildkite/test_areas/model_executor.yaml

@@ -9,9 +9,9 @@ steps:
   - vllm/config/model.py
   - vllm/model_executor
   - tests/model_executor
-  - tests/entrypoints/openai/test_tensorizer_entrypoint.py
+  - tests/entrypoints/openai/completion/test_tensorizer_entrypoint.py
   commands:
     - apt-get update && apt-get install -y curl libsodium23
     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
     - pytest -v -s model_executor
-    - pytest -v -s entrypoints/openai/test_tensorizer_entrypoint.py
+    - pytest -v -s entrypoints/openai/completion/test_tensorizer_entrypoint.py

.buildkite/test_areas/plugins.yaml

@@ -36,6 +36,6 @@ steps:
   - pytest -v -s plugins_tests/test_scheduler_plugins.py
   - pip install -e ./plugins/vllm_add_dummy_model
   - pytest -v -s distributed/test_distributed_oot.py
-  - pytest -v -s entrypoints/openai/test_oot_registration.py # it needs a clean process
+  - pytest -v -s entrypoints/openai/chat_completion/test_oot_registration.py # it needs a clean process
   - pytest -v -s models/test_oot_registration.py # it needs a clean process
   - pytest -v -s plugins/lora_resolvers # unit tests for in-tree lora resolver plugins

.buildkite/test_areas/pytorch.yaml

@@ -35,7 +35,7 @@ steps:
   # as it is a heavy test that is covered in other steps.
   # Use `find` to launch multiple instances of pytest so that
   # they do not suffer from https://github.com/vllm-project/vllm/issues/28965
-  - "find compile/fullgraph/ -name 'test_*.py' -not -name 'test_full_graph.py' -exec pytest -s -v {} \\;"
+  - "find compile/fullgraph/ -name 'test_*.py' -not -name 'test_full_graph.py' -print0 | xargs -0 -n1 -I{} pytest -s -v '{}'"
 
 - label: PyTorch Fullgraph
   timeout_in_minutes: 30

.github/mergify.yml

@@ -381,7 +381,7 @@ pull_request_rules:
     - or:
       - files~=^vllm/model_executor/model_loader/tensorizer.py
       - files~=^vllm/model_executor/model_loader/tensorizer_loader.py
-      - files~=^tests/entrypoints/openai/test_tensorizer_entrypoint.py
+      - files~=^tests/entrypoints/openai/completion/test_tensorizer_entrypoint.py
       - files~=^tests/model_executor/model_loader/tensorizer_loader/
   actions:
     assign:

benchmarks/attention_benchmarks/benchmark.py

@@ -47,6 +47,8 @@ from common import (
     is_mla_backend,
 )
 
+from vllm.v1.worker.workspace import init_workspace_manager
+
 
 def run_standard_attention_benchmark(config: BenchmarkConfig) -> BenchmarkResult:
     """Run standard attention benchmark (Flash/Triton/FlashInfer)."""
@@ -462,7 +464,7 @@ def main():
     parser.add_argument(
         "--batch-specs",
         nargs="+",
-        default=["q2k", "8q1s1k"],
+        default=None,
         help="Batch specifications using extended grammar",
     )
 
@@ -478,6 +480,21 @@ def main():
     parser.add_argument("--repeats", type=int, default=1, help="Repetitions")
     parser.add_argument("--warmup-iters", type=int, default=3, help="Warmup iterations")
     parser.add_argument("--profile-memory", action="store_true", help="Profile memory")
+    parser.add_argument(
+        "--kv-cache-dtype",
+        default="auto",
+        choices=["auto", "fp8"],
+        help="KV cache dtype: auto or fp8",
+    )
+    parser.add_argument(
+        "--cuda-graphs",
+        action=argparse.BooleanOptionalAction,
+        default=True,
+        help=(
+            "Launch kernels with CUDA graphs to eliminate CPU overhead"
+            "in measurements (default: True)"
+        ),
+    )
 
     # Parameter sweep (use YAML config for advanced sweeps)
     parser.add_argument(
@@ -536,21 +553,24 @@ def main():
 
         # Batch specs and sizes
         # Support both explicit batch_specs and generated batch_spec_ranges
-        if "batch_spec_ranges" in yaml_config:
-            # Generate batch specs from ranges
-            generated_specs = generate_batch_specs_from_ranges(
-                yaml_config["batch_spec_ranges"]
-            )
-            # Combine with any explicit batch_specs
-            if "batch_specs" in yaml_config:
-                args.batch_specs = yaml_config["batch_specs"] + generated_specs
-            else:
-                args.batch_specs = generated_specs
-            console.print(
-                f"[dim]Generated {len(generated_specs)} batch specs from ranges[/]"
-            )
-        elif "batch_specs" in yaml_config:
-            args.batch_specs = yaml_config["batch_specs"]
+        # CLI --batch-specs takes precedence over YAML when provided.
+        cli_batch_specs_provided = args.batch_specs is not None
+        if not cli_batch_specs_provided:
+            if "batch_spec_ranges" in yaml_config:
+                # Generate batch specs from ranges
+                generated_specs = generate_batch_specs_from_ranges(
+                    yaml_config["batch_spec_ranges"]
+                )
+                # Combine with any explicit batch_specs
+                if "batch_specs" in yaml_config:
+                    args.batch_specs = yaml_config["batch_specs"] + generated_specs
+                else:
+                    args.batch_specs = generated_specs
+                console.print(
+                    f"[dim]Generated {len(generated_specs)} batch specs from ranges[/]"
+                )
+            elif "batch_specs" in yaml_config:
+                args.batch_specs = yaml_config["batch_specs"]
 
         if "batch_sizes" in yaml_config:
             args.batch_sizes = yaml_config["batch_sizes"]
@@ -575,6 +595,10 @@ def main():
             args.warmup_iters = yaml_config["warmup_iters"]
         if "profile_memory" in yaml_config:
             args.profile_memory = yaml_config["profile_memory"]
+        if "kv_cache_dtype" in yaml_config:
+            args.kv_cache_dtype = yaml_config["kv_cache_dtype"]
+        if "cuda_graphs" in yaml_config:
+            args.cuda_graphs = yaml_config["cuda_graphs"]
 
         # Parameter sweep configuration
         if "parameter_sweep" in yaml_config:
@@ -629,12 +653,18 @@ def main():
     # Determine backends
     backends = args.backends or ([args.backend] if args.backend else ["flash"])
     prefill_backends = getattr(args, "prefill_backends", None)
+    if not args.batch_specs:
+        args.batch_specs = ["q2k", "8q1s1k"]
     console.print(f"Backends: {', '.join(backends)}")
     if prefill_backends:
         console.print(f"Prefill backends: {', '.join(prefill_backends)}")
     console.print(f"Batch specs: {', '.join(args.batch_specs)}")
+    console.print(f"KV cache dtype: {args.kv_cache_dtype}")
+    console.print(f"CUDA graphs: {args.cuda_graphs}")
     console.print()
 
+    init_workspace_manager(args.device)
+
     # Run benchmarks
     all_results = []
 
@@ -687,6 +717,8 @@ def main():
                         repeats=args.repeats,
                         warmup_iters=args.warmup_iters,
                         profile_memory=args.profile_memory,
+                        kv_cache_dtype=args.kv_cache_dtype,
+                        use_cuda_graphs=args.cuda_graphs,
                     )
 
                     # Add decode pipeline config
@@ -839,6 +871,8 @@ def main():
             "repeats": args.repeats,
             "warmup_iters": args.warmup_iters,
             "profile_memory": args.profile_memory,
+            "kv_cache_dtype": args.kv_cache_dtype,
+            "use_cuda_graphs": args.cuda_graphs,
         }
         all_results = run_model_parameter_sweep(
             backends,
@@ -861,6 +895,8 @@ def main():
             "repeats": args.repeats,
             "warmup_iters": args.warmup_iters,
             "profile_memory": args.profile_memory,
+            "kv_cache_dtype": args.kv_cache_dtype,
+            "use_cuda_graphs": args.cuda_graphs,
         }
         all_results = run_parameter_sweep(
             backends, args.batch_specs, base_config_args, args.parameter_sweep, console
@@ -891,6 +927,8 @@ def main():
                             repeats=args.repeats,
                             warmup_iters=args.warmup_iters,
                             profile_memory=args.profile_memory,
+                            kv_cache_dtype=args.kv_cache_dtype,
+                            use_cuda_graphs=args.cuda_graphs,
                         )
 
                         result = run_benchmark(config)

benchmarks/attention_benchmarks/common.py

@@ -213,6 +213,9 @@ class BenchmarkConfig:
     profile_memory: bool = False
     use_cuda_graphs: bool = False
 
+    # "auto" or "fp8"
+    kv_cache_dtype: str = "auto"
+
     # MLA-specific
     prefill_backend: str | None = None
     kv_lora_rank: int | None = None
@@ -369,6 +372,7 @@ class ResultsFormatter:
                     "backend",
                     "batch_spec",
                     "num_layers",
+                    "kv_cache_dtype",
                     "mean_time",
                     "std_time",
                     "throughput",
@@ -382,6 +386,7 @@ class ResultsFormatter:
                         "backend": r.config.backend,
                         "batch_spec": r.config.batch_spec,
                         "num_layers": r.config.num_layers,
+                        "kv_cache_dtype": r.config.kv_cache_dtype,
                         "mean_time": r.mean_time,
                         "std_time": r.std_time,
                         "throughput": r.throughput_tokens_per_sec or 0,

benchmarks/attention_benchmarks/configs/mla_mixed_batch.yaml

@@ -30,9 +30,9 @@ batch_specs:
   - "2q16k_32q1s4k"         # 2 very large prefill + 32 decode
 
   # Context extension + decode
-  - "2q1kkv2k_16q1s1k"       # 2 extend + 16 decode
-  - "4q2kkv4k_32q1s2k"       # 4 extend + 32 decode
-  - "2q1kkv8k_32q1s2k"       # 2 large extend + 32 decode
+  - "2q1ks2k_16q1s1k"       # 2 extend + 16 decode
+  - "4q2ks4k_32q1s2k"       # 4 extend + 32 decode
+  - "2q1ks8k_32q1s2k"       # 2 large extend + 32 decode
 
   # Explicitly chunked prefill
   - "q8k"           # 8k prefill with chunking hint

benchmarks/attention_benchmarks/configs/mla_sparse_decode.yaml

@@ -0,0 +1,58 @@
+# MLA decode-only benchmark configuration
+
+model:
+  name: "deepseek-v3"
+  num_layers: 60
+  num_q_heads: 128  # Base value, can be swept for TP simulation
+  num_kv_heads: 1  # MLA uses single latent KV
+  head_dim: 576
+  kv_lora_rank: 512
+  qk_nope_head_dim: 128
+  qk_rope_head_dim: 64
+  v_head_dim: 128
+  block_size: 128  # CUTLASS MLA and FlashAttn MLA use 128
+
+# Model parameter sweep: simulate tensor parallelism by varying num_q_heads
+# TP=1: 128 heads, TP=2: 64 heads, TP=4: 32 heads, TP=8: 16 heads
+model_parameter_sweep:
+  param_name: "num_q_heads"
+  values: [128, 64, 32, 16]
+  label_format: "{backend}_{value}h"
+
+batch_specs:
+  # Small batches, varying sequence lengths
+  - "16q1s512"     # 16 requests, 512 KV cache
+  - "16q1s1k"      # 16 requests, 1k KV cache
+  - "16q1s2k"      # 16 requests, 2k KV cache
+  - "16q1s4k"      # 16 requests, 4k KV cache
+
+  # Medium batches
+  - "32q1s1k"      # 32 requests, 1k KV cache
+  - "32q1s2k"      # 32 requests, 2k KV cache
+  - "32q1s4k"      # 32 requests, 4k KV cache
+  - "32q1s8k"      # 32 requests, 8k KV cache
+
+  # Large batches
+  - "64q1s1k"      # 64 requests, 1k KV cache
+  - "64q1s2k"      # 64 requests, 2k KV cache
+  - "64q1s4k"      # 64 requests, 4k KV cache
+  - "64q1s8k"      # 64 requests, 8k KV cache
+
+  # Very large batches
+  - "128q1s1k"     # 128 requests, 1k KV cache
+  - "128q1s2k"     # 128 requests, 2k KV cache
+  - "128q1s4k"     # 128 requests, 4k KV cache
+  - "128q1s8k"     # 128 requests, 8k KV cache
+
+  # Long context
+  - "32q1s16k"     # 32 requests, 16k KV cache
+  - "32q1s32k"     # 32 requests, 32k KV cache
+
+backends:
+  - FLASHMLA_SPARSE
+  - FLASHINFER_MLA_SPARSE
+
+device: "cuda:0"
+repeats: 100
+warmup_iters: 10
+profile_memory: true

benchmarks/attention_benchmarks/mla_runner.py

@@ -60,9 +60,11 @@ def create_minimal_vllm_config(
     model_name: str = "deepseek-v3",
     block_size: int = 128,
     max_num_seqs: int = 256,
+    max_num_batched_tokens: int = 8192,
     mla_dims: dict | None = None,
     index_topk: int | None = None,
     prefill_backend: str | None = None,
+    kv_cache_dtype: str = "auto",
 ) -> VllmConfig:
     """
     Create minimal VllmConfig for MLA benchmarks.
@@ -149,13 +151,13 @@ def create_minimal_vllm_config(
     cache_config = CacheConfig(
         block_size=block_size,
         gpu_memory_utilization=0.9,
-        cache_dtype="auto",
+        cache_dtype=kv_cache_dtype,
         enable_prefix_caching=False,
     )
 
     scheduler_config = SchedulerConfig(
         max_num_seqs=max_num_seqs,
-        max_num_batched_tokens=8192,
+        max_num_batched_tokens=max(max_num_batched_tokens, max_num_seqs),
         max_model_len=32768,
         is_encoder_decoder=False,
         enable_chunked_prefill=True,
@@ -535,6 +537,7 @@ def _create_backend_impl(
     device: torch.device,
     max_num_tokens: int = 8192,
     index_topk: int | None = None,
+    kv_cache_dtype: str = "auto",
 ):
     """
     Create backend implementation instance.
@@ -583,7 +586,7 @@ def _create_backend_impl(
         "num_kv_heads": mla_dims["num_kv_heads"],
         "alibi_slopes": None,
         "sliding_window": None,
-        "kv_cache_dtype": "auto",
+        "kv_cache_dtype": kv_cache_dtype,
         "logits_soft_cap": None,
         "attn_type": "decoder",
         "kv_sharing_target_layer_name": None,
@@ -701,6 +704,7 @@ def _run_single_benchmark(
     mla_dims: dict,
     device: torch.device,
     indexer=None,
+    kv_cache_dtype: str | None = None,
 ) -> BenchmarkResult:
     """
     Run a single benchmark iteration.
@@ -734,49 +738,124 @@ def _run_single_benchmark(
     )
 
     # Create KV cache
-    kv_cache = torch.zeros(
-        num_blocks,
-        block_size,
-        mla_dims["kv_lora_rank"] + mla_dims["qk_rope_head_dim"],
-        device=device,
-        dtype=torch.bfloat16,
-    )
+    if kv_cache_dtype is None:
+        kv_cache_dtype = getattr(config, "kv_cache_dtype", "auto")
+    head_size = mla_dims["kv_lora_rank"] + mla_dims["qk_rope_head_dim"]
+    if kv_cache_dtype == "fp8_ds_mla":
+        # FlashMLA sparse custom format: 656 bytes per token, stored as uint8.
+        # Layout: kv_lora_rank fp8 bytes + 4 float32 tile scales
+        #         + 2*rope_dim bf16 bytes
+        # = 512 + 16 + 128 = 656 bytes for DeepSeek dims.
+        kv_cache = torch.zeros(
+            num_blocks,
+            block_size,
+            656,
+            device=device,
+            dtype=torch.uint8,
+        )
+    elif kv_cache_dtype == "fp8":
+        from vllm.platforms import current_platform
 
-    # Create input tensors for both decode and prefill modes
-    decode_inputs, prefill_inputs = _create_input_tensors(
-        total_q,
-        mla_dims,
-        backend_cfg["query_format"],
-        device,
-        torch.bfloat16,
-    )
+        kv_cache = torch.zeros(
+            num_blocks,
+            block_size,
+            head_size,
+            device=device,
+            dtype=torch.uint8,
+        ).view(current_platform.fp8_dtype())
+    else:
+        kv_cache = torch.zeros(
+            num_blocks,
+            block_size,
+            head_size,
+            device=device,
+            dtype=torch.bfloat16,
+        )
 
     # Fill indexer with random indices for sparse backends
     is_sparse = backend_cfg.get("is_sparse", False)
     if is_sparse and indexer is not None:
         indexer.fill_random_indices(total_q, max_kv_len)
 
-    # Determine which forward method to use based on metadata
-    if metadata.decode is not None:
-        forward_fn = lambda: impl.forward_mqa(decode_inputs, kv_cache, metadata, layer)
-    elif metadata.prefill is not None:
-        forward_fn = lambda: impl.forward_mha(
-            prefill_inputs["q"],
-            prefill_inputs["k_c_normed"],
-            prefill_inputs["k_pe"],
-            kv_cache,
-            metadata,
-            prefill_inputs["k_scale"],
-            prefill_inputs["output"],
-        )
-    else:
+    # Determine which forward methods to use based on metadata.
+    # Sparse MLA backends always use forward_mqa
+    has_decode = is_sparse or getattr(metadata, "decode", None) is not None
+    has_prefill = not is_sparse and getattr(metadata, "prefill", None) is not None
+    if not has_decode and not has_prefill:
         raise RuntimeError("Metadata has neither decode nor prefill metadata")
 
+    num_decode = (
+        metadata.num_decode_tokens
+        if (has_decode and has_prefill)
+        else total_q
+        if has_decode
+        else 0
+    )
+    num_prefill = total_q - num_decode
+
+    # Some backends requires fp8 queries when using fp8 KV cache.
+    is_fp8_kvcache = kv_cache_dtype.startswith("fp8")
+    quantize_query = is_fp8_kvcache and getattr(
+        impl, "supports_quant_query_input", False
+    )
+
+    # quantize_query forces concat format
+    query_fmt = "concat" if quantize_query else backend_cfg["query_format"]
+
+    # Create decode query tensors
+    if has_decode:
+        decode_inputs, _ = _create_input_tensors(
+            num_decode, mla_dims, query_fmt, device, torch.bfloat16
+        )
+        # Cast decode query to fp8 if the backend supports it
+        if quantize_query:
+            from vllm.platforms import current_platform
+
+            if isinstance(decode_inputs, tuple):
+                decode_inputs = torch.cat(list(decode_inputs), dim=-1)
+            decode_inputs = decode_inputs.to(current_platform.fp8_dtype())
+
+    # Create prefill input tensors
+    if has_prefill:
+        _, prefill_inputs = _create_input_tensors(
+            num_prefill, mla_dims, query_fmt, device, torch.bfloat16
+        )
+
+    # Build forward function
+    def forward_fn():
+        results = []
+        if has_decode:
+            results.append(impl.forward_mqa(decode_inputs, kv_cache, metadata, layer))
+        if has_prefill:
+            results.append(
+                impl.forward_mha(
+                    prefill_inputs["q"],
+                    prefill_inputs["k_c_normed"],
+                    prefill_inputs["k_pe"],
+                    kv_cache,
+                    metadata,
+                    prefill_inputs["k_scale"],
+                    prefill_inputs["output"],
+                )
+            )
+        return results[0] if len(results) == 1 else tuple(results)
+
     # Warmup
     for _ in range(config.warmup_iters):
         forward_fn()
     torch.accelerator.synchronize()
 
+    # Optionally capture a CUDA graph after warmup.
+    # Graph replay eliminates CPU launch overhead so timings reflect pure
+    # kernel time.
+    if config.use_cuda_graphs:
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            forward_fn()
+        benchmark_fn = graph.replay
+    else:
+        benchmark_fn = forward_fn
+
     # Benchmark
     times = []
     for _ in range(config.repeats):
@@ -785,7 +864,7 @@ def _run_single_benchmark(
 
         start.record()
         for _ in range(config.num_layers):
-            forward_fn()
+            benchmark_fn()
         end.record()
 
         torch.accelerator.synchronize()
@@ -852,13 +931,30 @@ def _run_mla_benchmark_batched(
     # Determine if this is a sparse backend
     is_sparse = backend_cfg.get("is_sparse", False)
 
+    # Extract kv_cache_dtype from the first config
+    kv_cache_dtype = getattr(first_config, "kv_cache_dtype", "auto")
+
+    # FlashMLA sparse only supports "fp8_ds_mla" internally (not generic "fp8").
+    # Remap here so the user can pass --kv-cache-dtype fp8 regardless of backend.
+    if backend.upper() == "FLASHMLA_SPARSE" and kv_cache_dtype == "fp8":
+        kv_cache_dtype = "fp8_ds_mla"
+
+    # Compute max total_q across all configs so the metadata builder buffer
+    # and scheduler config are large enough for all batch specs.
+    max_total_q = max(
+        sum(r.q_len for r in parse_batch_spec(cfg.batch_spec))
+        for cfg, *_ in configs_with_params
+    )
+
     # Create and set vLLM config for MLA (reused across all benchmarks)
     vllm_config = create_minimal_vllm_config(
         model_name="deepseek-v3",  # Used only for model path
         block_size=block_size,
+        max_num_batched_tokens=max_total_q,
         mla_dims=mla_dims,  # Use custom dims from config or default
         index_topk=index_topk if is_sparse else None,
         prefill_backend=prefill_backend,
+        kv_cache_dtype=kv_cache_dtype,
     )
 
     results = []
@@ -883,7 +979,9 @@ def _run_mla_benchmark_batched(
             mla_dims,
             vllm_config,
             device,
+            max_num_tokens=max_total_q,
             index_topk=index_topk if is_sparse else None,
+            kv_cache_dtype=kv_cache_dtype,
         )
 
         # Verify the actual prefill backend matches what was requested
@@ -942,6 +1040,7 @@ def _run_mla_benchmark_batched(
                     mla_dims,
                     device,
                     indexer=indexer,
+                    kv_cache_dtype=kv_cache_dtype,
                 )
                 results.append(result)
 

benchmarks/attention_benchmarks/runner.py

@@ -140,7 +140,7 @@ def _create_vllm_config(
 
     cache_config = CacheConfig(
         block_size=config.block_size,
-        cache_dtype="auto",
+        cache_dtype=config.kv_cache_dtype,
     )
     cache_config.num_gpu_blocks = max_num_blocks
     cache_config.num_cpu_blocks = 0
@@ -215,7 +215,7 @@ def _create_backend_impl(
         num_kv_heads=config.num_kv_heads,
         alibi_slopes=None,
         sliding_window=None,
-        kv_cache_dtype="auto",
+        kv_cache_dtype=config.kv_cache_dtype,
     )
 
     kv_cache_spec = FullAttentionSpec(
@@ -288,12 +288,22 @@ def _create_input_tensors(
     total_q: int,
     device: torch.device,
     dtype: torch.dtype,
+    quantize_query: bool = False,
 ) -> tuple:
-    """Create Q, K, V input tensors for all layers."""
+    """Create Q, K, V input tensors for all layers.
+
+    When quantize_query is True, queries are cast to fp8 to match backends
+    that require query/key/value dtype consistency.
+    """
+    q_dtype = dtype
+    if quantize_query:
+        from vllm.platforms import current_platform
+
+        q_dtype = current_platform.fp8_dtype()
     q_list = [
         torch.randn(
             total_q, config.num_q_heads, config.head_dim, device=device, dtype=dtype
-        )
+        ).to(q_dtype)
         for _ in range(config.num_layers)
     ]
     k_list = [
@@ -344,10 +354,17 @@ def _create_kv_cache(
     # Compute inverse permutation to get back to logical view
     inv_order = [stride_order.index(i) for i in range(len(stride_order))]
 
+    # Use fp8 dtype for cache when requested.
+    cache_dtype = dtype
+    if config.kv_cache_dtype == "fp8":
+        from vllm.platforms import current_platform
+
+        cache_dtype = current_platform.fp8_dtype()
+
     cache_list = []
     for _ in range(config.num_layers):
         # Allocate in physical layout order (contiguous in memory)
-        cache = torch.zeros(*physical_shape, device=device, dtype=dtype)
+        cache = torch.zeros(*physical_shape, device=device, dtype=cache_dtype)
         # Permute to logical view
         cache = cache.permute(*inv_order)
         cache_list.append(cache)
@@ -392,6 +409,37 @@ def _run_single_benchmark(
             )
     torch.accelerator.synchronize()
 
+    # Optionally capture a CUDA graph after warmup.
+    # Graph replay eliminates CPU launch overhead so timings reflect pure
+    # kernel time.
+    if config.use_cuda_graphs:
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            for i in range(config.num_layers):
+                impl.forward(
+                    layer,
+                    q_list[i],
+                    k_list[i],
+                    v_list[i],
+                    cache_list[i],
+                    attn_metadata,
+                    output=out,
+                )
+        benchmark_fn = graph.replay
+    else:
+
+        def benchmark_fn():
+            for i in range(config.num_layers):
+                impl.forward(
+                    layer,
+                    q_list[i],
+                    k_list[i],
+                    v_list[i],
+                    cache_list[i],
+                    attn_metadata,
+                    output=out,
+                )
+
     # Benchmark
     times = []
     for _ in range(config.repeats):
@@ -399,16 +447,7 @@ def _run_single_benchmark(
         end = torch.cuda.Event(enable_timing=True)
 
         start.record()
-        for i in range(config.num_layers):
-            impl.forward(
-                layer,
-                q_list[i],
-                k_list[i],
-                v_list[i],
-                cache_list[i],
-                attn_metadata,
-                output=out,
-            )
+        benchmark_fn()
         end.record()
 
         torch.accelerator.synchronize()
@@ -502,8 +541,12 @@ def run_attention_benchmark(config: BenchmarkConfig) -> BenchmarkResult:
                 common_attn_metadata=common_metadata,
             )
 
+            # Only quantize queries when the impl supports it
+            quantize_query = config.kv_cache_dtype.startswith("fp8") and getattr(
+                impl, "supports_quant_query_input", False
+            )
             q_list, k_list, v_list = _create_input_tensors(
-                config, total_q, device, dtype
+                config, total_q, device, dtype, quantize_query=quantize_query
             )
 
             cache_list = _create_kv_cache(

benchmarks/kernels/cpu/benchmark_cpu_attn.py

@@ -27,7 +27,7 @@ def get_attn_isa(
     else:
         if current_platform.get_cpu_architecture() == CpuArchEnum.ARM:
             return "neon"
-        elif torch._C._cpu._is_amx_tile_supported():
+        elif torch.cpu._is_amx_tile_supported():
             return "amx"
         else:
             return "vec"

benchmarks/kernels/cpu/benchmark_cpu_fused_moe.py

@@ -24,7 +24,7 @@ except (ImportError, AttributeError) as e:
     sys.exit(1)
 
 # ISA selection following test_cpu_fused_moe.py pattern
-ISA_CHOICES = ["amx", "vec"] if torch._C._cpu._is_amx_tile_supported() else ["vec"]
+ISA_CHOICES = ["amx", "vec"] if torch.cpu._is_amx_tile_supported() else ["vec"]
 
 
 @torch.inference_mode()

cmake/external_projects/vllm_flash_attn.cmake

@@ -39,7 +39,7 @@ else()
   FetchContent_Declare(
           vllm-flash-attn
           GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git
-          GIT_TAG 1488682bb545f7d020e958a33116b1419d1cfc83
+          GIT_TAG 29210221863736a08f71a866459e368ad1ac4a95
           GIT_PROGRESS TRUE
           # Don't share the vllm-flash-attn build between build types
           BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn

csrc/ops.h

@@ -295,10 +295,14 @@ void cutlass_scaled_sparse_mm(torch::Tensor& out, torch::Tensor const& a,
 
 std::vector<torch::Tensor> cutlass_sparse_compress(torch::Tensor const& a);
 
-void scaled_fp4_quant(torch::Tensor& output, torch::Tensor const& input,
-                      torch::Tensor& output_scale,
-                      torch::Tensor const& input_scale,
-                      bool is_sf_swizzled_layout);
+std::tuple<torch::Tensor, torch::Tensor> scaled_fp4_quant_func(
+    torch::Tensor const& input, torch::Tensor const& input_scale,
+    bool is_sf_swizzled_layout);
+
+void scaled_fp4_quant_out(torch::Tensor const& input,
+                          torch::Tensor const& input_scale,
+                          bool is_sf_swizzled_layout, torch::Tensor& output,
+                          torch::Tensor& output_scale);
 
 void scaled_fp4_experts_quant(
     torch::Tensor& output, torch::Tensor& output_scale,

csrc/quantization/fp4/nvfp4_quant_entry.cu

@@ -16,6 +16,8 @@
 
 #include <torch/all.h>
 
+#include "nvfp4_utils.cuh"
+
 #if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
     (defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)
 void scaled_fp4_quant_sm1xxa(torch::Tensor const& output,
@@ -51,9 +53,10 @@ void silu_and_mul_scaled_fp4_experts_quant_sm1xxa(
     torch::Tensor const& output_scale_offset_by_experts);
 #endif
 
-void scaled_fp4_quant(torch::Tensor& output, torch::Tensor const& input,
-                      torch::Tensor& output_sf, torch::Tensor const& input_sf,
-                      bool is_sf_swizzled_layout) {
+void scaled_fp4_quant_out(torch::Tensor const& input,
+                          torch::Tensor const& input_sf,
+                          bool is_sf_swizzled_layout, torch::Tensor& output,
+                          torch::Tensor& output_sf) {
 #if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
     (defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)
   return scaled_fp4_quant_sm1xxa(output, input, output_sf, input_sf,
@@ -62,6 +65,34 @@ void scaled_fp4_quant(torch::Tensor& output, torch::Tensor const& input,
   TORCH_CHECK_NOT_IMPLEMENTED(false, "No compiled nvfp4 quantization kernel");
 }
 
+std::tuple<torch::Tensor, torch::Tensor> scaled_fp4_quant_func(
+    torch::Tensor const& input, torch::Tensor const& input_sf,
+    bool is_sf_swizzled_layout) {
+  int64_t n = input.size(-1);
+  int64_t m = input.numel() / n;
+  auto device = input.device();
+
+  // Two fp4 values packed into a uint8
+  auto output = torch::empty(
+      {m, n / 2}, torch::TensorOptions().device(device).dtype(torch::kUInt8));
+
+  torch::Tensor output_sf;
+  if (is_sf_swizzled_layout) {
+    auto [sf_m, sf_n] = vllm::computeSwizzledSFShape(m, n);
+    output_sf = torch::empty(
+        {sf_m, sf_n},
+        torch::TensorOptions().device(device).dtype(torch::kInt32));
+  } else {
+    output_sf = torch::empty(
+        {m, n / CVT_FP4_SF_VEC_SIZE},
+        torch::TensorOptions().device(device).dtype(torch::kUInt8));
+  }
+
+  scaled_fp4_quant_out(input, input_sf, is_sf_swizzled_layout, output,
+                       output_sf);
+  return {output, output_sf};
+}
+
 void scaled_fp4_experts_quant(
     torch::Tensor& output, torch::Tensor& output_scale,
     torch::Tensor const& input, torch::Tensor const& input_global_scale,

csrc/quantization/fp4/nvfp4_utils.cuh

@@ -18,6 +18,7 @@
 
 #include <cuda_runtime.h>
 #include <cuda_fp8.h>
+#include <utility>
 
 #include "../../cuda_vec_utils.cuh"
 
@@ -54,6 +55,18 @@ inline int computeEffectiveRows(int m) {
   return round_up(m, ROW_TILE);
 }
 
+// Compute the shape of the swizzled SF output tensor.
+// Returns (rounded_m, rounded_n / 4) where:
+//   rounded_m = round_up(m, 128)
+//   rounded_n = round_up(n / CVT_FP4_SF_VEC_SIZE, 4)
+inline std::pair<int64_t, int64_t> computeSwizzledSFShape(int64_t m,
+                                                          int64_t n) {
+  int64_t rounded_m = round_up(m, static_cast<int64_t>(128));
+  int64_t scale_n = n / CVT_FP4_SF_VEC_SIZE;
+  int64_t rounded_n = round_up(scale_n, static_cast<int64_t>(4));
+  return {rounded_m, rounded_n / 4};
+}
+
 // Convert 8 float32 values into 8 e2m1 values (represented as one uint32_t).
 inline __device__ uint32_t fp32_vec8_to_e2m1(float (&array)[8]) {
   uint32_t val;

csrc/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu

@@ -286,6 +286,15 @@ void rms_norm_per_block_quant(torch::Tensor& out, torch::Tensor const& input,
                 "Outer scale stride must be 1 when scales are not transposed");
   }
 
+  int64_t hidden_size = input.size(-1);
+  TORCH_CHECK(hidden_size > 0 && hidden_size % group_size == 0,
+              "hidden_size must be a positive multiple of group_size");
+  int64_t num_tokens = input.numel() / hidden_size;
+  int64_t num_groups = hidden_size / group_size;
+  TORCH_CHECK(scales.numel() >= num_tokens * num_groups,
+              "scales buffer too small: need ", num_tokens * num_groups,
+              " elements, got ", scales.numel());
+
   rms_norm_per_block_quant_dispatch(out, input, weight, scales, group_size,
                                     var_epsilon, scale_ub, residual,
                                     is_scale_transposed);

csrc/torch_bindings.cpp

@@ -564,10 +564,21 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
 
   // Compute NVFP4 block quantized tensor.
   ops.def(
-      "scaled_fp4_quant(Tensor! output, Tensor input,"
-      "                 Tensor! output_scale, Tensor input_scale, bool "
-      "is_sf_swizzled_layout) -> ()");
-  ops.impl("scaled_fp4_quant", torch::kCUDA, &scaled_fp4_quant);
+      "scaled_fp4_quant(Tensor input,"
+      "                 Tensor input_scale, bool "
+      "is_sf_swizzled_layout) -> (Tensor, Tensor)");
+  ops.impl("scaled_fp4_quant", torch::kCUDA, &scaled_fp4_quant_func);
+
+  // Out variant
+  // TODO: Add {at::Tag::out_variant} tag and update all call sites
+  // to use the functional variant once vLLM upgrades PyTorch.
+  // See pytorch/pytorch#176117.
+  ops.def(
+      "scaled_fp4_quant.out(Tensor input,"
+      "                     Tensor input_scale, bool "
+      "is_sf_swizzled_layout, *, Tensor(a!) output, Tensor(b!) output_scale) "
+      "-> ()");
+  ops.impl("scaled_fp4_quant.out", torch::kCUDA, &scaled_fp4_quant_out);
 
   // Compute NVFP4 experts quantization.
   ops.def(

docs/design/attention_backends.md

@@ -127,8 +127,8 @@ Priority is **1 = highest** (tried first).
 | 3 | `FLASH_ATTN_MLA` |
 | 4 | `FLASHMLA` |
 | 5 | `TRITON_MLA` |
-| 6 | `FLASHMLA_SPARSE` |
-| 7 | `FLASHINFER_MLA_SPARSE` |
+| 6 | `FLASHINFER_MLA_SPARSE`**\*** |
+| 7 | `FLASHMLA_SPARSE` |
 
 **Ampere/Hopper (SM 8.x-9.x):**
 
@@ -140,6 +140,8 @@ Priority is **1 = highest** (tried first).
 | 4 | `TRITON_MLA` |
 | 5 | `FLASHMLA_SPARSE` |
 
+> **\*** For sparse MLA, FP8 KV cache always prefers `FLASHINFER_MLA_SPARSE`. With BF16 KV cache, `FLASHINFER_MLA_SPARSE` is preferred for low query-head counts (<= 16), while `FLASHMLA_SPARSE` is preferred otherwise.
+>
 > **Note:** ROCm and CPU platforms have their own selection logic. See the platform-specific documentation for details.
 
 ## Legend

docs/features/lora.md

@@ -389,3 +389,17 @@ vllm serve model --enable-lora --max-lora-rank 64
 # Bad: unnecessarily high, wastes memory
 vllm serve model --enable-lora --max-lora-rank 256
 ```
+
+### Restricting LoRA to Specific Modules
+
+The `--lora-target-modules` parameter allows you to restrict which model modules have LoRA applied at deployment time. This is useful for performance tuning when you only need LoRA on specific layers:
+
+```bash
+# Apply LoRA only to output projection layers
+vllm serve model --enable-lora --lora-target-modules o_proj
+
+# Apply LoRA to multiple specific modules
+vllm serve model --enable-lora --lora-target-modules o_proj qkv_proj down_proj
+```
+
+When `--lora-target-modules` is not specified, LoRA will be applied to all supported modules in the model. This parameter accepts module suffixes (the last component of the module name), such as `o_proj`, `qkv_proj`, `gate_proj`, etc.

docs/features/tool_calling.md

@@ -107,6 +107,27 @@ vLLM supports the `tool_choice='none'` option in the chat completion API. When t
 !!! note
     When tools are specified in the request, vLLM includes tool definitions in the prompt by default, regardless of the `tool_choice` setting. To exclude tool definitions when `tool_choice='none'`, use the `--exclude-tools-when-tool-choice-none` option.
 
+## Constrained Decoding Behavior
+
+Whether vLLM enforces the tool parameter schema during generation depends on the `tool_choice` mode:
+
+| `tool_choice` value | Schema-constrained decoding | Behavior |
+| --- | --- | --- |
+| Named function | Yes (via structured outputs backend) | Arguments are guaranteed to be valid JSON conforming to the function's parameter schema. |
+| `"required"` | Yes (via structured outputs backend) | Same as named function. The model must produce at least one tool call. |
+| `"auto"` | No | The model generates freely. A tool-call parser extracts tool calls from the raw text. Arguments may be malformed or not match the schema. |
+| `"none"` | N/A | No tool calls are produced. |
+
+When schema conformance matters, prefer `tool_choice="required"` or named function calling over `"auto"`.
+
+### Strict Mode (`strict` parameter)
+
+The [OpenAI API](https://platform.openai.com/docs/guides/function-calling#strict-mode) supports a `strict` field on function definitions. When set to `true`, OpenAI uses constrained decoding to guarantee that tool-call arguments match the function schema, even in `tool_choice="auto"` mode.
+
+vLLM **does not implement** `strict` mode today. The `strict` field is accepted in requests (to avoid breaking clients that set it), but it has no effect on decoding behavior. In auto mode, argument validity depends entirely on the model's output quality and the parser's extraction logic.
+
+Tracking issues: [#15526](https://github.com/vllm-project/vllm/issues/15526), [#16313](https://github.com/vllm-project/vllm/issues/16313).
+
 ## Automatic Function Calling
 
 To enable this feature, you should set the following flags:
@@ -124,6 +145,9 @@ from HuggingFace; and you can find an example of this in a `tokenizer_config.jso
 
 If your favorite tool-calling model is not supported, please feel free to contribute a parser & tool use chat template!
 
+!!! note
+    With `tool_choice="auto"`, tool-call arguments are extracted from the model's raw text output by the selected parser. No schema-level constraint is applied during decoding, so arguments may occasionally be malformed or violate the function's parameter schema. See [Constrained Decoding Behavior](#constrained-decoding-behavior) for details.
+
 ### Hermes Models (`hermes`)
 
 All Nous Research Hermes-series models newer than Hermes 2 Pro should be supported.

docs/serving/openai_compatible_server.md

@@ -72,6 +72,9 @@ In addition, we have the following custom APIs:
     - Only applicable to [classification models](../models/pooling_models.md).
 - [Score API](#score-api) (`/score`)
     - Applicable to [embedding models and cross-encoder models](../models/pooling_models.md).
+- [Cohere Embed API](#cohere-embed-api) (`/v2/embed`)
+    - Compatible with [Cohere's Embed API](https://docs.cohere.com/reference/embed)
+    - Works with any [embedding model](../models/pooling_models.md), including multimodal models.
 - [Re-rank API](#re-rank-api) (`/rerank`, `/v1/rerank`, `/v2/rerank`)
     - Implements [Jina AI's v1 re-rank API](https://jina.ai/reranker/)
     - Also compatible with [Cohere's v1 & v2 re-rank APIs](https://docs.cohere.com/v2/reference/rerank)
@@ -429,6 +432,137 @@ these extra parameters are supported instead:
     --8<-- "vllm/entrypoints/pooling/base/protocol.py:embed-extra-params"
     ```
 
+### Cohere Embed API
+
+Our API is also compatible with [Cohere's Embed v2 API](https://docs.cohere.com/reference/embed) which adds support for some modern embedding feature such as truncation, output dimensions, embedding types, and input types. This endpoint works with any embedding model (including multimodal models).
+
+#### Cohere Embed API request parameters
+
+| Parameter | Type | Required | Description |
+| --------- | ---- | -------- | ----------- |
+| `model` | string | Yes | Model name |
+| `input_type` | string | No | Prompt prefix key (model-dependent, see below) |
+| `texts` | list[string] | No | Text inputs (use one of `texts`, `images`, or `inputs`) |
+| `images` | list[string] | No | Base64 data URI images |
+| `inputs` | list[object] | No | Mixed text and image content objects |
+| `embedding_types` | list[string] | No | Output types (default: `["float"]`) |
+| `output_dimension` | int | No | Truncate embeddings to this dimension (Matryoshka) |
+| `truncate` | string | No | `END`, `START`, or `NONE` (default: `END`) |
+
+#### Text embedding
+
+```bash
+curl -X POST "http://localhost:8000/v2/embed" \
+  -H "Content-Type: application/json" \
+  -d '{
+    "model": "Snowflake/snowflake-arctic-embed-m-v1.5",
+    "input_type": "query",
+    "texts": ["Hello world", "How are you?"],
+    "embedding_types": ["float"]
+  }'
+```
+
+??? console "Response"
+
+    ```json
+    {
+      "id": "embd-...",
+      "embeddings": {
+        "float": [
+          [0.012, -0.034, ...],
+          [0.056, 0.078, ...]
+        ]
+      },
+      "texts": ["Hello world", "How are you?"],
+      "meta": {
+        "api_version": {"version": "2"},
+        "billed_units": {"input_tokens": 12}
+      }
+    }
+    ```
+
+#### Mixed text and image inputs
+
+For multimodal models, you can embed images by passing base64 data URIs. The `inputs` field accepts a list of objects with mixed text and image content:
+
+```bash
+curl -X POST "http://localhost:8000/v2/embed" \
+  -H "Content-Type: application/json" \
+  -d '{
+    "model": "google/siglip-so400m-patch14-384",
+    "inputs": [
+      {
+        "content": [
+          {"type": "text", "text": "A photo of a cat"},
+          {"type": "image_url", "image_url": {"url": "data:image/png;base64,iVBOR..."}}
+        ]
+      }
+    ],
+    "embedding_types": ["float"]
+  }'
+```
+
+#### Embedding types
+
+The `embedding_types` parameter controls the output format. Multiple types can be requested in a single call:
+
+| Type | Description |
+| ---- | ----------- |
+| `float` | Raw float32 embeddings (default) |
+| `binary` | Bit-packed signed binary |
+| `ubinary` | Bit-packed unsigned binary |
+| `base64` | Little-endian float32 encoded as base64 |
+
+```bash
+curl -X POST "http://localhost:8000/v2/embed" \
+  -H "Content-Type: application/json" \
+  -d '{
+    "model": "Snowflake/snowflake-arctic-embed-m-v1.5",
+    "input_type": "query",
+    "texts": ["What is machine learning?"],
+    "embedding_types": ["float", "binary"]
+  }'
+```
+
+??? console "Response"
+
+    ```json
+    {
+      "id": "embd-...",
+      "embeddings": {
+        "float": [[0.012, -0.034, ...]],
+        "binary": [[42, -117, ...]]
+      },
+      "texts": ["What is machine learning?"],
+      "meta": {
+        "api_version": {"version": "2"},
+        "billed_units": {"input_tokens": 8}
+      }
+    }
+    ```
+
+#### Truncation
+
+The `truncate` parameter controls how inputs exceeding the model's maximum sequence length are handled:
+
+| Value | Behavior |
+| ----- | --------- |
+| `END` (default) | Keep the first tokens, drop the end |
+| `START` | Keep the last tokens, drop the beginning |
+| `NONE` | Return an error if the input is too long |
+
+#### Input type and prompt prefixes
+
+The `input_type` field selects a prompt prefix to prepend to each text input. The available values
+depend on the model:
+
+- **Models with `task_instructions` in `config.json`**: The keys from the `task_instructions` dict are
+  the valid `input_type` values and the corresponding value is prepended to each text.
+- **Models with `config_sentence_transformers.json` prompts**: The keys from the `prompts` dict are
+  the valid `input_type` values. For example, `Snowflake/snowflake-arctic-embed-xs` defines `"query"`,
+  so setting `input_type: "query"` prepends `"Represent this sentence for searching relevant passages: "`.
+- **Other models**: `input_type` is not accepted and will raise a validation error if passed.
+
 ### Transcriptions API
 
 Our Transcriptions API is compatible with [OpenAI's Transcriptions API](https://platform.openai.com/docs/api-reference/audio/createTranscription);

examples/online_serving/openai_chat_completion_client_for_multimodal.py

@@ -20,9 +20,9 @@ run the script with
 python openai_chat_completion_client_for_multimodal.py --chat-type audio
 """
 
-import base64
 import os
 
+import pybase64 as base64
 import requests
 from openai import OpenAI
 from utils import get_first_model

examples/online_serving/openai_realtime_client.py

@@ -24,11 +24,11 @@ The script:
 
 import argparse
 import asyncio
-import base64
 import json
 
 import librosa
 import numpy as np
+import pybase64 as base64
 import websockets
 
 from vllm.assets.audio import AudioAsset

examples/online_serving/openai_realtime_microphone_client.py

@@ -18,13 +18,13 @@ Requirements: websockets, numpy, gradio
 
 import argparse
 import asyncio
-import base64
 import json
 import queue
 import threading
 
 import gradio as gr
 import numpy as np
+import pybase64 as base64
 import websockets
 
 SAMPLE_RATE = 16_000

examples/pooling/embed/embedding_requests_base64_online.py

@@ -7,8 +7,8 @@ NOTE:
 """
 
 import argparse
-import base64
 
+import pybase64 as base64
 import requests
 import torch
 

examples/pooling/embed/vision_embedding_online.py

@@ -7,10 +7,10 @@ Refer to each `run_*` function for the command to run the server for that model.
 """
 
 import argparse
-import base64
 import io
 from typing import Literal
 
+import pybase64 as base64
 from openai import OpenAI
 from openai._types import NOT_GIVEN, NotGiven
 from openai.types.chat import ChatCompletionMessageParam

examples/pooling/plugin/prithvi_geospatial_mae_io_processor.py

@@ -1,8 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-import base64
 import os
 
+import pybase64 as base64
 import torch
 
 from vllm import LLM

examples/pooling/plugin/prithvi_geospatial_mae_online.py

@@ -1,9 +1,9 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-import base64
 import os
 
+import pybase64 as base64
 import requests
 
 # This example shows how to perform an online inference that generates

examples/pooling/score/colqwen3_rerank_online.py

@@ -15,9 +15,9 @@ Then run this script:
     python colqwen3_rerank_online.py
 """
 
-import base64
 from io import BytesIO
 
+import pybase64 as base64
 import requests
 from PIL import Image
 

examples/pooling/token_embed/colqwen3_token_embed_online.py

@@ -21,10 +21,10 @@ Then run this script:
 """
 
 import argparse
-import base64
 from io import BytesIO
 
 import numpy as np
+import pybase64 as base64
 import requests
 from PIL import Image
 

requirements/common.txt

@@ -37,7 +37,7 @@ pyyaml
 six>=1.16.0; python_version > '3.11' # transitive dependency of pandas that needs to be the latest version for python 3.12
 setuptools>=77.0.3,<81.0.0; python_version > '3.11' # Setuptools is used by triton, we need to ensure a modern version is installed for 3.12+ so that it does not try to import distutils, which was removed in 3.12
 einops # Required for Qwen2-VL.
-compressed-tensors == 0.13.0 # required for compressed-tensors
+compressed-tensors == 0.14.0.1 # required for compressed-tensors
 depyf==0.20.0 # required for profiling and debugging with compilation config
 cloudpickle # allows pickling lambda functions in model_executor/models/registry.py
 watchfiles # required for http server to monitor the updates of TLS files

requirements/rocm-test.txt

@@ -50,7 +50,7 @@ av==16.1.0
 blobfile==3.0.0
     # Multi-Modal Models Test
 decord==0.6.0
-    # video processing, required by entrypoints/openai/test_video.py
+    # video processing, required by entrypoints/openai/chat_completion/test_video.py
 rapidfuzz==3.12.1
 
 # OpenAI compatibility and testing

tests/benchmarks/test_random_multimodal_dataset_video.py

@@ -1,12 +1,12 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-import base64
 import os
 from tempfile import NamedTemporaryFile
 from typing import Any, cast
 
 import cv2
+import pybase64 as base64
 import pytest
 from transformers import AutoTokenizer, PreTrainedTokenizerBase
 

tests/compile/passes/distributed/test_fusion_all_reduce.py

@@ -179,7 +179,7 @@ class TestAllReduceFusedAddRMSNormStaticQuantFP4Model(torch.nn.Module):
     def ops_in_model_before(self):
         return [
             torch.ops.vllm.all_reduce.default,
-            torch.ops._C.scaled_fp4_quant.default,
+            torch.ops._C.scaled_fp4_quant.out,
         ]
 
 

tests/distributed/test_distributed_oot.py

@@ -1,7 +1,9 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-from ..entrypoints.openai.test_oot_registration import run_and_test_dummy_opt_api_server
+from tests.entrypoints.openai.chat_completion.test_oot_registration import (
+    run_and_test_dummy_opt_api_server,
+)
 
 
 def test_distributed_oot(dummy_opt_path: str):

tests/distributed/test_weight_transfer.py

@@ -6,10 +6,10 @@ Unit tests for engine classes (parsing, validation, registry).
 Integration tests for NCCL and IPC weight transfer between processes using Ray.
 """
 
-import base64
 import pickle
 from unittest.mock import MagicMock
 
+import pybase64 as base64
 import pytest
 import ray
 import torch

tests/entrypoints/llm/test_chat.py

@@ -4,12 +4,11 @@ import weakref
 
 import pytest
 
+from tests.entrypoints.openai.chat_completion.test_vision import TEST_IMAGE_ASSETS
 from vllm import LLM
 from vllm.distributed import cleanup_dist_env_and_memory
 from vllm.sampling_params import SamplingParams
 
-from ..openai.test_vision import TEST_IMAGE_ASSETS
-
 
 @pytest.fixture(scope="function")
 def text_llm():

tests/entrypoints/llm/test_mm_cache_stats.py

@@ -6,13 +6,12 @@ import logging
 import pytest
 import regex as re
 
+from tests.entrypoints.openai.chat_completion.test_vision import TEST_IMAGE_ASSETS
 from vllm import LLM
 from vllm.entrypoints.chat_utils import ChatCompletionMessageParam
 from vllm.v1.metrics import loggers as stat_loggers
 from vllm.v1.metrics.reader import Counter, Metric
 
-from ..openai.test_vision import TEST_IMAGE_ASSETS
-
 
 def _make_messages(image_url: str) -> list[ChatCompletionMessageParam]:
     return [

tests/entrypoints/openai/chat_completion/test_audio.py

@@ -7,11 +7,10 @@ import openai
 import pytest
 import pytest_asyncio
 
+from tests.utils import RemoteOpenAIServer
 from vllm.assets.audio import AudioAsset
 from vllm.multimodal.utils import encode_audio_base64, encode_audio_url, fetch_audio
 
-from ...utils import RemoteOpenAIServer
-
 MODEL_NAME = "fixie-ai/ultravox-v0_5-llama-3_2-1b"
 TEST_AUDIO_URLS = [
     AudioAsset("winning_call").url,

tests/entrypoints/openai/chat_completion/test_audio_in_video.py

@@ -1,15 +1,15 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-import base64
 import json
 
 import openai
+import pybase64 as base64
 import pytest
 import pytest_asyncio
 
-from ...conftest import VideoTestAssets
-from ...utils import RemoteOpenAIServer
+from tests.conftest import VideoTestAssets
+from tests.utils import RemoteOpenAIServer
 
 MODEL_NAME = "Qwen/Qwen2.5-Omni-3B"
 

tests/entrypoints/openai/chat_completion/test_chat_error.py

@@ -111,7 +111,7 @@ def _build_serving_chat(engine: AsyncLLM) -> OpenAIServingChat:
             [{"prompt_token_ids": [1, 2, 3]}],
         )
 
-    serving_chat.openai_serving_render._preprocess_chat = AsyncMock(
+    serving_chat.openai_serving_render.preprocess_chat = AsyncMock(
         side_effect=_fake_preprocess_chat
     )
     return serving_chat

tests/entrypoints/openai/chat_completion/test_default_mm_loras.py

@@ -8,8 +8,8 @@ import pytest
 import pytest_asyncio
 from huggingface_hub import snapshot_download
 
-from ...conftest import AudioTestAssets
-from ...utils import RemoteOpenAIServer
+from tests.conftest import AudioTestAssets
+from tests.utils import RemoteOpenAIServer
 
 # NOTE - the tests in this module are currently analogous to test_chat, but are
 # separated to avoid OOM killing due to module-scoped servers, since we

tests/entrypoints/openai/chat_completion/test_oot_registration.py

@@ -1,7 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-from ...utils import VLLM_PATH, RemoteOpenAIServer
+from tests.utils import VLLM_PATH, RemoteOpenAIServer
 
 chatml_jinja_path = VLLM_PATH / "examples/template_chatml.jinja"
 assert chatml_jinja_path.exists()

tests/entrypoints/openai/chat_completion/test_root_path.py

@@ -8,7 +8,7 @@ from typing import Any, NamedTuple
 import openai  # use the official client for correctness check
 import pytest
 
-from ...utils import RemoteOpenAIServer
+from tests.utils import RemoteOpenAIServer
 
 # # any model with a chat template should work here
 MODEL_NAME = "Qwen/Qwen2-1.5B-Instruct"

tests/entrypoints/openai/chat_completion/test_video.py

@@ -7,11 +7,10 @@ import openai
 import pytest
 import pytest_asyncio
 
+from tests.utils import RemoteOpenAIServer
 from vllm.multimodal.utils import encode_video_url, fetch_video
 from vllm.platforms import current_platform
 
-from ...utils import RemoteOpenAIServer
-
 MODEL_NAME = "llava-hf/llava-onevision-qwen2-0.5b-ov-hf"
 MAXIMUM_VIDEOS = 3
 

tests/entrypoints/openai/chat_completion/test_vision.py

@@ -8,12 +8,11 @@ import pytest
 import pytest_asyncio
 from transformers import AutoProcessor
 
+from tests.utils import ROCM_ENV_OVERRIDES, ROCM_EXTRA_ARGS, RemoteOpenAIServer
 from vllm.multimodal.media import MediaWithBytes
 from vllm.multimodal.utils import encode_image_url, fetch_image
 from vllm.platforms import current_platform
 
-from ...utils import ROCM_ENV_OVERRIDES, ROCM_EXTRA_ARGS, RemoteOpenAIServer
-
 MODEL_NAME = "microsoft/Phi-3.5-vision-instruct"
 MAXIMUM_IMAGES = 2
 

tests/entrypoints/openai/chat_completion/test_vision_embeds.py

@@ -1,17 +1,15 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-import base64
-
 import numpy as np
+import pybase64 as base64
 import pytest
 import requests
 import torch
 
+from tests.utils import RemoteOpenAIServer
 from vllm.utils.serial_utils import tensor2base64
 
-from ...utils import RemoteOpenAIServer
-
 
 @pytest.mark.parametrize(
     "model_name", ["ibm-nasa-geospatial/Prithvi-EO-2.0-300M-TL-Sen1Floods11"]

tests/entrypoints/openai/completion/test_completion_with_prompt_embeds.py

@@ -1,11 +1,11 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-import base64
 import io
 import json
 
 import openai  # use the official client for correctness check
+import pybase64 as base64
 import pytest
 import pytest_asyncio
 import torch
@@ -14,7 +14,7 @@ import torch
 from openai import BadRequestError
 from transformers import AutoConfig
 
-from ...utils import RemoteOpenAIServer
+from tests.utils import RemoteOpenAIServer
 
 # any model with a chat template should work here
 MODEL_NAME = "facebook/opt-125m"

tests/entrypoints/openai/completion/test_prompt_validation.py

@@ -11,11 +11,10 @@ import pytest
 import regex as re
 import torch
 
+from tests.utils import RemoteOpenAIServer
 from vllm.config import ModelConfig
 from vllm.renderers.embed_utils import safe_load_prompt_embeds
 
-from ...utils import RemoteOpenAIServer
-
 
 @pytest.mark.asyncio
 async def test_empty_prompt():

tests/entrypoints/openai/completion/test_tensorizer_entrypoint.py

@@ -9,6 +9,7 @@ import pytest
 import pytest_asyncio
 import torch.cuda
 
+from tests.utils import RemoteOpenAIServer
 from vllm.engine.arg_utils import EngineArgs
 from vllm.model_executor.model_loader.tensorizer import (
     TensorizerConfig,
@@ -17,8 +18,6 @@ from vllm.model_executor.model_loader.tensorizer import (
 )
 from vllm.platforms import current_platform
 
-from ...utils import RemoteOpenAIServer
-
 MODEL_NAME = "unsloth/llama-3.2-1b-Instruct"
 LORA_PATH = "davzoku/finqa_adapter_1b"
 

tests/entrypoints/openai/completion/test_token_in_token_out.py

@@ -6,11 +6,10 @@ import tempfile
 
 import pytest
 
+from tests.utils import RemoteOpenAIServer
 from vllm.model_executor.model_loader.weight_utils import download_weights_from_hf
 from vllm.tokenizers import get_tokenizer
 
-from ...utils import RemoteOpenAIServer
-
 MODEL_NAME = "Qwen/Qwen3-0.6B"
 MODEL_PATH = os.path.join(tempfile.gettempdir(), "qwen3_06b")
 

tests/entrypoints/openai/responses/conftest.py

@@ -8,6 +8,9 @@ from collections.abc import Callable
 from typing import Any
 
 import pytest
+import pytest_asyncio
+
+from tests.utils import RemoteOpenAIServer
 
 logger = logging.getLogger(__name__)
 
@@ -361,3 +364,38 @@ def log_response_diagnostics(
     )
 
     return diagnostics
+
+
+@pytest.fixture(scope="module")
+def default_server_args():
+    return [
+        "--max-model-len",
+        "8192",
+        "--enforce-eager",  # For faster startup.
+        "--enable-auto-tool-choice",
+        "--structured-outputs-config.backend",
+        "xgrammar",
+        "--tool-call-parser",
+        "hermes",
+        "--reasoning-parser",
+        "qwen3",
+    ]
+
+
+@pytest.fixture(scope="module")
+def server_with_store(default_server_args):
+    with RemoteOpenAIServer(
+        "Qwen/Qwen3-1.7B",
+        default_server_args,
+        env_dict={
+            "VLLM_ENABLE_RESPONSES_API_STORE": "1",
+            "VLLM_SERVER_DEV_MODE": "1",
+        },
+    ) as remote_server:
+        yield remote_server
+
+
+@pytest_asyncio.fixture
+async def client(server_with_store):
+    async with server_with_store.get_async_client() as async_client:
+        yield async_client

tests/entrypoints/openai/responses/test_function_call.py

@@ -134,6 +134,34 @@ async def test_function_tool_use(
     assert reasoning.type == "reasoning"
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("model_name", [MODEL_NAME])
+async def test_max_tokens_with_tool_choice_required(
+    client: openai.AsyncOpenAI, model_name: str
+):
+    prompt = [
+        {
+            "role": "user",
+            "content": "Can you tell me what the current weather is in Berlin and the "
+            "forecast for the next 5 days, in fahrenheit?",
+        },
+    ]
+    response = await client.responses.create(
+        model=model_name,
+        input=prompt,
+        tools=tools,
+        tool_choice="required",
+        max_output_tokens=10,
+    )
+    assert len(response.output) >= 1
+    for out in response.output:
+        # When `tool_choice="required"` and the tokens of `tools`
+        # exceed `max_output_tokens`,`function_call` should be empty.
+        # This behavior should be consistent with OpenAI
+        assert out.type != "function_call"
+    assert response.incomplete_details.reason == "max_output_tokens"
+
+
 @pytest.mark.asyncio
 async def test_named_tool_use(client: openai.AsyncOpenAI):
     def get_weather(latitude: float, longitude: float) -> str:

tests/entrypoints/openai/responses/test_harmony.py

@@ -16,7 +16,8 @@ import requests
 from openai import InternalServerError, NotFoundError, OpenAI
 from openai_harmony import Message
 
-from ....utils import RemoteOpenAIServer
+from tests.utils import RemoteOpenAIServer
+
 from .conftest import (
     BASE_TEST_ENV,
     events_contain_type,

tests/entrypoints/openai/responses/test_mcp_tools.py

@@ -9,9 +9,9 @@ import pytest_asyncio
 from openai import OpenAI
 from openai_harmony import ToolDescription, ToolNamespaceConfig
 
+from tests.utils import RemoteOpenAIServer
 from vllm.entrypoints.mcp.tool_server import MCPToolServer
 
-from ....utils import RemoteOpenAIServer
 from .conftest import (
     BASE_TEST_ENV,
     events_contain_type,

tests/entrypoints/openai/responses/test_parsable_context.py

@@ -9,7 +9,8 @@ import pytest
 import pytest_asyncio
 from openai import OpenAI
 
-from ....utils import RemoteOpenAIServer
+from tests.utils import RemoteOpenAIServer
+
 from .conftest import (
     BASE_TEST_ENV,
     has_output_type,

tests/entrypoints/openai/responses/test_simple.py

@@ -5,7 +5,8 @@ import pytest
 import pytest_asyncio
 from openai import OpenAI
 
-from ....utils import RemoteOpenAIServer
+from tests.utils import RemoteOpenAIServer
+
 from .conftest import validate_streaming_event_stack
 
 MODEL_NAME = "Qwen/Qwen3-8B"

tests/entrypoints/openai/test_cli_args.py

@@ -291,3 +291,32 @@ def test_served_model_name_parsing(tmp_path, vllm_parser, args, raises):
     else:
         with pytest.raises(raises):
             vllm_parser.parse_args(args=args)
+
+
+### Tests for LoRA target modules parsing
+def test_lora_target_modules_single(serve_parser):
+    """Test parsing single lora-target-modules argument"""
+    args = serve_parser.parse_args(
+        args=["--enable-lora", "--lora-target-modules", "o_proj"]
+    )
+    assert args.lora_target_modules == ["o_proj"]
+
+
+def test_lora_target_modules_multiple(serve_parser):
+    """Test parsing multiple lora-target-modules arguments"""
+    args = serve_parser.parse_args(
+        args=[
+            "--enable-lora",
+            "--lora-target-modules",
+            "o_proj",
+            "qkv_proj",
+            "down_proj",
+        ]
+    )
+    assert args.lora_target_modules == ["o_proj", "qkv_proj", "down_proj"]
+
+
+def test_lora_target_modules_default_none(serve_parser):
+    """Test that lora-target-modules defaults to None"""
+    args = serve_parser.parse_args(args=[])
+    assert args.lora_target_modules is None

tests/entrypoints/openai/test_realtime_validation.py

@@ -2,12 +2,12 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import asyncio
-import base64
 import json
 import warnings
 
 import librosa
 import numpy as np
+import pybase64 as base64
 import pytest
 import websockets
 

tests/entrypoints/openai/test_serving_responses.py

@@ -159,6 +159,7 @@ class TestInitializeToolSessions:
         instance = OpenAIServingResponses(
             engine_client=engine_client,
             models=models,
+            openai_serving_render=MagicMock(),
             request_logger=None,
             chat_template=None,
             chat_template_content_format="auto",
@@ -245,6 +246,7 @@ class TestValidateGeneratorInput:
         instance = OpenAIServingResponses(
             engine_client=engine_client,
             models=models,
+            openai_serving_render=MagicMock(),
             request_logger=None,
             chat_template=None,
             chat_template_content_format="auto",
@@ -308,6 +310,7 @@ async def test_reasoning_tokens_counted_for_text_reasoning_model(monkeypatch):
     serving = OpenAIServingResponses(
         engine_client=engine_client,
         models=models,
+        openai_serving_render=MagicMock(),
         request_logger=None,
         chat_template=None,
         chat_template_content_format="auto",
@@ -607,6 +610,7 @@ def _make_serving_instance_with_reasoning():
     serving = OpenAIServingResponses(
         engine_client=engine_client,
         models=models,
+        openai_serving_render=MagicMock(),
         request_logger=None,
         chat_template=None,
         chat_template_content_format="auto",

tests/entrypoints/openai/tool_parsers/test_gigachat3_tool_parser.py

@@ -5,7 +5,7 @@ import json
 
 import pytest
 
-from tests.entrypoints.openai.tool_parsers.utils import (
+from tests.tool_parsers.utils import (
     run_tool_extraction,
     run_tool_extraction_streaming,
 )
@@ -13,6 +13,13 @@ from vllm.entrypoints.openai.engine.protocol import FunctionCall
 from vllm.tokenizers import TokenizerLike
 from vllm.tool_parsers import ToolParser, ToolParserManager
 
+MSG_SEP_TOKEN = "<|message_sep|>\n\n"
+ROLE_SEP_TOKEN = "<|role_sep|>\n"
+EOS_TOKEN = "</s>"
+TOOL_HEADER_GIGACHAT3 = f"function call{ROLE_SEP_TOKEN}"
+TOOL_HEADER_GIGACHAT31 = "<|function_call|>"
+
+
 SIMPLE_ARGS_DICT = {
     "action": "create",
     "id": "preferences",
@@ -24,7 +31,10 @@ SIMPLE_FUNCTION_JSON = json.dumps(
     },
     ensure_ascii=False,
 )
-SIMPLE_FUNCTION_OUTPUT = "function call" + SIMPLE_FUNCTION_JSON
+SIMPLE_FUNCTION_OUTPUT_GIGACHAT3 = (
+    f"{MSG_SEP_TOKEN}{TOOL_HEADER_GIGACHAT3}{SIMPLE_FUNCTION_JSON}"
+)
+SIMPLE_FUNCTION_OUTPUT_GIGACHAT31 = f"{TOOL_HEADER_GIGACHAT31}{SIMPLE_FUNCTION_JSON}"
 SIMPLE_FUNCTION_CALL = FunctionCall(
     name="manage_user_memory",
     arguments=json.dumps(SIMPLE_ARGS_DICT, ensure_ascii=False),
@@ -38,7 +48,12 @@ PARAMETERLESS_FUNCTION_JSON = json.dumps(
     },
     ensure_ascii=False,
 )
-PARAMETERLESS_FUNCTION_OUTPUT = "function call" + PARAMETERLESS_FUNCTION_JSON
+PARAMETERLESS_FUNCTION_OUTPUT_GIGACHAT3 = (
+    f"{MSG_SEP_TOKEN}{TOOL_HEADER_GIGACHAT3}{PARAMETERLESS_FUNCTION_JSON}"
+)
+PARAMETERLESS_FUNCTION_OUTPUT_GIGACHAT31 = (
+    f"{TOOL_HEADER_GIGACHAT31}{PARAMETERLESS_FUNCTION_JSON}"
+)
 PARAMETERLESS_FUNCTION_CALL = FunctionCall(
     name="manage_user_memory",
     arguments=json.dumps({}, ensure_ascii=False),
@@ -62,17 +77,38 @@ COMPLEX_FUNCTION_JSON = json.dumps(
     },
     ensure_ascii=False,
 )
-COMPLEX_FUNCTION_OUTPUT = "function call" + COMPLEX_FUNCTION_JSON
+COMPLEX_FUNCTION_OUTPUT_GIGACHAT3 = (
+    f"{MSG_SEP_TOKEN}{TOOL_HEADER_GIGACHAT3}{COMPLEX_FUNCTION_JSON}"
+)
+COMPLEX_FUNCTION_OUTPUT_GIGACHAT31 = f"{TOOL_HEADER_GIGACHAT31}{COMPLEX_FUNCTION_JSON}"
 COMPLEX_FUNCTION_CALL = FunctionCall(
     name="manage_user_memory",
     arguments=json.dumps(COMPLEX_ARGS_DICT, ensure_ascii=False),
 )
 
 
+CONTENT_TEXT = "I'll check that for you."
+MIXED_OUTPUT_GIGACHAT3 = f"{CONTENT_TEXT}{SIMPLE_FUNCTION_OUTPUT_GIGACHAT3}"
+MIXED_OUTPUT_GIGACHAT31 = f"{CONTENT_TEXT}{SIMPLE_FUNCTION_OUTPUT_GIGACHAT31}"
+
+
+@pytest.fixture(name="gigachat_tokenizer")
+def fixture_gigachat_tokenizer(default_tokenizer: TokenizerLike):
+    default_tokenizer.add_tokens(
+        [
+            MSG_SEP_TOKEN,
+            ROLE_SEP_TOKEN,
+            TOOL_HEADER_GIGACHAT31,
+            EOS_TOKEN,
+        ]
+    )
+    return default_tokenizer
+
+
 @pytest.mark.parametrize("streaming", [True, False])
-def test_no_tool_call(streaming: bool, default_tokenizer: TokenizerLike):
+def test_no_tool_call(streaming: bool, gigachat_tokenizer: TokenizerLike):
     tool_parser: ToolParser = ToolParserManager.get_tool_parser("gigachat3")(
-        default_tokenizer
+        gigachat_tokenizer
     )
     model_output = "How can I help you today?"
     content, tool_calls = run_tool_extraction(
@@ -85,45 +121,143 @@ def test_no_tool_call(streaming: bool, default_tokenizer: TokenizerLike):
 TEST_CASES = [
     pytest.param(
         True,
-        SIMPLE_FUNCTION_OUTPUT,
+        SIMPLE_FUNCTION_OUTPUT_GIGACHAT3,
         [SIMPLE_FUNCTION_CALL],
         None,
-        id="simple_streaming",
+        id="simple_streaming_gigachat3",
     ),
     pytest.param(
         False,
-        SIMPLE_FUNCTION_OUTPUT,
+        SIMPLE_FUNCTION_OUTPUT_GIGACHAT3,
         [SIMPLE_FUNCTION_CALL],
         None,
-        id="simple_nonstreaming",
+        id="simple_nonstreaming_gigachat3",
     ),
     pytest.param(
         True,
-        PARAMETERLESS_FUNCTION_OUTPUT,
+        PARAMETERLESS_FUNCTION_OUTPUT_GIGACHAT3,
         [PARAMETERLESS_FUNCTION_CALL],
         None,
-        id="parameterless_streaming",
+        id="parameterless_streaming_gigachat3",
     ),
     pytest.param(
         False,
-        PARAMETERLESS_FUNCTION_OUTPUT,
+        PARAMETERLESS_FUNCTION_OUTPUT_GIGACHAT3,
         [PARAMETERLESS_FUNCTION_CALL],
         None,
-        id="parameterless_nonstreaming",
+        id="parameterless_nonstreaming_gigachat3",
     ),
     pytest.param(
         True,
-        COMPLEX_FUNCTION_OUTPUT,
+        COMPLEX_FUNCTION_OUTPUT_GIGACHAT3,
         [COMPLEX_FUNCTION_CALL],
         None,
-        id="complex_streaming",
+        id="complex_streaming_gigachat3",
     ),
     pytest.param(
         False,
-        COMPLEX_FUNCTION_OUTPUT,
+        COMPLEX_FUNCTION_OUTPUT_GIGACHAT3,
         [COMPLEX_FUNCTION_CALL],
         None,
-        id="complex_nonstreaming",
+        id="complex_nonstreaming_gigachat3",
+    ),
+    pytest.param(
+        True,
+        MIXED_OUTPUT_GIGACHAT3,
+        [SIMPLE_FUNCTION_CALL],
+        CONTENT_TEXT,
+        id="mixed_content_streaming_gigachat3",
+    ),
+    pytest.param(
+        False,
+        MIXED_OUTPUT_GIGACHAT3,
+        [SIMPLE_FUNCTION_CALL],
+        CONTENT_TEXT,
+        id="mixed_content_nonstreaming_gigachat3",
+    ),
+    pytest.param(
+        True,
+        MIXED_OUTPUT_GIGACHAT3 + EOS_TOKEN,
+        [SIMPLE_FUNCTION_CALL],
+        CONTENT_TEXT,
+        id="mixed_content_streaming_with_eos_gigachat3",
+    ),
+    pytest.param(
+        False,
+        MIXED_OUTPUT_GIGACHAT3 + EOS_TOKEN,
+        [SIMPLE_FUNCTION_CALL],
+        CONTENT_TEXT,
+        id="mixed_content_nonstreaming_with_eos_gigachat3",
+    ),
+    pytest.param(
+        True,
+        SIMPLE_FUNCTION_OUTPUT_GIGACHAT31,
+        [SIMPLE_FUNCTION_CALL],
+        None,
+        id="simple_streaming_gigachat31",
+    ),
+    pytest.param(
+        False,
+        SIMPLE_FUNCTION_OUTPUT_GIGACHAT31,
+        [SIMPLE_FUNCTION_CALL],
+        None,
+        id="simple_nonstreaming_gigachat31",
+    ),
+    pytest.param(
+        True,
+        PARAMETERLESS_FUNCTION_OUTPUT_GIGACHAT31,
+        [PARAMETERLESS_FUNCTION_CALL],
+        None,
+        id="parameterless_streaming_gigachat31",
+    ),
+    pytest.param(
+        False,
+        PARAMETERLESS_FUNCTION_OUTPUT_GIGACHAT31,
+        [PARAMETERLESS_FUNCTION_CALL],
+        None,
+        id="parameterless_nonstreaming_gigachat31",
+    ),
+    pytest.param(
+        True,
+        COMPLEX_FUNCTION_OUTPUT_GIGACHAT31,
+        [COMPLEX_FUNCTION_CALL],
+        None,
+        id="complex_streaming_gigachat31",
+    ),
+    pytest.param(
+        False,
+        COMPLEX_FUNCTION_OUTPUT_GIGACHAT31,
+        [COMPLEX_FUNCTION_CALL],
+        None,
+        id="complex_nonstreaming_gigachat31",
+    ),
+    pytest.param(
+        True,
+        MIXED_OUTPUT_GIGACHAT31,
+        [SIMPLE_FUNCTION_CALL],
+        CONTENT_TEXT,
+        id="mixed_content_streaming_gigachat31",
+    ),
+    pytest.param(
+        False,
+        MIXED_OUTPUT_GIGACHAT31,
+        [SIMPLE_FUNCTION_CALL],
+        CONTENT_TEXT,
+        id="mixed_content_nonstreaming_gigachat31",
+    ),
+    pytest.param(
+        True,
+        MIXED_OUTPUT_GIGACHAT31 + EOS_TOKEN,
+        [SIMPLE_FUNCTION_CALL],
+        CONTENT_TEXT,
+        id="mixed_content_streaming_with_eos_gigachat31",
+    ),
+    pytest.param(
+        False,
+        MIXED_OUTPUT_GIGACHAT31 + EOS_TOKEN,
+        [SIMPLE_FUNCTION_CALL],
+        CONTENT_TEXT,
+        id="mixed_content_nonstreaming_with_eos_gigachat31",
     ),
 ]
 
@@ -136,14 +270,16 @@ def test_tool_call(
     model_output: str,
     expected_tool_calls: list[FunctionCall],
     expected_content: str | None,
-    default_tokenizer: TokenizerLike,
+    gigachat_tokenizer: TokenizerLike,
 ):
     tool_parser: ToolParser = ToolParserManager.get_tool_parser("gigachat3")(
-        default_tokenizer
+        gigachat_tokenizer
     )
     content, tool_calls = run_tool_extraction(
         tool_parser, model_output, streaming=streaming
     )
+    if content == "":
+        content = None
     assert content == expected_content
     assert len(tool_calls) == len(expected_tool_calls)
     for actual, expected in zip(tool_calls, expected_tool_calls):
@@ -154,15 +290,46 @@ def test_tool_call(
         assert actual_args == expected_args
 
 
-def test_streaming_tool_call_with_large_steps(default_tokenizer: TokenizerLike):
+@pytest.mark.parametrize(
+    "model_output_deltas",
+    [
+        pytest.param(
+            [
+                CONTENT_TEXT[:3],
+                CONTENT_TEXT[3:5],
+                CONTENT_TEXT[5:],
+                MSG_SEP_TOKEN,
+                TOOL_HEADER_GIGACHAT3,
+                COMPLEX_FUNCTION_JSON[:40],
+                COMPLEX_FUNCTION_JSON[40:-1],
+                COMPLEX_FUNCTION_JSON[-1],
+            ],
+            id="gigachat3",
+        ),
+        pytest.param(
+            [
+                CONTENT_TEXT[:3],
+                CONTENT_TEXT[3:5],
+                CONTENT_TEXT[5:],
+                TOOL_HEADER_GIGACHAT31,
+                COMPLEX_FUNCTION_JSON[:40],
+                COMPLEX_FUNCTION_JSON[40:-1],
+                COMPLEX_FUNCTION_JSON[-1],
+            ],
+            id="gigachat31",
+        ),
+    ],
+)
+def test_streaming_tool_call_with_large_steps(
+    model_output_deltas: list[str],
+    gigachat_tokenizer: TokenizerLike,
+):
+    """
+    Test that the closing braces are streamed correctly.
+    """
     tool_parser: ToolParser = ToolParserManager.get_tool_parser("gigachat3")(
-        default_tokenizer
+        gigachat_tokenizer
     )
-    model_output_deltas = [
-        "function call",
-        COMPLEX_FUNCTION_JSON[:40],
-        COMPLEX_FUNCTION_JSON[40:],
-    ]
     reconstructor = run_tool_extraction_streaming(
         tool_parser,
         model_output_deltas,

tests/entrypoints/openai/tool_parsers/test_hunyuan_a13b_tool_parser.py

@@ -7,7 +7,7 @@ from unittest.mock import MagicMock
 
 import pytest
 
-from tests.entrypoints.openai.tool_parsers.utils import (
+from tests.tool_parsers.utils import (
     run_tool_extraction,
     run_tool_extraction_streaming,
 )

tests/entrypoints/openai/tool_parsers/test_llama4_pythonic_tool_parser.py

@@ -5,7 +5,7 @@ from unittest.mock import MagicMock, patch
 
 import pytest
 
-from tests.entrypoints.openai.tool_parsers.utils import (
+from tests.tool_parsers.utils import (
     run_tool_extraction,
     run_tool_extraction_streaming,
 )

tests/entrypoints/openai/tool_parsers/test_olmo3_tool_parser.py

@@ -5,7 +5,7 @@ from unittest.mock import MagicMock, patch
 
 import pytest
 
-from tests.entrypoints.openai.tool_parsers.utils import (
+from tests.tool_parsers.utils import (
     run_tool_extraction,
     run_tool_extraction_streaming,
 )

tests/entrypoints/openai/tool_parsers/test_pythonic_tool_parser.py

@@ -5,7 +5,7 @@ from unittest.mock import MagicMock, patch
 
 import pytest
 
-from tests.entrypoints.openai.tool_parsers.utils import (
+from tests.tool_parsers.utils import (
     run_tool_extraction,
     run_tool_extraction_streaming,
 )

tests/entrypoints/pooling/embed/test_cohere_online.py

@@ -0,0 +1,310 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for the Cohere /v2/embed API with generic (non-Cohere) models.
+
+Validates that the Cohere v2 embed endpoint works correctly with standard
+embedding models, covering text embedding, embedding type conversions,
+response structure, batching, normalisation, and semantic similarity.
+"""
+
+import struct
+
+import numpy as np
+import pybase64 as base64
+import pytest
+import requests
+
+from tests.utils import RemoteOpenAIServer
+
+DTYPE = "bfloat16"
+
+MODELS: list[tuple[str, list[str]]] = [
+    ("intfloat/multilingual-e5-small", []),
+    (
+        "Snowflake/snowflake-arctic-embed-m-v1.5",
+        [
+            "--trust_remote_code",
+            "--hf_overrides",
+            '{"matryoshka_dimensions":[256]}',
+        ],
+    ),
+]
+
+
+@pytest.fixture(scope="module", params=MODELS, ids=lambda m: m[0])
+def model_config(request):
+    return request.param
+
+
+@pytest.fixture(scope="module")
+def model_name(model_config):
+    return model_config[0]
+
+
+@pytest.fixture(scope="module")
+def server(model_config):
+    name, extra_args = model_config
+    args = [
+        "--runner",
+        "pooling",
+        "--dtype",
+        DTYPE,
+        "--enforce-eager",
+        "--max-model-len",
+        "512",
+        "--gpu-memory-utilization",
+        "0.02",
+    ] + extra_args
+    with RemoteOpenAIServer(name, args) as remote_server:
+        yield remote_server
+
+
+def _cohere_embed(
+    server: RemoteOpenAIServer,
+    model_name: str,
+    texts: list[str] | None = None,
+    images: list[str] | None = None,
+    input_type: str | None = None,
+    embedding_types: list[str] | None = None,
+) -> dict:
+    body: dict = {"model": model_name}
+    if input_type is not None:
+        body["input_type"] = input_type
+    if texts is not None:
+        body["texts"] = texts
+    if images is not None:
+        body["images"] = images
+    if embedding_types is not None:
+        body["embedding_types"] = embedding_types
+    resp = requests.post(server.url_for("/v2/embed"), json=body)
+    resp.raise_for_status()
+    return resp.json()
+
+
+def _openai_embed(
+    server: RemoteOpenAIServer, model_name: str, texts: list[str]
+) -> dict:
+    body = {"model": model_name, "input": texts, "encoding_format": "float"}
+    resp = requests.post(server.url_for("/v1/embeddings"), json=body)
+    resp.raise_for_status()
+    return resp.json()
+
+
+def _cosine_sim(a: list[float], b: list[float]) -> float:
+    va, vb = np.array(a), np.array(b)
+    return float(np.dot(va, vb) / (np.linalg.norm(va) * np.linalg.norm(vb)))
+
+
+# -----------------------------------------------------------
+# Text embedding tests
+# -----------------------------------------------------------
+
+
+def test_basic_embed(server: RemoteOpenAIServer, model_name: str):
+    r = _cohere_embed(
+        server, model_name, texts=["hello world"], embedding_types=["float"]
+    )
+    assert "embeddings" in r
+    assert len(r["embeddings"]["float"]) == 1
+    assert len(r["embeddings"]["float"][0]) > 0
+
+
+def test_unsupported_input_type_rejected(server: RemoteOpenAIServer, model_name: str):
+    """An input_type not defined in the model's prompt config should be
+    rejected with a 400 error."""
+    body = {
+        "model": model_name,
+        "input_type": "nonexistent_type",
+        "texts": ["hello world"],
+        "embedding_types": ["float"],
+    }
+    resp = requests.post(server.url_for("/v2/embed"), json=body)
+    assert resp.status_code == 400
+    assert "Unsupported input_type" in resp.json()["error"]["message"]
+
+
+def test_omitted_input_type_accepted(server: RemoteOpenAIServer, model_name: str):
+    """Omitting input_type should always work (no prompt prefix applied)."""
+    body = {
+        "model": model_name,
+        "texts": ["hello world"],
+        "embedding_types": ["float"],
+    }
+    resp = requests.post(server.url_for("/v2/embed"), json=body)
+    assert resp.status_code == 200
+    data = resp.json()
+    assert len(data["embeddings"]["float"]) == 1
+
+
+def test_v1_v2_parity(server: RemoteOpenAIServer, model_name: str):
+    """v1 (OpenAI) and v2 (Cohere) endpoints should produce the same
+    float embeddings for a generic model."""
+    texts = ["hello world"]
+    v2 = _cohere_embed(server, model_name, texts=texts, embedding_types=["float"])
+    v1 = _openai_embed(server, model_name, texts)
+    cos = _cosine_sim(v2["embeddings"]["float"][0], v1["data"][0]["embedding"])
+    assert cos > 0.9999, f"v1/v2 parity failed, cosine={cos}"
+
+
+def test_embedding_types(server: RemoteOpenAIServer, model_name: str):
+    r = _cohere_embed(
+        server,
+        model_name,
+        texts=["test"],
+        embedding_types=["float", "binary", "ubinary"],
+    )
+    dim = len(r["embeddings"]["float"][0])
+    assert len(r["embeddings"]["binary"][0]) == dim // 8
+    assert len(r["embeddings"]["ubinary"][0]) == dim // 8
+
+
+def test_response_structure(server: RemoteOpenAIServer, model_name: str):
+    r = _cohere_embed(server, model_name, texts=["test"], embedding_types=["float"])
+    assert "id" in r
+    assert "embeddings" in r
+    assert "texts" in r
+    assert r["texts"] == ["test"]
+    assert "meta" in r
+    assert r["meta"]["api_version"]["version"] == "2"
+    assert "billed_units" in r["meta"]
+    assert r["meta"]["billed_units"]["input_tokens"] > 0
+    assert r["meta"]["billed_units"]["image_tokens"] == 0
+
+
+def test_batch(server: RemoteOpenAIServer, model_name: str):
+    texts = ["apple", "banana", "cherry"]
+    r = _cohere_embed(server, model_name, texts=texts, embedding_types=["float"])
+    assert len(r["embeddings"]["float"]) == 3
+    dim = len(r["embeddings"]["float"][0])
+    for emb in r["embeddings"]["float"]:
+        assert len(emb) == dim
+
+
+def test_l2_normalized(server: RemoteOpenAIServer, model_name: str):
+    r = _cohere_embed(
+        server, model_name, texts=["hello world"], embedding_types=["float"]
+    )
+    emb = np.array(r["embeddings"]["float"][0])
+    assert abs(float(np.linalg.norm(emb)) - 1.0) < 0.01
+
+
+def test_semantic_similarity(server: RemoteOpenAIServer, model_name: str):
+    r = _cohere_embed(
+        server,
+        model_name,
+        texts=["machine learning", "deep learning", "chocolate cake recipe"],
+        embedding_types=["float"],
+    )
+    embs = r["embeddings"]["float"]
+    cos_related = _cosine_sim(embs[0], embs[1])
+    cos_unrelated = _cosine_sim(embs[0], embs[2])
+    assert cos_related > cos_unrelated
+
+
+def test_missing_input_returns_error(server: RemoteOpenAIServer, model_name: str):
+    body = {"model": model_name}
+    resp = requests.post(server.url_for("/v2/embed"), json=body)
+    assert resp.status_code == 400
+
+
+def test_base64_embedding_type(server: RemoteOpenAIServer, model_name: str):
+    r = _cohere_embed(
+        server,
+        model_name,
+        texts=["test encoding"],
+        embedding_types=["float", "base64"],
+    )
+    float_emb = r["embeddings"]["float"][0]
+    b64_str = r["embeddings"]["base64"][0]
+    decoded = struct.unpack(f"<{len(float_emb)}f", base64.b64decode(b64_str))
+    np.testing.assert_allclose(float_emb, decoded, rtol=1e-5)
+
+
+# -----------------------------------------------------------
+# Truncation tests
+# -----------------------------------------------------------
+
+
+def _cohere_embed_raw(
+    server: RemoteOpenAIServer,
+    body: dict,
+) -> requests.Response:
+    return requests.post(server.url_for("/v2/embed"), json=body)
+
+
+def test_truncate_end_succeeds(server: RemoteOpenAIServer, model_name: str):
+    """truncate=END should silently truncate long input."""
+    long_text = " ".join(["word"] * 2000)
+    body = {
+        "model": model_name,
+        "texts": [long_text],
+        "embedding_types": ["float"],
+        "truncate": "END",
+    }
+    resp = _cohere_embed_raw(server, body)
+    assert resp.status_code == 200
+    data = resp.json()
+    assert len(data["embeddings"]["float"]) == 1
+
+
+def test_truncate_start_succeeds(server: RemoteOpenAIServer, model_name: str):
+    """truncate=START should silently truncate long input from the start."""
+    long_text = " ".join(["word"] * 2000)
+    body = {
+        "model": model_name,
+        "texts": [long_text],
+        "embedding_types": ["float"],
+        "truncate": "START",
+    }
+    resp = _cohere_embed_raw(server, body)
+    assert resp.status_code == 200
+    data = resp.json()
+    assert len(data["embeddings"]["float"]) == 1
+
+
+def test_truncate_none_rejects_long_input(server: RemoteOpenAIServer, model_name: str):
+    """truncate=NONE should error when input exceeds model context."""
+    long_text = " ".join(["word"] * 2000)
+    body = {
+        "model": model_name,
+        "texts": [long_text],
+        "embedding_types": ["float"],
+        "truncate": "NONE",
+    }
+    resp = _cohere_embed_raw(server, body)
+    assert resp.status_code == 400
+
+
+def test_truncate_start_vs_end_differ(server: RemoteOpenAIServer, model_name: str):
+    """START and END truncation should produce different embeddings
+    when the input is long enough to actually be truncated.
+
+    We construct input with distinct tokens at the start vs end
+    so that keeping different halves produces different embeddings.
+    """
+    start_words = " ".join([f"alpha{i}" for i in range(300)])
+    end_words = " ".join([f"omega{i}" for i in range(300)])
+    long_text = start_words + " " + end_words
+
+    body_end = {
+        "model": model_name,
+        "texts": [long_text],
+        "embedding_types": ["float"],
+        "truncate": "END",
+    }
+    body_start = {
+        "model": model_name,
+        "texts": [long_text],
+        "embedding_types": ["float"],
+        "truncate": "START",
+    }
+    r_end = _cohere_embed_raw(server, body_end).json()
+    r_start = _cohere_embed_raw(server, body_start).json()
+
+    emb_end = r_end["embeddings"]["float"][0]
+    emb_start = r_start["embeddings"]["float"][0]
+    cos = _cosine_sim(emb_end, emb_start)
+    assert cos < 0.99, (
+        f"START and END truncation should produce different embeddings "
+        f"for long input, but cosine similarity was {cos}"
+    )

tests/entrypoints/pooling/embed/test_cohere_online_vision.py

@@ -0,0 +1,135 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for the Cohere /v2/embed API with a multimodal model (SigLIP).
+
+Validates image embedding, batching, normalisation, and embedding type
+conversions through the /v2/embed endpoint.
+"""
+
+import struct
+import zlib
+
+import numpy as np
+import pybase64 as base64
+import pytest
+import requests
+
+from tests.utils import RemoteOpenAIServer
+
+MODEL_NAME = "google/siglip-so400m-patch14-384"
+DTYPE = "bfloat16"
+
+
+@pytest.fixture(scope="module")
+def server():
+    args = [
+        "--runner",
+        "pooling",
+        "--dtype",
+        DTYPE,
+        "--enforce-eager",
+        "--max-model-len",
+        "64",
+        "--gpu-memory-utilization",
+        "0.3",
+    ]
+    with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
+        yield remote_server
+
+
+def _make_tiny_png(r: int, g: int, b: int, w: int = 2, h: int = 2) -> str:
+    raw = b""
+    for _ in range(h):
+        raw += b"\x00" + bytes([r, g, b]) * w
+    compressed = zlib.compress(raw)
+
+    def chunk(ctype: bytes, cdata: bytes) -> bytes:
+        c = ctype + cdata
+        return (
+            struct.pack(">I", len(cdata))
+            + c
+            + struct.pack(">I", zlib.crc32(c) & 0xFFFFFFFF)
+        )
+
+    ihdr = struct.pack(">IIBBBBB", w, h, 8, 2, 0, 0, 0)
+    png = (
+        b"\x89PNG\r\n\x1a\n"
+        + chunk(b"IHDR", ihdr)
+        + chunk(b"IDAT", compressed)
+        + chunk(b"IEND", b"")
+    )
+    return "data:image/png;base64," + base64.b64encode(png).decode()
+
+
+def _cohere_embed(
+    server: RemoteOpenAIServer,
+    texts: list[str] | None = None,
+    images: list[str] | None = None,
+    embedding_types: list[str] | None = None,
+) -> dict:
+    body: dict = {"model": MODEL_NAME}
+    if texts is not None:
+        body["texts"] = texts
+    if images is not None:
+        body["images"] = images
+    if embedding_types is not None:
+        body["embedding_types"] = embedding_types
+    resp = requests.post(server.url_for("/v2/embed"), json=body)
+    resp.raise_for_status()
+    return resp.json()
+
+
+def test_image_embed(server: RemoteOpenAIServer):
+    img_uri = _make_tiny_png(255, 0, 0)
+    r = _cohere_embed(
+        server,
+        images=[img_uri],
+        embedding_types=["float"],
+    )
+    assert "embeddings" in r
+    assert len(r["embeddings"]["float"]) == 1
+    assert len(r["embeddings"]["float"][0]) > 0
+    assert r["meta"]["billed_units"]["image_tokens"] > 0
+    assert r["meta"]["billed_units"]["input_tokens"] == 0
+
+
+def test_image_batch(server: RemoteOpenAIServer):
+    red = _make_tiny_png(255, 0, 0)
+    blue = _make_tiny_png(0, 0, 255)
+    r = _cohere_embed(
+        server,
+        images=[red, blue],
+        embedding_types=["float"],
+    )
+    assert len(r["embeddings"]["float"]) == 2
+
+
+def test_image_l2_normalized(server: RemoteOpenAIServer):
+    img_uri = _make_tiny_png(0, 255, 0)
+    r = _cohere_embed(
+        server,
+        images=[img_uri],
+        embedding_types=["float"],
+    )
+    emb = np.array(r["embeddings"]["float"][0])
+    assert abs(float(np.linalg.norm(emb)) - 1.0) < 0.01
+
+
+def test_image_embedding_types(server: RemoteOpenAIServer):
+    img_uri = _make_tiny_png(128, 128, 128)
+    r = _cohere_embed(
+        server,
+        images=[img_uri],
+        embedding_types=["float", "binary", "ubinary"],
+    )
+    dim = len(r["embeddings"]["float"][0])
+    assert len(r["embeddings"]["binary"][0]) == dim // 8
+    assert len(r["embeddings"]["ubinary"][0]) == dim // 8
+
+
+def test_text_embed_on_multimodal(server: RemoteOpenAIServer):
+    """SigLIP also supports text-only embedding via /v2/embed."""
+    r = _cohere_embed(server, texts=["hello world"], embedding_types=["float"])
+    assert "embeddings" in r
+    assert len(r["embeddings"]["float"]) == 1
+    assert len(r["embeddings"]["float"][0]) > 0

tests/entrypoints/pooling/embed/test_cohere_openai_parity.py

@@ -0,0 +1,102 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Parity test between Cohere /v2/embed and OpenAI /v1/embeddings.
+
+Verifies that both endpoints produce identical float embeddings when
+no prompt prefix is applied (input_type omitted for Cohere /v2/embed).
+"""
+
+import numpy as np
+import pytest
+import requests
+
+from tests.utils import RemoteOpenAIServer
+
+MODEL_NAME = "BAAI/bge-base-en-v1.5"
+DTYPE = "bfloat16"
+
+
+@pytest.fixture(scope="module")
+def server():
+    args = [
+        "--runner",
+        "pooling",
+        "--dtype",
+        DTYPE,
+        "--enforce-eager",
+        "--max-model-len",
+        "512",
+        "--gpu-memory-utilization",
+        "0.02",
+    ]
+    with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
+        yield remote_server
+
+
+def _cohere_embed(
+    server: RemoteOpenAIServer,
+    texts: list[str],
+) -> list[list[float]]:
+    body = {
+        "model": MODEL_NAME,
+        "texts": texts,
+        "embedding_types": ["float"],
+    }
+    resp = requests.post(server.url_for("/v2/embed"), json=body)
+    resp.raise_for_status()
+    return resp.json()["embeddings"]["float"]
+
+
+def _openai_embed(
+    server: RemoteOpenAIServer,
+    texts: list[str],
+) -> list[list[float]]:
+    body = {"model": MODEL_NAME, "input": texts, "encoding_format": "float"}
+    resp = requests.post(server.url_for("/v1/embeddings"), json=body)
+    resp.raise_for_status()
+    return [item["embedding"] for item in resp.json()["data"]]
+
+
+def test_single_text_parity(server: RemoteOpenAIServer):
+    """A single text should produce identical embeddings via both APIs."""
+    texts = ["the quick brown fox jumps over the lazy dog"]
+    v2 = _cohere_embed(server, texts)
+    v1 = _openai_embed(server, texts)
+    np.testing.assert_allclose(v2[0], v1[0], rtol=1e-5)
+
+
+def test_batch_parity(server: RemoteOpenAIServer):
+    """A batch of texts should produce identical embeddings via both APIs,
+    in the same order."""
+    texts = [
+        "machine learning",
+        "deep learning",
+        "natural language processing",
+    ]
+    v2 = _cohere_embed(server, texts)
+    v1 = _openai_embed(server, texts)
+    assert len(v2) == len(v1) == 3
+    for i in range(3):
+        np.testing.assert_allclose(v2[i], v1[i], rtol=1e-5, err_msg=f"index {i}")
+
+
+def test_token_count_parity(server: RemoteOpenAIServer):
+    """Both APIs should report the same prompt token count."""
+    texts = ["hello world"]
+    v2_resp = requests.post(
+        server.url_for("/v2/embed"),
+        json={
+            "model": MODEL_NAME,
+            "texts": texts,
+            "embedding_types": ["float"],
+        },
+    )
+    v1_resp = requests.post(
+        server.url_for("/v1/embeddings"),
+        json={"model": MODEL_NAME, "input": texts, "encoding_format": "float"},
+    )
+    v2_resp.raise_for_status()
+    v1_resp.raise_for_status()
+    v2_tokens = v2_resp.json()["meta"]["billed_units"]["input_tokens"]
+    v1_tokens = v1_resp.json()["usage"]["prompt_tokens"]
+    assert v2_tokens == v1_tokens

tests/entrypoints/pooling/embed/test_io_processor.py

@@ -0,0 +1,208 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for EmbedIOProcessor."""
+
+import pytest
+
+from vllm.entrypoints.pooling.embed.io_processor import EmbedIOProcessor
+from vllm.entrypoints.pooling.embed.protocol import (
+    CohereEmbedRequest,
+)
+
+
+class TestResolveTruncation:
+    """Unit tests for EmbedIOProcessor._resolve_cohere_truncation."""
+
+    @staticmethod
+    def _make_request(**kwargs) -> CohereEmbedRequest:
+        defaults = {
+            "model": "test",
+            "input_type": "search_document",
+            "texts": ["hello"],
+        }
+        return CohereEmbedRequest(**(defaults | kwargs))
+
+    def test_truncate_end_default(self):
+        req = self._make_request()
+        tokens, side = EmbedIOProcessor._resolve_cohere_truncation(req)
+        assert tokens == -1
+        assert side is None
+
+    def test_truncate_end_explicit(self):
+        req = self._make_request(truncate="END")
+        tokens, side = EmbedIOProcessor._resolve_cohere_truncation(req)
+        assert tokens == -1
+        assert side is None
+
+    def test_truncate_end_with_max_tokens(self):
+        req = self._make_request(truncate="END", max_tokens=128)
+        tokens, side = EmbedIOProcessor._resolve_cohere_truncation(req)
+        assert tokens == 128
+        assert side is None
+
+    def test_truncate_none(self):
+        req = self._make_request(truncate="NONE")
+        tokens, side = EmbedIOProcessor._resolve_cohere_truncation(req)
+        assert tokens is None
+        assert side is None
+
+    def test_truncate_none_with_max_tokens(self):
+        """truncate=NONE should NOT set truncate_prompt_tokens; the
+        max_tokens limit is enforced separately via _check_max_tokens."""
+        req = self._make_request(truncate="NONE", max_tokens=10)
+        tokens, side = EmbedIOProcessor._resolve_cohere_truncation(req)
+        assert tokens is None
+        assert side is None
+
+    def test_truncate_start(self):
+        req = self._make_request(truncate="START")
+        tokens, side = EmbedIOProcessor._resolve_cohere_truncation(req)
+        assert tokens == -1
+        assert side == "left"
+
+    def test_truncate_start_with_max_tokens(self):
+        req = self._make_request(truncate="START", max_tokens=64)
+        tokens, side = EmbedIOProcessor._resolve_cohere_truncation(req)
+        assert tokens == 64
+        assert side == "left"
+
+
+class TestApplyStPrompt:
+    """Unit tests for EmbedIOProcessor._apply_task_instruction."""
+
+    @staticmethod
+    def _make_handler(task_instructions: dict[str, str] | None):
+        handler = object.__new__(EmbedIOProcessor)
+        handler.task_instructions = task_instructions
+        return handler
+
+    def test_no_prompts_configured(self):
+        handler = self._make_handler(None)
+        texts = ["hello", "world"]
+        assert handler._apply_task_instruction(texts, "query") is texts
+
+    def test_matching_input_type(self):
+        handler = self._make_handler({"query": "search_query: "})
+        result = handler._apply_task_instruction(["hello"], "query")
+        assert result == ["search_query: hello"]
+
+    def test_non_matching_input_type(self):
+        handler = self._make_handler({"query": "search_query: "})
+        texts = ["hello"]
+        assert handler._apply_task_instruction(texts, "document") is texts
+
+    def test_multiple_texts(self):
+        handler = self._make_handler(
+            {"query": "Represent this sentence for searching: "}
+        )
+        result = handler._apply_task_instruction(["a", "b", "c"], "query")
+        assert result == [
+            "Represent this sentence for searching: a",
+            "Represent this sentence for searching: b",
+            "Represent this sentence for searching: c",
+        ]
+
+    def test_empty_prefix_returns_unchanged(self):
+        handler = self._make_handler({"passage": ""})
+        texts = ["hello"]
+        assert handler._apply_task_instruction(texts, "passage") is texts
+
+
+class TestLoadTaskInstructions:
+    """Unit tests for EmbedIOProcessor._load_task_instructions."""
+
+    def test_no_attribute(self):
+        class FakeConfig:
+            pass
+
+        assert EmbedIOProcessor._load_task_instructions(FakeConfig()) is None
+
+    def test_with_task_instructions(self):
+        class FakeConfig:
+            task_instructions = {
+                "retrieval.query": "Represent the query: ",
+                "retrieval.passage": "",
+            }
+
+        result = EmbedIOProcessor._load_task_instructions(FakeConfig())
+        assert result == {
+            "retrieval.query": "Represent the query: ",
+            "retrieval.passage": "",
+        }
+
+    def test_empty_dict(self):
+        class FakeConfig:
+            task_instructions = {}
+
+        assert EmbedIOProcessor._load_task_instructions(FakeConfig()) is None
+
+    def test_non_dict(self):
+        class FakeConfig:
+            task_instructions = "not a dict"
+
+        assert EmbedIOProcessor._load_task_instructions(FakeConfig()) is None
+
+
+class TestCheckMaxTokens:
+    """Unit tests for EmbedIOProcessor._check_cohere_max_tokens."""
+
+    @staticmethod
+    def _fake_output(n_tokens: int):
+        class _Out:
+            def __init__(self, n: int):
+                self.prompt_token_ids = list(range(n))
+
+        return _Out(n_tokens)
+
+    def test_none_check_is_noop(self):
+        outs = [self._fake_output(100)]
+        EmbedIOProcessor._check_cohere_max_tokens(outs, None)
+
+    def test_within_limit(self):
+        outs = [self._fake_output(5), self._fake_output(3)]
+        EmbedIOProcessor._check_cohere_max_tokens(outs, 5)
+
+    def test_exceeds_limit(self):
+        outs = [self._fake_output(3), self._fake_output(10)]
+        with pytest.raises(ValueError, match="exceeds max_tokens=5"):
+            EmbedIOProcessor._check_cohere_max_tokens(outs, 5)
+
+    def test_exact_limit(self):
+        outs = [self._fake_output(5)]
+        EmbedIOProcessor._check_cohere_max_tokens(outs, 5)
+
+
+class TestValidateInputType:
+    """Unit tests for EmbedIOProcessor._validate_input_type."""
+
+    @staticmethod
+    def _make_handler(task_instructions: dict[str, str] | None):
+        handler = object.__new__(EmbedIOProcessor)
+        handler.task_instructions = task_instructions
+        return handler
+
+    def test_none_input_type_always_accepted(self):
+        handler = self._make_handler(None)
+        handler._validate_input_type(None)
+        handler_with = self._make_handler({"query": "q: "})
+        handler_with._validate_input_type(None)
+
+    def test_no_prompts_rejects(self):
+        handler = self._make_handler(None)
+        with pytest.raises(ValueError, match="does not define any input_type"):
+            handler._validate_input_type("anything")
+
+    def test_known_type_accepted(self):
+        handler = self._make_handler({"query": "q: ", "document": "d: "})
+        handler._validate_input_type("query")
+        handler._validate_input_type("document")
+
+    def test_unknown_type_rejected(self):
+        handler = self._make_handler({"query": "q: ", "document": "d: "})
+        with pytest.raises(ValueError, match="Unsupported input_type 'other'"):
+            handler._validate_input_type("other")
+
+    def test_error_lists_supported(self):
+        handler = self._make_handler({"a": "", "b": ""})
+        with pytest.raises(ValueError, match="Supported values: a, b"):
+            handler._validate_input_type("z")

tests/entrypoints/pooling/embed/test_online.py

@@ -1,11 +1,11 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-import base64
 import json
 
 import numpy as np
 import openai
+import pybase64 as base64
 import pytest
 import pytest_asyncio
 import requests

tests/entrypoints/pooling/embed/test_protocol.py

@@ -0,0 +1,129 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for Cohere embed protocol: build_typed_embeddings and its
+underlying packing helpers, plus Cohere-specific serving helpers."""
+
+import struct
+
+import numpy as np
+import pybase64 as base64
+import pytest
+
+from vllm.entrypoints.pooling.embed.protocol import (
+    build_typed_embeddings,
+)
+
+
+@pytest.fixture
+def sample_embeddings() -> list[list[float]]:
+    return [
+        [0.1, -0.2, 0.3, -0.4, 0.5, -0.6, 0.7, -0.8],
+        [-0.05, 0.15, -0.25, 0.35, -0.45, 0.55, -0.65, 0.75],
+    ]
+
+
+class TestBuildTypedEmbeddingsFloat:
+    def test_float_passthrough(self, sample_embeddings: list[list[float]]):
+        result = build_typed_embeddings(sample_embeddings, ["float"])
+        assert result.float == sample_embeddings
+        assert result.binary is None
+
+    def test_empty_input(self):
+        result = build_typed_embeddings([], ["float"])
+        assert result.float == []
+
+
+class TestBuildTypedEmbeddingsBinary:
+    def test_binary_packing(self):
+        # 8 values: positive->1, negative->0 => bits: 10101010 = 0xAA = 170
+        # signed: 170 - 128 = 42
+        embs = [[1.0, -1.0, 1.0, -1.0, 1.0, -1.0, 1.0, -1.0]]
+        result = build_typed_embeddings(embs, ["binary"])
+        assert result.binary is not None
+        assert result.binary[0] == [42]
+
+    def test_ubinary_packing(self):
+        embs = [[1.0, -1.0, 1.0, -1.0, 1.0, -1.0, 1.0, -1.0]]
+        result = build_typed_embeddings(embs, ["ubinary"])
+        assert result.ubinary is not None
+        assert result.ubinary[0] == [170]  # 0b10101010
+
+    def test_binary_all_positive(self):
+        embs = [[0.1] * 8]
+        result = build_typed_embeddings(embs, ["binary"])
+        assert result.binary is not None
+        # all bits = 1 => 0xFF = 255, signed: 255 - 128 = 127
+        assert result.binary[0] == [127]
+
+    def test_binary_all_negative(self):
+        embs = [[-0.1] * 8]
+        result = build_typed_embeddings(embs, ["binary"])
+        assert result.binary is not None
+        # all bits = 0, signed: 0 - 128 = -128
+        assert result.binary[0] == [-128]
+
+    def test_binary_dimension_is_eighth(self, sample_embeddings: list[list[float]]):
+        result = build_typed_embeddings(sample_embeddings, ["binary"])
+        assert result.binary is not None
+        for orig, packed in zip(sample_embeddings, result.binary):
+            assert len(packed) == len(orig) // 8
+
+    def test_zero_treated_as_positive(self):
+        embs = [[0.0] * 8]
+        result = build_typed_embeddings(embs, ["binary"])
+        assert result.binary is not None
+        # 0.0 >= 0 is True, so bit=1 for all => 127 (signed)
+        assert result.binary[0] == [127]
+
+    def test_non_multiple_of_8_raises(self):
+        embs = [[0.1] * 7]
+        with pytest.raises(ValueError, match="multiple of 8"):
+            build_typed_embeddings(embs, ["binary"])
+
+    def test_ubinary_non_multiple_of_8_raises(self):
+        embs = [[0.1] * 10]
+        with pytest.raises(ValueError, match="multiple of 8"):
+            build_typed_embeddings(embs, ["ubinary"])
+
+
+class TestBuildTypedEmbeddingsBase64:
+    def test_base64_roundtrip(self, sample_embeddings: list[list[float]]):
+        result = build_typed_embeddings(sample_embeddings, ["base64"])
+        assert result.base64 is not None
+        assert len(result.base64) == 2
+
+        for orig, b64_str in zip(sample_embeddings, result.base64):
+            decoded = base64.b64decode(b64_str)
+            n = len(orig)
+            values = struct.unpack(f"<{n}f", decoded)
+            np.testing.assert_allclose(orig, values, rtol=1e-5)
+
+    def test_base64_byte_length(self):
+        embs = [[0.1, 0.2, 0.3]]
+        result = build_typed_embeddings(embs, ["base64"])
+        assert result.base64 is not None
+        raw = base64.b64decode(result.base64[0])
+        assert len(raw) == 3 * 4  # 3 floats * 4 bytes each
+
+
+class TestBuildTypedEmbeddingsMultiple:
+    def test_all_types_at_once(self, sample_embeddings: list[list[float]]):
+        result = build_typed_embeddings(
+            sample_embeddings,
+            ["float", "binary", "ubinary", "base64"],
+        )
+        assert result.float is not None
+        assert result.binary is not None
+        assert result.ubinary is not None
+        assert result.base64 is not None
+
+    def test_subset_types(self, sample_embeddings: list[list[float]]):
+        result = build_typed_embeddings(sample_embeddings, ["float", "binary"])
+        assert result.float is not None
+        assert result.binary is not None
+        assert result.ubinary is None
+        assert result.base64 is None
+
+    def test_unknown_type_ignored(self, sample_embeddings: list[list[float]]):
+        result = build_typed_embeddings(sample_embeddings, ["float", "unknown_type"])
+        assert result.float is not None

tests/entrypoints/pooling/pooling/test_online.py

@@ -1,10 +1,10 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-import base64
 import json
 
 import numpy as np
+import pybase64 as base64
 import pytest
 import requests
 import torch

tests/kernels/attention/test_cpu_attn.py

@@ -48,7 +48,7 @@ def get_attn_isa(
     else:
         if current_platform.get_cpu_architecture() == CpuArchEnum.ARM:
             return "neon"
-        elif torch._C._cpu._is_amx_tile_supported():
+        elif torch.cpu._is_amx_tile_supported():
             return "amx"
         else:
             return "vec"
@@ -400,9 +400,7 @@ def test_varlen_with_paged_kv_normal_vec(
 @pytest.mark.parametrize("use_alibi", [False])
 @pytest.mark.parametrize("use_sink", [False])
 @pytest.mark.parametrize("isa", ["amx"])
-@pytest.mark.skipif(
-    not torch._C._cpu._is_amx_tile_supported(), reason="no AMX support."
-)
+@pytest.mark.skipif(not torch.cpu._is_amx_tile_supported(), reason="no AMX support.")
 def test_varlen_with_paged_kv_normal_amx(
     seq_lens: list[tuple[int, int]],
     num_heads: tuple[int, int],

tests/kernels/attention/test_trtllm_kvfp8_dequant.py

@@ -0,0 +1,434 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Standalone unit tests for trtllm_prefill_attn_kvfp8_dequant.
+
+Tests both contiguous and non-contiguous (cross-layer unified) KV cache
+layouts against a pure-PyTorch reference implementation.
+"""
+
+import pytest
+import torch
+
+from vllm.platforms import current_platform
+
+FP8_DTYPE = current_platform.fp8_dtype()
+
+NUM_BLOCKS = 128
+
+
+def to_float8(x, dtype=None):
+    if dtype is None:
+        dtype = FP8_DTYPE
+    finfo = torch.finfo(dtype)
+    min_val, max_val = x.aminmax()
+    amax = torch.maximum(min_val.abs(), max_val.abs()).clamp(min=1e-12)
+    scale = finfo.max / amax * 0.1
+    x_scl_sat = (x * scale).clamp(min=finfo.min, max=finfo.max)
+    return x_scl_sat.to(dtype), scale.float().reciprocal()
+
+
+def make_contiguous_kv_cache(num_blocks, num_kv_heads, block_size, head_size):
+    """Create a standard contiguous fp8 KV cache (HND layout)."""
+    raw = torch.randn(
+        num_blocks,
+        2,
+        num_kv_heads,
+        block_size,
+        head_size,
+        dtype=torch.bfloat16,
+        device="cuda",
+    )
+    kv_cache, scale = to_float8(raw)
+    return kv_cache, scale
+
+
+def make_cross_layer_kv_cache(
+    num_blocks,
+    num_kv_heads,
+    block_size,
+    head_size,
+    num_layers=4,
+):
+    """
+    Create a non-contiguous per-layer view mimicking cross-layer allocation.
+
+    Physical layout: (num_blocks, 2, num_kv_heads, num_layers, block_size, head_size)
+    Returned view:   (num_blocks, 2, num_kv_heads, block_size, head_size)
+    with non-contiguous strides on dims 0, 1, 2 (they skip over num_layers).
+    """
+    raw = torch.randn(
+        num_blocks,
+        2,
+        num_kv_heads,
+        num_layers,
+        block_size,
+        head_size,
+        dtype=torch.bfloat16,
+        device="cuda",
+    )
+    fp8_full, scale = to_float8(raw)
+    layer_view = fp8_full[:, :, :, 0, :, :]
+    assert not layer_view.is_contiguous(), (
+        f"Expected non-contiguous view, got strides {layer_view.stride()}"
+    )
+    return layer_view, scale
+
+
+def ref_dequant(kv_cache, block_tables, k_scale, v_scale, dequant_dtype):
+    """Pure PyTorch reference: gather pages and dequantize fp8 -> dequant_dtype."""
+    batch_size, num_pages_per_seq = block_tables.shape
+    s = kv_cache.shape
+    out = torch.zeros(
+        batch_size * num_pages_per_seq + 1,
+        s[1],
+        s[2],
+        s[3],
+        s[4],
+        dtype=dequant_dtype,
+        device=kv_cache.device,
+    )
+    for b in range(batch_size):
+        for p in range(num_pages_per_seq):
+            page_idx = block_tables[b, p].item()
+            if page_idx <= 0:
+                continue
+            mock_idx = b * num_pages_per_seq + p + 1
+            out[mock_idx, 0] = (kv_cache[page_idx, 0].float() * k_scale.item()).to(
+                dequant_dtype
+            )
+            out[mock_idx, 1] = (kv_cache[page_idx, 1].float() * v_scale.item()).to(
+                dequant_dtype
+            )
+    return out
+
+
+@pytest.mark.parametrize("num_kv_heads", [1, 8])
+@pytest.mark.parametrize("head_size", [64, 128])
+@pytest.mark.parametrize("block_size", [16, 32])
+@pytest.mark.parametrize("batch_size", [1, 4])
+@pytest.mark.parametrize("num_pages_per_seq", [3, 8])
+@pytest.mark.parametrize("contiguous", [True, False])
+@torch.inference_mode()
+def test_trtllm_kvfp8_dequant(
+    num_kv_heads: int,
+    head_size: int,
+    block_size: int,
+    batch_size: int,
+    num_pages_per_seq: int,
+    contiguous: bool,
+):
+    from vllm.v1.attention.backends.flashinfer import (
+        trtllm_prefill_attn_kvfp8_dequant,
+    )
+
+    torch.set_default_device("cuda")
+
+    if contiguous:
+        kv_cache, scale = make_contiguous_kv_cache(
+            NUM_BLOCKS,
+            num_kv_heads,
+            block_size,
+            head_size,
+        )
+    else:
+        kv_cache, scale = make_cross_layer_kv_cache(
+            NUM_BLOCKS,
+            num_kv_heads,
+            block_size,
+            head_size,
+        )
+
+    k_scale = scale.clone()
+    v_scale = scale.clone()
+
+    block_tables = torch.randint(
+        1,
+        NUM_BLOCKS,
+        (batch_size, num_pages_per_seq),
+        dtype=torch.int32,
+    )
+
+    mock_kv_cache, mock_block_table = trtllm_prefill_attn_kvfp8_dequant(
+        kv_cache,
+        block_tables,
+        k_scale,
+        v_scale,
+        torch.bfloat16,
+    )
+
+    ref = ref_dequant(kv_cache, block_tables, k_scale, v_scale, torch.bfloat16)
+
+    expected_bt = torch.arange(
+        1,
+        batch_size * num_pages_per_seq + 1,
+        dtype=torch.int32,
+        device="cuda",
+    ).reshape(batch_size, num_pages_per_seq)
+    torch.testing.assert_close(mock_block_table, expected_bt)
+
+    # Page 0 is padding (never written), compare only pages 1+
+    torch.testing.assert_close(mock_kv_cache[1:], ref[1:], atol=1e-3, rtol=1e-3)
+
+
+@torch.inference_mode()
+def test_block_tables_with_zero_pages():
+    """Pages with index <= 0 must be skipped (early return in kernel)."""
+    from vllm.v1.attention.backends.flashinfer import (
+        trtllm_prefill_attn_kvfp8_dequant,
+    )
+
+    torch.set_default_device("cuda")
+    num_kv_heads, block_size, head_size = 8, 16, 64
+
+    kv_cache, scale = make_contiguous_kv_cache(
+        NUM_BLOCKS,
+        num_kv_heads,
+        block_size,
+        head_size,
+    )
+    k_scale = v_scale = scale.clone()
+
+    # Mix of valid pages and zeros (padding)
+    block_tables = torch.tensor(
+        [[5, 0, 10], [0, 0, 0], [3, 7, 0]],
+        dtype=torch.int32,
+        device="cuda",
+    )
+
+    mock_kv_cache, _ = trtllm_prefill_attn_kvfp8_dequant(
+        kv_cache,
+        block_tables,
+        k_scale,
+        v_scale,
+        torch.bfloat16,
+    )
+    ref = ref_dequant(kv_cache, block_tables, k_scale, v_scale, torch.bfloat16)
+
+    # Only compare pages that were actually written (non-zero page indices)
+    for b in range(block_tables.shape[0]):
+        for p in range(block_tables.shape[1]):
+            if block_tables[b, p].item() > 0:
+                idx = b * block_tables.shape[1] + p + 1
+                torch.testing.assert_close(
+                    mock_kv_cache[idx],
+                    ref[idx],
+                    atol=1e-3,
+                    rtol=1e-3,
+                )
+
+
+@torch.inference_mode()
+def test_all_zero_block_tables():
+    """All-zero block_tables: kernel should write nothing."""
+    from vllm.v1.attention.backends.flashinfer import (
+        trtllm_prefill_attn_kvfp8_dequant,
+    )
+
+    torch.set_default_device("cuda")
+    num_kv_heads, block_size, head_size = 4, 16, 64
+
+    kv_cache, scale = make_contiguous_kv_cache(
+        NUM_BLOCKS,
+        num_kv_heads,
+        block_size,
+        head_size,
+    )
+    k_scale = v_scale = scale.clone()
+
+    block_tables = torch.zeros(2, 4, dtype=torch.int32, device="cuda")
+
+    # Should not crash even though no pages are valid
+    mock_kv_cache, mock_block_table = trtllm_prefill_attn_kvfp8_dequant(
+        kv_cache,
+        block_tables,
+        k_scale,
+        v_scale,
+        torch.bfloat16,
+    )
+    assert mock_kv_cache.shape[0] == 2 * 4 + 1
+    assert mock_block_table.shape == (2, 4)
+
+
+@torch.inference_mode()
+def test_different_k_v_scales():
+    """Verify K and V are dequantized with independent scales."""
+    from vllm.v1.attention.backends.flashinfer import (
+        trtllm_prefill_attn_kvfp8_dequant,
+    )
+
+    torch.set_default_device("cuda")
+    num_kv_heads, block_size, head_size = 8, 16, 64
+
+    kv_cache, _ = make_contiguous_kv_cache(
+        NUM_BLOCKS,
+        num_kv_heads,
+        block_size,
+        head_size,
+    )
+    k_scale = torch.tensor([0.5], dtype=torch.float32, device="cuda")
+    v_scale = torch.tensor([2.0], dtype=torch.float32, device="cuda")
+
+    block_tables = torch.tensor([[1, 2]], dtype=torch.int32, device="cuda")
+
+    mock_kv_cache, _ = trtllm_prefill_attn_kvfp8_dequant(
+        kv_cache,
+        block_tables,
+        k_scale,
+        v_scale,
+        torch.bfloat16,
+    )
+    ref = ref_dequant(kv_cache, block_tables, k_scale, v_scale, torch.bfloat16)
+
+    torch.testing.assert_close(mock_kv_cache[1:], ref[1:], atol=1e-3, rtol=1e-3)
+
+
+@torch.inference_mode()
+def test_single_page_per_seq():
+    """Minimum grid dim 1 = 1 page per sequence."""
+    from vllm.v1.attention.backends.flashinfer import (
+        trtllm_prefill_attn_kvfp8_dequant,
+    )
+
+    torch.set_default_device("cuda")
+    num_kv_heads, block_size, head_size = 8, 16, 128
+
+    kv_cache, scale = make_contiguous_kv_cache(
+        NUM_BLOCKS,
+        num_kv_heads,
+        block_size,
+        head_size,
+    )
+    k_scale = v_scale = scale.clone()
+
+    block_tables = torch.tensor([[5], [10], [20]], dtype=torch.int32, device="cuda")
+
+    mock_kv_cache, _ = trtllm_prefill_attn_kvfp8_dequant(
+        kv_cache,
+        block_tables,
+        k_scale,
+        v_scale,
+        torch.bfloat16,
+    )
+    ref = ref_dequant(kv_cache, block_tables, k_scale, v_scale, torch.bfloat16)
+
+    torch.testing.assert_close(mock_kv_cache[1:], ref[1:], atol=1e-3, rtol=1e-3)
+
+
+@torch.inference_mode()
+def test_large_page_indices():
+    """Page indices near the top of the buffer stress offset arithmetic."""
+    from vllm.v1.attention.backends.flashinfer import (
+        trtllm_prefill_attn_kvfp8_dequant,
+    )
+
+    torch.set_default_device("cuda")
+    num_kv_heads, block_size, head_size = 8, 16, 128
+    large_num_blocks = 32768
+
+    kv_cache, scale = make_contiguous_kv_cache(
+        large_num_blocks,
+        num_kv_heads,
+        block_size,
+        head_size,
+    )
+    k_scale = v_scale = scale.clone()
+
+    # Use page indices near the top of the buffer
+    block_tables = torch.tensor(
+        [[large_num_blocks - 1, large_num_blocks - 2, 1]],
+        dtype=torch.int32,
+        device="cuda",
+    )
+
+    mock_kv_cache, _ = trtllm_prefill_attn_kvfp8_dequant(
+        kv_cache,
+        block_tables,
+        k_scale,
+        v_scale,
+        torch.bfloat16,
+    )
+    ref = ref_dequant(kv_cache, block_tables, k_scale, v_scale, torch.bfloat16)
+
+    torch.testing.assert_close(mock_kv_cache[1:], ref[1:], atol=1e-3, rtol=1e-3)
+
+
+@torch.inference_mode()
+def test_large_block_size():
+    """block_size=64 -> HEAD_STRIDE=8192, large tl.arange per thread block."""
+    from vllm.v1.attention.backends.flashinfer import (
+        trtllm_prefill_attn_kvfp8_dequant,
+    )
+
+    torch.set_default_device("cuda")
+    num_kv_heads, block_size, head_size = 4, 64, 128
+
+    kv_cache, scale = make_contiguous_kv_cache(
+        NUM_BLOCKS,
+        num_kv_heads,
+        block_size,
+        head_size,
+    )
+    k_scale = v_scale = scale.clone()
+
+    block_tables = torch.randint(
+        1,
+        NUM_BLOCKS,
+        (2, 4),
+        dtype=torch.int32,
+        device="cuda",
+    )
+
+    mock_kv_cache, _ = trtllm_prefill_attn_kvfp8_dequant(
+        kv_cache,
+        block_tables,
+        k_scale,
+        v_scale,
+        torch.bfloat16,
+    )
+    ref = ref_dequant(kv_cache, block_tables, k_scale, v_scale, torch.bfloat16)
+
+    torch.testing.assert_close(mock_kv_cache[1:], ref[1:], atol=1e-3, rtol=1e-3)
+
+
+@torch.inference_mode()
+def test_cross_layer_many_layers():
+    """
+    Non-contiguous with 36 layers -- matches real gpt-oss-120b.
+    Strides are far from contiguous (factor of 36 in the gaps).
+    """
+    from vllm.v1.attention.backends.flashinfer import (
+        trtllm_prefill_attn_kvfp8_dequant,
+    )
+
+    torch.set_default_device("cuda")
+    num_kv_heads, block_size, head_size = 8, 16, 64
+    num_layers = 36
+
+    kv_cache, scale = make_cross_layer_kv_cache(
+        NUM_BLOCKS,
+        num_kv_heads,
+        block_size,
+        head_size,
+        num_layers=num_layers,
+    )
+    k_scale = v_scale = scale.clone()
+
+    block_tables = torch.randint(
+        1,
+        NUM_BLOCKS,
+        (4, 6),
+        dtype=torch.int32,
+        device="cuda",
+    )
+
+    mock_kv_cache, _ = trtllm_prefill_attn_kvfp8_dequant(
+        kv_cache,
+        block_tables,
+        k_scale,
+        v_scale,
+        torch.bfloat16,
+    )
+    ref = ref_dequant(kv_cache, block_tables, k_scale, v_scale, torch.bfloat16)
+
+    torch.testing.assert_close(mock_kv_cache[1:], ref[1:], atol=1e-3, rtol=1e-3)

tests/kernels/core/test_fused_quant_layernorm.py

@@ -280,21 +280,22 @@ def test_rms_norm(
         assert torch.allclose(ref_residual, ops_residual)
 
     output = torch.empty(x.shape, dtype=quant_dtype, device=x.device)
-    scales = torch.empty(
-        (x.numel() // x.shape[-1], 1), device=x.device, dtype=torch.float32
-    )
-
     if group_size is None:
+        scales = torch.empty(
+            (x.numel() // x.shape[-1], 1), device=x.device, dtype=torch.float32
+        )
         opcheck(
             torch.ops._C.rms_norm_dynamic_per_token_quant,
             (output, x, layer.weight, scales, 1e-5, scale_ub, residual),
         )
     else:
-        # TODO(luka/eliza) opcheck is broken?
-        #  Somehow the cloned args are getting mutated in-place,
-        #  which causes the opcheck to fail.
-        # https://github.com/vllm-project/vllm/issues/36688
-        return
+        assert hidden_size % group_size[1] == 0
+        num_groups = hidden_size // group_size[1]
+        scales = torch.empty(
+            (num_groups, num_tokens),
+            device=x.device,
+            dtype=torch.float32,
+        ).transpose(0, 1)
         opcheck(
             torch.ops._C.rms_norm_per_block_quant,
             (

tests/kernels/moe/test_cpu_fused_moe.py

@@ -22,7 +22,7 @@ INTERMEDIATE_DIM = [128, 2880]
 BATCH_SIZE = [1, 64, 256]
 ACT = [MoEActivation.SILU, MoEActivation.SWIGLUOAI]
 USE_BIAS = [True, False]
-ISA = ["amx", "vec"] if torch._C._cpu._is_amx_tile_supported() else ["vec"]
+ISA = ["amx", "vec"] if torch.cpu._is_amx_tile_supported() else ["vec"]
 DTYPE = [torch.bfloat16]
 
 

tests/kernels/quantization/test_nvfp4_quant.py

@@ -159,6 +159,52 @@ def test_quantize_to_fp4(
     torch.testing.assert_close(scale_ans, scale_ref)
 
 
+@pytest.mark.parametrize(
+    "shape",
+    [(32, 4096), (128, 4096), (1, 64), (127, 1024), (256, 16384)],
+)
+@pytest.mark.parametrize("is_sf_swizzled_layout", [True, False])
+@torch.inference_mode()
+def test_python_util_matches_cpp_allocation(
+    shape: tuple[int, int],
+    is_sf_swizzled_layout: bool,
+) -> None:
+    """
+    Verify that the Python utility (create_fp4_output_tensors) allocates
+    tensors with the same shapes and dtypes as the C++ functional variant
+    (scaled_fp4_quant_func).
+    """
+    from vllm._custom_ops import create_fp4_output_tensors
+
+    torch.set_default_device("cuda:0")
+    m, n = shape
+    input_tensor = torch.randn((m, n), dtype=torch.bfloat16)
+    input_scale = torch.tensor([1.0], dtype=torch.float32, device="cuda:0")
+
+    # C++ functional variant allocates internally
+    cpp_out, cpp_scale = torch.ops._C.scaled_fp4_quant(
+        input_tensor, input_scale, is_sf_swizzled_layout
+    )
+
+    # Python utility
+    py_out, py_scale = create_fp4_output_tensors(
+        m, n, torch.device("cuda:0"), is_sf_swizzled_layout
+    )
+
+    assert py_out.shape == cpp_out.shape, (
+        f"Output shape mismatch: Python {py_out.shape} vs C++ {cpp_out.shape}"
+    )
+    assert py_out.dtype == cpp_out.dtype, (
+        f"Output dtype mismatch: Python {py_out.dtype} vs C++ {cpp_out.dtype}"
+    )
+    assert py_scale.shape == cpp_scale.shape, (
+        f"Scale shape mismatch: Python {py_scale.shape} vs C++ {cpp_scale.shape}"
+    )
+    assert py_scale.dtype == cpp_scale.dtype, (
+        f"Scale dtype mismatch: Python {py_scale.dtype} vs C++ {cpp_scale.dtype}"
+    )
+
+
 @pytest.mark.parametrize("pad_shape", PAD_SHAPES)
 @torch.inference_mode()
 def test_quantize_to_fp4_padded(pad_shape: tuple[int, int]) -> None:

tests/lora/test_lora_manager.py

@@ -711,3 +711,192 @@ def test_packed_loras(default_vllm_config, dist_init, dummy_model_gate_up, devic
     torch.testing.assert_close(
         packed_lora1.lora_b[1], model_lora_clone1.get_lora("up_proj").lora_b
     )
+
+
+def _test_target_modules(
+    model,
+    target_modules: list[str] | None,
+    device: str,
+    expected_lora: list[tuple[str, type]],
+    expected_no_lora: list[tuple[str, type]],
+):
+    """Create a LoRAModelManager and assert which modules have LoRA applied."""
+    LoRAModelManager(
+        model,
+        2,
+        2,
+        2,
+        LoRAConfig(
+            max_lora_rank=8,
+            max_cpu_loras=2,
+            max_loras=2,
+            lora_dtype=DEFAULT_DTYPE,
+            target_modules=target_modules,
+        ),
+        device=device,
+    )
+    for module_path, lora_cls in expected_lora:
+        assert isinstance(model.get_submodule(module_path), lora_cls)
+    for module_path, lora_cls in expected_no_lora:
+        assert not isinstance(model.get_submodule(module_path), lora_cls)
+
+
+@pytest.mark.parametrize("device", DEVICES)
+def test_target_modules_config(default_vllm_config, dist_init, dummy_model, device):
+    """Test that target_modules config restricts which modules get LoRA applied."""
+    _test_target_modules(
+        dummy_model,
+        ["dense1"],
+        device,
+        expected_lora=[
+            ("dense1", ColumnParallelLinearWithLoRA),
+            ("layer1.dense1", ColumnParallelLinearWithLoRA),
+        ],
+        expected_no_lora=[
+            ("dense2", RowParallelLinearWithLoRA),
+            ("layer1.dense2", RowParallelLinearWithLoRA),
+        ],
+    )
+
+
+@pytest.mark.parametrize("device", DEVICES)
+def test_target_modules_multiple(default_vllm_config, dist_init, dummy_model, device):
+    """Test that multiple target_modules work correctly."""
+    _test_target_modules(
+        dummy_model,
+        ["dense1", "dense2"],
+        device,
+        expected_lora=[
+            ("dense1", ColumnParallelLinearWithLoRA),
+            ("layer1.dense1", ColumnParallelLinearWithLoRA),
+            ("dense2", RowParallelLinearWithLoRA),
+            ("layer1.dense2", RowParallelLinearWithLoRA),
+        ],
+        expected_no_lora=[],
+    )
+
+
+@pytest.mark.parametrize("device", DEVICES)
+def test_target_modules_none_uses_all(
+    default_vllm_config, dist_init, dummy_model, device
+):
+    """Test that target_modules=None uses all supported modules."""
+    _test_target_modules(
+        dummy_model,
+        None,
+        device,
+        expected_lora=[
+            ("dense1", ColumnParallelLinearWithLoRA),
+            ("layer1.dense1", ColumnParallelLinearWithLoRA),
+            ("dense2", RowParallelLinearWithLoRA),
+            ("layer1.dense2", RowParallelLinearWithLoRA),
+        ],
+        expected_no_lora=[],
+    )
+
+
+@pytest.mark.parametrize("device", DEVICES)
+def test_load_adapter_warns_on_unsupported_modules(
+    default_vllm_config, dist_init, dummy_model_gate_up, device, tmp_path
+):
+    """Test that _load_adapter warns when a LoRA adapter contains modules
+    not in the model's supported LoRA target modules."""
+    from unittest.mock import patch
+
+    import vllm.lora.worker_manager as wm_module
+
+    lora_config = LoRAConfig(
+        max_lora_rank=8, max_cpu_loras=4, max_loras=4, lora_dtype=DEFAULT_DTYPE
+    )
+
+    dummy_lora_files = f"{tmp_path}/lora_adapter"
+    os.makedirs(dummy_lora_files, exist_ok=True)
+    create_peft_lora(
+        dummy_model_gate_up,
+        save_dir=dummy_lora_files,
+        target_modules=["layer1.dense1", "dense2"],
+        lora_dtype=DEFAULT_DTYPE,
+    )
+
+    model_config = ModelConfig(max_model_len=16)
+    vllm_config = VllmConfig(model_config=model_config, lora_config=lora_config)
+    vllm_config.scheduler_config.max_num_seqs = 4
+    vllm_config.scheduler_config.max_num_batched_tokens = 2
+
+    worker_manager = WorkerLoRAManager(vllm_config, device, EMBEDDING_MODULES)
+    worker_manager.vocab_size = dummy_model_gate_up.unpadded_vocab_size
+    worker_manager.create_lora_manager(dummy_model_gate_up)
+
+    # Patch from_local_checkpoint to inject an unsupported module
+    original_from_checkpoint = LoRAModel.from_local_checkpoint
+
+    def patched_from_checkpoint(*args, **kwargs):
+        lora = original_from_checkpoint(*args, **kwargs)
+        lora.loras["unsupported_module"] = LoRALayerWeights(
+            module_name="unsupported_module",
+            rank=8,
+            lora_alpha=16,
+            lora_a=torch.randn(8, 10),
+            lora_b=torch.randn(10, 8),
+        )
+        return lora
+
+    lora_request = LoRARequest("test", 1, dummy_lora_files)
+    with (
+        patch.object(LoRAModel, "from_local_checkpoint", patched_from_checkpoint),
+        patch.object(wm_module.logger, "warning_once") as mock_warning,
+    ):
+        worker_manager._load_adapter(lora_request)
+        warning_args = mock_warning.call_args_list
+        found = any("unsupported_module" in str(call) for call in warning_args)
+        assert found, (
+            f"Expected warning about 'unsupported_module', got: {warning_args}"
+        )
+
+
+@pytest.mark.parametrize("device", DEVICES)
+def test_load_adapter_warns_on_target_modules_restriction(
+    default_vllm_config, dist_init, dummy_model_gate_up, device, tmp_path
+):
+    """Test that _load_adapter warns when a LoRA adapter contains modules
+    excluded by the deployment-time target_modules restriction."""
+    from unittest.mock import patch
+
+    import vllm.lora.worker_manager as wm_module
+
+    # Restrict to only dense2 — adapter has dense1 which will be excluded
+    lora_config = LoRAConfig(
+        max_lora_rank=8,
+        max_cpu_loras=4,
+        max_loras=4,
+        lora_dtype=DEFAULT_DTYPE,
+        target_modules=["dense2"],
+    )
+
+    dummy_lora_files = f"{tmp_path}/lora_adapter"
+    os.makedirs(dummy_lora_files, exist_ok=True)
+    create_peft_lora(
+        dummy_model_gate_up,
+        save_dir=dummy_lora_files,
+        target_modules=["layer1.dense1", "dense2"],
+        lora_dtype=DEFAULT_DTYPE,
+    )
+
+    model_config = ModelConfig(max_model_len=16)
+    vllm_config = VllmConfig(model_config=model_config, lora_config=lora_config)
+    vllm_config.scheduler_config.max_num_seqs = 4
+    vllm_config.scheduler_config.max_num_batched_tokens = 2
+
+    worker_manager = WorkerLoRAManager(vllm_config, device, EMBEDDING_MODULES)
+    worker_manager.vocab_size = dummy_model_gate_up.unpadded_vocab_size
+    worker_manager.create_lora_manager(dummy_model_gate_up)
+
+    lora_request = LoRARequest("test", 1, dummy_lora_files)
+    with patch.object(wm_module.logger, "warning_once") as mock_warning:
+        worker_manager._load_adapter(lora_request)
+        warning_args = mock_warning.call_args_list
+        # dense1 is supported by the model but excluded by target_modules
+        found = any("target_modules" in str(call) for call in warning_args)
+        assert found, (
+            f"Expected warning about target_modules restriction, got: {warning_args}"
+        )

tests/lora/test_lora_utils.py

@@ -0,0 +1,60 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+
+from vllm.lora.utils import is_in_target_modules, is_supported_lora_module
+
+
+class TestIsSupportedLoraModule:
+    """Tests for is_supported_lora_module (model-definition check)."""
+
+    def test_suffix_match(self):
+        assert is_supported_lora_module(
+            "model.layers.0.self_attn.o_proj", ["o_proj", "q_proj"]
+        )
+
+    def test_no_match(self):
+        assert not is_supported_lora_module(
+            "model.layers.0.self_attn.o_proj", ["q_proj", "k_proj"]
+        )
+
+    def test_exact_match(self):
+        assert is_supported_lora_module("o_proj", ["o_proj"])
+
+    def test_regex_suffix_matching(self):
+        """Regex anchors to end — partial suffix should not match."""
+        assert not is_supported_lora_module("model.layers.0.self_attn.o_proj", ["proj"])
+
+    def test_empty_supported_modules(self):
+        assert not is_supported_lora_module("model.layers.0.self_attn.o_proj", [])
+
+    def test_multiple_supported_modules(self):
+        supported = ["q_proj", "k_proj", "v_proj", "o_proj"]
+        assert is_supported_lora_module("model.layers.0.self_attn.v_proj", supported)
+        assert not is_supported_lora_module("model.layers.0.mlp.gate_proj", supported)
+
+
+class TestIsInTargetModules:
+    """Tests for is_in_target_modules (deployment-time filter)."""
+
+    def test_none_allows_all(self):
+        assert is_in_target_modules("model.layers.0.self_attn.o_proj", None)
+
+    def test_suffix_in_target(self):
+        assert is_in_target_modules(
+            "model.layers.0.self_attn.o_proj", ["o_proj", "q_proj"]
+        )
+
+    def test_suffix_not_in_target(self):
+        assert not is_in_target_modules(
+            "model.layers.0.self_attn.o_proj", ["q_proj", "k_proj"]
+        )
+
+    def test_empty_target_modules(self):
+        assert not is_in_target_modules("model.layers.0.self_attn.o_proj", [])
+
+    def test_exact_name_match(self):
+        assert is_in_target_modules("dense1", ["dense1", "dense2"])
+
+    def test_exact_name_no_match(self):
+        assert not is_in_target_modules("dense3", ["dense1", "dense2"])

tests/models/multimodal/generation/vlm_utils/model_utils.py

@@ -1235,9 +1235,9 @@ def voxtral_patch_hf_runner(hf_model: "HfRunner") -> "HfRunner":
     generated).
     """
 
-    import base64
     import io
 
+    import pybase64 as base64
     import soundfile as sf
 
     processor = hf_model.processor

tests/models/multimodal/pooling/test_colpali.py

@@ -7,9 +7,9 @@ ColPali is a multi-vector retrieval model based on PaliGemma backbone
 It produces per-token embeddings for both text and image inputs.
 """
 
-import base64
 from io import BytesIO
 
+import pybase64 as base64
 import pytest
 import torch
 from PIL import Image

tests/models/multimodal/pooling/test_colqwen3.py

@@ -7,9 +7,9 @@ ColBERT-style late interaction scoring (MaxSim). It produces per-token
 embeddings for both text and image inputs.
 """
 
-import base64
 from io import BytesIO
 
+import pybase64 as base64
 import pytest
 import torch
 from PIL import Image

tests/models/multimodal/pooling/test_llama_nemotron_vl.py

@@ -9,10 +9,10 @@ Tests for the LlamaNemotronVL model family:
 Both variants share a SigLIP vision encoder with a bidirectional LLaMA backbone.
 """
 
-import base64
 from io import BytesIO
 from pathlib import Path
 
+import pybase64 as base64
 import pytest
 import torch
 from transformers import AutoModel, AutoModelForSequenceClassification, AutoProcessor

tests/models/multimodal/pooling/test_phi3v.py

@@ -3,6 +3,7 @@
 
 import pytest
 import torch.nn.functional as F
+import transformers.utils
 from PIL import Image
 
 from vllm.assets.base import get_vllm_public_assets
@@ -12,6 +13,12 @@ from ....conftest import IMAGE_ASSETS, HfRunner, PromptImageInput, VllmRunner
 from ....utils import large_gpu_test
 from ...utils import check_embeddings_close
 
+# BC for method that was deleted in Transformers v5.
+# Only needed for generating the HF reference.
+transformers.utils.is_flash_attn_greater_or_equal_2_10 = (
+    lambda: transformers.utils.is_flash_attn_greater_or_equal("2.1.0")
+)
+
 HF_TEXT_PROMPTS = [
     # T -> X
     "Find me an everyday image that matches the given caption: The label of the object is stop sign",  # noqa: E501

tests/models/multimodal/processing/test_h2ovl.py

@@ -23,7 +23,7 @@ def _get_expected_num_patches(
     min_num: int,
     max_num: int,
 ):
-    from vllm.model_executor.models.h2ovl import (
+    from vllm.transformers_utils.processors.h2ovl import (
         calculate_h2ovl_targets,
         get_h2ovl_target_ratios,
     )

tests/models/multimodal/processing/test_internvl.py

@@ -23,7 +23,7 @@ def _get_expected_num_patches(
     min_num: int,
     max_num: int,
 ):
-    from vllm.model_executor.models.internvl import (
+    from vllm.transformers_utils.processors.internvl import (
         calculate_internvl_targets,
         get_internvl_target_ratios,
     )