diff --git a/.buildkite/intel_jobs/test-intel.yaml b/.buildkite/intel_jobs/test-intel.yaml index 3aa75f4754f..295840440ad 100644 --- a/.buildkite/intel_jobs/test-intel.yaml +++ b/.buildkite/intel_jobs/test-intel.yaml @@ -56,7 +56,7 @@ steps: 'cd tests && pytest -v -s v1/core --ignore=v1/core/test_reset_prefix_cache_e2e.py --ignore=v1/core/test_scheduler_e2e.py && pytest -v -s v1/engine --ignore=v1/engine/test_output_processor.py && - pytest -v -s v1/sample --ignore=v1/sample/test_logprobs.py --ignore=v1/sample/test_logprobs_e2e.py && + pytest -v -s v1/sample --ignore=v1/sample/test_logprobs.py --ignore=v1/sample/test_logprobs_e2e.py -k "not test_topk_only and not test_topp_only and not test_topk_and_topp" && pytest -v -s v1/worker --ignore=v1/worker/test_gpu_model_runner.py --ignore=v1/worker/test_worker_memory_snapshot.py && pytest -v -s v1/structured_output && pytest -v -s v1/test_serial_utils.py && diff --git a/.buildkite/scripts/hardware_ci/run-cpu-distributed-smoke-test.sh b/.buildkite/scripts/hardware_ci/run-cpu-distributed-smoke-test.sh index d905403164a..4ff15067aa0 100644 --- a/.buildkite/scripts/hardware_ci/run-cpu-distributed-smoke-test.sh +++ b/.buildkite/scripts/hardware_ci/run-cpu-distributed-smoke-test.sh @@ -23,22 +23,22 @@ if [ "$failed_req" -ne 0 ]; then exit 1 fi -echo "--- DP+TP" -vllm serve meta-llama/Llama-3.2-3B-Instruct -tp=2 -dp=2 --max-model-len=4096 & -server_pid=$! -timeout 600 bash -c "until curl localhost:8000/v1/models > /dev/null 2>&1; do sleep 1; done" || exit 1 -vllm bench serve \ - --backend vllm \ - --dataset-name random \ - --model meta-llama/Llama-3.2-3B-Instruct \ - --num-prompts 20 \ - --result-dir ./test_results \ - --result-filename dp_pp.json \ - --save-result \ - --endpoint /v1/completions -kill -s SIGTERM $server_pid; wait $server_pid || true -failed_req=$(jq '.failed' ./test_results/dp_pp.json) -if [ "$failed_req" -ne 0 ]; then - echo "Some requests were failed!" - exit 1 -fi +#echo "--- DP+TP" +#vllm serve meta-llama/Llama-3.2-3B-Instruct -tp=2 -dp=2 --max-model-len=4096 & +#server_pid=$! +#timeout 600 bash -c "until curl localhost:8000/v1/models > /dev/null 2>&1; do sleep 1; done" || exit 1 +#vllm bench serve \ +# --backend vllm \ +# --dataset-name random \ +# --model meta-llama/Llama-3.2-3B-Instruct \ +# --num-prompts 20 \ +# --result-dir ./test_results \ +# --result-filename dp_pp.json \ +# --save-result \ +# --endpoint /v1/completions +#kill -s SIGTERM $server_pid; wait $server_pid || true +#failed_req=$(jq '.failed' ./test_results/dp_pp.json) +#if [ "$failed_req" -ne 0 ]; then +# echo "Some requests were failed!" +# exit 1 +#fi diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index b7254efd2dc..f42c495f864 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -790,7 +790,7 @@ steps: - tests/kernels/helion/ - vllm/platforms/rocm.py commands: - - pip install helion + - pip install helion==0.3.3 - pytest -v -s kernels/helion/ diff --git a/.buildkite/test_areas/compile.yaml b/.buildkite/test_areas/compile.yaml index c21b6655249..aa46447c24a 100644 --- a/.buildkite/test_areas/compile.yaml +++ b/.buildkite/test_areas/compile.yaml @@ -72,6 +72,7 @@ steps: - vllm/v1/attention/backends/flashinfer.py - vllm/compilation/ # TODO(luka) limit to vllm/compilation/passes - tests/compile/passes/test_fusion_attn.py + - tests/compile/passes/test_mla_attn_quant_fusion.py - tests/compile/passes/test_silu_mul_quant_fusion.py - tests/compile/passes/distributed/test_fusion_all_reduce.py - tests/compile/fullgraph/test_full_graph.py @@ -79,6 +80,7 @@ steps: # b200 runners are limited, so we limit the tests to the minimum set only supported on Blackwell - nvidia-smi - pytest -v -s tests/compile/passes/test_fusion_attn.py -k FLASHINFER + - pytest -v -s tests/compile/passes/test_mla_attn_quant_fusion.py - pytest -v -s tests/compile/passes/test_silu_mul_quant_fusion.py # this runner has 2 GPUs available even though num_devices=2 is not set - pytest -v -s tests/compile/passes/distributed/test_fusion_all_reduce.py diff --git a/.buildkite/test_areas/distributed.yaml b/.buildkite/test_areas/distributed.yaml index cfa9b848e34..92648c2f24a 100644 --- a/.buildkite/test_areas/distributed.yaml +++ b/.buildkite/test_areas/distributed.yaml @@ -224,6 +224,20 @@ steps: commands: - ./.buildkite/scripts/run-multi-node-test.sh /vllm-workspace/tests 2 2 $IMAGE_TAG "VLLM_TEST_SAME_HOST=0 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_same_node.py | grep 'Same node test passed' && NUM_NODES=2 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_node_count.py | grep 'Node count test passed' && python3 ../examples/offline_inference/data_parallel.py -dp=2 -tp=1 --dp-num-nodes=2 --dp-node-rank=0 --dp-master-addr=192.168.10.10 --dp-master-port=12345 --enforce-eager --trust-remote-code && VLLM_MULTI_NODE=1 pytest -v -s distributed/test_multi_node_assignment.py && VLLM_MULTI_NODE=1 pytest -v -s distributed/test_pipeline_parallel.py" "VLLM_TEST_SAME_HOST=0 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_same_node.py | grep 'Same node test passed' && NUM_NODES=2 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_node_count.py | grep 'Node count test passed' && python3 ../examples/offline_inference/data_parallel.py -dp=2 -tp=1 --dp-num-nodes=2 --dp-node-rank=1 --dp-master-addr=192.168.10.10 --dp-master-port=12345 --enforce-eager --trust-remote-code" +- label: MessageQueue TCP Multi-Node (2 GPUs) + timeout_in_minutes: 10 + working_dir: "/vllm-workspace/tests" + num_devices: 1 + num_nodes: 2 + no_plugin: true + optional: true + source_file_dependencies: + - vllm/distributed/device_communicators/shm_broadcast.py + - vllm/distributed/parallel_state.py + - tests/distributed/test_mq_tcp_multinode.py + commands: + - ./.buildkite/scripts/run-multi-node-test.sh /vllm-workspace/tests 2 1 $IMAGE_TAG "torchrun --nnodes 2 --nproc-per-node=1 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_mq_tcp_multinode.py" "torchrun --nnodes 2 --nproc-per-node=1 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_mq_tcp_multinode.py" + - label: Distributed NixlConnector PD accuracy (4 GPUs) timeout_in_minutes: 30 working_dir: "/vllm-workspace/tests" @@ -294,3 +308,23 @@ steps: commands: - pytest -v -s distributed/test_pp_cudagraph.py - pytest -v -s distributed/test_pipeline_parallel.py + +- label: RayExecutorV2 (4 GPUs) + timeout_in_minutes: 60 + working_dir: "/vllm-workspace/tests" + num_devices: 4 + source_file_dependencies: + - vllm/v1/executor/ray_executor_v2.py + - vllm/v1/executor/abstract.py + - vllm/v1/executor/multiproc_executor.py + - tests/distributed/test_ray_v2_executor.py + - tests/distributed/test_ray_v2_executor_e2e.py + - tests/distributed/test_pipeline_parallel.py + - tests/basic_correctness/test_basic_correctness.py + commands: + - export VLLM_USE_RAY_V2_EXECUTOR_BACKEND=1 + - export NCCL_CUMEM_HOST_ENABLE=0 + - pytest -v -s distributed/test_ray_v2_executor.py + - pytest -v -s distributed/test_ray_v2_executor_e2e.py + - pytest -v -s distributed/test_pipeline_parallel.py -k "ray" + - TARGET_TEST_SUITE=L4 pytest -v -s basic_correctness/test_basic_correctness.py -k "ray" diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index 69abc69b0fb..da26caf72ef 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -29,6 +29,7 @@ steps: - vllm/v1/attention # TODO: remove this dependency (https://github.com/vllm-project/vllm/issues/32267) - vllm/model_executor/layers/attention + - vllm/utils/flashinfer.py - tests/kernels/attention commands: - pytest -v -s kernels/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT @@ -139,7 +140,7 @@ steps: - vllm/utils/import_utils.py - tests/kernels/helion/ commands: - - pip install helion + - pip install helion==0.3.3 - pytest -v -s kernels/helion/ diff --git a/.buildkite/test_areas/models_distributed.yaml b/.buildkite/test_areas/models_distributed.yaml index 9df1bf830c1..55e7410b8af 100644 --- a/.buildkite/test_areas/models_distributed.yaml +++ b/.buildkite/test_areas/models_distributed.yaml @@ -18,5 +18,6 @@ steps: # Avoid importing model tests that cause CUDA reinitialization error - pytest models/test_transformers.py -v -s -m 'distributed(num_gpus=2)' - pytest models/language -v -s -m 'distributed(num_gpus=2)' - - pytest models/multimodal -v -s -m 'distributed(num_gpus=2)' --ignore models/multimodal/generation/test_whisper.py + - pytest models/multimodal/generation/test_phi4siglip.py -v -s -m 'distributed(num_gpus=2)' + - pytest models/multimodal -v -s -m 'distributed(num_gpus=2)' --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_phi4siglip.py - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest models/multimodal/generation/test_whisper.py -v -s -m 'distributed(num_gpus=2)' diff --git a/CMakeLists.txt b/CMakeLists.txt index 3db7ff0bbda..dd6ebce34be 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -340,7 +340,8 @@ if(VLLM_GPU_LANG STREQUAL "CUDA") list(APPEND VLLM_EXT_SRC "csrc/quantization/awq/gemm_kernels.cu" - "csrc/cutlass_extensions/common.cpp") + "csrc/cutlass_extensions/common.cpp" + "csrc/quantization/fused_kernels/fused_silu_mul_block_quant.cu") set_gencode_flags_for_srcs( SRCS "${VLLM_EXT_SRC}" @@ -1029,7 +1030,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA") list(APPEND VLLM_MOE_EXT_SRC "csrc/moe/moe_wna16.cu" "csrc/moe/grouped_topk_kernels.cu" - "csrc/moe/gpt_oss_router_gemm.cu" "csrc/moe/router_gemm.cu") endif() diff --git a/benchmarks/fused_kernels/merge_attn_states_benchmarks.py b/benchmarks/fused_kernels/merge_attn_states_benchmarks.py new file mode 100644 index 00000000000..26b04299b35 --- /dev/null +++ b/benchmarks/fused_kernels/merge_attn_states_benchmarks.py @@ -0,0 +1,264 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Benchmark: Fused FP8 output quantization in merge_attn_states + +Compares fused vs unfused approaches for producing FP8-quantized merged +attention output: + 1. Fused CUDA -- single CUDA kernel (merge + FP8 quant) + 2. Fused Triton -- single Triton kernel (merge + FP8 quant) + 3. Unfused CUDA -- CUDA merge + torch.compiled FP8 quant + 4. Unfused Triton -- Triton merge + torch.compiled FP8 quant + +Usage: + python benchmarks/fused_kernels/merge_attn_states_benchmarks.py + python benchmarks/fused_kernels/merge_attn_states_benchmarks.py --tp 1 4 8 + python benchmarks/fused_kernels/merge_attn_states_benchmarks.py --dtype bfloat16 +""" + +import argparse +import itertools + +import torch + +from vllm._custom_ops import merge_attn_states as merge_attn_states_cuda +from vllm.benchmarks.lib.utils import default_vllm_config +from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8 +from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape +from vllm.platforms import current_platform +from vllm.triton_utils import triton +from vllm.v1.attention.ops.triton_merge_attn_states import ( + merge_attn_states as merge_attn_states_triton, +) + +# --------------------------------------------------------------------------- +# Configuration defaults +# --------------------------------------------------------------------------- + +NUM_TOKENS_LIST = [1, 16, 64, 256, 1024, 4096] + +# (label, num_heads, head_size) — num_heads is for TP=1 +HEAD_CONFIGS = [ + ("DeepSeek-V3 MLA", 128, 128), + ("Llama-70B", 64, 128), + ("Llama-8B", 32, 128), +] + +TP_SIZES = [1, 2, 4, 8] + +INPUT_DTYPES = [torch.float32, torch.float16, torch.bfloat16] + +QUANTILES = [0.5, 0.2, 0.8] + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def short_dtype(dtype: torch.dtype) -> str: + return str(dtype).removeprefix("torch.") + + +def make_inputs( + num_tokens: int, + num_heads: int, + head_size: int, + dtype: torch.dtype, +): + """Create random prefix/suffix outputs and LSEs.""" + prefix_output = torch.randn( + (num_tokens, num_heads, head_size), dtype=dtype, device="cuda" + ) + suffix_output = torch.randn( + (num_tokens, num_heads, head_size), dtype=dtype, device="cuda" + ) + prefix_lse = torch.randn(num_heads, num_tokens, dtype=torch.float32, device="cuda") + suffix_lse = torch.randn(num_heads, num_tokens, dtype=torch.float32, device="cuda") + # Sprinkle some inf values to exercise edge-case paths + mask = torch.rand(num_heads, num_tokens, device="cuda") < 0.05 + prefix_lse[mask] = float("inf") + mask2 = torch.rand(num_heads, num_tokens, device="cuda") < 0.05 + suffix_lse[mask2] = float("inf") + return prefix_output, suffix_output, prefix_lse, suffix_lse + + +def build_configs(head_configs, num_tokens_list, input_dtypes, tp_sizes): + """Build (num_tokens, num_heads, head_size, dtype_str) config tuples, + applying TP division to num_heads and skipping invalid combos.""" + configs = [] + for (_, nh, hs), nt, dtype, tp in itertools.product( + head_configs, num_tokens_list, input_dtypes, tp_sizes + ): + nh_tp = nh // tp + if nh_tp >= 1: + configs.append((nt, nh_tp, hs, short_dtype(dtype))) + return configs + + +def parse_args(): + parser = argparse.ArgumentParser( + description="Benchmark merge_attn_states fused FP8 quantization" + ) + parser.add_argument( + "--num-tokens", + type=int, + nargs="+", + default=None, + help=f"Override token counts (default: {NUM_TOKENS_LIST})", + ) + parser.add_argument( + "--tp", + type=int, + nargs="+", + default=None, + help=f"TP sizes to simulate (divides num_heads) (default: {TP_SIZES})", + ) + parser.add_argument( + "--dtype", + type=str, + nargs="+", + default=None, + help="Input dtypes (e.g. bfloat16 float16 float32). " + f"Default: {[short_dtype(d) for d in INPUT_DTYPES]}", + ) + return parser.parse_args() + + +# --------------------------------------------------------------------------- +# Parse args and build configs before decorators +# --------------------------------------------------------------------------- + +args = parse_args() + +num_tokens_list = args.num_tokens if args.num_tokens else NUM_TOKENS_LIST +tp_sizes = args.tp if args.tp else TP_SIZES + +if args.dtype: + from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE + + input_dtypes = [STR_DTYPE_TO_TORCH_DTYPE[d] for d in args.dtype] +else: + input_dtypes = INPUT_DTYPES + +configs = build_configs(HEAD_CONFIGS, num_tokens_list, input_dtypes, tp_sizes) + +torch._dynamo.config.recompile_limit = 8888 + + +# --------------------------------------------------------------------------- +# Benchmark function +# --------------------------------------------------------------------------- + + +@triton.testing.perf_report( + triton.testing.Benchmark( + x_names=["num_tokens", "num_heads", "head_size", "dtype_str"], + x_vals=configs, + line_arg="provider", + line_vals=["fused_cuda", "fused_triton", "unfused_cuda", "unfused_triton"], + line_names=["Fused CUDA", "Fused Triton", "Unfused CUDA", "Unfused Triton"], + styles=[("blue", "-"), ("green", "-"), ("blue", "--"), ("green", "--")], + ylabel="us", + plot_name="merge_attn_states FP8 (fused vs unfused)", + args={}, + ) +) +@default_vllm_config() +def benchmark(num_tokens, num_heads, head_size, dtype_str, provider): + input_dtype = getattr(torch, dtype_str) + fp8_dtype = current_platform.fp8_dtype() + prefix_out, suffix_out, prefix_lse, suffix_lse = make_inputs( + num_tokens, num_heads, head_size, input_dtype + ) + output_scale = torch.tensor([0.1], dtype=torch.float32, device="cuda") + + if provider == "fused_cuda": + output = torch.empty( + (num_tokens, num_heads, head_size), dtype=fp8_dtype, device="cuda" + ) + fn = lambda: merge_attn_states_cuda( + output, + prefix_out, + prefix_lse, + suffix_out, + suffix_lse, + output_scale=output_scale, + ) + elif provider == "fused_triton": + output = torch.empty( + (num_tokens, num_heads, head_size), dtype=fp8_dtype, device="cuda" + ) + fn = lambda: merge_attn_states_triton( + output, + prefix_out, + prefix_lse, + suffix_out, + suffix_lse, + output_scale=output_scale, + ) + elif provider == "unfused_cuda": + merge_buf = torch.empty( + (num_tokens, num_heads, head_size), dtype=input_dtype, device="cuda" + ) + quant_fp8 = QuantFP8( + static=True, + group_shape=GroupShape.PER_TENSOR, + column_major_scales=False, + ) + quant_input = merge_buf.view(-1, head_size) + compiled_quant = torch.compile( + quant_fp8.forward_native, fullgraph=True, dynamic=False + ) + + def unfused_fn(): + merge_attn_states_cuda( + merge_buf, prefix_out, prefix_lse, suffix_out, suffix_lse + ) + compiled_quant(quant_input, output_scale) + + fn = unfused_fn + else: # unfused_triton + merge_buf = torch.empty( + (num_tokens, num_heads, head_size), dtype=input_dtype, device="cuda" + ) + quant_fp8 = QuantFP8( + static=True, + group_shape=GroupShape.PER_TENSOR, + column_major_scales=False, + ) + quant_input = merge_buf.view(-1, head_size) + compiled_quant = torch.compile( + quant_fp8.forward_native, fullgraph=True, dynamic=False + ) + + def unfused_fn(): + merge_attn_states_triton( + merge_buf, prefix_out, prefix_lse, suffix_out, suffix_lse + ) + compiled_quant(quant_input, output_scale) + + fn = unfused_fn + + ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(fn, quantiles=QUANTILES) + return 1000 * ms, 1000 * max_ms, 1000 * min_ms # us + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + + +def main(): + device_name = current_platform.get_device_name() + print(f"Device: {device_name}") + print(f"Token counts: {num_tokens_list}") + print(f"TP sizes: {tp_sizes}") + print(f"Input dtypes: {[short_dtype(d) for d in input_dtypes]}") + print(f"Head configs: {[(c[0], c[1], c[2]) for c in HEAD_CONFIGS]}") + benchmark.run(print_data=True) + + +if __name__ == "__main__": + with torch.inference_mode(): + main() diff --git a/benchmarks/fused_kernels/silu_mul_block_quant_benchmark.py b/benchmarks/fused_kernels/silu_mul_block_quant_benchmark.py new file mode 100644 index 00000000000..4e8d787bf9c --- /dev/null +++ b/benchmarks/fused_kernels/silu_mul_block_quant_benchmark.py @@ -0,0 +1,211 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from collections.abc import Callable, Iterable +from dataclasses import dataclass +from itertools import product + +import torch +import torch.nn.functional as F +import torch.utils.benchmark as TBenchmark +from torch.utils.benchmark import Measurement as TMeasurement +from tqdm import tqdm + +import vllm._custom_ops as ops +from vllm.model_executor.layers.quantization.utils.fp8_utils import ( + per_token_group_quant_fp8, +) + + +@dataclass +class bench_params_t: + num_tokens: int + hidden_size: int + dtype: torch.dtype + group_size: int # Changed from list[int] to int + + def description(self): + return ( + f"N {self.num_tokens} " + f"x D {self.hidden_size} " + f"x DT {self.dtype} " + f"x GS {self.group_size}" + ) + + +def get_bench_params() -> list[bench_params_t]: + """Test configurations covering common model sizes.""" + NUM_TOKENS = [16, 128, 512, 2048] + HIDDEN_SIZES = [1024, 2048, 4096, 5120, 14336] # Common FFN sizes + DTYPES = [torch.float16, torch.bfloat16] + GROUP_SIZES = [64, 128] # Changed from [[1, 64], [1, 128]] + + combinations = product(NUM_TOKENS, HIDDEN_SIZES, DTYPES, GROUP_SIZES) + bench_params = list( + map(lambda x: bench_params_t(x[0], x[1], x[2], x[3]), combinations) + ) + return bench_params + + +# Reference implementations +def unfused_fp8_impl( + x: torch.Tensor, + quant_dtype: torch.dtype, + group_size: int, # Changed from list[int] +): + """Unfused: SiLU+Mul then per-tensor quantize.""" + hidden = x.shape[-1] // 2 + gate, up = x.split(hidden, dim=-1) + + # SiLU(gate) * up + silu_out = F.silu(gate) * up + + # Per-tensor quantize (no group_size used here) + silu_out, _ = ops.scaled_fp8_quant(silu_out) + + +def unfused_groupwise_fp8_impl( + x: torch.Tensor, + quant_dtype: torch.dtype, + group_size: int, # Changed from list[int] +): + """Unfused: SiLU+Mul then group-wise quantize.""" + hidden = x.shape[-1] // 2 + gate, up = x.split(hidden, dim=-1) + + # SiLU(gate) * up + silu_out = F.silu(gate) * up + + # Group quantize - use group_size directly + silu_out, _ = per_token_group_quant_fp8( + silu_out, group_size=group_size, use_ue8m0=False + ) + + +def fused_impl( + x: torch.Tensor, + quant_dtype: torch.dtype, + group_size: int, +): + """Fused: SiLU+Mul+Block Quantization in single kernel.""" + out, _ = ops.silu_and_mul_per_block_quant( + x, + group_size=group_size, + quant_dtype=quant_dtype, + is_scale_transposed=False, + ) + + +# Bench functions +def bench_fn( + x: torch.Tensor, + quant_dtype: torch.dtype, + group_size: int, + label: str, + sub_label: str, + fn: Callable, + description: str, +) -> TMeasurement: + min_run_time = 1 + + globals = { + "x": x, + "quant_dtype": quant_dtype, + "group_size": group_size, + "fn": fn, + } + return TBenchmark.Timer( + stmt="fn(x, quant_dtype, group_size)", + globals=globals, + label=label, + sub_label=sub_label, + description=description, + ).blocked_autorange(min_run_time=min_run_time) + + +def bench(params: bench_params_t, label: str, sub_label: str) -> Iterable[TMeasurement]: + """Run benchmarks for all implementations.""" + # Make inputs: [num_tokens, hidden_size * 2] for [gate || up] + scale = 1 / params.hidden_size + x = ( + torch.randn( + params.num_tokens, + params.hidden_size * 2, + dtype=params.dtype, + device="cuda", + ) + * scale + ) + + timers = [] + + # Unfused per-tensor FP8 + timers.append( + bench_fn( + x, + torch.float8_e4m3fn, + params.group_size, + label, + sub_label, + unfused_fp8_impl, + "unfused_fp8_impl", + ) + ) + + # Unfused group-wise FP8 + timers.append( + bench_fn( + x, + torch.float8_e4m3fn, + params.group_size, + label, + sub_label, + unfused_groupwise_fp8_impl, + "unfused_groupwise_fp8_impl", + ) + ) + + # Fused group-wise FP8 + timers.append( + bench_fn( + x, + torch.float8_e4m3fn, + params.group_size, + label, + sub_label, + fused_impl, + "fused_groupwise_fp8_impl", + ) + ) + + return timers + + +def print_timers(timers: Iterable[TMeasurement]): + compare = TBenchmark.Compare(timers) + compare.print() + + +def main(): + torch.set_default_device("cuda") + bench_params = get_bench_params() + + print(f"Running {len(bench_params)} benchmark configurations...") + print( + f"This will take approximately {len(bench_params) * 3} seconds (1s per variant)" + ) + print() + + timers = [] + for bp in tqdm(bench_params): + result_timers = bench(bp, "silu-mul-block-quant", bp.description()) + timers.extend(result_timers) + + print("\n" + "=" * 80) + print("FINAL COMPARISON - ALL RESULTS") + print("=" * 80) + print_timers(timers) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/kernels/benchmark_router_gemm.py b/benchmarks/kernels/benchmark_router_gemm.py deleted file mode 100644 index cc63f8904c2..00000000000 --- a/benchmarks/kernels/benchmark_router_gemm.py +++ /dev/null @@ -1,134 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import torch -import torch.nn.functional as F - -from vllm import _custom_ops as ops -from vllm.platforms import current_platform -from vllm.transformers_utils.config import get_config -from vllm.triton_utils import triton -from vllm.utils.argparse_utils import FlexibleArgumentParser - -# Dimensions supported by the DSV3 specialized kernel -DSV3_SUPPORTED_NUM_EXPERTS = [256, 384] -DSV3_SUPPORTED_HIDDEN_SIZES = [7168] - -# Dimensions supported by the gpt-oss specialized kernel -GPT_OSS_SUPPORTED_NUM_EXPERTS = [32, 128] -GPT_OSS_SUPPORTED_HIDDEN_SIZES = [2880] - - -def get_batch_size_range(max_batch_size): - return [2**x for x in range(14) if 2**x <= max_batch_size] - - -def get_model_params(config): - if config.architectures[0] in ( - "DeepseekV2ForCausalLM", - "DeepseekV3ForCausalLM", - "DeepseekV32ForCausalLM", - ): - num_experts = config.n_routed_experts - hidden_size = config.hidden_size - elif config.architectures[0] in ("GptOssForCausalLM",): - num_experts = config.num_local_experts - hidden_size = config.hidden_size - else: - raise ValueError(f"Unsupported architecture: {config.architectures}") - return num_experts, hidden_size - - -def get_benchmark(model, max_batch_size, trust_remote_code): - @triton.testing.perf_report( - triton.testing.Benchmark( - x_names=["batch_size"], - x_vals=get_batch_size_range(max_batch_size), - x_log=False, - line_arg="provider", - line_vals=[ - "torch", - "vllm", - ], - line_names=["PyTorch", "vLLM"], - styles=([("blue", "-"), ("red", "-")]), - ylabel="TFLOPs", - plot_name=f"{model} router gemm throughput", - args={}, - ) - ) - def benchmark(batch_size, provider): - config = get_config(model=model, trust_remote_code=trust_remote_code) - num_experts, hidden_size = get_model_params(config) - - mat_a = torch.randn( - (batch_size, hidden_size), dtype=torch.bfloat16, device="cuda" - ).contiguous() - mat_b = torch.randn( - (num_experts, hidden_size), dtype=torch.bfloat16, device="cuda" - ).contiguous() - bias = torch.randn( - num_experts, dtype=torch.bfloat16, device="cuda" - ).contiguous() - - is_hopper_or_blackwell = current_platform.is_device_capability( - 90 - ) or current_platform.is_device_capability_family(100) - allow_dsv3_router_gemm = ( - is_hopper_or_blackwell - and num_experts in DSV3_SUPPORTED_NUM_EXPERTS - and hidden_size in DSV3_SUPPORTED_HIDDEN_SIZES - ) - allow_gpt_oss_router_gemm = ( - is_hopper_or_blackwell - and num_experts in GPT_OSS_SUPPORTED_NUM_EXPERTS - and hidden_size in GPT_OSS_SUPPORTED_HIDDEN_SIZES - ) - - has_bias = False - if allow_gpt_oss_router_gemm: - has_bias = True - - quantiles = [0.5, 0.2, 0.8] - - if provider == "torch": - - def runner(): - if has_bias: - F.linear(mat_a, mat_b, bias) - else: - F.linear(mat_a, mat_b) - elif provider == "vllm": - - def runner(): - if allow_dsv3_router_gemm: - ops.dsv3_router_gemm(mat_a, mat_b, torch.bfloat16) - elif allow_gpt_oss_router_gemm: - ops.gpt_oss_router_gemm(mat_a, mat_b, bias) - else: - raise ValueError("Unsupported router gemm") - - ms, min_ms, max_ms = triton.testing.do_bench_cudagraph( - runner, quantiles=quantiles - ) - - def tflops(t_ms): - flops = 2 * batch_size * hidden_size * num_experts - return flops / (t_ms * 1e-3) / 1e12 - - return tflops(ms), tflops(max_ms), tflops(min_ms) - - return benchmark - - -if __name__ == "__main__": - parser = FlexibleArgumentParser() - parser.add_argument("--model", type=str, default="openai/gpt-oss-20b") - parser.add_argument("--max-batch-size", default=16, type=int) - parser.add_argument("--trust-remote-code", action="store_true") - args = parser.parse_args() - - # Get the benchmark function - benchmark = get_benchmark(args.model, args.max_batch_size, args.trust_remote_code) - # Run performance benchmark - benchmark.run(print_data=True) diff --git a/cmake/external_projects/vllm_flash_attn.cmake b/cmake/external_projects/vllm_flash_attn.cmake index 443d41d5a21..9414f5af372 100644 --- a/cmake/external_projects/vllm_flash_attn.cmake +++ b/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 29210221863736a08f71a866459e368ad1ac4a95 + GIT_TAG c0ec424fd8a546d0cbbf4bf050bbcfe837c55afb GIT_PROGRESS TRUE # Don't share the vllm-flash-attn build between build types BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn diff --git a/csrc/attention/merge_attn_states.cu b/csrc/attention/merge_attn_states.cu index f6c1bf61761..75f066e8091 100644 --- a/csrc/attention/merge_attn_states.cu +++ b/csrc/attention/merge_attn_states.cu @@ -7,19 +7,29 @@ #include "attention_dtypes.h" #include "attention_utils.cuh" +#include "../quantization/w8a8/fp8/common.cuh" +#include "../dispatch_utils.h" namespace vllm { // Implements section 2.2 of https://www.arxiv.org/pdf/2501.01005 // can be used to combine partial attention results (in the split-KV case) -template +template __global__ void merge_attn_states_kernel( - scalar_t* output, float* output_lse, const scalar_t* prefix_output, + output_t* output, float* output_lse, const scalar_t* prefix_output, const float* prefix_lse, const scalar_t* suffix_output, const float* suffix_lse, const uint num_tokens, const uint num_heads, const uint head_size, const uint prefix_head_stride, - const uint output_head_stride, const uint prefix_num_tokens) { - using pack_128b_t = uint4; + const uint output_head_stride, const uint prefix_num_tokens, + const float* output_scale) { + // Inputs always load 128-bit packs (pack_size elements of scalar_t). + // Outputs store pack_size elements of output_t, which is smaller for FP8. + using input_pack_t = uint4; + using output_pack_t = + std::conditional_t, + uint4>; const uint pack_size = 16 / sizeof(scalar_t); const uint threads_per_head = head_size / pack_size; @@ -42,15 +52,36 @@ __global__ void merge_attn_states_kernel( head_idx * output_head_stride; const scalar_t* prefix_head_ptr = prefix_output + src_head_offset; const scalar_t* suffix_head_ptr = suffix_output + src_head_offset; - scalar_t* output_head_ptr = output + dst_head_offset; + output_t* output_head_ptr = output + dst_head_offset; + + // Pre-invert scale: multiplication is faster than division + float fp8_scale_inv = 1.0f; + if constexpr (USE_FP8_OUTPUT) { + fp8_scale_inv = 1.0f / *output_scale; + } // If token_idx >= prefix_num_tokens, just copy from suffix if (token_idx >= prefix_num_tokens) { if (pack_offset < head_size) { - pack_128b_t s_out_pack = reinterpret_cast( + input_pack_t s_out_pack = reinterpret_cast( suffix_head_ptr)[pack_offset / pack_size]; - reinterpret_cast(output_head_ptr)[pack_offset / pack_size] = - s_out_pack; + + if constexpr (USE_FP8_OUTPUT) { + output_t o_out_pack[pack_size]; +#pragma unroll + for (uint i = 0; i < pack_size; ++i) { + const float val = + vllm::to_float(reinterpret_cast(&s_out_pack)[i]); + o_out_pack[i] = + vllm::scaled_fp8_conversion(val, fp8_scale_inv); + } + reinterpret_cast( + output_head_ptr)[pack_offset / pack_size] = + *reinterpret_cast(o_out_pack); + } else { + reinterpret_cast( + output_head_ptr)[pack_offset / pack_size] = s_out_pack; + } } if (output_lse != nullptr && pack_idx == 0) { float s_lse = suffix_lse[head_idx * num_tokens + token_idx]; @@ -70,20 +101,34 @@ __global__ void merge_attn_states_kernel( /* In certain edge cases, MLA can produce p_lse = s_lse = -inf; continuing the pipeline then yields NaN. Root cause: with chunked prefill a batch may be split into two chunks; if a request in that batch has no - prefix hit, every LSE entry for that request’s position is -inf, and at + prefix hit, every LSE entry for that request's position is -inf, and at this moment we merge cross-attention at first. For now we simply emit prefix_output (expected to be all zeros) and prefix_lse (-inf) to fix this problem. */ if (std::isinf(max_lse)) { if (pack_offset < head_size) { - // Pack 128b load - pack_128b_t p_out_pack = reinterpret_cast( + input_pack_t p_out_pack = reinterpret_cast( prefix_head_ptr)[pack_offset / pack_size]; - // Pack 128b storage - reinterpret_cast(output_head_ptr)[pack_offset / pack_size] = - p_out_pack; + if constexpr (USE_FP8_OUTPUT) { + // Convert prefix values to FP8 (since -inf means no data, + // prefix_output is expected to be zeros) + output_t o_out_pack[pack_size]; +#pragma unroll + for (uint i = 0; i < pack_size; ++i) { + const float val = + vllm::to_float(reinterpret_cast(&p_out_pack)[i]); + o_out_pack[i] = + vllm::scaled_fp8_conversion(val, fp8_scale_inv); + } + reinterpret_cast( + output_head_ptr)[pack_offset / pack_size] = + *reinterpret_cast(o_out_pack); + } else { + reinterpret_cast( + output_head_ptr)[pack_offset / pack_size] = p_out_pack; + } } // We only need to write to output_lse once per head. if (output_lse != nullptr && pack_idx == 0) { @@ -101,30 +146,43 @@ __global__ void merge_attn_states_kernel( const float s_scale = s_se / out_se; if (pack_offset < head_size) { - // Pack 128b load - pack_128b_t p_out_pack = reinterpret_cast( + input_pack_t p_out_pack = reinterpret_cast( prefix_head_ptr)[pack_offset / pack_size]; - pack_128b_t s_out_pack = reinterpret_cast( + input_pack_t s_out_pack = reinterpret_cast( suffix_head_ptr)[pack_offset / pack_size]; - pack_128b_t o_out_pack; + // Compute merged values in float32 + float o_out_f[pack_size]; #pragma unroll for (uint i = 0; i < pack_size; ++i) { - // Always use float for FMA to keep high precision. - // half(uint16_t), bfloat16, float -> float. const float p_out_f = vllm::to_float(reinterpret_cast(&p_out_pack)[i]); const float s_out_f = vllm::to_float(reinterpret_cast(&s_out_pack)[i]); - // fma: a * b + c = p_out_f * p_scale + (s_out_f * s_scale) - const float o_out_f = p_out_f * p_scale + (s_out_f * s_scale); - // float -> half(uint16_t), bfloat16, float. - vllm::from_float(reinterpret_cast(&o_out_pack)[i], o_out_f); + o_out_f[i] = p_out_f * p_scale + (s_out_f * s_scale); } - // Pack 128b storage - reinterpret_cast(output_head_ptr)[pack_offset / pack_size] = - o_out_pack; + // Convert and store + if constexpr (USE_FP8_OUTPUT) { + output_t o_out_pack[pack_size]; +#pragma unroll + for (uint i = 0; i < pack_size; ++i) { + o_out_pack[i] = vllm::scaled_fp8_conversion( + o_out_f[i], fp8_scale_inv); + } + reinterpret_cast( + output_head_ptr)[pack_offset / pack_size] = + *reinterpret_cast(o_out_pack); + } else { + output_pack_t o_out_pack; +#pragma unroll + for (uint i = 0; i < pack_size; ++i) { + vllm::from_float(reinterpret_cast(&o_out_pack)[i], + o_out_f[i]); + } + reinterpret_cast( + output_head_ptr)[pack_offset / pack_size] = o_out_pack; + } } // We only need to write to output_lse once per head. if (output_lse != nullptr && pack_idx == 0) { @@ -151,24 +209,26 @@ __global__ void merge_attn_states_kernel( } \ } -#define LAUNCH_MERGE_ATTN_STATES(scalar_t, NUM_THREADS) \ +#define LAUNCH_MERGE_ATTN_STATES(scalar_t, output_t, NUM_THREADS, \ + USE_FP8_OUTPUT) \ { \ - vllm::merge_attn_states_kernel \ + vllm::merge_attn_states_kernel \ <<>>( \ - reinterpret_cast(output.data_ptr()), output_lse_ptr, \ + reinterpret_cast(output.data_ptr()), output_lse_ptr, \ reinterpret_cast(prefix_output.data_ptr()), \ reinterpret_cast(prefix_lse.data_ptr()), \ reinterpret_cast(suffix_output.data_ptr()), \ reinterpret_cast(suffix_lse.data_ptr()), num_tokens, \ num_heads, head_size, prefix_head_stride, output_head_stride, \ - prefix_num_tokens); \ + prefix_num_tokens, output_scale_ptr); \ } /*@brief Merges the attention states from prefix and suffix * into the output tensor. NUM_TOKENS: n, NUM_HEADS: h, HEAD_SIZE: d * * @param output [n,h,d] The output tensor to store the merged attention states. - * @param output_lse [h,d] Optional tensor to store the log-sum-exp values. + * @param output_lse [h,n] Optional tensor to store the log-sum-exp values. * @param prefix_output [n,h,d] The prefix attention states. * @param prefix_lse [h,n] The log-sum-exp values for the prefix attention * states. @@ -180,19 +240,23 @@ __global__ void merge_attn_states_kernel( * is computed by merging prefix_output and suffix_output. For remaining tokens * (prefill_tokens_with_context <= token_idx < n), output is copied directly * from suffix_output. + * @param output_scale Optional scalar tensor for FP8 static quantization. + * When provided, output must be FP8 dtype. */ template void merge_attn_states_launcher( torch::Tensor& output, std::optional output_lse, const torch::Tensor& prefix_output, const torch::Tensor& prefix_lse, const torch::Tensor& suffix_output, const torch::Tensor& suffix_lse, - const std::optional prefill_tokens_with_context) { + const std::optional prefill_tokens_with_context, + const std::optional& output_scale) { constexpr uint NUM_THREADS = 128; const uint num_tokens = output.size(0); const uint num_heads = output.size(1); const uint head_size = output.size(2); const uint prefix_head_stride = prefix_output.stride(1); const uint output_head_stride = output.stride(1); + // Thread mapping is based on input BF16 pack_size const uint pack_size = 16 / sizeof(scalar_t); TORCH_CHECK(head_size % pack_size == 0, "headsize must be multiple of pack_size:", pack_size); @@ -208,6 +272,10 @@ void merge_attn_states_launcher( if (output_lse.has_value()) { output_lse_ptr = output_lse.value().data_ptr(); } + float* output_scale_ptr = nullptr; + if (output_scale.has_value()) { + output_scale_ptr = output_scale.value().data_ptr(); + } // Process one pack elements per thread. for float, the // pack_size is 4 for half/bf16, the pack_size is 8. const uint threads_per_head = head_size / pack_size; @@ -219,20 +287,44 @@ void merge_attn_states_launcher( const c10::cuda::OptionalCUDAGuard device_guard(prefix_output.device()); auto stream = at::cuda::getCurrentCUDAStream(); - LAUNCH_MERGE_ATTN_STATES(scalar_t, NUM_THREADS); + if (output_scale.has_value()) { + // FP8 output path - dispatch on output FP8 type + VLLM_DISPATCH_FP8_TYPES(output.scalar_type(), "merge_attn_states_fp8", [&] { + LAUNCH_MERGE_ATTN_STATES(scalar_t, fp8_t, NUM_THREADS, true); + }); + } else { + // Original BF16/FP16/FP32 output path + LAUNCH_MERGE_ATTN_STATES(scalar_t, scalar_t, NUM_THREADS, false); + } } #define CALL_MERGE_ATTN_STATES_LAUNCHER(scalar_t) \ { \ merge_attn_states_launcher( \ output, output_lse, prefix_output, prefix_lse, suffix_output, \ - suffix_lse, prefill_tokens_with_context); \ + suffix_lse, prefill_tokens_with_context, output_scale); \ } -void merge_attn_states( - torch::Tensor& output, std::optional output_lse, - const torch::Tensor& prefix_output, const torch::Tensor& prefix_lse, - const torch::Tensor& suffix_output, const torch::Tensor& suffix_lse, - std::optional prefill_tokens_with_context = std::nullopt) { - DISPATCH_BY_SCALAR_DTYPE(output.dtype(), CALL_MERGE_ATTN_STATES_LAUNCHER); +void merge_attn_states(torch::Tensor& output, + std::optional output_lse, + const torch::Tensor& prefix_output, + const torch::Tensor& prefix_lse, + const torch::Tensor& suffix_output, + const torch::Tensor& suffix_lse, + std::optional prefill_tokens_with_context, + const std::optional& output_scale) { + if (output_scale.has_value()) { + TORCH_CHECK(output.scalar_type() == at::ScalarType::Float8_e4m3fn || + output.scalar_type() == at::ScalarType::Float8_e4m3fnuz, + "output must be FP8 when output_scale is provided, got: ", + output.scalar_type()); + } else { + TORCH_CHECK(output.scalar_type() == prefix_output.scalar_type(), + "output dtype (", output.scalar_type(), + ") must match prefix_output dtype (", + prefix_output.scalar_type(), ") when output_scale is not set"); + } + // Always dispatch on prefix_output (input) dtype + DISPATCH_BY_SCALAR_DTYPE(prefix_output.dtype(), + CALL_MERGE_ATTN_STATES_LAUNCHER); } diff --git a/csrc/cache.h b/csrc/cache.h index 0188a568edc..821d5e719a4 100644 --- a/csrc/cache.h +++ b/csrc/cache.h @@ -10,6 +10,10 @@ void swap_blocks(torch::Tensor& src, torch::Tensor& dst, int64_t block_size_in_bytes, const torch::Tensor& block_mapping); +void swap_blocks_batch(const torch::Tensor& src_ptrs, + const torch::Tensor& dst_ptrs, + const torch::Tensor& sizes); + void reshape_and_cache(torch::Tensor& key, torch::Tensor& value, torch::Tensor& key_cache, torch::Tensor& value_cache, torch::Tensor& slot_mapping, diff --git a/csrc/cache_kernels.cu b/csrc/cache_kernels.cu index 2b3906df9ec..c59da4379da 100644 --- a/csrc/cache_kernels.cu +++ b/csrc/cache_kernels.cu @@ -24,6 +24,8 @@ #ifdef USE_ROCM #include typedef __hip_bfloat16 __nv_bfloat16; +#else + #include #endif #if defined(__gfx942__) @@ -73,6 +75,59 @@ void swap_blocks(torch::Tensor& src, torch::Tensor& dst, } } +void swap_blocks_batch(const torch::Tensor& src_ptrs, + const torch::Tensor& dst_ptrs, + const torch::Tensor& sizes) { + TORCH_CHECK(src_ptrs.device().is_cpu(), "src_ptrs must be on CPU"); + TORCH_CHECK(dst_ptrs.device().is_cpu(), "dst_ptrs must be on CPU"); + TORCH_CHECK(sizes.device().is_cpu(), "sizes must be on CPU"); + TORCH_CHECK(src_ptrs.dtype() == torch::kInt64, "src_ptrs must be int64"); + TORCH_CHECK(dst_ptrs.dtype() == torch::kInt64, "dst_ptrs must be int64"); + TORCH_CHECK(sizes.dtype() == torch::kInt64, "sizes must be int64"); + + const int64_t n = src_ptrs.size(0); + TORCH_CHECK(dst_ptrs.size(0) == n, "dst_ptrs length must match src_ptrs"); + TORCH_CHECK(sizes.size(0) == n, "sizes length must match src_ptrs"); + + if (n == 0) return; + + const int64_t* src_data = src_ptrs.data_ptr(); + const int64_t* dst_data = dst_ptrs.data_ptr(); + const int64_t* size_data = sizes.data_ptr(); + + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + + // Use cuMemcpyBatchAsync (CUDA 12.8+) to submit all copies in a single + // driver call, amortizing per-copy submission overhead. + // int64_t and CUdeviceptr/size_t are both 8 bytes on 64-bit platforms, + // so we reinterpret_cast the tensor data directly to avoid copies. + static_assert(sizeof(CUdeviceptr) == sizeof(int64_t)); + static_assert(sizeof(size_t) == sizeof(int64_t)); +#if !defined(USE_ROCM) && defined(CUDA_VERSION) && CUDA_VERSION >= 12080 + CUmemcpyAttributes attr = {}; + attr.srcAccessOrder = CU_MEMCPY_SRC_ACCESS_ORDER_STREAM; + size_t attrs_idx = 0; + size_t fail_idx = 0; + CUresult result = cuMemcpyBatchAsync( + reinterpret_cast(const_cast(dst_data)), + reinterpret_cast(const_cast(src_data)), + reinterpret_cast(const_cast(size_data)), + static_cast(n), &attr, &attrs_idx, 1, &fail_idx, + static_cast(stream)); + TORCH_CHECK(result == CUDA_SUCCESS, "cuMemcpyBatchAsync failed at index ", + fail_idx, " with error ", result); +#else + // Fallback for CUDA < 12.8 and ROCm: individual async copies. + // cudaMemcpyDefault lets the driver infer direction from pointer types. + for (int64_t i = 0; i < n; i++) { + cudaMemcpyAsync(reinterpret_cast(dst_data[i]), + reinterpret_cast(src_data[i]), + static_cast(size_data[i]), cudaMemcpyDefault, + stream); + } +#endif +} + namespace vllm { // Grid: (num_layers, num_pairs) diff --git a/csrc/cpu/cpu_fused_moe.cpp b/csrc/cpu/cpu_fused_moe.cpp index 1a82645397b..0dc5060fe99 100644 --- a/csrc/cpu/cpu_fused_moe.cpp +++ b/csrc/cpu/cpu_fused_moe.cpp @@ -30,13 +30,15 @@ }() namespace { -enum class FusedMOEAct { SiluAndMul, SwigluOAIAndMul }; +enum class FusedMOEAct { SiluAndMul, SwigluOAIAndMul, GeluAndMul }; FusedMOEAct get_act_type(const std::string& act) { if (act == "silu") { return FusedMOEAct::SiluAndMul; } else if (act == "swigluoai") { return FusedMOEAct::SwigluOAIAndMul; + } else if (act == "gelu") { + return FusedMOEAct::GeluAndMul; } else { TORCH_CHECK(false, "Invalid act type: " + act); } @@ -104,6 +106,43 @@ void silu_and_mul(float* __restrict__ input, scalar_t* __restrict__ output, } } +template +void gelu_and_mul(float* __restrict__ input, scalar_t* __restrict__ output, + const int32_t m_size, const int32_t n_size, + const int32_t input_stride, const int32_t output_stride) { + using scalar_vec_t = typename cpu_utils::VecTypeTrait::vec_t; + const int32_t dim = n_size / 2; + float* __restrict__ gate = input; + float* __restrict__ up = input + dim; + vec_op::FP32Vec16 one_vec(1.0); + vec_op::FP32Vec16 w1_vec(M_SQRT1_2); + vec_op::FP32Vec16 w2_vec(0.5); + alignas(64) float temp[16]; + + DEFINE_FAST_EXP + + for (int32_t m = 0; m < m_size; ++m) { + for (int32_t n = 0; n < dim; n += 16) { + vec_op::FP32Vec16 gate_vec(gate + n); + vec_op::FP32Vec16 up_vec(up + n); + auto er_input_vec = gate_vec * w1_vec; + + er_input_vec.save(temp); + for (int32_t i = 0; i < 16; ++i) { + temp[i] = std::erf(temp[i]); + } + vec_op::FP32Vec16 er_vec(temp); + auto gelu = gate_vec * w2_vec * (one_vec + er_vec); + auto gated_output_fp32 = up_vec * gelu; + scalar_vec_t gated_output = scalar_vec_t(gated_output_fp32); + gated_output.save(output + n); + } + gate += input_stride; + up += input_stride; + output += output_stride; + } +} + template FORCE_INLINE void apply_gated_act(const FusedMOEAct act, float* __restrict__ input, @@ -118,6 +157,9 @@ FORCE_INLINE void apply_gated_act(const FusedMOEAct act, case FusedMOEAct::SiluAndMul: silu_and_mul(input, output, m, n, input_stride, output_stride); return; + case FusedMOEAct::GeluAndMul: + gelu_and_mul(input, output, m, n, input_stride, output_stride); + return; default: TORCH_CHECK(false, "Unsupported act type."); } diff --git a/csrc/cpu/torch_bindings.cpp b/csrc/cpu/torch_bindings.cpp index a1d7d361d19..fbc9c65241b 100644 --- a/csrc/cpu/torch_bindings.cpp +++ b/csrc/cpu/torch_bindings.cpp @@ -8,8 +8,6 @@ // libraries use different ISAs. #define TORCH_EXTENSION_NAME _C -std::string init_cpu_threads_env(const std::string& cpu_ids); - void release_dnnl_matmul_handler(int64_t handler); int64_t create_onednn_scaled_mm_handler(const torch::Tensor& b, @@ -354,7 +352,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { "str act, str isa) -> ()"); ops.impl("cpu_fused_moe", torch::kCPU, &cpu_fused_moe); #endif - ops.def("init_cpu_threads_env(str cpu_ids) -> str", &init_cpu_threads_env); ops.def( "mla_decode_kvcache(" " Tensor! out, Tensor query, Tensor kv_cache," diff --git a/csrc/cpu/utils.cpp b/csrc/cpu/utils.cpp index 3c133a0c59c..55d0a3d1757 100644 --- a/csrc/cpu/utils.cpp +++ b/csrc/cpu/utils.cpp @@ -21,150 +21,6 @@ std::string init_cpu_threads_env(const std::string& cpu_ids) { #endif -#ifndef VLLM_NUMA_DISABLED -std::string init_cpu_threads_env(const std::string& cpu_ids) { - bitmask* omp_cpu_mask = numa_parse_cpustring_all(cpu_ids.c_str()); - TORCH_CHECK(omp_cpu_mask != nullptr, - "Failed to parse CPU string: " + cpu_ids); - TORCH_CHECK(omp_cpu_mask->size > 0); - std::vector omp_cpu_ids; - omp_cpu_ids.reserve(omp_cpu_mask->size); - - constexpr int group_size = 8 * sizeof(*omp_cpu_mask->maskp); - - for (int offset = 0; offset < omp_cpu_mask->size; offset += group_size) { - unsigned long group_mask = omp_cpu_mask->maskp[offset / group_size]; - int i = 0; - while (group_mask) { - if (group_mask & 1) { - omp_cpu_ids.emplace_back(offset + i); - } - ++i; - group_mask >>= 1; - } - } - - // Memory node binding - if (numa_available() != -1) { - std::set node_ids; - for (const auto& cpu_id : omp_cpu_ids) { - int node_id = numa_node_of_cpu(cpu_id); - if (node_id != -1) { - node_ids.insert(node_id); - } - } - // Concatenate all node_ids into a single comma-separated string - if (!node_ids.empty()) { - std::string node_ids_str; - for (const int node_id : node_ids) { - if (!node_ids_str.empty()) { - node_ids_str += ","; - } - node_ids_str += std::to_string(node_id); - } - - bitmask* mask = numa_parse_nodestring(node_ids_str.c_str()); - bitmask* src_mask = numa_get_mems_allowed(); - - int pid = getpid(); - - if (mask && src_mask) { - // move all existing pages to the specified numa node. - *(src_mask->maskp) = *(src_mask->maskp) ^ *(mask->maskp); - int page_num = numa_migrate_pages(pid, src_mask, mask); - if (page_num == -1) { - TORCH_WARN("numa_migrate_pages failed. errno: " + - std::to_string(errno)); - } - - // Restrict memory allocation to the selected NUMA node(s). - // Enhances memory locality for the threads bound to those NUMA CPUs. - if (node_ids.size() > 1) { - errno = 0; - numa_set_interleave_mask(mask); - if (errno != 0) { - TORCH_WARN("numa_set_interleave_mask failed. errno: " + - std::to_string(errno)); - } else { - TORCH_WARN( - "NUMA binding: Using INTERLEAVE policy for memory " - "allocation across multiple NUMA nodes (nodes: " + - node_ids_str + - "). Memory allocations will be " - "interleaved across the specified NUMA nodes."); - } - } else { - errno = 0; - numa_set_membind(mask); - if (errno != 0) { - TORCH_WARN("numa_set_membind failed. errno: " + - std::to_string(errno)); - } else { - TORCH_WARN( - "NUMA binding: Using MEMBIND policy for memory " - "allocation on the NUMA nodes (" + - node_ids_str + - "). Memory allocations will be " - "strictly bound to these NUMA nodes."); - } - } - - numa_set_strict(1); - - numa_free_nodemask(mask); - numa_free_nodemask(src_mask); - } else { - TORCH_WARN( - "numa_parse_nodestring or numa_get_run_node_mask failed. errno: " + - std::to_string(errno)); - } - } - } - - // OMP threads binding - omp_set_num_threads((int)omp_cpu_ids.size()); - torch::set_num_threads((int)omp_cpu_ids.size()); - TORCH_CHECK_EQ(omp_cpu_ids.size(), torch::get_num_threads()); - TORCH_CHECK_EQ(omp_cpu_ids.size(), omp_get_max_threads()); - - std::vector> thread_core_mapping; - thread_core_mapping.reserve(omp_cpu_ids.size()); - omp_lock_t writelock; - omp_init_lock(&writelock); - - #pragma omp parallel for schedule(static, 1) - for (size_t i = 0; i < omp_cpu_ids.size(); ++i) { - cpu_set_t mask; - CPU_ZERO(&mask); - CPU_SET(omp_cpu_ids[i], &mask); - int ret = sched_setaffinity(0, sizeof(cpu_set_t), &mask); - if (ret == -1) { - TORCH_CHECK(false, - "sched_setaffinity failed. errno: " + std::to_string(errno)); - } - - omp_set_lock(&writelock); - thread_core_mapping.emplace_back(gettid(), omp_cpu_ids[i]); - omp_unset_lock(&writelock); - } - - omp_destroy_lock(&writelock); - - numa_free_nodemask(omp_cpu_mask); - - std::stringstream ss; - ss << "OMP threads binding of Process " << getpid() << ":\n"; - std::sort(thread_core_mapping.begin(), thread_core_mapping.end(), - [](auto&& a, auto&& b) { return a.second < b.second; }); - for (auto&& item : thread_core_mapping) { - ss << "\t" - << "OMP tid: " << item.first << ", core " << item.second << "\n"; - } - - return ss.str(); -} -#endif // VLLM_NUMA_DISABLED - namespace cpu_utils { ScratchPadManager::ScratchPadManager() : size_(0), ptr_(nullptr) { this->realloc(allocation_unit * 128); diff --git a/csrc/libtorch_stable/quantization/fp4/nvfp4_quant_entry.cu b/csrc/libtorch_stable/quantization/fp4/nvfp4_quant_entry.cu index 4adc6243765..8d4ba1accc7 100644 --- a/csrc/libtorch_stable/quantization/fp4/nvfp4_quant_entry.cu +++ b/csrc/libtorch_stable/quantization/fp4/nvfp4_quant_entry.cu @@ -58,16 +58,19 @@ void silu_and_mul_scaled_fp4_experts_quant_sm1xxa( torch::stable::Tensor const& output_scale_offset_by_experts); #endif +#if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \ + (defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120) static bool nvfp4_quant_sm_supported() { const int32_t sm = get_sm_version_num(); -#if defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100 + #if defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100 if (sm >= 100 && sm < 120) return true; -#endif -#if defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120 + #endif + #if defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120 if (sm >= 120 && sm < 130) return true; -#endif + #endif return false; } +#endif void scaled_fp4_quant_out(torch::stable::Tensor const& input, torch::stable::Tensor const& input_sf, diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm120_fp8_dispatch.cuh b/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm120_fp8_dispatch.cuh index 14de9b1e4fd..a9008ce4424 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm120_fp8_dispatch.cuh +++ b/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm120_fp8_dispatch.cuh @@ -26,8 +26,10 @@ using namespace cute; template + class EpilogueScheduler, class MainloopScheduler, + bool swap_ab_ = false> struct cutlass_3x_gemm_fp8_blockwise { + static constexpr bool swap_ab = swap_ab_; using ElementAB = cutlass::float_e4m3_t; using ElementA = ElementAB; @@ -55,9 +57,13 @@ struct cutlass_3x_gemm_fp8_blockwise { using ElementCompute = float; using ElementBlockScale = float; - using ScaleConfig = cutlass::detail::Sm120BlockwiseScaleConfig< + using ScaleConfig = conditional_t; + cute::UMMA::Major::K, cute::UMMA::Major::MN>, + cutlass::detail::Sm120BlockwiseScaleConfig< + ScaleGranularityM, ScaleGranularityN, ScaleGranularityK, + cute::UMMA::Major::MN, cute::UMMA::Major::K>>; // layout_SFA and layout_SFB cannot be swapped since they are deduced. using LayoutSFA = decltype(ScaleConfig::deduce_layoutSFA()); @@ -78,17 +84,32 @@ struct cutlass_3x_gemm_fp8_blockwise { ElementAccumulator, ElementCompute, ElementC, - LayoutC, + conditional_t, AlignmentC, ElementD, - LayoutD, + conditional_t, AlignmentD, EpilogueScheduler, DefaultOperation >::CollectiveOp; using StageCountType = cutlass::gemm::collective::StageCountAuto; - using CollectiveMainloop = + using CollectiveMainloop = conditional_t, + AlignmentB, + ElementA, + cute::tuple, + AlignmentA, + ElementAccumulator, + MmaTileShape, + ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + MainloopScheduler + >::CollectiveOp, typename cutlass::gemm::collective::CollectiveBuilder< ArchTag, OperatorClass, @@ -103,7 +124,7 @@ struct cutlass_3x_gemm_fp8_blockwise { ClusterShape, cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, MainloopScheduler - >::CollectiveOp; + >::CollectiveOp>; // SM12x family to support both SM120 (RTX 5090) and SM121 (DGX Spark) using KernelType = enable_sm120_family struct sm120_blockwise_fp8_config_default { - // M > 256: use 128x128x128 tile with Cooperative (Auto) schedule + // use 128x128x128 tile with Cooperative (Auto) schedule using KernelSchedule = cutlass::gemm::collective::KernelScheduleAuto; using EpilogueSchedule = cutlass::epilogue::collective::EpilogueScheduleAuto; using TileShape = Shape<_128, _128, _128>; @@ -127,8 +148,8 @@ struct sm120_blockwise_fp8_config_default { }; template -struct sm120_blockwise_fp8_config_M64 { - // M in [1, 256]: use 64x128x128 tile with Pingpong schedule +struct sm120_blockwise_fp8_config_pingpong { + // use 64x128x128 tile with Pingpong schedule using KernelSchedule = cutlass::gemm::KernelTmaWarpSpecializedBlockwisePingpongSm120; using EpilogueSchedule = cutlass::epilogue::collective::EpilogueScheduleAuto; using TileShape = Shape<_64, _128, _128>; @@ -139,11 +160,24 @@ struct sm120_blockwise_fp8_config_M64 { EpilogueSchedule, KernelSchedule>; }; +template +struct sm120_blockwise_fp8_config_swapab { + // use 128x32x128 tile with Cooperative schedule + using KernelSchedule = cutlass::gemm::KernelTmaWarpSpecializedBlockwiseCooperativeSm120; + using EpilogueSchedule = cutlass::epilogue::collective::EpilogueScheduleAuto; + using TileShape = Shape<_128, _32, _128>; + using ClusterShape = Shape<_1, _1, _1>; + using Gemm = cutlass_3x_gemm_fp8_blockwise< + OutType, 128, 1, 128, TileShape, ClusterShape, + EpilogueSchedule, KernelSchedule, true>; +}; + template void cutlass_gemm_caller_blockwise(torch::stable::Tensor& out, torch::stable::Tensor const& a, torch::stable::Tensor const& b, torch::stable::Tensor const& a_scales, torch::stable::Tensor const& b_scales) { + static constexpr bool swap_ab = Gemm::swap_ab; using GemmKernel = typename Gemm::GemmKernel; using StrideA = typename Gemm::GemmKernel::StrideA; using StrideB = typename Gemm::GemmKernel::StrideB; @@ -167,11 +201,13 @@ void cutlass_gemm_caller_blockwise(torch::stable::Tensor& out, torch::stable::Te b_stride = cutlass::make_cute_packed_stride(StrideB{}, cute::make_shape(n, k, 1)); c_stride = - cutlass::make_cute_packed_stride(StrideC{}, cute::make_shape(m, n, 1)); + cutlass::make_cute_packed_stride(StrideC{}, swap_ab ? cute::make_shape(n, m, 1) : cute::make_shape(m, n, 1)); - LayoutSFA layout_SFA = + LayoutSFA layout_SFA = swap_ab ? + ScaleConfig::tile_atom_to_shape_SFA(make_shape(n, m, k, 1)) : ScaleConfig::tile_atom_to_shape_SFA(make_shape(m, n, k, 1)); - LayoutSFB layout_SFB = + LayoutSFB layout_SFB = swap_ab ? + ScaleConfig::tile_atom_to_shape_SFB(make_shape(n, m, k, 1)) : ScaleConfig::tile_atom_to_shape_SFB(make_shape(m, n, k, 1)); auto a_ptr = static_cast(a.data_ptr()); @@ -180,15 +216,24 @@ void cutlass_gemm_caller_blockwise(torch::stable::Tensor& out, torch::stable::Te auto b_scales_ptr = static_cast(b_scales.data_ptr()); typename GemmKernel::MainloopArguments mainloop_args{}; - mainloop_args.ptr_A = a_ptr; - mainloop_args.dA = a_stride; - mainloop_args.ptr_B = b_ptr; - mainloop_args.dB = b_stride; - mainloop_args.ptr_SFA = a_scales_ptr; mainloop_args.layout_SFA = layout_SFA; - mainloop_args.ptr_SFB = b_scales_ptr; mainloop_args.layout_SFB = layout_SFB; - auto prob_shape = cute::make_shape(m, n, k, 1); + if (swap_ab) { + mainloop_args.ptr_A = b_ptr; + mainloop_args.dA = b_stride; + mainloop_args.ptr_B = a_ptr; + mainloop_args.dB = a_stride; + mainloop_args.ptr_SFA = b_scales_ptr; + mainloop_args.ptr_SFB = a_scales_ptr; + } else { + mainloop_args.ptr_A = a_ptr; + mainloop_args.dA = a_stride; + mainloop_args.ptr_B = b_ptr; + mainloop_args.dB = b_stride; + mainloop_args.ptr_SFA = a_scales_ptr; + mainloop_args.ptr_SFB = b_scales_ptr; + } + auto prob_shape = swap_ab ? cute::make_shape(n, m, k, 1) : cute::make_shape(m, n, k, 1); auto c_ptr = static_cast(out.data_ptr()); typename GemmKernel::EpilogueArguments epilogue_args{ @@ -204,15 +249,26 @@ void cutlass_gemm_blockwise_sm120_fp8_dispatch(torch::stable::Tensor& out, torch::stable::Tensor const& a_scales, torch::stable::Tensor const& b_scales) { int M = a.size(0); - if (M <= 256) { - using Gemm = typename sm120_blockwise_fp8_config_M64::Gemm; + // more heuristic tuning can be done here by checking N/K dimensions as well + bool swap_ab = (M <= 64) || (M % 4 != 0); + + if (!swap_ab) { + if (M <= 256) { + using Gemm = typename sm120_blockwise_fp8_config_pingpong::Gemm; + return cutlass_gemm_caller_blockwise( + out, a, b, a_scales, b_scales); + } + // M > 256: use default 128x128x128 config with Cooperative (Auto) schedule + using Gemm = typename sm120_blockwise_fp8_config_default::Gemm; + return cutlass_gemm_caller_blockwise( + out, a, b, a_scales, b_scales); + } else { + // Swap A/B for small M to improve performance + // Use TILE_N=32 as the minimum compatible tile size. + using Gemm = typename sm120_blockwise_fp8_config_swapab::Gemm; return cutlass_gemm_caller_blockwise( out, a, b, a_scales, b_scales); } - // M > 256: use default 128x128x128 config with Cooperative (Auto) schedule - using Gemm = typename sm120_blockwise_fp8_config_default::Gemm; - return cutlass_gemm_caller_blockwise( - out, a, b, a_scales, b_scales); } } // namespace vllm diff --git a/csrc/moe/gpt_oss_router_gemm.cu b/csrc/moe/gpt_oss_router_gemm.cu deleted file mode 100644 index 0294cd36aa8..00000000000 --- a/csrc/moe/gpt_oss_router_gemm.cu +++ /dev/null @@ -1,144 +0,0 @@ -/* - * Adapted from - * https://github.com/NVIDIA/TensorRT-LLM/blob/v1.3.0rc7/cpp/tensorrt_llm/kernels/tinygemm2/tinygemm2_cuda.cu - * Copyright (c) 2025, The vLLM team. - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. - * All rights reserved. SPDX-License-Identifier: Apache-2.0 - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#include -#include -#include -#include -#include -#include "gpt_oss_router_gemm.cuh" - -void launch_gpt_oss_router_gemm(__nv_bfloat16* gA, __nv_bfloat16* gB, - __nv_bfloat16* gC, __nv_bfloat16* bias, - int batch_size, int output_features, - int input_features, cudaStream_t stream) { - static int const WARP_TILE_M = 16; - static int const TILE_M = WARP_TILE_M; - static int const TILE_N = 8; - static int const TILE_K = 64; - static int const STAGES = 16; - static int const STAGE_UNROLL = 4; - static bool const PROFILE = false; - - CUtensorMap weight_map{}; - CUtensorMap activation_map{}; - - constexpr uint32_t rank = 2; - uint64_t size[rank] = {(uint64_t)input_features, (uint64_t)output_features}; - uint64_t stride[rank - 1] = {input_features * sizeof(__nv_bfloat16)}; - uint32_t box_size[rank] = {TILE_K, TILE_M}; - uint32_t elem_stride[rank] = {1, 1}; - - CUresult res = cuTensorMapEncodeTiled( - &weight_map, CUtensorMapDataType::CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, rank, - gB, size, stride, box_size, elem_stride, - CUtensorMapInterleave::CU_TENSOR_MAP_INTERLEAVE_NONE, - CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_128B, - CUtensorMapL2promotion::CU_TENSOR_MAP_L2_PROMOTION_NONE, - CUtensorMapFloatOOBfill::CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TORCH_CHECK(res == CUDA_SUCCESS, - "cuTensorMapEncodeTiled failed for weight_map, error code=", - static_cast(res)); - - size[1] = batch_size; - box_size[1] = TILE_N; - - res = cuTensorMapEncodeTiled( - &activation_map, CUtensorMapDataType::CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, - rank, gA, size, stride, box_size, elem_stride, - CUtensorMapInterleave::CU_TENSOR_MAP_INTERLEAVE_NONE, - CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_128B, - CUtensorMapL2promotion::CU_TENSOR_MAP_L2_PROMOTION_NONE, - CUtensorMapFloatOOBfill::CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - TORCH_CHECK(res == CUDA_SUCCESS, - "cuTensorMapEncodeTiled failed for activation_map, error code=", - static_cast(res)); - - int smem_size = STAGES * STAGE_UNROLL * - (TILE_M * TILE_K * sizeof(__nv_bfloat16) + - TILE_N * TILE_K * sizeof(__nv_bfloat16)); - - gpuErrChk(cudaFuncSetAttribute( - gpt_oss_router_gemm_kernel, - cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); - - int tiles_m = (output_features + TILE_M - 1) / TILE_M; - int tiles_n = (batch_size + TILE_N - 1) / TILE_N; - - dim3 grid(tiles_m, tiles_n); - dim3 block(384); - - cudaLaunchConfig_t config; - cudaLaunchAttribute attrs[1]; - config.gridDim = grid; - config.blockDim = block; - config.dynamicSmemBytes = smem_size; - config.stream = stream; - config.attrs = attrs; - attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; - attrs[0].val.programmaticStreamSerializationAllowed = 1; - config.numAttrs = 1; - - cudaLaunchKernelEx( - &config, - &gpt_oss_router_gemm_kernel, - gC, gA, gB, bias, output_features, batch_size, input_features, weight_map, - activation_map, nullptr); -} - -void gpt_oss_router_gemm_cuda_forward(torch::Tensor& output, - torch::Tensor input, torch::Tensor weight, - torch::Tensor bias) { - auto const batch_size = input.size(0); - auto const input_dim = input.size(1); - auto const output_dim = weight.size(0); - - auto stream = at::cuda::getCurrentCUDAStream(); - - if (input.scalar_type() == at::ScalarType::BFloat16) { - launch_gpt_oss_router_gemm((__nv_bfloat16*)input.data_ptr(), - (__nv_bfloat16*)weight.data_ptr(), - (__nv_bfloat16*)output.mutable_data_ptr(), - (__nv_bfloat16*)bias.data_ptr(), batch_size, - output_dim, input_dim, stream); - } else { - throw std::invalid_argument("Unsupported dtype, only supports bfloat16"); - } -} - -void gpt_oss_router_gemm(torch::Tensor& output, torch::Tensor input, - torch::Tensor weight, torch::Tensor bias) { - TORCH_CHECK(input.dim() == 2, "input must be 2D"); - TORCH_CHECK(weight.dim() == 2, "weight must be 2D"); - TORCH_CHECK(bias.dim() == 1, "bias must be 1D"); - TORCH_CHECK(input.sizes()[1] == weight.sizes()[1], - "input.size(1) must match weight.size(1)"); - TORCH_CHECK(weight.sizes()[0] == bias.sizes()[0], - "weight.size(0) must match bias.size(0)"); - TORCH_CHECK(input.scalar_type() == at::ScalarType::BFloat16, - "input tensor must be bfloat16"); - TORCH_CHECK(weight.scalar_type() == at::ScalarType::BFloat16, - "weight tensor must be bfloat16"); - TORCH_CHECK(bias.scalar_type() == at::ScalarType::BFloat16, - "bias tensor must be bfloat16"); - gpt_oss_router_gemm_cuda_forward(output, input, weight, bias); -} diff --git a/csrc/moe/gpt_oss_router_gemm.cuh b/csrc/moe/gpt_oss_router_gemm.cuh deleted file mode 100644 index 5cc653f19cf..00000000000 --- a/csrc/moe/gpt_oss_router_gemm.cuh +++ /dev/null @@ -1,447 +0,0 @@ -/* - * Adapted from - * https://github.com/NVIDIA/TensorRT-LLM/blob/v1.3.0rc7/cpp/tensorrt_llm/kernels/tinygemm2/tinygemm2_kernel.cuh - * Copyright (c) 2025, The vLLM team. - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. - * All rights reserved. SPDX-License-Identifier: Apache-2.0 - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#include "cuda_bf16.h" -#include -#include -#include - -#include "cuda_pipeline.h" -#include -#include -#include -#include - -using barrier = cuda::barrier; -namespace cde = cuda::device::experimental; -namespace ptx = cuda::ptx; - -#define gpuErrChk(ans) \ - { \ - gpuAssert((ans), __FILE__, __LINE__); \ - } - -inline void gpuAssert(cudaError_t code, char const* file, int line, - bool abort = true) { - if (code != cudaSuccess) { - fprintf(stderr, "GPUassert: %s %s %d\n", cudaGetErrorString(code), file, - line); - if (abort) { - throw std::runtime_error(cudaGetErrorString(code)); - } - } -} - -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) -__device__ uint64_t gclock64() { - unsigned long long int rv; - asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(rv)); - return rv; -} - -__device__ void ldmatrix(__nv_bfloat16 rv[2], uint32_t smem_ptr) { - int dst; - asm volatile("ldmatrix.sync.aligned.x1.m8n8.shared.b16 {%0}, [%1];\n" - : "=r"(dst) - : "r"(smem_ptr)); - int* rvi = reinterpret_cast(&rv[0]); - rvi[0] = dst; -} - -__device__ void ldmatrix2(__nv_bfloat16 rv[4], uint32_t smem_ptr) { - int x, y; - asm volatile("ldmatrix.sync.aligned.x2.m8n8.shared.b16 {%0, %1}, [%2];\n" - : "=r"(x), "=r"(y) - : "r"(smem_ptr)); - - int* rvi = reinterpret_cast(&rv[0]); - rvi[0] = x; - rvi[1] = y; -} - -__device__ void ldmatrix4(__nv_bfloat16 rv[8], uint32_t smem_ptr) { - int x, y, z, w; - asm volatile( - "ldmatrix.sync.aligned.x4.m8n8.shared.b16 {%0, %1, %2, %3}, [%4];" - : "=r"(x), "=r"(y), "=r"(z), "=r"(w) - : "r"(smem_ptr)); - int* rvi = reinterpret_cast(&rv[0]); - rvi[0] = x; - rvi[1] = y; - rvi[2] = z; - rvi[3] = w; -} - -__device__ void HMMA_1688(float d[4], __nv_bfloat16 a[4], __nv_bfloat16 b[2], - float c[4]) { - uint32_t const* A = reinterpret_cast(&a[0]); - uint32_t const* B = reinterpret_cast(&b[0]); - float const* C = reinterpret_cast(&c[0]); - float* D = reinterpret_cast(&d[0]); - - asm volatile( - "mma.sync.aligned.m16n8k8.row.col.f32.bf16.bf16.f32 " - "{%0,%1,%2,%3}, {%4,%5}, {%6}, {%7,%8,%9,%10};\n" - : "=f"(D[0]), "=f"(D[1]), "=f"(D[2]), "=f"(D[3]) - : "r"(A[0]), "r"(A[1]), "r"(B[0]), "f"(C[0]), "f"(C[1]), "f"(C[2]), - "f"(C[3])); -} - -__device__ void HMMA_16816(float d[4], __nv_bfloat16 a[8], __nv_bfloat16 b[4], - float c[4]) { - uint32_t const* A = reinterpret_cast(&a[0]); - uint32_t const* B = reinterpret_cast(&b[0]); - float const* C = reinterpret_cast(&c[0]); - float* D = reinterpret_cast(&d[0]); - - asm volatile( - "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 " - "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n" - : "=f"(D[0]), "=f"(D[1]), "=f"(D[2]), "=f"(D[3]) - : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), - "f"(C[0]), "f"(C[1]), "f"(C[2]), "f"(C[3])); -} - -__device__ void bar_wait(uint32_t bar_ptr, int phase) { - asm volatile( - "{\n" - ".reg .pred P1;\n" - "LAB_WAIT:\n" - "mbarrier.try_wait.parity.shared::cta.b64 P1, [%0], %1;\n" - "@P1 bra.uni DONE;\n" - "bra.uni LAB_WAIT;\n" - "DONE:\n" - "}\n" ::"r"(bar_ptr), - "r"(phase)); -} - -__device__ bool bar_try_wait(uint32_t bar_ptr, int phase) { - uint32_t success; - #ifdef INTERNAL - asm volatile(".pragma \"set knob DontInsertYield\";\n" : : : "memory"); - #endif - asm volatile( - "{\n\t" - ".reg .pred P1; \n\t" - "mbarrier.try_wait.parity.shared::cta.b64 P1, [%1], %2; \n\t" - "selp.b32 %0, 1, 0, P1; \n\t" - "}" - : "=r"(success) - : "r"(bar_ptr), "r"(phase)); - return success; -} - -__device__ uint32_t elect_one_sync() { - uint32_t pred = 0; - uint32_t laneid = 0; - asm volatile( - "{\n" - ".reg .b32 %%rx;\n" - ".reg .pred %%px;\n" - " elect.sync %%rx|%%px, %2;\n" - "@%%px mov.s32 %1, 1;\n" - " mov.s32 %0, %%rx;\n" - "}\n" - : "+r"(laneid), "+r"(pred) - : "r"(0xFFFFFFFF)); - return pred; -} -#endif - -struct Profile { - uint64_t start; - uint64_t weight_load_start; - uint64_t act_load_start; - uint64_t compute_start; - uint64_t complete; -}; - -template -__global__ __launch_bounds__(384, 1) void gpt_oss_router_gemm_kernel( - __nv_bfloat16* output, __nv_bfloat16* weights, __nv_bfloat16* activations, - __nv_bfloat16* bias, int M, int N, int K, - const __grid_constant__ CUtensorMap weight_map, - const __grid_constant__ CUtensorMap activation_map, - Profile* profile = nullptr) { -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) - - if (PROFILE && threadIdx.x == 0 && blockIdx.y == 0) - profile[blockIdx.x].start = gclock64(); - - extern __shared__ __align__(128) char smem[]; - - __nv_bfloat16* sh_weights = (__nv_bfloat16*)&smem[0]; - __nv_bfloat16* sh_activations = - (__nv_bfloat16*)&smem[STAGES * STAGE_UNROLL * TILE_M * TILE_K * - sizeof(__nv_bfloat16)]; - - #pragma nv_diag_suppress static_var_with_dynamic_init - __shared__ barrier bar_wt_ready[STAGES]; - __shared__ barrier bar_act_ready[STAGES]; - __shared__ barrier bar_data_consumed[STAGES]; - - __shared__ float4 reduction_buffer[128]; - - __shared__ nv_bfloat16 sh_bias[TILE_M]; - - if (threadIdx.x == 0) { - for (int i = 0; i < STAGES; i++) { - init(&bar_wt_ready[i], 1); - init(&bar_act_ready[i], 1); - init(&bar_data_consumed[i], 32); - } - ptx::fence_proxy_async(ptx::space_shared); - asm volatile("prefetch.tensormap [%0];" - : - : "l"(reinterpret_cast(&weight_map)) - : "memory"); - asm volatile("prefetch.tensormap [%0];" - : - : "l"(reinterpret_cast(&activation_map)) - : "memory"); - } - __syncthreads(); - - int warp_id = threadIdx.x / 32; - int lane_id = threadIdx.x % 32; - - int phase = 0; - - int mib = blockIdx.x * TILE_M; - int ni = blockIdx.y * TILE_N; - - float accum[4]; - for (int i = 0; i < 4; i++) accum[i] = 0.f; - - int const K_LOOPS_DMA = - (K + 4 * TILE_K * STAGE_UNROLL - 1) / (4 * (TILE_K * STAGE_UNROLL)); - int const K_LOOPS_COMPUTE = K_LOOPS_DMA; - - // Data loading thread - if (warp_id >= 4 && elect_one_sync()) { - int stage = warp_id % 4; - - bool weight_warp = warp_id < 8; - if (!weight_warp) { - cudaGridDependencySynchronize(); - cudaTriggerProgrammaticLaunchCompletion(); - } - - for (int ki = 0; ki < K_LOOPS_DMA; ki++) { - int k = (ki * 4 + (warp_id % 4)) * TILE_K * STAGE_UNROLL; - - uint64_t desc_ptr_wt = reinterpret_cast(&weight_map); - uint64_t desc_ptr_act = reinterpret_cast(&activation_map); - - uint32_t bar_ptr_wt = __cvta_generic_to_shared(&bar_wt_ready[stage]); - uint32_t bar_ptr_act = __cvta_generic_to_shared(&bar_act_ready[stage]); - int bytes_wt = TILE_M * TILE_K * sizeof(__nv_bfloat16); - int bytes_act = TILE_N * TILE_K * sizeof(__nv_bfloat16); - - bar_wait(__cvta_generic_to_shared(&bar_data_consumed[stage]), phase ^ 1); - - if (weight_warp) - asm volatile("mbarrier.arrive.expect_tx.shared.b64 _, [%0], %1;" - : - : "r"(bar_ptr_wt), "r"(STAGE_UNROLL * bytes_wt)); - if (!weight_warp) - asm volatile("mbarrier.arrive.expect_tx.shared.b64 _, [%0], %1;" - : - : "r"(bar_ptr_act), "r"(STAGE_UNROLL * bytes_act)); - - if (PROFILE && blockIdx.y == 0 && ki == 0 && weight_warp) - profile[blockIdx.x].weight_load_start = gclock64(); - if (PROFILE && blockIdx.y == 0 && ki == 0 && !weight_warp) - profile[blockIdx.x].act_load_start = gclock64(); - - for (int i = 0; i < STAGE_UNROLL; i++) { - uint32_t smem_ptr_wt = __cvta_generic_to_shared( - &sh_weights[(stage * STAGE_UNROLL + i) * TILE_M * TILE_K]); - uint32_t crd0 = k + i * TILE_K; - uint32_t crd1 = mib; - if (weight_warp) - asm volatile( - "cp.async.bulk.tensor.2d.shared::cta.global.mbarrier::complete_" - "tx::bytes [%0], [%1, {%3,%4}], " - "[%2];" - : - : "r"(smem_ptr_wt), "l"(desc_ptr_wt), "r"(bar_ptr_wt), "r"(crd0), - "r"(crd1) - : "memory"); - - uint32_t smem_ptr_act = __cvta_generic_to_shared( - &sh_activations[(stage * STAGE_UNROLL + i) * TILE_N * TILE_K]); - crd0 = k + i * TILE_K; - crd1 = ni; - if (!weight_warp) - asm volatile( - "cp.async.bulk.tensor.2d.shared::cta.global.mbarrier::complete_" - "tx::bytes [%0], [%1, {%3,%4}], " - "[%2];" - : - : "r"(smem_ptr_act), "l"(desc_ptr_act), "r"(bar_ptr_act), - "r"(crd0), "r"(crd1) - : "memory"); - } - - stage += 4; - if (stage >= STAGES) { - stage = warp_id % 4; - phase ^= 1; - } - } - // Wait for pending loads to be consumed before exiting, to avoid race - for (int i = 0; i < (STAGES / 4) - 1; i++) { - bar_wait(__cvta_generic_to_shared(&bar_data_consumed[stage]), phase ^ 1); - stage += 4; - if (stage >= STAGES) { - stage = warp_id % 4; - phase ^= 1; - } - } - } - // Compute threads - else if (warp_id < 4) { - // Sneak the bias load into the compute warps since they're just waiting for - // stuff anyway - if (threadIdx.x < TILE_M) sh_bias[threadIdx.x] = bias[mib + threadIdx.x]; - - int stage = warp_id; - - int phase = 0; - int lane_id_div8 = lane_id / 8; - int lane_id_mod8 = lane_id % 8; - - int lane_row_offset_wt = (lane_id_div8 % 2) ? 8 : 0; - int lane_col_offset_wt = (lane_id_div8 / 2) ? 1 : 0; - - int row_wt = lane_id_mod8 + lane_row_offset_wt; - int row_act = lane_id_mod8; - - int row_offset_wt = (reinterpret_cast(sh_weights) / 128) % 8; - int row_offset_act = row_offset_wt; - - uint32_t bar_ptr_wt = __cvta_generic_to_shared(&bar_wt_ready[stage]); - uint32_t bar_ptr_act = __cvta_generic_to_shared(&bar_act_ready[stage]); - - bool weight_ready = bar_try_wait(bar_ptr_wt, phase); - bool act_ready = bar_try_wait(bar_ptr_act, phase); - - #pragma unroll 2 - for (int ki = 0; ki < K_LOOPS_COMPUTE; ki++) { - int next_stage = stage + 4; - int next_phase = phase; - if (next_stage >= STAGES) { - next_stage = warp_id; - next_phase ^= 1; - } - - while (!weight_ready || !act_ready) { - weight_ready = bar_try_wait(bar_ptr_wt, phase); - act_ready = bar_try_wait(bar_ptr_act, phase); - } - - if (PROFILE && blockIdx.y == 0 && threadIdx.x == 0 && ki == 0) - profile[blockIdx.x].compute_start = gclock64(); - - if (ki + 1 < K_LOOPS_COMPUTE) { - weight_ready = bar_try_wait( - __cvta_generic_to_shared(&bar_wt_ready[next_stage]), next_phase); - act_ready = bar_try_wait( - __cvta_generic_to_shared(&bar_act_ready[next_stage]), next_phase); - } - - #pragma unroll - for (int su = 0; su < STAGE_UNROLL; su++) { - __nv_bfloat16* ptr_weights = - &sh_weights[(stage * STAGE_UNROLL + su) * TILE_M * TILE_K]; - __nv_bfloat16* ptr_act = - &sh_activations[(stage * STAGE_UNROLL + su) * TILE_N * TILE_K]; - - #pragma unroll - for (int kii = 0; kii < TILE_K / 16; kii++) { - __nv_bfloat16 a[8]; - __nv_bfloat16 b[4]; - - int col = 2 * kii + lane_col_offset_wt; - int col_sw = ((row_wt + row_offset_wt) % 8) ^ col; - - ldmatrix4(a, __cvta_generic_to_shared( - &ptr_weights[row_wt * TILE_K + col_sw * 8])); - - col = 2 * kii + lane_id_div8; - col_sw = ((row_act + row_offset_act) % 8) ^ col; - - ldmatrix2(b, __cvta_generic_to_shared( - &ptr_act[row_act * TILE_K + 8 * col_sw])); - - HMMA_16816(accum, a, b, accum); - } - } - - uint32_t bar_c = __cvta_generic_to_shared(&bar_data_consumed[stage]); - asm volatile("mbarrier.arrive.shared::cta.b64 _, [%0];" : : "r"(bar_c)); - - stage = next_stage; - phase = next_phase; - } - - float4 accum4; - accum4.x = accum[0]; - accum4.y = accum[1]; - accum4.z = accum[2]; - accum4.w = accum[3]; - reduction_buffer[threadIdx.x] = accum4; - - __syncthreads(); - - if (warp_id == 0) { - int mi = mib + warp_id * WARP_TILE_M; - int tm = mi + lane_id / 4; - int tn = ni + 2 * (lane_id % 4); - - float4 accum1 = reduction_buffer[32 + threadIdx.x]; - float4 accum2 = reduction_buffer[64 + threadIdx.x]; - float4 accum3 = reduction_buffer[96 + threadIdx.x]; - - accum[0] = accum[0] + accum1.x + accum2.x + accum3.x; - accum[1] = accum[1] + accum1.y + accum2.y + accum3.y; - accum[2] = accum[2] + accum1.z + accum2.z + accum3.z; - accum[3] = accum[3] + accum1.w + accum2.w + accum3.w; - - float bias_lo = __bfloat162float(sh_bias[tm - mib]); - float bias_hi = __bfloat162float(sh_bias[tm + 8 - mib]); - - if (tn < N && tm < M) - output[tn * M + tm] = __float2bfloat16(accum[0] + bias_lo); - if (tn + 1 < N && tm < M) - output[(tn + 1) * M + tm] = __float2bfloat16(accum[1] + bias_lo); - if (tn < N && tm + 8 < M) - output[tn * M + tm + 8] = __float2bfloat16(accum[2] + bias_hi); - if (tn + 1 < N && tm + 8 < M) - output[(tn + 1) * M + tm + 8] = __float2bfloat16(accum[3] + bias_hi); - - if (PROFILE && blockIdx.y == 0 && threadIdx.x == 0) - profile[blockIdx.x].complete = gclock64(); - } - } -#endif // end if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) -} diff --git a/csrc/moe/moe_ops.h b/csrc/moe/moe_ops.h index de931dc7646..d8d962887da 100644 --- a/csrc/moe/moe_ops.h +++ b/csrc/moe/moe_ops.h @@ -70,8 +70,4 @@ torch::Tensor router_gemm_bf16_fp32(torch::Tensor const& input, // Supports num_tokens in [1, 16], num_experts in {256, 384}, hidden_dim = 7168 void dsv3_router_gemm(torch::Tensor& output, const torch::Tensor& mat_a, const torch::Tensor& mat_b); - -// gpt-oss optimized router GEMM kernel for SM90+ -void gpt_oss_router_gemm(torch::Tensor& output, torch::Tensor input, - torch::Tensor weight, torch::Tensor bias); #endif diff --git a/csrc/moe/torch_bindings.cpp b/csrc/moe/torch_bindings.cpp index 4cd74366ea4..7b627a6f876 100644 --- a/csrc/moe/torch_bindings.cpp +++ b/csrc/moe/torch_bindings.cpp @@ -132,12 +132,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) { // DeepSeek V3 optimized router GEMM for SM90+ m.def("dsv3_router_gemm(Tensor! output, Tensor mat_a, Tensor mat_b) -> ()"); // conditionally compiled so impl registration is in source file - - // gpt-oss optimized router GEMM kernel for SM90+ - m.def( - "gpt_oss_router_gemm(Tensor! output, Tensor input, Tensor weights, " - "Tensor bias) -> ()"); - m.impl("gpt_oss_router_gemm", torch::kCUDA, &gpt_oss_router_gemm); #endif } diff --git a/csrc/ops.h b/csrc/ops.h index e7886633edd..cc58422231f 100644 --- a/csrc/ops.h +++ b/csrc/ops.h @@ -57,7 +57,8 @@ void merge_attn_states( torch::Tensor& output, std::optional output_lse, const torch::Tensor& prefix_output, const torch::Tensor& prefix_lse, const torch::Tensor& suffix_output, const torch::Tensor& suffix_lse, - const std::optional prefill_tokens_with_context); + const std::optional prefill_tokens_with_context, + const std::optional& output_scale = std::nullopt); #ifndef USE_ROCM void convert_vertical_slash_indexes( torch::Tensor& block_count, // [BATCH, N_HEADS, NUM_ROWS] @@ -142,6 +143,14 @@ void rms_norm_per_block_quant(torch::Tensor& out, torch::Tensor const& input, std::optional residual, int64_t group_size, bool is_scale_transposed); +#ifndef USE_ROCM +void silu_and_mul_per_block_quant(torch::Tensor& out, + torch::Tensor const& input, + torch::Tensor& scales, int64_t group_size, + std::optional scale_ub, + bool is_scale_transposed); +#endif + void rotary_embedding(torch::Tensor& positions, torch::Tensor& query, std::optional key, int64_t head_size, torch::Tensor& cos_sin_cache, bool is_neox); diff --git a/csrc/quantization/fused_kernels/fused_silu_mul_block_quant.cu b/csrc/quantization/fused_kernels/fused_silu_mul_block_quant.cu new file mode 100644 index 00000000000..993ee641b5d --- /dev/null +++ b/csrc/quantization/fused_kernels/fused_silu_mul_block_quant.cu @@ -0,0 +1,169 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +#include +#include + +#include "../../dispatch_utils.h" +#include "quant_conversions.cuh" +#include "../w8a8/fp8/common.cuh" + +namespace vllm { + +// Logic: one thread block per (token, group) pair + +template +__global__ void silu_and_mul_per_block_quant_kernel( + scalar_out_t* __restrict__ out, // Output: [num_tokens, hidden_size] in + // FP8/INT8 + float* __restrict__ scales, // Output: [num_tokens, hidden_size / + // group_size] or [hidden_size / group_size, + // num_tokens] + scalar_t const* __restrict__ input, // Input: [num_tokens, hidden_size * 2] + float const* scale_ub, // Optional scale upper bound + int32_t const hidden_size // Output hidden size (input is 2x this) +) { + static_assert((group_size & (group_size - 1)) == 0, + "group_size must be a power of 2 for correct reduction"); + + // Grid: (num_tokens, num_groups) + int const token_idx = blockIdx.x; + int const group_idx = blockIdx.y; + int const tid = threadIdx.x; // tid in [0, group_size) + int const num_tokens = gridDim.x; + + // Input layout: [gate || up] concatenated along last dimension + int const input_stride = hidden_size * 2; + int const group_start = group_idx * group_size; + + // Pointers to this token's data + scalar_t const* token_input_gate = + input + token_idx * input_stride + group_start; + scalar_t const* token_input_up = token_input_gate + hidden_size; + scalar_out_t* token_output = out + token_idx * hidden_size + group_start; + + // Scale pointer for this group + int const num_groups = gridDim.y; + float* group_scale_ptr = is_scale_transposed + ? scales + group_idx * num_tokens + token_idx + : scales + token_idx * num_groups + group_idx; + + // Shared memory for reduction (compile-time sized) + __shared__ float shared_max[group_size]; + + // Step 1: Each thread loads one element, computes SiLU, stores in register + float gate = static_cast(token_input_gate[tid]); + float up = static_cast(token_input_up[tid]); + + // Compute SiLU(gate) * up + float sigmoid_gate = 1.0f / (1.0f + expf(-gate)); + float silu_gate = gate * sigmoid_gate; + float result = silu_gate * up; // Keep in register + + // Step 2: Reduce to find group max + shared_max[tid] = fabsf(result); + __syncthreads(); + +// Power-of-2 reduction (group_size guaranteed to be power of 2) +#pragma unroll + for (int stride = group_size / 2; stride > 0; stride >>= 1) { + if (tid < stride) { + shared_max[tid] = fmaxf(shared_max[tid], shared_max[tid + stride]); + } + __syncthreads(); + } + + // Step 3: Compute scale (thread 0), broadcast via shared memory + if (tid == 0) { + float group_max = shared_max[0]; + + float const quant_range = quant_type_max_v; + float group_scale = group_max / quant_range; + + // Apply scale upper bound if provided + if (scale_ub != nullptr) { + group_scale = fminf(group_scale, *scale_ub); + } + + // Use minimum safe scaling factor + group_scale = fmaxf(group_scale, min_scaling_factor::val()); + + // Store scale to global memory + *group_scale_ptr = group_scale; + + // Reuse shared_max[0] to broadcast scale + shared_max[0] = group_scale; + } + __syncthreads(); + + float group_scale = shared_max[0]; + + // Step 4: Quantize and write output + token_output[tid] = + vllm::ScaledQuant::quant_fn(result, group_scale); +} + +} // namespace vllm + +void silu_and_mul_per_block_quant(torch::Tensor& out, + torch::Tensor const& input, + torch::Tensor& scales, int64_t group_size, + std::optional scale_ub, + bool is_scale_transposed) { + static c10::ScalarType kFp8Type = is_fp8_ocp() + ? c10::ScalarType::Float8_e4m3fn + : c10::ScalarType::Float8_e4m3fnuz; + + TORCH_CHECK(out.dtype() == kFp8Type || out.dtype() == torch::kInt8); + TORCH_CHECK(out.is_contiguous() && input.is_contiguous()); + TORCH_CHECK( + input.dtype() == torch::kFloat16 || input.dtype() == torch::kBFloat16, + "Input must be FP16 or BF16"); + TORCH_CHECK(scales.dtype() == torch::kFloat32, "Scales must be FP32"); + TORCH_CHECK(group_size == 128 || group_size == 64, + "Unsupported group size: ", group_size); + + if (scale_ub.has_value()) { + TORCH_CHECK(out.dtype() == kFp8Type); + } + + int32_t hidden_size = out.size(-1); + auto num_tokens = input.size(0); + int32_t num_groups = hidden_size / group_size; + + TORCH_CHECK(input.size(-1) == hidden_size * 2, + "input last dim must be 2x output hidden_size"); + TORCH_CHECK(hidden_size % group_size == 0, + "hidden_size must be divisible by group_size"); + + const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + + dim3 grid(num_tokens, num_groups); + dim3 block(group_size); + + VLLM_DISPATCH_FLOATING_TYPES( + input.scalar_type(), "silu_and_mul_per_block_quant", [&] { + using scalar_in_t = scalar_t; + + VLLM_DISPATCH_QUANT_TYPES( + out.scalar_type(), "silu_and_mul_per_block_quant", [&] { + using scalar_out_t = scalar_t; + + VLLM_DISPATCH_GROUP_SIZE(group_size, gs, [&] { + VLLM_DISPATCH_BOOL(is_scale_transposed, transpose_scale, [&] { + vllm::silu_and_mul_per_block_quant_kernel< + scalar_in_t, scalar_out_t, transpose_scale, gs> + <<>>( + out.data_ptr(), + scales.data_ptr(), + input.data_ptr(), + scale_ub.has_value() ? scale_ub->data_ptr() + : nullptr, + hidden_size); + }); + }); + }); + }); +} \ No newline at end of file diff --git a/csrc/torch_bindings.cpp b/csrc/torch_bindings.cpp index 4f42477b2f6..3593f1d2225 100644 --- a/csrc/torch_bindings.cpp +++ b/csrc/torch_bindings.cpp @@ -2,7 +2,6 @@ #include "cuda_utils.h" #include "ops.h" #include "core/registration.h" - #include #include @@ -74,7 +73,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { " Tensor prefix_lse," " Tensor suffix_output," " Tensor suffix_lse," - " int!? prefill_tokens_with_context) -> ()"); + " int!? prefill_tokens_with_context," + " Tensor? output_scale=None) -> ()"); ops.impl("merge_attn_states", torch::kCUDA, &merge_attn_states); #ifndef USE_ROCM ops.def( @@ -233,6 +233,17 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { // Quantization ops #ifndef USE_ROCM + // Fused SiLU+Mul + per-block quantization + ops.def( + "silu_and_mul_per_block_quant(" + "Tensor! out, " + "Tensor input, " + "Tensor! scales, " + "int group_size, " + "Tensor? scale_ub=None, " + "bool is_scale_transposed=False) -> ()"); + ops.impl("silu_and_mul_per_block_quant", torch::kCUDA, + &silu_and_mul_per_block_quant); // DeepSeek V3 fused A GEMM (SM 9.0+, bf16 only, 1-16 tokens). ops.def( "dsv3_fused_a_gemm(Tensor! output, Tensor mat_a, Tensor mat_b) -> ()"); @@ -497,6 +508,12 @@ TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _cache_ops), cache_ops) { " int block_size_in_bytes, Tensor block_mapping) -> ()"); cache_ops.impl("swap_blocks", torch::kCUDA, &swap_blocks); + // Batch swap: submit all block copies in a single driver call. + cache_ops.def( + "swap_blocks_batch(Tensor src_ptrs, Tensor dst_ptrs," + " Tensor sizes) -> ()"); + cache_ops.impl("swap_blocks_batch", torch::kCPU, &swap_blocks_batch); + // Reshape the key and value tensors and cache them. cache_ops.def( "reshape_and_cache(Tensor key, Tensor value," diff --git a/docker/Dockerfile.cpu b/docker/Dockerfile.cpu index abae5d1bedc..11585343304 100644 --- a/docker/Dockerfile.cpu +++ b/docker/Dockerfile.cpu @@ -203,7 +203,8 @@ WORKDIR /vllm-workspace RUN --mount=type=cache,target=/root/.cache/uv \ --mount=type=cache,target=/root/.cache/ccache \ --mount=type=bind,from=vllm-build,src=/vllm-workspace/dist,target=dist \ - uv pip install dist/*.whl + uv pip install dist/*.whl && \ + uv pip install "vllm[audio]" # Add labels to document build configuration LABEL org.opencontainers.image.title="vLLM CPU" diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm index 6db6d8b8359..01a215a293d 100644 --- a/docker/Dockerfile.rocm +++ b/docker/Dockerfile.rocm @@ -390,7 +390,20 @@ ENV MIOPEN_DEBUG_CONV_GEMM=0 RUN mkdir src && mv vllm src/vllm # This is a workaround to ensure pytest exits with the correct status code in CI tests. -RUN echo "import os\n\ndef pytest_sessionfinish(session, exitstatus):\n os._exit(int(exitstatus))" > /vllm-workspace/conftest.py +RUN cat << 'EOF' > /vllm-workspace/conftest.py +import os + +_exit_code = 1 + +def pytest_sessionfinish(session, exitstatus): + global _exit_code + _exit_code = int(exitstatus) + +def pytest_unconfigure(config): + sys.stdout.flush() + sys.stderr.flush() + os._exit(_exit_code) +EOF # ----------------------- # Final vLLM image diff --git a/docs/design/attention_backends.md b/docs/design/attention_backends.md index 6c0ed46720e..242cc6b3b1e 100644 --- a/docs/design/attention_backends.md +++ b/docs/design/attention_backends.md @@ -167,7 +167,7 @@ Priority is **1 = highest** (tried first). | ------- | ------- | ------ | --------- | ----------- | ---------- | ---- | --------- | --- | --------------- | ------------ | | `CPU_ATTN` | | fp16, bf16, fp32 | `auto` | Any | 32, 64, 80, 96, 112, 128, 160, 192, 224, 256, 512 | ❌ | ❌ | ❌ | All | N/A | | `FLASHINFER` | Native† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ❌ | ❌ | ✅ | Decoder | 7.x-9.x | -| `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ✅ | ❌ | ✅ | Decoder | 10.x | +| `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ✅ | ❌ | ✅ | Decoder | 10.0 | | `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ✅ | All | ≥8.0 | | `FLASH_ATTN` | FA3* | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ❌ | ✅ | All | 9.x | | `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ✅ | All | ≥10.0 | @@ -177,7 +177,7 @@ Priority is **1 = highest** (tried first). | `ROCM_AITER_UNIFIED_ATTN` | | fp16, bf16 | `auto` | %16 | Any | ✅ | ✅ | ❌ | All | N/A | | `ROCM_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 32, 64, 80, 96, 128, 160, 192, 224, 256 | ❌ | ✅ | ❌ | Decoder, Encoder, Encoder Only | N/A | | `TREE_ATTN` | | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | 32, 64, 96, 128, 160, 192, 224, 256 | ❌ | ❌ | ❌ | Decoder | Any | -| `TRITON_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ✅ | ❌ | All | Any | +| `TRITON_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `int8_per_token_head`, `fp8_per_token_head` | %16 | Any | ✅ | ✅ | ❌ | All | Any | > **†** FlashInfer uses TRTLLM attention on Blackwell (SM100), which supports sinks. Disable via `--attention-config.use_trtllm_attention=0`. > diff --git a/docs/design/fusions.md b/docs/design/fusions.md index 28a29a7f351..3fe1f769bb7 100644 --- a/docs/design/fusions.md +++ b/docs/design/fusions.md @@ -22,6 +22,7 @@ or just on the low or high end. | ------------------------------------------------------------------------------ | ---------------------------- | ---------------------------------------------- | ------------------------------ | ------------------ | --------- | ------------ | | [AllReduce + RMSNorm](#allreduce--rmsnorm-fuse_allreduce_rms) | `fuse_allreduce_rms` | All-reduce → RMSNorm (+residual_add) (→ quant) | O2 (Hopper/Blackwell + TP > 1) | 5-20% | No | Low | | [Attention + Quant](#attention--quantization-fuse_attn_quant) | `fuse_attn_quant` | Attention output → FP8/NVFP4 quant | Off by default | 3-7% | Yes | Always | +| [MLA Attention + Quant](#attention--quantization-fuse_attn_quant) | `fuse_attn_quant` | MLA Attention output → FP8/NVFP4 quant | Off by default | TBD | Yes | Always | | [RoPE + KV-Cache Update](#rope--kv-cache-update-fuse_rope_kvcache) | `fuse_rope_kvcache` | Rotary embedding → KV cache write | O2 (ROCm/AITER only) | 2-4% | No | Low | | [QK Norm + RoPE](#qk-norm--rope-enable_qk_norm_rope_fusion) | `enable_qk_norm_rope_fusion` | Q/K RMSNorm → rotary embedding | Off by default | 2-3% | No | Low | | [Sequence Parallelism](#sequence-parallelism-enable_sp) | `enable_sp` | AllReduce → ReduceScatter + AllGather | Off by default | Prereq for AsyncTP | Yes | High | @@ -40,12 +41,13 @@ The table below lists the quantization schemes supported by each fusion on each | ---------------------------- | ---------------------------------------- | ---------------------------------------- | ---------------------------------------- | ------------- | ---------------------------------------- | | `fuse_allreduce_rms` | FP16/BF16, FP8 static, NVFP4 | FP16/BF16, FP8 static | — | — | — | | `fuse_attn_quant`\* | FP8 static\*, NVFP4\* | FP8 static\* | FP8 static\* | — | FP8 static\* | +| `fuse_attn_quant` (MLA)\* | FP8 static\*, NVFP4\* | FP8 static\* | FP8 static\* | — | FP8 static(untested)\* | | `fuse_rope_kvcache` | — | — | — | — | FP16/BF16 | | `enable_qk_norm_rope_fusion` | FP16/BF16 | FP16/BF16 | FP16/BF16† | FP16/BF16† | — | | `enable_sp` | FP16/BF16, FP8 static† | FP16/BF16, FP8 static | FP16/BF16† | FP16/BF16† | — | | `fuse_gemm_comms` | FP16/BF16, FP8 static† | FP16/BF16, FP8 static | FP16/BF16† | FP16/BF16† | — | | `fuse_norm_quant` | FP8 static, FP8 per-token, FP8 per-group | FP8 static, FP8 per-token, FP8 per-group | FP8 static, FP8 per-token, FP8 per-group | — | FP8 static, FP8 per-token, FP8 per-group | -| `fuse_act_quant` | FP8 static, NVFP4 | FP8 static | FP8 static | — | FP8 per-group | +| `fuse_act_quant` | FP8 static, NVFP4 | FP8 static, FP8 per-group (128/64) | FP8 static, FP8 per-group (128/64) | — | FP8 per-group | | `fuse_act_padding` | — | — | — | — | FP16/BF16 | \* `fuse_attn_quant` support depends on the attention backend in use; not all backends support @@ -129,7 +131,8 @@ on SM90/SM100) and configurable via `PassConfig.fi_allreduce_fusion_max_size_mb` explicitly. It requires the full model graph to be visible (Inductor partition or `splitting_ops=[]`). **What it fuses.** Fuses the attention output quantization directly after the attention computation, -eliminating a full-precision memory round-trip of the attention output. Patterns covered: +eliminating a full-precision memory round-trip of the attention output. This fusion supports both +standard `Attention` and `MLAAttention` (used by DeepSeek-V2/V3/R1 models). Patterns covered: `Attention → FP8 static quant`: @@ -142,11 +145,24 @@ eliminating a full-precision memory round-trip of the attention output. Patterns - `FLASHINFER`: CUDA sm100+ with FlashInfer installed +`MLAAttention → FP8 static quant` / `MLAAttention → NVFP4 dynamic quant`: + +The MLA fusion operates at the graph level on the `unified_mla_attention_with_output` op and works +with all MLA decode and prefill backend combinations. Unlike standard `Attention` backends (where +the kernel writes FP8 output directly), no MLA prefill or decode backend currently supports direct +FP8/FP4 output. The fusion writes to an intermediate buffer and quantizes in a separate step, so +there is no memory round-trip elimination yet. + +!!! info + The MLA attention fusion is not expected to yield a measurable speedup yet. + This will improve once MLA prefill/decode kernels support direct FP8/FP4 output. + Other attention backends do not support fused output quantization yet. **Code locations.** -- Pass: [`vllm/compilation/passes/fusion/attn_quant_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/attn_quant_fusion.py) +- Pass (Attention): [`vllm/compilation/passes/fusion/attn_quant_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/attn_quant_fusion.py) +- Pass (MLAAttention): [`vllm/compilation/passes/fusion/mla_attn_quant_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/mla_attn_quant_fusion.py) - Attention backends: [`vllm/v1/attention/backends/`](https://github.com/vllm-project/vllm/blob/main/vllm/v1/attention/backends/) ### RoPE + KV-Cache Update (`fuse_rope_kvcache`) @@ -305,6 +321,7 @@ Note that AITER fusions are in a separate pass in `vllm.compilation.passes.fusio Supported quantization scheme/hardware combinations: - FP8 static per-tensor: CUDA & HIP kernel +- FP8 dynamic per-group (128/64): CUDA kernel (sm89+, not active when DeepGemm is used on sm100+) - NVFP4 dynamic: CUDA sm100+ only with FlashInfer - FP8 per-token-group (128): ROCm AITER only @@ -313,6 +330,7 @@ Supported quantization scheme/hardware combinations: - Pass: [`vllm/compilation/passes/fusion/act_quant_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/act_quant_fusion.py) - ROCm AITER pass: [`vllm/compilation/passes/fusion/rocm_aiter_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/rocm_aiter_fusion.py) - CUDA/HIP kernels: [`csrc/quantization/`](https://github.com/vllm-project/vllm/blob/main/csrc/quantization/) +- Fused SiLU+Mul+BlockQuant kernel: [`csrc/quantization/fused_kernels/fused_silu_mul_block_quant.cu`](https://github.com/vllm-project/vllm/blob/main/csrc/quantization/fused_kernels/fused_silu_mul_block_quant.cu) ### RMSNorm + Padding (`fuse_act_padding`) diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index b5c329ef5e1..c987acfa3f9 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -481,6 +481,7 @@ th { | `Step3p5ForCausalLM` | Step-3.5-flash | `stepfun-ai/Step-3.5-Flash`, etc. | | ✅︎ | | `TeleChatForCausalLM` | TeleChat | `chuhac/TeleChat2-35B`, etc. | ✅︎ | ✅︎ | | `TeleChat2ForCausalLM` | TeleChat2 | `Tele-AI/TeleChat2-3B`, `Tele-AI/TeleChat2-7B`, `Tele-AI/TeleChat2-35B`, etc. | ✅︎ | ✅︎ | +| `TeleChat3ForCausalLM` | TeleChat3 | `Tele-AI/TeleChat3-36B-Thinking`, `Tele-AI/TeleChat3-Coder-36B-Thinking`, etc. | ✅︎ | ✅︎ | | `TeleFLMForCausalLM` | TeleFLM | `CofeAI/FLM-2-52B-Instruct-2407`, `CofeAI/Tele-FLM`, etc. | ✅︎ | ✅︎ | | `XverseForCausalLM` | XVERSE | `xverse/XVERSE-7B-Chat`, `xverse/XVERSE-13B-Chat`, `xverse/XVERSE-65B-Chat`, etc. | ✅︎ | ✅︎ | | `MiniMaxM1ForCausalLM` | MiniMax-Text | `MiniMaxAI/MiniMax-M1-40k`, `MiniMaxAI/MiniMax-M1-80k`, etc. | | | @@ -541,6 +542,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen | `BeeForConditionalGeneration` | Bee-8B | T + IE+ | `Open-Bee/Bee-8B-RL`, `Open-Bee/Bee-8B-SFT` | | ✅︎ | | `Blip2ForConditionalGeneration` | BLIP-2 | T + IE | `Salesforce/blip2-opt-2.7b`, `Salesforce/blip2-opt-6.7b`, etc. | ✅︎ | ✅︎ | | `ChameleonForConditionalGeneration` | Chameleon | T + I | `facebook/chameleon-7b`, etc. | | ✅︎ | +| `CheersForConditionalGeneration` | Cheers | T + I | `ai9stars/Cheers` | | ✅︎ | | `Cohere2VisionForConditionalGeneration` | Command A Vision | T + I+ | `CohereLabs/command-a-vision-07-2025`, etc. | | ✅︎ | | `DeepseekVLV2ForCausalLM` | DeepSeek-VL2 | T + I+ | `deepseek-ai/deepseek-vl2-tiny`, `deepseek-ai/deepseek-vl2-small`, `deepseek-ai/deepseek-vl2`, etc. | | ✅︎ | | `DeepseekOCRForCausalLM` | DeepSeek-OCR | T + I+ | `deepseek-ai/DeepSeek-OCR`, etc. | ✅︎ | ✅︎ | diff --git a/examples/offline_inference/vision_language.py b/examples/offline_inference/vision_language.py index 56154c12212..cc97126c2ac 100755 --- a/examples/offline_inference/vision_language.py +++ b/examples/offline_inference/vision_language.py @@ -179,6 +179,33 @@ def run_chameleon(questions: list[str], modality: str) -> ModelRequestData: ) +# Cheers +def run_cheers(questions: list[str], modality: str) -> ModelRequestData: + assert modality == "image" + model_name = "ai9stars/Cheers" + + engine_args = EngineArgs( + model=model_name, + trust_remote_code=True, + max_model_len=4096, + limit_mm_per_prompt={modality: 1}, + ) + + prompts = [ + ( + f"<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n" + f"<|im_start|>user\n<|image_pad|>{question}<|im_end|>\n" + f"<|im_start|>assistant\n" + ) + for question in questions + ] + + return ModelRequestData( + engine_args=engine_args, + prompts=prompts, + ) + + def run_command_a_vision(questions: list[str], modality: str) -> ModelRequestData: assert modality == "image" @@ -2140,6 +2167,7 @@ model_example_map = { "aria": run_aria, "aya_vision": run_aya_vision, "bagel": run_bagel, + "cheers": run_cheers, "bee": run_bee, "blip-2": run_blip2, "chameleon": run_chameleon, diff --git a/requirements/cuda.txt b/requirements/cuda.txt index 6d7f9693f75..cfee494b5a6 100644 --- a/requirements/cuda.txt +++ b/requirements/cuda.txt @@ -16,5 +16,5 @@ flashinfer-cubin==0.6.7 nvidia-cudnn-frontend>=1.13.0,<1.19.0 # QuACK and Cutlass DSL for FA4 (cute-DSL implementation) -nvidia-cutlass-dsl>=4.4.0.dev1 -quack-kernels>=0.2.7 +nvidia-cutlass-dsl>=4.4.2 +quack-kernels>=0.3.3 diff --git a/requirements/xpu.txt b/requirements/xpu.txt index 0cddd6dc6ab..26ba38f3efa 100644 --- a/requirements/xpu.txt +++ b/requirements/xpu.txt @@ -15,4 +15,4 @@ torch==2.10.0+xpu torchaudio torchvision -vllm_xpu_kernels @ https://github.com/vllm-project/vllm-xpu-kernels/releases/download/v0.1.4/vllm_xpu_kernels-0.1.4-cp38-abi3-manylinux_2_28_x86_64.whl +vllm_xpu_kernels @ https://github.com/vllm-project/vllm-xpu-kernels/releases/download/v0.1.5/vllm_xpu_kernels-0.1.5-cp38-abi3-manylinux_2_28_x86_64.whl diff --git a/setup.py b/setup.py index 74997702950..0d478c59a23 100644 --- a/setup.py +++ b/setup.py @@ -1063,7 +1063,10 @@ setup( # Optional deps for AMD FP4 quantization support "petit-kernel": ["petit-kernel"], # Optional deps for Helion kernel development - "helion": ["helion==0.3.2"], + # NOTE: When updating helion version, also update CI files: + # - .buildkite/test_areas/kernels.yaml + # - .buildkite/test-amd.yaml + "helion": ["helion==0.3.3"], # Optional deps for gRPC server (vllm serve --grpc) "grpc": ["smg-grpc-servicer[vllm] >= 0.5.0"], # Optional deps for OpenTelemetry tracing diff --git a/tests/compile/fullgraph/test_full_cudagraph.py b/tests/compile/fullgraph/test_full_cudagraph.py index c7c737371fc..95306e2062f 100644 --- a/tests/compile/fullgraph/test_full_cudagraph.py +++ b/tests/compile/fullgraph/test_full_cudagraph.py @@ -170,14 +170,3 @@ class TestFullCUDAGraph: piecewise_res.outputs[0].text.lower() == full_res.outputs[0].text.lower() ) - - -@pytest.mark.skipif(not current_platform.is_cuda(), reason="Skip if not cuda") -def test_full_cudagraph_with_invalid_backend(): - # Flex_Attention is not supported with full cuda graph - with pytest.raises(RuntimeError): - LLM( - model="Qwen/Qwen2-1.5B-Instruct", - compilation_config=CompilationConfig(cudagraph_mode="FULL"), - attention_config={"backend": "FLEX_ATTENTION"}, - ) diff --git a/tests/compile/fusions_e2e/conftest.py b/tests/compile/fusions_e2e/conftest.py index ca67d90d202..adc569192d1 100644 --- a/tests/compile/fusions_e2e/conftest.py +++ b/tests/compile/fusions_e2e/conftest.py @@ -84,10 +84,14 @@ def run_e2e_fusion_test(monkeypatch, caplog_mp_spawn): rocm_aiter_ops.refresh_env_variables() # Filter here to reduce code duplication + backend_name = attn_backend.backend.name.lower() requires_mla = "deepseek" in model_name.lower() - is_mla = "mla" in attn_backend.backend.name.lower() + is_mla = "mla" in backend_name + # DeepSeek V3.2 uses sparse MLA + requires_sparse = "v3.2" in model_name.lower() + is_sparse = "sparse" in backend_name - if requires_mla != is_mla: + if requires_mla != is_mla or requires_sparse != is_sparse: pytest.skip( f"Incompatible model '{model_name}' and " f"attention backend '{attn_backend.backend.name}'" @@ -231,7 +235,9 @@ def run_e2e_fusion_test(monkeypatch, caplog_mp_spawn): ) elif match_name == "attn_quant_fusion": - actual_match = match_table.get(match_name, 0) + actual_match = match_table.get( + "attn_quant_fusion", 0 + ) + match_table.get("mla_attn_quant_fusion", 0) assert actual_match == expected_matches * n_expected, ( f"Could not find {expected_matches * n_expected} " f"{match_name} (found {actual_match})." diff --git a/tests/compile/fusions_e2e/models.py b/tests/compile/fusions_e2e/models.py index 1a5f18cc0d5..8d830e88406 100644 --- a/tests/compile/fusions_e2e/models.py +++ b/tests/compile/fusions_e2e/models.py @@ -58,6 +58,15 @@ TRITON_MLA_ATTN = pytest.param( id="TRITON_MLA", ) +FLASHMLA_SPARSE_ATTN = pytest.param( + AttentionBackendCase(backend=AttentionBackendEnum.FLASHMLA_SPARSE), + id="FLASHMLA_SPARSE", + marks=pytest.mark.skipif( + not is_blackwell(), + reason="FlashMLA Sparse requires Blackwell", + ), +) + # Models llama3_8b = ModelFusionInfo( model_name="meta-llama/Llama-3.1-8B-Instruct", @@ -141,6 +150,18 @@ qwen3_a3b_fp8 = ModelFusionInfo( ), ) +deepseek_coder_v2_lite_fp8 = ModelFusionInfo( + model_name="RedHatAI/DeepSeek-Coder-V2-Lite-Instruct-FP8", + matches=lambda n_layers: Matches( + # first_k_dense_replace=1; MoE hides most rms+quant sites + rms_quant_fusion=1, + act_quant_fusion=min(1, n_layers), # dense layers only + # MLA attn + static FP8 quant + attn_quant_fusion=n_layers, + ar_rms_fusion=n_layers * 2 + 1, + ), +) + deepseek_v3_fp8 = ModelFusionInfo( model_name="deepseek-ai/DeepSeek-V3", matches=lambda n_layers: Matches( @@ -150,10 +171,9 @@ deepseek_v3_fp8 = ModelFusionInfo( # - post_attn_layernorm + MLP # 2 per MoE layer (remaining) due to MoE wrapping rms_quant_fusion=n_layers * 2 + min(3, n_layers), # add for 3 dense layers - # TODO silu+block quant - # act_quant_fusion=min(3, n_layers), # dense layers only - act_quant_fusion=0, - # MLA attn + quant not supported yet: + # silu+block quant + act_quant_fusion=min(3, n_layers), # dense layers only + # MLA attn + per-group FP8 quant not supported yet: # https://github.com/vllm-project/vllm/issues/35792 attn_quant_fusion=0, ar_rms_fusion=n_layers * 2 + 1, @@ -163,6 +183,16 @@ deepseek_v3_fp8 = ModelFusionInfo( ), ) +deepseek_v32_fp4 = ModelFusionInfo( + model_name="nvidia/DeepSeek-V3.2-NVFP4", + matches=lambda n_layers: Matches( + rms_quant_fusion=0, + act_quant_fusion=0, + attn_quant_fusion=n_layers, + ar_rms_fusion=n_layers * 2 + 1, + ), +) + gpt_oss_20b = ModelFusionInfo( model_name="openai/gpt-oss-20b", matches=lambda n_layers: Matches( diff --git a/tests/compile/fusions_e2e/test_tp1_quant.py b/tests/compile/fusions_e2e/test_tp1_quant.py index 8186ecbb49f..ded39939e16 100644 --- a/tests/compile/fusions_e2e/test_tp1_quant.py +++ b/tests/compile/fusions_e2e/test_tp1_quant.py @@ -18,11 +18,14 @@ from .common import ( from .models import ( FLASHINFER_ATTN, FLASHINFER_MLA_ATTN, + FLASHMLA_SPARSE_ATTN, ROCM_AITER_UNIFIED_ATTN, ROCM_ATTN, TRITON_ATTN, TRITON_MLA_ATTN, + deepseek_coder_v2_lite_fp8, deepseek_v3_fp8, + deepseek_v32_fp4, llama3_8b_fp4, llama3_8b_fp8, llama4_scout_fp4, @@ -37,6 +40,7 @@ from .models import ( (*llama3_8b_fp8, False), (*qwen3_a3b_fp8, False), (*qwen3_a3b_fp8, True), + (*deepseek_coder_v2_lite_fp8, False), (*deepseek_v3_fp8, False), (*deepseek_v3_fp8, True), pytest.param( @@ -144,9 +148,12 @@ def test_tp1_fp8_fusions( @pytest.mark.parametrize( "model_name, matches_fn, model_kwargs, hf_overrides", - [llama3_8b_fp4, llama4_scout_fp4], + [llama3_8b_fp4, llama4_scout_fp4, deepseek_v32_fp4], +) +@pytest.mark.parametrize( + "attn_backend", + [FLASHINFER_ATTN, FLASHMLA_SPARSE_ATTN], ) -@pytest.mark.parametrize("attn_backend", [FLASHINFER_ATTN]) @pytest.mark.parametrize("n_layers", [6]) @pytest.mark.parametrize("custom_ops", custom_ops_combos("rms_norm")) @pytest.mark.parametrize("inductor_graph_partition", INDUCTOR_GRAPH_PARTITION) diff --git a/tests/compile/fusions_e2e/test_tp2_ar_rms.py b/tests/compile/fusions_e2e/test_tp2_ar_rms.py index fa1ceb7f011..4b0a0859b02 100644 --- a/tests/compile/fusions_e2e/test_tp2_ar_rms.py +++ b/tests/compile/fusions_e2e/test_tp2_ar_rms.py @@ -18,8 +18,11 @@ from .common import ( from .models import ( FLASHINFER_ATTN, FLASHINFER_MLA_ATTN, + FLASHMLA_SPARSE_ATTN, TRITON_ATTN, + deepseek_coder_v2_lite_fp8, deepseek_v3_fp8, + deepseek_v32_fp4, gpt_oss_20b, llama3_8b, llama3_8b_fp4, @@ -37,7 +40,13 @@ pytestmark = pytest.mark.skipif(not current_platform.is_cuda(), reason="Only tes @pytest.mark.parametrize( "model_name, matches_fn, model_kwargs, hf_overrides", # qwen3 & dsv3 should still fuse AR+rms even though group quant is not yet supported - [llama3_8b_fp8, llama4_scout_fp8, qwen3_a3b_fp8, deepseek_v3_fp8], + [ + llama3_8b_fp8, + llama4_scout_fp8, + qwen3_a3b_fp8, + deepseek_coder_v2_lite_fp8, + deepseek_v3_fp8, + ], ) @pytest.mark.parametrize( "attn_backend", [TRITON_ATTN, FLASHINFER_ATTN, FLASHINFER_MLA_ATTN] @@ -104,9 +113,12 @@ def test_tp2_ar_rms_fp8_fusions( @multi_gpu_test(num_gpus=2) @pytest.mark.parametrize( "model_name, matches_fn, model_kwargs, hf_overrides", - [llama3_8b_fp4, llama4_scout_fp4], + [llama3_8b_fp4, llama4_scout_fp4, deepseek_v32_fp4], +) +@pytest.mark.parametrize( + "attn_backend", + [FLASHINFER_ATTN, FLASHMLA_SPARSE_ATTN], ) -@pytest.mark.parametrize("attn_backend", [FLASHINFER_ATTN]) @pytest.mark.parametrize("n_layers", [4]) @pytest.mark.parametrize("custom_ops", custom_ops_combos("rms_norm")) @pytest.mark.parametrize("inductor_graph_partition", INDUCTOR_GRAPH_PARTITION) diff --git a/tests/compile/passes/test_mla_attn_quant_fusion.py b/tests/compile/passes/test_mla_attn_quant_fusion.py new file mode 100644 index 00000000000..426fbb6a7e5 --- /dev/null +++ b/tests/compile/passes/test_mla_attn_quant_fusion.py @@ -0,0 +1,508 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import copy + +import pytest +import torch._dynamo + +from tests.compile.backend import LazyInitPass, TestBackend +from tests.utils import TestFP8Layer, flat_product +from tests.v1.attention.utils import BatchSpec, create_common_attn_metadata +from vllm._custom_ops import cutlass_scaled_fp4_mm, scaled_fp4_quant +from vllm.compilation.passes.fusion.matcher_utils import QUANT_OPS +from vllm.compilation.passes.fusion.mla_attn_quant_fusion import ( + MLA_ATTN_OP, + MLAAttnQuantFusionPass, +) +from vllm.compilation.passes.fx_utils import find_op_nodes +from vllm.compilation.passes.utility.noop_elimination import NoOpEliminationPass +from vllm.compilation.passes.utility.post_cleanup import PostCleanupPass +from vllm.config import ( + AttentionConfig, + CacheConfig, + CompilationConfig, + CompilationMode, + ModelConfig, + PassConfig, + SchedulerConfig, + VllmConfig, + set_current_vllm_config, +) +from vllm.forward_context import get_forward_context, set_forward_context +from vllm.model_executor.layers.attention import MLAAttention +from vllm.model_executor.layers.linear import ColumnParallelLinear +from vllm.model_executor.layers.quantization.fp8 import Fp8Config +from vllm.model_executor.layers.quantization.modelopt import ModelOptNvFp4Config +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + QuantKey, + kFp8StaticTensorSym, + kNvfp4Dynamic, +) +from vllm.platforms import current_platform +from vllm.v1.attention.backend import AttentionMetadata +from vllm.v1.attention.backends.registry import AttentionBackendEnum +from vllm.v1.kv_cache_interface import MLAAttentionSpec + +FP8_DTYPE = current_platform.fp8_dtype() +FP4_DTYPE = torch.uint8 + + +class MLAAttentionQuantPatternModel(torch.nn.Module): + """Base model for MLA AttentionQuantPattern fusion.""" + + def __init__( + self, + num_heads: int, + qk_nope_head_dim: int, + qk_rope_head_dim: int, + v_head_dim: int, + kv_lora_rank: int, + kv_cache_dtype: torch.dtype, + device: torch.device, + vllm_config: VllmConfig, + **kwargs, + ): + super().__init__() + self.num_heads = num_heads + self.qk_nope_head_dim = qk_nope_head_dim + self.qk_rope_head_dim = qk_rope_head_dim + self.qk_head_dim = qk_nope_head_dim + qk_rope_head_dim + self.v_head_dim = v_head_dim + self.kv_lora_rank = kv_lora_rank + self.output_dim = num_heads * v_head_dim + self.head_size = kv_lora_rank + qk_rope_head_dim + self.kv_cache_dtype = kv_cache_dtype + self.device = device + self.vllm_config = vllm_config + + # Create kv_b_proj (ColumnParallelLinear) on device. + # Reuse weights from prior model instance when available, because + # ColumnParallelLinear may get NaN from recycled CUDA memory after + # torch.compile runs in the same process. + kv_b_proj = ColumnParallelLinear( + input_size=kv_lora_rank, + output_size=num_heads * (qk_nope_head_dim + v_head_dim), + bias=False, + prefix="model.layers.0.self_attn.kv_b_proj", + ).to(device) + kv_b_proj_weight = kwargs.get("kv_b_proj_weight") + if kv_b_proj_weight is not None: + kv_b_proj.weight.data.copy_(kv_b_proj_weight) + elif kv_b_proj.weight.data.isnan().any(): + # Sanitize NaN from recycled CUDA memory + kv_b_proj.weight.data.normal_() + + # Create MLAAttention + self.mla_attn = MLAAttention( + num_heads=num_heads, + scale=1.0 / (self.qk_head_dim**0.5), + qk_nope_head_dim=qk_nope_head_dim, + qk_rope_head_dim=qk_rope_head_dim, + v_head_dim=v_head_dim, + q_lora_rank=None, + kv_lora_rank=kv_lora_rank, + kv_b_proj=kv_b_proj, + cache_config=vllm_config.cache_config, + quant_config=self.quant_config, + prefix="model.layers.0.self_attn.attn", + ) + self.mla_attn._k_scale = self.mla_attn._k_scale.to(device) + self.mla_attn._v_scale = self.mla_attn._v_scale.to(device) + + # Initialize W_UK_T and W_UV from kv_b_proj weights + self.mla_attn.process_weights_after_loading(torch.get_default_dtype()) + self.kv_b_proj_weight = kv_b_proj.weight.data.clone() + + self.block_size = 16 + + # Initialize MLA MetadataBuilder + self.builder = self.mla_attn.attn_backend.get_builder_cls()( + kv_cache_spec=MLAAttentionSpec( + block_size=self.block_size, + num_kv_heads=1, + head_size=self.head_size, + dtype=self.kv_cache_dtype, + ), + layer_names=[self.mla_attn.layer_name], + vllm_config=self.vllm_config, + device=self.device, + ) + + def build_attn_metadata(self, batch_size: int) -> AttentionMetadata: + """Initialize MLA attention metadata. + + NOTE: Uses decode-only batch (query_len=1 per request). The prefill + (forward_mha) path is not separately tested here because it requires + FlashAttention availability and different input tensor shapes. The + quant logic in forward_impl is identical for both paths — it quantizes + the full output[:num_actual_toks] buffer after both forward_mha and + forward_mqa have written their results. + """ + + batch_spec = BatchSpec(seq_lens=[1] * batch_size, query_lens=[1] * batch_size) + common_attn_metadata = create_common_attn_metadata( + batch_spec, self.block_size, self.device, arange_block_indices=True + ) + + max_blocks = (max(batch_spec.seq_lens) + self.block_size - 1) // self.block_size + num_blocks = batch_size * max_blocks + + # MLA KV cache is 3D: (num_blocks, block_size, head_size) + attn_backend = self.mla_attn.attn_backend + kv_cache_shape = attn_backend.get_kv_cache_shape( + num_blocks, self.block_size, 1, self.head_size + ) + try: + kv_cache_stride_order = attn_backend.get_kv_cache_stride_order() + except (AttributeError, NotImplementedError): + kv_cache_stride_order = tuple(range(len(kv_cache_shape))) + + ordered_shape = tuple(kv_cache_shape[i] for i in kv_cache_stride_order) + inv_order = [ + kv_cache_stride_order.index(i) for i in range(len(kv_cache_stride_order)) + ] + + raw_tensor = torch.zeros( + ordered_shape, dtype=self.kv_cache_dtype, device=self.device + ) + kv_cache = raw_tensor.permute(*inv_order) + + self.mla_attn.kv_cache = kv_cache + + self.attn_metadata = self.builder.build( + common_prefix_len=0, common_attn_metadata=common_attn_metadata + ) + + return self.attn_metadata + + +class TestMLAAttentionFp8StaticQuantPatternModel(MLAAttentionQuantPatternModel): + """Test model for MLA Attention + FP8 static quant fusion.""" + + quant_key = kFp8StaticTensorSym + quant_config = Fp8Config() + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + self.fp8_linear = TestFP8Layer( + weight_shape=(self.output_dim, self.output_dim), + activation_quant_key=self.quant_key, + weight_quant_key=self.quant_key, + device=self.device, + ) + + w = kwargs.get("w") + if w is not None: + self.fp8_linear.weight = w["weight"] + self.fp8_linear.weight_scale = w["wscale"] + self.fp8_linear.input_scale = w["scale"] + + self.w = { + "weight": self.fp8_linear.weight, + "wscale": self.fp8_linear.weight_scale, + "scale": self.fp8_linear.input_scale, + } + + def forward( + self, + q: torch.Tensor, + kv_c_normed: torch.Tensor, + k_pe: torch.Tensor, + ): + """Forward pass that creates the MLA attention + FP8 quant pattern.""" + attn_output = self.mla_attn( + q, + kv_c_normed, + k_pe, + output_shape=(q.shape[0], self.output_dim), + ) + return self.fp8_linear(attn_output) + + +class TestMLAAttentionNvfp4QuantPatternModel(MLAAttentionQuantPatternModel): + """Test model for MLA Attention + NVFP4 quant fusion.""" + + quant_key = kNvfp4Dynamic + quant_config = ModelOptNvFp4Config( + is_checkpoint_nvfp4_serialized=False, + kv_cache_quant_algo=None, + exclude_modules=[], + ) + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + self.w = kwargs.get( + "w", + { + "weight": torch.randint( + 256, + (self.output_dim, self.output_dim // 2), + dtype=FP4_DTYPE, + device=self.device, + ), + "wscale_swizzled": torch.randn( + self.output_dim, self.output_dim // 16 + ).to(dtype=FP8_DTYPE, device=self.device), + "wscale": torch.tensor([500], dtype=torch.float32, device=self.device), + "scale": torch.tensor([0.002], dtype=torch.float32, device=self.device), + }, + ) + + def forward( + self, + q: torch.Tensor, + kv_c_normed: torch.Tensor, + k_pe: torch.Tensor, + ): + """Forward pass that creates the MLA attention + NVFP4 quant pattern.""" + attn_output = self.mla_attn( + q, + kv_c_normed, + k_pe, + output_shape=(q.shape[0], self.output_dim), + ) + quant_output, output_block_scale = scaled_fp4_quant( + attn_output, 1 / self.w["scale"] + ) + return cutlass_scaled_fp4_mm( + a=quant_output, + b=self.w["weight"], + block_scale_a=output_block_scale, + block_scale_b=self.w["wscale_swizzled"], + alpha=self.w["scale"] * self.w["wscale"], + out_dtype=attn_output.dtype, + ) + + +def is_nvfp4_supported(): + return current_platform.has_device_capability(100) + + +# MLA test configuration +MLA_DIMS: list[tuple[int, int, int, int, int]] = [] +PATTERN_TEST_MODELS_MLA_FP8: list[tuple[str, type]] = [] +PATTERN_TEST_MODELS_MLA_FP4: list[tuple[str, type]] = [] +BACKENDS_MLA_FP8: list[AttentionBackendEnum] = [] +BACKENDS_MLA_FP4: list[AttentionBackendEnum] = [] + +if current_platform.is_cuda(): + # (num_heads, qk_nope_head_dim, qk_rope_head_dim, v_head_dim, kv_lora_rank) + MLA_DIMS = [(16, 128, 64, 128, 512)] + PATTERN_TEST_MODELS_MLA_FP8 = [ + ( + "deepseek-ai/DeepSeek-V2-Lite", + TestMLAAttentionFp8StaticQuantPatternModel, + ) + ] + PATTERN_TEST_MODELS_MLA_FP4 = [ + ( + "deepseek-ai/DeepSeek-V2-Lite", + TestMLAAttentionNvfp4QuantPatternModel, + ) + ] + BACKENDS_MLA_FP8 = [AttentionBackendEnum.TRITON_MLA] + BACKENDS_MLA_FP4 = [AttentionBackendEnum.TRITON_MLA] + + +@pytest.mark.parametrize( + "num_heads, qk_nope_head_dim, qk_rope_head_dim, v_head_dim, kv_lora_rank", + MLA_DIMS, +) +@pytest.mark.parametrize("batch_size", [7, 256] if current_platform.is_cuda() else [8]) +@pytest.mark.parametrize("dtype", [torch.bfloat16]) +@pytest.mark.parametrize( + "backend, model_name, model_class, custom_ops", + list( + flat_product( + BACKENDS_MLA_FP8, + PATTERN_TEST_MODELS_MLA_FP8, + ["+quant_fp8", "-quant_fp8"], + ) + ) + + list(flat_product(BACKENDS_MLA_FP4, PATTERN_TEST_MODELS_MLA_FP4, [""])), +) +@pytest.mark.skipif( + not current_platform.is_cuda_alike(), reason="Only test ROCm or CUDA" +) +@pytest.mark.skipif(not current_platform.supports_fp8(), reason="Need FP8") +def test_mla_attention_quant_pattern( + num_heads: int, + qk_nope_head_dim: int, + qk_rope_head_dim: int, + v_head_dim: int, + kv_lora_rank: int, + batch_size: int, + dtype: torch.dtype, + custom_ops: str, + model_name: str, + model_class: type[MLAAttentionQuantPatternModel], + backend: AttentionBackendEnum, + dist_init, + monkeypatch, + use_fresh_inductor_cache, +): + """Test MLA AttentionQuantPattern fusion pass""" + if ( + model_class is TestMLAAttentionNvfp4QuantPatternModel + and not is_nvfp4_supported() + ): + pytest.skip("NVFP4 is not supported on this GPU (requires SM 100+).") + + monkeypatch.setenv("VLLM_DISABLE_COMPILE_CACHE", "1") + + custom_ops_list = custom_ops.split(",") if custom_ops else [] + + device = torch.device("cuda:0") + torch.set_default_dtype(dtype) + torch.manual_seed(42) + + model_config = ModelConfig( + model=model_name, + max_model_len=2048, + dtype=dtype, + ) + vllm_config = VllmConfig( + model_config=model_config, + scheduler_config=SchedulerConfig( + max_num_seqs=1024, + max_model_len=model_config.max_model_len, + is_encoder_decoder=model_config.is_encoder_decoder, + ), + compilation_config=CompilationConfig( + mode=CompilationMode.VLLM_COMPILE, + custom_ops=custom_ops_list, + ), + cache_config=CacheConfig(cache_dtype="auto"), + attention_config=AttentionConfig(backend=backend), + ) + + # MLA inputs: q(B, N, qk_head_dim), kv_c_normed(B, L), k_pe(B, 1, R) + qk_head_dim = qk_nope_head_dim + qk_rope_head_dim + q = torch.randn(batch_size, num_heads, qk_head_dim, dtype=dtype, device=device) + kv_c_normed = torch.randn(batch_size, kv_lora_rank, dtype=dtype, device=device) + k_pe = torch.randn(batch_size, 1, qk_rope_head_dim, dtype=dtype, device=device) + + # Mark first dimension as dynamic + torch._dynamo.mark_dynamic(q, 0) + torch._dynamo.mark_dynamic(kv_c_normed, 0) + torch._dynamo.mark_dynamic(k_pe, 0) + + # Run model without fusion + vllm_config_unfused = copy.deepcopy(vllm_config) + with ( + set_current_vllm_config(vllm_config_unfused), + set_forward_context(attn_metadata=None, vllm_config=vllm_config_unfused), + ): + model_unfused = model_class( + num_heads=num_heads, + qk_nope_head_dim=qk_nope_head_dim, + qk_rope_head_dim=qk_rope_head_dim, + v_head_dim=v_head_dim, + kv_lora_rank=kv_lora_rank, + kv_cache_dtype=dtype, + device=device, + vllm_config=vllm_config_unfused, + ) + model_unfused = model_unfused.to(device) + # HACK: See #131044 + result_unfused_0 = model_unfused(q, kv_c_normed, k_pe) # noqa: F841 + + forward_ctx = get_forward_context() + forward_ctx.attn_metadata = model_unfused.build_attn_metadata(batch_size) + + compiled_unfused = torch.compile(model_unfused, fullgraph=True) + result_unfused = compiled_unfused(q, kv_c_normed, k_pe) + + # Run model with attn fusion enabled + vllm_config.compilation_config.pass_config = PassConfig( + fuse_attn_quant=True, eliminate_noops=True + ) + with ( + set_current_vllm_config(vllm_config), + set_forward_context(attn_metadata=None, vllm_config=vllm_config), + ): + model_fused = model_class( + num_heads=num_heads, + qk_nope_head_dim=qk_nope_head_dim, + qk_rope_head_dim=qk_rope_head_dim, + v_head_dim=v_head_dim, + kv_lora_rank=kv_lora_rank, + kv_cache_dtype=dtype, + device=device, + vllm_config=vllm_config, + w=model_unfused.w, + kv_b_proj_weight=model_unfused.kv_b_proj_weight, + ) + model_fused = model_fused.to(device) + + forward_ctx = get_forward_context() + forward_ctx.attn_metadata = model_fused.build_attn_metadata(batch_size) + + # Create test backend with fusion passes + noop_pass = NoOpEliminationPass(vllm_config) + attn_pass = LazyInitPass(MLAAttnQuantFusionPass, vllm_config) + cleanup_pass = PostCleanupPass(vllm_config) + + test_backend = TestBackend(noop_pass, attn_pass, cleanup_pass) + # HACK: See https://github.com/vllm-project/vllm/issues/31044 + result_fused_0 = model_fused(q, kv_c_normed, k_pe) # noqa: F841 + + compiled_fused = torch.compile( + model_fused, backend=test_backend, fullgraph=True + ) + + result_fused = compiled_fused(q, kv_c_normed, k_pe) + + # Check attn fusion support + quant_key: QuantKey = model_class.quant_key + attn_fusion_supported = [ + layer.impl.fused_output_quant_supported(quant_key) + for key, layer in vllm_config.compilation_config.static_forward_context.items() + if isinstance(layer, MLAAttention) + ] + assert sum(attn_fusion_supported) == len(attn_fusion_supported), ( + "All MLA layers should support attention fusion" + ) + + # Check quantization ops in the graph + quant_op = ( + torch.ops.aten.reciprocal + if "-quant_fp8" in custom_ops_list + else QUANT_OPS[quant_key] + ) + test_backend.check_before_ops([quant_op], fully_replaced=quant_key is kNvfp4Dynamic) + + assert attn_pass.pass_.matched_count == sum(attn_fusion_supported) + + # Check MLA attention ops in the graph + attn_nodes_pre = list(find_op_nodes(MLA_ATTN_OP, test_backend.graph_pre_pass)) + attn_nodes_post = list(find_op_nodes(MLA_ATTN_OP, test_backend.graph_post_pass)) + + assert len(attn_nodes_pre) > 0, "Should have MLA attention nodes before fusion" + assert len(attn_nodes_pre) == len(attn_nodes_post), ( + "Should have same number of MLA attention nodes before and after fusion" + ) + assert attn_nodes_pre[0].kwargs.get("output_scale") is None, ( + "MLA attention should not have output_scale before fusion" + ) + assert attn_nodes_post[0].kwargs.get("output_scale") is not None, ( + "MLA attention should have output_scale after fusion" + ) + + assert attn_nodes_pre[0].kwargs.get("output_block_scale") is None, ( + "MLA attention should not have output_block_scale before fusion" + ) + + if quant_key.dtype == FP8_DTYPE: + assert attn_nodes_post[0].kwargs.get("output_block_scale") is None, ( + "MLA attention should not have output_block_scale after FP8 fusion" + ) + elif quant_key.dtype == FP4_DTYPE: + assert attn_nodes_post[0].kwargs.get("output_block_scale") is not None, ( + "MLA attention should have output_block_scale after FP4 fusion" + ) + + # Check numerical correctness + torch.testing.assert_close(result_unfused, result_fused, atol=1e-2, rtol=1e-2) diff --git a/tests/compile/passes/test_silu_mul_quant_fusion.py b/tests/compile/passes/test_silu_mul_quant_fusion.py index a77b4e6de7b..383d59d03a7 100644 --- a/tests/compile/passes/test_silu_mul_quant_fusion.py +++ b/tests/compile/passes/test_silu_mul_quant_fusion.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import itertools +from functools import partial import pytest import torch @@ -34,13 +35,16 @@ from vllm.model_executor.kernels.linear import ( ROCmFP8ScaledMMLinearKernel, ) from vllm.model_executor.layers.activation import SiluAndMul +from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8 from vllm.model_executor.layers.quantization.utils.fp8_utils import W8A8BlockFp8LinearOp from vllm.model_executor.layers.quantization.utils.quant_utils import ( GroupShape, + kFp8Dynamic128Sym, kFp8StaticTensorSym, kNvfp4Dynamic, ) from vllm.platforms import current_platform +from vllm.utils.deep_gemm import is_deep_gemm_supported FP8_DTYPE = current_platform.fp8_dtype() FP4_DTYPE = torch.uint8 @@ -165,6 +169,48 @@ class TestSiluMulGroupFp8QuantModel(torch.nn.Module): return [torch.ops.vllm.rocm_aiter_act_mul_and_fp8_group_quant] +class TestSiluMulBlockQuantModel(torch.nn.Module): + quant_key = kFp8Dynamic128Sym + + def __init__(self, hidden_size: int, is_scale_transposed: bool = False, **kwargs): + super().__init__() + self.silu_and_mul = SiluAndMul() + self.is_scale_transposed = is_scale_transposed + self.quant_fp8 = QuantFP8( + static=False, + group_shape=GroupShape(1, 128), + column_major_scales=is_scale_transposed, + compile_native=False, + ) + + self.enable_silu_mul_custom_op = self.silu_and_mul.enabled() + self.enable_quant_fp8_custom_op = self.quant_fp8.enabled() + + def forward(self, x): + y = self.silu_and_mul(x) + out, scale = self.quant_fp8(y) + group_size = self.quant_key.scale.group_shape[1] + scale_expanded = scale.repeat_interleave(group_size, dim=1) + dequant = out.to(dtype=torch.float32) * scale_expanded + return (dequant,) + + def ops_in_model_before(self): + ops = [] + if self.enable_silu_mul_custom_op: + ops.append(SILU_MUL_OP) + # When silu custom op is disabled, aten.mul.Tensor also appears + # in dequant code, so we skip checking it to avoid false positives. + ops.append( + QUANT_OPS[self.quant_key] + if self.enable_quant_fp8_custom_op + else torch.ops.aten.reciprocal.default + ) + return ops + + def ops_in_model_after(self): + return [FUSED_OPS[self.quant_key]] + + ROCM_KERNELS = [ROCmFP8ScaledMMLinearKernel, PerTensorTorchFP8ScaledMMLinearKernel] CUDA_KERNELS = [ FlashInferFP8ScaledMMLinearKernel, @@ -200,6 +246,19 @@ TEST_KERNELS = ROCM_KERNELS if current_platform.is_rocm() else CUDA_KERNELS not current_platform.is_rocm(), reason="ROCm only" ), ), + # Block quant fusion for per-group FP8 (CUDA only). + *[ + pytest.param( + partial(TestSiluMulBlockQuantModel, is_scale_transposed=transposed), + True, + None, + marks=pytest.mark.skipif( + not current_platform.is_cuda(), reason="CUDA only" + ), + id=f"TestSiluMulBlockQuant-transposed={transposed}", + ) + for transposed in [False, True] + ], ], ) @pytest.mark.skipif( @@ -213,6 +272,7 @@ def test_fusion_silu_and_mul_quant( TestSiluMulFp8QuantModel | TestSiluMulNvfp4QuantModel | TestSiluMulGroupFp8QuantModel + | TestSiluMulBlockQuantModel ], enable_silu_mul_custom_op: bool, enable_quant_fp8_custom_op: bool, @@ -223,6 +283,12 @@ def test_fusion_silu_and_mul_quant( pytest.skip("NVFP4 is not supported on this GPU.") if model_class is TestSiluMulGroupFp8QuantModel and not IS_AITER_FOUND: pytest.skip("AITER is not supported on this GPU.") + if ( + isinstance(model_class, partial) + and model_class.func is TestSiluMulBlockQuantModel + and is_deep_gemm_supported() + ): + pytest.skip("SiluMul+BlockQuant fusion not applicable with DeepGemm") torch.set_default_device("cuda") torch.set_default_dtype(dtype) @@ -269,11 +335,13 @@ def test_fusion_silu_and_mul_quant( result2 = model2(x) # Check that it gives the same answer - if model_class == TestSiluMulFp8QuantModel: + if isinstance(model, TestSiluMulFp8QuantModel): atol, rtol = 1e-3, 1e-3 - elif model_class == TestSiluMulNvfp4QuantModel: + elif isinstance(model, TestSiluMulNvfp4QuantModel): atol, rtol = 1e-1, 1e-1 - elif model_class == TestSiluMulGroupFp8QuantModel: + elif isinstance( + model, (TestSiluMulGroupFp8QuantModel, TestSiluMulBlockQuantModel) + ): atol, rtol = 5e-2, 5e-2 torch.testing.assert_close( diff --git a/tests/distributed/conftest.py b/tests/distributed/conftest.py index 9c146a3323d..da661c5e13b 100644 --- a/tests/distributed/conftest.py +++ b/tests/distributed/conftest.py @@ -1,5 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import os import random import msgspec @@ -166,3 +167,31 @@ class MockSubscriber: self.sub.close() for replay in self.replay_sockets: replay.close() + + +@pytest.fixture +def enable_ray_v2_backend(): + """Set env vars for the Ray V2 executor backend and shut down Ray + between tests.""" + import ray + + saved = { + "VLLM_USE_RAY_V2_EXECUTOR_BACKEND": os.environ.get( + "VLLM_USE_RAY_V2_EXECUTOR_BACKEND" + ), + "VLLM_ENABLE_V1_MULTIPROCESSING": os.environ.get( + "VLLM_ENABLE_V1_MULTIPROCESSING" + ), + } + os.environ["VLLM_USE_RAY_V2_EXECUTOR_BACKEND"] = "1" + os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0" + if ray.is_initialized(): + ray.shutdown() + try: + yield + finally: + if ray.is_initialized(): + ray.shutdown() + os.environ.update({k: v for k, v in saved.items() if v is not None}) + for key in (k for k, v in saved.items() if v is None): + os.environ.pop(key, None) diff --git a/tests/distributed/test_mq_tcp_multinode.py b/tests/distributed/test_mq_tcp_multinode.py new file mode 100644 index 00000000000..135ef11d7fa --- /dev/null +++ b/tests/distributed/test_mq_tcp_multinode.py @@ -0,0 +1,119 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Multi-node integration test for MessageQueue TCP fallback. + +Verifies that when writer and readers span separate nodes (Docker containers +with isolated /dev/shm), `create_from_process_group` correctly detects +cross-node ranks via `in_the_same_node_as()` and falls back to ZMQ TCP +transport — and that data actually arrives. +""" + +import numpy as np +import torch.distributed as dist + +from vllm.distributed.device_communicators.shm_broadcast import MessageQueue +from vllm.distributed.parallel_state import in_the_same_node_as + + +def main(): + dist.init_process_group(backend="gloo") + + rank = dist.get_rank() + world_size = dist.get_world_size() + assert world_size >= 2, ( + f"Need at least 2 ranks across nodes, got world_size={world_size}" + ) + + # Verify that in_the_same_node_as detects cross-node correctly + status = in_the_same_node_as(dist.group.WORLD, source_rank=0) + local_count = sum(status) + print( + f"[Rank {rank}] in_the_same_node_as(source=0): {status} " + f"(local={local_count}/{world_size})" + ) + # With 2 Docker containers (1 proc each), rank 0 and rank 1 + # should be on different nodes. + assert local_count < world_size, ( + f"Expected cross-node ranks but all {world_size} ranks appear local." + ) + + # Create MessageQueue + writer_rank = 0 + mq = MessageQueue.create_from_process_group( + dist.group.WORLD, + max_chunk_bytes=1024 * 1024, # 1 MiB + max_chunks=10, + writer_rank=writer_rank, + ) + + # Verify the transport path selection + if rank == writer_rank: + print( + f"[Rank {rank}] Writer: n_local_reader={mq.n_local_reader}, " + f"n_remote_reader={mq.n_remote_reader}" + ) + assert mq.n_remote_reader > 0, ( + "Writer should have at least 1 remote (TCP) reader in a multi-node setup." + ) + else: + if status[rank]: + assert mq._is_local_reader, ( + f"Rank {rank} is on the same node as writer but is not a local reader." + ) + print(f"[Rank {rank}] Reader: local (shared memory)") + else: + assert mq._is_remote_reader, ( + f"Rank {rank} is on a different node but is not a remote (TCP) reader." + ) + print(f"[Rank {rank}] Reader: remote (TCP)") + + # Test data transfer: simple objects + dist.barrier() + if rank == writer_rank: + mq.enqueue("hello_from_node0") + else: + msg = mq.dequeue(timeout=10) + assert msg == "hello_from_node0" + dist.barrier() + print(f"[Rank {rank}] Simple object test passed") + + # Test data transfer: numpy arrays + np.random.seed(42) + arrays = [ + np.random.randint(0, 100, size=np.random.randint(100, 5000)) for _ in range(100) + ] + + dist.barrier() + if rank == writer_rank: + for arr in arrays: + mq.enqueue(arr) + else: + for i, expected in enumerate(arrays): + received = mq.dequeue(timeout=10) + assert np.array_equal(expected, received), ( + f"Array mismatch at index {i}: " + f"expected shape {expected.shape}, got shape {received.shape}" + ) + dist.barrier() + print(f"[Rank {rank}] Numpy array test passed") + + # Test data transfer: large payload (> max_chunk_bytes) + dist.barrier() + big_array = np.zeros(200_000, dtype=np.int64) # ~1.6 MiB > 1 MiB chunk + if rank == writer_rank: + mq.enqueue(big_array) + else: + received = mq.dequeue(timeout=10) + assert np.array_equal(big_array, received) + dist.barrier() + print(f"[Rank {rank}] Large payload test passed") + + # Done -- cleanup + dist.barrier() + print(f"[Rank {rank}] All MessageQueue TCP multi-node tests passed!") + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/tests/distributed/test_ray_v2_executor.py b/tests/distributed/test_ray_v2_executor.py new file mode 100644 index 00000000000..5daec22df6f --- /dev/null +++ b/tests/distributed/test_ray_v2_executor.py @@ -0,0 +1,345 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +""" +Integration tests for RayExecutorV2 at the executor level. +Validates executor initialization, placement group support, RPC calls, +and distributed execution with various TP/PP configurations. +""" + +import gc +import threading +from unittest.mock import patch + +import pytest +import ray + +from vllm import LLM +from vllm.config import VllmConfig +from vllm.engine.arg_utils import EngineArgs +from vllm.v1.executor.ray_executor_v2 import RayExecutorV2 + +pytestmark = pytest.mark.usefixtures("enable_ray_v2_backend") + +MODEL = "facebook/opt-125m" + + +def create_vllm_config( + tensor_parallel_size: int = 1, + pipeline_parallel_size: int = 1, + max_model_len: int = 256, + gpu_memory_utilization: float = 0.3, + placement_group=None, +) -> VllmConfig: + engine_args = EngineArgs( + model=MODEL, + tensor_parallel_size=tensor_parallel_size, + pipeline_parallel_size=pipeline_parallel_size, + max_model_len=max_model_len, + gpu_memory_utilization=gpu_memory_utilization, + distributed_executor_backend="ray", + enforce_eager=True, + ) + vllm_config = engine_args.create_engine_config() + + if placement_group is not None: + vllm_config.parallel_config.placement_group = placement_group + + return vllm_config + + +def ensure_ray_initialized(): + if not ray.is_initialized(): + ray.init(ignore_reinit_error=True) + + +@pytest.fixture +def create_placement_group(request): + ensure_ray_initialized() + num_gpus = request.param + bundles = [{"GPU": 1, "CPU": 1} for _ in range(num_gpus)] + pg = ray.util.placement_group(bundles, strategy="PACK") + ray.get(pg.ready()) + yield pg + ray.util.remove_placement_group(pg) + + +@pytest.fixture +def executor(request): + """Create a RayExecutorV2 and shut it down after the test.""" + executor = RayExecutorV2(vllm_config=request.param) + yield executor + executor.shutdown() + + +def assert_executor(executor, tp_size, pp_size): + """Common assertions for executor initialization tests.""" + world_size = tp_size * pp_size + expected_output_rank = (pp_size - 1) * tp_size + + assert executor.world_size == world_size + assert len(executor.ray_worker_handles) == world_size + assert len(executor.response_mqs) == world_size + assert executor._get_output_rank() == expected_output_rank + + if pp_size > 1: + assert executor.max_concurrent_batches == pp_size + + executor.check_health() + assert not executor.is_failed + + ranks = sorted(h.rank for h in executor.ray_worker_handles) + assert ranks == list(range(world_size)) + + for handle in executor.ray_worker_handles: + assert handle.node_id is not None + + +@pytest.mark.parametrize("tp_size, pp_size", [(1, 1), (2, 1), (4, 1), (2, 2)]) +def test_ray_v2_executor(tp_size, pp_size): + """Validate RayExecutorV2 with various TP/PP configs.""" + vllm_config = create_vllm_config( + tensor_parallel_size=tp_size, + pipeline_parallel_size=pp_size, + ) + executor = RayExecutorV2(vllm_config=vllm_config) + try: + assert_executor(executor, tp_size, pp_size) + finally: + executor.shutdown() + + +@pytest.mark.parametrize( + "tp_size, pp_size, create_placement_group", + [(2, 1, 2), (4, 1, 4), (2, 2, 4)], + indirect=["create_placement_group"], +) +def test_ray_v2_executor_pg(tp_size, pp_size, create_placement_group): + """Validate RayExecutorV2 with various TP/PP configs using external PG.""" + vllm_config = create_vllm_config( + tensor_parallel_size=tp_size, + pipeline_parallel_size=pp_size, + placement_group=create_placement_group, + ) + executor = RayExecutorV2(vllm_config=vllm_config) + try: + assert_executor(executor, tp_size, pp_size) + finally: + executor.shutdown() + + +@pytest.mark.parametrize( + "executor", + [create_vllm_config(tensor_parallel_size=2)], + indirect=True, +) +def test_ray_v2_executor_failure_callback(executor): + """Validate failure callback registration.""" + callback_invoked = False + + def test_callback(): + nonlocal callback_invoked + callback_invoked = True + + executor.register_failure_callback(test_callback) + assert not callback_invoked + + executor.is_failed = True + executor.register_failure_callback(test_callback) + assert callback_invoked + + +@pytest.mark.parametrize( + "executor", + [create_vllm_config(tensor_parallel_size=2)], + indirect=True, +) +def test_ray_v2_executor_collective_rpc(executor): + """Validate collective RPC calls through MessageQueue.""" + executor.check_health() + assert not executor.is_failed + assert executor.rpc_broadcast_mq is not None + + +@pytest.mark.parametrize( + "executor", + [create_vllm_config(tensor_parallel_size=2)], + indirect=True, +) +def test_ray_v2_executor_driver_node_rank_0(executor): + """Validate that driver node workers get the lowest ranks.""" + driver_node = ray.get_runtime_context().get_node_id() + + for handle in executor.ray_worker_handles: + assert handle.node_id == driver_node + + rank0_handle = next(h for h in executor.ray_worker_handles if h.rank == 0) + assert rank0_handle.node_id == driver_node + + +@pytest.mark.parametrize( + "executor", + [create_vllm_config(tensor_parallel_size=2)], + indirect=True, +) +def test_ray_v2_executor_worker_death(executor): + """Validate executor detects worker death via ray.wait().""" + callback_event = threading.Event() + + def on_failure(): + callback_event.set() + + executor.register_failure_callback(on_failure) + assert not executor.is_failed + + # Kill one worker actor externally + victim = executor.ray_worker_handles[1].actor + ray.kill(victim, no_restart=True) + + # Monitor thread should detect the death and invoke callback + assert callback_event.wait(timeout=30) + assert executor.is_failed + assert executor.shutting_down + + +def test_ray_v2_executor_shutdown(): + """Validate graceful shutdown: ray.kill() terminates all worker actors.""" + executor = RayExecutorV2(vllm_config=create_vllm_config(tensor_parallel_size=2)) + assert executor.rpc_broadcast_mq is not None + assert len(executor.response_mqs) == executor.world_size + + actors = [h.actor for h in executor.ray_worker_handles] + executor.shutdown() + + for actor in actors: + with pytest.raises(ray.exceptions.RayActorError): + ray.get(actor.wait_for_init.remote(), timeout=5) + + assert executor.rpc_broadcast_mq is None + assert len(executor.response_mqs) == 0 + + +@pytest.mark.parametrize( + "executor", + [create_vllm_config(tensor_parallel_size=2)], + indirect=True, +) +def test_ray_v2_run_refs_stored_for_monitoring(executor): + """Validate worker handles store run_ref for monitoring.""" + for handle in executor.ray_worker_handles: + assert handle.run_ref is not None + ready, _ = ray.wait([handle.run_ref], timeout=0) + assert len(ready) == 0, "run_ref should be pending" + + +@pytest.mark.parametrize("tp_size, pp_size", [(2, 1), (2, 2)]) +def test_ray_v2_single_node_generation(tp_size, pp_size): + """End-to-end LLM generation with RayExecutorV2.""" + + llm = LLM( + model=MODEL, + tensor_parallel_size=tp_size, + pipeline_parallel_size=pp_size, + distributed_executor_backend="ray", + enforce_eager=True, + max_model_len=256, + gpu_memory_utilization=0.3, + ) + try: + prompts = [ + "Hello, my name is", + "The capital of France is", + "The future of AI is", + ] + outputs = llm.generate(prompts) + + assert len(outputs) == len(prompts) + for output in outputs: + assert len(output.outputs) > 0 + assert len(output.outputs[0].text) > 0 + finally: + llm.llm_engine.model_executor.shutdown() + del llm + gc.collect() + + +@pytest.mark.parametrize( + "bundle_indices, expected_bundle_ids, create_placement_group", + [("2,3", [2, 3], 4), ("3,2", [3, 2], 4)], + indirect=["create_placement_group"], +) +def test_ray_v2_bundle_indices_env( + bundle_indices, expected_bundle_ids, create_placement_group, monkeypatch +): + """Validate explicit VLLM_RAY_BUNDLE_INDICES bundle placement.""" + monkeypatch.setenv("VLLM_RAY_BUNDLE_INDICES", bundle_indices) + vllm_config = create_vllm_config( + tensor_parallel_size=2, + placement_group=create_placement_group, + ) + executor = RayExecutorV2(vllm_config=vllm_config) + try: + actual = [ + h.bundle_id_idx + for h in sorted(executor.ray_worker_handles, key=lambda h: h.rank) + ] + assert actual == expected_bundle_ids + assert_executor(executor, tp_size=2, pp_size=1) + finally: + executor.shutdown() + + +@pytest.mark.parametrize( + "bundle_indices, expected_error, create_placement_group", + [ + ("1,1", "cannot have duplicate values,", 4), + ("0,1,2", "must have the same size", 4), + ], + indirect=["create_placement_group"], +) +def test_ray_v2_invalid_bundle_indices( + bundle_indices, expected_error, create_placement_group, monkeypatch +): + """Validate invalid bundle indices are rejected.""" + monkeypatch.setenv("VLLM_RAY_BUNDLE_INDICES", bundle_indices) + vllm_config = create_vllm_config( + tensor_parallel_size=2, placement_group=create_placement_group + ) + with pytest.raises(AssertionError, match=expected_error): + RayExecutorV2(vllm_config=vllm_config) + + +@pytest.mark.parametrize("tp_size, pp_size", [(2, 1), (2, 2)]) +def test_ray_v2_single_node_generation_with_pg(tp_size, pp_size): + """E2E LLM generation with a user-provided placement group.""" + ensure_ray_initialized() + bundles = [{"GPU": 1, "CPU": 1} for _ in range(tp_size * pp_size)] + pg = ray.util.placement_group(bundles, strategy="PACK") + ray.get(pg.ready()) + + try: + with patch.object(ray.util, "get_current_placement_group", return_value=pg): + llm = LLM( + model=MODEL, + tensor_parallel_size=tp_size, + pipeline_parallel_size=pp_size, + distributed_executor_backend="ray", + enforce_eager=True, + max_model_len=256, + gpu_memory_utilization=0.3, + ) + prompts = [ + "Hello, my name is", + "The capital of France is", + "The future of AI is", + ] + outputs = llm.generate(prompts) + + assert len(outputs) == len(prompts) + for output in outputs: + assert len(output.outputs) > 0 + assert len(output.outputs[0].text) > 0 + finally: + llm.llm_engine.model_executor.shutdown() + del llm + gc.collect() diff --git a/tests/distributed/test_ray_v2_executor_e2e.py b/tests/distributed/test_ray_v2_executor_e2e.py new file mode 100644 index 00000000000..fb583013269 --- /dev/null +++ b/tests/distributed/test_ray_v2_executor_e2e.py @@ -0,0 +1,209 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Orchestration-level integration tests for RayExecutorV2. +""" + +import gc +import os +import pathlib + +import pytest +import ray + +pytestmark = pytest.mark.usefixtures("enable_ray_v2_backend") + +MODEL = "facebook/opt-125m" + + +def _get_env_var(worker, name): + return os.environ.get(name) + + +def _ray_init(): + """Start Ray with the project root on workers' PYTHONPATH. + + Without this, workers cannot unpickle actor classes defined in the + ``tests`` package, causing FunctionActorManager to fall back to + TemporaryActor which drops async method signatures.""" + project_root = str(pathlib.Path(__file__).resolve().parents[2]) + ray.init( + ignore_reinit_error=True, + runtime_env={"env_vars": {"PYTHONPATH": project_root}}, + ) + + +@pytest.fixture +def ray_init(): + _ray_init() + + +class _AsyncLLMActor: + def start(self, pg, bundle_indices=None, ray_runtime_env=None): + os.environ["VLLM_USE_RAY_V2_EXECUTOR_BACKEND"] = "1" + # Needed so collective_rpc can pickle _get_env_var over the + # AsyncLLM -> EngineCore ZMQ boundary. + os.environ["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1" + if bundle_indices is not None: + os.environ["VLLM_RAY_BUNDLE_INDICES"] = bundle_indices + else: + os.environ.pop("VLLM_RAY_BUNDLE_INDICES", None) + + from vllm.engine.arg_utils import AsyncEngineArgs + from vllm.v1.engine.async_llm import AsyncLLM + from vllm.v1.executor.abstract import Executor + + engine_args = AsyncEngineArgs( + model=MODEL, + tensor_parallel_size=2, + distributed_executor_backend="ray", + enforce_eager=True, + max_model_len=256, + gpu_memory_utilization=0.8, + ) + vllm_config = engine_args.create_engine_config() + vllm_config.parallel_config.placement_group = pg + if ray_runtime_env is not None: + vllm_config.parallel_config.ray_runtime_env = ray_runtime_env + + executor_class = Executor.get_class(vllm_config) + self.engine = AsyncLLM( + vllm_config=vllm_config, + executor_class=executor_class, + log_stats=False, + log_requests=False, + ) + + async def generate(self, prompt): + from vllm.sampling_params import SamplingParams + + params = SamplingParams(max_tokens=16) + result = None + async for output in self.engine.generate( + prompt, params, request_id="test_request_id" + ): + result = output + assert result is not None + return result.outputs[0].text + + async def generate_and_get_worker_envs(self, prompt, env_names): + from vllm.sampling_params import SamplingParams + + params = SamplingParams(max_tokens=16) + result = None + async for output in self.engine.generate( + prompt, params, request_id="test_request_id" + ): + result = output + assert result is not None + text = result.outputs[0].text + + env_results = {} + for name in env_names: + vals = await self.engine.collective_rpc( + _get_env_var, timeout=10, args=(name,) + ) + env_results[name] = vals + return text, env_results + + def shutdown(self): + if engine := getattr(self, "engine", None): + engine.shutdown() + del self.engine + gc.collect() + + +AsyncLLMActor = ray.remote(num_cpus=0, max_concurrency=1)(_AsyncLLMActor) + + +def test_multi_replicas(ray_init): + pg1 = ray.util.placement_group([{"GPU": 1, "CPU": 1}] * 2, strategy="PACK") + pg2 = ray.util.placement_group([{"GPU": 1, "CPU": 1}] * 2, strategy="PACK") + ray.get([pg1.ready(), pg2.ready()]) + + actor1 = AsyncLLMActor.remote() + actor2 = AsyncLLMActor.remote() + + ray.get(actor1.start.remote(pg1)) + ray.get(actor2.start.remote(pg2)) + + out1, out2 = ray.get( + [ + actor1.generate.remote("Hello world"), + actor2.generate.remote("Hello world"), + ] + ) + assert len(out1) > 0 + assert len(out2) > 0 + + +def test_multi_replicas_with_bundle_indices(ray_init): + pg = ray.util.placement_group([{"GPU": 1, "CPU": 1}] * 4, strategy="PACK") + ray.get(pg.ready()) + + actor1 = AsyncLLMActor.remote() + actor2 = AsyncLLMActor.remote() + + ray.get(actor1.start.remote(pg, bundle_indices="2,1")) + ray.get(actor2.start.remote(pg, bundle_indices="0,3")) + + out1, out2 = ray.get( + [ + actor1.generate.remote("Hello world"), + actor2.generate.remote("Hello world"), + ] + ) + assert len(out1) > 0 + assert len(out2) > 0 + + +def test_env_var_and_runtime_env_propagation(): + """ + Verify env vars (NCCL_, HF_) and parallel_config.ray_runtime_env + propagate to RayWorkerProc actors. + """ + sentinel_vars = { + "NCCL_DEBUG": "INFO", + "HF_TOKEN": "test_sentinel_token", + } + for k, v in sentinel_vars.items(): + os.environ[k] = v + + try: + # Called directly (not via the ray_init fixture) because sentinel + # env vars must be in os.environ before ray.init() so that Ray + # worker processes inherit them. + _ray_init() + + pg = ray.util.placement_group([{"GPU": 1, "CPU": 1}] * 2, strategy="PACK") + ray.get(pg.ready()) + + # Include the project root so that RayWorkerProc actors can + # unpickle _get_env_var. + project_root = str(pathlib.Path(__file__).resolve().parents[2]) + ray_runtime_env = { + "env_vars": { + "RAY_RUNTIME_ENV_TEST": "ray_runtime_env", + "PYTHONPATH": project_root, + }, + } + + actor = AsyncLLMActor.remote() + ray.get(actor.start.remote(pg, ray_runtime_env=ray_runtime_env)) + + all_env_names = list(sentinel_vars) + ["RAY_RUNTIME_ENV_TEST"] + text, env_results = ray.get( + actor.generate_and_get_worker_envs.remote("Hello world", all_env_names) + ) + assert len(text) > 0 + + for name, expected in sentinel_vars.items(): + for val in env_results[name]: + assert val == expected + + for val in env_results["RAY_RUNTIME_ENV_TEST"]: + assert val == "ray_runtime_env" + + finally: + for k in sentinel_vars: + os.environ.pop(k, None) diff --git a/tests/entrypoints/pooling/scoring/test_late_interaction_online.py b/tests/entrypoints/pooling/scoring/test_late_interaction_online.py index 9eedec6d2b9..7e4501fe850 100644 --- a/tests/entrypoints/pooling/scoring/test_late_interaction_online.py +++ b/tests/entrypoints/pooling/scoring/test_late_interaction_online.py @@ -26,13 +26,18 @@ TEXTS_2 = [ ] -@pytest.fixture(scope="module") -def server(): +@pytest.fixture(scope="module", params=[True, False]) +def server(request): args = [ "--max-model-len", str(MAX_MODEL_LEN), ] + # Test run pooling score MaxSim on worker side (GPU) + # aka flash-late-interaction + if not request.param: + args += ["--no-enable-flash-late-interaction"] + with RemoteOpenAIServer(MODEL_NAME, args) as remote_server: yield remote_server diff --git a/tests/entrypoints/serve/disagg/test_generate_stream.py b/tests/entrypoints/serve/disagg/test_generate_stream.py new file mode 100644 index 00000000000..a9ca026306f --- /dev/null +++ b/tests/entrypoints/serve/disagg/test_generate_stream.py @@ -0,0 +1,474 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import json +from dataclasses import dataclass, field +from typing import Any +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from vllm.config.multimodal import MultiModalConfig +from vllm.entrypoints.openai.engine.protocol import StreamOptions +from vllm.entrypoints.openai.models.protocol import BaseModelPath +from vllm.entrypoints.openai.models.serving import OpenAIServingModels +from vllm.entrypoints.serve.disagg.protocol import GenerateRequest +from vllm.entrypoints.serve.disagg.serving import ServingTokens +from vllm.entrypoints.serve.render.serving import OpenAIServingRender +from vllm.logprobs import Logprob +from vllm.outputs import CompletionOutput, RequestOutput +from vllm.renderers import renderer_from_config +from vllm.sampling_params import SamplingParams +from vllm.v1.engine.async_llm import AsyncLLM + +MODEL_NAME = "openai-community/gpt2" +BASE_MODEL_PATHS = [ + BaseModelPath(name=MODEL_NAME, model_path=MODEL_NAME), +] + + +@dataclass +class MockHFConfig: + model_type: str = "any" + + +@dataclass +class MockModelConfig: + task = "generate" + runner_type = "generate" + model = MODEL_NAME + tokenizer = MODEL_NAME + trust_remote_code = False + tokenizer_mode = "auto" + max_model_len = 100 + tokenizer_revision = None + multimodal_config = MultiModalConfig() + hf_config = MockHFConfig() + hf_text_config = MockHFConfig() + logits_processors: list[str] | None = None + diff_sampling_param: dict | None = None + allowed_local_media_path: str = "" + allowed_media_domains: list[str] | None = None + encoder_config = None + generation_config: str = "auto" + media_io_kwargs: dict[str, dict[str, Any]] = field(default_factory=dict) + skip_tokenizer_init = False + is_encoder_decoder: bool = False + is_multimodal_model: bool = False + renderer_num_workers: int = 1 + + def get_diff_sampling_param(self): + return self.diff_sampling_param or {} + + +@dataclass +class MockParallelConfig: + _api_process_rank: int = 0 + + +@dataclass +class MockVllmConfig: + model_config: MockModelConfig + parallel_config: MockParallelConfig + + +def _build_renderer(model_config: MockModelConfig): + return renderer_from_config( + MockVllmConfig(model_config, parallel_config=MockParallelConfig()), + ) + + +def _build_serving_tokens(engine: AsyncLLM, **kwargs) -> ServingTokens: + models = OpenAIServingModels( + engine_client=engine, + base_model_paths=BASE_MODEL_PATHS, + ) + serving_render = OpenAIServingRender( + model_config=engine.model_config, + renderer=engine.renderer, + io_processor=engine.io_processor, + model_registry=models.registry, + request_logger=None, + chat_template=None, + chat_template_content_format="auto", + ) + serving = ServingTokens( + engine, + models, + openai_serving_render=serving_render, + request_logger=None, + **kwargs, + ) + + async def _fake_preprocess(*args, **kwargs): + return [{"prompt_token_ids": [1, 2, 3]}] + + serving.openai_serving_render.preprocess_completion = AsyncMock( + side_effect=_fake_preprocess + ) + return serving + + +def _make_request_output( + request_id: str, + token_ids: list[int], + finish_reason: str | None = None, + finished: bool = False, + prompt_token_ids: list[int] | None = None, + logprobs: list[dict[int, Any] | None] | None = None, + num_cached_tokens: int | None = None, + index: int = 0, +) -> RequestOutput: + return RequestOutput( + request_id=request_id, + prompt=None, + prompt_token_ids=prompt_token_ids or [1, 2, 3], + prompt_logprobs=None, + outputs=[ + CompletionOutput( + index=index, + text="", + token_ids=token_ids, + cumulative_logprob=None, + logprobs=logprobs, + finish_reason=finish_reason, + ) + ], + finished=finished, + metrics=None, + lora_request=None, + encoder_prompt=None, + encoder_prompt_token_ids=None, + num_cached_tokens=num_cached_tokens, + ) + + +def _mock_engine() -> MagicMock: + engine = MagicMock(spec=AsyncLLM) + engine.errored = False + engine.model_config = MockModelConfig() + engine.input_processor = MagicMock() + engine.io_processor = MagicMock() + engine.renderer = _build_renderer(engine.model_config) + return engine + + +def _parse_sse_chunks(chunks: list[str]) -> list[Any]: + """Parse SSE chunks into dicts (JSON) or raw strings ([DONE]).""" + parsed: list[Any] = [] + for chunk in chunks: + assert chunk.startswith("data: ") and chunk.endswith("\n\n") + payload = chunk[len("data: ") : -len("\n\n")] + if payload == "[DONE]": + parsed.append("[DONE]") + else: + parsed.append(json.loads(payload)) + return parsed + + +@pytest.mark.asyncio +async def test_stream_basic(): + """Streaming returns SSE chunks with correct token_ids and ends with [DONE].""" + engine = _mock_engine() + + async def mock_generate(*args, **kwargs): + yield _make_request_output("req-1", token_ids=[10]) + yield _make_request_output("req-1", token_ids=[20, 30]) + yield _make_request_output( + "req-1", token_ids=[40], finish_reason="stop", finished=True + ) + + engine.generate = MagicMock(side_effect=mock_generate) + serving = _build_serving_tokens(engine) + + request = GenerateRequest( + token_ids=[1, 2, 3], + sampling_params=SamplingParams(max_tokens=10), + model=MODEL_NAME, + stream=True, + ) + + response = await serving.serve_tokens(request) + chunks = [] + async for chunk in response: + chunks.append(chunk) + + parsed = _parse_sse_chunks(chunks) + + # 3 data chunks + [DONE] + assert parsed[-1] == "[DONE]" + data_chunks = [c for c in parsed if c != "[DONE]"] + assert len(data_chunks) == 3 + + assert data_chunks[0]["choices"][0]["token_ids"] == [10] + assert data_chunks[1]["choices"][0]["token_ids"] == [20, 30] + assert data_chunks[2]["choices"][0]["token_ids"] == [40] + assert data_chunks[2]["choices"][0]["finish_reason"] == "stop" + + +@pytest.mark.asyncio +async def test_stream_error_mid_generation(): + """finish_reason='error' mid-stream yields error chunk then [DONE].""" + engine = _mock_engine() + + async def mock_generate(*args, **kwargs): + yield _make_request_output("req-1", token_ids=[10]) + yield _make_request_output( + "req-1", token_ids=[20], finish_reason="error", finished=True + ) + + engine.generate = MagicMock(side_effect=mock_generate) + serving = _build_serving_tokens(engine) + + request = GenerateRequest( + token_ids=[1, 2, 3], + sampling_params=SamplingParams(max_tokens=10), + model=MODEL_NAME, + stream=True, + ) + + response = await serving.serve_tokens(request) + chunks = [] + async for chunk in response: + chunks.append(chunk) + + assert len(chunks) >= 2 + assert any("Internal server error" in chunk for chunk in chunks), ( + f"Expected error message in chunks: {chunks}" + ) + assert chunks[-1] == "data: [DONE]\n\n" + + +@pytest.mark.asyncio +async def test_stream_error_with_empty_delta(): + """finish_reason='error' with empty delta_token_ids still raises.""" + engine = _mock_engine() + + async def mock_generate(*args, **kwargs): + yield _make_request_output("req-1", token_ids=[10]) + yield _make_request_output( + "req-1", token_ids=[], finish_reason="error", finished=True + ) + + engine.generate = MagicMock(side_effect=mock_generate) + serving = _build_serving_tokens(engine) + + request = GenerateRequest( + token_ids=[1, 2, 3], + sampling_params=SamplingParams(max_tokens=10), + model=MODEL_NAME, + stream=True, + ) + + response = await serving.serve_tokens(request) + chunks = [] + async for chunk in response: + chunks.append(chunk) + + assert any("Internal server error" in chunk for chunk in chunks), ( + f"Expected error message in chunks: {chunks}" + ) + assert chunks[-1] == "data: [DONE]\n\n" + + +@pytest.mark.asyncio +async def test_stream_skips_empty_token_output(): + """Outputs with empty token_ids are skipped (no chunk emitted).""" + engine = _mock_engine() + + async def mock_generate(*args, **kwargs): + yield _make_request_output("req-1", token_ids=[10]) + yield _make_request_output("req-1", token_ids=[]) + yield _make_request_output( + "req-1", token_ids=[20], finish_reason="stop", finished=True + ) + + engine.generate = MagicMock(side_effect=mock_generate) + serving = _build_serving_tokens(engine) + + request = GenerateRequest( + token_ids=[1, 2, 3], + sampling_params=SamplingParams(max_tokens=10), + model=MODEL_NAME, + stream=True, + ) + + response = await serving.serve_tokens(request) + chunks = [] + async for chunk in response: + chunks.append(chunk) + + parsed = _parse_sse_chunks(chunks) + assert parsed[-1] == "[DONE]" + data_chunks = [c for c in parsed if c != "[DONE]"] + + # Only 2 data chunks — the empty one is skipped + assert len(data_chunks) == 2 + assert data_chunks[0]["choices"][0]["token_ids"] == [10] + assert data_chunks[1]["choices"][0]["token_ids"] == [20] + + +@pytest.mark.asyncio +async def test_stream_include_usage(): + """stream_options.include_usage emits a final usage-only chunk.""" + engine = _mock_engine() + + async def mock_generate(*args, **kwargs): + yield _make_request_output("req-1", token_ids=[10]) + yield _make_request_output( + "req-1", token_ids=[20], finish_reason="stop", finished=True + ) + + engine.generate = MagicMock(side_effect=mock_generate) + serving = _build_serving_tokens(engine) + + request = GenerateRequest( + token_ids=[1, 2, 3], + sampling_params=SamplingParams(max_tokens=10), + model=MODEL_NAME, + stream=True, + stream_options=StreamOptions(include_usage=True), + ) + + response = await serving.serve_tokens(request) + chunks = [] + async for chunk in response: + chunks.append(chunk) + + parsed = _parse_sse_chunks(chunks) + assert parsed[-1] == "[DONE]" + + # The chunk before [DONE] should be the usage-only chunk + usage_chunk = parsed[-2] + assert usage_chunk["choices"] == [] + assert usage_chunk["usage"]["prompt_tokens"] == 3 + assert usage_chunk["usage"]["completion_tokens"] == 2 + assert usage_chunk["usage"]["total_tokens"] == 5 + + +@pytest.mark.asyncio +async def test_stream_continuous_usage(): + """continuous_usage_stats adds usage to every data chunk.""" + engine = _mock_engine() + + async def mock_generate(*args, **kwargs): + yield _make_request_output("req-1", token_ids=[10]) + yield _make_request_output( + "req-1", token_ids=[20], finish_reason="stop", finished=True + ) + + engine.generate = MagicMock(side_effect=mock_generate) + serving = _build_serving_tokens(engine) + + request = GenerateRequest( + token_ids=[1, 2, 3], + sampling_params=SamplingParams(max_tokens=10), + model=MODEL_NAME, + stream=True, + stream_options=StreamOptions( + include_usage=True, + continuous_usage_stats=True, + ), + ) + + response = await serving.serve_tokens(request) + chunks = [] + async for chunk in response: + chunks.append(chunk) + + parsed = _parse_sse_chunks(chunks) + data_chunks = [c for c in parsed if isinstance(c, dict) and c.get("choices")] + + # Every data chunk should have usage + for i, dc in enumerate(data_chunks): + assert dc["usage"] is not None, f"chunk {i} missing usage" + assert dc["usage"]["prompt_tokens"] == 3 + + # First chunk: 1 completion token + assert data_chunks[0]["usage"]["completion_tokens"] == 1 + assert data_chunks[0]["usage"]["total_tokens"] == 4 + + # Second chunk: 2 completion tokens (cumulative) + assert data_chunks[1]["usage"]["completion_tokens"] == 2 + assert data_chunks[1]["usage"]["total_tokens"] == 5 + + +@pytest.mark.asyncio +async def test_stream_with_logprobs(): + """Streaming with logprobs includes logprob data in each chunk.""" + engine = _mock_engine() + + async def mock_generate(*args, **kwargs): + yield _make_request_output( + "req-1", + token_ids=[10], + logprobs=[{10: Logprob(logprob=-0.5)}], + ) + yield _make_request_output( + "req-1", + token_ids=[20], + logprobs=[{20: Logprob(logprob=-1.0)}], + finish_reason="stop", + finished=True, + ) + + engine.generate = MagicMock(side_effect=mock_generate) + serving = _build_serving_tokens(engine) + + request = GenerateRequest( + token_ids=[1, 2, 3], + sampling_params=SamplingParams(max_tokens=10, logprobs=1), + model=MODEL_NAME, + stream=True, + ) + + response = await serving.serve_tokens(request) + chunks = [] + async for chunk in response: + chunks.append(chunk) + + parsed = _parse_sse_chunks(chunks) + data_chunks = [c for c in parsed if isinstance(c, dict) and c.get("choices")] + + for dc in data_chunks: + lp = dc["choices"][0]["logprobs"] + assert lp is not None + assert len(lp["content"]) == 1 + assert lp["content"][0]["token"].startswith("token_id:") + + +@pytest.mark.asyncio +async def test_stream_prompt_tokens_details(): + """enable_prompt_tokens_details includes cached_tokens in final usage.""" + engine = _mock_engine() + + async def mock_generate(*args, **kwargs): + yield _make_request_output( + "req-1", + token_ids=[10], + finish_reason="stop", + finished=True, + num_cached_tokens=2, + ) + + engine.generate = MagicMock(side_effect=mock_generate) + serving = _build_serving_tokens(engine, enable_prompt_tokens_details=True) + + request = GenerateRequest( + token_ids=[1, 2, 3], + sampling_params=SamplingParams(max_tokens=10), + model=MODEL_NAME, + stream=True, + stream_options=StreamOptions(include_usage=True), + ) + + response = await serving.serve_tokens(request) + chunks = [] + async for chunk in response: + chunks.append(chunk) + + parsed = _parse_sse_chunks(chunks) + # Usage-only chunk (before [DONE]) + usage_chunk = parsed[-2] + assert usage_chunk["choices"] == [] + assert usage_chunk["usage"]["prompt_tokens_details"]["cached_tokens"] == 2 diff --git a/tests/entrypoints/serve/disagg/test_serving_tokens.py b/tests/entrypoints/serve/disagg/test_serving_tokens.py index b62cb01bb45..4ae7e0494a0 100644 --- a/tests/entrypoints/serve/disagg/test_serving_tokens.py +++ b/tests/entrypoints/serve/disagg/test_serving_tokens.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import json import os import httpx @@ -113,6 +114,54 @@ async def test_generate_endpoint(client): assert "choices" in data +@pytest.mark.asyncio +async def test_generate_stream(client): + payload = { + "model": MODEL_NAME, + "token_ids": [1, 2, 3], + "sampling_params": {"max_tokens": 5}, + "stream": True, + } + async with client.stream("POST", GEN_ENDPOINT, json=payload) as resp: + resp.raise_for_status() + chunks = [] + async for line in resp.aiter_lines(): + if not line.startswith("data: "): + continue + payload_str = line[len("data: ") :] + if payload_str == "[DONE]": + break + chunks.append(json.loads(payload_str)) + + assert len(chunks) > 0 + # Every chunk has choices with token_ids + all_token_ids = [] + for chunk in chunks: + assert "choices" in chunk + assert len(chunk["choices"]) == 1 + choice = chunk["choices"][0] + assert "token_ids" in choice + assert len(choice["token_ids"]) > 0 + all_token_ids.extend(choice["token_ids"]) + + # Last chunk should have a finish_reason + assert chunks[-1]["choices"][0]["finish_reason"] is not None + + # Streaming should produce the same tokens as non-streaming + non_stream_resp = await client.post( + GEN_ENDPOINT, + json={ + "model": MODEL_NAME, + "token_ids": [1, 2, 3], + "sampling_params": {"max_tokens": 5, "temperature": 0.0}, + "stream": False, + }, + ) + non_stream_data = non_stream_resp.json() + # Just verify we got the right number of tokens + assert len(all_token_ids) == len(non_stream_data["choices"][0]["token_ids"]) + + @pytest.mark.asyncio @pytest.mark.parametrize("logprobs_value", [0, 1, 5]) async def test_generate_logprobs(client, logprobs_value): diff --git a/tests/entrypoints/serve/tokenize/test_tokenize_then_chat_vlm.py b/tests/entrypoints/serve/tokenize/test_tokenize_then_chat_vlm.py new file mode 100644 index 00000000000..50083e34553 --- /dev/null +++ b/tests/entrypoints/serve/tokenize/test_tokenize_then_chat_vlm.py @@ -0,0 +1,85 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Regression test: calling ``/tokenize`` with multimodal data followed by +``/v1/chat/completions`` with the same data must not cause an error. + +Ensures that the ``/tokenize`` endpoint does not pollute internal caches +(e.g. multimodal feature caches) and that a subsequent +``/v1/chat/completions`` request with the same multimodal payload +completes successfully. +""" + +import json + +import openai +import pytest +import pytest_asyncio +import requests + +from tests.utils import RemoteOpenAIServer + +MODEL_NAME = "Qwen/Qwen2.5-VL-3B-Instruct" + + +@pytest.fixture(scope="module") +def server(): + args = [ + "--dtype", + "bfloat16", + "--max-model-len", + "4096", + "--max-num-seqs", + "5", + "--enforce-eager", + "--limit-mm-per-prompt", + json.dumps({"image": 1}), + ] + with RemoteOpenAIServer(MODEL_NAME, args) as remote_server: + yield remote_server + + +@pytest_asyncio.fixture +async def client(server): + async with server.get_async_client() as async_client: + yield async_client + + +@pytest.mark.asyncio +async def test_tokenize_then_chat_completion_with_image( + client: openai.AsyncOpenAI, + server: RemoteOpenAIServer, + local_asset_server, +): + """Tokenize a multimodal message, then send the same message to chat + completions. The chat completion must succeed (not 500).""" + + image_url = local_asset_server.url_for("stop_sign.jpg") + messages = [ + { + "role": "user", + "content": [ + {"type": "image_url", "image_url": {"url": image_url}}, + {"type": "text", "text": "Describe this image briefly."}, + ], + } + ] + + tok_resp = requests.post( + server.url_for("tokenize"), + json={"model": MODEL_NAME, "messages": messages}, + ) + tok_resp.raise_for_status() + tok_data = tok_resp.json() + assert tok_data["count"] > 0, "Tokenization must return tokens" + + chat_completion = await client.chat.completions.create( + model=MODEL_NAME, + messages=messages, + max_tokens=10, + temperature=0.0, + ) + + assert chat_completion.choices[0].message.content, ( + "Chat completion must produce non-empty content after tokenize" + ) diff --git a/tests/evals/gpt_oss/configs/gpt-oss-20b-rocm-quark-mxfp4-bf16-aiter.yaml b/tests/evals/gpt_oss/configs/gpt-oss-20b-rocm-quark-mxfp4-bf16-aiter.yaml new file mode 100644 index 00000000000..850a6d28be0 --- /dev/null +++ b/tests/evals/gpt_oss/configs/gpt-oss-20b-rocm-quark-mxfp4-bf16-aiter.yaml @@ -0,0 +1,8 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +model_name: amd/gpt-oss-20b-w-mxfp4-a-bf16 +metric_threshold: 0.568 +reasoning_effort: low +server_args: "--attention-backend ROCM_AITER_UNIFIED_ATTN --moe-backend aiter" +env: + VLLM_ROCM_USE_AITER: "1" diff --git a/tests/evals/gpt_oss/configs/gpt-oss-20b-rocm-quark-mxfp4-bf16-triton.yaml b/tests/evals/gpt_oss/configs/gpt-oss-20b-rocm-quark-mxfp4-bf16-triton.yaml new file mode 100644 index 00000000000..903f30e59e7 --- /dev/null +++ b/tests/evals/gpt_oss/configs/gpt-oss-20b-rocm-quark-mxfp4-bf16-triton.yaml @@ -0,0 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +model_name: amd/gpt-oss-20b-w-mxfp4-a-bf16 +metric_threshold: 0.568 +reasoning_effort: low +server_args: "--attention-backend ROCM_AITER_UNIFIED_ATTN --moe-backend triton" \ No newline at end of file diff --git a/tests/evals/gpt_oss/configs/gpt-oss-20b-rocm-quark-mxfp4-fp8-triton.yaml b/tests/evals/gpt_oss/configs/gpt-oss-20b-rocm-quark-mxfp4-fp8-triton.yaml new file mode 100644 index 00000000000..f7dd14784a1 --- /dev/null +++ b/tests/evals/gpt_oss/configs/gpt-oss-20b-rocm-quark-mxfp4-fp8-triton.yaml @@ -0,0 +1,8 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +model_name: amd/gpt-oss-20b-MoE-Quant-W-MXFP4-A-FP8-KV-FP8 +metric_threshold: 0.568 +reasoning_effort: low +server_args: "--attention-backend ROCM_AITER_UNIFIED_ATTN" +env: + VLLM_ROCM_USE_AITER: "1" \ No newline at end of file diff --git a/tests/evals/gpt_oss/configs/models-gfx950.txt b/tests/evals/gpt_oss/configs/models-gfx950.txt index 2b6ff4f4a8d..d25f4e3a5e2 100644 --- a/tests/evals/gpt_oss/configs/models-gfx950.txt +++ b/tests/evals/gpt_oss/configs/models-gfx950.txt @@ -1,3 +1,6 @@ # GFX950 model configurations for GPQA evaluation # Tests different environment variable combinations -gpt-oss-20b-rocm-baseline.yaml \ No newline at end of file +gpt-oss-20b-rocm-baseline.yaml +gpt-oss-20b-rocm-quark-mxfp4-bf16-aiter.yaml +gpt-oss-20b-rocm-quark-mxfp4-bf16-triton.yaml +gpt-oss-20b-rocm-quark-mxfp4-fp8-triton.yaml diff --git a/tests/evals/gsm8k/configs/Qwen3.5-35B-A3B-MXFP4-TP2.yaml b/tests/evals/gsm8k/configs/Qwen3.5-35B-A3B-MXFP4-TP2.yaml new file mode 100644 index 00000000000..67eda114155 --- /dev/null +++ b/tests/evals/gsm8k/configs/Qwen3.5-35B-A3B-MXFP4-TP2.yaml @@ -0,0 +1,8 @@ +model_name: "amd/Qwen3.5-35B-A3B-MXFP4" +accuracy_threshold: 0.82 +tolerance: 0.03 +num_questions: 1319 +num_fewshot: 5 +server_args: >- + --max-model-len 4096 + --tensor-parallel-size 2 diff --git a/tests/evals/gsm8k/configs/models-qwen35-mi355.txt b/tests/evals/gsm8k/configs/models-qwen35-mi355.txt index 4e7af71c7f4..db8e88e2735 100644 --- a/tests/evals/gsm8k/configs/models-qwen35-mi355.txt +++ b/tests/evals/gsm8k/configs/models-qwen35-mi355.txt @@ -1 +1,2 @@ Qwen3.5-35B-A3B-DEP2.yaml +Qwen3.5-35B-A3B-MXFP4-TP2.yaml diff --git a/tests/kernels/attention/test_merge_attn_states.py b/tests/kernels/attention/test_merge_attn_states.py index c1b71d93e4d..40af84887a9 100644 --- a/tests/kernels/attention/test_merge_attn_states.py +++ b/tests/kernels/attention/test_merge_attn_states.py @@ -4,7 +4,12 @@ import pytest import torch -from vllm._custom_ops import merge_attn_states as merge_attn_states_cuda +from vllm._custom_ops import ( + merge_attn_states as merge_attn_states_cuda, +) +from vllm._custom_ops import ( + scaled_fp8_quant, +) from vllm.platforms import current_platform from vllm.v1.attention.ops.triton_merge_attn_states import ( merge_attn_states as merge_attn_states_triton, @@ -21,6 +26,7 @@ def merge_attn_states_torch( suffix_lse: torch.Tensor, # [NUM_HEADS, NUM_TOKENS] output_lse: torch.Tensor | None = None, # [NUM_HEADS, NUM_TOKENS] prefill_tokens_with_context: int | None = None, + output_scale: torch.Tensor | None = None, # scalar, per-tensor FP8 scale ): # Apply prefill_tokens_with_context mask if needed if prefill_tokens_with_context is None: @@ -49,9 +55,13 @@ def merge_attn_states_torch( s_scale = s_lse_exp / out_se # [NUM_HEADS, NUM_TOKENS] p_scale = torch.transpose(p_scale, 0, 1).unsqueeze(2) # [NUM_TOKENS, NUM_HEADS, 1] s_scale = torch.transpose(s_scale, 0, 1).unsqueeze(2) # [NUM_TOKENS, NUM_HEADS, 1] - output.copy_( - prefix_output * p_scale * mask + suffix_output * (s_scale * mask + (1 - mask)) + output = prefix_output * p_scale * mask + suffix_output * ( + s_scale * mask + (1 - mask) ) + if output_scale is not None: + shape = output.shape + output, _ = scaled_fp8_quant(output.float().view(-1, shape[-1]), output_scale) + output = output.view(shape) return output, output_lse @@ -102,18 +112,20 @@ def generate_markdown_table(): ) +@pytest.mark.parametrize("use_fp8", [False, True]) @pytest.mark.parametrize("prefill_tokens_with_context", [None, 128]) @pytest.mark.parametrize("num_tokens", NUM_BATCH_TOKENS) @pytest.mark.parametrize("num_query_heads", NUM_QUERY_HEADS) @pytest.mark.parametrize("head_size", HEAD_SIZES) -@pytest.mark.parametrize("output_dtype", DTYPES) +@pytest.mark.parametrize("input_dtype", DTYPES) @torch.inference_mode() def test_merge_attn_states( prefill_tokens_with_context: int | None, num_tokens: int, num_query_heads: int, head_size: int, - output_dtype: torch.dtype, + input_dtype: torch.dtype, + use_fp8: bool, ): if not current_platform.is_cuda(): pytest.skip( @@ -125,9 +137,18 @@ def test_merge_attn_states( NUM_HEADS = num_query_heads HEAD_SIZE = head_size + # When use_fp8 is set, inputs stay as input_dtype (bf16/fp16/fp32) + # and output becomes FP8. + output_dtype = input_dtype + output_scale = None + if use_fp8: + output_dtype = current_platform.fp8_dtype() + output_scale = torch.tensor([0.05], dtype=torch.float32, device="cuda") + print( f"\nNUM_TOKENS:{NUM_TOKENS}, NUM_HEADS:{NUM_HEADS}, " - f"HEAD_SIZE:{HEAD_SIZE}, DTYPE: {output_dtype}, " + f"HEAD_SIZE:{HEAD_SIZE}, input_dtype: {input_dtype}, " + f"output_dtype: {output_dtype}, use_fp8: {use_fp8}, " f"prefill_tokens_with_context: {prefill_tokens_with_context}, " f"Device: {current_platform.get_device_name()}" ) @@ -156,10 +177,10 @@ def test_merge_attn_states( (NUM_HEADS, NUM_TOKENS), dtype=torch.float32, device="cuda" ) prefix_output = torch.randn( - (NUM_TOKENS, NUM_HEADS, HEAD_SIZE), dtype=output_dtype, device="cuda" + (NUM_TOKENS, NUM_HEADS, HEAD_SIZE), dtype=input_dtype, device="cuda" ) suffix_output = torch.randn( - (NUM_TOKENS, NUM_HEADS, HEAD_SIZE), dtype=output_dtype, device="cuda" + (NUM_TOKENS, NUM_HEADS, HEAD_SIZE), dtype=input_dtype, device="cuda" ) warmup_times = 2 @@ -183,6 +204,7 @@ def test_merge_attn_states( suffix_lse_torch, output_lse_torch, prefill_tokens_with_context, + output_scale, ) torch.accelerator.synchronize() @@ -196,6 +218,7 @@ def test_merge_attn_states( suffix_lse_torch, output_lse_torch, prefill_tokens_with_context, + output_scale, ) end.record() torch.accelerator.synchronize() @@ -220,6 +243,7 @@ def test_merge_attn_states( suffix_lse, output_lse_ref_triton, prefill_tokens_with_context, + output_scale, ) torch.accelerator.synchronize() @@ -233,6 +257,7 @@ def test_merge_attn_states( suffix_lse, output_lse_ref_triton, prefill_tokens_with_context, + output_scale, ) end.record() torch.accelerator.synchronize() @@ -254,6 +279,7 @@ def test_merge_attn_states( suffix_lse, output_lse_cuda, prefill_tokens_with_context, + output_scale, ) torch.accelerator.synchronize() @@ -267,6 +293,7 @@ def test_merge_attn_states( suffix_lse, output_lse_cuda, prefill_tokens_with_context, + output_scale, ) end.record() torch.accelerator.synchronize() @@ -288,7 +315,19 @@ def test_merge_attn_states( # Liger Kernel: Efficient Triton Kernels for LLM Training # https://arxiv.org/pdf/2410.10989, 3.3 Correctness # use rtol = 1e-2 for bfloat16. - rtol = 1e-2 if output_dtype == torch.bfloat16 else 1e-3 + if use_fp8: + # Compare in dequantized space (multiply back by scale) so that + # absolute differences reflect real precision, not amplified FP8 + # quantization steps. + atol, rtol = 1e-1, 1e-1 + assert output_scale is not None + scale = output_scale.item() + elif output_dtype == torch.bfloat16: + atol, rtol = 1e-3, 1e-2 + scale = 1.0 + else: + atol, rtol = 1e-3, 1e-3 + scale = 1.0 def diff(a: torch.Tensor, b: torch.Tensor): max_diff = torch.max(torch.abs(a.float() - b.float())) @@ -300,16 +339,26 @@ def test_merge_attn_states( output_ref = output_ref_triton output_lse_ref = output_lse_ref_triton torch.testing.assert_close( - output_cuda.float(), output_ref.float(), atol=1e-3, rtol=rtol + output_cuda.float() * scale, + output_ref.float() * scale, + atol=atol, + rtol=rtol, ) - print("Output all match, max abs diff:") - print(f"(Triton vs Torch) : {diff(output_torch, output_ref)}") - print(f" (CUDA vs Torch) : {diff(output_torch, output_cuda)}") - print(f" (CUDA vs Triton): {diff(output_ref, output_cuda)}") + print( + "Output all match, max abs diff (dequantized):" + if use_fp8 + else "Output all match, max abs diff:" + ) + _diff = diff(output_ref.float() * scale, output_torch.float() * scale) + print(f"(Triton vs Torch) : {_diff}") + _diff = diff(output_torch.float() * scale, output_cuda.float() * scale) + print(f" (CUDA vs Torch) : {_diff}") + _diff = diff(output_ref.float() * scale, output_cuda.float() * scale) + print(f" (CUDA vs Triton): {_diff}") print("-" * 100) torch.testing.assert_close( - output_lse_cuda.float(), output_lse_ref.float(), atol=1e-3, rtol=rtol + output_lse_cuda.float(), output_lse_ref.float(), atol=atol, rtol=rtol ) print("Output LSE all match, max abs diff:") print(f"(Triton vs Torch) : {diff(output_lse_torch, output_lse_ref)}") diff --git a/tests/kernels/attention/test_use_trtllm_attention.py b/tests/kernels/attention/test_use_trtllm_attention.py index fba18fe46e3..12ab146a983 100644 --- a/tests/kernels/attention/test_use_trtllm_attention.py +++ b/tests/kernels/attention/test_use_trtllm_attention.py @@ -62,7 +62,7 @@ def test_supports_batch_invariant_disables(): @patch("vllm.envs.VLLM_BATCH_INVARIANT", False) @patch( - "vllm.utils.flashinfer.current_platform.is_device_capability_family", + "vllm.utils.flashinfer.current_platform.is_device_capability", return_value=True, ) @patch("vllm.utils.flashinfer.has_nvidia_artifactory", return_value=True) @@ -72,7 +72,7 @@ def test_supports_sm100_with_artifactory(_art, _cap): @patch("vllm.envs.VLLM_BATCH_INVARIANT", False) @patch( - "vllm.utils.flashinfer.current_platform.is_device_capability_family", + "vllm.utils.flashinfer.current_platform.is_device_capability", return_value=False, ) def test_supports_non_sm100_platform(_cap): @@ -81,7 +81,7 @@ def test_supports_non_sm100_platform(_cap): @patch("vllm.envs.VLLM_BATCH_INVARIANT", False) @patch( - "vllm.utils.flashinfer.current_platform.is_device_capability_family", + "vllm.utils.flashinfer.current_platform.is_device_capability", return_value=True, ) @patch("vllm.utils.flashinfer.has_nvidia_artifactory", return_value=False) diff --git a/tests/kernels/core/test_fused_silu_mul_block_quant.py b/tests/kernels/core/test_fused_silu_mul_block_quant.py new file mode 100644 index 00000000000..1878390ac2f --- /dev/null +++ b/tests/kernels/core/test_fused_silu_mul_block_quant.py @@ -0,0 +1,189 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch +import torch.nn.functional as F + +import vllm._custom_ops as ops +from tests.kernels.utils import opcheck +from vllm.model_executor.layers.quantization.utils.fp8_utils import ( + per_token_group_quant_fp8, +) +from vllm.model_executor.layers.quantization.utils.int8_utils import ( + per_token_group_quant_int8, +) +from vllm.platforms import current_platform + +DTYPES = [torch.float16, torch.bfloat16] +QUANT_DTYPES = [torch.float8_e4m3fn, torch.int8] +VEC_HIDDEN_SIZES = [1024, 1025, 1027, 1029] +NUM_TOKENS_HIDDEN_SIZES = [ + *[(1, i) for i in [64, *VEC_HIDDEN_SIZES, 2048, 5120]], + *[(16, i) for i in [64, *VEC_HIDDEN_SIZES, 5120]], + *[(128, i) for i in [64, *VEC_HIDDEN_SIZES]], + *[(512, i) for i in [64, 5120]], +] +SCALE_UBS = [False] +GROUP_SIZES = [64, 128] +IS_SCALE_TRANSPOSED = [False, True] +SEEDS = [0] +CUDA_DEVICES = [ + f"cuda:{i}" for i in range(1 if torch.accelerator.device_count() == 1 else 2) +] + + +def ref_silu_and_mul_per_block_quant( + x: torch.Tensor, + quant_dtype: torch.dtype, + group_size: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Reference implementation: unfused SiLU+Mul then group quantization.""" + hidden = x.shape[-1] // 2 + gate, up = x.split(hidden, dim=-1) + silu_out = F.silu(gate) * up + + if quant_dtype == current_platform.fp8_dtype(): + return per_token_group_quant_fp8( + silu_out, group_size=group_size, use_ue8m0=False + ) + elif quant_dtype == torch.int8: + return per_token_group_quant_int8(silu_out, group_size=group_size) + else: + raise ValueError(f"Unsupported quant_dtype: {quant_dtype}") + + +@pytest.mark.parametrize("num_tokens, hidden_size", NUM_TOKENS_HIDDEN_SIZES) +@pytest.mark.parametrize("has_scale_ub", SCALE_UBS) +@pytest.mark.parametrize("dtype", DTYPES) +@pytest.mark.parametrize("quant_dtype", QUANT_DTYPES) +@pytest.mark.parametrize("group_size", GROUP_SIZES) +@pytest.mark.parametrize("is_scale_transposed", IS_SCALE_TRANSPOSED) +@pytest.mark.parametrize("seed", SEEDS) +@pytest.mark.parametrize("device", CUDA_DEVICES) +@torch.inference_mode() +def test_silu_and_mul_per_block_quant( + default_vllm_config, + num_tokens: int, + hidden_size: int, + has_scale_ub: bool, + dtype: torch.dtype, + quant_dtype: torch.dtype, + group_size: int, + is_scale_transposed: bool, + seed: int, + device: str, +) -> None: + """Test SiLU+Mul+Block Quantization kernel correctness.""" + torch.random.manual_seed(seed) + torch.set_default_device(device) + + if hidden_size % group_size != 0: + return + + if has_scale_ub: + pytest.skip("Scale upper bound not yet supported") + + scale = 1 / hidden_size + x = torch.randn(num_tokens, hidden_size * 2, dtype=dtype, device=device) * scale + + # Reference implementation + ref_out, ref_scales = ref_silu_and_mul_per_block_quant(x, quant_dtype, group_size) + + # Fused kernel implementation + ops_out, ops_scales = ops.silu_and_mul_per_block_quant( + x, group_size, quant_dtype, None, is_scale_transposed + ) + + # Check for NaN/Inf + assert not torch.isnan(ops_out.float()).any(), "Kernel output contains NaN" + assert not torch.isinf(ops_out.float()).any(), "Kernel output contains Inf" + assert not torch.isnan(ops_scales).any(), "Kernel scales contain NaN" + assert not torch.isinf(ops_scales).any(), "Kernel scales contain Inf" + + # Check dtypes + assert ref_out.dtype == quant_dtype + assert ops_out.dtype == quant_dtype + + # Check scales match + torch.testing.assert_close(ref_scales, ops_scales, rtol=1e-5, atol=1e-5) + + # Check output correctness via dequantized values + ref_scales_expanded = ref_scales.repeat_interleave(group_size, dim=1) + ops_scales_expanded = ops_scales.repeat_interleave(group_size, dim=1) + ref_deq = ref_out.to(dtype=torch.float32) * ref_scales_expanded + ops_deq = ops_out.to(dtype=torch.float32) * ops_scales_expanded + torch.testing.assert_close(ref_deq, ops_deq, atol=5e-2, rtol=5e-2) + + # opcheck + output = torch.empty(num_tokens, hidden_size, device=device, dtype=quant_dtype) + num_groups = hidden_size // group_size + if is_scale_transposed: + scales = torch.empty(num_groups, num_tokens, device=device, dtype=torch.float32) + else: + scales = torch.empty(num_tokens, num_groups, device=device, dtype=torch.float32) + opcheck( + torch.ops._C.silu_and_mul_per_block_quant, + (output, x, scales, group_size, None, is_scale_transposed), + ) + + +@pytest.mark.parametrize("dtype", [torch.float16]) +@pytest.mark.parametrize("hidden_size", [4096]) +@pytest.mark.parametrize("num_tokens", [128]) +@pytest.mark.parametrize("group_size", [128]) +def test_silu_block_quant_shapes( + default_vllm_config, + dtype: torch.dtype, + hidden_size: int, + num_tokens: int, + group_size: int, +): + """Test that output shapes are correct.""" + torch.set_default_device("cuda") + x = torch.randn(num_tokens, hidden_size * 2, dtype=dtype, device="cuda") + + # Row-major scales + out, scales = ops.silu_and_mul_per_block_quant( + x, + group_size=group_size, + quant_dtype=torch.float8_e4m3fn, + is_scale_transposed=False, + ) + assert out.shape == (num_tokens, hidden_size) + assert scales.shape == (num_tokens, hidden_size // group_size) + + # Column-major scales (logical shape same after .t() in _custom_ops) + out, scales = ops.silu_and_mul_per_block_quant( + x, + group_size=group_size, + quant_dtype=torch.float8_e4m3fn, + is_scale_transposed=True, + ) + assert out.shape == (num_tokens, hidden_size) + assert scales.shape == (num_tokens, hidden_size // group_size) + + +@pytest.mark.parametrize("dtype", [torch.float16]) +@pytest.mark.parametrize("batch_size", [1, 16, 256]) +@pytest.mark.parametrize("hidden_size", [1024, 5120, 14336]) +def test_silu_block_quant_edge_cases( + default_vllm_config, dtype: torch.dtype, batch_size: int, hidden_size: int +): + """Test edge cases: single token, large batch, large hidden size.""" + torch.set_default_device("cuda") + x = torch.randn(batch_size, hidden_size * 2, dtype=dtype, device="cuda") + + out, scales = ops.silu_and_mul_per_block_quant( + x, + group_size=128, + quant_dtype=torch.float8_e4m3fn, + is_scale_transposed=False, + ) + + assert out.shape == (batch_size, hidden_size) + assert out.dtype == torch.float8_e4m3fn + assert scales.dtype == torch.float32 + assert not torch.isnan(out.float()).any() + assert not torch.isnan(scales).any() + assert not torch.isinf(scales).any() diff --git a/tests/kernels/moe/test_cpu_fused_moe.py b/tests/kernels/moe/test_cpu_fused_moe.py index 467ba3c5f69..73859175cd1 100644 --- a/tests/kernels/moe/test_cpu_fused_moe.py +++ b/tests/kernels/moe/test_cpu_fused_moe.py @@ -20,7 +20,7 @@ EXPERT_NUM = [ HIDDEN_DIM = [128, 2880] INTERMEDIATE_DIM = [128, 2880] BATCH_SIZE = [1, 64, 256] -ACT = [MoEActivation.SILU, MoEActivation.SWIGLUOAI] +ACT = [MoEActivation.SILU, MoEActivation.SWIGLUOAI, MoEActivation.GELU] USE_BIAS = [True, False] ISA = ["amx", "vec"] if torch.cpu._is_amx_tile_supported() else ["vec"] DTYPE = [torch.bfloat16] diff --git a/tests/kernels/moe/test_router_gemm.py b/tests/kernels/moe/test_router_gemm.py deleted file mode 100644 index 906e47708f2..00000000000 --- a/tests/kernels/moe/test_router_gemm.py +++ /dev/null @@ -1,37 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Tests for optimized router GEMM kernel - -Run `pytest tests/kernels/moe/test_router_gemm.py`. -""" - -import pytest -import torch - -import vllm._custom_ops as ops -from vllm.platforms import current_platform -from vllm.utils.torch_utils import set_random_seed - - -@pytest.mark.skipif( - not ( - current_platform.is_cuda() - and ( - current_platform.is_device_capability(90) - or current_platform.is_device_capability_family(100) - ) - ), - reason="This test only runs on Hopper or Blackwell GPUs.", -) -@pytest.mark.parametrize("batch_size", [1, 2, 4, 8]) -@pytest.mark.parametrize("input_dim", [360, 720, 1440, 2880]) -@pytest.mark.parametrize("output_dim", [32, 64, 128]) -def test_gpt_oss_router_gemm(batch_size, input_dim, output_dim): - set_random_seed(0) - x = torch.randn(batch_size, input_dim, device="cuda", dtype=torch.bfloat16) - weight = torch.randn(output_dim, input_dim, device="cuda", dtype=torch.bfloat16) - bias = torch.randn(output_dim, device="cuda", dtype=torch.bfloat16) - - output = ops.gpt_oss_router_gemm(x, weight, bias) - output_ref = torch.nn.functional.linear(x, weight, bias) - torch.testing.assert_close(output, output_ref, atol=1e-2, rtol=1e-2) diff --git a/tests/kernels/test_flex_attention.py b/tests/kernels/test_flex_attention.py index 69113b57c74..41d29813476 100644 --- a/tests/kernels/test_flex_attention.py +++ b/tests/kernels/test_flex_attention.py @@ -26,6 +26,59 @@ MINIMUM_TORCH_VERSION = version.parse("2.7.0") DIRECT_BUILD_VERSION = version.parse("2.9.dev0") +@pytest.mark.skipif( + not torch.cuda.is_available() or TORCH_VERSION < MINIMUM_TORCH_VERSION, + reason="CUDA not available or PyTorch version < 2.7", +) +def test_flex_attention_full_cudagraphs(vllm_runner): + """Test the numerics for flex attention full cudagraphs support.""" + model_name = "Qwen/Qwen2.5-1.5B-Instruct" + seed = 42 + max_tokens = 24 + num_logprobs = 5 + prompts = [ + "Hello, my name is", + "The president of the United States is", + "The capital of France is", + ] + + # Run with flex attention eager + set_random_seed(seed) + with vllm_runner( + model_name, + runner="generate", + tensor_parallel_size=1, + num_gpu_blocks_override=128, + enforce_eager=True, + attention_config={"backend": "FLEX_ATTENTION"}, + ) as llm_flex: + output_eager = llm_flex.generate_greedy_logprobs( + prompts, max_tokens, num_logprobs + ) + + # Run with flex attention compiled + set_random_seed(seed) + with vllm_runner( + model_name, + runner="generate", + tensor_parallel_size=1, + num_gpu_blocks_override=128, + enforce_eager=False, + gpu_memory_utilization=0.85, + attention_config={"backend": "FLEX_ATTENTION"}, + ) as llm_default: + output_compile = llm_default.generate_greedy_logprobs( + prompts, max_tokens, num_logprobs + ) + + check_logprobs_close( + outputs_0_lst=output_eager, + outputs_1_lst=output_compile, + name_0="eager", + name_1="compile", + ) + + @pytest.mark.skipif( not torch.cuda.is_available() or TORCH_VERSION < MINIMUM_TORCH_VERSION, reason="CUDA not available or PyTorch version < 2.7", diff --git a/tests/kernels/test_fused_gdn_post_conv.py b/tests/kernels/test_fused_gdn_post_conv.py new file mode 100644 index 00000000000..ffc8ce281f9 --- /dev/null +++ b/tests/kernels/test_fused_gdn_post_conv.py @@ -0,0 +1,209 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for fused_gdn_prefill_post_conv kernel. + +Verifies that the fused kernel matches the reference: + split → rearrange → contiguous → l2norm → gating +""" + +import pytest +import torch +import torch.nn.functional as F + +from vllm.model_executor.layers.fla.ops.fused_gdn_prefill_post_conv import ( + fused_post_conv_prep, +) + + +def reference_post_conv( + conv_output: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + H: int, + K: int, + V: int, + apply_l2norm: bool = True, + output_g_exp: bool = False, +): + """Reference implementation using individual ops.""" + L = conv_output.shape[0] + HV = A_log.shape[0] + + # Split + q_flat, k_flat, v_flat = torch.split(conv_output, [H * K, H * K, HV * V], dim=-1) + + # Rearrange + contiguous + q = q_flat.view(L, H, K).contiguous() + k = k_flat.view(L, H, K).contiguous() + v = v_flat.view(L, HV, V).contiguous() + + # L2 norm + if apply_l2norm: + q = F.normalize(q.float(), p=2, dim=-1, eps=1e-6).to(conv_output.dtype) + k = F.normalize(k.float(), p=2, dim=-1, eps=1e-6).to(conv_output.dtype) + + # Gating + x = a.float() + dt_bias.float() + sp = F.softplus(x, beta=1.0, threshold=20.0) + g = -torch.exp(A_log.float()) * sp + + if output_g_exp: + g = torch.exp(g) + + beta_out = torch.sigmoid(b.float()) + + return q, k, v, g, beta_out + + +# Qwen3.5-35B config: H=16, HV=32, K=128, V=128 +# Qwen3.5-397B config: H=16, HV=64, K=128, V=128 +@pytest.mark.parametrize( + "H, HV, K, V", + [ + (16, 32, 128, 128), # 35B + (16, 64, 128, 128), # 397B + (4, 8, 64, 64), # small + ], +) +@pytest.mark.parametrize("L", [1, 16, 128, 512, 2048]) +@pytest.mark.parametrize("apply_l2norm", [True, False]) +@pytest.mark.parametrize("output_g_exp", [True, False]) +@pytest.mark.parametrize("dtype", [torch.bfloat16]) +def test_fused_post_conv_correctness(H, HV, K, V, L, apply_l2norm, output_g_exp, dtype): + """Test fused kernel matches reference for all configs.""" + torch.manual_seed(42) + device = "cuda" + qkv_dim = 2 * H * K + HV * V + + conv_output = torch.randn(L, qkv_dim, dtype=dtype, device=device) + a = torch.randn(L, HV, dtype=dtype, device=device) + b = torch.randn(L, HV, dtype=dtype, device=device) + A_log = torch.randn(HV, dtype=torch.float32, device=device) - 2.0 + dt_bias = torch.randn(HV, dtype=torch.float32, device=device) * 0.1 + + # Reference + ref_q, ref_k, ref_v, ref_g, ref_beta = reference_post_conv( + conv_output, + a, + b, + A_log, + dt_bias, + H, + K, + V, + apply_l2norm, + output_g_exp, + ) + + # Fused kernel + fused_q, fused_k, fused_v, fused_g, fused_beta = fused_post_conv_prep( + conv_output, + a, + b, + A_log, + dt_bias, + num_k_heads=H, + head_k_dim=K, + head_v_dim=V, + apply_l2norm=apply_l2norm, + output_g_exp=output_g_exp, + ) + + # Check shapes + assert fused_q.shape == (L, H, K), f"q shape: {fused_q.shape}" + assert fused_k.shape == (L, H, K), f"k shape: {fused_k.shape}" + assert fused_v.shape == (L, HV, V), f"v shape: {fused_v.shape}" + assert fused_g.shape == (L, HV), f"g shape: {fused_g.shape}" + assert fused_beta.shape == (L, HV), f"beta shape: {fused_beta.shape}" + + # Check dtypes + assert fused_q.dtype == dtype + assert fused_k.dtype == dtype + assert fused_v.dtype == dtype + assert fused_g.dtype == torch.float32 + assert fused_beta.dtype == torch.float32 + + # Check contiguity + assert fused_q.is_contiguous() + assert fused_k.is_contiguous() + assert fused_v.is_contiguous() + + # Check values + atol_qkv = 1e-2 if apply_l2norm else 1e-3 + rtol_qkv = 1e-2 if apply_l2norm else 1e-3 + + torch.testing.assert_close(fused_q, ref_q, atol=atol_qkv, rtol=rtol_qkv) + torch.testing.assert_close(fused_k, ref_k, atol=atol_qkv, rtol=rtol_qkv) + torch.testing.assert_close(fused_v, ref_v, atol=1e-3, rtol=1e-3) + torch.testing.assert_close(fused_g, ref_g, atol=1e-4, rtol=1e-4) + torch.testing.assert_close(fused_beta, ref_beta, atol=1e-4, rtol=1e-4) + + +@pytest.mark.parametrize("L", [1, 64, 256]) +def test_fused_post_conv_sanity(L): + """Sanity checks: no NaN, unit-norm q/k, beta in (0,1).""" + torch.manual_seed(0) + device = "cuda" + H, HV, K, V = 16, 32, 128, 128 + qkv_dim = 2 * H * K + HV * V + + conv_output = torch.randn(L, qkv_dim, dtype=torch.bfloat16, device=device) + a = torch.randn(L, HV, dtype=torch.bfloat16, device=device) + b = torch.randn(L, HV, dtype=torch.bfloat16, device=device) + A_log = torch.randn(HV, dtype=torch.float32, device=device) - 2.0 + dt_bias = torch.randn(HV, dtype=torch.float32, device=device) + + q, k, v, g, beta = fused_post_conv_prep( + conv_output, + a, + b, + A_log, + dt_bias, + num_k_heads=H, + head_k_dim=K, + head_v_dim=V, + ) + + # Basic sanity + assert not torch.isnan(q).any(), "NaN in q" + assert not torch.isnan(k).any(), "NaN in k" + assert not torch.isnan(v).any(), "NaN in v" + assert not torch.isnan(g).any(), "NaN in g" + assert not torch.isnan(beta).any(), "NaN in beta" + + # L2 norm check: each head vector should have unit norm + q_norms = torch.norm(q.float(), dim=-1) + k_norms = torch.norm(k.float(), dim=-1) + torch.testing.assert_close(q_norms, torch.ones_like(q_norms), atol=1e-3, rtol=1e-3) + torch.testing.assert_close(k_norms, torch.ones_like(k_norms), atol=1e-3, rtol=1e-3) + + # Beta should be in (0, 1) + assert (beta >= 0).all() and (beta <= 1).all(), "beta out of range" + + +def test_fused_post_conv_l0(): + """Test L=0 edge case.""" + device = "cuda" + H, HV, K, V = 16, 32, 128, 128 + qkv_dim = 2 * H * K + HV * V + + conv_output = torch.empty(0, qkv_dim, dtype=torch.bfloat16, device=device) + a = torch.empty(0, HV, dtype=torch.bfloat16, device=device) + b = torch.empty(0, HV, dtype=torch.bfloat16, device=device) + A_log = torch.randn(HV, dtype=torch.float32, device=device) + dt_bias = torch.randn(HV, dtype=torch.float32, device=device) + + q, k, v, g, beta = fused_post_conv_prep( + conv_output, + a, + b, + A_log, + dt_bias, + num_k_heads=H, + head_k_dim=K, + head_v_dim=V, + ) + assert q.shape == (0, H, K) + assert g.shape == (0, HV) diff --git a/tests/lora/test_fused_moe_lora_kernel.py b/tests/lora/test_fused_moe_lora_kernel.py index 66a985a067e..8adc2086575 100644 --- a/tests/lora/test_fused_moe_lora_kernel.py +++ b/tests/lora/test_fused_moe_lora_kernel.py @@ -637,7 +637,7 @@ def use_fused_moe_lora_kernel_tensor_parallel( set_random_seed(seed) - device = torch.device(f"cuda:{local_rank}") + device = torch.device(f"{DEVICE_TYPE}:{local_rank}") torch.accelerator.set_device_index(device) torch.set_default_device(device) torch.set_default_dtype(dtype) diff --git a/tests/lora/test_layers.py b/tests/lora/test_layers.py index 08fd037249b..2a37abac6d7 100644 --- a/tests/lora/test_layers.py +++ b/tests/lora/test_layers.py @@ -60,8 +60,12 @@ pytestmark = pytest.mark.skipif( reason="Backend not supported", ) +DEVICE_TYPE = current_platform.device_type DEVICES = ( - [f"cuda:{i}" for i in range(1 if torch.accelerator.device_count() == 1 else 2)] + [ + f"{DEVICE_TYPE}:{i}" + for i in range(1 if torch.accelerator.device_count() == 1 else 2) + ] if current_platform.is_cuda_alike() else ["cpu"] ) @@ -196,7 +200,7 @@ def create_random_inputs( input_size: tuple[int, ...], input_range: tuple[float, float], input_type: torch.dtype = torch.int, - device: torch.device = "cuda", + device: torch.device = DEVICE_TYPE, ) -> tuple[list[torch.Tensor], list[int], list[int]]: """Creates random inputs. diff --git a/tests/lora/test_lora_manager.py b/tests/lora/test_lora_manager.py index e7addab119d..e80d96f00e7 100644 --- a/tests/lora/test_lora_manager.py +++ b/tests/lora/test_lora_manager.py @@ -35,9 +35,9 @@ EMBEDDING_MODULES = { "lm_head": "output_embeddings", } - +DEVICE_TYPE = current_platform.device_type DEVICES = ( - [f"cuda:{i}" for i in range(1 if torch.accelerator.device_count() == 1 else 2)] + [f"{DEVICE_TYPE}:{i}" for i in range(min(torch.accelerator.device_count(), 2))] if current_platform.is_cuda_alike() else ["cpu"] ) diff --git a/tests/lora/test_moe_lora_align_sum.py b/tests/lora/test_moe_lora_align_sum.py index bb46b4d8680..1451ec162b7 100644 --- a/tests/lora/test_moe_lora_align_sum.py +++ b/tests/lora/test_moe_lora_align_sum.py @@ -6,6 +6,9 @@ import pytest import torch from vllm import _custom_ops as ops +from vllm.platforms import current_platform + +DEVICE_TYPE = current_platform.device_type def round_up(x, base): @@ -27,7 +30,7 @@ def sample_data(num_experts, max_loras, num_tokens, topk_num): topk_ids[i, j] = pool[j] token_lora_mapping[i] = random.randint(0, max_loras - 1) - return topk_ids.to("cuda"), token_lora_mapping.to("cuda") + return topk_ids.to(DEVICE_TYPE), token_lora_mapping.to(DEVICE_TYPE) @pytest.mark.parametrize("num_tokens", [100, 200, 1024, 4096]) # 81920 @@ -56,14 +59,21 @@ def test_moe_lora_align_block_size( (max_loras * max_num_tokens_padded,), topk_ids.numel(), dtype=torch.int32, - device="cuda", + device=DEVICE_TYPE, ) expert_ids = torch.full( - (max_loras * max_num_m_blocks,), num_experts, dtype=torch.int32, device="cuda" + (max_loras * max_num_m_blocks,), + num_experts, + dtype=torch.int32, + device=DEVICE_TYPE, ) - num_tokens_post_pad = torch.zeros((max_loras,), dtype=torch.int32, device="cuda") - adapter_enabled = torch.ones((max_loras + 1,), dtype=torch.int32, device="cuda") - lora_ids = torch.arange(max_loras + 2, dtype=torch.int32, device="cuda") + num_tokens_post_pad = torch.zeros( + (max_loras,), dtype=torch.int32, device=DEVICE_TYPE + ) + adapter_enabled = torch.ones( + (max_loras + 1,), dtype=torch.int32, device=DEVICE_TYPE + ) + lora_ids = torch.arange(max_loras + 2, dtype=torch.int32, device=DEVICE_TYPE) # call kernel ops.moe_lora_align_block_size( diff --git a/tests/lora/test_punica_ops.py b/tests/lora/test_punica_ops.py index 8a2634e82ba..7706d0e2aab 100644 --- a/tests/lora/test_punica_ops.py +++ b/tests/lora/test_punica_ops.py @@ -9,10 +9,13 @@ import vllm.lora.ops.torch_ops as torch_ops import vllm.lora.ops.triton_ops as triton_ops from vllm.lora.ops.triton_ops import LoRAKernelMeta from vllm.lora.ops.triton_ops.utils import _LORA_A_PTR_DICT, _LORA_B_PTR_DICT +from vllm.platforms import current_platform from vllm.utils.torch_utils import set_random_seed from .utils import PunicaTensors, assert_close, generate_data_for_nslices +DEVICE_TYPE = current_platform.device_type + @pytest.fixture(autouse=True) def reset_device(reset_default_device): @@ -146,7 +149,9 @@ def check_lora_shrink_kernel( # Setup metadata information for the LoRA kernel. lora_meta = LoRAKernelMeta.make( - max_loras=num_loras, max_num_tokens=token_nums, device="cuda" + max_loras=num_loras, + max_num_tokens=token_nums, + device=DEVICE_TYPE, ) lora_meta.prepare_tensors(data.token_lora_mapping) @@ -219,7 +224,9 @@ def check_lora_expand_kernel( # Setup metadata information for the LoRA kernel. lora_meta = LoRAKernelMeta.make( - max_loras=num_loras, max_num_tokens=token_nums, device="cuda" + max_loras=num_loras, + max_num_tokens=token_nums, + device=DEVICE_TYPE, ) lora_meta.prepare_tensors(data.token_lora_mapping) @@ -367,7 +374,7 @@ test_params = { } DTYPES = [torch.float16, torch.bfloat16] -DEVICES = [f"cuda:{0}"] +DEVICES = [f"{DEVICE_TYPE}:{0}"] SEED = [0] diff --git a/tests/lora/test_punica_ops_fp8.py b/tests/lora/test_punica_ops_fp8.py index 04231333642..3e7fe7b2758 100644 --- a/tests/lora/test_punica_ops_fp8.py +++ b/tests/lora/test_punica_ops_fp8.py @@ -28,9 +28,11 @@ from vllm.lora.ops.triton_ops.lora_shrink_fp8_op import ( _SHRINK_LORA_SCALE_PTR_DICT, ) from vllm.lora.ops.triton_ops.utils import _LORA_A_PTR_DICT, _LORA_B_PTR_DICT +from vllm.platforms import current_platform from vllm.utils.torch_utils import set_random_seed -DEVICES = [f"cuda:{0}"] +DEVICE_TYPE = current_platform.device_type +DEVICES = [f"{DEVICE_TYPE}:{0}"] SEED = [0] _dict_lock = Lock() diff --git a/tests/lora/test_worker.py b/tests/lora/test_worker.py index 4af3ccf893f..88763551cf1 100644 --- a/tests/lora/test_worker.py +++ b/tests/lora/test_worker.py @@ -19,11 +19,14 @@ from vllm.config.load import LoadConfig from vllm.config.lora import LoRAConfig from vllm.lora.model_manager import LoRAMapping from vllm.lora.request import LoRARequest +from vllm.platforms import current_platform from vllm.v1.worker.gpu_worker import Worker MODEL_PATH = "Qwen/Qwen3-0.6B" NUM_LORAS = 16 +DEVICE_TYPE = current_platform.device_type + @patch.dict(os.environ, {"RANK": "0"}) def test_worker_apply_lora(qwen3_lora_files): @@ -61,7 +64,7 @@ def test_worker_apply_lora(qwen3_lora_files): max_num_seqs=32, max_num_partial_prefills=32, ), - device_config=DeviceConfig("cuda"), + device_config=DeviceConfig(DEVICE_TYPE), cache_config=CacheConfig( block_size=16, cache_dtype="auto", diff --git a/tests/lora/utils.py b/tests/lora/utils.py index 6aba5299b58..e5ce7a88464 100644 --- a/tests/lora/utils.py +++ b/tests/lora/utils.py @@ -9,10 +9,13 @@ import torch from safetensors.torch import save_file from vllm.lora.lora_weights import LoRALayerWeights, PackedLoRALayerWeights +from vllm.platforms import current_platform + +DEVICE_TYPE = current_platform.device_type class DummyLoRAManager: - def __init__(self, device: torch.device = "cuda:0"): + def __init__(self, device: torch.device = f"{DEVICE_TYPE}:0"): super().__init__() self._loras: dict[str, LoRALayerWeights] = {} self._device = device @@ -57,8 +60,8 @@ class DummyLoRAManager: module_name, rank=rank, lora_alpha=1, - lora_a=torch.rand([rank, input_dim], device="cuda"), - lora_b=torch.rand([output_dim, input_dim], device="cuda"), + lora_a=torch.rand([rank, input_dim], device=DEVICE_TYPE), + lora_b=torch.rand([output_dim, input_dim], device=DEVICE_TYPE), embeddings_tensor=embeddings_tensor, ) self.set_module_lora(module_name, lora) diff --git a/tests/models/language/generation/test_hybrid.py b/tests/models/language/generation/test_hybrid.py index 0f587558b50..01d395b1e0d 100644 --- a/tests/models/language/generation/test_hybrid.py +++ b/tests/models/language/generation/test_hybrid.py @@ -60,6 +60,14 @@ MAX_NUM_SEQS = 4 ATTN_BACKEND = "TRITON_ATTN" if current_platform.is_rocm() else "auto" +def _set_conv_state_layout(monkeypatch, layout: str) -> None: + """Set conv state layout env var and clear cache to pick up new value.""" + from vllm.model_executor.layers.mamba import mamba_utils + + monkeypatch.setenv("VLLM_SSM_CONV_STATE_LAYOUT", layout) + mamba_utils.get_conv_state_layout.cache_clear() + + @pytest.mark.parametrize("model", SSM_MODELS + HYBRID_MODELS) @pytest.mark.parametrize("max_tokens", [64]) @pytest.mark.parametrize("num_logprobs", [5]) @@ -102,12 +110,15 @@ def test_models( @pytest.mark.parametrize("model", [SSM_MODELS[0], HYBRID_MODELS[0]]) @pytest.mark.parametrize("max_tokens", [64]) @pytest.mark.parametrize("num_logprobs", [5]) +@pytest.mark.parametrize("conv_state_layout", ["SD", "DS"]) def test_batching( vllm_runner, example_prompts, + monkeypatch, model: str, max_tokens: int, num_logprobs: int, + conv_state_layout: str, ) -> None: try: model_info = HF_EXAMPLE_MODELS.find_hf_info(model) @@ -116,6 +127,8 @@ def test_batching( except ValueError: pass + _set_conv_state_layout(monkeypatch, conv_state_layout) + for_loop_outputs = [] with vllm_runner(model, max_num_seqs=MAX_NUM_SEQS) as vllm_model: for prompt in example_prompts: @@ -138,11 +151,14 @@ def test_batching( @pytest.mark.parametrize("model", [SSM_MODELS[0], HYBRID_MODELS[0]]) @pytest.mark.parametrize("max_tokens", [10]) +@pytest.mark.parametrize("conv_state_layout", ["SD", "DS"]) def test_chunked_prefill_with_parallel_sampling( vllm_runner, example_prompts, + monkeypatch, model: str, max_tokens: int, + conv_state_layout: str, ) -> None: """ Tests chunked prefill in conjunction with n > 1. @@ -154,6 +170,8 @@ def test_chunked_prefill_with_parallel_sampling( decoding steps inside a chunked prefill forward pass (where we have both prefill and decode together) """ + _set_conv_state_layout(monkeypatch, conv_state_layout) + sampling_params = SamplingParams(n=3, temperature=1, seed=0, max_tokens=max_tokens) with vllm_runner( model, @@ -168,17 +186,22 @@ def test_chunked_prefill_with_parallel_sampling( @pytest.mark.parametrize("model", [SSM_MODELS[0], HYBRID_MODELS[0]]) @pytest.mark.parametrize("max_tokens", [20]) +@pytest.mark.parametrize("conv_state_layout", ["SD", "DS"]) def test_mamba_cache_cg_padding( vllm_runner, example_prompts, + monkeypatch, model: str, max_tokens: int, + conv_state_layout: str, ) -> None: """ This test is for verifying that mamba cache is padded to CG captured batch size. If it's not, a torch RuntimeError will be raised because tensor dimensions aren't compatible. """ + _set_conv_state_layout(monkeypatch, conv_state_layout) + vllm_config = EngineArgs(model=model, trust_remote_code=True).create_engine_config() cudagraph_dispatcher = CudagraphDispatcher(vllm_config) cudagraph_dispatcher.initialize_cudagraph_keys( diff --git a/tests/models/multimodal/generation/test_common.py b/tests/models/multimodal/generation/test_common.py index 1404d9628fa..bf5119cf44f 100644 --- a/tests/models/multimodal/generation/test_common.py +++ b/tests/models/multimodal/generation/test_common.py @@ -394,6 +394,22 @@ VLM_TEST_SETTINGS = { vllm_runner_kwargs={"mm_processor_kwargs": {"do_pan_and_scan": True}}, patch_hf_runner=model_utils.gemma3_patch_hf_runner, ), + "gemma4": VLMTestInfo( + models=["google/gemma-4-E2B-it"], + test_type=(VLMTestType.IMAGE, VLMTestType.MULTI_IMAGE), + prompt_formatter=lambda img_prompt: f"user\n{img_prompt}\nmodel\n", # noqa: E501 + single_image_prompts=IMAGE_ASSETS.prompts( + { + "stop_sign": "What's the content in the center of the image?", + "cherry_blossom": "What is the season?", + } + ), + multi_image_prompt="Describe the two images in detail.", + max_model_len=4096, + max_num_seqs=2, + auto_cls=AutoModelForImageTextToText, + vllm_runner_kwargs={"limit_mm_per_prompt": {"image": 4}}, + ), "granite_vision": VLMTestInfo( models=["ibm-granite/granite-vision-3.3-2b"], test_type=(VLMTestType.IMAGE), diff --git a/tests/models/multimodal/generation/test_phi4siglip.py b/tests/models/multimodal/generation/test_phi4siglip.py new file mode 100644 index 00000000000..e8f4ba82925 --- /dev/null +++ b/tests/models/multimodal/generation/test_phi4siglip.py @@ -0,0 +1,187 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from collections.abc import Sequence + +import pytest +import regex as re +from transformers import AutoModelForCausalLM, AutoTokenizer + +from vllm.logprobs import SampleLogprobs +from vllm.multimodal.image import rescale_image_size + +from ....conftest import ( + IMAGE_ASSETS, + HfRunner, + PromptImageInput, + VllmRunner, +) +from ....utils import multi_gpu_test +from ...utils import check_logprobs_close + +MODEL_ID = "microsoft/Phi-4-reasoning-vision-15B" + +HF_IMAGE_PROMPTS = IMAGE_ASSETS.prompts( + { + "stop_sign": "<|user|>\n\nWhat's the content of the image?<|end|>\n<|assistant|>\n", # noqa: E501 + "cherry_blossom": "<|user|>\n\nPlease infer the season with reason in details.<|end|>\n<|assistant|>\n", # noqa: E501 + } +) +HF_MULTIIMAGE_IMAGE_PROMPT = ( + "<|user|>\n\n\nDescribe these images.<|end|>\n<|assistant|>\n" # noqa: E501 +) + +DTYPE = "half" +MAX_TOKENS = 128 +NUM_LOGPROBS = 10 + + +def vllm_to_hf_output( + vllm_output: tuple[list[int], str, SampleLogprobs | None], model: str +): + """Sanitize vllm output to be comparable with hf output.""" + _, output_str, out_logprobs = vllm_output + + output_str_without_image = re.sub(r"()+", "", output_str) + if output_str_without_image and output_str_without_image[0] == " ": + output_str_without_image = output_str_without_image[1:] + + hf_output_str = output_str_without_image + "<|end|><|endoftext|>" + + tokenizer = AutoTokenizer.from_pretrained(model, trust_remote_code=True) + hf_output_ids = tokenizer.encode(output_str_without_image) + if hf_output_ids and hf_output_ids[0] == tokenizer.bos_token_id: + hf_output_ids = hf_output_ids[1:] + + return hf_output_ids, hf_output_str, out_logprobs + + +def _build_single_image_inputs( + image_assets, +) -> list[tuple[list[str], PromptImageInput]]: + """Build single-image inputs for all size_factors at once.""" + images = [asset.pil_image for asset in image_assets] + all_inputs: list[tuple[list[str], PromptImageInput]] = [] + for size_factors in [[1.0], [0.25, 0.5, 1.0]]: + for image, prompt in zip(images, HF_IMAGE_PROMPTS): + all_inputs.append( + ( + [prompt for _ in size_factors], + [rescale_image_size(image, f) for f in size_factors], + ) + ) + return all_inputs + + +def _build_multi_image_inputs( + image_assets, +) -> list[tuple[list[str], PromptImageInput]]: + """Build multi-image inputs for all size_factors at once.""" + images = [asset.pil_image for asset in image_assets] + all_inputs: list[tuple[list[str], PromptImageInput]] = [] + for size_factors in [[0.5], [0.15, 0.30]]: + all_inputs.append( + ( + [HF_MULTIIMAGE_IMAGE_PROMPT for _ in size_factors], + [ + [rescale_image_size(image, factor) for image in images] + for factor in size_factors + ], + ) + ) + return all_inputs + + +def _run_and_compare( + hf_runner: type[HfRunner], + vllm_runner: type[VllmRunner], + all_inputs: Sequence[tuple[list[str], PromptImageInput]], + model: str, + max_model_len: int, + max_num_seqs: int, + mm_limit: int, + gpu_memory_utilization: float, +): + """Load each runner once, run all inputs, then compare.""" + # NOTE: run vLLM first, then HF. vLLM needs a fresh process without + # cuda initialization; running HF first would break the multiprocessing + # backend with fork method. + with vllm_runner( + model, + runner="generate", + max_model_len=max_model_len, + max_num_seqs=max_num_seqs, + gpu_memory_utilization=gpu_memory_utilization, + dtype=DTYPE, + limit_mm_per_prompt={"image": mm_limit}, + tensor_parallel_size=2, + trust_remote_code=True, + enforce_eager=True, + ) as vllm_model: + vllm_outputs_per_case = [ + vllm_model.generate_greedy_logprobs( + prompts, + MAX_TOKENS, + num_logprobs=NUM_LOGPROBS, + images=images, + ) + for prompts, images in all_inputs + ] + + hf_model_kwargs = {"_attn_implementation": "sdpa", "device_map": "auto"} + with hf_runner( + model, + dtype=DTYPE, + model_kwargs=hf_model_kwargs, + auto_cls=AutoModelForCausalLM, + trust_remote_code=True, + ) as hf_model: + hf_outputs_per_case = [ + hf_model.generate_greedy_logprobs_limit( + prompts, + MAX_TOKENS, + num_logprobs=NUM_LOGPROBS, + images=images, + ) + for prompts, images in all_inputs + ] + + for hf_outputs, vllm_outputs in zip(hf_outputs_per_case, vllm_outputs_per_case): + check_logprobs_close( + outputs_0_lst=hf_outputs, + outputs_1_lst=vllm_outputs, + name_0="hf", + name_1="vllm", + ) + + +@multi_gpu_test(num_gpus=2) +@pytest.mark.parametrize("model", [MODEL_ID]) +def test_models(hf_runner, vllm_runner, image_assets, model) -> None: + all_inputs = _build_single_image_inputs(image_assets) + _run_and_compare( + hf_runner, + vllm_runner, + all_inputs, + model, + max_model_len=8192, + max_num_seqs=2, + mm_limit=1, + gpu_memory_utilization=0.80, + ) + + +@multi_gpu_test(num_gpus=2) +@pytest.mark.parametrize("model", [MODEL_ID]) +def test_multi_images_models(hf_runner, vllm_runner, image_assets, model) -> None: + all_inputs = _build_multi_image_inputs(image_assets) + _run_and_compare( + hf_runner, + vllm_runner, + all_inputs, + model, + max_model_len=8192, + max_num_seqs=2, + mm_limit=2, + gpu_memory_utilization=0.80, + ) diff --git a/tests/models/multimodal/processing/test_gemma4.py b/tests/models/multimodal/processing/test_gemma4.py new file mode 100644 index 00000000000..808fab6a030 --- /dev/null +++ b/tests/models/multimodal/processing/test_gemma4.py @@ -0,0 +1,44 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest + +from vllm.multimodal import MULTIMODAL_REGISTRY + +from ....conftest import ImageTestAssets +from ...utils import build_model_context + +# TODO: to be updated to "google/gemma-4-e2b-it" once the models are available +GEMMA4_MODEL_ID = "google/gemma-4-E2B-it" + + +@pytest.mark.parametrize("model_id", [GEMMA4_MODEL_ID]) +def test_limit_mm_per_prompt( + image_assets: ImageTestAssets, + model_id: str, +): + """Test that limit_mm_per_prompt accurately restricts multiple images.""" + # We only allow 1 image + ctx = build_model_context( + model_id, + mm_processor_kwargs={}, + limit_mm_per_prompt={"image": 1}, + ) + processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config) + + # Provide 2 images in the prompt + prompt = "" + # image_assets usually has multiple images + images = [asset.pil_image for asset in image_assets][:2] + if len(images) < 2: + images = [images[0], images[0]] + + mm_data = {"image": images} + + # Expect ValueError when exceeding limit + with pytest.raises(ValueError, match="At most 1 image"): + processor( + prompt, + mm_items=processor.info.parse_mm_data(mm_data), + hf_processor_mm_kwargs={}, + ) diff --git a/tests/models/quantization/test_per_token_kv_cache.py b/tests/models/quantization/test_per_token_kv_cache.py new file mode 100644 index 00000000000..c581f01eb92 --- /dev/null +++ b/tests/models/quantization/test_per_token_kv_cache.py @@ -0,0 +1,94 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""End-to-end accuracy tests for per-token-head KV cache quantization. + +Compares logprobs between a baseline bf16 model and the same model with +per-token-head quantized KV cache (int8 or fp8) using the Triton attention +backend. + +Run: pytest tests/models/quantization/test_per_token_kv_cache.py -v -s +""" + +import pytest + +from vllm.platforms import current_platform + +from ..utils import check_logprobs_close + + +@pytest.mark.skipif( + not current_platform.is_cuda_alike(), + reason="Per-token-head KV cache requires CUDA or ROCm GPU.", +) +@pytest.mark.parametrize( + "base_model,test_model", + [ + ( + "meta-llama/Llama-3.2-1B-Instruct", + "meta-llama/Llama-3.2-1B-Instruct", + ), + ], +) +@pytest.mark.parametrize( + "kv_cache_dtype", ["int8_per_token_head", "fp8_per_token_head"] +) +@pytest.mark.parametrize("max_tokens", [4]) +@pytest.mark.parametrize("enforce_eager", [True]) +@pytest.mark.parametrize("backend", ["TRITON_ATTN"]) +@pytest.mark.parametrize("tensor_parallel_size", [1]) +def test_per_token_head_kv_cache_accuracy( + vllm_runner, + example_prompts, + base_model: str, + test_model: str, + kv_cache_dtype: str, + max_tokens: int, + enforce_eager: bool, + backend: str, + tensor_parallel_size: int, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Compare logprobs between bf16 baseline and per-token-head quantized KV + cache. + + Uses calculate_kv_scales (dynamic scale computation) since there are + no per-token-head calibrated checkpoints available yet. + """ + with monkeypatch.context() as m: + m.setenv("TOKENIZERS_PARALLELISM", "true") + + MAX_MODEL_LEN = 1024 + NUM_LOG_PROBS = 8 + + with vllm_runner( + base_model, + max_model_len=MAX_MODEL_LEN, + tensor_parallel_size=tensor_parallel_size, + enforce_eager=enforce_eager, + kv_cache_dtype="auto", + attention_config={"backend": backend}, + ) as vllm_model: + baseline_outputs = vllm_model.generate_greedy_logprobs( + example_prompts, max_tokens, NUM_LOG_PROBS + ) + + with vllm_runner( + test_model, + max_model_len=MAX_MODEL_LEN, + tensor_parallel_size=tensor_parallel_size, + enforce_eager=enforce_eager, + kv_cache_dtype=kv_cache_dtype, + calculate_kv_scales=True, + attention_config={"backend": backend}, + ) as vllm_model: + test_outputs = vllm_model.generate_greedy_logprobs( + example_prompts, max_tokens, NUM_LOG_PROBS + ) + + check_logprobs_close( + outputs_0_lst=baseline_outputs, + outputs_1_lst=test_outputs, + name_0="bf16_kv_cache", + name_1=f"{kv_cache_dtype}_kv_cache", + ) diff --git a/tests/models/registry.py b/tests/models/registry.py index 98c2a041068..895dc457927 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -7,6 +7,7 @@ from typing import Any, Literal import pytest from packaging.version import Version +from transformers import PretrainedConfig from transformers import __version__ as TRANSFORMERS_VERSION from vllm.config.model import ModelDType, TokenizerMode @@ -277,6 +278,10 @@ _TEXT_GENERATION_EXAMPLE_MODELS = { "google/gemma-2-9b", extras={"tiny": "google/gemma-2-2b-it"} ), "Gemma3ForCausalLM": _HfExamplesInfo("google/gemma-3-1b-it"), + "Gemma4ForCausalLM": _HfExamplesInfo( + "google/gemma-4-E2B-it", + min_transformers_version="5.0.0", + ), "Gemma3nForCausalLM": _HfExamplesInfo("google/gemma-3n-E2B-it"), "GlmForCausalLM": _HfExamplesInfo("zai-org/glm-4-9b-chat-hf"), "Glm4ForCausalLM": _HfExamplesInfo("zai-org/GLM-4-9B-0414"), @@ -533,6 +538,9 @@ _TEXT_GENERATION_EXAMPLE_MODELS = { "TeleChat2ForCausalLM": _HfExamplesInfo( "Tele-AI/TeleChat2-3B", trust_remote_code=True ), + "TeleChat3ForCausalLM": _HfExamplesInfo( + "Tele-AI/TeleChat3-36B-Thinking", trust_remote_code=True + ), "TeleFLMForCausalLM": _HfExamplesInfo( "CofeAI/FLM-2-52B-Instruct-2407", trust_remote_code=True ), @@ -766,6 +774,14 @@ _MULTIMODAL_EXAMPLE_MODELS = { extras={"6b": "Salesforce/blip2-opt-6.7b"}, ), "ChameleonForConditionalGeneration": _HfExamplesInfo("facebook/chameleon-7b"), + "Cheers": _HfExamplesInfo( + "ai9stars/Cheers", + trust_remote_code=True, + ), + "CheersForConditionalGeneration": _HfExamplesInfo( + "ai9stars/Cheers", + trust_remote_code=True, + ), "Cohere2VisionForConditionalGeneration": _HfExamplesInfo( "CohereLabs/command-a-vision-07-2025" ), @@ -805,6 +821,10 @@ _MULTIMODAL_EXAMPLE_MODELS = { ), "FuyuForCausalLM": _HfExamplesInfo("adept/fuyu-8b"), "Gemma3ForConditionalGeneration": _HfExamplesInfo("google/gemma-3-4b-it"), + "Gemma4ForConditionalGeneration": _HfExamplesInfo( + "google/gemma-4-E2B-it", + min_transformers_version="5.5.0", + ), "Gemma3nForConditionalGeneration": _HfExamplesInfo("google/gemma-3n-E2B-it"), "GlmAsrForConditionalGeneration": _HfExamplesInfo( "zai-org/GLM-ASR-Nano-2512", @@ -985,7 +1005,26 @@ _MULTIMODAL_EXAMPLE_MODELS = { trust_remote_code=True, ), "NemotronH_Nano_VL_V2": _HfExamplesInfo( - "nano_vl_dummy", is_available_online=False, trust_remote_code=True + "nvidia/NVIDIA-Nemotron-Nano-12B-v2-VL-BF16", + max_model_len=4096, + # NemotronH layers are constructed via `hybrid_override_pattern`: + use_original_num_layers=True, + hf_overrides={ + "vision_config": PretrainedConfig( + args={ + "min_num_patches": 1, # Trigger image dynamic res + "max_num_patches": 12, + "model": "vit_huge_patch16_224", + }, + # Trigger conv3d: + video_temporal_patch_size=2, + ), + "text_config": { + "num_hidden_layers": 2, + "hybrid_override_pattern": "M*", + }, + }, + trust_remote_code=True, ), "OpenCUAForConditionalGeneration": _HfExamplesInfo( "xlangai/OpenCUA-7B", trust_remote_code=True @@ -1030,6 +1069,9 @@ _MULTIMODAL_EXAMPLE_MODELS = { }, # noqa: E501 extras={"phi3.5": "microsoft/Phi-3.5-vision-instruct"}, ), + "Phi4ForCausalLMV": _HfExamplesInfo( + "microsoft/Phi-4-reasoning-vision-15B", trust_remote_code=True + ), "Phi4MMForCausalLM": _HfExamplesInfo( "microsoft/Phi-4-multimodal-instruct", trust_remote_code=True ), diff --git a/tests/models/utils.py b/tests/models/utils.py index 4830f18dccf..6d6636c96f9 100644 --- a/tests/models/utils.py +++ b/tests/models/utils.py @@ -447,9 +447,16 @@ def dummy_hf_overrides( Dummy HF overrides function used to create dummy model with only minimum nums of layer. """ - hf_config.update(exist_overrides or {}) + # Copy because this helper is called more than once + # while loading config, and we `.pop()` + exist_overrides = (exist_overrides or {}).copy() + text_config_override = exist_overrides.pop("text_config", None) + hf_config.update(exist_overrides) text_config = hf_config.get_text_config() + if text_config_override is not None: + # multimodal test models may override *some* text-model fields + text_config.update(text_config_override) # Ensure at least 2 expert per group # Since `grouped_topk` assumes top-2 diff --git a/tests/quantization/test_per_token_kv_cache.py b/tests/quantization/test_per_token_kv_cache.py new file mode 100644 index 00000000000..3e660e6b00d --- /dev/null +++ b/tests/quantization/test_per_token_kv_cache.py @@ -0,0 +1,560 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for per-token-head KV cache quantization (INT8 and FP8). + +Covers: +- Per-token-head Triton reshape-and-cache kernel +- Round-trip quantize/dequantize accuracy +- process_weights_after_loading early-return path +- End-to-end integration with Triton unified attention kernel + +Run: pytest tests/quantization/test_per_token_kv_cache.py -v -s +""" + +import random +from dataclasses import dataclass +from unittest.mock import MagicMock + +import pytest +import torch + +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + get_fp8_min_max, +) +from vllm.platforms import current_platform +from vllm.utils.torch_utils import set_random_seed +from vllm.v1.kv_cache_interface import KVQuantMode, is_quantized_kv_cache + +# Skip entire module if no CUDA/ROCm GPU available +pytestmark = [ + pytest.mark.skipif( + not current_platform.is_cuda_alike(), + reason="Per-token-head KV cache tests require CUDA or ROCm GPU.", + ), +] + +# --------------------------------------------------------------------------- +# Test parameters +# --------------------------------------------------------------------------- +NUM_TOKENS = [1, 7, 42] +NUM_KV_HEADS = [1, 4, 8] +HEAD_SIZES = [64, 128] +BLOCK_SIZES = [16] +SEEDS = [0] + +# Platform-dependent FP8 dtype and range +FP8_DTYPE = current_platform.fp8_dtype() +FP8_MIN, FP8_MAX = get_fp8_min_max() + + +# --------------------------------------------------------------------------- +# Per-dtype quantization config +# --------------------------------------------------------------------------- +@dataclass(frozen=True) +class QuantConfig: + """Quantization parameters for a given cache dtype.""" + + cache_dtype: torch.dtype # torch.int8 or FP8_DTYPE + kv_cache_dtype_str: str # "int8_per_token_head" or "fp8_per_token_head" + quant_max: float + quant_min: float + kv_quant_mode: KVQuantMode + # INT8 Triton stores truncate; FP8 hardware casts round. + uses_trunc: bool + + +INT8_CONFIG = QuantConfig( + cache_dtype=torch.int8, + kv_cache_dtype_str="int8_per_token_head", + quant_max=127.0, + quant_min=-128.0, + kv_quant_mode=KVQuantMode.INT8_PER_TOKEN_HEAD, + uses_trunc=True, +) +FP8_CONFIG = QuantConfig( + cache_dtype=FP8_DTYPE, + kv_cache_dtype_str="fp8_per_token_head", + quant_max=FP8_MAX, + quant_min=FP8_MIN, + kv_quant_mode=KVQuantMode.FP8_PER_TOKEN_HEAD, + uses_trunc=False, +) + +QUANT_CONFIGS = [INT8_CONFIG, FP8_CONFIG] + + +@pytest.fixture(params=QUANT_CONFIGS, ids=["int8", "fp8"]) +def qcfg(request) -> QuantConfig: + return request.param + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- +def _quantize_per_token_head_ref( + data: torch.Tensor, # [num_tokens, num_heads, head_size] + cfg: QuantConfig, +) -> tuple[torch.Tensor, torch.Tensor]: + """Reference per-token-head quantization (one scale per token per head). + + Returns (quantized, scales) where scales is [num_tokens, num_heads]. + """ + absmax = data.float().abs().amax(dim=2) # [num_tokens, num_heads] + scales = (absmax / cfg.quant_max).clamp(min=1e-6) + scaled = data.float() * (1.0 / scales[:, :, None]) + if cfg.uses_trunc: + q = scaled.round().clamp(cfg.quant_min, cfg.quant_max).to(cfg.cache_dtype) + else: + q = scaled.clamp(cfg.quant_min, cfg.quant_max).to(cfg.cache_dtype) + return q, scales + + +# =========================================================================== +# 1. is_quantized_kv_cache / get_kv_quant_mode +# =========================================================================== +class TestIsQuantizedKvCache: + def test_fp8_variants(self): + assert is_quantized_kv_cache("fp8") + assert is_quantized_kv_cache("fp8_e4m3") + assert is_quantized_kv_cache("fp8_e5m2") + + def test_int8_per_token_head(self): + assert is_quantized_kv_cache("int8_per_token_head") + + def test_fp8_per_token_head(self): + assert is_quantized_kv_cache("fp8_per_token_head") + + def test_auto(self): + assert not is_quantized_kv_cache("auto") + + def test_bfloat16(self): + assert not is_quantized_kv_cache("bfloat16") + + def test_kv_quant_mode_int8(self): + from vllm.v1.kv_cache_interface import get_kv_quant_mode + + assert ( + get_kv_quant_mode("int8_per_token_head") == KVQuantMode.INT8_PER_TOKEN_HEAD + ) + + def test_kv_quant_mode_fp8(self): + from vllm.v1.kv_cache_interface import get_kv_quant_mode + + assert get_kv_quant_mode("fp8_per_token_head") == KVQuantMode.FP8_PER_TOKEN_HEAD + + +# =========================================================================== +# 2. Triton per-token-head kernel (reshape-and-cache) +# =========================================================================== +@pytest.mark.parametrize("num_tokens", NUM_TOKENS) +@pytest.mark.parametrize("num_heads", NUM_KV_HEADS) +@pytest.mark.parametrize("head_size", HEAD_SIZES) +@pytest.mark.parametrize("block_size", BLOCK_SIZES) +@pytest.mark.parametrize("seed", SEEDS) +@torch.inference_mode() +def test_reshape_and_cache_per_token_head( + qcfg: QuantConfig, + num_tokens: int, + num_heads: int, + head_size: int, + block_size: int, + seed: int, +): + """Test triton_reshape_and_cache_flash_per_token_head_quant kernel.""" + from vllm.v1.attention.ops.triton_reshape_and_cache_flash import ( + triton_reshape_and_cache_flash_per_token_head_quant, + ) + + set_random_seed(seed) + torch.set_default_device("cuda") + + num_blocks = (num_tokens + block_size - 1) // block_size + 4 + + key = torch.randn(num_tokens, num_heads, head_size, dtype=torch.bfloat16) + value = torch.randn(num_tokens, num_heads, head_size, dtype=torch.bfloat16) + + key_cache = torch.zeros( + num_blocks, block_size, num_heads, head_size, dtype=qcfg.cache_dtype + ) + value_cache = torch.zeros( + num_blocks, block_size, num_heads, head_size, dtype=qcfg.cache_dtype + ) + k_scale_cache = torch.ones(num_blocks, block_size, num_heads, dtype=torch.float32) + v_scale_cache = torch.ones(num_blocks, block_size, num_heads, dtype=torch.float32) + + num_slots = block_size * num_blocks + slot_mapping = torch.tensor( + random.sample(range(num_slots), num_tokens), dtype=torch.long + ) + + triton_reshape_and_cache_flash_per_token_head_quant( + key, + value, + key_cache, + value_cache, + k_scale_cache, + v_scale_cache, + slot_mapping, + ) + + # Reference + ref_k_quant, ref_k_scales = _quantize_per_token_head_ref(key, qcfg) + ref_v_quant, ref_v_scales = _quantize_per_token_head_ref(value, qcfg) + + # Compare dequantized values rather than raw quantized values. + # Triton and PyTorch reductions can differ at FP8 rounding boundaries + # (up to 32 in quantized domain for fp8_e4m3), but the dequantized + # error is bounded by the scale. + for i, slot in enumerate(slot_mapping.tolist()): + blk = slot // block_size + off = slot % block_size + + actual_k_scale = k_scale_cache[blk, off] # [num_heads] + k_deq = key_cache[blk, off].float() * actual_k_scale[:, None] + k_ref_deq = key[i].float() + torch.testing.assert_close( + k_deq, + k_ref_deq, + atol=0.1, + rtol=0.1, + ) + actual_v_scale = v_scale_cache[blk, off] # [num_heads] + v_deq = value_cache[blk, off].float() * actual_v_scale[:, None] + v_ref_deq = value[i].float() + torch.testing.assert_close( + v_deq, + v_ref_deq, + atol=0.1, + rtol=0.1, + ) + # Per-head scales: [num_heads] + torch.testing.assert_close( + k_scale_cache[blk, off], ref_k_scales[i], atol=1e-4, rtol=1e-3 + ) + torch.testing.assert_close( + v_scale_cache[blk, off], ref_v_scales[i], atol=1e-4, rtol=1e-3 + ) + + +# =========================================================================== +# 3. Per-token-head round-trip accuracy (quantize -> dequantize) +# =========================================================================== +@pytest.mark.parametrize("num_tokens", [1, 16]) +@pytest.mark.parametrize("num_heads", [4]) +@pytest.mark.parametrize("head_size", [128]) +@pytest.mark.parametrize("block_size", [16]) +@torch.inference_mode() +def test_per_token_head_round_trip_accuracy( + qcfg: QuantConfig, + num_tokens: int, + num_heads: int, + head_size: int, + block_size: int, +): + """Verify per-token-head round-trip: kernel dequant matches reference. + + INT8: Triton truncates on float->int8 store. + FP8: hardware cast (clamp then cast). + """ + from vllm.v1.attention.ops.triton_reshape_and_cache_flash import ( + triton_reshape_and_cache_flash_per_token_head_quant, + ) + + torch.set_default_device("cuda") + set_random_seed(42) + + num_blocks = (num_tokens + block_size - 1) // block_size + 2 + + key = torch.randn(num_tokens, num_heads, head_size, dtype=torch.bfloat16) * 0.5 + value = torch.randn(num_tokens, num_heads, head_size, dtype=torch.bfloat16) * 0.5 + + key_cache = torch.zeros( + num_blocks, block_size, num_heads, head_size, dtype=qcfg.cache_dtype + ) + value_cache = torch.zeros( + num_blocks, block_size, num_heads, head_size, dtype=qcfg.cache_dtype + ) + k_scale_cache = torch.ones(num_blocks, block_size, num_heads, dtype=torch.float32) + v_scale_cache = torch.ones(num_blocks, block_size, num_heads, dtype=torch.float32) + + slot_mapping = torch.arange(num_tokens, dtype=torch.long) + + triton_reshape_and_cache_flash_per_token_head_quant( + key, + value, + key_cache, + value_cache, + k_scale_cache, + v_scale_cache, + slot_mapping, + ) + + for i in range(num_tokens): + blk = i // block_size + off = i % block_size + + for label, data, cache, sc in [ + ("key", key, key_cache, k_scale_cache), + ("val", value, value_cache, v_scale_cache), + ]: + for h in range(num_heads): + orig = data[i, h].float() # [head_size] + + actual_q = cache[blk, off, h] + actual_sc = sc[blk, off, h] + actual_deq = actual_q.float() * actual_sc + + # Round-trip: dequantized should be close to original + torch.testing.assert_close( + actual_deq, + orig, + atol=0.1, + rtol=0.1, + ) + + +# =========================================================================== +# 4. Negative slot mapping (padding tokens should be skipped) +# =========================================================================== +@torch.inference_mode() +def test_per_token_head_negative_slot_skipped(qcfg: QuantConfig): + """Tokens with slot_mapping=-1 should leave the cache unchanged.""" + from vllm.v1.attention.ops.triton_reshape_and_cache_flash import ( + triton_reshape_and_cache_flash_per_token_head_quant, + ) + + torch.set_default_device("cuda") + num_tokens = 4 + num_heads = 2 + head_size = 64 + block_size = 16 + num_blocks = 2 + + key = torch.randn(num_tokens, num_heads, head_size, dtype=torch.bfloat16) + value = torch.randn(num_tokens, num_heads, head_size, dtype=torch.bfloat16) + + key_cache = torch.zeros( + num_blocks, block_size, num_heads, head_size, dtype=qcfg.cache_dtype + ) + value_cache = torch.zeros( + num_blocks, block_size, num_heads, head_size, dtype=qcfg.cache_dtype + ) + k_scale_cache = torch.ones(num_blocks, block_size, num_heads, dtype=torch.float32) + v_scale_cache = torch.ones(num_blocks, block_size, num_heads, dtype=torch.float32) + + slot_mapping = torch.tensor([0, -1, 1, -1], dtype=torch.long) + + key_cache_before = key_cache.clone() + val_cache_before = value_cache.clone() + + triton_reshape_and_cache_flash_per_token_head_quant( + key, + value, + key_cache, + value_cache, + k_scale_cache, + v_scale_cache, + slot_mapping, + ) + + # Slots 0 and 1 should have been written (tokens 0 and 2) + assert not torch.equal(key_cache[0, 0], key_cache_before[0, 0]) + assert not torch.equal(key_cache[0, 1], key_cache_before[0, 1]) + assert not torch.equal(value_cache[0, 0], val_cache_before[0, 0]) + + # All other slots should be unchanged + assert torch.equal(key_cache[0, 2:], key_cache_before[0, 2:]) + assert torch.equal(key_cache[1], key_cache_before[1]) + assert torch.equal(value_cache[0, 2:], val_cache_before[0, 2:]) + + +# =========================================================================== +# 5. process_weights_after_loading -- per-token-head early return +# =========================================================================== +@pytest.mark.parametrize( + "kv_cache_dtype", ["int8_per_token_head", "fp8_per_token_head"] +) +def test_process_weights_sets_placeholder_scales(kv_cache_dtype: str): + """Per-token-head should set _k_scale=1.0, _v_scale=1.0 + and delete checkpoint attrs.""" + from vllm.model_executor.layers.quantization.kv_cache import ( + BaseKVCacheMethod, + ) + + layer = MagicMock() + layer.kv_cache_dtype = kv_cache_dtype + layer.calculate_kv_scales = False + layer.k_scale = torch.nn.Parameter(torch.tensor(-1.0), requires_grad=False) + layer.v_scale = torch.nn.Parameter(torch.tensor(-1.0), requires_grad=False) + layer.q_scale = torch.nn.Parameter(torch.tensor(-1.0), requires_grad=False) + layer.prob_scale = torch.nn.Parameter(torch.tensor(-1.0), requires_grad=False) + layer._k_scale = torch.tensor(0.0) + layer._v_scale = torch.tensor(0.0) + layer._k_scale_float = 0.0 + layer._v_scale_float = 0.0 + + method = BaseKVCacheMethod.__new__(BaseKVCacheMethod) + method.quant_config = MagicMock() + method.process_weights_after_loading(layer) + + assert layer._k_scale_float == 1.0 + assert layer._v_scale_float == 1.0 + assert not hasattr(layer, "k_scale") + assert not hasattr(layer, "v_scale") + assert not hasattr(layer, "q_scale") + assert not hasattr(layer, "prob_scale") + + +# =========================================================================== +# 6. Triton unified_attention -- per-token-head scale cache (INT8 and FP8) +# =========================================================================== +@pytest.mark.parametrize( + "seq_lens", + [ + [(1, 128)], + [(1, 64), (1, 32)], + ], +) +@pytest.mark.parametrize("num_heads", [(4, 4)]) +@pytest.mark.parametrize("head_size", [128]) +@pytest.mark.parametrize("block_size", [16]) +@torch.inference_mode() +def test_triton_unified_attention_per_token_head_scale( + qcfg: QuantConfig, + seq_lens: list[tuple[int, int]], + num_heads: tuple[int, int], + head_size: int, + block_size: int, +): + """End-to-end: quantized KV with per-token-head scale caches.""" + from vllm.utils.math_utils import next_power_of_2 + from vllm.v1.attention.ops.triton_unified_attention import unified_attention + + torch.set_default_device("cuda") + set_random_seed(0) + + num_seqs = len(seq_lens) + query_lens = [s[0] for s in seq_lens] + kv_lens = [s[1] for s in seq_lens] + num_query_heads, num_kv_heads = num_heads + max_query_len = max(query_lens) + max_kv_len = max(kv_lens) + scale = head_size**-0.5 + num_blocks = 2048 + + query = torch.randn( + sum(query_lens), num_query_heads, head_size, dtype=torch.bfloat16 + ) + + key_cache_bf16 = torch.randn( + num_blocks, block_size, num_kv_heads, head_size, dtype=torch.bfloat16 + ) + value_cache_bf16 = torch.randn_like(key_cache_bf16) + + # Per-token-head quantization: one scale per (block, slot, head) + k_absmax = key_cache_bf16.float().abs().amax(dim=-1) # [..., num_kv_heads] + v_absmax = value_cache_bf16.float().abs().amax(dim=-1) + k_scale_cache = (k_absmax / qcfg.quant_max).clamp(min=1e-6).to(torch.float32) + v_scale_cache = (v_absmax / qcfg.quant_max).clamp(min=1e-6).to(torch.float32) + + scaled_k = key_cache_bf16.float() / k_scale_cache[:, :, :, None] + scaled_v = value_cache_bf16.float() / v_scale_cache[:, :, :, None] + if qcfg.uses_trunc: + key_cache_q = ( + scaled_k.round().clamp(qcfg.quant_min, qcfg.quant_max).to(qcfg.cache_dtype) + ) + value_cache_q = ( + scaled_v.round().clamp(qcfg.quant_min, qcfg.quant_max).to(qcfg.cache_dtype) + ) + else: + key_cache_q = scaled_k.clamp(qcfg.quant_min, qcfg.quant_max).to( + qcfg.cache_dtype + ) + value_cache_q = scaled_v.clamp(qcfg.quant_min, qcfg.quant_max).to( + qcfg.cache_dtype + ) + + # Dequantized reference + key_cache_deq = key_cache_q.float() * k_scale_cache[:, :, :, None] + value_cache_deq = value_cache_q.float() * v_scale_cache[:, :, :, None] + + cu_query_lens = torch.tensor([0] + query_lens, dtype=torch.int32).cumsum( + dim=0, dtype=torch.int32 + ) + kv_lens_t = torch.tensor(kv_lens, dtype=torch.int32) + + max_num_blocks_per_seq = (max_kv_len + block_size - 1) // block_size + block_tables = torch.randint( + 0, num_blocks, (num_seqs, max_num_blocks_per_seq), dtype=torch.int32 + ) + + head_size_padded = next_power_of_2(head_size) + seq_threshold_3D = 0 + num_par_softmax_segments = 16 + softmax_segm_output = torch.empty( + (seq_threshold_3D, num_query_heads, num_par_softmax_segments, head_size_padded), + dtype=torch.float32, + ) + softmax_segm_max = torch.empty( + (seq_threshold_3D, num_query_heads, num_par_softmax_segments), + dtype=torch.float32, + ) + softmax_segm_expsum = torch.empty( + (seq_threshold_3D, num_query_heads, num_par_softmax_segments), + dtype=torch.float32, + ) + + output_q = torch.empty_like(query) + unified_attention( + q=query, + k=key_cache_q, + v=value_cache_q, + out=output_q, + cu_seqlens_q=cu_query_lens, + seqused_k=kv_lens_t, + max_seqlen_q=max_query_len, + max_seqlen_k=max_kv_len, + softmax_scale=scale, + causal=True, + window_size=(-1, -1), + block_table=block_tables, + softcap=0, + q_descale=None, + k_descale=None, + v_descale=None, + seq_threshold_3D=seq_threshold_3D, + num_par_softmax_segments=num_par_softmax_segments, + softmax_segm_output=softmax_segm_output, + softmax_segm_max=softmax_segm_max, + softmax_segm_expsum=softmax_segm_expsum, + kv_quant_mode=qcfg.kv_quant_mode, + k_scale_cache=k_scale_cache, + v_scale_cache=v_scale_cache, + ) + + output_ref = torch.empty_like(query) + unified_attention( + q=query, + k=key_cache_deq.to(torch.bfloat16), + v=value_cache_deq.to(torch.bfloat16), + out=output_ref, + cu_seqlens_q=cu_query_lens, + seqused_k=kv_lens_t, + max_seqlen_q=max_query_len, + max_seqlen_k=max_kv_len, + softmax_scale=scale, + causal=True, + window_size=(-1, -1), + block_table=block_tables, + softcap=0, + q_descale=None, + k_descale=None, + v_descale=None, + seq_threshold_3D=seq_threshold_3D, + num_par_softmax_segments=num_par_softmax_segments, + softmax_segm_output=softmax_segm_output, + softmax_segm_max=softmax_segm_max, + softmax_segm_expsum=softmax_segm_expsum, + ) + + torch.testing.assert_close(output_q, output_ref, atol=5e-2, rtol=5e-2) diff --git a/tests/reasoning/test_gemma4_reasoning_parser.py b/tests/reasoning/test_gemma4_reasoning_parser.py new file mode 100644 index 00000000000..cdda7dea51d --- /dev/null +++ b/tests/reasoning/test_gemma4_reasoning_parser.py @@ -0,0 +1,196 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest + +from tests.reasoning.utils import run_reasoning_extraction +from vllm.reasoning import ReasoningParser, ReasoningParserManager + +# Using mistral tokenizer as a generic mock since the actual model is not on HF +from vllm.tokenizers.registry import get_tokenizer + +parser_name = "gemma4" + + +@pytest.fixture(scope="module") +def generic_tokenizer(): + return get_tokenizer("google/gemma-4-E2B-it") + + +INVALID_SIMPLE_NONSTREAMING = { + "output": "This is a reasoning sectionThis is the rest", + "reasoning": "This is a reasoning section", + "content": "This is the rest", + "is_reasoning_end": True, +} +INVALID_SIMPLE_STREAMING = { + "output": "This is a reasoning sectionThis is the rest", + "reasoning": None, + "content": "This is a reasoning sectionThis is the rest", + "is_reasoning_end": True, +} +INVALID_COMPLETE_NONSTREAMING = { + "output": "This is a reasoning section", + "reasoning": "This is a reasoning section", + "content": None, + "is_reasoning_end": True, +} +INVALID_COMPLETE_STREAMING = { + "output": "This is a reasoning section", + "reasoning": None, + "content": "This is a reasoning section", + "is_reasoning_end": True, +} +NO_CONTENT = { + "output": "<|channel>This is reasoning", + "reasoning": "This is reasoning", + "content": None, + "is_reasoning_end": False, +} +NO_REASONING = { + "output": "This is content", + "reasoning": None, + "content": "This is content", + "is_reasoning_end": False, +} +REASONING_WITH_CHANNEL = { + "output": "<|channel>This is a reasoning sectionThis is the rest", + "reasoning": "This is a reasoning section", + "content": "This is the rest", + "is_reasoning_end": True, +} +COMPLETE_REASONING_WITH_CHANNEL = { + "output": "<|channel>This is a reasoning section", + "reasoning": "This is a reasoning section", + "content": None, + "is_reasoning_end": True, +} +MULTIPLE_LINES_WITH_CHANNEL = { + "output": "<|channel>This\nThatThis is the rest\nThat", + "reasoning": "This\nThat", + "content": "This is the rest\nThat", + "is_reasoning_end": True, +} +CHANNEL_NO_END = { + "output": "<|channel>This is a reasoning section", + "reasoning": "This is a reasoning section", + "content": None, + "is_reasoning_end": False, +} +EMPTY = { + "output": "", + "reasoning": None, + "content": "", + "is_reasoning_end": False, +} +NEW_LINE_NONSTREAMING = { + "output": ( + "Before\n<|channel>This is a reasoning section\nThis is the rest" + ), + "reasoning": "This is a reasoning section", + "content": "\nThis is the rest", + "is_reasoning_end": True, +} +NEW_LINE_STREAMING = { + "output": ( + "Before\n<|channel>This is a reasoning section\nThis is the rest" + ), + "reasoning": "This is a reasoning section", + "content": "Before\n\nThis is the rest", + "is_reasoning_end": True, +} + +TEST_CASES = [ + pytest.param(False, INVALID_SIMPLE_NONSTREAMING, id="invalid_simple"), + pytest.param(True, INVALID_SIMPLE_STREAMING, id="invalid_simple_streaming"), + pytest.param(False, INVALID_COMPLETE_NONSTREAMING, id="invalid_complete"), + pytest.param(True, INVALID_COMPLETE_STREAMING, id="invalid_complete_streaming"), + pytest.param(False, NO_CONTENT, id="no_content"), + pytest.param(False, NO_REASONING, id="no_reasoning"), + pytest.param(False, REASONING_WITH_CHANNEL, id="reasoning"), + pytest.param(True, REASONING_WITH_CHANNEL, id="reasoning_streaming"), + pytest.param(False, COMPLETE_REASONING_WITH_CHANNEL, id="complete_reasoning"), + pytest.param( + True, COMPLETE_REASONING_WITH_CHANNEL, id="complete_reasoning_streaming" + ), + pytest.param(False, MULTIPLE_LINES_WITH_CHANNEL, id="multiple_lines"), + pytest.param(True, MULTIPLE_LINES_WITH_CHANNEL, id="multiple_lines_streaming"), + pytest.param(False, CHANNEL_NO_END, id="no_end"), + pytest.param(True, CHANNEL_NO_END, id="no_end_streaming"), + pytest.param(False, EMPTY, id="empty"), + pytest.param(False, NEW_LINE_NONSTREAMING, id="new_line"), + pytest.param(True, NEW_LINE_STREAMING, id="new_line_streaming"), +] + + +@pytest.mark.parametrize("streaming, param_dict", TEST_CASES) +def test_gemma4_reasoning( + streaming: bool, + param_dict: dict, + generic_tokenizer, +): + output = param_dict["output"] + + # Resolve token IDs dynamically from the real tokenizer + vocab = generic_tokenizer.get_vocab() + start_token_id = vocab["<|channel>"] + end_token_id = vocab[""] + + index_start = output.find("<|channel>") + len_start = len("<|channel>") + index_end = output.find("") + len_end = len("") + + output_tokens = [] + + def _encode(text: str) -> list[int]: + if not text: + return [] + # Handle both raw transformers and vLLM wrappers + enc = getattr(generic_tokenizer, "tokenizer", generic_tokenizer) + try: + return enc.encode(text, add_special_tokens=False) + except TypeError: + return enc.encode(text) + + if index_start != -1: + output_before = output[:index_start] + output_tokens += _encode(output_before) + output_tokens += [start_token_id] + + if index_end != -1: + output_middle = output[index_start + len_start : index_end] + output_after = output[index_end + len_end :] + output_tokens += _encode(output_middle) + output_tokens += [end_token_id] + output_tokens += _encode(output_after) + else: + output_middle = output[index_start + len_start :] + output_tokens += _encode(output_middle) + elif index_end != -1: + output_before = output[:index_end] + output_after = output[index_end + len_end :] + output_tokens += _encode(output_before) + output_tokens += [end_token_id] + output_tokens += _encode(output_after) + else: + output_tokens += _encode(output) + + parser: ReasoningParser = ReasoningParserManager.get_reasoning_parser(parser_name)( + generic_tokenizer + ) + + # We use the generic run_reasoning_extraction from utils + # Use decode per token to get standard spaces instead of + # SentencePiece space characters + output_token_strings = [generic_tokenizer.decode([t]) for t in output_tokens] + reasoning, content = run_reasoning_extraction( + parser, output_token_strings, streaming=streaming + ) + + assert reasoning == param_dict["reasoning"] + assert content == param_dict["content"] + + # Test is_reasoning_end + is_reasoning_end = parser.is_reasoning_end(output_tokens) + assert is_reasoning_end == param_dict["is_reasoning_end"] diff --git a/tests/test_ray_env_utils.py b/tests/test_ray_env_utils.py new file mode 100644 index 00000000000..d311de41ba9 --- /dev/null +++ b/tests/test_ray_env_utils.py @@ -0,0 +1,51 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for vllm.v1.executor.ray_env_utils.""" + +import os +from unittest.mock import patch + +from vllm.v1.executor.ray_env_utils import get_driver_env_vars + +WORKER_VARS: set[str] = { + "CUDA_VISIBLE_DEVICES", + "LOCAL_RANK", +} + + +class TestDefaultPropagation: + """All env vars are propagated unless explicitly excluded.""" + + @patch.dict(os.environ, {"NCCL_DEBUG": "INFO"}, clear=False) + def test_nccl_prefix(self): + assert get_driver_env_vars(WORKER_VARS)["NCCL_DEBUG"] == "INFO" + + @patch.dict(os.environ, {"HF_TOKEN": "secret"}, clear=False) + def test_hf_token(self): + assert "HF_TOKEN" in get_driver_env_vars(WORKER_VARS) + + @patch.dict(os.environ, {"LMCACHE_LOCAL_CPU": "True"}, clear=False) + def test_lmcache_prefix(self): + assert "LMCACHE_LOCAL_CPU" in get_driver_env_vars(WORKER_VARS) + + @patch.dict(os.environ, {"PYTHONHASHSEED": "42"}, clear=False) + def test_pythonhashseed(self): + assert get_driver_env_vars(WORKER_VARS)["PYTHONHASHSEED"] == "42" + + @patch.dict(os.environ, {"MYLIB_FOO": "bar"}, clear=False) + def test_arbitrary_var_propagated(self): + assert get_driver_env_vars(WORKER_VARS)["MYLIB_FOO"] == "bar" + + +class TestExclusion: + @patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "0,1"}, clear=False) + def test_worker_specific_excluded(self): + assert "CUDA_VISIBLE_DEVICES" not in get_driver_env_vars(WORKER_VARS) + + @patch.dict(os.environ, {"LMCACHE_LOCAL_CPU": "True"}, clear=False) + @patch( + "vllm.v1.executor.ray_env_utils.RAY_NON_CARRY_OVER_ENV_VARS", + {"LMCACHE_LOCAL_CPU"}, + ) + def test_non_carry_over_blacklist(self): + assert "LMCACHE_LOCAL_CPU" not in get_driver_env_vars(WORKER_VARS) diff --git a/tests/tool_parsers/test_gemma4_tool_parser.py b/tests/tool_parsers/test_gemma4_tool_parser.py new file mode 100644 index 00000000000..80cf70d6c7d --- /dev/null +++ b/tests/tool_parsers/test_gemma4_tool_parser.py @@ -0,0 +1,504 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import json +from typing import Any +from unittest.mock import MagicMock + +import pytest + +from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest +from vllm.tool_parsers.gemma4_tool_parser import ( + TOOL_CALL_END, + TOOL_CALL_START, + Gemma4ToolParser, + _parse_gemma4_args, + _parse_gemma4_array, +) + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + + +@pytest.fixture +def mock_tokenizer(): + tokenizer = MagicMock() + tokenizer.encode.return_value = [1, 2, 3] + # Include the tool call start token in the vocab for the parser + tokenizer.get_vocab.return_value = {TOOL_CALL_START: 48, TOOL_CALL_END: 49} + return tokenizer + + +@pytest.fixture +def parser(mock_tokenizer): + return Gemma4ToolParser(mock_tokenizer) + + +@pytest.fixture +def mock_request(): + request = MagicMock(spec=ChatCompletionRequest) + request.tools = [] + request.tool_choice = "auto" + return request + + +# --------------------------------------------------------------------------- +# Unit tests for _parse_gemma4_args (shared parser logic) +# --------------------------------------------------------------------------- + + +class TestParseGemma4Args: + def test_empty_string(self): + assert _parse_gemma4_args("") == {} + + def test_whitespace_only(self): + assert _parse_gemma4_args(" ") == {} + + def test_single_string_value(self): + result = _parse_gemma4_args('location:<|"|>Paris<|"|>') + assert result == {"location": "Paris"} + + def test_string_value_with_comma(self): + result = _parse_gemma4_args('location:<|"|>Paris, France<|"|>') + assert result == {"location": "Paris, France"} + + def test_multiple_string_values(self): + result = _parse_gemma4_args( + 'location:<|"|>San Francisco<|"|>,unit:<|"|>celsius<|"|>' + ) + assert result == {"location": "San Francisco", "unit": "celsius"} + + def test_integer_value(self): + result = _parse_gemma4_args("count:42") + assert result == {"count": 42} + + def test_float_value(self): + result = _parse_gemma4_args("score:3.14") + assert result == {"score": 3.14} + + def test_boolean_true(self): + result = _parse_gemma4_args("flag:true") + assert result == {"flag": True} + + def test_boolean_false(self): + result = _parse_gemma4_args("flag:false") + assert result == {"flag": False} + + def test_mixed_types(self): + result = _parse_gemma4_args( + 'name:<|"|>test<|"|>,count:42,active:true,score:3.14' + ) + assert result == { + "name": "test", + "count": 42, + "active": True, + "score": 3.14, + } + + def test_nested_object(self): + result = _parse_gemma4_args('nested:{inner:<|"|>value<|"|>}') + assert result == {"nested": {"inner": "value"}} + + def test_array_of_strings(self): + result = _parse_gemma4_args('items:[<|"|>a<|"|>,<|"|>b<|"|>]') + assert result == {"items": ["a", "b"]} + + def test_unterminated_string(self): + """Unterminated strings should take everything after the delimiter.""" + result = _parse_gemma4_args('key:<|"|>unterminated') + assert result == {"key": "unterminated"} + + def test_empty_value(self): + """Key with no value after colon.""" + result = _parse_gemma4_args("key:") + assert result == {"key": ""} + + +class TestParseGemma4Array: + def test_string_array(self): + result = _parse_gemma4_array('<|"|>a<|"|>,<|"|>b<|"|>') + assert result == ["a", "b"] + + def test_empty_array(self): + result = _parse_gemma4_array("") + assert result == [] + + def test_bare_values(self): + result = _parse_gemma4_array("42,true,3.14") + assert result == [42, True, 3.14] + + +# --------------------------------------------------------------------------- +# Non-streaming extraction tests +# --------------------------------------------------------------------------- + + +class TestExtractToolCalls: + def test_no_tool_calls(self, parser, mock_request): + model_output = "Hello, how can I help you today?" + result = parser.extract_tool_calls(model_output, mock_request) + + assert result.tools_called is False + assert result.tool_calls == [] + assert result.content == model_output + + def test_single_tool_call(self, parser, mock_request): + model_output = ( + '<|tool_call>call:get_weather{location:<|"|>London<|"|>}' + ) + result = parser.extract_tool_calls(model_output, mock_request) + + assert result.tools_called is True + assert len(result.tool_calls) == 1 + assert result.tool_calls[0].function.name == "get_weather" + args = json.loads(result.tool_calls[0].function.arguments) + assert args == {"location": "London"} + + def test_multiple_arguments(self, parser, mock_request): + model_output = ( + "<|tool_call>call:get_weather{" + 'location:<|"|>San Francisco<|"|>,' + 'unit:<|"|>celsius<|"|>}' + "" + ) + result = parser.extract_tool_calls(model_output, mock_request) + + assert result.tools_called is True + assert len(result.tool_calls) == 1 + assert result.tool_calls[0].function.name == "get_weather" + args = json.loads(result.tool_calls[0].function.arguments) + assert args == {"location": "San Francisco", "unit": "celsius"} + + def test_text_before_tool_call(self, parser, mock_request): + model_output = ( + "Let me check the weather for you. " + '<|tool_call>call:get_weather{location:<|"|>Paris<|"|>}' + "" + ) + result = parser.extract_tool_calls(model_output, mock_request) + + assert result.tools_called is True + assert result.content == "Let me check the weather for you." + assert len(result.tool_calls) == 1 + assert result.tool_calls[0].function.name == "get_weather" + + def test_multiple_tool_calls(self, parser, mock_request): + model_output = ( + '<|tool_call>call:get_weather{location:<|"|>London<|"|>}' + "" + '<|tool_call>call:get_time{location:<|"|>London<|"|>}' + "" + ) + result = parser.extract_tool_calls(model_output, mock_request) + + assert result.tools_called is True + assert len(result.tool_calls) == 2 + assert result.tool_calls[0].function.name == "get_weather" + assert result.tool_calls[1].function.name == "get_time" + + def test_nested_arguments(self, parser, mock_request): + model_output = ( + "<|tool_call>call:complex_function{" + 'nested:{inner:<|"|>value<|"|>},' + 'list:[<|"|>a<|"|>,<|"|>b<|"|>]}' + "" + ) + result = parser.extract_tool_calls(model_output, mock_request) + + assert result.tools_called is True + assert len(result.tool_calls) == 1 + assert result.tool_calls[0].function.name == "complex_function" + args = json.loads(result.tool_calls[0].function.arguments) + assert args == {"nested": {"inner": "value"}, "list": ["a", "b"]} + + def test_tool_call_with_number_and_boolean(self, parser, mock_request): + model_output = ( + "<|tool_call>call:set_status{" + "is_active:true," + "count:42," + "score:3.14}" + "" + ) + result = parser.extract_tool_calls(model_output, mock_request) + + assert result.tools_called is True + assert len(result.tool_calls) == 1 + assert result.tool_calls[0].function.name == "set_status" + args = json.loads(result.tool_calls[0].function.arguments) + assert args == {"is_active": True, "count": 42, "score": 3.14} + + def test_incomplete_tool_call(self, parser, mock_request): + model_output = '<|tool_call>call:get_weather{location:<|"|>London' + result = parser.extract_tool_calls(model_output, mock_request) + + # Incomplete — no end marker, regex won't match + assert result.tools_called is False + assert result.content == model_output + + def test_hyphenated_function_name(self, parser, mock_request): + """Ensure function names with hyphens are parsed correctly.""" + model_output = ( + '<|tool_call>call:get-weather{location:<|"|>London<|"|>}' + ) + result = parser.extract_tool_calls(model_output, mock_request) + + assert result.tools_called is True + assert result.tool_calls[0].function.name == "get-weather" + + def test_dotted_function_name(self, parser, mock_request): + """Ensure function names with dots are parsed correctly.""" + model_output = ( + '<|tool_call>call:weather.get{location:<|"|>London<|"|>}' + ) + result = parser.extract_tool_calls(model_output, mock_request) + + assert result.tools_called is True + assert result.tool_calls[0].function.name == "weather.get" + + def test_no_arguments(self, parser, mock_request): + """Tool calls with empty arguments.""" + model_output = "<|tool_call>call:get_status{}" + result = parser.extract_tool_calls(model_output, mock_request) + + assert result.tools_called is True + assert result.tool_calls[0].function.name == "get_status" + args = json.loads(result.tool_calls[0].function.arguments) + assert args == {} + + +# --------------------------------------------------------------------------- +# Streaming extraction tests +# --------------------------------------------------------------------------- + + +class TestStreamingExtraction: + """Tests for the streaming tool call extraction. + + These simulate the token-by-token streaming that vLLM performs, + feeding incremental text to extract_tool_calls_streaming() and + verifying that the accumulated argument deltas form valid JSON. + """ + + def _simulate_streaming( + self, parser: Gemma4ToolParser, mock_request: Any, chunks: list[str] + ) -> list[tuple[Any, str]]: + """Feed chunks through the streaming parser and collect results. + + Returns a list of (delta_message, accumulated_text) tuples. + """ + results: list[tuple[Any, str]] = [] + previous_text: str = "" + previous_token_ids: list[int] = [] + + for chunk in chunks: + current_text = previous_text + chunk + # Use token ID 48 for tool_call start, 49 for end, 0 otherwise + delta_token_ids: list[int] = [] + if TOOL_CALL_START in chunk: + delta_token_ids.append(48) + elif TOOL_CALL_END in chunk: + delta_token_ids.append(49) + else: + delta_token_ids.append(0) + + current_token_ids = previous_token_ids + delta_token_ids + + delta = parser.extract_tool_calls_streaming( + previous_text=previous_text, + current_text=current_text, + delta_text=chunk, + previous_token_ids=tuple(previous_token_ids), + current_token_ids=tuple(current_token_ids), + delta_token_ids=tuple(delta_token_ids), + request=mock_request, + ) + results.append((delta, current_text)) + previous_text = current_text + previous_token_ids = list(current_token_ids) + + return results + + def _collect_arguments(self, results): + """Collect all argument deltas from streaming results into one string.""" + args_text = "" + for delta, _ in results: + if delta and delta.tool_calls: + for tc in delta.tool_calls: + func = tc.function if isinstance(tc.function, dict) else tc.function + if isinstance(func, dict): + arg = func.get("arguments", "") + else: + arg = getattr(func, "arguments", "") or "" + if arg: + args_text += arg + return args_text + + def _collect_function_name(self, results): + """Extract the function name from streaming results.""" + for delta, _ in results: + if delta and delta.tool_calls: + for tc in delta.tool_calls: + func = tc.function if isinstance(tc.function, dict) else tc.function + if isinstance(func, dict): + name = func.get("name") + else: + name = getattr(func, "name", None) + if name: + return name + return None + + def test_basic_streaming_single_tool(self, parser, mock_request): + """Simulate the exact streaming scenario from the bug report. + + Model generates: + <|tool_call>call:get_weather{location:<|"|>Paris, France<|"|>} + + Expected: arguments should be valid JSON {"location": "Paris, France"} + """ + chunks = [ + "<|tool_call>", + "call:get_weather{", + 'location:<|"|>Paris', + ", France", + '<|"|>}', + "", + ] + + results = self._simulate_streaming(parser, mock_request, chunks) + + # Verify function name + name = self._collect_function_name(results) + assert name == "get_weather", f"Expected 'get_weather', got '{name}'" + + # Verify arguments form valid JSON + args_text = self._collect_arguments(results) + assert args_text, "No arguments were streamed" + parsed_args = json.loads(args_text) + assert parsed_args == {"location": "Paris, France"} + + def test_streaming_multi_arg(self, parser, mock_request): + """Streaming with multiple arguments.""" + chunks = [ + "<|tool_call>", + "call:get_weather{", + 'location:<|"|>Tokyo<|"|>,', + 'unit:<|"|>celsius<|"|>}', + "", + ] + + results = self._simulate_streaming(parser, mock_request, chunks) + + name = self._collect_function_name(results) + assert name == "get_weather" + + args_text = self._collect_arguments(results) + assert args_text + parsed_args = json.loads(args_text) + assert parsed_args == {"location": "Tokyo", "unit": "celsius"} + + def test_streaming_no_extra_brace(self, parser, mock_request): + """Verify the closing } is NOT leaked into arguments (Bug #2).""" + chunks = [ + "<|tool_call>", + "call:get_weather{", + 'location:<|"|>London<|"|>}', + "", + ] + + results = self._simulate_streaming(parser, mock_request, chunks) + args_text = self._collect_arguments(results) + assert args_text + + # The args text must be valid JSON (no extra }) + parsed = json.loads(args_text) + assert parsed == {"location": "London"} + + # Specifically assert no double-brace + assert args_text.count("}") <= 1, ( + f"Arguments contain extra closing brace: {args_text!r}" + ) + + def test_streaming_no_unquoted_keys(self, parser, mock_request): + """Verify keys are properly quoted in JSON (Bug #1).""" + chunks = [ + "<|tool_call>", + "call:get_weather{", + 'location:<|"|>Paris<|"|>}', + "", + ] + + results = self._simulate_streaming(parser, mock_request, chunks) + args_text = self._collect_arguments(results) + + # Must start with { and contain quoted key + assert args_text.lstrip().startswith("{"), ( + f"Arguments don't start with '{{': {args_text!r}" + ) + assert '"location"' in args_text, ( + f"Key 'location' not properly quoted: {args_text!r}" + ) + + def test_streaming_name_no_call_prefix(self, parser, mock_request): + """Verify function name has no 'call:' prefix.""" + chunks = [ + "<|tool_call>", + "call:get_weather{", + 'location:<|"|>Paris<|"|>}', + "", + ] + + results = self._simulate_streaming(parser, mock_request, chunks) + name = self._collect_function_name(results) + assert name == "get_weather" + assert not name.startswith("call:"), f"Name has 'call:' prefix: {name!r}" + + def test_streaming_text_before_tool_call(self, parser, mock_request): + """Text before tool call should be emitted as content.""" + chunks = [ + "Let me check ", + "the weather. ", + "<|tool_call>", + "call:get_weather{", + 'location:<|"|>London<|"|>}', + "", + ] + + results = self._simulate_streaming(parser, mock_request, chunks) + + # First chunks should be content + content_parts = [] + for delta, _ in results: + if delta and delta.content: + content_parts.append(delta.content) + + assert "".join(content_parts).strip().startswith("Let me check") + + def test_streaming_numeric_args(self, parser, mock_request): + """Streaming with numeric and boolean argument values.""" + chunks = [ + "<|tool_call>", + "call:set_config{", + "count:42,", + "active:true}", + "", + ] + + results = self._simulate_streaming(parser, mock_request, chunks) + args_text = self._collect_arguments(results) + if args_text: + parsed_args = json.loads(args_text) + assert parsed_args["count"] == 42 + assert parsed_args["active"] is True + + def test_streaming_empty_args(self, parser, mock_request): + """Tool call with no arguments.""" + chunks = [ + "<|tool_call>", + "call:get_status{}", + "", + ] + + results = self._simulate_streaming(parser, mock_request, chunks) + name = self._collect_function_name(results) + assert name == "get_status" diff --git a/tests/tool_parsers/test_qwen3coder_tool_parser.py b/tests/tool_parsers/test_qwen3coder_tool_parser.py index 8f92d5b3745..7db1b685715 100644 --- a/tests/tool_parsers/test_qwen3coder_tool_parser.py +++ b/tests/tool_parsers/test_qwen3coder_tool_parser.py @@ -429,208 +429,6 @@ hello world assert args["obj_param"] == {"key": "value"} -def test_extract_tool_calls_anyof_type_conversion(qwen3_tool_parser): - """Test type conversion for anyOf/oneOf nullable schemas (Pydantic v2). - - Pydantic v2 emits anyOf for Optional[T] fields, e.g.: - Optional[int] -> {"anyOf": [{"type": "integer"}, {"type": "null"}]} - The parser must extract the non-null type and apply the correct - conversion (int(), float(), etc.) instead of returning a raw string. - """ - tools = [ - ChatCompletionToolsParam( - type="function", - function={ - "name": "test_anyof", - "parameters": { - "type": "object", - "properties": { - "anyof_int": { - "anyOf": [ - {"type": "integer"}, - {"type": "null"}, - ], - "default": 5, - }, - "anyof_str": { - "anyOf": [ - {"type": "string"}, - {"type": "null"}, - ], - }, - "anyof_array": { - "anyOf": [ - {"type": "array", "items": {"type": "string"}}, - {"type": "null"}, - ], - }, - "anyof_obj": { - "anyOf": [ - {"type": "object"}, - {"type": "null"}, - ], - }, - "type_as_array": { - "type": ["integer", "null"], - }, - "multi_non_null": { - "anyOf": [ - {"type": "string"}, - {"type": "integer"}, - {"type": "null"}, - ], - }, - "ref_param": { - "$ref": "#/$defs/ToolInput", - }, - }, - }, - }, - ) - ] - - model_output = """ - - -5 - - -hello - - -["a", "b", "c"] - - -{"key": "value"} - - -42 - - -some text - - -{"city": "Paris"} - - -""" - - request = ChatCompletionRequest(model=MODEL, messages=[], tools=tools) - extracted = qwen3_tool_parser.extract_tool_calls(model_output, request=request) - - args = json.loads(extracted.tool_calls[0].function.arguments) - assert args["anyof_int"] == 5 - assert isinstance(args["anyof_int"], int) - assert args["anyof_str"] == "hello" - assert isinstance(args["anyof_str"], str) - assert args["anyof_array"] == ["a", "b", "c"] - assert isinstance(args["anyof_array"], list) - assert args["anyof_obj"] == {"key": "value"} - assert isinstance(args["anyof_obj"], dict) - assert args["type_as_array"] == 42 - assert isinstance(args["type_as_array"], int) - # Multi non-null: anyOf[string, integer, null] → first non-null is string - assert args["multi_non_null"] == "some text" - assert isinstance(args["multi_non_null"], str) - # $ref: treated as object, parsed via json.loads - assert args["ref_param"] == {"city": "Paris"} - assert isinstance(args["ref_param"], dict) - - -def test_extract_tool_calls_anyof_type_conversion_streaming( - qwen3_tool_parser, qwen3_tokenizer -): - """Test streaming e2e for anyOf/oneOf nullable schemas (Pydantic v2). - - Verifies that the full streaming pipeline — tokenize, incrementally - decode, extract_tool_calls_streaming — correctly resolves types from - anyOf schemas and produces valid JSON with properly typed values. - """ - tools = [ - ChatCompletionToolsParam( - type="function", - function={ - "name": "search_web", - "parameters": { - "type": "object", - "properties": { - "query": { - "anyOf": [ - {"type": "string"}, - {"type": "null"}, - ], - }, - "count": { - "anyOf": [ - {"type": "integer"}, - {"type": "null"}, - ], - "default": 5, - }, - "verbose": { - "anyOf": [ - {"type": "boolean"}, - {"type": "null"}, - ], - }, - "filters": { - "$ref": "#/$defs/SearchFilters", - }, - }, - }, - }, - ) - ] - - model_output = """ - - -vllm tool parser - - -10 - - -true - - -{"lang": "en", "year": 2025} - - -""" - - request = ChatCompletionRequest(model=MODEL, messages=[], tools=tools) - - tool_states = {} - for delta_message in stream_delta_message_generator( - qwen3_tool_parser, qwen3_tokenizer, model_output, request - ): - if delta_message.tool_calls: - for tool_call in delta_message.tool_calls: - idx = tool_call.index - if idx not in tool_states: - tool_states[idx] = {"name": None, "arguments": ""} - if tool_call.function: - if tool_call.function.name: - tool_states[idx]["name"] = tool_call.function.name - if tool_call.function.arguments is not None: - tool_states[idx]["arguments"] += tool_call.function.arguments - - assert len(tool_states) == 1 - assert tool_states[0]["name"] == "search_web" - assert tool_states[0]["arguments"] is not None - args = json.loads(tool_states[0]["arguments"]) - assert args["query"] == "vllm tool parser" - assert isinstance(args["query"], str) - assert args["count"] == 10 - assert isinstance(args["count"], int) - assert args["verbose"] is True - assert isinstance(args["verbose"], bool) - # $ref: treated as object, parsed via json.loads - assert args["filters"] == {"lang": "en", "year": 2025} - assert isinstance(args["filters"], dict) - - @pytest.mark.parametrize( ids=[ "no_tools", diff --git a/tests/utils_/test_ray_utils.py b/tests/utils_/test_ray_utils.py new file mode 100644 index 00000000000..0872ae9413f --- /dev/null +++ b/tests/utils_/test_ray_utils.py @@ -0,0 +1,100 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from unittest.mock import MagicMock, patch + +import pytest + +from vllm.v1.executor.ray_utils import get_bundles_sorted_by_node + +NODE_A = "node_a" +NODE_B = "node_b" +NODE_C = "node_c" + +IP_A = "10.0.0.1" +IP_B = "10.0.0.2" +IP_C = "10.0.0.3" + +NODE_ID_TO_IP = {NODE_A: IP_A, NODE_B: IP_B, NODE_C: IP_C} + +MOCK_RAY_NODES = [ + {"NodeID": NODE_A, "NodeManagerAddress": IP_A, "Alive": True}, + {"NodeID": NODE_B, "NodeManagerAddress": IP_B, "Alive": True}, + {"NodeID": NODE_C, "NodeManagerAddress": IP_C, "Alive": True}, +] + + +@pytest.mark.parametrize( + "bundles_to_node_id,bundle_specs,expected", + [ + pytest.param( + {0: NODE_C, 1: NODE_A, 2: NODE_B, 3: NODE_C, 4: NODE_A, 5: NODE_B}, + [{"GPU": 1}] * 6, + [ + (1, NODE_A, IP_A), + (4, NODE_A, IP_A), + (2, NODE_B, IP_B), + (5, NODE_B, IP_B), + (0, NODE_C, IP_C), + (3, NODE_C, IP_C), + ], + ), + pytest.param( + {0: NODE_B, 1: NODE_B, 2: NODE_A, 3: NODE_A}, + [{"GPU": 1}] * 4, + [ + (2, NODE_A, IP_A), + (3, NODE_A, IP_A), + (0, NODE_B, IP_B), + (1, NODE_B, IP_B), + ], + ), + pytest.param( + {0: NODE_C, 1: NODE_B, 2: NODE_C, 3: NODE_B}, + [{"GPU": 1}] * 4, + [ + (1, NODE_B, IP_B), + (3, NODE_B, IP_B), + (0, NODE_C, IP_C), + (2, NODE_C, IP_C), + ], + ), + pytest.param( + {0: NODE_A, 1: NODE_A, 2: NODE_A}, + [{"GPU": 1}] * 3, + [(0, NODE_A, IP_A), (1, NODE_A, IP_A), (2, NODE_A, IP_A)], + ), + pytest.param( + {}, + [], + [], + ), + pytest.param( + {0: NODE_A, 1: NODE_B, 2: NODE_A}, + [{"CPU": 1}, {"GPU": 1}, {"GPU": 1}], + [(2, NODE_A, IP_A), (1, NODE_B, IP_B)], + ), + ], +) +def test_get_bundles_sorted_by_node(bundles_to_node_id, bundle_specs, expected): + mock_pg = MagicMock() + mock_pg.bundle_specs = bundle_specs + + mock_ctx = MagicMock() + mock_ctx.get_node_id.return_value = NODE_A + + with ( + patch( + "vllm.v1.executor.ray_utils.placement_group_table", + return_value={"bundles_to_node_id": bundles_to_node_id}, + ), + patch("vllm.v1.executor.ray_utils.ray") as mock_ray, + patch("vllm.v1.executor.ray_utils.current_platform") as mock_platform, + ): + mock_ray.get_runtime_context.return_value = mock_ctx + mock_ray.nodes.return_value = MOCK_RAY_NODES + mock_platform.ray_device_key = "GPU" + + result = get_bundles_sorted_by_node(mock_pg) + + assert result == expected diff --git a/tests/v1/attention/test_attention_backends.py b/tests/v1/attention/test_attention_backends.py index 8c3a62b6ea5..06095b87e59 100644 --- a/tests/v1/attention/test_attention_backends.py +++ b/tests/v1/attention/test_attention_backends.py @@ -40,6 +40,8 @@ BACKENDS_TO_TEST = [ "FLEX_ATTENTION_SLOW", ] +DEVICE_TYPE = current_platform.device_type + # Remove flashinfer from the list if it's not available try: import flashinfer # noqa: F401 @@ -366,7 +368,7 @@ def _test_backend_correctness( num_gpu_blocks=8192, hf_config_override=hf_config_override, ) - device = torch.device("cuda:0") + device = torch.device(f"{DEVICE_TYPE}:0") kv_cache_spec = create_standard_kv_cache_spec(vllm_config) diff --git a/tests/v1/attention/test_chunked_local_attention.py b/tests/v1/attention/test_chunked_local_attention.py index 4529c2cfc29..c2798c8f2fa 100644 --- a/tests/v1/attention/test_chunked_local_attention.py +++ b/tests/v1/attention/test_chunked_local_attention.py @@ -7,6 +7,7 @@ import pytest import torch from tests.v1.attention.utils import BatchSpec, create_common_attn_metadata +from vllm.platforms import current_platform from vllm.v1.attention.backends.utils import make_local_attention_virtual_batches @@ -22,6 +23,8 @@ class LocalAttentionTestData: expected_local_block_table: list[list[int]] +DEVICE_TYPE = current_platform.device_type + test_data_list = [ # Same as example in docstring of make_local_attention_virtual_batches # except block table has 9 columns instead of 10 @@ -151,7 +154,7 @@ test_data_list = [ @pytest.mark.parametrize("test_data", test_data_list) def test_local_attention_virtual_batches(test_data: LocalAttentionTestData): - device = torch.device("cuda:0") + device = torch.device(f"{DEVICE_TYPE}:0") batch_spec = test_data.batch_spec attn_chunk_size = test_data.attn_chunk_size block_size = test_data.block_size diff --git a/tests/v1/attention/test_mla_backends.py b/tests/v1/attention/test_mla_backends.py index 796912a6806..e65d1d60402 100644 --- a/tests/v1/attention/test_mla_backends.py +++ b/tests/v1/attention/test_mla_backends.py @@ -42,6 +42,8 @@ BACKENDS_TO_TEST = [ AttentionBackendEnum.TRITON_MLA, ] +DEVICE_TYPE = current_platform.device_type + # Remove sm100 backends from the list if not using sm100 if not torch.cuda.is_available() or torch.cuda.get_device_properties(0).major < 10: BACKENDS_TO_TEST.remove(AttentionBackendEnum.CUTLASS_MLA) @@ -763,7 +765,7 @@ def test_backend_correctness( method="ngram", num_speculative_tokens=query_len - 1 ) - device = torch.device("cuda:0") + device = torch.device(f"{DEVICE_TYPE}:0") # 1. Setup batch_size = batch_spec.batch_size diff --git a/tests/v1/attention/test_sparse_mla_backends.py b/tests/v1/attention/test_sparse_mla_backends.py index c49ccd24e3a..22acc748d24 100644 --- a/tests/v1/attention/test_sparse_mla_backends.py +++ b/tests/v1/attention/test_sparse_mla_backends.py @@ -64,6 +64,8 @@ SPARSE_BACKEND_BATCH_SPECS["large_q_pure_prefill"] = BatchSpec( seq_lens=[256] * 2, query_lens=[256] * 2 ) +DEVICE_TYPE = current_platform.device_type + def _float_to_e8m0_truncate(f: float) -> float: """Simulate SM100's float -> e8m0 -> bf16 scale conversion. @@ -222,7 +224,7 @@ def test_sparse_backend_decode_correctness( batch_spec = SPARSE_BACKEND_BATCH_SPECS[batch_name] use_fp8_ds_mla_quantization = kv_cache_dtype == "fp8_ds_mla" - device = torch.device("cuda") + device = torch.device(DEVICE_TYPE) dtype = torch.bfloat16 # Model hyper-parameters (kept intentionally small for the unit test) @@ -586,7 +588,7 @@ def _triton_convert_reference_impl( def test_triton_convert_req_index_to_global_index_decode_only( block_size, num_topk_tokens ): - device = torch.device("cuda") + device = torch.device(DEVICE_TYPE) num_tokens = 8 num_requests = 4 max_blocks_per_req = 10 @@ -639,7 +641,7 @@ def test_triton_convert_req_index_to_global_index_decode_only( reason="FlashMLASparseBackend requires CUDA 9.0 or higher", ) def test_triton_convert_req_index_to_global_index_with_prefill_workspace(block_size): - device = torch.device("cuda") + device = torch.device(DEVICE_TYPE) num_requests = 4 max_blocks_per_req = 8 num_topk_tokens = 128 @@ -794,7 +796,7 @@ def test_split_indexer_prefill_chunks_single_request_overflow(): def test_triton_convert_returns_valid_counts(): """Test that return_valid_counts correctly counts non-negative indices.""" - device = torch.device("cuda") + device = torch.device(DEVICE_TYPE) num_tokens = 8 num_requests = 2 max_blocks_per_req = 10 diff --git a/tests/v1/attention/test_trtllm_attention_integration.py b/tests/v1/attention/test_trtllm_attention_integration.py index 113442bf6e4..12af0773cb3 100644 --- a/tests/v1/attention/test_trtllm_attention_integration.py +++ b/tests/v1/attention/test_trtllm_attention_integration.py @@ -55,6 +55,7 @@ class MockAttentionLayer: MODEL = "Qwen/Qwen2.5-0.5B" BLOCK_SIZE = 16 NUM_GPU_BLOCKS = 8192 +DEVICE_TYPE = current_platform.device_type BATCH_SPECS = { "decode_only": BatchSpec( @@ -172,7 +173,7 @@ def _run_trtllm_integration(batch_spec): """Run TRTLLM attention through the full FlashInfer pipeline and compare against an SDPA reference.""" set_random_seed(42) - device = torch.device("cuda:0") + device = torch.device(f"{DEVICE_TYPE}:0") vllm_config = create_vllm_config( model_name=MODEL, diff --git a/tests/v1/cudagraph/test_cudagraph_dispatch.py b/tests/v1/cudagraph/test_cudagraph_dispatch.py index 52e927cee8e..66e6d7dd460 100644 --- a/tests/v1/cudagraph/test_cudagraph_dispatch.py +++ b/tests/v1/cudagraph/test_cudagraph_dispatch.py @@ -23,6 +23,8 @@ from vllm.forward_context import BatchDescriptor, set_forward_context from vllm.platforms import current_platform from vllm.v1.cudagraph_dispatcher import CudagraphDispatcher +DEVICE_TYPE = current_platform.device_type + # Helper MLP for testing class SimpleMLP(nn.Module): @@ -269,9 +271,9 @@ class TestCudagraphDispatcher: class TestCUDAGraphWrapper: def setup_method(self): self.vllm_config = _create_vllm_config(CompilationConfig()) - self.model = SimpleMLP().to("cuda") - self.persistent_input_buffer = torch.zeros(1, 10, device="cuda") - self.input_tensor = torch.randn(1, 10, device="cuda") + self.model = SimpleMLP().to(DEVICE_TYPE) + self.persistent_input_buffer = torch.zeros(1, 10, device=DEVICE_TYPE) + self.input_tensor = torch.randn(1, 10, device=DEVICE_TYPE) def test_capture_and_replay(self): wrapper = CUDAGraphWrapper( @@ -428,10 +430,10 @@ class TestCudagraphIntegration: @create_new_process_for_each_test("spawn") def test_capture_replay_bypass_logic(self): - model = SimpleMLP().to("cuda") + model = SimpleMLP().to(DEVICE_TYPE) full_wrapper = CUDAGraphWrapper(model, self.vllm_config, CUDAGraphMode.FULL) max_bs = 16 - persistent_input_buffer = torch.zeros(max_bs, 10, device="cuda") + persistent_input_buffer = torch.zeros(max_bs, 10, device=DEVICE_TYPE) input_1 = persistent_input_buffer[:1] input_2 = persistent_input_buffer[:2] input_3 = persistent_input_buffer[:3] @@ -486,17 +488,17 @@ class TestCudagraphIntegration: @create_new_process_for_each_test("spawn") def test_nested_wrappers(self): """Tests a scenario with a PIECEWISE wrapper inside a FULL one.""" - model = SimpleMLP().to("cuda") + model = SimpleMLP().to(DEVICE_TYPE) full_wrapper = CUDAGraphWrapper(model, self.vllm_config, CUDAGraphMode.FULL) - input_1 = torch.randn(1, 10, device="cuda") + input_1 = torch.randn(1, 10, device=DEVICE_TYPE) # Setup: Inner model is wrapped with PIECEWISE, outer with FULL - inner_model = SimpleMLP().to("cuda") + inner_model = SimpleMLP().to(DEVICE_TYPE) piecewise_wrapper = CUDAGraphWrapper( inner_model, self.vllm_config, CUDAGraphMode.PIECEWISE ) inner_model.forward = MagicMock(wraps=inner_model.forward) - outer_model = SimpleMLP().to("cuda") + outer_model = SimpleMLP().to(DEVICE_TYPE) # When outer model is called, it calls the piecewise_wrapper outer_model.forward = MagicMock( wraps=outer_model.forward, side_effect=piecewise_wrapper diff --git a/tests/v1/determinism/test_batch_invariance.py b/tests/v1/determinism/test_batch_invariance.py index 6465985f0e9..b162b469bcd 100644 --- a/tests/v1/determinism/test_batch_invariance.py +++ b/tests/v1/determinism/test_batch_invariance.py @@ -187,7 +187,7 @@ def test_logprobs_bitwise_batch_invariance_bs1_vs_bsN( tensor_parallel_size=tp_size, max_num_seqs=128, max_model_len=8192, - dtype="bfloat16", # not everything is supported + dtype="auto", # not everything is supported gpu_memory_utilization=0.9, enforce_eager=IS_DEVICE_CAPABILITY_BELOW_90, attention_config={"backend": backend}, @@ -400,7 +400,7 @@ def test_simple_generation(backend): tensor_parallel_size=int(os.getenv("VLLM_TP_SIZE", "1")), gpu_memory_utilization=0.9, max_model_len=2048, - dtype="bfloat16", + dtype="auto", enable_prefix_caching=False, enforce_eager=IS_DEVICE_CAPABILITY_BELOW_90, attention_config={"backend": backend}, @@ -466,7 +466,7 @@ def test_logprobs_without_batch_invariance_should_fail( tensor_parallel_size=tp_size, max_num_seqs=32, max_model_len=8192, - dtype="bfloat16", + dtype="auto", enforce_eager=IS_DEVICE_CAPABILITY_BELOW_90, attention_config={"backend": backend}, ) @@ -686,7 +686,7 @@ def test_decode_logprobs_match_prefill_logprobs( tensor_parallel_size=tp_size, max_num_seqs=32, max_model_len=8192, - dtype="bfloat16", + dtype="auto", enforce_eager=IS_DEVICE_CAPABILITY_BELOW_90, attention_config={"backend": backend}, ) @@ -931,7 +931,7 @@ def LLM_with_max_seqs( max_num_seqs=max_num_seqs, gpu_memory_utilization=gpu_memory_utilization, max_model_len=max_model_len, - dtype="bfloat16", + dtype="auto", tensor_parallel_size=int(os.getenv("VLLM_TP_SIZE", "1")), enable_prefix_caching=False, enforce_eager=IS_DEVICE_CAPABILITY_BELOW_90, diff --git a/tests/v1/determinism/test_rms_norm_batch_invariant.py b/tests/v1/determinism/test_rms_norm_batch_invariant.py index 5e5b40d09c2..7d3b8437a93 100644 --- a/tests/v1/determinism/test_rms_norm_batch_invariant.py +++ b/tests/v1/determinism/test_rms_norm_batch_invariant.py @@ -13,6 +13,9 @@ from utils import skip_unsupported from vllm.model_executor.layers.batch_invariant import rms_norm as triton_rms_norm from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.platforms import current_platform + +DEVICE_TYPE = current_platform.device_type @skip_unsupported @@ -34,7 +37,7 @@ def test_rms_norm_batch_invariant_vs_standard( equivalent results to the standard CUDA implementation across various configurations. """ - device = torch.device("cuda") + device = torch.device(DEVICE_TYPE) # Create test input and weight torch.manual_seed(42) @@ -81,7 +84,7 @@ def test_rms_norm_3d_input( Ensures that the batch-invariant RMS norm correctly handles multi-dimensional inputs that are common in transformer models. """ - device = torch.device("cuda") + device = torch.device(DEVICE_TYPE) dtype = torch.bfloat16 eps = 1e-6 @@ -120,7 +123,7 @@ def test_rms_norm_numerical_stability(default_vllm_config): Ensures that both implementations handle edge cases like very small or large values without producing NaN or Inf. """ - device = torch.device("cuda") + device = torch.device(DEVICE_TYPE) dtype = torch.float16 eps = 1e-6 hidden_size = 2048 @@ -179,7 +182,7 @@ def test_rms_norm_formula(default_vllm_config): Verifies: output = input / sqrt(mean(input^2) + eps) * weight """ - device = torch.device("cuda") + device = torch.device(DEVICE_TYPE) dtype = torch.float32 # Use float32 for higher precision in formula check eps = 1e-6 hidden_size = 1024 @@ -214,7 +217,7 @@ def test_rms_norm_different_hidden_sizes(default_vllm_config, hidden_size: int): The Triton kernel uses a fixed BLOCK_SIZE=1024, so this tests that it correctly handles hidden sizes both smaller and larger than the block size. """ - device = torch.device("cuda") + device = torch.device(DEVICE_TYPE) dtype = torch.bfloat16 eps = 1e-6 batch_size = 16 @@ -251,7 +254,7 @@ def test_rms_norm_determinism(default_vllm_config): Runs the same input through the kernel multiple times and verifies identical outputs. """ - device = torch.device("cuda") + device = torch.device(DEVICE_TYPE) dtype = torch.bfloat16 eps = 1e-6 hidden_size = 4096 @@ -283,7 +286,7 @@ if __name__ == "__main__": # Run a quick smoke test print("Running quick smoke test of RMS norm implementations...") - device = torch.device("cuda") + device = torch.device(DEVICE_TYPE) batch_size = 8 hidden_size = 4096 dtype = torch.bfloat16 diff --git a/tests/v1/e2e/general/test_mamba_prefix_cache.py b/tests/v1/e2e/general/test_mamba_prefix_cache.py index 747c5defebd..8b9f7bb6c5a 100644 --- a/tests/v1/e2e/general/test_mamba_prefix_cache.py +++ b/tests/v1/e2e/general/test_mamba_prefix_cache.py @@ -16,6 +16,7 @@ from vllm import LLM, SamplingParams, TokensPrompt from vllm.config import CacheConfig from vllm.distributed import cleanup_dist_env_and_memory from vllm.model_executor.layers.mamba.mamba_utils import MambaStateCopyFunc +from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors from vllm.v1.attention.backends.utils import CommonAttentionMetadata from vllm.v1.core.kv_cache_manager import KVCacheBlocks, KVCacheManager @@ -48,6 +49,7 @@ num_accepted_tokens = 1 prompt_token_ids: list[int] = [] MODEL = "Qwen/Qwen3-Next-80B-A3B-Instruct-FP8" BLOCK_SIZE = 560 +DEVICE_TYPE = current_platform.device_type NUM_HIDDEN_LAYERS = 1 cur_step_action_idx = 0 cur_step_action: StepAction | None = None @@ -71,7 +73,7 @@ def get_fake_sample_fn() -> SamplerOutput: return SamplerOutput( sampled_token_ids=torch.tensor( [[prompt_token_ids[first_token_id_index]]], - device="cuda", + device=DEVICE_TYPE, dtype=torch.int32, ), logprobs_tensors=None, @@ -83,7 +85,9 @@ def get_fake_sample_fn() -> SamplerOutput: sampled_token_ids = accepted_tokens return SamplerOutput( sampled_token_ids=torch.tensor( - [sampled_token_ids], device="cuda", dtype=torch.int32 + [sampled_token_ids], + device=DEVICE_TYPE, + dtype=torch.int32, ), logprobs_tensors=None, ) @@ -128,17 +132,23 @@ def get_fake_propose_draft_token_ids_fn(): - 1 + num_accepted_tokens ], - device="cuda", + device=DEVICE_TYPE, dtype=torch.int32, ) valid_sampled_tokens_count = torch.tensor( - [num_accepted_tokens], device="cuda", dtype=torch.int32 + [num_accepted_tokens], + device=DEVICE_TYPE, + dtype=torch.int32, ) self._copy_valid_sampled_token_count(next_token_ids, valid_sampled_tokens_count) - return torch.tensor(proposed_draft_token_ids, device="cuda", dtype=torch.int32) + return torch.tensor( + proposed_draft_token_ids, + device=DEVICE_TYPE, + dtype=torch.int32, + ) return fake_propose_draft_token_ids_fn diff --git a/tests/v1/kv_connector/unit/test_nixl_connector.py b/tests/v1/kv_connector/unit/test_nixl_connector.py index 17a70a3a5f9..f39b78dd2b7 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector.py @@ -1613,7 +1613,7 @@ def test_register_kv_caches( ) ] ], - cache_dtype=torch.bfloat16, + cache_dtype="bfloat16", device=torch.accelerator.current_device_index(), kernel_block_sizes=[block_size], ) diff --git a/tests/v1/kv_offload/test_cpu_gpu.py b/tests/v1/kv_offload/test_cpu_gpu.py index 1983cca22d8..2da3a5e56b1 100644 --- a/tests/v1/kv_offload/test_cpu_gpu.py +++ b/tests/v1/kv_offload/test_cpu_gpu.py @@ -6,6 +6,7 @@ import time import pytest import torch +from vllm.platforms import current_platform from vllm.utils.torch_utils import set_random_seed from vllm.v1.kv_offload.mediums import CPULoadStoreSpec, GPULoadStoreSpec from vllm.v1.kv_offload.spec import ( @@ -21,7 +22,8 @@ GPU_PAGE_SIZES = [512, 1024] BLOCK_SIZE_FACTORS = [1, 3] NUM_TENSORS = [4] SEEDS = [0] -CUDA_DEVICES = ["cuda:0"] +DEVICE_TYPE = current_platform.device_type +DEVICES = [f"{DEVICE_TYPE}:0"] NUM_MAPPINGS = [3] @@ -33,7 +35,7 @@ NUM_MAPPINGS = [3] @pytest.mark.parametrize("num_cpu_blocks", NUM_CPU_BLOCKS) @pytest.mark.parametrize("num_tensors", NUM_TENSORS) @pytest.mark.parametrize("seed", SEEDS) -@pytest.mark.parametrize("device", CUDA_DEVICES) +@pytest.mark.parametrize("device", DEVICES) @torch.inference_mode() def test_transfer( default_vllm_config, diff --git a/tests/v1/logits_processors/test_correctness.py b/tests/v1/logits_processors/test_correctness.py index bc2cc1720fb..bf29793710a 100644 --- a/tests/v1/logits_processors/test_correctness.py +++ b/tests/v1/logits_processors/test_correctness.py @@ -39,8 +39,9 @@ PIN_MEMORY_AVAILABLE = is_pin_memory_available() MAX_NUM_REQS = 256 VOCAB_SIZE = 1024 NUM_OUTPUT_TOKENS = 20 -CUDA_DEVICES = [ - f"{current_platform.device_type}:{i}" +DEVICE_TYPE = current_platform.device_type +DEVICES = [ + f"{DEVICE_TYPE}:{i}" for i in range(1 if current_platform.device_count() == 1 else 2) ] MAX_NUM_PROMPT_TOKENS = 64 @@ -801,7 +802,7 @@ def _assert_valid( @create_new_process_for_each_test() -@pytest.mark.parametrize("device", CUDA_DEVICES) +@pytest.mark.parametrize("device", DEVICES) @pytest.mark.parametrize("reqs_per_logitproc", [REQS_PER_LOGITPROC]) @pytest.mark.parametrize("logitsprocs_under_test", _get_test_cases()) def test_logitsprocs( diff --git a/tests/v1/sample/test_rejection_sampler.py b/tests/v1/sample/test_rejection_sampler.py index 552a27fe22d..ecfcade2b61 100644 --- a/tests/v1/sample/test_rejection_sampler.py +++ b/tests/v1/sample/test_rejection_sampler.py @@ -19,7 +19,7 @@ from vllm.v1.sample.rejection_sampler import ( from vllm.v1.sample.sampler import Sampler, SamplerOutput from vllm.v1.spec_decode.metadata import SpecDecodeMetadata -DEVICE = current_platform.device_type +DEVICE_TYPE = current_platform.device_type @pytest.fixture @@ -57,7 +57,7 @@ def create_logits_tensor( will produce desired token ids on argmax""" token_ids = [tokens[:-1] for tokens in output_token_ids] num_total_tokens = sum(len(tokens) for tokens in token_ids) - logits = torch.full((num_total_tokens, vocab_size), -100.0, device=DEVICE) + logits = torch.full((num_total_tokens, vocab_size), -100.0, device=DEVICE_TYPE) start_loc = 0 for tokens in token_ids: for j, token_id in enumerate(tokens): @@ -99,9 +99,9 @@ def create_sampling_metadata( assert output_token_ids assert len(output_token_ids) > 0 - frequency_penalties = torch.tensor(frequency_penalties, device=DEVICE) - presence_penalties = torch.tensor(presence_penalties, device=DEVICE) - repetition_penalties = torch.tensor(repetition_penalties, device=DEVICE) + frequency_penalties = torch.tensor(frequency_penalties, device=DEVICE_TYPE) + presence_penalties = torch.tensor(presence_penalties, device=DEVICE_TYPE) + repetition_penalties = torch.tensor(repetition_penalties, device=DEVICE_TYPE) else: no_penalties = True frequency_penalties = torch.tensor([]) @@ -320,14 +320,27 @@ def test_deterministic_when_seeded( n_rep: int, ): num_tokens = batch_size * k - draft_probs = torch.rand(num_tokens, vocab_size, dtype=torch.float32, device=DEVICE) + draft_probs = torch.rand( + num_tokens, + vocab_size, + dtype=torch.float32, + device=DEVICE_TYPE, + ) draft_probs = F.softmax(draft_probs, dim=-1) target_logits = torch.rand_like(draft_probs) bonus_token_ids = torch.randint( - low=0, high=vocab_size, size=(batch_size, 1), dtype=torch.int64, device=DEVICE + low=0, + high=vocab_size, + size=(batch_size, 1), + dtype=torch.int64, + device=DEVICE_TYPE, ) draft_token_ids = torch.randint( - low=0, high=vocab_size, size=(batch_size, k), dtype=torch.int64, device=DEVICE + low=0, + high=vocab_size, + size=(batch_size, k), + dtype=torch.int64, + device=DEVICE_TYPE, ) seeded_mask = torch.rand(batch_size, dtype=torch.float32) <= frac_seeded @@ -335,12 +348,12 @@ def test_deterministic_when_seeded( results = [] for _ in range(n_rep): seeded_seqs = { - i: torch.Generator(device=DEVICE).manual_seed(i) + i: torch.Generator(device=DEVICE_TYPE).manual_seed(i) for i in range(batch_size) if seeded_mask[i] } - temperature = torch.ones(batch_size, dtype=torch.float32, device=DEVICE) + temperature = torch.ones(batch_size, dtype=torch.float32, device=DEVICE_TYPE) sampling_metadata = create_sampling_metadata( all_greedy=False, temperature=temperature, generators=seeded_seqs ) @@ -387,7 +400,7 @@ def test_rejection_sampling_approximates_target_distribution(): much more than the distance improvement between the observed distribution and the random distribution. """ - torch.set_default_device(DEVICE) + torch.set_default_device(DEVICE_TYPE) vocab_size = 10 k = 2 num_reference_probs = 100 @@ -410,7 +423,7 @@ def test_rejection_sampling_approximates_target_distribution(): rej_sample_probs = estimate_rejection_sampling_pdf( draft_probs, target_logits, k, vocab_size, num_samples ) - rej_sample_probs = rej_sample_probs.to(DEVICE) + rej_sample_probs = rej_sample_probs.to(DEVICE_TYPE) # Average distance from reference probs. reference_vs_rejsample_dist = ( @@ -491,11 +504,11 @@ def estimate_rejection_sampling_pdf( draft_probs = draft_probs.view(num_tokens, vocab_size) # Bonus tokens not used but required. - bonus_token_ids = torch.zeros((1, 1), dtype=torch.int64, device=DEVICE).repeat( + bonus_token_ids = torch.zeros((1, 1), dtype=torch.int64, device=DEVICE_TYPE).repeat( num_samples, 1 ) - temperature = torch.ones(num_samples, dtype=torch.float32, device=DEVICE) + temperature = torch.ones(num_samples, dtype=torch.float32, device=DEVICE_TYPE) sampling_metadata = create_sampling_metadata( all_greedy=False, temperature=temperature ) @@ -600,7 +613,7 @@ def _test_masked_logits( # Create random draft probabilities. draft_probs = torch.rand( - (num_tokens, vocab_size), dtype=torch.float32, device=DEVICE + (num_tokens, vocab_size), dtype=torch.float32, device=DEVICE_TYPE ) draft_probs = F.softmax(draft_probs, dim=-1) @@ -610,7 +623,11 @@ def _test_masked_logits( draft_token_ids = draft_token_ids.tolist() # Bonus tokens not used but required - bonus_token_ids = torch.zeros((batch_size, 1), dtype=torch.int64, device=DEVICE) + bonus_token_ids = torch.zeros( + (batch_size, 1), + dtype=torch.int64, + device=DEVICE_TYPE, + ) # Create spec decode metadata spec_decode_metadata = create_spec_decode_metadata(draft_token_ids, target_logits) @@ -645,12 +662,13 @@ def test_top_k(rejection_sampler, top_k): # Randomly create top-k indices. top_k_indices = [ - torch.randperm(vocab_size, device=DEVICE)[:top_k] for _ in range(num_tokens) + torch.randperm(vocab_size, device=DEVICE_TYPE)[:top_k] + for _ in range(num_tokens) ] top_k_indices = torch.stack(top_k_indices) # Create logits with the uniform distribution. - target_logits = torch.zeros((num_tokens, vocab_size), device=DEVICE) + target_logits = torch.zeros((num_tokens, vocab_size), device=DEVICE_TYPE) # Increment the logits for top-k indices, a little bit more than the other # ones. If the masking is effective, the non-topk indices will never be @@ -659,11 +677,11 @@ def test_top_k(rejection_sampler, top_k): target_logits[i, top_k_indices[i]] += 0.1 # Create sampling metadata - temperature = torch.ones(batch_size, dtype=torch.float32, device=DEVICE) + temperature = torch.ones(batch_size, dtype=torch.float32, device=DEVICE_TYPE) sampling_metadata = create_sampling_metadata( all_greedy=False, temperature=temperature, - top_k=torch.tensor([top_k] * batch_size, device=DEVICE, dtype=torch.int64), + top_k=torch.tensor([top_k] * batch_size, device=DEVICE_TYPE, dtype=torch.int64), ) _test_masked_logits( @@ -686,8 +704,8 @@ def test_top_p(rejection_sampler, top_p): num_tokens = batch_size * num_draft_tokens # Create logits with the uniform distribution. - target_logits = torch.randn((num_tokens, vocab_size), device=DEVICE) - temperature = torch.ones(batch_size, dtype=torch.float32, device=DEVICE) + target_logits = torch.randn((num_tokens, vocab_size), device=DEVICE_TYPE) + temperature = torch.ones(batch_size, dtype=torch.float32, device=DEVICE_TYPE) rescaled_logits = target_logits / temperature logits_sort, logits_idx = rescaled_logits.sort(dim=-1, descending=False) @@ -706,7 +724,11 @@ def test_top_p(rejection_sampler, top_p): sampling_metadata = create_sampling_metadata( all_greedy=False, temperature=temperature, - top_p=torch.tensor([top_p] * batch_size, device=DEVICE, dtype=torch.float32), + top_p=torch.tensor( + [top_p] * batch_size, + device=DEVICE_TYPE, + dtype=torch.float32, + ), ) _test_masked_logits( @@ -732,7 +754,10 @@ def test_frequency_penalties(rejection_sampler): all_greedy=True, output_token_ids=[[2], [3], [4]], spec_token_ids=spec_tokens, - prompt_token_ids=torch.tensor([[5, 6, 7], [6, 7, 8], [7, 8, 9]], device=DEVICE), + prompt_token_ids=torch.tensor( + [[5, 6, 7], [6, 7, 8], [7, 8, 9]], + device=DEVICE_TYPE, + ), frequency_penalties=[1.5, 1.5, 0.7], presence_penalties=[0.0] * num_requests, repetition_penalties=[1.0] * num_requests, @@ -858,21 +883,26 @@ def test_sample_recovered_tokens( num_tokens = batch_size * max_spec_len # Create random draft probabilities. - draft_probs = torch.rand(num_tokens, vocab_size, dtype=torch.float32, device=DEVICE) + draft_probs = torch.rand( + num_tokens, + vocab_size, + dtype=torch.float32, + device=DEVICE_TYPE, + ) draft_probs = F.softmax(draft_probs, dim=-1) # Create random target probabilities. target_logits = torch.rand( - num_tokens, vocab_size, dtype=torch.float32, device=DEVICE + num_tokens, vocab_size, dtype=torch.float32, device=DEVICE_TYPE ) target_probs = F.softmax(target_logits, dim=-1) # Randomly sample draft token ids from draft probs draft_token_ids = torch.multinomial(draft_probs, num_samples=1).to(torch.int32) - temperature = torch.ones(batch_size, dtype=torch.float32, device=DEVICE) + temperature = torch.ones(batch_size, dtype=torch.float32, device=DEVICE_TYPE) generators = { - i: torch.Generator(device=DEVICE).manual_seed(i) for i in range(batch_size) + i: torch.Generator(device=DEVICE_TYPE).manual_seed(i) for i in range(batch_size) } sampling_metadata = create_sampling_metadata( all_greedy=False, temperature=temperature, generators=generators @@ -890,7 +920,7 @@ def test_sample_recovered_tokens( None if no_draft_probs else draft_probs, target_probs, sampling_metadata, - device=DEVICE, + device=DEVICE_TYPE, ) recovered_token_ids = sample_recovered_tokens( max_spec_len, @@ -900,6 +930,6 @@ def test_sample_recovered_tokens( None if no_draft_probs else draft_probs, target_probs, sampling_metadata, - device=DEVICE, + device=DEVICE_TYPE, ) assert torch.equal(recovered_token_ids, ref_recovered_token_ids) diff --git a/tests/v1/sample/test_sampler.py b/tests/v1/sample/test_sampler.py index 51f2bf5e753..c67199fa407 100644 --- a/tests/v1/sample/test_sampler.py +++ b/tests/v1/sample/test_sampler.py @@ -17,8 +17,9 @@ PIN_MEMORY_AVAILABLE = is_pin_memory_available() MAX_NUM_REQS = 256 VOCAB_SIZE = 1024 NUM_OUTPUT_TOKENS = 20 -CUDA_DEVICES = [ - f"{current_platform.device_type}:{i}" +DEVICE_TYPE = current_platform.device_type +DEVICES = [ + f"{DEVICE_TYPE}:{i}" for i in range(1 if current_platform.device_count() == 1 else 2) ] MAX_NUM_PROMPT_TOKENS = 64 @@ -199,7 +200,7 @@ def _create_weighted_output_token_list( return output_token_ids, sorted_token_ids_in_output -@pytest.mark.parametrize("device", CUDA_DEVICES) +@pytest.mark.parametrize("device", DEVICES) @pytest.mark.parametrize("batch_size", [1, 2, 32]) @pytest.mark.parametrize("presence_penalty", [-2.0, 2.0]) def test_sampler_presence_penalty( @@ -249,7 +250,7 @@ def test_sampler_presence_penalty( assert penalized_token_id not in output_token_ids[batch_idx] -@pytest.mark.parametrize("device", CUDA_DEVICES) +@pytest.mark.parametrize("device", DEVICES) @pytest.mark.parametrize("batch_size", [1, 2, 32]) @pytest.mark.parametrize("frequency_penalty", [-2.0, 2.0]) def test_sampler_frequency_penalty( @@ -305,7 +306,7 @@ def test_sampler_frequency_penalty( assert penalized_token_id not in distinct_sorted_token_ids_in_output -@pytest.mark.parametrize("device", CUDA_DEVICES) +@pytest.mark.parametrize("device", DEVICES) @pytest.mark.parametrize("batch_size", [1, 2, 32]) @pytest.mark.parametrize("repetition_penalty", [0.1, 1.9]) def test_sampler_repetition_penalty( @@ -363,7 +364,7 @@ def test_sampler_repetition_penalty( ) -@pytest.mark.parametrize("device", CUDA_DEVICES) +@pytest.mark.parametrize("device", DEVICES) @pytest.mark.parametrize("batch_size", [1, 2, 32]) @pytest.mark.parametrize("num_allowed_token_ids", [0, 1, 2]) def test_sampler_allowed_token_ids( @@ -409,7 +410,7 @@ def test_sampler_allowed_token_ids( assert logits_for_req[token_id] != -float("inf") -@pytest.mark.parametrize("device", CUDA_DEVICES) +@pytest.mark.parametrize("device", DEVICES) @pytest.mark.parametrize("batch_size", [1, 2, 32]) @pytest.mark.parametrize("bad_words_lengths", [(1,), (1, 3), (2, 2)]) def test_sampler_bad_words( diff --git a/tests/v1/sample/test_topk_topp_sampler.py b/tests/v1/sample/test_topk_topp_sampler.py index ce1e288a241..511f2668075 100644 --- a/tests/v1/sample/test_topk_topp_sampler.py +++ b/tests/v1/sample/test_topk_topp_sampler.py @@ -7,8 +7,7 @@ from torch import Generator from vllm.platforms import current_platform from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p_pytorch -CUDA_DEVICE = "cuda" if current_platform.is_cuda() else None -DEVICE = current_platform.device_type +DEVICE_TYPE = current_platform.device_type BATCH_SIZE = 1024 VOCAB_SIZE = 128 * 1024 @@ -26,8 +25,8 @@ def reset_default_device(): def test_topk_impl_equivalence(): - torch.set_default_device(DEVICE) - generator = Generator(device=DEVICE).manual_seed(33) + torch.set_default_device(DEVICE_TYPE) + generator = Generator(device=DEVICE_TYPE).manual_seed(33) logits = torch.rand((BATCH_SIZE, VOCAB_SIZE), generator=generator) @@ -76,8 +75,8 @@ def test_flashinfer_sampler(): if not FLASHINFER_ENABLED: pytest.skip("FlashInfer not installed or not available on this platform.") - torch.set_default_device(DEVICE) - generator = Generator(device=DEVICE).manual_seed(42) + torch.set_default_device(DEVICE_TYPE) + generator = Generator(device=DEVICE_TYPE).manual_seed(42) # Generate random logits logits = torch.rand((BATCH_SIZE, VOCAB_SIZE), generator=generator) @@ -128,15 +127,15 @@ def test_flashinfer_sampler(): # ============================================================================= -@pytest.mark.skipif(CUDA_DEVICE is None, reason="CUDA not available") +@pytest.mark.skipif("CPU" in DEVICE_TYPE, reason="CUDA/XPU not available") class TestTritonTopkTopp: """Tests for the Triton top-k/top-p kernel.""" @pytest.fixture(autouse=True) def setup(self): """Set up test fixtures.""" - torch.set_default_device(CUDA_DEVICE) - self.generator = Generator(device=CUDA_DEVICE).manual_seed(42) + torch.set_default_device(DEVICE_TYPE) + self.generator = Generator(device=DEVICE_TYPE).manual_seed(42) def _compare_results( self, diff --git a/tests/v1/spec_decode/test_eagle.py b/tests/v1/spec_decode/test_eagle.py index c1082202448..5d587fa3ec1 100644 --- a/tests/v1/spec_decode/test_eagle.py +++ b/tests/v1/spec_decode/test_eagle.py @@ -42,6 +42,7 @@ dflash_target_dir = "Qwen/Qwen3-8B" dflash_dir = "z-lab/Qwen3-8B-DFlash-b16" BLOCK_SIZE = 16 +DEVICE_TYPE = current_platform.device_type def _create_proposer( @@ -92,7 +93,7 @@ def _create_proposer( # Overwrite pard_token to avoid crash during init speculative_config.draft_model_config.hf_config.pard_token = 0 - device = current_platform.device_type + device = DEVICE_TYPE vllm_config = VllmConfig( model_config=model_config, cache_config=CacheConfig(block_size=16), @@ -124,7 +125,7 @@ def test_prepare_next_token_ids(): either the GPU tensor of sampled_token_ids with -1 for rejected tokens, or the CPU python list[list[int]] with the rejected tokens removed. """ - device = torch.device(current_platform.device_type) + device = torch.device(DEVICE_TYPE) num_requests = 4 num_speculative_tokens = 4 @@ -207,7 +208,7 @@ def test_prepare_inputs(): a, a + 1, ..., a + b - n2 - 1, a + b, a + b + 1, ..., a + b + c - n3 - 1] """ - device = torch.device(current_platform.device_type) + device = torch.device(DEVICE_TYPE) # q1 = 4, q2 = 7, q3 = 5 # n1 = 1, n2 = 3, n3 = 2 @@ -300,7 +301,7 @@ def test_prepare_inputs_padded(): from the original indices to sample from. """ - device = torch.device(current_platform.device_type) + device = torch.device(DEVICE_TYPE) expected_token_indices_to_sample = torch.tensor( [1, 5, 6], dtype=torch.int32, device=device @@ -370,7 +371,7 @@ def test_set_inputs_first_pass_default_eagle(): - After inserting next_tokens [100, 200, 300]: [a2, a3, 100, b2, 200, c2, c3, c4, 300] """ - device = torch.device(current_platform.device_type) + device = torch.device(DEVICE_TYPE) num_speculative_tokens = 3 proposer = _create_proposer("eagle", num_speculative_tokens) @@ -471,7 +472,7 @@ def test_set_inputs_first_pass_draft_model(): - idx 5: token 21, pos 1 - idx 6: token 200, pos 2 (bonus token) """ - device = torch.device(current_platform.device_type) + device = torch.device(DEVICE_TYPE) num_speculative_tokens = 2 block_size = BLOCK_SIZE @@ -609,7 +610,7 @@ def test_set_inputs_first_pass_parallel_drafting(): - idx 9: bonus token 200 - idx 10-11: parallel_drafting_tokens, is_masked=True """ - device = torch.device(current_platform.device_type) + device = torch.device(DEVICE_TYPE) num_speculative_tokens = 3 block_size = BLOCK_SIZE @@ -859,7 +860,7 @@ def test_propose(method, attn_backend, num_speculative_tokens, monkeypatch): monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1") # Use GPU device - device = torch.device(current_platform.device_type) + device = torch.device(DEVICE_TYPE) # Setup test parameters batch_size = 2 @@ -1030,7 +1031,7 @@ def test_propose(method, attn_backend, num_speculative_tokens, monkeypatch): ) def test_propose_tree(spec_token_tree): # Get GPU device. - device = torch.device(current_platform.device_type) + device = torch.device(DEVICE_TYPE) # Setup test parameters. batch_size = 2 diff --git a/tests/v1/spec_decode/test_eagle_step_kernel.py b/tests/v1/spec_decode/test_eagle_step_kernel.py index 319ab4a33ad..275a157d1be 100644 --- a/tests/v1/spec_decode/test_eagle_step_kernel.py +++ b/tests/v1/spec_decode/test_eagle_step_kernel.py @@ -5,11 +5,14 @@ import pytest import torch +from vllm.platforms import current_platform from vllm.v1.spec_decode.utils import ( PADDING_SLOT_ID, eagle_step_update_slot_mapping_and_metadata, ) +DEVICE_TYPE = current_platform.device_type + # Skip if no CUDA - Triton kernel requires GPU pytest.importorskip("triton") if not torch.cuda.is_available(): @@ -47,7 +50,7 @@ def _reference_eagle_step_slot_mapping( def test_eagle_step_slot_mapping_kernel(): """Test fused kernel matches Python reference for slot mapping and metadata.""" - device = torch.device("cuda") + device = torch.device(DEVICE_TYPE) batch_size = 32 block_size = 16 max_model_len = 4096 @@ -93,7 +96,7 @@ def test_eagle_step_slot_mapping_kernel(): def test_eagle_step_slot_mapping_kernel_exceeds_max(): """Test fused kernel when position exceeds max_model_len.""" - device = torch.device("cuda") + device = torch.device(DEVICE_TYPE) batch_size = 4 block_size = 16 max_model_len = 100 @@ -130,7 +133,7 @@ def test_eagle_step_slot_mapping_kernel_exceeds_max(): def test_eagle_step_slot_mapping_kernel_cudagraph_padding(): """Test that padding threads write PADDING_SLOT_ID when input_batch_size > batch_size (cudagraph padding).""" - device = torch.device("cuda") + device = torch.device(DEVICE_TYPE) batch_size = 4 input_batch_size = 8 block_size = 16 diff --git a/tests/v1/spec_decode/test_extract_hidden_states.py b/tests/v1/spec_decode/test_extract_hidden_states.py index 9f9758b829b..95004eb65da 100644 --- a/tests/v1/spec_decode/test_extract_hidden_states.py +++ b/tests/v1/spec_decode/test_extract_hidden_states.py @@ -27,6 +27,7 @@ from vllm.v1.spec_decode.extract_hidden_states import ExtractHiddenStatesPropose from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch model_dir = "TinyLlama/TinyLlama-1.1B-Chat-v1.0" +DEVICE_TYPE = current_platform.device_type def _create_proposer( @@ -51,7 +52,7 @@ def _create_proposer( }, ) - device = current_platform.device_type + device = DEVICE_TYPE vllm_config = VllmConfig( model_config=model_config, cache_config=CacheConfig(), @@ -101,7 +102,7 @@ def test_proposer_initialization_missing_layer_ids(): }, ) - device = current_platform.device_type + device = DEVICE_TYPE vllm_config = VllmConfig( model_config=model_config, cache_config=CacheConfig(), @@ -130,7 +131,7 @@ def test_prepare_next_token_ids_padded(): For each request we either use the sampled token (if valid and not discarded) or a backup token from the request state. """ - device = torch.device(current_platform.device_type) + device = torch.device(DEVICE_TYPE) num_requests = 4 req_ids = [f"req_{i + 1}" for i in range(num_requests)] @@ -197,7 +198,7 @@ def test_propose(): 2. Return the sampled tokens as "draft" tokens (shape [batch_size, 1]) 3. Cache the hidden states in the model's KV cache """ - device = torch.device(current_platform.device_type) + device = torch.device(DEVICE_TYPE) # Setup test parameters batch_size = 2 @@ -273,7 +274,7 @@ def test_propose(): @pytest.mark.parametrize("num_hidden_layers", [1, 4, 8]) def test_propose_different_layer_counts(num_hidden_layers): """Test that propose works correctly with different numbers of hidden layers.""" - device = torch.device(current_platform.device_type) + device = torch.device(DEVICE_TYPE) batch_size = 2 num_tokens = 5 diff --git a/tests/v1/spec_decode/test_mtp.py b/tests/v1/spec_decode/test_mtp.py index 0a48b0e7b98..094611e05c1 100644 --- a/tests/v1/spec_decode/test_mtp.py +++ b/tests/v1/spec_decode/test_mtp.py @@ -28,6 +28,7 @@ from vllm.v1.attention.backends.registry import AttentionBackendEnum from vllm.v1.spec_decode.eagle import EagleProposer mimo_7b_dir = "XiaomiMiMo/MiMo-7B-Base" +DEVICE_TYPE = current_platform.device_type def _create_mtp_proposer(num_speculative_tokens: int) -> EagleProposer: @@ -48,7 +49,7 @@ def _create_mtp_proposer(num_speculative_tokens: int) -> EagleProposer: model_config=model_config, cache_config=CacheConfig(), speculative_config=speculative_config, - device_config=DeviceConfig(device=current_platform.device_type), + device_config=DeviceConfig(device=DEVICE_TYPE), parallel_config=ParallelConfig(), load_config=LoadConfig(), scheduler_config=SchedulerConfig( @@ -57,7 +58,7 @@ def _create_mtp_proposer(num_speculative_tokens: int) -> EagleProposer: ), ) - return EagleProposer(vllm_config=vllm_config, device=current_platform.device_type) + return EagleProposer(vllm_config=vllm_config, device=DEVICE_TYPE) @mock.patch("vllm.v1.spec_decode.eagle.get_pp_group") @@ -118,7 +119,7 @@ def test_mtp_load_model_unified(mock_get_model, mock_get_layers, mock_get_pp_gro def test_mtp_propose(num_speculative_tokens, monkeypatch): """Test that MTP's forward method returns hidden states directly""" - device = torch.device(current_platform.device_type) + device = torch.device(DEVICE_TYPE) batch_size = 2 seq_lens = [5, 3] total_tokens = sum(seq_lens) diff --git a/tests/v1/spec_decode/test_tree_attention.py b/tests/v1/spec_decode/test_tree_attention.py index 52bc722cfcb..cb487acec0a 100644 --- a/tests/v1/spec_decode/test_tree_attention.py +++ b/tests/v1/spec_decode/test_tree_attention.py @@ -18,6 +18,8 @@ from vllm.v1.attention.backend import CommonAttentionMetadata from vllm.v1.attention.backends.fa_utils import is_flash_attn_varlen_func_available from vllm.v1.attention.backends.registry import AttentionBackendEnum +DEVICE_TYPE = current_platform.device_type + if not is_flash_attn_varlen_func_available(): pytest.skip( "This test requires flash_attn_varlen_func, but it's not available.", @@ -170,9 +172,9 @@ def _get_available_reference_backends() -> list[AttentionBackendEnum]: class MockAttentionLayer(torch.nn.Module): - _q_scale = torch.tensor(1.0, dtype=torch.float32, device="cuda") - _k_scale = torch.tensor(1.0, dtype=torch.float32, device="cuda") - _v_scale = torch.tensor(1.0, dtype=torch.float32, device="cuda") + _q_scale = torch.tensor(1.0, dtype=torch.float32, device=DEVICE_TYPE) + _k_scale = torch.tensor(1.0, dtype=torch.float32, device=DEVICE_TYPE) + _v_scale = torch.tensor(1.0, dtype=torch.float32, device=DEVICE_TYPE) layer_name = "mock_layer" def __init__(self): diff --git a/tests/v1/worker/test_gpu_input_batch.py b/tests/v1/worker/test_gpu_input_batch.py index d4eee19adab..3a478d21013 100644 --- a/tests/v1/worker/test_gpu_input_batch.py +++ b/tests/v1/worker/test_gpu_input_batch.py @@ -22,10 +22,8 @@ from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch VOCAB_SIZE = 1024 NUM_OUTPUT_TOKENS = 20 MAX_PROMPT_SIZE = 100 -CUDA_DEVICES = [ - f"{current_platform.device_type}:{i}" - for i in range(min(current_platform.device_count(), 2)) -] +DEVICE_TYPE = current_platform.device_type +DEVICES = [f"{DEVICE_TYPE}:{i}" for i in range(min(current_platform.device_count(), 2))] MAX_NUM_PROMPT_TOKENS = 64 @@ -219,7 +217,7 @@ def _construct_cached_request_state(req_id_suffix: int): ) -@pytest.mark.parametrize("device", CUDA_DEVICES) +@pytest.mark.parametrize("device", DEVICES) @pytest.mark.parametrize("batch_size", [1, 2, 32, 64]) def test_sampling_metadata_in_input_batch(device: str, batch_size: int): """ @@ -313,7 +311,7 @@ def test_sampling_metadata_in_input_batch(device: str, batch_size: int): ) -@pytest.mark.parametrize("device", CUDA_DEVICES) +@pytest.mark.parametrize("device", DEVICES) @pytest.mark.parametrize("batch_size", [32]) @pytest.mark.parametrize("swap_list", [((0, 1),)]) def test_swap_states_in_input_batch(device: str, batch_size: int, swap_list: list): @@ -400,7 +398,7 @@ def _construct_pooling_request(req_id_suffix: int, pooling_params=None): ) -@pytest.mark.parametrize("device", CUDA_DEVICES) +@pytest.mark.parametrize("device", DEVICES) def test_pooling_prompt_lens_not_aliased(device: str): """Verify that prompt_lens in PoolingMetadata does not share memory with the internal num_prompt_tokens pinned buffer. Guards against possible diff --git a/tests/v1/worker/test_gpu_model_runner.py b/tests/v1/worker/test_gpu_model_runner.py index d7695027a28..0de443858c9 100644 --- a/tests/v1/worker/test_gpu_model_runner.py +++ b/tests/v1/worker/test_gpu_model_runner.py @@ -45,7 +45,7 @@ from vllm.v1.worker.utils import AttentionGroup, select_common_block_size BLOCK_SIZE = 16 NUM_BLOCKS = 10 -DEVICE = current_platform.device_type +DEVICE_TYPE = current_platform.device_type def initialize_kv_cache(runner: GPUModelRunner): @@ -121,7 +121,7 @@ def model_runner(): vllm_config.compilation_config.static_forward_context["layer.0"] = Attention( num_heads, head_size, 0.1 ) - runner = GPUModelRunner(vllm_config, DEVICE) + runner = GPUModelRunner(vllm_config, DEVICE_TYPE) initialize_kv_cache(runner) yield runner @@ -340,7 +340,7 @@ def test_get_nans_in_logits(model_runner, dist_init): [1.0, 2.0, 3.0], [3.0, 2.0, 1.0], ], - device=DEVICE, + device=DEVICE_TYPE, ) result = model_runner._get_nans_in_logits(logits) assert result == {"req_0": 0, "req_1": 0} @@ -350,7 +350,7 @@ def test_get_nans_in_logits(model_runner, dist_init): [1.0, float("nan"), 3.0], [4.0, float("nan"), float("nan")], ], - device=DEVICE, + device=DEVICE_TYPE, ) result = model_runner._get_nans_in_logits(logits) assert result == {"req_0": 1, "req_1": 2} @@ -360,7 +360,7 @@ def test_get_nans_in_logits(model_runner, dist_init): [1.0, 2.0, 3.0], [4.0, float("nan"), float("nan")], ], - device=DEVICE, + device=DEVICE_TYPE, ) result = model_runner._get_nans_in_logits(logits) assert result == {"req_0": 0, "req_1": 2} @@ -372,7 +372,7 @@ def test_get_nans_in_logits(model_runner, dist_init): [ [1.0, float("nan"), 3.0], ], - device=DEVICE, + device=DEVICE_TYPE, ) result = model_runner._get_nans_in_logits(logits) assert result == {"req_0": 1, "req_1": 0} @@ -383,7 +383,7 @@ def test_get_nans_in_logits(model_runner, dist_init): [1.0, 2.0, 3.0], [float("nan"), 2.0, 3.0], ], - device=DEVICE, + device=DEVICE_TYPE, ) result = model_runner._get_nans_in_logits(logits) assert result == {"req_0": 2, "req_1": 0} @@ -643,7 +643,7 @@ def test_init_kv_cache_without_kv_sharing(default_vllm_config): # Set high context length to test max context length estimation vllm_config.model_config.max_model_len = 3_000_000 vllm_ctx = vllm_config.compilation_config.static_forward_context - runner = GPUModelRunner(vllm_config, DEVICE) + runner = GPUModelRunner(vllm_config, DEVICE_TYPE) kv_cache_spec = runner.get_kv_cache_spec() assert len(kv_cache_spec) == 2 assert len(runner.shared_kv_cache_layers) == 0 @@ -711,7 +711,7 @@ def test_init_kv_cache_with_kv_sharing_valid(default_vllm_config): # Set high context length to test max context length estimation vllm_config.model_config.max_model_len = 3_000_000 vllm_ctx = vllm_config.compilation_config.static_forward_context - runner = GPUModelRunner(vllm_config, DEVICE) + runner = GPUModelRunner(vllm_config, DEVICE_TYPE) kv_cache_spec = runner.get_kv_cache_spec() assert len(kv_cache_spec) == 1 assert layer_0 in kv_cache_spec @@ -850,7 +850,7 @@ def test_hybrid_attention_mamba_tensor_shapes(): assert fwd_context is not None vllm_ctx = vllm_config.compilation_config.static_forward_context - runner = GPUModelRunner(vllm_config, DEVICE) + runner = GPUModelRunner(vllm_config, DEVICE_TYPE) current_platform.update_block_size_for_backend(vllm_config) kv_cache_spec = runner.get_kv_cache_spec() @@ -896,13 +896,13 @@ def test_hybrid_attention_mamba_tensor_shapes(): ssm_constant_shape = ssm_shape[1:] attn_blocks_constant = torch.full( - (test_block_size, *attn_constant_shape), device=DEVICE, fill_value=3.33 + (test_block_size, *attn_constant_shape), device=DEVICE_TYPE, fill_value=3.33 ) conv_blocks_constant = torch.full( - (test_block_size, *conv_constant_shape), device=DEVICE, fill_value=6.66 + (test_block_size, *conv_constant_shape), device=DEVICE_TYPE, fill_value=6.66 ) ssm_blocks_constant = torch.full( - (test_block_size, *ssm_constant_shape), device=DEVICE, fill_value=9.99 + (test_block_size, *ssm_constant_shape), device=DEVICE_TYPE, fill_value=9.99 ) # Fill attention blocks with constants using kv block indices @@ -997,7 +997,7 @@ def test_hybrid_block_table_initialization(): max_num_blocks_per_req=max_num_blocks_per_req, max_num_batched_tokens=max_num_batched_tokens, pin_memory=False, - device=torch.device(DEVICE), + device=torch.device(DEVICE_TYPE), kernel_block_size=kernel_block_sizes[0], cp_kv_cache_interleave_size=cp_kv_cache_interleave_size, ) @@ -1036,7 +1036,7 @@ def test_input_batch_with_kernel_block_sizes(): max_num_reqs = 10 max_model_len = 512 max_num_batched_tokens = 512 - device = torch.device(DEVICE) + device = torch.device(DEVICE_TYPE) pin_memory = False vocab_size = 50272 @@ -1083,7 +1083,7 @@ def test_hybrid_cache_integration(default_vllm_config, dist_init): num_heads, head_size, 0.1 ) - runner = GPUModelRunner(vllm_config, DEVICE) + runner = GPUModelRunner(vllm_config, DEVICE_TYPE) # Initialize KV cache with configuration attn_spec = FullAttentionSpec( @@ -1306,7 +1306,7 @@ def test_mamba_cache_raises_when_max_num_seqs_exceeds_blocks(): ) assert fwd_context is not None - runner = GPUModelRunner(vllm_config, DEVICE) + runner = GPUModelRunner(vllm_config, DEVICE_TYPE) current_platform.update_block_size_for_backend(vllm_config) kv_cache_spec = runner.get_kv_cache_spec() diff --git a/tests/v1/worker/test_late_interaction_runner.py b/tests/v1/worker/test_late_interaction_runner.py index 5be3f6e6f10..9719485cd54 100644 --- a/tests/v1/worker/test_late_interaction_runner.py +++ b/tests/v1/worker/test_late_interaction_runner.py @@ -4,12 +4,12 @@ import pytest import torch +from vllm.entrypoints.pooling.scoring.utils import compute_maxsim_score from vllm.pooling_params import LateInteractionParams, PoolingParams from vllm.v1.pool.late_interaction import ( LATE_INTERACTION_MODE_CACHE_QUERY, build_late_interaction_doc_params, build_late_interaction_query_params, - compute_maxsim_score, ) from vllm.v1.worker.gpu.pool.late_interaction_runner import LateInteractionRunner diff --git a/tools/pre_commit/generate_attention_backend_docs.py b/tools/pre_commit/generate_attention_backend_docs.py index bbbf4f4b64f..9e14f8739dc 100644 --- a/tools/pre_commit/generate_attention_backend_docs.py +++ b/tools/pre_commit/generate_attention_backend_docs.py @@ -235,10 +235,11 @@ def _resolve_import_to_file( def _find_cc_in_function(tree: ast.AST, func_name: str) -> str | None: - """Find a compute capability from is_device_capability_family() calls in a function. + """Find a compute capability from is_device_capability*() calls in a function. - Looks for the pattern: current_platform.is_device_capability_family(N) - and converts N (e.g. 100) to a CC string (e.g. "10.x"). + Handles two patterns: + - is_device_capability_family(N): "M.x" (e.g. 100 -> "10.x") + - is_device_capability(N): "M.m" (e.g. 100 -> "10.0") """ for node in ast.walk(tree): if not isinstance(node, ast.FunctionDef) or node.name != func_name: @@ -247,12 +248,15 @@ def _find_cc_in_function(tree: ast.AST, func_name: str) -> str | None: if ( isinstance(n, ast.Call) and isinstance(n.func, ast.Attribute) - and n.func.attr == "is_device_capability_family" and n.args and isinstance(n.args[0], ast.Constant) and isinstance(n.args[0].value, int) ): - return f"{n.args[0].value // 10}.x" + val = n.args[0].value + if n.func.attr == "is_device_capability_family": + return f"{val // 10}.x" + elif n.func.attr == "is_device_capability": + return f"{val // 10}.{val % 10}" return None diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index 55e50e4ec10..0c2a53ec02e 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -265,6 +265,7 @@ def merge_attn_states( suffix_lse: torch.Tensor, output_lse: torch.Tensor | None = None, prefill_tokens_with_context: int | None = None, + output_scale: torch.Tensor | None = None, ) -> None: torch.ops._C.merge_attn_states( output, @@ -274,6 +275,7 @@ def merge_attn_states( suffix_output, suffix_lse, prefill_tokens_with_context, + output_scale, ) @@ -579,6 +581,56 @@ def rms_norm_per_block_quant( return output, scales +# fused silu_and_mul + block quant +def silu_and_mul_per_block_quant( + input: torch.Tensor, + group_size: int, # Changed from list[int] + quant_dtype: torch.dtype, + scale_ub: torch.Tensor | None = None, + is_scale_transposed: bool = False, +) -> tuple[torch.Tensor, torch.Tensor]: + assert input.ndim == 2, f"input must be 2D [batch, hidden*2], got {input.shape}" + assert input.shape[-1] % 2 == 0, ( + f"input last dim must be even (gate||up layout), got {input.shape[-1]}" + ) + + # Output is half the width of input (after silu_and_mul) + num_tokens = input.shape[0] + hidden_size = input.shape[-1] // 2 # Divide by 2 because input is [gate || up] + + # Allocate output tensor (FP8 or INT8) + output = torch.empty( + (num_tokens, hidden_size), device=input.device, dtype=quant_dtype + ) + + # Allocate scales tensor + num_groups = hidden_size // group_size # Directly use group_size + if is_scale_transposed: + scales = torch.empty( + (num_groups, num_tokens), + device=input.device, + dtype=torch.float32, + ).t() + else: + scales = torch.empty( + (num_tokens, num_groups), + device=input.device, + dtype=torch.float32, + ) + + # Call the C++ kernel + torch.ops._C.silu_and_mul_per_block_quant( + output, + input, + scales, + group_size, # Pass directly as int + scale_ub, + is_scale_transposed, + ) + + return output, scales + + # quantization ops # awq def awq_dequantize( @@ -2277,19 +2329,6 @@ def dsv3_router_gemm( return output -def gpt_oss_router_gemm( - hidden_states: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor -) -> torch.Tensor: - output = torch.empty( - hidden_states.shape[0], - weight.shape[0], - device=hidden_states.device, - dtype=hidden_states.dtype, - ) - torch.ops._moe_C.gpt_oss_router_gemm(output, hidden_states, weight, bias) - return output - - def topk_softmax( topk_weights: torch.Tensor, topk_ids: torch.Tensor, @@ -2602,6 +2641,22 @@ def swap_blocks( torch.ops._C_cache_ops.swap_blocks(src, dst, block_size_in_bytes, block_mapping) +def swap_blocks_batch( + src_ptrs: torch.Tensor, + dst_ptrs: torch.Tensor, + sizes: torch.Tensor, +) -> None: + """ + Batch version of swap_blocks: submit all copies in a single driver call. + + Each entry specifies a raw pointer copy: src_ptrs[i] -> dst_ptrs[i] + of sizes[i] bytes. All three tensors must be int64 CPU tensors. + On CUDA 12.8+ this uses cuMemcpyBatchAsync for minimal submission + overhead; on older CUDA it falls back to a loop of cudaMemcpyAsync. + """ + torch.ops._C_cache_ops.swap_blocks_batch(src_ptrs, dst_ptrs, sizes) + + def convert_fp8( output: torch.Tensor, input: torch.Tensor, scale: float = 1.0, kv_dtype: str = "fp8" ) -> None: diff --git a/vllm/compilation/decorators.py b/vllm/compilation/decorators.py index ab52d544c61..9c55a42a492 100644 --- a/vllm/compilation/decorators.py +++ b/vllm/compilation/decorators.py @@ -205,6 +205,8 @@ def support_torch_compile( if v.annotation in [ torch.Tensor, torch.Tensor | None, + torch.FloatTensor, + torch.FloatTensor | None, IntermediateTensors, IntermediateTensors | None, ]: @@ -346,7 +348,7 @@ def _support_torch_compile( def __init__( self: _T, - *, + *args, vllm_config: VllmConfig | None = None, prefix: str = "", **kwargs: Any, @@ -357,11 +359,24 @@ def _support_torch_compile( # NOTE: to support multimodal models (such as encoder), # we may not have vllm_config so we may need to patch it sig = inspect.signature(old_init) + # Check that any positional arguments match the old_init method signature + annotations = [p.annotation for p in sig.parameters.values()] + for arg, annotation in zip(args, annotations): + if annotation is inspect._empty: + continue + if not isinstance(arg, annotation): + init = f"'{type(self).__name__}.__init__'" + arg_type = f"'{type(arg).__name__}'" + raise TypeError( + f"{init} received a positional argument of type {arg_type}, " + "but no parameter of that type was found in the method signature. " + f"Please either annotate {init} or pass it as a keyword argument." + ) if "vllm_config" in sig.parameters: kwargs["vllm_config"] = vllm_config if "prefix" in sig.parameters: kwargs["prefix"] = prefix - old_init(self, **kwargs) + old_init(self, *args, **kwargs) self.vllm_config = vllm_config self.compilation_config = self.vllm_config.compilation_config diff --git a/vllm/compilation/passes/fusion/act_quant_fusion.py b/vllm/compilation/passes/fusion/act_quant_fusion.py index 911775f6996..2a1d37a1dae 100644 --- a/vllm/compilation/passes/fusion/act_quant_fusion.py +++ b/vllm/compilation/passes/fusion/act_quant_fusion.py @@ -17,6 +17,8 @@ from vllm.config import VllmConfig from vllm.logger import init_logger from vllm.model_executor.layers.quantization.utils.quant_utils import ( QuantKey, + kFp8Dynamic64Sym, + kFp8Dynamic128Sym, kFp8StaticTensorSym, kNvfp4Dynamic, ) @@ -43,6 +45,10 @@ silu_and_mul_nvfp4_quant_supported = current_platform.is_cuda() and hasattr( if silu_and_mul_nvfp4_quant_supported: FUSED_OPS[kNvfp4Dynamic] = torch.ops._C.silu_and_mul_nvfp4_quant.default # noqa: E501 +if current_platform.is_cuda(): + FUSED_OPS[kFp8Dynamic128Sym] = torch.ops._C.silu_and_mul_per_block_quant.default + FUSED_OPS[kFp8Dynamic64Sym] = torch.ops._C.silu_and_mul_per_block_quant.default + class ActivationQuantPattern(ABC): """ @@ -174,6 +180,102 @@ class SiluMulNvfp4QuantPattern(ActivationQuantPattern): register_replacement(pattern, replacement, self.get_inputs(), fwd_only, pm_pass) +class SiluMulBlockQuantPattern(ActivationQuantPattern): + """ + Fusion for SiluMul+BlockQuant (FP8 dynamic per-group) Pattern. + Supports group_size 128 and 64 via QuantKey. + Parameterized on is_scale_transposed for different scale layouts. + """ + + def __init__( + self, + quant_key: QuantKey, + is_scale_transposed: bool = False, + is_e8m0: bool = False, + is_tma_aligned: bool = False, + ) -> None: + super().__init__(quant_key) + self.quant_matcher = MatcherQuantFP8( + quant_key, + has_col_major_scales=is_scale_transposed, + is_e8m0=is_e8m0, + is_tma_aligned=is_tma_aligned, + ) + self.group_size = quant_key.scale.group_shape[1] + self.is_scale_transposed = is_scale_transposed + self.is_e8m0 = is_e8m0 + self.is_tma_aligned = is_tma_aligned + + def get_inputs(self) -> list[torch.Tensor]: + scale = self.quant_matcher.empty_f32(1, 1) + return self.silu_and_mul_matcher.inputs() + [scale] + + def register(self, pm_pass: PatternMatcherPass) -> None: + is_scale_transposed = self.is_scale_transposed + + def pattern( + input: torch.Tensor, + scale: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + silu_out = self.silu_and_mul_matcher(input) + result = torch.empty( + silu_out.shape, + device=silu_out.device, + dtype=self.quant_dtype, + ) + assert scale is not None + finfo = torch.finfo(self.quant_dtype) + _, result, scale = auto_functionalized( + self.quant_matcher.QUANT_OP, + input=silu_out, + output_q=result, + output_s=scale, + group_size=self.group_size, + eps=1e-10, + fp8_min=finfo.min, + fp8_max=finfo.max, + scale_ue8m0=self.is_e8m0, + dummy_is_scale_transposed=is_scale_transposed, + dummy_is_tma_aligned=self.is_tma_aligned, + ) + return result, scale + + def replacement( + input: torch.Tensor, + scale: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + d = input.shape[-1] // 2 + output_shape = input.shape[:-1] + (d,) + result = torch.empty( + output_shape, device=input.device, dtype=self.quant_dtype + ) + if is_scale_transposed: + scale = torch.empty( + (d // self.group_size, input.shape[0]), + device=input.device, + dtype=torch.float32, + ).permute(-1, -2) + else: + scale = torch.empty( + (input.shape[0], d // self.group_size), + device=input.device, + dtype=torch.float32, + ) + at = auto_functionalized( + self.FUSED_OP, + out=result, + input=input, + scales=scale, + group_size=self.group_size, + scale_ub=None, + is_scale_transposed=is_scale_transposed, + ) + return at[1], at[2] + + inps = self.get_inputs() + register_replacement(pattern, replacement, inps, fwd_only, pm_pass) + + class ActivationQuantFusionPass(VllmPatternMatcherPass): """ This pass fuses a pre-defined set of custom ops into fused ops. @@ -199,6 +301,18 @@ class ActivationQuantFusionPass(VllmPatternMatcherPass): pattern_silu_mul_nvfp4 = SiluMulNvfp4QuantPattern() pattern_silu_mul_nvfp4.register(self.patterns) + if current_platform.is_cuda(): + for quant_key in [kFp8Dynamic128Sym, kFp8Dynamic64Sym]: + for is_scale_transposed in [False, True]: + for is_e8m0 in [True, False]: + for is_tma_aligned in [False, True]: + SiluMulBlockQuantPattern( + quant_key, + is_scale_transposed=is_scale_transposed, + is_e8m0=is_e8m0, + is_tma_aligned=is_tma_aligned, + ).register(self.patterns) + self.dump_patterns(config, self.patterns) @VllmInductorPass.time_and_log @@ -212,4 +326,5 @@ class ActivationQuantFusionPass(VllmPatternMatcherPass): ActivationQuantPattern, SiluMulFp8StaticQuantPattern, SiluMulNvfp4QuantPattern, + SiluMulBlockQuantPattern, ) diff --git a/vllm/compilation/passes/fusion/mla_attn_quant_fusion.py b/vllm/compilation/passes/fusion/mla_attn_quant_fusion.py new file mode 100644 index 00000000000..5a9ef46a0fc --- /dev/null +++ b/vllm/compilation/passes/fusion/mla_attn_quant_fusion.py @@ -0,0 +1,262 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from collections.abc import Callable + +import torch +from torch._higher_order_ops.auto_functionalize import auto_functionalized + +from vllm._custom_ops import create_fp4_output_tensors +from vllm.config import VllmConfig, get_layers_from_vllm_config +from vllm.logger import init_logger +from vllm.model_executor.layers.attention.mla_attention import MLAAttention +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + kFp8StaticTensorSym, + kNvfp4Dynamic, +) +from vllm.platforms import current_platform + +from ..vllm_inductor_pass import VllmFusionPatternMatcherPass, VllmPatternReplacement +from .matcher_utils import MatcherQuantFP8 +from .rms_quant_fusion import QUANT_OPS + +logger = init_logger(__name__) + +FP8_DTYPE = current_platform.fp8_dtype() +FP4_DTYPE = torch.uint8 + +MLA_ATTN_OP = torch.ops.vllm.unified_mla_attention_with_output.default + + +class MLAAttnFp8StaticQuantPattern(VllmPatternReplacement[..., torch.Tensor]): + """ + Fusion for MLA Attention+Fp8StaticQuant. + + Matches the pattern: MLA attention -> static FP8 quant, and replaces + it with MLA attention(output_scale=scale, output=fp8_buffer). + """ + + def __init__(self, layer: MLAAttention, dtype: torch.dtype) -> None: + self._layer_name = layer.layer_name + self._num_heads = layer.num_heads + self._v_head_dim = layer.v_head_dim + self._kv_lora_rank = layer.kv_lora_rank + self._qk_rope_head_dim = layer.qk_rope_head_dim + self._qk_head_dim = layer.qk_nope_head_dim + layer.qk_rope_head_dim + self._output_dim = layer.num_heads * layer.v_head_dim + self._dtype = dtype + self._quant_matcher = MatcherQuantFP8(kFp8StaticTensorSym) + + @property + def pattern(self) -> Callable[..., torch.Tensor]: + def _pattern( + q: torch.Tensor, + kv_c_normed: torch.Tensor, + k_pe: torch.Tensor, + output_attn: torch.Tensor, + scale: torch.Tensor, + kv_cache_dummy_dep: torch.Tensor, + ) -> torch.Tensor: + at1 = auto_functionalized( + MLA_ATTN_OP, + q=q, + kv_c_normed=kv_c_normed, + k_pe=k_pe, + output=output_attn, + layer_name=self._layer_name, + output_scale=None, + output_block_scale=None, + kv_cache_dummy_dep=kv_cache_dummy_dep, + ) + # MLA output is already 2D (T, N*V), no reshape needed + return self._quant_matcher(at1[1], scale)[0] + + return _pattern + + @property + def replacement(self) -> Callable[..., torch.Tensor]: + def _replacement( + q: torch.Tensor, + kv_c_normed: torch.Tensor, + k_pe: torch.Tensor, + output_attn: torch.Tensor, + scale: torch.Tensor, + kv_cache_dummy_dep: torch.Tensor, + ) -> torch.Tensor: + # MLA output in quant_dtype + output_attn = torch.empty( + [q.shape[0], self._output_dim], + dtype=FP8_DTYPE, + device=q.device, + ) + at1 = auto_functionalized( + MLA_ATTN_OP, + q=q, + kv_c_normed=kv_c_normed, + k_pe=k_pe, + output=output_attn, + layer_name=self._layer_name, + output_scale=scale, + output_block_scale=None, + kv_cache_dummy_dep=kv_cache_dummy_dep, + ) + return at1[1] + + return _replacement + + def get_inputs(self) -> list[torch.Tensor]: + return [ + self.empty(5, self._num_heads, self._qk_head_dim, dtype=self._dtype), + self.empty(5, self._kv_lora_rank, dtype=self._dtype), + self.empty(5, 1, self._qk_rope_head_dim, dtype=self._dtype), + self.empty(5, self._output_dim, dtype=self._dtype), + self.empty_fp32(1, 1), + self.empty(0, dtype=self._dtype), + ] + + +class MLAAttnNvfp4QuantPattern( + VllmPatternReplacement[..., tuple[torch.Tensor, torch.Tensor]] +): + """ + Fusion for MLA Attention+Nvfp4Quant. + + Matches the pattern: MLA attention -> NVFP4 quant, and replaces + it with MLA attention(output_scale=scale, output_block_scale=block_scale, + output=fp4_buffer). + """ + + def __init__(self, layer: MLAAttention, dtype: torch.dtype) -> None: + self._layer_name = layer.layer_name + self._num_heads = layer.num_heads + self._v_head_dim = layer.v_head_dim + self._kv_lora_rank = layer.kv_lora_rank + self._qk_rope_head_dim = layer.qk_rope_head_dim + self._qk_head_dim = layer.qk_nope_head_dim + layer.qk_rope_head_dim + self._output_dim = layer.num_heads * layer.v_head_dim + self._dtype = dtype + self._QUANT_OP = QUANT_OPS[kNvfp4Dynamic] + + @property + def pattern( + self, + ) -> Callable[..., tuple[torch.Tensor, torch.Tensor]]: + def _pattern( + q: torch.Tensor, + kv_c_normed: torch.Tensor, + k_pe: torch.Tensor, + output_attn: torch.Tensor, + input_scale: torch.Tensor, + kv_cache_dummy_dep: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + at1 = auto_functionalized( + MLA_ATTN_OP, + q=q, + kv_c_normed=kv_c_normed, + k_pe=k_pe, + output=output_attn, + layer_name=self._layer_name, + output_scale=None, + output_block_scale=None, + kv_cache_dummy_dep=kv_cache_dummy_dep, + ) + # Replicate what scaled_fp4_quant() does: allocate output + # tensors inline then call the .out variant. + output_quant, output_scale = create_fp4_output_tensors( + at1[1].shape[0], at1[1].shape[1], at1[1].device, True + ) + at2 = auto_functionalized( + self._QUANT_OP, + input=at1[1], + input_scale=input_scale, + is_sf_swizzled_layout=True, + output=output_quant, + output_scale=output_scale, + ) + output_scale_view = torch.ops.aten.view.dtype(at2[2], FP8_DTYPE) + return at2[1], output_scale_view + + return _pattern + + @property + def replacement( + self, + ) -> Callable[..., tuple[torch.Tensor, torch.Tensor]]: + def _replacement( + q: torch.Tensor, + kv_c_normed: torch.Tensor, + k_pe: torch.Tensor, + output_attn: torch.Tensor, + input_scale: torch.Tensor, + kv_cache_dummy_dep: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + # MLA output in quant_dtype (FP4 packed as uint8) + output_attn = torch.empty( + [q.shape[0], self._output_dim // 2], + dtype=FP4_DTYPE, + device=q.device, + ) + # attention output block scale + output_scale = create_fp4_output_tensors( + q.shape[0], self._output_dim, q.device, True + )[1] + output_scale_view = torch.ops.aten.view.dtype(output_scale, FP8_DTYPE) + at2 = auto_functionalized( + MLA_ATTN_OP, + q=q, + kv_c_normed=kv_c_normed, + k_pe=k_pe, + output=output_attn, + layer_name=self._layer_name, + output_scale=input_scale, + output_block_scale=output_scale_view, + kv_cache_dummy_dep=kv_cache_dummy_dep, + ) + return at2[1], at2[2] + + return _replacement + + def get_inputs(self) -> list[torch.Tensor]: + return [ + self.empty(5, self._num_heads, self._qk_head_dim, dtype=self._dtype), + self.empty(5, self._kv_lora_rank, dtype=self._dtype), + self.empty(5, 1, self._qk_rope_head_dim, dtype=self._dtype), + self.empty(5, self._output_dim, dtype=self._dtype), + self.empty_fp32(1, 1), + self.empty(0, dtype=self._dtype), + ] + + +class MLAAttnQuantFusionPass(VllmFusionPatternMatcherPass): + """ + This pass fuses post-attention quantization onto MLA attention if supported. + + It uses the pattern matcher and matches each MLA layer manually, as strings + cannot be wildcarded. This also lets us check support on attention layers + upon registration instead of during pattern matching. + """ + + def __init__(self, config: VllmConfig) -> None: + super().__init__(config, "mla_attn_quant_fusion") + + dtype = config.model_config.dtype + layers = list(get_layers_from_vllm_config(config, MLAAttention).values()) + + if len(layers) == 0: + logger.warning( + "MLA attention + quant fusion is enabled, but no MLA " + "attention layers were found in " + "CompilationConfig.static_forward_context " + "so no fusion patterns were registered." + ) + + for layer in layers: + if layer.impl.fused_output_quant_supported(kFp8StaticTensorSym): + self.register(MLAAttnFp8StaticQuantPattern(layer, dtype)) + + if current_platform.is_cuda() and hasattr(torch.ops._C, "scaled_fp4_quant"): + for layer in layers: + if layer.impl.fused_output_quant_supported(kNvfp4Dynamic): + self.register(MLAAttnNvfp4QuantPattern(layer, dtype)) + + self.dump_patterns(config, self.pm_pass) diff --git a/vllm/compilation/passes/ir/lowering_pass.py b/vllm/compilation/passes/ir/lowering_pass.py index 474a09ae22b..02acdd1a298 100644 --- a/vllm/compilation/passes/ir/lowering_pass.py +++ b/vllm/compilation/passes/ir/lowering_pass.py @@ -149,9 +149,8 @@ class VllmIRLoweringPass(VllmInductorPass): ) impl_uuids_str = ";".join( - f"{name}={ - ','.join(IrOp.registry[name].impls[provider].uuid() for provider in p) - }" + f"{name}=" + + ",".join(IrOp.registry[name].impls[provider].uuid() for provider in p) for name, p in priorities.items() ) diff --git a/vllm/compilation/passes/pass_manager.py b/vllm/compilation/passes/pass_manager.py index 0571741419f..b4823a0afde 100644 --- a/vllm/compilation/passes/pass_manager.py +++ b/vllm/compilation/passes/pass_manager.py @@ -27,6 +27,7 @@ if rocm_aiter_ops.is_enabled(): if current_platform.is_cuda_alike(): from .fusion.act_quant_fusion import ActivationQuantFusionPass from .fusion.attn_quant_fusion import AttnQuantFusionPass + from .fusion.mla_attn_quant_fusion import MLAAttnQuantFusionPass from .fusion.qk_norm_rope_fusion import QKNormRoPEFusionPass from .fusion.rms_quant_fusion import RMSNormQuantFusionPass from .fusion.rope_kvcache_fusion import RopeKVCacheFusionPass @@ -157,6 +158,7 @@ class PostGradPassManager(CustomGraphPass): # type: ignore[misc] if self.pass_config.fuse_attn_quant: self.passes += [AttnQuantFusionPass(config)] + self.passes += [MLAAttnQuantFusionPass(config)] if self.pass_config.enable_qk_norm_rope_fusion: self.passes += [SplitCoalescingPass(config)] diff --git a/vllm/config/cache.py b/vllm/config/cache.py index 1fdce002e9e..cd1554590ea 100644 --- a/vllm/config/cache.py +++ b/vllm/config/cache.py @@ -8,7 +8,10 @@ from pydantic import Field, SkipValidation, field_validator, model_validator from vllm.config.utils import config from vllm.logger import init_logger -from vllm.utils.torch_utils import is_quantized_kv_cache +from vllm.utils.torch_utils import ( + is_quantized_kv_cache, + kv_cache_uses_per_token_head_scales, +) logger = init_logger(__name__) @@ -21,6 +24,8 @@ CacheDType = Literal[ "fp8_e5m2", "fp8_inc", "fp8_ds_mla", + "int8_per_token_head", + "fp8_per_token_head", ] MambaDType = Literal["auto", "float32", "float16"] MambaCacheMode = Literal["all", "align", "none"] @@ -237,12 +242,20 @@ class CacheConfig: @field_validator("cache_dtype", mode="after") @classmethod def _validate_cache_dtype(cls, cache_dtype: CacheDType) -> CacheDType: - if is_quantized_kv_cache(cache_dtype): + if kv_cache_uses_per_token_head_scales(cache_dtype): logger.info( - "Using fp8 data type to store kv cache. It reduces the GPU " + "Using %s data type to store kv cache. It reduces the GPU " + "memory footprint and boosts the performance. " + "Dynamic per-token-head scales will be computed at runtime.", + str(cache_dtype), + ) + elif is_quantized_kv_cache(cache_dtype): + logger.info( + "Using %s data type to store kv cache. It reduces the GPU " "memory footprint and boosts the performance. " "Meanwhile, it may cause accuracy drop without a proper " - "scaling factor." + "scaling factor", + str(cache_dtype), ) return cache_dtype diff --git a/vllm/config/compilation.py b/vllm/config/compilation.py index 916c5a00205..716c208a904 100644 --- a/vllm/config/compilation.py +++ b/vllm/config/compilation.py @@ -121,7 +121,7 @@ class PassConfig: fuse_act_quant: bool = None # type: ignore[assignment] """Fuse the custom SiluMul + quant ops.""" fuse_attn_quant: bool = None # type: ignore[assignment] - """Fuse the custom attention + quant ops.""" + """Fuse the custom Attention and MLAAttention + quant ops.""" eliminate_noops: bool = Field(default=True) """Eliminate no-op ops.""" enable_sp: bool = None # type: ignore[assignment] @@ -495,9 +495,10 @@ class CompilationConfig: If empty list [], no ops are excluded (suitable for full cudagraphs).""" compile_mm_encoder: bool = False """Whether or not to compile the multimodal encoder. - Currently, this only works for `Qwen2_5_vl` and `mLLaMa4` models - on selected platforms. Disabled by default until more models - are supported/tested to work.""" + Currently, this only works for `Qwen2_5_vl` and `mLLaMa4` models on selected + platforms. It may also work for models loaded with the Transformers modeling backend + if the encoder is compilable. Disabled by default until more models are + supported/tested to work.""" # Vision encoder CUDA graph cudagraph_mm_encoder: bool = False diff --git a/vllm/config/model.py b/vllm/config/model.py index c4ee654fe8b..7bb3655f29e 100644 --- a/vllm/config/model.py +++ b/vllm/config/model.py @@ -1203,6 +1203,7 @@ class ModelConfig: "gemma3", "molmo2", "paligemma", + "umm", ) if not hasattr(self.hf_config, "model_type"): return False diff --git a/vllm/entrypoints/grpc_server.py b/vllm/entrypoints/grpc_server.py index aec9b15bb81..ddd8a5c50e4 100644 --- a/vllm/entrypoints/grpc_server.py +++ b/vllm/entrypoints/grpc_server.py @@ -29,11 +29,13 @@ try: from grpc_reflection.v1alpha import reflection from smg_grpc_proto import vllm_engine_pb2, vllm_engine_pb2_grpc from smg_grpc_servicer.vllm.servicer import VllmEngineServicer -except ImportError: +except ImportError as e: raise ImportError( - "smg-grpc-servicer is required for gRPC mode. " - "Install it with: pip install vllm[grpc]" - ) from None + "gRPC mode requires smg-grpc-servicer. " + "If not installed, run: pip install vllm[grpc]. " + "If already installed, there may be a broken import due to a " + "version mismatch — see the chained exception above for details." + ) from e import uvloop diff --git a/vllm/entrypoints/llm.py b/vllm/entrypoints/llm.py index a576a3f28a8..b9eea87451b 100644 --- a/vllm/entrypoints/llm.py +++ b/vllm/entrypoints/llm.py @@ -1449,14 +1449,14 @@ class LLM: pooling_task = io_processor.pooling_task scoring_data = io_processor.valid_inputs(data_1, data_2) - offset = len(scoring_data.data_1) + n_queries = len(scoring_data.data_1) ctx = OfflineInputsContext( prompts=scoring_data, pooling_params=pooling_params, tokenization_kwargs=tokenization_kwargs, chat_template=chat_template, - offset=offset, + n_queries=n_queries, ) processor_inputs = io_processor.pre_process_offline(ctx) @@ -1487,7 +1487,7 @@ class LLM: outputs = self._run_engine(use_tqdm=use_tqdm, output_type=PoolingRequestOutput) outputs = io_processor.post_process_offline( - ctx=OfflineOutputsContext(outputs=outputs, offset=offset), + ctx=OfflineOutputsContext(outputs=outputs, n_queries=n_queries), ) return [ScoringRequestOutput.from_base(item) for item in outputs] diff --git a/vllm/entrypoints/openai/cli_args.py b/vllm/entrypoints/openai/cli_args.py index 7491c41c271..898e62f7713 100644 --- a/vllm/entrypoints/openai/cli_args.py +++ b/vllm/entrypoints/openai/cli_args.py @@ -278,6 +278,9 @@ class FrontendArgs(BaseFrontendArgs): Enable offline FastAPI documentation for air-gapped environments. Uses vendored static assets bundled with vLLM. """ + enable_flash_late_interaction: bool = True + """If set, run pooling score MaxSim on GPU in the API server process. + Can significantly improve late-interaction scoring performance.""" @classmethod def _customize_cli_kwargs( diff --git a/vllm/entrypoints/pooling/__init__.py b/vllm/entrypoints/pooling/__init__.py index b843c791319..fb0c10e6f4f 100644 --- a/vllm/entrypoints/pooling/__init__.py +++ b/vllm/entrypoints/pooling/__init__.py @@ -123,6 +123,9 @@ def init_pooling_state( chat_template=resolved_chat_template, chat_template_content_format=args.chat_template_content_format, trust_request_chat_template=args.trust_request_chat_template, + enable_flash_late_interaction=getattr( + args, "enable_flash_late_interaction", True + ), ) if enable_scoring_api(supported_tasks, model_config) else None diff --git a/vllm/entrypoints/pooling/base/serving.py b/vllm/entrypoints/pooling/base/serving.py index cf6e01742ce..90554aa634b 100644 --- a/vllm/entrypoints/pooling/base/serving.py +++ b/vllm/entrypoints/pooling/base/serving.py @@ -82,9 +82,20 @@ class PoolingServing: request: AnyPoolingRequest, raw_request: Request | None = None, ) -> Response: + ctx = await self._init_ctx(request, raw_request) + await self.io_processor.pre_process_online_async(ctx) + await self._prepare_generators(ctx) + await self._collect_batch(ctx) + await self.io_processor.post_process_online_async(ctx) + return await self._build_response(ctx) + + async def _init_ctx( + self, + request: AnyPoolingRequest, + raw_request: Request | None = None, + ): model_name = self.models.model_name() request_id = f"{self.request_id_prefix}-{self._base_request_id(raw_request)}" - await self._check_model(request) ctx = PoolingServeContext( @@ -96,11 +107,7 @@ class PoolingServing: self._validate_request(ctx) self._maybe_get_adapters(ctx) - await self.io_processor.pre_process_online_async(ctx) - await self._prepare_generators(ctx) - await self._collect_batch(ctx) - await self.io_processor.post_process_online_async(ctx) - return await self._build_response(ctx) + return ctx async def _prepare_generators( self, diff --git a/vllm/entrypoints/pooling/scoring/io_processor.py b/vllm/entrypoints/pooling/scoring/io_processor.py index 70fe1b22141..c520eb5ceb3 100644 --- a/vllm/entrypoints/pooling/scoring/io_processor.py +++ b/vllm/entrypoints/pooling/scoring/io_processor.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import time from collections.abc import Sequence -from typing import Any, TypeAlias, cast +from typing import Any, TypeAlias import torch.nn.functional as F @@ -16,7 +16,7 @@ from vllm.entrypoints.pooling.typing import ( from vllm.inputs import EngineInput from vllm.renderers import TokenizeParams from vllm.renderers.hf import safe_apply_chat_template -from vllm.tasks import PoolingTask, ScoreType +from vllm.tasks import PoolingTask from vllm.utils.mistral import is_mistral_tokenizer from ...chat_utils import ChatTemplateResolutionError @@ -34,7 +34,7 @@ ScoringServeContext: TypeAlias = PoolingServeContext[ScoringRequest] class ScoringIOProcessor(PoolingIOProcessor): - name: ScoreType + name: str pooling_task: PoolingTask def __init__(self, *args, **kwargs): @@ -63,7 +63,7 @@ class ScoringIOProcessor(PoolingIOProcessor): class BiEncoderIOProcessor(ScoringIOProcessor): - name: ScoreType = "bi-encoder" + name = "bi-encoder" pooling_task: PoolingTask = "embed" ####################################### @@ -94,20 +94,17 @@ class BiEncoderIOProcessor(ScoringIOProcessor): ) ctx.engine_inputs = engine_inputs - ctx.intermediates = len(scoring_data.data_1) + ctx.n_queries = len(scoring_data.data_1) def post_process_online( self, ctx: ScoringServeContext, ): - if ctx.final_res_batch is None: - raise ValueError("Final response batch not available") - - if ctx.intermediates is None: - raise ValueError("data_1 len not available") + assert ctx.final_res_batch is not None + assert isinstance(ctx.n_queries, int) ctx.final_res_batch = self._post_process( - outputs=ctx.final_res_batch, offset=cast(int, ctx.intermediates) + outputs=ctx.final_res_batch, n_queries=ctx.n_queries ) ####################################### @@ -124,8 +121,8 @@ class BiEncoderIOProcessor(ScoringIOProcessor): self, ctx: OfflineOutputsContext, ) -> list[PoolingRequestOutput]: - assert ctx.offset is not None - return self._post_process(outputs=ctx.outputs, offset=ctx.offset) + assert ctx.n_queries is not None + return self._post_process(outputs=ctx.outputs, n_queries=ctx.n_queries) ####################################### # helpers @@ -145,9 +142,9 @@ class BiEncoderIOProcessor(ScoringIOProcessor): prompts=data_1 + data_2, tok_params=tok_params, prompt_extras=prompt_extras ) - def _post_process(self, outputs: list[PoolingRequestOutput], offset: int): - emb_data_1 = outputs[:offset] - emb_data_2 = outputs[offset:] + def _post_process(self, outputs: list[PoolingRequestOutput], n_queries: int): + emb_data_1 = outputs[:n_queries] + emb_data_2 = outputs[n_queries:] if len(emb_data_1) == 1: emb_data_1 = emb_data_1 * len(emb_data_2) @@ -177,13 +174,13 @@ class BiEncoderIOProcessor(ScoringIOProcessor): class LateInteractionIOProcessor(BiEncoderIOProcessor): - name: ScoreType = "late-interaction" + name = "late-interaction" pooling_task: PoolingTask = "token_embed" - def _post_process(self, outputs: list[PoolingRequestOutput], offset: int): + def _post_process(self, outputs: list[PoolingRequestOutput], n_queries: int): # Split into query and document embeddings - emb_data_1 = outputs[:offset] - emb_data_2 = outputs[offset:] + emb_data_1 = outputs[:n_queries] + emb_data_2 = outputs[n_queries:] # Expand queries if 1:N scoring if len(emb_data_1) == 1: @@ -217,8 +214,15 @@ class LateInteractionIOProcessor(BiEncoderIOProcessor): return final_res_batch +class FlashLateInteractionIOProcessor(LateInteractionIOProcessor): + name = "flash-late-interaction" + + def _post_process(self, outputs: list[PoolingRequestOutput], n_queries: int): + return outputs + + class CrossEncoderIOProcessor(ScoringIOProcessor): - name: ScoreType = "cross-encoder" + name = "cross-encoder" pooling_task: PoolingTask = "classify" def __init__(self, *args, **kwargs): @@ -412,8 +416,12 @@ class CrossEncoderIOProcessor(ScoringIOProcessor): return full_prompt, engine_prompt -ScoringIOProcessors: dict[ScoreType, type[ScoringIOProcessor]] = { - "bi-encoder": BiEncoderIOProcessor, - "late-interaction": LateInteractionIOProcessor, - "cross-encoder": CrossEncoderIOProcessor, +ScoringIOProcessors: dict[str, type[ScoringIOProcessor]] = { + p.name: p + for p in [ + BiEncoderIOProcessor, + LateInteractionIOProcessor, + FlashLateInteractionIOProcessor, + CrossEncoderIOProcessor, + ] } diff --git a/vllm/entrypoints/pooling/scoring/serving.py b/vllm/entrypoints/pooling/scoring/serving.py index 57e5684e4e3..de5b5797ce4 100644 --- a/vllm/entrypoints/pooling/scoring/serving.py +++ b/vllm/entrypoints/pooling/scoring/serving.py @@ -1,9 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from fastapi.responses import JSONResponse +from fastapi.responses import JSONResponse, Response +from vllm import PoolingParams from vllm.config import ModelConfig +from vllm.engine.protocol import EngineClient from vllm.entrypoints.chat_utils import ChatTemplateConfig from vllm.entrypoints.openai.engine.protocol import UsageInfo from vllm.entrypoints.pooling.base.io_processor import PoolingIOProcessor @@ -11,6 +13,10 @@ from vllm.entrypoints.pooling.base.serving import PoolingServing from vllm.logger import init_logger from vllm.outputs import PoolingRequestOutput, ScoringRequestOutput from vllm.renderers import BaseRenderer +from vllm.v1.pool.late_interaction import ( + build_late_interaction_doc_params, + build_late_interaction_query_params, +) from .io_processor import ScoringIOProcessors, ScoringServeContext from .protocol import ( @@ -31,13 +37,30 @@ logger = init_logger(__name__) class ServingScores(PoolingServing): request_id_prefix = "score" + def __init__( + self, + engine_client: EngineClient, + *args, + enable_flash_late_interaction: bool = True, + **kwargs, + ): + self.score_type = engine_client.model_config.score_type + self.enable_flash_late_interaction = ( + self.score_type == "late-interaction" and enable_flash_late_interaction + ) + + super().__init__(engine_client, *args, **kwargs) + def init_io_processor( self, model_config: ModelConfig, renderer: BaseRenderer, chat_template_config: ChatTemplateConfig, ) -> PoolingIOProcessor: - score_type = model_config.score_type + score_type: str = model_config.score_type + if self.enable_flash_late_interaction: + score_type = "flash-late-interaction" + assert score_type in ScoringIOProcessors processor_cls = ScoringIOProcessors[score_type] return processor_cls( @@ -46,6 +69,12 @@ class ServingScores(PoolingServing): chat_template_config=chat_template_config, ) + async def __call__(self, *args, **kwargs) -> Response: + if not self.enable_flash_late_interaction: + return await super().__call__(*args, **kwargs) + + return await self.flash_late_interaction(*args, **kwargs) + async def _build_response( self, ctx: ScoringServeContext, @@ -158,3 +187,106 @@ class ServingScores(PoolingServing): ) return JSONResponse(content=response.model_dump()) + + ################################################################################### + ### Run pooling score MaxSim on worker side (GPU) in the API server process + ### Can significantly improve late-interaction scoring performance. + + async def flash_late_interaction(self, *args, **kwargs) -> Response: + ctx = await self._init_ctx(*args, **kwargs) + ctx.pooling_params = self.io_processor.create_pooling_params(ctx.request) + await self.io_processor.pre_process_online_async(ctx) + + # stage 1: encode queries and cache token embeddings on workers. + await self._flash_late_interaction_encode_queries(ctx) + # stage 2: encode docs and return scalar scores from workers. + await self._flash_late_interaction_encode_docs(ctx) + + await self.io_processor.post_process_online_async(ctx) + return await self._build_response(ctx) + + async def _flash_late_interaction_encode_queries(self, ctx: ScoringServeContext): + assert ctx.n_queries is not None + assert ctx.engine_inputs is not None + assert isinstance(ctx.pooling_params, PoolingParams) + + n_queries = ctx.n_queries + n_docs = len(ctx.engine_inputs) - n_queries + query_engine_inputs = ctx.engine_inputs[:n_queries] + + query_keys = [f"{ctx.request_id}-query-{i}" for i in range(n_queries)] + query_uses = [n_docs if n_queries == 1 else 1] * n_queries + + query_pooling_params_list = [] + for i in range(n_queries): + pooling_params = ctx.pooling_params.clone() + pooling_params.late_interaction_params = ( + build_late_interaction_query_params( + query_key=query_keys[i], + query_uses=query_uses[i], + ) + ) + query_pooling_params_list.append(pooling_params) + + assert ( + n_queries + == len(query_pooling_params_list) + == len(query_engine_inputs) + == len(query_keys) + ) + + query_ctx = ScoringServeContext( + request=ctx.request, + raw_request=ctx.raw_request, + model_name=ctx.model_name, + request_id=ctx.request_id, + pooling_params=query_pooling_params_list, + prompt_request_ids=query_keys, + engine_inputs=query_engine_inputs, + ) + + await self._prepare_generators(query_ctx) + await self._collect_batch(query_ctx) + + async def _flash_late_interaction_encode_docs(self, ctx: ScoringServeContext): + assert ctx.n_queries is not None + assert ctx.engine_inputs is not None + assert isinstance(ctx.pooling_params, PoolingParams) + + n_queries = ctx.n_queries + n_docs = len(ctx.engine_inputs) - n_queries + doc_engine_inputs = ctx.engine_inputs[n_queries:] + + query_keys = [f"{ctx.request_id}-query-{i}" for i in range(n_queries)] + doc_keys = [f"{ctx.request_id}-doc-{i}" for i in range(n_docs)] + + doc_pooling_params_list = [] + for i in range(n_docs): + query_idx = 0 if n_queries == 1 else i + pooling_params = ctx.pooling_params.clone() + pooling_params.late_interaction_params = build_late_interaction_doc_params( + query_key=query_keys[query_idx] + ) + doc_pooling_params_list.append(pooling_params) + + assert ( + n_docs + == len(doc_pooling_params_list) + == len(doc_engine_inputs) + == len(doc_keys) + ) + + doc_ctx = ScoringServeContext( + request=ctx.request, + raw_request=ctx.raw_request, + model_name=ctx.model_name, + request_id=ctx.request_id, + pooling_params=doc_pooling_params_list, + prompt_request_ids=doc_keys, + engine_inputs=doc_engine_inputs, + ) + + await self._prepare_generators(doc_ctx) + await self._collect_batch(doc_ctx) + + ctx.final_res_batch = doc_ctx.final_res_batch diff --git a/vllm/entrypoints/pooling/scoring/utils.py b/vllm/entrypoints/pooling/scoring/utils.py index 812a75ab806..01b8514eb7b 100644 --- a/vllm/entrypoints/pooling/scoring/utils.py +++ b/vllm/entrypoints/pooling/scoring/utils.py @@ -36,8 +36,9 @@ def compute_maxsim_score(q_emb: torch.Tensor, d_emb: torch.Tensor) -> torch.Tens Returns: MaxSim score (sum over query tokens of max similarity to any doc token) """ + # compute in float32 for numerical stability # [query_len, doc_len] - token_scores = torch.matmul(q_emb, d_emb.T) + token_scores = torch.matmul(q_emb.float(), d_emb.float().T) # Max over document tokens, sum over query tokens return token_scores.amax(dim=-1).sum() diff --git a/vllm/entrypoints/pooling/typing.py b/vllm/entrypoints/pooling/typing.py index 8ccc5d49c2d..66dd9dd4d2b 100644 --- a/vllm/entrypoints/pooling/typing.py +++ b/vllm/entrypoints/pooling/typing.py @@ -83,6 +83,9 @@ class PoolingServeContext(Generic[PoolingRequestT]): model_config = ConfigDict(arbitrary_types_allowed=True) + ## for bi-encoder & late-interaction + n_queries: int | None = None + @dataclass class OfflineInputsContext: @@ -92,7 +95,7 @@ class OfflineInputsContext: chat_template: str | None = None ## for bi-encoder & late-interaction - offset: int | None = None + n_queries: int | None = None @dataclass @@ -100,4 +103,4 @@ class OfflineOutputsContext: outputs: list[PoolingRequestOutput] ## for bi-encoder & late-interaction - offset: int | None = None + n_queries: int | None = None diff --git a/vllm/entrypoints/serve/disagg/protocol.py b/vllm/entrypoints/serve/disagg/protocol.py index af4e8c20c14..345992d3b0b 100644 --- a/vllm/entrypoints/serve/disagg/protocol.py +++ b/vllm/entrypoints/serve/disagg/protocol.py @@ -6,7 +6,7 @@ from pydantic import BaseModel, Field, field_validator from vllm.config import ModelConfig from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionLogProbs -from vllm.entrypoints.openai.engine.protocol import StreamOptions +from vllm.entrypoints.openai.engine.protocol import StreamOptions, UsageInfo from vllm.logprobs import Logprob from vllm.renderers import TokenizeParams from vllm.sampling_params import SamplingParams @@ -122,6 +122,26 @@ class GenerateResponseChoice(BaseModel): token_ids: list[int] | None = None +class GenerateResponseStreamChoice(BaseModel): + index: int + logprobs: ChatCompletionLogProbs | None = None + finish_reason: str | None = None + token_ids: list[int] | None = None + + +class GenerateStreamResponse(BaseModel): + request_id: str = Field( + default_factory=lambda: f"{random_uuid()}", + description=( + "The request_id related to this request. If the caller does " + "not set it, a random_uuid will be generated. This id is used " + "through out the inference process and return in response." + ), + ) + choices: list[GenerateResponseStreamChoice] + usage: UsageInfo | None = Field(default=None) + + class GenerateResponse(BaseModel): request_id: str = Field( default_factory=lambda: f"{random_uuid()}", diff --git a/vllm/entrypoints/serve/disagg/serving.py b/vllm/entrypoints/serve/disagg/serving.py index 79367622c36..14ba85ecf8c 100644 --- a/vllm/entrypoints/serve/disagg/serving.py +++ b/vllm/entrypoints/serve/disagg/serving.py @@ -18,6 +18,7 @@ from vllm.entrypoints.openai.chat_completion.protocol import ( ) from vllm.entrypoints.openai.engine.protocol import ( ErrorResponse, + GenerationError, PromptTokenUsageInfo, RequestResponseMetadata, UsageInfo, @@ -28,12 +29,15 @@ from vllm.entrypoints.serve.disagg.protocol import ( GenerateRequest, GenerateResponse, GenerateResponseChoice, + GenerateResponseStreamChoice, + GenerateStreamResponse, ) from vllm.entrypoints.serve.render.serving import OpenAIServingRender +from vllm.entrypoints.utils import should_include_usage from vllm.logger import init_logger from vllm.logprobs import Logprob from vllm.outputs import RequestOutput -from vllm.sampling_params import SamplingParams +from vllm.sampling_params import RequestOutputKind, SamplingParams from vllm.utils.collection_utils import as_list logger = init_logger(__name__) @@ -74,7 +78,7 @@ class ServingTokens(OpenAIServing): self, request: GenerateRequest, raw_request: Request | None = None, - ) -> GenerateResponse | ErrorResponse: + ) -> GenerateResponse | ErrorResponse | AsyncGenerator[str, None]: error_check_ret = await self._check_model(request) if error_check_ret is not None: logger.error("Error with model %s", error_check_ret) @@ -110,6 +114,8 @@ class ServingTokens(OpenAIServing): sampling_params = request.sampling_params if self.force_no_detokenize: sampling_params.detokenize = False + if request.stream: + sampling_params.output_kind = RequestOutputKind.DELTA self._log_inputs( request_id, @@ -133,9 +139,17 @@ class ServingTokens(OpenAIServing): priority=request.priority, ) - # TODO(NickLucche): Implement streaming response - assert result_generator is not None + + if request.stream: + return self.serve_tokens_stream_generator( + request, + result_generator, + request_id, + model_name, + request_metadata, + ) + return await self.serve_tokens_full_generator( request, result_generator, request_id, model_name, request_metadata ) @@ -236,6 +250,109 @@ class ServingTokens(OpenAIServing): return response + async def serve_tokens_stream_generator( + self, + request: GenerateRequest, + result_generator: AsyncGenerator[RequestOutput, None], + request_id: str, + model_name: str, + request_metadata: RequestResponseMetadata, + ) -> AsyncGenerator[str, None]: + num_prompt_tokens = 0 + num_generated_tokens: list[int] = [] + first_iteration = True + num_cached_tokens = None + sampling_params: SamplingParams = request.sampling_params + + include_usage, include_continuous_usage = should_include_usage( + request.stream_options, False + ) + + try: + async for res in result_generator: + if first_iteration: + if res.prompt_token_ids is not None: + num_prompt_tokens = len(res.prompt_token_ids) + if res.encoder_prompt_token_ids is not None: + num_prompt_tokens += len(res.encoder_prompt_token_ids) + num_cached_tokens = res.num_cached_tokens + num_generated_tokens = [0] * len(res.outputs) + first_iteration = False + + for output in res.outputs: + i = output.index + delta_token_ids = output.token_ids + num_generated_tokens[i] += len(delta_token_ids) + + finish_reason = output.finish_reason + self._raise_if_error(finish_reason, request_id) + + if not delta_token_ids: + continue + + if sampling_params.logprobs is not None: + out_logprobs = output.logprobs + assert out_logprobs is not None, "Did not output logprobs" + logprobs = self._create_tokens_logprobs( + token_ids=delta_token_ids, + top_logprobs=out_logprobs, + num_output_top_logprobs=sampling_params.logprobs, + ) + else: + logprobs = None + + chunk = GenerateStreamResponse( + request_id=request_id, + choices=[ + GenerateResponseStreamChoice( + index=i, + logprobs=logprobs, + finish_reason=finish_reason, + token_ids=as_list(delta_token_ids), + ) + ], + ) + if include_continuous_usage: + chunk.usage = UsageInfo( + prompt_tokens=num_prompt_tokens, + completion_tokens=num_generated_tokens[i], + total_tokens=(num_prompt_tokens + num_generated_tokens[i]), + ) + + yield f"data: {chunk.model_dump_json()}\n\n" + + total_completion_tokens = sum(num_generated_tokens) + final_usage_info = UsageInfo( + prompt_tokens=num_prompt_tokens, + completion_tokens=total_completion_tokens, + total_tokens=num_prompt_tokens + total_completion_tokens, + ) + + if self.enable_prompt_tokens_details and num_cached_tokens: + final_usage_info.prompt_tokens_details = PromptTokenUsageInfo( + cached_tokens=num_cached_tokens + ) + + if include_usage: + final_chunk = GenerateStreamResponse( + request_id=request_id, + choices=[], + usage=final_usage_info, + ) + yield f"data: {final_chunk.model_dump_json(exclude_none=True)}\n\n" + + request_metadata.final_usage_info = final_usage_info + + except GenerationError as e: + yield ( + f"data: {self._convert_generation_error_to_streaming_response(e)}\n\n" + ) + except Exception as e: + logger.exception("Error in token generation stream.") + data = self.create_streaming_error_response(e) + yield f"data: {data}\n\n" + yield "data: [DONE]\n\n" + def _create_tokens_logprobs( self, token_ids: GenericSequence[int], diff --git a/vllm/entrypoints/serve/render/serving.py b/vllm/entrypoints/serve/render/serving.py index 83b41bbda2d..43ea5127b3b 100644 --- a/vllm/entrypoints/serve/render/serving.py +++ b/vllm/entrypoints/serve/render/serving.py @@ -451,6 +451,8 @@ class OpenAIServingRender: request: Any, prompt_input: str | list[str] | list[int] | list[list[int]] | None, prompt_embeds: bytes | list[bytes] | None, + *, + skip_mm_cache: bool = False, ) -> list[EngineInput]: """Copied from OpenAIServing._preprocess_completion.""" prompts = list[SingletonPrompt | bytes]() @@ -458,12 +460,14 @@ class OpenAIServingRender: prompts.extend(prompt_to_seq(prompt_embeds)) if prompt_input is not None: prompts.extend(prompt_to_seq(prompt_input)) - return await self.preprocess_cmpl(request, prompts) + return await self.preprocess_cmpl(request, prompts, skip_mm_cache=skip_mm_cache) async def preprocess_cmpl( self, request: Any, prompts: Sequence[PromptType | bytes], + *, + skip_mm_cache: bool = False, ) -> list[EngineInput]: """Copied from OpenAIServing._preprocess_cmpl.""" renderer = self.renderer @@ -487,6 +491,7 @@ class OpenAIServingRender: for k in ("mm_processor_kwargs", "cache_salt") if (v := getattr(request, k, None)) is not None }, + skip_mm_cache=skip_mm_cache, ) async def preprocess_chat( @@ -498,6 +503,8 @@ class OpenAIServingRender: default_template_kwargs: dict[str, Any] | None, tool_dicts: list[dict[str, Any]] | None = None, tool_parser: type[ToolParser] | None = None, + *, + skip_mm_cache: bool = False, ) -> tuple[list[ConversationMessage], list[EngineInput]]: """Copied from OpenAIServing._preprocess_chat.""" renderer = self.renderer @@ -529,6 +536,7 @@ class OpenAIServingRender: for k in ("mm_processor_kwargs", "cache_salt") if (v := getattr(request, k, None)) is not None }, + skip_mm_cache=skip_mm_cache, ) # tool parsing is done only if a tool_parser has been set and if diff --git a/vllm/entrypoints/serve/tokenize/serving.py b/vllm/entrypoints/serve/tokenize/serving.py index 22b852d2778..9b573b69eb8 100644 --- a/vllm/entrypoints/serve/tokenize/serving.py +++ b/vllm/entrypoints/serve/tokenize/serving.py @@ -86,12 +86,14 @@ class OpenAIServingTokenization(OpenAIServing): default_template_content_format=self.chat_template_content_format, default_template_kwargs=self.default_chat_template_kwargs, tool_dicts=tool_dicts, + skip_mm_cache=True, ) else: engine_inputs = await self.openai_serving_render.preprocess_completion( request, prompt_input=request.prompt, prompt_embeds=None, + skip_mm_cache=True, ) input_ids: list[int] = [] diff --git a/vllm/envs.py b/vllm/envs.py index ec8d663141a..c2f8ca8c580 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -51,6 +51,7 @@ if TYPE_CHECKING: VLLM_CPU_OMP_THREADS_BIND: str = "auto" VLLM_CPU_NUM_OF_RESERVED_CPU: int | None = None VLLM_CPU_SGL_KERNEL: bool = False + VLLM_CPU_ATTN_SPLIT_KV: bool = True VLLM_ZENTORCH_WEIGHT_PREPACK: bool = True VLLM_CPU_INT4_W4A8: bool = True VLLM_XLA_CACHE_PATH: str = os.path.join(VLLM_CACHE_ROOT, "xla_cache") @@ -59,6 +60,7 @@ if TYPE_CHECKING: VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE: Literal["auto", "nccl", "shm"] = "auto" VLLM_USE_RAY_COMPILED_DAG_OVERLAP_COMM: bool = False VLLM_USE_RAY_WRAPPED_PP_COMM: bool = True + VLLM_USE_RAY_V2_EXECUTOR_BACKEND: bool = False VLLM_XLA_USE_SPMD: bool = False VLLM_WORKER_MULTIPROC_METHOD: Literal["fork", "spawn"] = "fork" VLLM_ASSETS_CACHE: str = os.path.join(VLLM_CACHE_ROOT, "assets") @@ -189,6 +191,7 @@ if TYPE_CHECKING: VLLM_MQ_MAX_CHUNK_BYTES_MB: int = 16 VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS: int = 300 VLLM_KV_CACHE_LAYOUT: Literal["NHD", "HND"] | None = None + VLLM_SSM_CONV_STATE_LAYOUT: Literal["SD", "DS"] | None = None VLLM_COMPUTE_NANS_IN_LOGITS: bool = False VLLM_USE_NVFP4_CT_EMULATIONS: bool = False VLLM_ROCM_QUICK_REDUCE_QUANTIZATION: Literal[ @@ -725,6 +728,10 @@ environment_variables: dict[str, Callable[[], Any]] = { else None, # (CPU backend only) whether to use SGL kernels, optimized for small batch. "VLLM_CPU_SGL_KERNEL": lambda: bool(int(os.getenv("VLLM_CPU_SGL_KERNEL", "0"))), + # (CPU backend only) whether to enable attention spilt KV. + "VLLM_CPU_ATTN_SPLIT_KV": lambda: bool( + int(os.getenv("VLLM_CPU_ATTN_SPLIT_KV", "1")) + ), # (Zen CPU backend) eagerly prepack weights into ZenDNN blocked layout # at model load time. Eliminates per-inference layout conversion overhead. "VLLM_ZENTORCH_WEIGHT_PREPACK": lambda: bool( @@ -753,6 +760,12 @@ environment_variables: dict[str, Callable[[], Any]] = { "VLLM_USE_RAY_WRAPPED_PP_COMM": lambda: bool( int(os.getenv("VLLM_USE_RAY_WRAPPED_PP_COMM", "1")) ), + # When True and distributed_executor_backend="ray", use RayExecutorV2 + # (MQ-based) instead of RayDistributedExecutor (compiled-graph backend). + # TODO (jeffreywang): Enabled by default in vLLM 0.20.0. + "VLLM_USE_RAY_V2_EXECUTOR_BACKEND": lambda: bool( + int(os.getenv("VLLM_USE_RAY_V2_EXECUTOR_BACKEND", "0")) + ), # Use dedicated multiprocess context for workers. # Both spawn and fork work "VLLM_WORKER_MULTIPROC_METHOD": env_with_choices( @@ -1397,6 +1410,13 @@ environment_variables: dict[str, Callable[[], Any]] = { "VLLM_KV_CACHE_LAYOUT": env_with_choices( "VLLM_KV_CACHE_LAYOUT", None, ["NHD", "HND"] ), + # SSM conv state layout used for Mamba models. + # - SD: (state_len, dim) — dim contiguous (default) + # - DS: (dim, state_len) — TP-sharded dim on dim1, + # consistent with SSM temporal state and HND KV cache layout. + "VLLM_SSM_CONV_STATE_LAYOUT": env_with_choices( + "VLLM_SSM_CONV_STATE_LAYOUT", None, ["SD", "DS"] + ), # Enable checking whether the generated logits contain NaNs, # indicating corrupted output. Useful for debugging low level bugs # or bad hardware but it may add compute overhead. diff --git a/vllm/kernels/helion/utils.py b/vllm/kernels/helion/utils.py index 5ff8046c73c..130d79093b7 100644 --- a/vllm/kernels/helion/utils.py +++ b/vllm/kernels/helion/utils.py @@ -2,11 +2,10 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Utility functions for Helion kernel management.""" -import logging - +from vllm.logger import init_logger from vllm.platforms import current_platform -logger = logging.getLogger(__name__) +logger = init_logger(__name__) # Maps known variant GPU names (after lowercase/underscore normalization) # to their canonical form. @@ -49,7 +48,7 @@ _GPU_NAME_ALIASES: dict[str, str] = { def get_gpu_name(device_id: int | None = None) -> str: if device_id is None: - logger.warning( + logger.warning_once( "get_gpu_name() called without device_id, defaulting to 0. " "This may return the wrong device name in multi-node setups." ) diff --git a/vllm/lora/layers/__init__.py b/vllm/lora/layers/__init__.py index 235f40b7385..1f3fdea2cda 100644 --- a/vllm/lora/layers/__init__.py +++ b/vllm/lora/layers/__init__.py @@ -13,7 +13,6 @@ from vllm.lora.layers.column_parallel_linear import ( QKVParallelLinearWithShardedLoRA, ) from vllm.lora.layers.fused_moe import FusedMoE3DWithLoRA, FusedMoEWithLoRA -from vllm.lora.layers.gate_linear import GateLinearWithLoRA from vllm.lora.layers.logits_processor import LogitsProcessorWithLoRA from vllm.lora.layers.replicated_linear import ReplicatedLinearWithLoRA from vllm.lora.layers.row_parallel_linear import ( @@ -39,7 +38,6 @@ __all__ = [ "RowParallelLinearWithLoRA", "RowParallelLinearWithShardedLoRA", "ReplicatedLinearWithLoRA", - "GateLinearWithLoRA", "LoRAMapping", "LoRAMappingType", "FusedMoEWithLoRA", diff --git a/vllm/lora/layers/gate_linear.py b/vllm/lora/layers/gate_linear.py deleted file mode 100644 index 9bcaaa5b8e2..00000000000 --- a/vllm/lora/layers/gate_linear.py +++ /dev/null @@ -1,30 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import torch.nn as nn -from transformers import PretrainedConfig - -from vllm.config.lora import LoRAConfig -from vllm.model_executor.custom_op import maybe_get_oot_by_class -from vllm.model_executor.layers.fused_moe.router.gate_linear import GateLinear - -from .replicated_linear import ReplicatedLinearWithLoRA - - -class GateLinearWithLoRA(ReplicatedLinearWithLoRA): - def __init__(self, base_layer: GateLinear) -> None: - super().__init__( - base_layer, - ) - - # GateLinearWithLoRA should always be replaced, regardless of the fully - # sharded LoRAs setting, because it is, by definition, copied per GPU. - @classmethod - def can_replace_layer( - cls, - source_layer: nn.Module, - lora_config: LoRAConfig, - packed_modules_list: list, - model_config: PretrainedConfig | None = None, - ) -> bool: - return type(source_layer) is maybe_get_oot_by_class(GateLinear) diff --git a/vllm/lora/utils.py b/vllm/lora/utils.py index 75ed9674af5..2349ace7084 100644 --- a/vllm/lora/utils.py +++ b/vllm/lora/utils.py @@ -21,7 +21,6 @@ from vllm.lora.layers import ( ColumnParallelLinearWithShardedLoRA, FusedMoE3DWithLoRA, FusedMoEWithLoRA, - GateLinearWithLoRA, LogitsProcessorWithLoRA, MergedColumnParallelLinearVariableSliceWithLoRA, MergedColumnParallelLinearWithLoRA, @@ -82,7 +81,6 @@ _all_lora_classes: set[type[BaseLayerWithLoRA]] = { MergedQKVParallelLinearWithLoRA, RowParallelLinearWithLoRA, ReplicatedLinearWithLoRA, - GateLinearWithLoRA, LogitsProcessorWithLoRA, ColumnParallelLinearWithShardedLoRA, QKVParallelLinearWithShardedLoRA, diff --git a/vllm/model_executor/layers/attention/attention.py b/vllm/model_executor/layers/attention/attention.py index 7610030f3ed..3ff4ec62a6b 100644 --- a/vllm/model_executor/layers/attention/attention.py +++ b/vllm/model_executor/layers/attention/attention.py @@ -38,6 +38,7 @@ from vllm.v1.kv_cache_interface import ( FullAttentionSpec, KVCacheSpec, SlidingWindowSpec, + get_kv_quant_mode, ) if TYPE_CHECKING: @@ -381,8 +382,10 @@ class Attention(nn.Module, AttentionLayerBase): # for attn backends supporting query quantization self.query_quant = None - if self.impl.supports_quant_query_input and self.kv_cache_dtype.startswith( - "fp8" + if ( + self.impl.supports_quant_query_input + and self.kv_cache_dtype.startswith("fp8") + and not self.kv_cache_dtype.endswith("per_token_head") ): is_per_head = ( hasattr(self, "q_scale") and self.q_scale.numel() == self.num_kv_heads @@ -539,6 +542,7 @@ class Attention(nn.Module, AttentionLayerBase): block_size = vllm_config.cache_config.block_size # Should not be called for enc-dec or encoder-only attention. assert self.attn_type == AttentionType.DECODER + quant_mode = get_kv_quant_mode(self.kv_cache_dtype) if self.sliding_window is not None: assert not vllm_config.model_config.use_mla, ( "MLA is not supported for slidingwindow" @@ -548,6 +552,7 @@ class Attention(nn.Module, AttentionLayerBase): num_kv_heads=self.num_kv_heads, head_size=self.head_size, dtype=self.kv_cache_torch_dtype, + kv_quant_mode=quant_mode, sliding_window=self.sliding_window, ) else: @@ -557,6 +562,7 @@ class Attention(nn.Module, AttentionLayerBase): head_size=self.head_size, head_size_v=self.head_size_v, dtype=self.kv_cache_torch_dtype, + kv_quant_mode=quant_mode, ) diff --git a/vllm/model_executor/layers/attention/chunked_local_attention.py b/vllm/model_executor/layers/attention/chunked_local_attention.py index b747304acd0..136574d9752 100644 --- a/vllm/model_executor/layers/attention/chunked_local_attention.py +++ b/vllm/model_executor/layers/attention/chunked_local_attention.py @@ -23,6 +23,7 @@ from vllm.v1.kv_cache_interface import ( AttentionSpec, ChunkedLocalAttentionSpec, KVCacheSpec, + get_kv_quant_mode, ) @@ -123,5 +124,6 @@ class ChunkedLocalAttention(Attention): num_kv_heads=self.num_kv_heads, head_size=self.head_size, dtype=self.kv_cache_torch_dtype, + kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype), attention_chunk_size=self.attention_chunk_size, ) diff --git a/vllm/model_executor/layers/attention/cross_attention.py b/vllm/model_executor/layers/attention/cross_attention.py index 5bd8e163f4a..31ac7fa1bd5 100644 --- a/vllm/model_executor/layers/attention/cross_attention.py +++ b/vllm/model_executor/layers/attention/cross_attention.py @@ -18,7 +18,11 @@ from vllm.v1.attention.backend import ( subclass_attention_backend_with_overrides, ) from vllm.v1.attention.selector import get_attn_backend -from vllm.v1.kv_cache_interface import CrossAttentionSpec, KVCacheSpec +from vllm.v1.kv_cache_interface import ( + CrossAttentionSpec, + KVCacheSpec, + get_kv_quant_mode, +) logger = init_logger(__name__) @@ -220,4 +224,5 @@ class CrossAttention(Attention): num_kv_heads=self.num_kv_heads, head_size=self.head_size, dtype=self.kv_cache_torch_dtype, + kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype), ) diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 0be46fbbc5a..699238b4875 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -449,6 +449,11 @@ class MLAAttention(nn.Module, AttentionLayerBase): group_shape=GroupShape.PER_TENSOR, compile_native=True, ) + self._quant_fp8_op = QuantFP8( + static=True, + group_shape=GroupShape.PER_TENSOR, + compile_native=True, + ) @property def chunked_prefill_workspace_size(self) -> int: @@ -545,9 +550,19 @@ class MLAAttention(nn.Module, AttentionLayerBase): ) -> torch.Tensor: assert output is not None, "Output tensor must be provided." - if output_scale is not None or output_block_scale is not None: - raise NotImplementedError( - "fused output quantization is not yet supported for MLA" + use_quant = output_scale is not None or output_block_scale is not None + if use_quant: + # The fusion pass has allocated output with quantized dtype + # (FP8 or uint8 for FP4). We can't write into it directly, + # so we swap in a temp buffer for computation, then quantize + # into the real output at the end. + # NOTE(carlyou): this is temporary until kernels support fp8 output + quant_output = output + output = torch.empty( + output.shape[0], + self.num_heads * self.v_head_dim, + dtype=q.dtype, + device=output.device, ) if attn_metadata is None: @@ -567,6 +582,8 @@ class MLAAttention(nn.Module, AttentionLayerBase): # The zero fill is required when used with DP + EP # to ensure all ranks within a DP group compute the # same expert outputs. + if use_quant: + return quant_output.fill_(0) return output.fill_(0) if self.impl.dcp_world_size == -1: @@ -706,6 +723,21 @@ class MLAAttention(nn.Module, AttentionLayerBase): # v_up projection self._v_up_proj(attn_out, out=mqa_output_slice) + + if use_quant: + # Quantize the BF16 computation result into the quantized output + actual = output[:num_actual_toks] + if output_block_scale is not None: + # NVFP4: two FP4 values packed into one uint8 + fp4_data, fp4_scales = ops.scaled_fp4_quant(actual, output_scale) + quant_output[:num_actual_toks].copy_(fp4_data) + output_block_scale.copy_(fp4_scales) + else: + # Static FP8 quantization + fp8_data, _ = self._quant_fp8_op(actual, output_scale) + quant_output[:num_actual_toks].copy_(fp8_data) + return quant_output + return output_padded def process_weights_after_loading(self, act_dtype: torch.dtype): @@ -2069,6 +2101,14 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]): understand this class """ + def fused_output_quant_supported(self, quant_key): + from vllm.model_executor.layers.quantization.utils.quant_utils import ( + kFp8StaticTensorSym, + kNvfp4Dynamic, + ) + + return quant_key in (kFp8StaticTensorSym, kNvfp4Dynamic) + def __init__( self, num_heads: int, @@ -2513,8 +2553,12 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]): if hasattr(self.kv_b_proj, "weight") else self.kv_b_proj.params_dtype ) - if use_fp8_prefill or _kv_b_proj_w_dtype != current_platform.fp8_dtype(): - kv_c_normed = kv_c_normed.to(_kv_b_proj_w_dtype) + # For NVFP4, weights are packed uint8 — keep input in model dtype + # since the NVFP4 linear layer quantizes internally. + if ( + use_fp8_prefill or _kv_b_proj_w_dtype != current_platform.fp8_dtype() + ) and _kv_b_proj_w_dtype != torch.uint8: + kv_c_normed = kv_c_normed.to(self.kv_b_proj.weight.dtype) k_pe = workspace[:toks][..., self.kv_lora_rank :].unsqueeze(1) kv_nope = self.kv_b_proj(kv_c_normed)[0].view( diff --git a/vllm/model_executor/layers/attention/static_sink_attention.py b/vllm/model_executor/layers/attention/static_sink_attention.py index 913d73a16c2..263d873218f 100644 --- a/vllm/model_executor/layers/attention/static_sink_attention.py +++ b/vllm/model_executor/layers/attention/static_sink_attention.py @@ -26,6 +26,7 @@ from vllm.v1.kv_cache_interface import ( AttentionSpec, KVCacheSpec, SinkFullAttentionSpec, + get_kv_quant_mode, ) logger = init_logger(__name__) @@ -217,6 +218,7 @@ class StaticSinkAttention(Attention, CustomOp): head_size_v=self.head_size_v, sink_len=self.sink_len, dtype=self.kv_cache_torch_dtype, + kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype), ) diff --git a/vllm/model_executor/layers/batch_invariant.py b/vllm/model_executor/layers/batch_invariant.py index 2f945024400..98fd6be8f6a 100644 --- a/vllm/model_executor/layers/batch_invariant.py +++ b/vllm/model_executor/layers/batch_invariant.py @@ -10,6 +10,7 @@ import vllm.envs as envs from vllm.logger import init_logger from vllm.platforms import current_platform from vllm.triton_utils import tl, triton +from vllm.utils.mem_utils import get_max_shared_memory_bytes from vllm.utils.platform_utils import num_compute_units from vllm.utils.torch_utils import is_torch_equal_or_newer from vllm.v1.attention.backends.registry import AttentionBackendEnum @@ -177,7 +178,7 @@ def matmul_persistent( }, torch.float16: { "BLOCK_SIZE_M": 128, - "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_N": _fp16_block_size_n, "BLOCK_SIZE_K": 64, "GROUP_SIZE_M": 8, "num_stages": 3, @@ -700,7 +701,7 @@ def bmm_batch_invariant(a, b, *, out=None): }, torch.float16: { "BLOCK_SIZE_M": 128, - "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_N": _fp16_block_size_n, "BLOCK_SIZE_K": 64, "num_stages": 3, "num_warps": 8, @@ -752,7 +753,8 @@ def addmm_batch_invariant(bias, a, b): def _log_softmax_batch_invariant(input, dim, _half_to_float): - assert not _half_to_float, "not implemented" + if _half_to_float: + return log_softmax(input.float(), dim=dim) return log_softmax(input, dim=dim) @@ -923,29 +925,34 @@ _original_fp16_reduction_precision = None _original_bf16_reduction_precision = None _original_cublas_workspace_cfg = None _original_cublaslt_workspace_size = None +_fp16_block_size_n = 256 def enable_batch_invariant_mode(): global _batch_invariant_MODE, _batch_invariant_LIB, _original_torch_bmm global _original_fp16_reduction_precision, _original_bf16_reduction_precision global _original_cublas_workspace_cfg, _original_cublaslt_workspace_size + global _fp16_block_size_n + if _batch_invariant_MODE: return _batch_invariant_MODE = True _batch_invariant_LIB = torch.library.Library("aten", "IMPL") - if ( - current_platform.is_device_capability_family(100) - or current_platform.is_device_capability(80) - or current_platform.is_device_capability(89) - ): + if current_platform.is_device_capability_family( + 100 + ) or current_platform.is_device_capability_family(80): # For PyTorch 2.9, B200 uses GEMV for bs=1 # Requires https://github.com/pytorch/pytorch/pull/166735 _batch_invariant_LIB.impl("aten::mm", mm_batch_invariant, "CUDA") _batch_invariant_LIB.impl("aten::addmm", addmm_batch_invariant, "CUDA") _batch_invariant_LIB.impl("aten::matmul", matmul_batch_invariant, "CUDA") _batch_invariant_LIB.impl("aten::linear", linear_batch_invariant, "CUDA") + + # Query the shared memory size and set block size + # accordingly to avoid triton OutOfResources + _fp16_block_size_n = 256 if get_max_shared_memory_bytes() > 106496 else 128 else: # Only source of batch invariance for Hopper is split-k, can disable through # cuBLAS workspace config diff --git a/vllm/model_executor/layers/fla/ops/__init__.py b/vllm/model_executor/layers/fla/ops/__init__.py index e52387a20b4..1942d8980bc 100644 --- a/vllm/model_executor/layers/fla/ops/__init__.py +++ b/vllm/model_executor/layers/fla/ops/__init__.py @@ -7,6 +7,7 @@ # the following copyright notice: # Copyright (c) 2023-2025, Songlin Yang, Yu Zhang from .chunk import chunk_gated_delta_rule +from .fused_gdn_prefill_post_conv import fused_post_conv_prep from .fused_recurrent import ( fused_recurrent_gated_delta_rule, fused_recurrent_gated_delta_rule_packed_decode, @@ -19,5 +20,6 @@ __all__ = [ "chunk_gated_delta_rule", "fused_recurrent_gated_delta_rule", "fused_recurrent_gated_delta_rule_packed_decode", + "fused_post_conv_prep", "fused_sigmoid_gating_delta_rule_update", ] diff --git a/vllm/model_executor/layers/fla/ops/chunk.py b/vllm/model_executor/layers/fla/ops/chunk.py index 73cba7f9035..02e48921d41 100644 --- a/vllm/model_executor/layers/fla/ops/chunk.py +++ b/vllm/model_executor/layers/fla/ops/chunk.py @@ -16,7 +16,7 @@ from .chunk_scaled_dot_kkt import chunk_scaled_dot_kkt_fwd from .cumsum import chunk_local_cumsum from .l2norm import l2norm_fwd from .solve_tril import solve_tril -from .utils import SUPPRESS_LEVEL, input_guard +from .utils import FLA_CHUNK_SIZE, SUPPRESS_LEVEL, input_guard from .wy_fast import recompute_w_u_fwd @@ -30,13 +30,24 @@ def chunk_gated_delta_rule_fwd( initial_state: torch.Tensor, output_final_state: bool, cu_seqlens: torch.Tensor | None = None, + chunk_indices: torch.Tensor | None = None, + chunk_offsets: torch.Tensor | None = None, ): - g = chunk_local_cumsum(g, chunk_size=64, cu_seqlens=cu_seqlens) + g = chunk_local_cumsum( + g, chunk_size=FLA_CHUNK_SIZE, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices + ) # obtain WY representation. u is actually the new v. A = chunk_scaled_dot_kkt_fwd( - k=k, beta=beta, g=g, cu_seqlens=cu_seqlens, output_dtype=torch.float32 + k=k, + beta=beta, + g=g, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + output_dtype=torch.float32, + ) + A = solve_tril( + A=A, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices, output_dtype=k.dtype ) - A = solve_tril(A=A, cu_seqlens=cu_seqlens, output_dtype=k.dtype) w, u = recompute_w_u_fwd( k=k, v=v, @@ -44,6 +55,7 @@ def chunk_gated_delta_rule_fwd( A=A, g_cumsum=g, cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, ) h, v_new, final_state = chunk_gated_delta_rule_fwd_h( k=k, @@ -53,6 +65,8 @@ def chunk_gated_delta_rule_fwd( initial_state=initial_state, output_final_state=output_final_state, cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_offsets=chunk_offsets, ) o = chunk_fwd_o( q=q, @@ -62,6 +76,7 @@ def chunk_gated_delta_rule_fwd( g=g, scale=scale, cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, ) if SUPPRESS_LEVEL < 3: return g, o, A, final_state, None, None, None @@ -84,6 +99,8 @@ class ChunkGatedDeltaRuleFunction(torch.autograd.Function): initial_state: torch.Tensor, output_final_state: bool, cu_seqlens: torch.Tensor | None = None, + chunk_indices: torch.Tensor | None = None, + chunk_offsets: torch.Tensor | None = None, use_qk_l2norm_in_kernel: bool = False, ): if use_qk_l2norm_in_kernel: @@ -100,6 +117,8 @@ class ChunkGatedDeltaRuleFunction(torch.autograd.Function): initial_state=initial_state, output_final_state=output_final_state, cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_offsets=chunk_offsets, ) ctx.scale = scale ctx.use_qk_l2norm_in_kernel = use_qk_l2norm_in_kernel @@ -117,6 +136,8 @@ def chunk_gated_delta_rule( initial_state: torch.Tensor = None, output_final_state: bool = False, cu_seqlens: torch.Tensor | None = None, + chunk_indices: torch.Tensor | None = None, + chunk_offsets: torch.Tensor | None = None, use_qk_l2norm_in_kernel: bool = False, ): r""" @@ -206,6 +227,8 @@ def chunk_gated_delta_rule( initial_state, output_final_state, cu_seqlens, + chunk_indices, + chunk_offsets, use_qk_l2norm_in_kernel, ) return o, final_state diff --git a/vllm/model_executor/layers/fla/ops/chunk_delta_h.py b/vllm/model_executor/layers/fla/ops/chunk_delta_h.py index ce60ca46f6c..574f6f25173 100644 --- a/vllm/model_executor/layers/fla/ops/chunk_delta_h.py +++ b/vllm/model_executor/layers/fla/ops/chunk_delta_h.py @@ -14,7 +14,7 @@ from vllm.triton_utils import tl, triton from .index import prepare_chunk_indices, prepare_chunk_offsets from .op import exp -from .utils import use_cuda_graph +from .utils import FLA_CHUNK_SIZE, use_cuda_graph NUM_WARPS = [2, 4, 8, 16] @@ -286,9 +286,11 @@ def chunk_gated_delta_rule_fwd_h( gk: torch.Tensor | None = None, initial_state: torch.Tensor | None = None, output_final_state: bool = False, - chunk_size: int = 64, # SY: remove this argument and force chunk size 64? + chunk_size: int = FLA_CHUNK_SIZE, save_new_value: bool = True, cu_seqlens: torch.Tensor | None = None, + chunk_indices: torch.Tensor | None = None, + chunk_offsets: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: # This kernel is slightly different from fla to support Q/K with different head numbers. # In fla, Q/K always have the same head number, so Hg is always equal to H. @@ -296,20 +298,15 @@ def chunk_gated_delta_rule_fwd_h( H = u.shape[-2] BT = chunk_size - chunk_indices = ( - prepare_chunk_indices(cu_seqlens, chunk_size) - if cu_seqlens is not None - else None - ) + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) # N: the actual number of sequences in the batch with either equal or variable lengths if cu_seqlens is None: N, NT, chunk_offsets = B, triton.cdiv(T, BT), None else: - N, NT, chunk_offsets = ( - len(cu_seqlens) - 1, - len(chunk_indices), - prepare_chunk_offsets(cu_seqlens, BT), - ) + N, NT = len(cu_seqlens) - 1, len(chunk_indices) + if chunk_offsets is None: + chunk_offsets = prepare_chunk_offsets(cu_seqlens, BT) assert K <= 256, "current kernel does not support head dimension larger than 256." h = k.new_empty(B, NT, H, V, K) diff --git a/vllm/model_executor/layers/fla/ops/chunk_o.py b/vllm/model_executor/layers/fla/ops/chunk_o.py index aab1ee006d4..d812ec43372 100644 --- a/vllm/model_executor/layers/fla/ops/chunk_o.py +++ b/vllm/model_executor/layers/fla/ops/chunk_o.py @@ -146,14 +146,14 @@ def chunk_fwd_o( g: torch.Tensor | None = None, # cumsum of log decay scale: float | None = None, cu_seqlens: torch.Tensor | None = None, + chunk_indices: torch.Tensor | None = None, chunk_size: int = FLA_CHUNK_SIZE, ) -> torch.Tensor: B, T, Hg, K, V = *q.shape, v.shape[-1] H = v.shape[-2] BT = chunk_size - chunk_indices = ( - prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None - ) + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) if scale is None: scale = k.shape[-1] ** -0.5 diff --git a/vllm/model_executor/layers/fla/ops/chunk_scaled_dot_kkt.py b/vllm/model_executor/layers/fla/ops/chunk_scaled_dot_kkt.py index 31bd489ebd8..3f7628487d6 100644 --- a/vllm/model_executor/layers/fla/ops/chunk_scaled_dot_kkt.py +++ b/vllm/model_executor/layers/fla/ops/chunk_scaled_dot_kkt.py @@ -14,6 +14,7 @@ from vllm.triton_utils import tl, triton from .index import prepare_chunk_indices from .op import exp +from .utils import FLA_CHUNK_SIZE @triton.heuristics( @@ -103,7 +104,8 @@ def chunk_scaled_dot_kkt_fwd( g: torch.Tensor | None = None, beta: torch.Tensor | None = None, cu_seqlens: torch.Tensor | None = None, - chunk_size: int = 64, + chunk_indices: torch.Tensor | None = None, + chunk_size: int = FLA_CHUNK_SIZE, output_dtype: torch.dtype = torch.float32, ) -> torch.Tensor: r""" @@ -119,6 +121,9 @@ def chunk_scaled_dot_kkt_fwd( cu_seqlens (torch.Tensor): The cumulative sequence lengths of the input tensor. Default: None + chunk_indices (torch.Tensor): + Pre-computed chunk indices. If None and cu_seqlens is provided, + computed internally. Default: None chunk_size (int): The chunk size. Default: 64. output_dtype (torch.dtype): @@ -132,9 +137,8 @@ def chunk_scaled_dot_kkt_fwd( B, T, Hg, K = k.shape H = beta.shape[-1] BT = chunk_size - chunk_indices = ( - prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None - ) + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) A = torch.empty(B, T, H, BT, device=k.device, dtype=output_dtype) diff --git a/vllm/model_executor/layers/fla/ops/cumsum.py b/vllm/model_executor/layers/fla/ops/cumsum.py index 13238020cbd..b0820104b1a 100644 --- a/vllm/model_executor/layers/fla/ops/cumsum.py +++ b/vllm/model_executor/layers/fla/ops/cumsum.py @@ -162,6 +162,7 @@ def chunk_local_cumsum_scalar( chunk_size: int, reverse: bool = False, cu_seqlens: torch.Tensor | None = None, + chunk_indices: torch.Tensor | None = None, head_first: bool = False, output_dtype: torch.dtype | None = torch.float, ) -> torch.Tensor: @@ -172,10 +173,9 @@ def chunk_local_cumsum_scalar( assert chunk_size == 2 ** (chunk_size.bit_length() - 1), ( "chunk_size must be a power of 2" ) + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) BT = chunk_size - chunk_indices = ( - prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None - ) NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) g_org, g = g, torch.empty_like(g, dtype=output_dtype or g.dtype) grid = (NT, B * H) @@ -199,6 +199,7 @@ def chunk_local_cumsum_vector( chunk_size: int, reverse: bool = False, cu_seqlens: torch.Tensor | None = None, + chunk_indices: torch.Tensor | None = None, head_first: bool = False, output_dtype: torch.dtype | None = torch.float, ) -> torch.Tensor: @@ -206,16 +207,13 @@ def chunk_local_cumsum_vector( B, H, T, S = g.shape else: B, T, H, S = g.shape - BT = chunk_size - chunk_indices = ( - prepare_chunk_indices(cu_seqlens, chunk_size) - if cu_seqlens is not None - else None - ) - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) assert chunk_size == 2 ** (chunk_size.bit_length() - 1), ( "chunk_size must be a power of 2" ) + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) + BT = chunk_size + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) g_org, g = g, torch.empty_like(g, dtype=output_dtype or g.dtype) @@ -247,6 +245,7 @@ def chunk_local_cumsum( chunk_size: int, reverse: bool = False, cu_seqlens: torch.Tensor | None = None, + chunk_indices: torch.Tensor | None = None, head_first: bool = False, output_dtype: torch.dtype | None = torch.float, **kwargs, @@ -257,11 +256,23 @@ def chunk_local_cumsum( ) if len(g.shape) == 3: return chunk_local_cumsum_scalar( - g, chunk_size, reverse, cu_seqlens, head_first, output_dtype + g, + chunk_size, + reverse, + cu_seqlens, + chunk_indices, + head_first, + output_dtype, ) elif len(g.shape) == 4: return chunk_local_cumsum_vector( - g, chunk_size, reverse, cu_seqlens, head_first, output_dtype + g, + chunk_size, + reverse, + cu_seqlens, + chunk_indices, + head_first, + output_dtype, ) else: raise ValueError( diff --git a/vllm/model_executor/layers/fla/ops/fused_gdn_prefill_post_conv.py b/vllm/model_executor/layers/fla/ops/fused_gdn_prefill_post_conv.py new file mode 100644 index 00000000000..4807c78e7b1 --- /dev/null +++ b/vllm/model_executor/layers/fla/ops/fused_gdn_prefill_post_conv.py @@ -0,0 +1,248 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Fused post-conv1d preparation for GDN prefill. + +Replaces the chain: + split → rearrange → contiguous * 3 → l2norm * 2 → gating +with a **single Triton kernel** that reads the conv'd mixed_qkv output +and writes directly to q/k/v/g/beta in the target contiguous layout. + +""" + +from __future__ import annotations + +import torch + +from vllm.triton_utils import tl, triton + + +@triton.jit +def _fused_post_conv_kernel( + # ---- inputs ---- + mixed_qkv_ptr, # [L, qkv_dim] conv'd output (contiguous) + a_ptr, # [L, HV] + b_ptr, # [L, HV] + # ---- params ---- + A_log_ptr, # [HV] + dt_bias_ptr, # [HV] + # ---- outputs ---- + q_ptr, # [L, H, K] contiguous + k_ptr, # [L, H, K] contiguous + v_ptr, # [L, HV, V] contiguous + g_ptr, # [L, HV] float32 + beta_ptr, # [L, HV] float32 + # ---- strides ---- + stride_x_tok, # qkv_dim + stride_a_tok, # HV + stride_b_tok, # HV + stride_q_tok, # H * K + stride_k_tok, # H * K + stride_v_tok, # HV * V + # ---- dims ---- + L, + H: tl.constexpr, + HV: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + APPLY_L2NORM: tl.constexpr, + L2NORM_EPS: tl.constexpr, + OUTPUT_G_EXP: tl.constexpr, + SOFTPLUS_THRESHOLD: tl.constexpr, + BLOCK_T: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, +): + """Single fused kernel for post-conv1d preparation. + + Grid: (ceil(L, BLOCK_T), H + HV) + - program_id(1) in [0, H): Q/K head processing + l2norm + - program_id(1) in [H, H+HV): V head processing + gating + """ + i_tb = tl.program_id(0) + i_head = tl.program_id(1) + + HK: tl.constexpr = H * K + + offs_t = i_tb * BLOCK_T + tl.arange(0, BLOCK_T) # [BLOCK_T] + mask_t = offs_t < L + + if i_head < H: + # ============ Q/K head processing ============ + i_h = i_head + offs_k = tl.arange(0, BK) # [BK] + mask_k = offs_k < K + mask_2d = mask_t[:, None] & mask_k[None, :] # [BLOCK_T, BK] + + # Load Q features: mixed_qkv[t, i_h*K + k] + q_offsets = offs_t[:, None] * stride_x_tok + i_h * K + offs_k[None, :] + q_f32 = tl.load(mixed_qkv_ptr + q_offsets, mask=mask_2d, other=0).to(tl.float32) + + # Load K features: mixed_qkv[t, HK + i_h*K + k] + k_offsets = offs_t[:, None] * stride_x_tok + HK + i_h * K + offs_k[None, :] + k_f32 = tl.load(mixed_qkv_ptr + k_offsets, mask=mask_2d, other=0).to(tl.float32) + + if APPLY_L2NORM: + q_sq_sum = tl.sum(q_f32 * q_f32, axis=1) # [BLOCK_T] + q_inv = 1.0 / tl.sqrt(q_sq_sum + L2NORM_EPS) + q_f32 = q_f32 * q_inv[:, None] + + k_sq_sum = tl.sum(k_f32 * k_f32, axis=1) + k_inv = 1.0 / tl.sqrt(k_sq_sum + L2NORM_EPS) + k_f32 = k_f32 * k_inv[:, None] + + # Store Q + q_out = offs_t[:, None] * stride_q_tok + i_h * K + offs_k[None, :] + tl.store( + q_ptr + q_out, + q_f32.to(q_ptr.dtype.element_ty), + mask=mask_2d, + ) + + # Store K + k_out = offs_t[:, None] * stride_k_tok + i_h * K + offs_k[None, :] + tl.store( + k_ptr + k_out, + k_f32.to(k_ptr.dtype.element_ty), + mask=mask_2d, + ) + else: + # ============ V head + gating processing ============ + i_hv = i_head - H + offs_v = tl.arange(0, BV) # [BV] + mask_v = offs_v < V + mask_2d = mask_t[:, None] & mask_v[None, :] # [BLOCK_T, BV] + + V_OFFSET: tl.constexpr = 2 * H * K + + # Load V features: mixed_qkv[t, 2*H*K + i_hv*V + v] + v_offsets = ( + offs_t[:, None] * stride_x_tok + V_OFFSET + i_hv * V + offs_v[None, :] + ) + v_vals = tl.load(mixed_qkv_ptr + v_offsets, mask=mask_2d, other=0) + + # Store V + v_out = offs_t[:, None] * stride_v_tok + i_hv * V + offs_v[None, :] + tl.store(v_ptr + v_out, v_vals, mask=mask_2d) + + # Gating: one scalar per (token, v-head) + A_log_val = tl.load(A_log_ptr + i_hv).to(tl.float32) + dt_bias_val = tl.load(dt_bias_ptr + i_hv).to(tl.float32) + + a_offsets = offs_t * stride_a_tok + i_hv + b_offsets = offs_t * stride_b_tok + i_hv + a_vals = tl.load(a_ptr + a_offsets, mask=mask_t, other=0).to(tl.float32) + b_vals = tl.load(b_ptr + b_offsets, mask=mask_t, other=0).to(tl.float32) + + # g = -exp(A_log) * softplus(a + dt_bias) + x = a_vals + dt_bias_val + sp = tl.where(x > 0, x + tl.log(1.0 + tl.exp(-x)), tl.log(1.0 + tl.exp(x))) + sp = tl.where(x <= SOFTPLUS_THRESHOLD, sp, x) + g_vals = -tl.exp(A_log_val) * sp + + if OUTPUT_G_EXP: + g_vals = tl.exp(g_vals) + + beta_vals = tl.sigmoid(b_vals) + + gb_offsets = offs_t * HV + i_hv + tl.store(g_ptr + gb_offsets, g_vals, mask=mask_t) + tl.store(beta_ptr + gb_offsets, beta_vals, mask=mask_t) + + +def fused_post_conv_prep( + conv_output: torch.Tensor, # [L, qkv_dim] conv'd mixed_qkv + a: torch.Tensor, # [L, HV] + b: torch.Tensor, # [L, HV] + A_log: torch.Tensor, # [HV] + dt_bias: torch.Tensor, # [HV] + num_k_heads: int, + head_k_dim: int, + head_v_dim: int, + apply_l2norm: bool = True, + output_g_exp: bool = False, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Fused post-conv1d prep: split + l2norm + gating in one kernel. + + Args: + conv_output: [L, qkv_dim] contiguous conv'd mixed_qkv + a: [L, HV] gating input + b: [L, HV] gating input + A_log: [HV] log decay parameter + dt_bias: [HV] dt bias parameter + num_k_heads: number of K heads (H) + head_k_dim: dimension per K head (K) + head_v_dim: dimension per V head (V) + apply_l2norm: whether to L2-normalize q and k + output_g_exp: if True, output exp(g) instead of g (for FlashInfer) + + Returns: + q: [L, H, K] contiguous, optionally l2-normalized + k: [L, H, K] contiguous, optionally l2-normalized + v: [L, HV, V] contiguous + g: [L, HV] float32 + beta: [L, HV] float32 + """ + L = conv_output.shape[0] + qkv_dim = conv_output.shape[1] + H = num_k_heads + K = head_k_dim + V = head_v_dim + HV = A_log.shape[0] + dtype = conv_output.dtype + device = conv_output.device + + assert qkv_dim == 2 * H * K + HV * V, ( + f"qkv_dim={qkv_dim} != 2*H*K + HV*V = {2 * H * K + HV * V}" + ) + + # Allocate outputs in target contiguous layout + q = torch.empty(L, H, K, dtype=dtype, device=device) + k = torch.empty(L, H, K, dtype=dtype, device=device) + v = torch.empty(L, HV, V, dtype=dtype, device=device) + g = torch.empty(L, HV, dtype=torch.float32, device=device) + beta = torch.empty(L, HV, dtype=torch.float32, device=device) + + if L == 0: + return q, k, v, g, beta + + # ---- Kernel config ---- + BK = triton.next_power_of_2(K) + BV = triton.next_power_of_2(V) + BLOCK_T = 16 # tokens per block + + # Single kernel: blocks [0,H) do Q/K, blocks [H, H+HV) do V+gating + grid = (triton.cdiv(L, BLOCK_T), H + HV) + _fused_post_conv_kernel[grid]( + mixed_qkv_ptr=conv_output, + a_ptr=a, + b_ptr=b, + A_log_ptr=A_log, + dt_bias_ptr=dt_bias, + q_ptr=q, + k_ptr=k, + v_ptr=v, + g_ptr=g, + beta_ptr=beta, + stride_x_tok=conv_output.stride(0), + stride_a_tok=a.stride(0), + stride_b_tok=b.stride(0), + stride_q_tok=q.stride(0), + stride_k_tok=k.stride(0), + stride_v_tok=v.stride(0), + L=L, + H=H, + HV=HV, + K=K, + V=V, + APPLY_L2NORM=apply_l2norm, + L2NORM_EPS=1e-6, + OUTPUT_G_EXP=output_g_exp, + SOFTPLUS_THRESHOLD=20.0, + BLOCK_T=BLOCK_T, + BK=BK, + BV=BV, + num_warps=4, + num_stages=2, + ) + + return q, k, v, g, beta diff --git a/vllm/model_executor/layers/fla/ops/kda.py b/vllm/model_executor/layers/fla/ops/kda.py index b8c07d1dc89..67cd0231d6e 100644 --- a/vllm/model_executor/layers/fla/ops/kda.py +++ b/vllm/model_executor/layers/fla/ops/kda.py @@ -23,7 +23,7 @@ from .index import prepare_chunk_indices from .l2norm import l2norm_fwd from .op import exp, log from .solve_tril import solve_tril -from .utils import is_amd +from .utils import FLA_CHUNK_SIZE, is_amd BT_LIST_AUTOTUNE = [32, 64, 128] NUM_WARPS_AUTOTUNE = [2, 4, 8, 16] if is_amd else [4, 8, 16, 32] @@ -721,7 +721,7 @@ def chunk_kda_scaled_dot_kkt_fwd( beta: torch.Tensor | None = None, scale: float | None = None, cu_seqlens: torch.Tensor | None = None, - chunk_size: int = 64, + chunk_size: int = FLA_CHUNK_SIZE, output_dtype: torch.dtype = torch.float32, ) -> tuple[torch.Tensor, torch.Tensor]: r""" @@ -1178,7 +1178,7 @@ def chunk_kda_fwd( output_final_state: bool, cu_seqlens: torch.Tensor | None = None, ): - chunk_size = 64 + chunk_size = FLA_CHUNK_SIZE g = chunk_local_cumsum(g, chunk_size=chunk_size, cu_seqlens=cu_seqlens) # the intra Aqk is kept in fp32 # the computation has very marginal effect on the entire throughput @@ -1189,6 +1189,7 @@ def chunk_kda_fwd( beta=beta, scale=scale, cu_seqlens=cu_seqlens, + chunk_size=chunk_size, output_dtype=torch.float32, ) A = solve_tril(A=A, cu_seqlens=cu_seqlens, output_dtype=k.dtype) diff --git a/vllm/model_executor/layers/fla/ops/solve_tril.py b/vllm/model_executor/layers/fla/ops/solve_tril.py index da85aab1920..8d3811ca4c1 100644 --- a/vllm/model_executor/layers/fla/ops/solve_tril.py +++ b/vllm/model_executor/layers/fla/ops/solve_tril.py @@ -507,6 +507,7 @@ def merge_16x16_to_64x64_inverse_kernel( def solve_tril( A: torch.Tensor, cu_seqlens: torch.Tensor | None = None, + chunk_indices: torch.Tensor | None = None, output_dtype: torch.dtype = torch.float, ) -> torch.Tensor: """ @@ -518,6 +519,8 @@ def solve_tril( [B, T, H, BT], where BT should only be 16, 32, or 64. cu_seqlens (torch.Tensor): The cumulative sequence lengths of the input tensor. Default: `None`. + chunk_indices (torch.Tensor): + Pre-computed chunk indices. Default: `None`. output_dtype (torch.dtype): The dtype of the output tensor. Default: `torch.float`. If `None`, the output dtype will be the same as the input dtype. @@ -529,9 +532,8 @@ def solve_tril( output_dtype = A.dtype if output_dtype is None else output_dtype B, T, H, BT = A.shape - chunk_indices = ( - prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None - ) + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) NT = len(chunk_indices) if cu_seqlens is not None else triton.cdiv(T, BT) Ai = torch.zeros_like(A, dtype=output_dtype) diff --git a/vllm/model_executor/layers/fla/ops/wy_fast.py b/vllm/model_executor/layers/fla/ops/wy_fast.py index 6baa08ab499..52d2b28195a 100644 --- a/vllm/model_executor/layers/fla/ops/wy_fast.py +++ b/vllm/model_executor/layers/fla/ops/wy_fast.py @@ -123,14 +123,14 @@ def recompute_w_u_fwd( g_cumsum: torch.Tensor, A: torch.Tensor, cu_seqlens: torch.Tensor | None, + chunk_indices: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: B, T, Hg, K, V = *k.shape, v.shape[-1] H = v.shape[-2] BT = A.shape[-1] - chunk_indices = ( - prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None - ) + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) BK = 64 BV = 64 diff --git a/vllm/model_executor/layers/fused_moe/configs/E=128,N=512,device_name=AMD_Radeon_R9700,dtype=fp8_w8a8,block_shape=[128,128].json b/vllm/model_executor/layers/fused_moe/configs/E=128,N=512,device_name=AMD_Radeon_R9700,dtype=fp8_w8a8,block_shape=[128,128].json new file mode 100644 index 00000000000..d7e503b3615 --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/configs/E=128,N=512,device_name=AMD_Radeon_R9700,dtype=fp8_w8a8,block_shape=[128,128].json @@ -0,0 +1,165 @@ +{ + "triton_version": "3.6.0", + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 1 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 2 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 1 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 4, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 4 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 8, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 4 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 1 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 2 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 2 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 4 + }, + "256": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 2 + }, + "512": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 4 + }, + "1024": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 2 + }, + "1536": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 2 + }, + "2048": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "3072": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4096": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + } +} diff --git a/vllm/model_executor/layers/fused_moe/configs/E=64,N=768,device_name=AMD_Radeon_R9700,dtype=fp8_w8a8,block_shape=[128,128].json b/vllm/model_executor/layers/fused_moe/configs/E=64,N=768,device_name=AMD_Radeon_R9700,dtype=fp8_w8a8,block_shape=[128,128].json new file mode 100644 index 00000000000..a5541722d4c --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/configs/E=64,N=768,device_name=AMD_Radeon_R9700,dtype=fp8_w8a8,block_shape=[128,128].json @@ -0,0 +1,165 @@ +{ + "triton_version": "3.6.0", + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 4 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 1 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 8, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 0 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 4 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 4 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 1 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 2 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 4 + }, + "128": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 2 + }, + "256": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 2 + }, + "512": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 2 + }, + "1024": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 1 + }, + "1536": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "2048": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 1 + }, + "3072": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + }, + "4096": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 4, + "num_warps": 4, + "num_stages": 2, + "waves_per_eu": 0 + } +} diff --git a/vllm/model_executor/layers/fused_moe/cpu_fused_moe.py b/vllm/model_executor/layers/fused_moe/cpu_fused_moe.py index 72e9db514a8..e1bedd6f45b 100644 --- a/vllm/model_executor/layers/fused_moe/cpu_fused_moe.py +++ b/vllm/model_executor/layers/fused_moe/cpu_fused_moe.py @@ -34,12 +34,20 @@ def _swigluoai_forward_native( return gated_output +def _gelu_and_mul( + x: torch.Tensor, +) -> torch.Tensor: + d = x.shape[-1] // 2 + return F.gelu(x[..., :d], approximate="none") * x[..., d:] + + # Map activation names to their native forward functions. # Uses static methods or standalone functions to avoid instantiating CustomOp # classes, which would call get_current_vllm_config() before config is set. _CPU_MOE_ACT_FN: dict[MoEActivation, Callable[[torch.Tensor], torch.Tensor]] = { MoEActivation.SILU: SiluAndMul.forward_native, MoEActivation.SWIGLUOAI: _swigluoai_forward_native, + MoEActivation.GELU: _gelu_and_mul, } diff --git a/vllm/model_executor/layers/fused_moe/fused_moe.py b/vllm/model_executor/layers/fused_moe/fused_moe.py index 370059f729e..6de25da051a 100644 --- a/vllm/model_executor/layers/fused_moe/fused_moe.py +++ b/vllm/model_executor/layers/fused_moe/fused_moe.py @@ -1923,14 +1923,17 @@ class TritonExperts(mk.FusedMoEExpertsModular): ) -> bool: p = current_platform if p.is_rocm(): - from vllm.platforms.rocm import on_gfx9 + from vllm.platforms.rocm import on_gfx9, on_gfx12x is_rocm_on_gfx9 = on_gfx9() + is_rocm_on_gfx12x = on_gfx12x() else: is_rocm_on_gfx9 = False + is_rocm_on_gfx12x = False device_supports_fp8 = ( is_rocm_on_gfx9 + or is_rocm_on_gfx12x or (p.is_cuda() and p.has_device_capability((8, 9))) or p.is_xpu() ) diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py index 7249b425fcb..d4a0817e0be 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py @@ -54,8 +54,8 @@ class Mxfp4MoeBackend(Enum): # Marlin BATCHED_MARLIN = "BATCHED_MARLIN" MARLIN = "MARLIN" - # ROCm AITER (CK) - CK = "CK" + # ROCm AITER + AITER = "AITER" # Triton TRITON = "TRITON" TRITON_UNFUSED = "TRITON_UNFUSED" @@ -130,7 +130,7 @@ def backend_to_kernel_cls( return [BatchedMarlinExperts] - elif backend == Mxfp4MoeBackend.CK: + elif backend == Mxfp4MoeBackend.AITER: from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import ( AiterExperts, ) @@ -155,7 +155,7 @@ def map_mxfp4_backend(runner_backend: str) -> Mxfp4MoeBackend: "flashinfer_cutlass_afp8": Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_MXFP8, "triton": Mxfp4MoeBackend.TRITON, "marlin": Mxfp4MoeBackend.MARLIN, - "ck": Mxfp4MoeBackend.CK, + "aiter": Mxfp4MoeBackend.AITER, "xpu": Mxfp4MoeBackend.XPU, } if backend := mapping.get(runner_backend): @@ -173,7 +173,7 @@ def _get_priority_backends() -> list[Mxfp4MoeBackend]: """ _AVAILABLE_BACKENDS = [ Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_BF16, - Mxfp4MoeBackend.CK, + Mxfp4MoeBackend.AITER, Mxfp4MoeBackend.TRITON, Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_BF16, Mxfp4MoeBackend.TRITON_UNFUSED, @@ -656,7 +656,7 @@ def convert_to_mxfp4_moe_kernel_format( w2_bias, ) - elif mxfp4_backend == Mxfp4MoeBackend.CK: + elif mxfp4_backend == Mxfp4MoeBackend.AITER: from vllm._aiter_ops import rocm_aiter_ops if w13_bias is not None: @@ -794,7 +794,7 @@ def make_mxfp4_moe_quant_config( Mxfp4MoeBackend.TRITON_UNFUSED, Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_BF16, Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_BF16, - Mxfp4MoeBackend.CK, + Mxfp4MoeBackend.AITER, ): return mxfp4_w4a16_moe_quant_config( w1_bias=w1_bias, diff --git a/vllm/model_executor/layers/fused_moe/router/gate_linear.py b/vllm/model_executor/layers/fused_moe/router/gate_linear.py index e8ed8a5249d..77d8e756026 100644 --- a/vllm/model_executor/layers/fused_moe/router/gate_linear.py +++ b/vllm/model_executor/layers/fused_moe/router/gate_linear.py @@ -3,11 +3,9 @@ import torch from torch.nn.parameter import Parameter -import vllm._custom_ops as ops from vllm.model_executor.custom_op import PluggableLayer from vllm.model_executor.layers.linear import ReplicatedLinear from vllm.platforms import current_platform -from vllm.utils.torch_utils import direct_register_custom_op @PluggableLayer.register("gate_linear") @@ -15,9 +13,8 @@ class GateLinear(ReplicatedLinear): """MoE gate linear layer with three-tier GEMM dispatch: 1. DSV3 specialized kernel (SM90+, batch<=16, supported dims) - 2. gpt-oss specialized kernel (SM90+, batch<=128, supported dims) - 3. cuBLAS bf16×bf16→fp32 (SM90+ + bf16 + fp32 out_dtype) - 4. F.linear via ReplicatedLinear (ultimate fallback) + 2. cuBLAS bf16×bf16→fp32 (SM90+ + bf16 + fp32 out_dtype) + 3. F.linear via ReplicatedLinear (ultimate fallback) The ``out_dtype`` attribute is mutable and can be set after init (e.g. when the required dtype depends on the expert quantization @@ -28,10 +25,6 @@ class GateLinear(ReplicatedLinear): DSV3_SUPPORTED_NUM_EXPERTS = [256, 384] DSV3_SUPPORTED_HIDDEN_SIZES = [7168] - # Dimensions supported by the gpt-oss specialized kernel - GPT_OSS_SUPPORTED_NUM_EXPERTS = [32, 128] - GPT_OSS_SUPPORTED_HIDDEN_SIZES = [2880] - def __init__( self, input_size: int, @@ -72,15 +65,6 @@ class GateLinear(ReplicatedLinear): and input_size in self.DSV3_SUPPORTED_HIDDEN_SIZES ) - # gpt-oss specialized kernel eligibility (SM90+, exact dims) - self.allow_gpt_oss_router_gemm = ( - self.weight.dtype == torch.bfloat16 - and current_platform.is_cuda() - and is_hopper_or_blackwell - and output_size in self.GPT_OSS_SUPPORTED_NUM_EXPERTS - and input_size in self.GPT_OSS_SUPPORTED_HIDDEN_SIZES - ) - # cuBLAS bf16→fp32 eligibility self.allow_cublas_router_gemm = ( self.allow_specialized_router_gemm @@ -108,6 +92,8 @@ class GateLinear(ReplicatedLinear): def forward( self, x: torch.Tensor ) -> torch.Tensor | tuple[torch.Tensor, Parameter | None]: + import vllm._custom_ops as ops + # Tier 1: DSV3 specialized kernel if self.allow_dsv3_router_gemm and x.shape[0] <= 16: output = ops.dsv3_router_gemm( @@ -117,47 +103,15 @@ class GateLinear(ReplicatedLinear): ) return output, None - # Tier 2: gpt-oss specialized kernel - if self.allow_gpt_oss_router_gemm: - output = torch.ops.vllm.gpt_oss_router_gemm(x, self.weight, self.bias) - return output, None - - # Tier 3: cuBLAS bf16→fp32 + # Tier 2: cuBLAS bf16→fp32 if self.allow_cublas_router_gemm and x.dtype == torch.bfloat16: output = ops.router_gemm_bf16_fp32(x, self.weight) return output, None - # Tier 4: F.linear (ReplicatedLinear) + # Tier 3: F.linear (ReplicatedLinear) if self.out_dtype is not None and x.dtype != self.weight.dtype: x = x.to(self.weight.dtype) output, output_bias = super().forward(x) if self.out_dtype is not None and output.dtype != self.out_dtype: output = output.to(self.out_dtype) return output, output_bias - - -def gpt_oss_router_gemm_impl( - x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor -) -> torch.Tensor: - """ - Dynamically run min-latency gemm if num_tokens <= 128. - This must be wrapped in a custom op because our torch.compile integration - does not support runtime dispatching on num_tokens. - """ - if x.shape[0] <= 128: - return ops.gpt_oss_router_gemm(x, weight, bias) - else: - return torch.nn.functional.linear(x, weight, bias) - - -def gpt_oss_router_gemm_fake( - x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor -) -> torch.Tensor: - return x.new_empty((x.shape[0], weight.shape[0])) - - -direct_register_custom_op( - op_name="gpt_oss_router_gemm", - op_func=gpt_oss_router_gemm_impl, - fake_impl=gpt_oss_router_gemm_fake, -) diff --git a/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py b/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py index f0d79982441..759f77b3657 100644 --- a/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py +++ b/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py @@ -222,6 +222,18 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp): self.cpu_fused_moe = cpu_fused_moe.CPUFusedMOE(layer) else: self.cpu_fused_moe = cpu_fused_moe.CPUFusedMOE(layer) + elif current_platform.is_xpu(): + w13 = layer.w13_weight + w2 = layer.w2_weight + + w13.data = w13.transpose(-1, -2).contiguous() + w2.data = w2.transpose(-1, -2).contiguous() + + self._setup_kernel( + layer=layer, + w13=w13, + w2=w2, + ) else: self._setup_kernel( layer=layer, diff --git a/vllm/model_executor/layers/kda.py b/vllm/model_executor/layers/kda.py index 46db5dc321d..b09f980c7e6 100644 --- a/vllm/model_executor/layers/kda.py +++ b/vllm/model_executor/layers/kda.py @@ -31,7 +31,11 @@ from .linear import ( RowParallelLinear, ) from .mamba.abstract import MambaBase -from .mamba.mamba_utils import MambaStateDtypeCalculator, MambaStateShapeCalculator +from .mamba.mamba_utils import ( + MambaStateDtypeCalculator, + MambaStateShapeCalculator, + is_conv_state_dim_first, +) from .mamba.ops.causal_conv1d import causal_conv1d_fn, causal_conv1d_update from .quantization.base_config import QuantizationConfig @@ -315,10 +319,12 @@ class KimiDeltaAttention(nn.Module, MambaBase): beta = beta[:num_actual_tokens] (conv_state_q, conv_state_k, conv_state_v, recurrent_state) = constant_caches - # deal with strides - conv_state_q = conv_state_q.transpose(-1, -2) - conv_state_k = conv_state_k.transpose(-1, -2) - conv_state_v = conv_state_v.transpose(-1, -2) + # conv_state must be (..., dim, width-1) for the conv kernels. + # DS layout stores it that way directly; SD layout needs a transpose. + if not is_conv_state_dim_first(): + conv_state_q = conv_state_q.transpose(-1, -2) + conv_state_k = conv_state_k.transpose(-1, -2) + conv_state_v = conv_state_v.transpose(-1, -2) q_conv_weights = self.q_conv1d.weight.view( self.q_conv1d.weight.size(0), self.q_conv1d.weight.size(2) diff --git a/vllm/model_executor/layers/layernorm.py b/vllm/model_executor/layers/layernorm.py index 500370d9f20..766bc46cebf 100644 --- a/vllm/model_executor/layers/layernorm.py +++ b/vllm/model_executor/layers/layernorm.py @@ -560,6 +560,11 @@ class RMSNormGated(CustomOp): activation=self.activation, ) + def forward_xpu( + self, x: torch.Tensor, z: torch.Tensor | None = None + ) -> torch.Tensor: + return self.forward_cuda(x, z) + class LayerNorm(nn.Module): """ diff --git a/vllm/model_executor/layers/mamba/gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn_linear_attn.py index 2b952e10e6e..aec855d9aeb 100644 --- a/vllm/model_executor/layers/mamba/gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn_linear_attn.py @@ -24,6 +24,7 @@ from vllm.model_executor.layers.fla.ops import ( chunk_gated_delta_rule as fla_chunk_gated_delta_rule, ) from vllm.model_executor.layers.fla.ops import ( + fused_post_conv_prep, fused_recurrent_gated_delta_rule_packed_decode, fused_sigmoid_gating_delta_rule_update, ) @@ -40,6 +41,7 @@ from vllm.model_executor.layers.mamba.mamba_mixer2 import mamba_v2_sharded_weigh from vllm.model_executor.layers.mamba.mamba_utils import ( MambaStateDtypeCalculator, MambaStateShapeCalculator, + is_conv_state_dim_first, ) from vllm.model_executor.layers.mamba.ops.causal_conv1d import ( causal_conv1d_fn, @@ -161,6 +163,8 @@ class ChunkGatedDeltaRule(CustomOp): initial_state: torch.Tensor, output_final_state: bool, cu_seqlens: torch.Tensor | None = None, + chunk_indices: torch.Tensor | None = None, + chunk_offsets: torch.Tensor | None = None, use_qk_l2norm_in_kernel: bool = True, ): return fi_chunk_gated_delta_rule( @@ -185,6 +189,8 @@ class ChunkGatedDeltaRule(CustomOp): initial_state: torch.Tensor, output_final_state: bool, cu_seqlens: torch.Tensor | None = None, + chunk_indices: torch.Tensor | None = None, + chunk_offsets: torch.Tensor | None = None, use_qk_l2norm_in_kernel: bool = True, ): return fla_chunk_gated_delta_rule( @@ -196,6 +202,8 @@ class ChunkGatedDeltaRule(CustomOp): initial_state=initial_state, output_final_state=output_final_state, cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_offsets=chunk_offsets, use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, ) @@ -260,6 +268,9 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase): else 0 ) self.gqa_interleaved_layout = gqa_interleaved_layout + self._forward_method = ( + self.forward_xpu if current_platform.is_xpu() else self.forward_cuda + ) # QKV self.conv_dim = self.key_dim * 2 + self.value_dim @@ -491,6 +502,13 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase): self, hidden_states: torch.Tensor, output: torch.Tensor, + ): + self._forward_method(hidden_states, output) + + def forward_cuda( + self, + hidden_states: torch.Tensor, + output: torch.Tensor, ): """ Forward pass with three parts: @@ -565,6 +583,90 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase): core_attn_out = rearrange(core_attn_out, "... h d -> ... (h d)") output[:num_tokens], _ = self.out_proj(core_attn_out) + def forward_xpu( + self, + hidden_states: torch.Tensor, + output: torch.Tensor, + ): + """ + Forward pass with three parts: + 1. Input projection + 2. Core attention (custom op) + 3. Output projection + """ + num_tokens = hidden_states.size(0) + + assert not hasattr(self, "in_proj_qkv"), "lora isn't supported on XPU." + + # ============================================================ + # Part 1: Input Projection + # ============================================================ + projected_states_qkvz, _ = self.in_proj_qkvz(hidden_states) + projected_states_ba, _ = self.in_proj_ba(hidden_states) + + # ============================================================ + # Part 2: Core Attention + # ============================================================ + forward_context = get_forward_context() + attn_metadata: AttentionMetadata = forward_context.attn_metadata + core_attn_out = torch.zeros( + (num_tokens, self.num_v_heads // self.tp_size, self.head_v_dim), + dtype=hidden_states.dtype, + device=hidden_states.device, + ) + z = torch.empty_like(core_attn_out) + if attn_metadata is not None: + attn_metadata = attn_metadata[self.prefix] + + # TODO: xpu does not support this param yet + spec_sequence_masks = attn_metadata.spec_sequence_masks + assert spec_sequence_masks is None + + conv_weights = self.conv1d.weight.view( + self.conv1d.weight.size(0), self.conv1d.weight.size(2) + ) + + conv_state = self.kv_cache[0] + ssm_state = self.kv_cache[1] + + torch.ops._xpu_C.gdn_attention( + core_attn_out, + z, + projected_states_qkvz, + projected_states_ba, + self.num_k_heads, + self.num_v_heads, + self.head_k_dim, + self.head_v_dim, + conv_state=conv_state, + ssm_state=ssm_state, + conv_weights=conv_weights, + conv_bias=self.conv1d.bias, + activation=self.activation, + A_log=self.A_log, + dt_bias=self.dt_bias, + num_prefills=attn_metadata.num_prefills, + num_decodes=attn_metadata.num_decodes, + has_initial_state=attn_metadata.has_initial_state, + non_spec_query_start_loc=attn_metadata.non_spec_query_start_loc, + non_spec_state_indices_tensor=attn_metadata.non_spec_state_indices_tensor, + num_actual_tokens=attn_metadata.num_actual_tokens, + tp_size=self.tp_size, + reorder_input=not self.gqa_interleaved_layout, + ) + + # ============================================================ + # Part 3: Output Projection + # ============================================================ + z_shape_og = z.shape + # Reshape input data into 2D tensor + core_attn_out = core_attn_out.reshape(-1, core_attn_out.shape[-1]) + z = z.reshape(-1, z.shape[-1]) + core_attn_out = self.norm(core_attn_out, z) + core_attn_out = core_attn_out.reshape(z_shape_og) + core_attn_out = rearrange(core_attn_out, "... h d -> ... (h d)") + output[:num_tokens], _ = self.out_proj(core_attn_out) + def _warmup_prefill_kernels(self, mixed_qkv: torch.Tensor) -> None: """Warm up GDN prefill kernels during V1 profiling. @@ -698,7 +800,13 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase): spec_state_indices_tensor = attn_metadata.spec_state_indices_tensor # noqa: E501 non_spec_state_indices_tensor = attn_metadata.non_spec_state_indices_tensor # noqa: E501 self_kv_cache = self.kv_cache - conv_state = self_kv_cache[0].transpose(-1, -2) + # conv_state must be (..., dim, width-1) for the conv kernels. + # DS layout stores it that way directly; SD layout needs a transpose. + conv_state = ( + self_kv_cache[0] + if is_conv_state_dim_first() + else self_kv_cache[0].transpose(-1, -2) + ) ssm_state = self_kv_cache[1] num_actual_tokens = attn_metadata.num_actual_tokens num_accepted_tokens = attn_metadata.num_accepted_tokens @@ -774,19 +882,44 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase): mixed_qkv_non_spec = None query_spec, key_spec, value_spec = self.rearrange_mixed_qkv(mixed_qkv_spec) - query_non_spec, key_non_spec, value_non_spec = self.rearrange_mixed_qkv( - mixed_qkv_non_spec - ) - if attn_metadata.num_prefills > 0: - g, beta = fused_gdn_gating(self.A_log, a, b, self.dt_bias) + assert mixed_qkv_non_spec is not None, ( + "mixed_qkv_non_spec must be provided for prefill path" + ) if spec_sequence_masks is not None: - g_non_spec = g.index_select(1, non_spec_token_indx) - beta_non_spec = beta.index_select(1, non_spec_token_indx) + a_non_spec = a.index_select(0, non_spec_token_indx) + b_non_spec = b.index_select(0, non_spec_token_indx) else: - g_non_spec = g - beta_non_spec = beta + a_non_spec = a + b_non_spec = b + + ( + query_non_spec, + key_non_spec, + value_non_spec, + g_non_spec, + beta_non_spec, + ) = fused_post_conv_prep( + conv_output=mixed_qkv_non_spec, + a=a_non_spec, + b=b_non_spec, + A_log=self.A_log, + dt_bias=self.dt_bias, + num_k_heads=self.num_k_heads // self.tp_size, + head_k_dim=self.head_k_dim, + head_v_dim=self.head_v_dim, + apply_l2norm=True, + output_g_exp=False, + ) + query_non_spec = query_non_spec.unsqueeze(0) + key_non_spec = key_non_spec.unsqueeze(0) + value_non_spec = value_non_spec.unsqueeze(0) + g_non_spec = g_non_spec.unsqueeze(0) + beta_non_spec = beta_non_spec.unsqueeze(0) else: + query_non_spec, key_non_spec, value_non_spec = self.rearrange_mixed_qkv( + mixed_qkv_non_spec + ) g_non_spec = None beta_non_spec = None @@ -832,7 +965,9 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase): initial_state=initial_state, output_final_state=True, cu_seqlens=non_spec_query_start_loc, - use_qk_l2norm_in_kernel=True, + chunk_indices=attn_metadata.chunk_indices, + chunk_offsets=attn_metadata.chunk_offsets, + use_qk_l2norm_in_kernel=False, ) # Init cache ssm_state[non_spec_state_indices_tensor] = last_recurrent_state.to( @@ -888,7 +1023,13 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase): """ non_spec_state_indices_tensor = attn_metadata.non_spec_state_indices_tensor # noqa: E501 self_kv_cache = self.kv_cache - conv_state = self_kv_cache[0].transpose(-1, -2) + # conv_state must be (..., dim, width-1) for the conv kernels. + # DS layout stores it that way directly; SD layout needs a transpose. + conv_state = ( + self_kv_cache[0] + if is_conv_state_dim_first() + else self_kv_cache[0].transpose(-1, -2) + ) ssm_state = self_kv_cache[1] num_actual_tokens = attn_metadata.num_actual_tokens diff --git a/vllm/model_executor/layers/mamba/mamba_mixer.py b/vllm/model_executor/layers/mamba/mamba_mixer.py index d79af2e2787..fd83d4b8322 100644 --- a/vllm/model_executor/layers/mamba/mamba_mixer.py +++ b/vllm/model_executor/layers/mamba/mamba_mixer.py @@ -24,6 +24,7 @@ from vllm.model_executor.layers.mamba.abstract import MambaBase from vllm.model_executor.layers.mamba.mamba_utils import ( MambaStateDtypeCalculator, MambaStateShapeCalculator, + is_conv_state_dim_first, ) from vllm.model_executor.layers.mamba.ops.causal_conv1d import ( causal_conv1d_fn, @@ -267,9 +268,12 @@ class MambaMixer(MambaBase, PluggableLayer): query_start_loc_p = attn_metadata.query_start_loc_p state_indices_tensor_p = attn_metadata.state_indices_tensor_p state_indices_tensor_d = attn_metadata.state_indices_tensor_d - self_kv_cache = self.kv_cache - conv_state = self_kv_cache[0].transpose(-1, -2) - ssm_state = self_kv_cache[1] + conv_state = ( + self.kv_cache[0] + if is_conv_state_dim_first() + else self.kv_cache[0].transpose(-1, -2) + ) + ssm_state = self.kv_cache[1] has_initial_states_p = attn_metadata.has_initial_states_p cu_chunk_seqlen_p = attn_metadata.cu_chunk_seqlen_p last_chunk_indices_p = attn_metadata.last_chunk_indices_p diff --git a/vllm/model_executor/layers/mamba/mamba_mixer2.py b/vllm/model_executor/layers/mamba/mamba_mixer2.py index 041405b05a0..01ea3fdca57 100644 --- a/vllm/model_executor/layers/mamba/mamba_mixer2.py +++ b/vllm/model_executor/layers/mamba/mamba_mixer2.py @@ -24,6 +24,7 @@ from vllm.model_executor.layers.mamba.abstract import MambaBase from vllm.model_executor.layers.mamba.mamba_utils import ( MambaStateDtypeCalculator, MambaStateShapeCalculator, + is_conv_state_dim_first, ) from vllm.model_executor.layers.mamba.ops.causal_conv1d import ( causal_conv1d_fn, @@ -575,10 +576,15 @@ class MambaMixer2(MambaBase, PluggableLayer): assert isinstance(attn_metadata, dict) attn_metadata = attn_metadata[self.prefix] assert isinstance(attn_metadata, Mamba2AttentionMetadata) - self_kv_cache = self.kv_cache - # conv_state = (..., dim, width-1) yet contiguous along 'dim' - conv_state = self_kv_cache[0].transpose(-1, -2) - ssm_state = self_kv_cache[1] + # conv_state must be (..., dim, width-1) for the conv kernels. + # DS layout stores it that way directly; SD layout needs a + # transpose (which keeps dim contiguous via stride tricks). + conv_state = ( + self.kv_cache[0] + if is_conv_state_dim_first() + else self.kv_cache[0].transpose(-1, -2) + ) + ssm_state = self.kv_cache[1] has_initial_states_p = attn_metadata.has_initial_states_p prep_initial_states = attn_metadata.prep_initial_states chunk_size = attn_metadata.chunk_size diff --git a/vllm/model_executor/layers/mamba/mamba_utils.py b/vllm/model_executor/layers/mamba/mamba_utils.py index 1f6751f6c8b..a5a30502b21 100644 --- a/vllm/model_executor/layers/mamba/mamba_utils.py +++ b/vllm/model_executor/layers/mamba/mamba_utils.py @@ -1,20 +1,52 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import functools from collections.abc import Callable from dataclasses import dataclass -from typing import TypeAlias +from typing import Literal, TypeAlias import torch +import vllm.envs as envs from vllm.config.cache import MambaDType from vllm.config.model import ModelDType from vllm.distributed import divide +from vllm.logger import init_logger from vllm.utils.torch_utils import ( STR_DTYPE_TO_TORCH_DTYPE, get_kv_cache_torch_dtype, ) +logger = init_logger(__name__) + +ConvStateLayoutType = Literal["SD", "DS"] + + +@functools.lru_cache +def get_conv_state_layout() -> ConvStateLayoutType: + """Return the SSM conv state layout. + + SD = (state_len, dim) — dim is the innermost contiguous dimension. + DS = (dim, state_len) — TP-sharded dim is on dim-1 (like HND for KV + cache), consistent with SSM temporal state layout. + """ + layout: ConvStateLayoutType | None = envs.VLLM_SSM_CONV_STATE_LAYOUT + if layout is not None: + logger.info_once( + "VLLM_SSM_CONV_STATE_LAYOUT env detected. " + "Setting SSM conv state layout to %s.", + layout, + ) + return layout + + return "SD" + + +def is_conv_state_dim_first() -> bool: + """True when the conv state is stored as (dim, state_len) per block.""" + return get_conv_state_layout() == "DS" + class MambaStateDtypeCalculator: @classmethod @@ -107,6 +139,13 @@ class MambaStateShapeCalculator: state_shape = (num_heads // tp_size, head_dim, head_dim) return (state_shape,) + @staticmethod + def _orient_conv_shape(dim: int, state_len: int) -> tuple[int, int]: + """Return (dim, state_len) for DS layout, (state_len, dim) for SD.""" + if is_conv_state_dim_first(): + return (dim, state_len) + return (state_len, dim) + @classmethod def mamba1_state_shape( cls, @@ -115,12 +154,11 @@ class MambaStateShapeCalculator: state_size: int, conv_kernel: int, ) -> tuple[tuple[int, int], tuple[int, int]]: - conv_state_shape = (divide(intermediate_size, tp_world_size), conv_kernel - 1) + conv_dim = divide(intermediate_size, tp_world_size) + conv_state_shape = cls._orient_conv_shape(conv_dim, conv_kernel - 1) temporal_state_shape = (divide(intermediate_size, tp_world_size), state_size) - conv_state_shape = conv_state_shape[1], conv_state_shape[0] - return conv_state_shape, temporal_state_shape @classmethod @@ -141,8 +179,9 @@ class MambaStateShapeCalculator: # heads and n_groups are TP-ed conv_dim = intermediate_size + 2 * n_groups * state_size - # contiguous along 'dim' axis - conv_state_shape = (conv_kernel - 1 + num_spec, divide(conv_dim, tp_world_size)) + conv_state_shape = cls._orient_conv_shape( + divide(conv_dim, tp_world_size), conv_kernel - 1 + num_spec + ) # These are not TP-ed as they depend on A, dt_bias, D # - they are typically small @@ -158,7 +197,7 @@ class MambaStateShapeCalculator: conv_kernel: int, ) -> tuple[tuple[int, int]]: conv_dim = divide(intermediate_size, tp_world_size) - conv_state_shape = (conv_kernel - 1, conv_dim) + conv_state_shape = cls._orient_conv_shape(conv_dim, conv_kernel - 1) return (conv_state_shape,) @classmethod @@ -185,13 +224,11 @@ class MambaStateShapeCalculator: num_spec: int = 0, ): conv_dim = head_k_dim * num_k_heads * 2 + head_v_dim * num_v_heads - conv_state_shape = ( + conv_state_shape = cls._orient_conv_shape( divide(conv_dim, tp_world_size), conv_kernel_size - 1 + num_spec, ) - conv_state_shape = conv_state_shape[1], conv_state_shape[0] - temporal_state_shape = ( divide(num_v_heads, tp_world_size), head_v_dim, @@ -218,12 +255,13 @@ class MambaStateShapeCalculator: proj_size = num_heads * head_dim proj_k_size = num_k_heads * head_k_dim - conv_state_shape = (divide(proj_size, tp_world_size), conv_kernel_size - 1) - conv_state_k_shape = (divide(proj_k_size, tp_world_size), conv_kernel_size - 1) + conv_state_shape = cls._orient_conv_shape( + divide(proj_size, tp_world_size), conv_kernel_size - 1 + ) + conv_state_k_shape = cls._orient_conv_shape( + divide(proj_k_size, tp_world_size), conv_kernel_size - 1 + ) recurrent_state_shape = (divide(num_heads, tp_world_size), head_dim, head_dim) - - conv_state_shape = conv_state_shape[1], conv_state_shape[0] - conv_state_k_shape = conv_state_k_shape[1], conv_state_k_shape[0] return ( conv_state_shape, conv_state_k_shape, @@ -267,9 +305,27 @@ def get_conv_copy_spec( cur_block_idx: int, num_accepted_tokens: int, ) -> MambaCopySpec: - """Return a MambaCopySpec for copying a convolutional state slice.""" + """Return a MambaCopySpec for copying a convolutional state slice. + + Works for both SD layout ``(num_blocks, state_len, dim)`` and + DS layout ``(num_blocks, dim, state_len)``. + """ src_block_id = block_ids[cur_block_idx] - src_state = state[src_block_id, num_accepted_tokens - 1 :] + offset = num_accepted_tokens - 1 + if is_conv_state_dim_first(): + # DS layout: (num_blocks, dim, state_len) — state_len is last. + if offset > 0: + # Slicing along the last dim yields a non-contiguous view + # because features (dim) are strided by state_len. + raise NotImplementedError( + "DS conv state layout does not yet support speculative " + "decoding with mamba_cache_mode='align' " + "(num_accepted_tokens > 1)." + ) + src_state = state[src_block_id] + else: + # SD layout: (num_blocks, state_len, dim) — dim contiguous. + src_state = state[src_block_id, offset:] return MambaCopySpec( start_addr=src_state.data_ptr(), num_elements=src_state.numel() ) diff --git a/vllm/model_executor/layers/mamba/ops/causal_conv1d.py b/vllm/model_executor/layers/mamba/ops/causal_conv1d.py index a8efdc9f18b..1160105ad10 100644 --- a/vllm/model_executor/layers/mamba/ops/causal_conv1d.py +++ b/vllm/model_executor/layers/mamba/ops/causal_conv1d.py @@ -592,7 +592,6 @@ def causal_conv1d_fn( stride_istate_seq = conv_states.stride(0) stride_istate_dim = conv_states.stride(1) stride_istate_token = conv_states.stride(2) - assert stride_istate_dim == 1 if out.dim() == 2: stride_o_dim = out.stride(0) stride_o_token = out.stride(1) @@ -1149,9 +1148,6 @@ def causal_conv1d_update( if validate_data: assert dim == weight.size(0) - assert conv_state.stride(-2) == 1, ( - f"ERROR: expect contiguous along feat-dim of conv_state (currently stride={conv_state.stride()})" - ) assert state_len >= width - 1 # when above happens, we don't shift-left to keep any records in conv_state assert dim == conv_state.size(1) diff --git a/vllm/model_executor/layers/mamba/short_conv.py b/vllm/model_executor/layers/mamba/short_conv.py index d36dc00964a..11e9b590f86 100644 --- a/vllm/model_executor/layers/mamba/short_conv.py +++ b/vllm/model_executor/layers/mamba/short_conv.py @@ -17,6 +17,7 @@ from vllm.model_executor.layers.mamba.abstract import MambaBase from vllm.model_executor.layers.mamba.mamba_utils import ( MambaStateDtypeCalculator, MambaStateShapeCalculator, + is_conv_state_dim_first, ) from vllm.model_executor.layers.mamba.ops.causal_conv1d import ( causal_conv1d_fn, @@ -117,8 +118,11 @@ class ShortConv(MambaBase, CustomOp): assert isinstance(attn_metadata, dict) attn_metadata = attn_metadata[self.prefix] assert isinstance(attn_metadata, ShortConvAttentionMetadata) - self_kv_cache = self.kv_cache - conv_state = self_kv_cache[0].transpose(-1, -2) + conv_state = ( + self.kv_cache[0] + if is_conv_state_dim_first() + else self.kv_cache[0].transpose(-1, -2) + ) state_indices_tensor_p = attn_metadata.state_indices_tensor_p state_indices_tensor_d = attn_metadata.state_indices_tensor_d has_initial_states_p = attn_metadata.has_initial_states_p diff --git a/vllm/model_executor/layers/quantization/awq.py b/vllm/model_executor/layers/quantization/awq.py index 58bb75d0a9e..37cffcb3da2 100644 --- a/vllm/model_executor/layers/quantization/awq.py +++ b/vllm/model_executor/layers/quantization/awq.py @@ -8,6 +8,7 @@ from safetensors.torch import _TYPES as _SAFETENSORS_TO_TORCH_DTYPE from transformers import PretrainedConfig from vllm import _custom_ops as ops +from vllm import envs from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe.layer import FusedMoE from vllm.model_executor.layers.linear import ( @@ -273,8 +274,9 @@ class AWQLinearMethod(LinearMethodBase): # num_tokens >= threshold FP16_MATMUL_HEURISTIC_CONDITION = x.shape[:-1].numel() >= 256 - - if FP16_MATMUL_HEURISTIC_CONDITION: + # Batch invariant mode requires torch.matmul path + # for Triton override + if FP16_MATMUL_HEURISTIC_CONDITION or envs.VLLM_BATCH_INVARIANT: out = ops.awq_dequantize(qweight, scales, qzeros, 0, 0, 0) out = torch.matmul(reshaped_x, out) else: diff --git a/vllm/model_executor/layers/quantization/awq_marlin.py b/vllm/model_executor/layers/quantization/awq_marlin.py index eff571ef2a7..be3001a7fa1 100644 --- a/vllm/model_executor/layers/quantization/awq_marlin.py +++ b/vllm/model_executor/layers/quantization/awq_marlin.py @@ -10,6 +10,7 @@ from transformers import PretrainedConfig import vllm.model_executor.layers.fused_moe # noqa from vllm import _custom_ops as ops +from vllm import envs from vllm.logger import init_logger from vllm.model_executor.kernels.linear import ( MPLinearLayerConfig, @@ -233,6 +234,11 @@ class AWQMarlinConfig(QuantizationConfig): def override_quantization_method( cls, hf_quant_cfg, user_quant ) -> "QuantizationMethods | None": + # Skip override to marlin kernels, as they are not + # batch invariant + if envs.VLLM_BATCH_INVARIANT: + return None + can_convert = cls.is_awq_marlin_compatible(hf_quant_cfg) is_valid_user_quant = ( user_quant is None or user_quant == "marlin" or user_quant == "awq_marlin" diff --git a/vllm/model_executor/layers/quantization/fp8.py b/vllm/model_executor/layers/quantization/fp8.py index 965e1af72e3..259a7d1f6c2 100644 --- a/vllm/model_executor/layers/quantization/fp8.py +++ b/vllm/model_executor/layers/quantization/fp8.py @@ -1028,6 +1028,10 @@ class Fp8OnlineMoEMethod(Fp8MoEMethod): layer.w2_weight[expert, :, :] ) + if current_platform.is_xpu(): + w13.data = w13.transpose(-1, -2).contiguous() + w2.data = w2.transpose(-1, -2).contiguous() + # Shuffle weights to runtime format and setup kernel. self._setup_kernel( layer, diff --git a/vllm/model_executor/layers/quantization/kv_cache.py b/vllm/model_executor/layers/quantization/kv_cache.py index 2fb67aacc54..726ac2232af 100644 --- a/vllm/model_executor/layers/quantization/kv_cache.py +++ b/vllm/model_executor/layers/quantization/kv_cache.py @@ -10,6 +10,7 @@ from vllm.model_executor.layers.quantization.base_config import ( ) from vllm.platforms import current_platform from vllm.utils.torch_utils import is_quantized_kv_cache +from vllm.v1.kv_cache_interface import kv_cache_uses_per_token_head_scales logger = init_logger(__name__) @@ -53,6 +54,20 @@ class BaseKVCacheMethod(QuantizeMethodBase): assert not hasattr(layer, "prob_scale") return + # Per-token-head quantized KV cache: scales are computed dynamically + # per (token, head) in the kernel at cache-write time. Checkpoint + # scales are never used regardless of calculate_kv_scales. + if kv_cache_uses_per_token_head_scales(layer.kv_cache_dtype): + layer._k_scale.copy_(1.0) + layer._v_scale.copy_(1.0) + layer._k_scale_float = 1.0 + layer._v_scale_float = 1.0 + del layer.k_scale + del layer.v_scale + del layer.q_scale + del layer.prob_scale + return + # If the kv-cache is not quantized, we enforce the k/v_scale to be 1.0 # regardless whether the kv-scale is available in the checkpoint. # No need to process kv scales after loading if we are going to diff --git a/vllm/model_executor/layers/quantization/quark/quark_moe.py b/vllm/model_executor/layers/quantization/quark/quark_moe.py index c48e49fe86c..3f7ddbfd756 100644 --- a/vllm/model_executor/layers/quantization/quark/quark_moe.py +++ b/vllm/model_executor/layers/quantization/quark/quark_moe.py @@ -5,6 +5,7 @@ from typing import Any import torch +import vllm.model_executor.layers.fused_moe.modular_kernel as mk from vllm import _custom_ops as ops from vllm import envs from vllm._aiter_ops import rocm_aiter_ops @@ -27,7 +28,11 @@ from vllm.model_executor.layers.fused_moe.config import ( ) from vllm.model_executor.layers.fused_moe.fused_marlin_moe import fused_marlin_moe from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import ( + TRITON_BACKENDS, Mxfp4MoeBackend, + convert_to_mxfp4_moe_kernel_format, + make_mxfp4_moe_kernel, + make_mxfp4_moe_quant_config, mxfp4_round_up_hidden_size_and_intermediate_size, select_mxfp4_moe_backend, ) @@ -47,7 +52,7 @@ from vllm.model_executor.layers.quantization.utils.w8a8_utils import ( normalize_e4m3fn_to_e4m3fnuz, per_tensor_dequantize, ) -from vllm.model_executor.utils import set_weight_attrs +from vllm.model_executor.utils import replace_parameter, set_weight_attrs from vllm.platforms import current_platform from vllm.scalar_type import scalar_types @@ -699,9 +704,16 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod): f"Please check that the combination is supported in OCP_MX_Scheme." ) - self.mxfp4_backend: Mxfp4MoeBackend | None = None + self.mxfp4_backend: Mxfp4MoeBackend = Mxfp4MoeBackend.NONE + self.experts_cls: type[mk.FusedMoEExperts] | None = None + self.moe_kernel: mk.FusedMoEKernel | None = None + + # Used for triton kernel precision configs + self.w13_precision_config = None + self.w2_precision_config = None + if self.ocp_mx_scheme == "w_mxfp4": - self.mxfp4_backend, _ = select_mxfp4_moe_backend(moe) + self.mxfp4_backend, self.experts_cls = select_mxfp4_moe_backend(moe) elif self.ocp_mx_scheme.startswith("w_mxfp4"): # TODO(bowenbao): refactor and introduce backends for other OCP MX schemes. self.mxfp4_backend = Mxfp4MoeBackend.NONE @@ -738,9 +750,7 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod): not current_platform.supports_mx() or not self.ocp_mx_scheme.startswith("w_mxfp4") ) and ( - self.mxfp4_backend is None - or self.mxfp4_backend is Mxfp4MoeBackend.NONE - or not self.use_rocm_aiter_moe + self.mxfp4_backend is Mxfp4MoeBackend.NONE or not self.use_rocm_aiter_moe ) if self.emulate: @@ -944,11 +954,23 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod): w2_input_scale, requires_grad=False ) - # secondly, process mxfp weights + # For w_mxfp4, use oracle functions + if ( + self.ocp_mx_scheme == "w_mxfp4" + and self.mxfp4_backend != Mxfp4MoeBackend.NONE + ): + self._setup_kernel_via_oracle(layer) + return + + # TODO(bowenbao): gradually migrate to oracles. + # secondly, process mxfp weights for other schemes if self.emulate: + # Build quant config for emulation path + self.moe_quant_config = self.get_fused_moe_quant_config(layer) torch.accelerator.empty_cache() return + # Existing AITER path for w_mxfp4_a_mxfp4 and other schemes from aiter.utility.fp4_utils import e8m0_shuffle # Pre-shuffle weight scales @@ -980,11 +1002,87 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod): layer.w2_weight = torch.nn.Parameter(shuffled_w2, requires_grad=False) layer.w13_weight.is_shuffled = True layer.w2_weight.is_shuffled = True + + # Build quant config for AITER path + self.moe_quant_config = self.get_fused_moe_quant_config(layer) torch.accelerator.empty_cache() + def _setup_kernel_via_oracle(self, layer: FusedMoE): + """Setup kernel using oracle functions for w_mxfp4 scheme.""" + w13 = layer.w13_weight + w2 = layer.w2_weight + w13_scale = layer.w13_weight_scale + w2_scale = layer.w2_weight_scale + w13_bias = getattr(layer, "w13_bias", None) + w2_bias = getattr(layer, "w2_bias", None) + + # Convert weights to kernel format + w13, w2, w13_scale, w2_scale, w13_bias, w2_bias = ( + convert_to_mxfp4_moe_kernel_format( + mxfp4_backend=self.mxfp4_backend, + layer=layer, + w13_weight=w13, + w2_weight=w2, + w13_weight_scale=w13_scale, + w2_weight_scale=w2_scale, + w13_bias=w13_bias, + w2_bias=w2_bias, + ) + ) + + # For TRITON backends, weights are wrapped tensors from triton_kernels + # that don't support .detach(). Manually assign parameters. + if self.mxfp4_backend not in TRITON_BACKENDS: + replace_parameter(layer, "w13_weight", w13) + replace_parameter(layer, "w2_weight", w2) + replace_parameter(layer, "w13_weight_scale", w13_scale) + replace_parameter(layer, "w2_weight_scale", w2_scale) + else: + layer.w13_weight = w13 + layer.w2_weight = w2 + self.w13_precision_config = w13_scale + self.w2_precision_config = w2_scale + + if w13_bias is not None and w2_bias is not None: + replace_parameter(layer, "w13_bias", w13_bias) + replace_parameter(layer, "w2_bias", w2_bias) + + # Build quant config and kernel + self.moe_quant_config = self.get_fused_moe_quant_config(layer) + if self.moe_quant_config is not None and self.experts_cls is not None: + self.moe_kernel = make_mxfp4_moe_kernel( + moe_quant_config=self.moe_quant_config, + moe_config=self.moe, + mxfp4_backend=self.mxfp4_backend, + experts_cls=self.experts_cls, + routing_tables=layer._maybe_init_expert_routing_tables(), + shared_experts=layer.shared_experts, + ) + def get_fused_moe_quant_config( self, layer: torch.nn.Module ) -> FusedMoEQuantConfig | None: + # For w_mxfp4 with oracle backend, use oracle function + if ( + self.ocp_mx_scheme == "w_mxfp4" + and self.mxfp4_backend != Mxfp4MoeBackend.NONE + ): + w1_scale = layer.w13_weight_scale + w2_scale = layer.w2_weight_scale + if self.mxfp4_backend in TRITON_BACKENDS: + w1_scale = self.w13_precision_config + w2_scale = self.w2_precision_config + return make_mxfp4_moe_quant_config( + mxfp4_backend=self.mxfp4_backend, + w1_scale=w1_scale, + w2_scale=w2_scale, + w1_bias=getattr(layer, "w13_bias", None), + w2_bias=getattr(layer, "w2_bias", None), + ) + + # Existing code for other schemes + # TODO(bowenbao): kept for emulation fallback, to be refactored into + # dedicated emulation backend. if self.ocp_mx_scheme == "w_mxfp4": return mxfp4_w4a16_moe_quant_config( w1_scale=layer.w13_weight_scale, @@ -1020,6 +1118,12 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod): block_shape=None, ) + @property + def is_monolithic(self) -> bool: + if self.moe_kernel is not None: + return self.moe_kernel.is_monolithic + return False + def apply( self, layer: FusedMoE, @@ -1028,6 +1132,22 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod): topk_ids: torch.Tensor, shared_experts_input: torch.Tensor | None, ) -> torch.Tensor: + # For w_mxfp4 with oracle kernel + if self.moe_kernel is not None: + return self.moe_kernel.apply( + hidden_states=x, + w1=layer.w13_weight, + w2=layer.w2_weight, + topk_weights=topk_weights, + topk_ids=topk_ids, + activation=layer.activation, + global_num_experts=layer.global_num_experts, + apply_router_weight_on_input=layer.apply_router_weight_on_input, + expert_map=layer.expert_map, + shared_experts_input=shared_experts_input, + ) + + # Existing code for emulation/AITER paths if not self.emulate: from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import ( rocm_aiter_fused_experts, @@ -1061,6 +1181,25 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod): quant_config=self.moe_quant_config, ) + def apply_monolithic( + self, + layer: FusedMoE, + x: torch.Tensor, + router_logits: torch.Tensor, + ) -> torch.Tensor: + assert self.is_monolithic + assert self.moe_kernel is not None + return self.moe_kernel.apply_monolithic( + hidden_states=x, + w1=layer.w13_weight, + w2=layer.w2_weight, + router_logits=router_logits, + activation=layer.activation, + global_num_experts=layer.global_num_experts, + expert_map=layer.expert_map, + apply_router_weight_on_input=layer.apply_router_weight_on_input, + ) + class QuarkOCP_MX_MoEMethod_OSS(QuarkOCP_MX_MoEMethod): def __init__( diff --git a/vllm/model_executor/layers/quantization/utils/int8_utils.py b/vllm/model_executor/layers/quantization/utils/int8_utils.py index 020098dffc3..a98e29ffd57 100644 --- a/vllm/model_executor/layers/quantization/utils/int8_utils.py +++ b/vllm/model_executor/layers/quantization/utils/int8_utils.py @@ -88,6 +88,13 @@ if current_platform.is_rocm(): def round_int8(x): return tl.extra.hip.libdevice.round(x).to(tl.int8) + +elif current_platform.is_xpu(): + + @triton.jit + def round_int8(x): + return tl.extra.intel.libdevice.round(x).to(tl.int8) + else: @triton.jit diff --git a/vllm/model_executor/layers/rotary_embedding/__init__.py b/vllm/model_executor/layers/rotary_embedding/__init__.py index 9ad7c9cdafd..9a541877575 100644 --- a/vllm/model_executor/layers/rotary_embedding/__init__.py +++ b/vllm/model_executor/layers/rotary_embedding/__init__.py @@ -12,6 +12,7 @@ from .dual_chunk_rope import DualChunkRotaryEmbedding from .dynamic_ntk_alpha_rope import DynamicNTKAlphaRotaryEmbedding from .dynamic_ntk_scaling_rope import DynamicNTKScalingRotaryEmbedding from .fope import FourierRotaryEmbedding +from .gemma4_rope import Gemma4RotaryEmbedding from .linear_scaling_rope import LinearScalingRotaryEmbedding from .llama3_rope import Llama3RotaryEmbedding from .llama4_vision_rope import Llama4VisionRotaryEmbedding @@ -19,6 +20,7 @@ from .mrope import MRotaryEmbedding from .mrope_interleaved import MRotaryEmbeddingInterleaved from .ntk_scaling_rope import NTKScalingRotaryEmbedding from .phi3_long_rope_scaled_rope import Phi3LongRoPEScaledRotaryEmbedding +from .telechat3_scaling_rope import TeleChat3RoPEScaledRotaryEmbedding from .xdrope import XDRotaryEmbedding from .yarn_scaling_rope import YaRNScalingRotaryEmbedding @@ -134,6 +136,17 @@ def get_rope( is_neox_style, dtype, ) + elif scaling_type == "proportional": + # Proportional RoPE is used by Gemma4 for global (full) attention. + # Gemma4 uses a sparse/fractional RoPE with cross-mixing between halves. + rotary_emb = Gemma4RotaryEmbedding( + head_size, + rotary_dim, + max_position, + base, + is_neox_style, + dtype, + ) elif scaling_type == "llama3": scaling_factor = rope_parameters["factor"] low_freq_factor = rope_parameters["low_freq_factor"] @@ -322,6 +335,36 @@ def get_rope( ) else: raise ValueError("Pangu mrope lacks necessary parameters.") + elif scaling_type == "telechat3-yarn": + scaling_factor = rope_parameters["factor"] + if "original_max_position_embeddings" in rope_parameters: + original_max_position = rope_parameters["original_max_position_embeddings"] + scaling_factor = max_position / original_max_position + else: + original_max_position = max_position + extra_kwargs = { + k: v + for k, v in rope_parameters.items() + if k + in ( + "extrapolation_factor", + "attn_factor", + "beta_fast", + "beta_slow", + "mscale", + "mscale_all_dim", + ) + } + rotary_emb = TeleChat3RoPEScaledRotaryEmbedding( + head_size, + rotary_dim, + original_max_position, + base, + is_neox_style, + scaling_factor, + dtype, + **extra_kwargs, + ) else: raise ValueError(f"Unknown RoPE scaling type {scaling_type}") _ROPE_DICT[key] = rotary_emb diff --git a/vllm/model_executor/layers/rotary_embedding/gemma4_rope.py b/vllm/model_executor/layers/rotary_embedding/gemma4_rope.py new file mode 100644 index 00000000000..48253f469cc --- /dev/null +++ b/vllm/model_executor/layers/rotary_embedding/gemma4_rope.py @@ -0,0 +1,84 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Gemma4-specific Rotary Positional Embeddings (proportional scaling). + +Gemma4 uses "proportional" RoPE which computes inv_freq frequencies scaled +by head_dim (not rotary_dim), and zero-pads for non-rotated dimensions when +partial_rotary_factor < 1. The actual rotation uses standard neox-style +rotate_half, matching HF transformers' apply_rotary_pos_emb. +""" + +import torch + +from .base import RotaryEmbedding + + +class Gemma4RotaryEmbedding(RotaryEmbedding): + """Gemma4 proportional RoPE. + + Extends RotaryEmbedding (which provides standard neox-style rotation + via ops.rotary_embedding CUDA kernel) but overrides the inv_freq + computation to match HF's _compute_proportional_rope_parameters: + - Frequency exponents use head_dim (not rotary_dim) as denominator + - Non-rotated dims are zero-padded (cos=1, sin=0 = identity rotation) + + When partial_rotary_factor=1.0 (the default for some variants), ALL dims are + rotated and this is equivalent to standard RotaryEmbedding with + head_dim-scaled frequencies. + """ + + def __init__( + self, + head_size: int, + rotary_dim: int, + max_position_embeddings: int, + base: float, + is_neox_style: bool, + dtype: torch.dtype, + ) -> None: + # Number of rotation angle pairs (from partial_rotary_factor) + self.rope_angles = rotary_dim // 2 + # Non-rotated angle pairs per half + self.nope_angles = (head_size // 2) - self.rope_angles + + # Important: set rotary_dim = head_size so the base class's + # forward_static applies rotation to ALL dims of the cos/sin cache. + # The non-rotated dims will have cos=1, sin=0 (identity) thanks + # to our _compute_inv_freq zero-padding. + super().__init__( + head_size, + head_size, # rotary_dim = head_size (full application) + max_position_embeddings, + base, + is_neox_style, + dtype, + ) + + def _compute_inv_freq(self, base: float) -> torch.Tensor: + """Compute frequencies matching HF proportional RoPE. + + Key difference from base: exponent denominator is head_size (not + rotary_dim), and non-rotated dims are zero-padded. + """ + # HF formula: base ** (arange(0, 2*rope_angles, 2) / head_dim) + freq_exponents = ( + torch.arange(0, 2 * self.rope_angles, 2, dtype=torch.float) / self.head_size + ) + inv_freq = 1.0 / (base**freq_exponents) + + # Zero-pad for non-rotated dims (identity rotation: cos=1, sin=0) + if self.nope_angles > 0: + inv_freq = torch.cat( + [ + inv_freq, + torch.zeros(self.nope_angles, dtype=torch.float), + ] + ) + return inv_freq + + def extra_repr(self) -> str: + s = f"head_size={self.head_size}, rotary_dim={self.rotary_dim}" + s += f", rope_angles={self.rope_angles}, nope_angles={self.nope_angles}" + s += f", max_position_embeddings={self.max_position_embeddings}" + s += f", base={self.base}, is_neox_style={self.is_neox_style}" + return s diff --git a/vllm/model_executor/layers/rotary_embedding/telechat3_scaling_rope.py b/vllm/model_executor/layers/rotary_embedding/telechat3_scaling_rope.py new file mode 100644 index 00000000000..dd2fb9c320b --- /dev/null +++ b/vllm/model_executor/layers/rotary_embedding/telechat3_scaling_rope.py @@ -0,0 +1,57 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import math + +import torch + +from .base import RotaryEmbedding +from .yarn_scaling_rope import YaRNScalingRotaryEmbedding + + +class TeleChat3RoPEScaledRotaryEmbedding(YaRNScalingRotaryEmbedding): + """TeleChat3 uses a variant of YaRN method. + + To achieve code reuse as much as possible, we have rewritten the + `get_mscale` method in the initialization function + """ + + def __init__( + self, + head_size: int, + rotary_dim: int, + max_position_embeddings: int, + base: int, + is_neox_style: bool, + scaling_factor: float, + dtype: torch.dtype, + *, + extrapolation_factor: float = 1, + attn_factor: float = 1, + beta_fast: int = 32, + beta_slow: int = 1, + truncate: bool = True, + ) -> None: + self.scaling_factor = scaling_factor + self.extrapolation_factor = extrapolation_factor + self.attn_factor = attn_factor + self.beta_fast = beta_fast + self.beta_slow = beta_slow + self.truncate = truncate + + def get_mscale(scale, mscale=1): + if scale <= 1: + return 1.0 + return 0.07 * mscale * math.log(scale) + 1.0 + + self.mscale = float(get_mscale(self.scaling_factor) * attn_factor) + # Initialization must be performed after mscale, otherwise mscale is useless + RotaryEmbedding.__init__( + self, + head_size, + rotary_dim, + max_position_embeddings, + base, + is_neox_style, + dtype, + ) diff --git a/vllm/model_executor/model_loader/reload/layerwise.py b/vllm/model_executor/model_loader/reload/layerwise.py index 5d4af7d1f59..2934b8b5ad5 100644 --- a/vllm/model_executor/model_loader/reload/layerwise.py +++ b/vllm/model_executor/model_loader/reload/layerwise.py @@ -14,6 +14,7 @@ from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBa from vllm.model_executor.model_loader.weight_utils import default_weight_loader from .meta import ( + SKIP_TENSORS, capture_layer_to_meta, get_numel_loaded, materialize_layer, @@ -124,6 +125,8 @@ def initialize_online_processing(layer: torch.nn.Module): # Wrap each parameter's weight loader # Note that nested wrapping will occur for shared tensors for name, tensor in get_layer_tensors(layer).items(): + if name in SKIP_TENSORS: + continue if _get_weight_loader(tensor).__name__ != "online_process_loader": tensor.weight_loader = make_online_process_loader(layer, name) diff --git a/vllm/model_executor/model_loader/reload/meta.py b/vllm/model_executor/model_loader/reload/meta.py index 82bf9ce3d2a..91fce6f57b3 100644 --- a/vllm/model_executor/model_loader/reload/meta.py +++ b/vllm/model_executor/model_loader/reload/meta.py @@ -27,6 +27,7 @@ SKIP_TENSORS: set[str] = { "expert_global_to_physical", "expert_physical_to_global", "expert_local_to_global", + "e_score_correction_bias", } diff --git a/vllm/model_executor/models/cheers.py b/vllm/model_executor/models/cheers.py new file mode 100644 index 00000000000..5f74c6771e4 --- /dev/null +++ b/vllm/model_executor/models/cheers.py @@ -0,0 +1,753 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Inference-only Cheers (UMM) model compatible with HuggingFace weights. + +Cheers is a unified multimodal model for image understanding and generation. +For vLLM, we focus on the image understanding (vision-to-text) capabilities. +The image generation part (gen_projector, hi_gate, etc.) is not supported, +but the VAE encoder + decoder projector are required for image understanding. +""" + +import math +from collections.abc import Iterable, Mapping, Sequence +from typing import Any, Literal, TypeAlias + +import torch +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange +from transformers import BatchFeature + +from vllm.config import VllmConfig +from vllm.config.multimodal import BaseDummyOptions +from vllm.inputs import MultiModalDataDict +from vllm.logger import init_logger +from vllm.model_executor.layers.quantization import QuantizationConfig +from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.multimodal.inputs import ( + MultiModalFieldConfig, + MultiModalKwargsItems, +) +from vllm.multimodal.parse import MultiModalDataItems +from vllm.multimodal.processing import ( + BaseDummyInputsBuilder, + BaseMultiModalProcessor, + BaseProcessingInfo, + PromptReplacement, +) +from vllm.sequence import IntermediateTensors +from vllm.transformers_utils.processors.cheers import CheersProcessor +from vllm.utils.tensor_schema import TensorSchema + +from .interfaces import ( + MultiModalEmbeddings, + SupportsLoRA, + SupportsMultiModal, + SupportsPP, +) +from .siglip import SiglipVisionModel +from .utils import ( + AutoWeightsLoader, + WeightsMapper, + init_vllm_registered_model, + maybe_prefix, +) + +logger = init_logger(__name__) + + +# ── VAE components (needed for image understanding pipeline) ──────── + + +def _swish(x: torch.Tensor) -> torch.Tensor: + return x * torch.sigmoid(x) + + +class _AttnBlock(nn.Module): + def __init__(self, in_channels: int): + super().__init__() + self.norm = nn.GroupNorm(32, in_channels, eps=1e-6, affine=True) + self.q = nn.Conv2d(in_channels, in_channels, 1) + self.k = nn.Conv2d(in_channels, in_channels, 1) + self.v = nn.Conv2d(in_channels, in_channels, 1) + self.proj_out = nn.Conv2d(in_channels, in_channels, 1) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + h_ = self.norm(x) + q = self.q(h_) + k = self.k(h_) + v = self.v(h_) + b, c, h, w = q.shape + q = rearrange(q, "b c h w -> b 1 (h w) c").contiguous() + k = rearrange(k, "b c h w -> b 1 (h w) c").contiguous() + v = rearrange(v, "b c h w -> b 1 (h w) c").contiguous() + h_ = F.scaled_dot_product_attention(q, k, v) + h_ = rearrange(h_, "b 1 (h w) c -> b c h w", h=h, w=w, c=c, b=b) + return x + self.proj_out(h_) + + +class _ResnetBlock(nn.Module): + def __init__(self, in_channels: int, out_channels: int): + super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + self.norm1 = nn.GroupNorm(32, in_channels, eps=1e-6, affine=True) + self.conv1 = nn.Conv2d(in_channels, out_channels, 3, 1, 1) + self.norm2 = nn.GroupNorm(32, out_channels, eps=1e-6, affine=True) + self.conv2 = nn.Conv2d(out_channels, out_channels, 3, 1, 1) + if in_channels != out_channels: + self.nin_shortcut = nn.Conv2d(in_channels, out_channels, 1) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + h = _swish(self.norm1(x)) + h = self.conv1(h) + h = _swish(self.norm2(h)) + h = self.conv2(h) + if self.in_channels != self.out_channels: + x = self.nin_shortcut(x) + return x + h + + +class _Downsample(nn.Module): + def __init__(self, in_channels: int): + super().__init__() + self.conv = nn.Conv2d(in_channels, in_channels, 3, stride=2, padding=0) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = F.pad(x, (0, 1, 0, 1), mode="constant", value=0) + return self.conv(x) + + +class _Upsample(nn.Module): + def __init__(self, in_channels: int): + super().__init__() + self.conv = nn.Conv2d(in_channels, in_channels, 3, 1, 1) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = F.interpolate(x, scale_factor=2.0, mode="nearest") + return self.conv(x) + + +_VAE_ENCODER_DEFAULTS = { + "in_channels": 3, + "ch": 128, + "ch_mult": [1, 2, 4, 4], + "num_res_blocks": 2, + "z_channels": 32, +} +_VAE_DECODER_DEFAULTS = { + "in_channels": 3, + "out_ch": 3, + "ch": 128, + "ch_mult": [1, 2, 4, 4], + "num_res_blocks": 2, + "z_channels": 32, +} + + +def _cfg(config, key, defaults=None): + """Access config attribute whether it's a dict or namespace object.""" + if isinstance(config, dict): + if key in config: + return config[key] + if defaults and key in defaults: + return defaults[key] + raise KeyError(f"Key '{key}' not found in config dict: {list(config.keys())}") + return getattr(config, key) + + +class CheersVAEEncoder(nn.Module): + """VAE encoder from the Cheers/UMM model.""" + + def __init__(self, config): + super().__init__() + d = _VAE_ENCODER_DEFAULTS + ch = _cfg(config, "ch", d) + ch_mult = _cfg(config, "ch_mult", d) + num_res_blocks = _cfg(config, "num_res_blocks", d) + z_channels = _cfg(config, "z_channels", d) + in_channels = _cfg(config, "in_channels", d) + num_resolutions = len(ch_mult) + + self.quant_conv = nn.Conv2d(2 * z_channels, 2 * z_channels, 1) + self.conv_in = nn.Conv2d(in_channels, ch, 3, 1, 1) + + in_ch_mult = (1,) + tuple(ch_mult) + self.down = nn.ModuleList() + block_in = ch + for i_level in range(num_resolutions): + block = nn.ModuleList() + attn = nn.ModuleList() + block_in = ch * in_ch_mult[i_level] + block_out = ch * ch_mult[i_level] + for _ in range(num_res_blocks): + block.append(_ResnetBlock(block_in, block_out)) + block_in = block_out + down = nn.Module() + down.block = block + down.attn = attn + if i_level != num_resolutions - 1: + down.downsample = _Downsample(block_in) + self.down.append(down) + + self.mid = nn.Module() + self.mid.block_1 = _ResnetBlock(block_in, block_in) + self.mid.attn_1 = _AttnBlock(block_in) + self.mid.block_2 = _ResnetBlock(block_in, block_in) + + self.norm_out = nn.GroupNorm(32, block_in, eps=1e-6, affine=True) + self.conv_out = nn.Conv2d(block_in, 2 * z_channels, 3, 1, 1) + self._num_resolutions = num_resolutions + self._num_res_blocks = num_res_blocks + + def forward(self, x: torch.Tensor) -> torch.Tensor: + hs = [self.conv_in(x)] + for i_level in range(self._num_resolutions): + for i_block in range(self._num_res_blocks): + h = self.down[i_level].block[i_block](hs[-1]) + if len(self.down[i_level].attn) > 0: + h = self.down[i_level].attn[i_block](h) + hs.append(h) + if hasattr(self.down[i_level], "downsample"): + hs.append(self.down[i_level].downsample(hs[-1])) + h = hs[-1] + h = self.mid.block_1(h) + h = self.mid.attn_1(h) + h = self.mid.block_2(h) + h = _swish(self.norm_out(h)) + h = self.conv_out(h) + h = self.quant_conv(h) + return h + + +class CheersVAEDecoder(nn.Module): + """VAE decoder (used inside VAEDecoderProjector).""" + + def __init__(self, config): + super().__init__() + d = _VAE_DECODER_DEFAULTS + ch = _cfg(config, "ch", d) + ch_mult = _cfg(config, "ch_mult", d) + num_res_blocks = _cfg(config, "num_res_blocks", d) + z_channels = _cfg(config, "z_channels", d) + out_ch = _cfg(config, "out_ch", d) + num_resolutions = len(ch_mult) + + self.post_quant_conv = nn.Conv2d(z_channels, z_channels, 1) + block_in = ch * ch_mult[num_resolutions - 1] + self.conv_in = nn.Conv2d(z_channels, block_in, 3, 1, 1) + + self.mid = nn.Module() + self.mid.block_1 = _ResnetBlock(block_in, block_in) + self.mid.attn_1 = _AttnBlock(block_in) + self.mid.block_2 = _ResnetBlock(block_in, block_in) + + self.up = nn.ModuleList() + for i_level in reversed(range(num_resolutions)): + block = nn.ModuleList() + attn = nn.ModuleList() + block_out = ch * ch_mult[i_level] + for _ in range(num_res_blocks + 1): + block.append(_ResnetBlock(block_in, block_out)) + block_in = block_out + up = nn.Module() + up.block = block + up.attn = attn + if i_level != 0: + up.upsample = _Upsample(block_in) + self.up.insert(0, up) + + self.norm_out = nn.GroupNorm(32, block_in, eps=1e-6, affine=True) + self.conv_out = nn.Conv2d(block_in, out_ch, 3, 1, 1) + self._num_resolutions = num_resolutions + self._num_res_blocks = num_res_blocks + + def forward(self, z: torch.Tensor) -> torch.Tensor: + z = self.post_quant_conv(z) + upscale_dtype = next(self.up.parameters()).dtype + h = self.conv_in(z) + h = self.mid.block_1(h) + h = self.mid.attn_1(h) + h = self.mid.block_2(h) + h = h.to(upscale_dtype) + for i_level in reversed(range(self._num_resolutions)): + for i_block in range(self._num_res_blocks + 1): + h = self.up[i_level].block[i_block](h) + if len(self.up[i_level].attn) > 0: + h = self.up[i_level].attn[i_block](h) + if i_level != 0: + h = self.up[i_level].upsample(h) + h = _swish(self.norm_out(h)) + return self.conv_out(h) + + +class CheersVAEModel(nn.Module): + """VAE model with encoder only (for image understanding).""" + + def __init__(self, config): + super().__init__() + enc_cfg = _cfg(config, "vae_encoder_config") + self.encoder = CheersVAEEncoder(enc_cfg) + self.ps = [2, 2] + z_ch = _cfg(enc_cfg, "z_channels", _VAE_ENCODER_DEFAULTS) + self.bn = nn.BatchNorm2d( + math.prod(self.ps) * z_ch, + eps=1e-4, + momentum=0.1, + affine=False, + track_running_stats=True, + ) + + def encode(self, x: torch.Tensor) -> torch.Tensor: + self.bn.eval() + moments = self.encoder(x) + mean = torch.chunk(moments, 2, dim=1)[0] + z = rearrange( + mean, + "... c (i pi) (j pj) -> ... (c pi pj) i j", + pi=self.ps[0], + pj=self.ps[1], + ) + return self.bn(z) + + +class CheersVAEDecoderProjector(nn.Module): + """VAE decoder projector that converts latent back to pixel-like space.""" + + def __init__(self, config): + super().__init__() + dec_cfg = _cfg(config, "vae_decoder_config") + enc_cfg = _cfg(config, "vae_encoder_config") + self.decoder = CheersVAEDecoder(dec_cfg) + self.ps = [2, 2] + z_ch = _cfg(enc_cfg, "z_channels", _VAE_ENCODER_DEFAULTS) + self.bn = nn.BatchNorm2d( + math.prod(self.ps) * z_ch, + eps=1e-4, + momentum=0.1, + affine=False, + track_running_stats=True, + ) + + def forward(self, z: torch.Tensor) -> torch.Tensor: + self.bn.eval() + s = torch.sqrt(self.bn.running_var.view(1, -1, 1, 1) + 1e-4) + m = self.bn.running_mean.view(1, -1, 1, 1) + z = z * s + m + z = rearrange( + z, + "... (c pi pj) i j -> ... c (i pi) (j pj)", + pi=self.ps[0], + pj=self.ps[1], + ) + return self.decoder(z) + + +class CheersImagePixelInputs(TensorSchema): + """ + Dimensions: + - bn: Batch size * number of images + - c: Number of channels (3) + - h: Height of each image + - w: Width of each image + """ + + type: Literal["pixel_values"] + pixel_values: torch.Tensor # Shape: (bn, 3, h, w) + + +CheersImageInputs: TypeAlias = CheersImagePixelInputs + + +class CheersUndProjector(nn.Module): + """Understanding projector that maps vision features to LLM dimension + with 2x2 spatial compression (4x token reduction).""" + + def __init__( + self, + image_embed_dim: int, + text_embed_dim: int, + compression_factor: tuple[int, int] = (2, 2), + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ): + super().__init__() + self.image_embed_dim = image_embed_dim + self.text_embed_dim = text_embed_dim + self.compression_factor = compression_factor + self.layernorm = nn.LayerNorm(image_embed_dim) + hidden_size = image_embed_dim * (compression_factor[0] * compression_factor[1]) + self.mlp = nn.Sequential( + nn.Linear(hidden_size, hidden_size), + nn.GELU(), + nn.Linear(hidden_size, text_embed_dim), + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.layernorm(x) + height = width = int(x.size(1) ** 0.5) + x = x.permute(0, 2, 1).unflatten(-1, (height, width)) + batch_size, dim, height, width = x.shape + unfolded = x.unfold( + 2, self.compression_factor[0], self.compression_factor[0] + ).unfold(3, self.compression_factor[1], self.compression_factor[1]) + unfolded = unfolded.contiguous().view( + batch_size, + dim, + -1, + self.compression_factor[0] * self.compression_factor[1], + ) + unfolded = ( + unfolded.permute(0, 2, 3, 1) + .contiguous() + .view( + batch_size, + -1, + dim * self.compression_factor[0] * self.compression_factor[1], + ) + ) + return self.mlp(unfolded) + + +class CheersProcessingInfo(BaseProcessingInfo): + """Processing information for Cheers model.""" + + def get_hf_processor(self, **kwargs: object) -> CheersProcessor: + from vllm.transformers_utils.processor import cached_get_image_processor + + image_processor = cached_get_image_processor( + self.ctx.model_config.model, + revision=self.ctx.model_config.revision, + trust_remote_code=self.ctx.model_config.trust_remote_code, + ) + + tokenizer = self.get_tokenizer() + + return CheersProcessor( + image_processor=image_processor, + tokenizer=tokenizer, + **kwargs, + ) + + def get_supported_mm_limits(self) -> Mapping[str, int | None]: + return {"image": None} + + def get_mm_max_tokens_per_item( + self, + seq_len: int, + mm_counts: Mapping[str, int], + ) -> Mapping[str, int]: + hf_config = self.get_hf_config() + vit_config = hf_config.vision_representation_config + patch_size = vit_config.patch_size + image_size = vit_config.image_size + num_patches = (image_size // patch_size) ** 2 + # After 2x2 compression, tokens reduce by 4x + num_tokens = num_patches // 4 + return {"image": num_tokens} + + def get_num_image_tokens( + self, + *, + image_width: int, + image_height: int, + ) -> int: + hf_config = self.get_hf_config() + vit_config = hf_config.vision_representation_config + patch_size = vit_config.patch_size + image_size = vit_config.image_size + num_patches = (image_size // patch_size) ** 2 + return num_patches // 4 + + +class CheersDummyInputsBuilder(BaseDummyInputsBuilder[CheersProcessingInfo]): + """Build dummy inputs for Cheers model profiling.""" + + def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str: + num_images = mm_counts.get("image", 0) + return "<|image_pad|>" * num_images + + def get_dummy_mm_data( + self, + seq_len: int, + mm_counts: Mapping[str, int], + mm_options: Mapping[str, BaseDummyOptions] | None = None, + ) -> MultiModalDataDict: + num_images = mm_counts.get("image", 0) + hf_config = self.info.get_hf_config() + vit_config = hf_config.vision_representation_config + image_size = vit_config.image_size + image_overrides = mm_options.get("image") if mm_options else None + + return { + "image": self._get_dummy_images( + width=image_size, + height=image_size, + num_images=num_images, + overrides=image_overrides, + ), + } + + +class CheersMultiModalProcessor(BaseMultiModalProcessor[CheersProcessingInfo]): + """Multimodal processor for Cheers model.""" + + def _call_hf_processor( + self, + prompt: str, + mm_data: Mapping[str, object], + mm_kwargs: Mapping[str, object], + tok_kwargs: Mapping[str, object], + ) -> BatchFeature: + return super()._call_hf_processor(prompt, mm_data, mm_kwargs, tok_kwargs) + + def _hf_processor_applies_updates( + self, + prompt_text: str, + mm_items: MultiModalDataItems, + hf_processor_mm_kwargs: Mapping[str, object], + tokenization_kwargs: Mapping[str, object], + ) -> bool: + return False + + def _get_prompt_updates( + self, + mm_items: MultiModalDataItems, + hf_processor_mm_kwargs: Mapping[str, Any], + out_mm_kwargs: MultiModalKwargsItems, + ) -> Sequence[PromptReplacement]: + hf_config = self.info.get_hf_config() + vit_config = hf_config.vision_representation_config + patch_size = vit_config.patch_size + image_size = vit_config.image_size + + tokenizer = self.info.get_tokenizer() + image_token_id = tokenizer.get_vocab().get("<|image_pad|>") + if image_token_id is None: + raise ValueError( + "Image token '<|image_pad|>' not found in tokenizer vocabulary" + ) + + def get_replacement_cheers(item_idx: int): + num_patches = (image_size // patch_size) ** 2 + num_tokens = num_patches // 4 + return [image_token_id] * num_tokens + + return [ + PromptReplacement( + modality="image", + target=[image_token_id], + replacement=get_replacement_cheers, + ) + ] + + def _get_mm_fields_config( + self, + hf_inputs: Any, + hf_processor_mm_kwargs: Mapping[str, object], + ) -> Mapping[str, MultiModalFieldConfig]: + return { + "pixel_values": MultiModalFieldConfig.batched("image"), + } + + +@MULTIMODAL_REGISTRY.register_processor( + CheersMultiModalProcessor, + info=CheersProcessingInfo, + dummy_inputs=CheersDummyInputsBuilder, +) +class CheersForConditionalGeneration( + nn.Module, SupportsMultiModal, SupportsLoRA, SupportsPP +): + """ + Cheers: A unified multimodal model for image understanding and generation. + + For vLLM, we focus on the image understanding (vision-to-text) capabilities. + The image generation part is not supported in vLLM. + """ + + requires_raw_input_tokens = True + + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_prefix={ + "model.language_model.": "language_model.model.", + "model.vision_representation.": "vision_representation.vision_model.", + "model.und_projector.": "und_projector.", + "model.vae_model.": "vae_model.", + "model.vae_decoder_projector.": "vae_decoder_projector.", + "lm_head.": "language_model.lm_head.", + } + ) + + @classmethod + def get_placeholder_str(cls, modality: str, i: int) -> str | None: + if modality.startswith("image"): + return "<|image_pad|>" + raise ValueError("Only image modality is supported") + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + + config = vllm_config.model_config.hf_config + quant_config = vllm_config.quant_config + multimodal_config = vllm_config.model_config.multimodal_config + + if type(config).__name__ not in ("CheersConfig", "UMMConfig"): + raise ValueError( + f"Expected CheersConfig or UMMConfig, got {type(config).__name__}." + ) + + self.config = config + self.multimodal_config = multimodal_config + + # The Cheers model's custom Qwen2Config defaults rope_theta to + # 1_000_000, but this isn't stored in the JSON. vLLM's standard + # Qwen2Config defaults to 10_000, causing a 100× mismatch. + # We must patch BOTH the attribute AND rope_parameters (which + # patch_rope_parameters may have already populated from the wrong + # default before __init__ runs). + _CHEERS_ROPE_THETA = 1_000_000.0 + tc = config.text_config + old_theta = getattr(tc, "rope_theta", None) + if old_theta != _CHEERS_ROPE_THETA: + logger.info( + "Overriding text_config.rope_theta from %s to %s", + old_theta, + _CHEERS_ROPE_THETA, + ) + tc.rope_theta = _CHEERS_ROPE_THETA + rp = getattr(tc, "rope_parameters", None) + if rp is not None and rp.get("rope_theta") != _CHEERS_ROPE_THETA: + logger.info( + "Overriding rope_parameters.rope_theta from %s to %s", + rp.get("rope_theta"), + _CHEERS_ROPE_THETA, + ) + rp["rope_theta"] = _CHEERS_ROPE_THETA + + with self._mark_language_model(vllm_config): + self.language_model = init_vllm_registered_model( + vllm_config=vllm_config, + hf_config=config.text_config, + prefix=maybe_prefix(prefix, "language_model"), + architectures=["Qwen2ForCausalLM"], + ) + + vit_config = config.vision_representation_config + + with self._mark_tower_model(vllm_config, "image"): + self.vae_model = CheersVAEModel(config) + self.vae_decoder_projector = CheersVAEDecoderProjector(config) + + self.vision_representation = SiglipVisionModel( + config=vit_config, + quant_config=quant_config, + prefix=maybe_prefix(prefix, "vision_representation"), + ) + + vit_hidden_size = vit_config.hidden_size + llm_hidden_size = config.text_config.hidden_size + + self.und_projector = CheersUndProjector( + image_embed_dim=vit_hidden_size, + text_embed_dim=llm_hidden_size, + compression_factor=(2, 2), + quant_config=quant_config, + prefix=maybe_prefix(prefix, "und_projector"), + ) + + self.make_empty_intermediate_tensors = ( + self.language_model.make_empty_intermediate_tensors + ) + + def _parse_and_validate_image_input( + self, **kwargs: object + ) -> CheersImageInputs | None: + pixel_values = kwargs.pop("pixel_values", None) + if pixel_values is None: + return None + return CheersImagePixelInputs( + type="pixel_values", + pixel_values=pixel_values, + ) + + def _process_image_input( + self, image_input: CheersImageInputs + ) -> tuple[torch.Tensor, ...]: + """Process image inputs through VAE → SigLIP → projector pipeline. + + HF native path: pixel_values → VAE.encode(t=1.0) → vae_decoder_projector + → SigLIP → und_projector → text-space embeddings + """ + pixel_values = image_input["pixel_values"] + + if pixel_values.ndim == 5: + batch_size, num_images, channels, height, width = pixel_values.shape + pixel_values = pixel_values.reshape( + batch_size * num_images, channels, height, width + ) + + with torch.no_grad(): + vae_dtype = next(self.vae_model.parameters()).dtype + image_latent = self.vae_model.encode(pixel_values.to(dtype=vae_dtype)) + image_pixel_hat = self.vae_decoder_projector(image_latent) + + vision_features = self.vision_representation(image_pixel_hat) + vision_embeds = self.und_projector(vision_features) + + return tuple(vision_embeds) + + def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings: + image_input = self._parse_and_validate_image_input(**kwargs) + if image_input is None: + return [] + return self._process_image_input(image_input) + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + **kwargs: object, + ) -> torch.Tensor | IntermediateTensors: + if intermediate_tensors is not None: + inputs_embeds = None + + hidden_states = self.language_model.model( + input_ids=input_ids, + positions=positions, + intermediate_tensors=intermediate_tensors, + inputs_embeds=inputs_embeds, + ) + return hidden_states + + def compute_logits( + self, + hidden_states: torch.Tensor, + ) -> torch.Tensor | None: + return self.language_model.compute_logits(hidden_states) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + """Load weights, keeping VAE encoder/decoder projector for understanding.""" + skip_prefixes = [ + "model.time_embed.", + "model.gen_projector.", + "model.hi_gate.", + "model.hi_projector.", + "model.vae_model.decoder.", + ] + skip_keywords = [ + "text_loss_fc", + ] + + filtered_weights = [] + for name, tensor in weights: + if any(name.startswith(p) for p in skip_prefixes): + continue + if any(kw in name for kw in skip_keywords): + continue + filtered_weights.append((name, tensor)) + + loader = AutoWeightsLoader(self) + return loader.load_weights(filtered_weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/config.py b/vllm/model_executor/models/config.py index 03b147e5c25..7b4fa9252b5 100644 --- a/vllm/model_executor/models/config.py +++ b/vllm/model_executor/models/config.py @@ -9,6 +9,7 @@ from vllm.utils.math_utils import round_up if TYPE_CHECKING: from vllm.config import ModelConfig, VllmConfig + logger = init_logger(__name__) @@ -52,6 +53,58 @@ class Gemma3TextModelConfig(VerifyAndUpdateConfig): hf_config.is_causal = not hf_config.use_bidirectional_attention +class Gemma4Config(VerifyAndUpdateConfig): + @staticmethod + def verify_and_update_config(vllm_config: "VllmConfig") -> None: + """Force unified attention backend for models with heterogeneous + head dimensions. + + Some Gemma4 variants use different head dimensions for + sliding window (head_dim) vs full attention (global_head_dim) layers. + When global_head_dim > 256, FlashAttention rejects those layers + (head_size <= 256 kernel limit), causing vLLM to select a different + backend for each layer type. This mixed-backend execution produces + numerical divergence and output corruption. + + The fix detects heterogeneous head dimensions from the model config + and forces TRITON_ATTN (which has no head_size ceiling) for all + layers when the user hasn't explicitly chosen a backend. + + TODO: Heterogeneous head_sizes (head_dim != global_head_dim) + require NixlConnector changes to support per-layer KV transfer + with different head dimensions for prefill-decode disaggregation. + """ + hf_text_config = vllm_config.model_config.hf_text_config + head_dim = getattr(hf_text_config, "head_dim", None) + global_head_dim = getattr(hf_text_config, "global_head_dim", None) + + # Only force Triton when head dimensions actually differ AND the + # larger one exceeds FlashAttention's kernel limit (head_size <= 256). + # This avoids unnecessary backend forcing on smaller models where + # the config carries global_head_dim but all layers can still use + # the same FA backend. + max_head_dim = max(head_dim or 0, global_head_dim or 0) + if ( + head_dim is not None + and global_head_dim is not None + and head_dim != global_head_dim + and max_head_dim > 256 + and vllm_config.attention_config.backend is None + ): + from vllm.v1.attention.backends.registry import ( + AttentionBackendEnum, + ) + + vllm_config.attention_config.backend = AttentionBackendEnum.TRITON_ATTN + logger.info( + "Gemma4 model has heterogeneous head dimensions " + "(head_dim=%d, global_head_dim=%d). Forcing TRITON_ATTN " + "backend to prevent mixed-backend numerical divergence.", + head_dim, + global_head_dim, + ) + + class GptOssForCausalLMConfig(VerifyAndUpdateConfig): @staticmethod def verify_and_update_config(vllm_config: "VllmConfig") -> None: @@ -533,6 +586,8 @@ MODELS_CONFIG_MAP: dict[str, type[VerifyAndUpdateConfig]] = { "Ernie4_5_VLMoeForConditionalGeneration": Ernie4_5_VLMoeForConditionalGenerationConfig, # noqa: E501 "FalconMambaForCausalLM": MambaModelConfig, "Gemma3TextModel": Gemma3TextModelConfig, + "Gemma4ForCausalLM": Gemma4Config, + "Gemma4ForConditionalGeneration": Gemma4Config, "GptOssForCausalLM": GptOssForCausalLMConfig, "GteModel": SnowflakeGteNewModelConfig, "GteNewForSequenceClassification": GteNewModelConfig, diff --git a/vllm/model_executor/models/deepseek_mtp.py b/vllm/model_executor/models/deepseek_mtp.py index c75ee1a1bbf..87364b1f85c 100644 --- a/vllm/model_executor/models/deepseek_mtp.py +++ b/vllm/model_executor/models/deepseek_mtp.py @@ -241,6 +241,9 @@ class DeepSeekMTP(nn.Module, DeepseekV2MixtureOfExperts): ("gate_up_proj", "up_proj", 1), ("fused_qkv_a_proj", "q_a_proj", 0), ("fused_qkv_a_proj", "kv_a_proj_with_mqa", 1), + # Fused indexer wk + weights_proj + ("wk_weights_proj", "wk", 0), + ("wk_weights_proj", "weights_proj", 1), ] expert_params_mapping = SharedFusedMoE.make_expert_params_mapping( diff --git a/vllm/model_executor/models/deepseek_v2.py b/vllm/model_executor/models/deepseek_v2.py index f1c4a7b2199..f50e38b60f8 100644 --- a/vllm/model_executor/models/deepseek_v2.py +++ b/vllm/model_executor/models/deepseek_v2.py @@ -639,21 +639,19 @@ class Indexer(nn.Module): quant_config=quant_config, prefix=f"{prefix}.wq_b", ) - self.wk = ReplicatedLinear( + # Fused wk + weights_proj: single GEMM producing [head_dim + n_head]. + # weights_proj does not get quantized, so we run both with quant_config=None + # wk may be upcasted from the default quant; experiments show fusion is always + # faster unless WK proj is in FP4, which is not the case for all known quants. + self.wk_weights_proj = MergedColumnParallelLinear( hidden_size, - self.head_dim, - bias=False, - quant_config=quant_config, - prefix=f"{prefix}.wk", - ) - self.k_norm = LayerNorm(self.head_dim, eps=1e-6) - self.weights_proj = ReplicatedLinear( - hidden_size, - self.n_head, + [self.head_dim, self.n_head], bias=False, quant_config=None, - prefix=f"{prefix}.weights_proj", + disable_tp=True, + prefix=f"{prefix}.wk_weights_proj", ) + self.k_norm = LayerNorm(self.head_dim, eps=1e-6) self.softmax_scale = self.head_dim**-0.5 self.scale_fmt = "ue8m0" @@ -694,7 +692,11 @@ class Indexer(nn.Module): q, [self.rope_dim, self.head_dim - self.rope_dim], dim=-1 ) - k, _ = self.wk(hidden_states) + # Fused wk + weights_proj: one GEMM, then split + kw, _ = self.wk_weights_proj(hidden_states) + k = kw[:, : self.head_dim] + weights_raw = kw[:, self.head_dim :] + k = self.k_norm(k) k_pe, k_nope = torch.split( k, [self.rope_dim, self.head_dim - self.rope_dim], dim=-1 @@ -723,9 +725,8 @@ class Indexer(nn.Module): q_fp8 = q_fp8.view(-1, self.n_head, self.head_dim) q_scale = q_scale.view(-1, self.n_head, 1) - weights, _ = self.weights_proj(hidden_states) weights = ( - weights.unsqueeze(-1) * q_scale * self.softmax_scale * self.n_head**-0.5 + weights_raw.unsqueeze(-1) * q_scale * self.softmax_scale * self.n_head**-0.5 ) weights = weights.squeeze(-1) @@ -1438,6 +1439,13 @@ class DeepseekV2ForCausalLM( ("qkv_proj", "k_proj", "k"), ("qkv_proj", "v_proj", "v"), ] + # Fused indexer wk + weights_proj (shard 0 = wk, shard 1 = weights_proj) + indexer_fused_mapping = [ + ("wk_weights_proj", "wk", 0), + ("wk_weights_proj", "weights_proj", 1), + ] + stacked_params_mapping.extend(indexer_fused_mapping) + if self.use_mha: stacked_params_mapping.extend(mha_params_mapping) else: diff --git a/vllm/model_executor/models/gemma4.py b/vllm/model_executor/models/gemma4.py new file mode 100644 index 00000000000..edb53313499 --- /dev/null +++ b/vllm/model_executor/models/gemma4.py @@ -0,0 +1,1239 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Copyright 2025 The vLLM team. +# Copyright 2025 Google Inc. HuggingFace Inc. team. All rights reserved. +# +# + +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Gemma 4 model implementation for vLLM.""" + +from collections.abc import Iterable +from itertools import islice + +import regex as re +import torch +from torch import nn + +from vllm.compilation.decorators import support_torch_compile +from vllm.config import CacheConfig, VllmConfig +from vllm.distributed import ( + get_pp_group, + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, +) +from vllm.logger import init_logger +from vllm.model_executor.layers.activation import GeluAndMul +from vllm.model_executor.layers.attention import Attention +from vllm.model_executor.layers.fused_moe import FusedMoE, GateLinear +from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.model_executor.layers.linear import ( + ColumnParallelLinear, + MergedColumnParallelLinear, + QKVParallelLinear, + ReplicatedLinear, + RowParallelLinear, +) +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.quantization import QuantizationConfig +from vllm.model_executor.layers.rotary_embedding import get_rope +from vllm.model_executor.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from vllm.model_executor.model_loader.weight_utils import ( + default_weight_loader, + maybe_remap_kv_scale_name, +) +from vllm.sequence import IntermediateTensors + +from .interfaces import MixtureOfExperts, SupportsLoRA, SupportsPP +from .utils import ( + AutoWeightsLoader, + extract_layer_index, + is_pp_missing_parameter, + make_layers, + maybe_prefix, +) + +logger = init_logger(__name__) + + +def _get_text_config(config): + """Dereference text_config if config is a nested Gemma4Config. + + Gemma4 checkpoints use architectures=["Gemma4ForConditionalGeneration"] + which yields a Gemma4Config with nested text_config. This function + transparently returns the text config regardless of nesting. + """ + if hasattr(config, "text_config"): + return config.text_config + return config + + +class Gemma4MLP(nn.Module): + def __init__( + self, + hidden_size: int, + intermediate_size: int, + hidden_activation: str, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + self.gate_up_proj = MergedColumnParallelLinear( + hidden_size, + [intermediate_size] * 2, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.gate_up_proj", + ) + self.down_proj = RowParallelLinear( + intermediate_size, + hidden_size, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.down_proj", + ) + if hidden_activation != "gelu_pytorch_tanh": + raise ValueError( + "Gemma4 uses `gelu_pytorch_tanh` as the hidden activation " + "function. Please set `hidden_act` and `hidden_activation` to " + "`gelu_pytorch_tanh`." + ) + self.act_fn = GeluAndMul(approximate="tanh") + + def forward(self, x: torch.Tensor) -> torch.Tensor: + gate_up, _ = self.gate_up_proj(x) + x = self.act_fn(gate_up) + x, _ = self.down_proj(x) + return x + + +class Gemma4Router(nn.Module): + """Router for Gemma4 MoE that preprocesses input before projection. + + Applies RMSNorm (no learned weight), root_size scaling + (hidden_size^{-0.5}), then a learned per-dimension scale before + projecting to expert logits. + + This preprocessing is applied ONLY to the router's input, not to + the expert MLPs' input. + """ + + def __init__( + self, + config, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + self.hidden_size = config.hidden_size + + # RMSNorm without learned weight — pure normalization only + self.norm = RMSNorm(self.hidden_size, eps=config.rms_norm_eps, has_weight=False) + # Per-dimension learned scale, applied after norm + root_size + self.scale = nn.Parameter(torch.ones(self.hidden_size)) + # Constant 1/sqrt(hidden_size) scaling factor + self.register_buffer( + "root_size", + torch.tensor(self.hidden_size**-0.5), + persistent=False, + ) + # Project to expert logits; replicated across TP for consistent routing + # GateLinear supports bf16 W/A → fp32 output, which is important + # because the topk kernel often needs fp32 for stable routing. + self.proj = GateLinear( + self.hidden_size, + config.num_experts, + bias=False, + out_dtype=torch.float32, + prefix=f"{prefix}.proj", + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Returns raw router logits [T, E].""" + x = self.norm(x) + x = x * self.root_size.to(x.dtype) + x = x * self.scale.to(x.dtype) + router_logits, _ = self.proj(x) + return router_logits + + +class Gemma4MoE(nn.Module): + """Mixture of Experts for Gemma4 using vLLM's FusedMoE. + + Wraps FusedMoE with custom routing. The router projection is + external (Gemma4Router) — this class only handles expert dispatch. + + Gemma4 routing: softmax over ALL experts → top-k → renormalize. + per_expert_scale is folded into routing weights for mathematical + correctness with FusedMoE's fused kernel. + """ + + def __init__( + self, + config, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + self.hidden_size = config.hidden_size + self.num_experts = config.num_experts + + # Per-expert output scale folded into routing weights so that + # FusedMoE's fused kernel computes: Σ_e (expert_e * w_e * scale_e) + self.per_expert_scale = nn.Parameter(torch.ones(config.num_experts)) + + # Gemma4 routing: softmax over ALL experts → top-k → renormalize. + # FusedMoE's built-in fused_topk scopes softmax differently, so + # a custom routing function is needed for numerical correctness. + per_expert_scale = self.per_expert_scale + + def routing_function( + hidden_states: torch.Tensor, + gating_output: torch.Tensor, + topk: int, + renormalize: bool, + ) -> tuple[torch.Tensor, torch.Tensor]: + _, topk_ids = torch.topk(gating_output, k=topk, dim=-1) + router_probabilities = torch.nn.functional.softmax(gating_output, dim=-1) + indicator = torch.nn.functional.one_hot( + topk_ids, num_classes=gating_output.size(-1) + ).sum(dim=-2) + gate_weights = indicator * router_probabilities + renorm_factor = torch.sum(gate_weights, dim=-1, keepdim=True) + renorm_factor = torch.where(renorm_factor > 0.0, renorm_factor, 1.0) + dispatch_weights = gate_weights / renorm_factor + + topk_weights = dispatch_weights.gather(1, topk_ids) + + # Fold per_expert_scale into routing weights + expert_scales = per_expert_scale[topk_ids].to(topk_weights.dtype) + topk_weights = topk_weights * expert_scales + return topk_weights.to(torch.float32), topk_ids.to(torch.int32) + + # FusedMoE experts with custom Gemma4 routing + self.experts = FusedMoE( + num_experts=config.num_experts, + top_k=config.top_k_experts, + hidden_size=config.hidden_size, + intermediate_size=getattr( + config, + "moe_intermediate_size", + getattr(config, "expert_intermediate_size", None), + ), + reduce_results=True, + renormalize=True, + quant_config=quant_config, + prefix=f"{prefix}.experts", + custom_routing_function=routing_function, + activation="gelu", + ) + + def forward(self, x: torch.Tensor, router_logits: torch.Tensor) -> torch.Tensor: + return self.experts(x, router_logits) + + +class Gemma4Attention(nn.Module): + def __init__( + self, + config, + hidden_size: int, + num_heads: int, + num_kv_heads: int, + head_dim: int, + max_position_embeddings: int, + use_k_eq_v: bool = False, + cache_config: CacheConfig | None = None, + quant_config: QuantizationConfig | None = None, + attn_logits_soft_cap: float | None = None, + prefix: str = "", + ) -> None: + super().__init__() + self.config = config + self.hidden_size = hidden_size + self.use_k_eq_v = use_k_eq_v + + tp_size = get_tensor_model_parallel_world_size() + self.tp_rank = get_tensor_model_parallel_rank() + self.total_num_heads = num_heads + assert self.total_num_heads % tp_size == 0 + self.num_heads = self.total_num_heads // tp_size + self.total_num_kv_heads = num_kv_heads + if self.total_num_kv_heads >= tp_size: + assert self.total_num_kv_heads % tp_size == 0 + else: + assert tp_size % self.total_num_kv_heads == 0 + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) + self.head_dim = head_dim + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + # Gemma4 uses scaling=1.0. + # Unlike Gemma2/3, query_pre_attn_scalar is NOT used here; + # Q/K norms with learnable weights handle scaling implicitly. + self.scaling = 1.0 + + # QKVParallelLinear handles GQA correctly for all layer types. + # k_eq_v layers load K weights into both K and V slots via + # _weight_iterator remapping — no structural difference needed. + self.qkv_proj = QKVParallelLinear( + hidden_size, + self.head_dim, + self.total_num_heads, + self.total_num_kv_heads, + bias=config.attention_bias, + quant_config=quant_config, + prefix=f"{prefix}.qkv_proj", + ) + self.o_proj = RowParallelLinear( + self.total_num_heads * self.head_dim, + hidden_size, + bias=config.attention_bias, + quant_config=quant_config, + prefix=f"{prefix}.o_proj", + ) + + # Q/K norms: output = norm(x) * weight (learnable per-head scale) + self.q_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.k_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) + # V norm: no learnable scale (pure normalization only) + self.v_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps, has_weight=False) + + # Determine layer type and sliding window + layer_idx = extract_layer_index(prefix) + layer_type = config.layer_types[layer_idx] + self.is_sliding = layer_type == "sliding_attention" + sliding_window = config.sliding_window if self.is_sliding else None + + # Initialize RoPE based on layer type. + # Gemma4 uses different RoPE parameters for sliding vs full attention. + if layer_type in config.rope_parameters: + # Per-layer-type rope config (dict format). + # rope_parameters already contains the correct + # partial_rotary_factor per layer type (1.0 for full + # attention, 1.0 for sliding). Do NOT override with + # global_partial_rotary_factor — that config key is + # not needed for Gemma4 — config uses per-layer rope_parameters. + rope_parameters = dict(config.rope_parameters[layer_type]) + else: + # Legacy config format fallback. + rope_parameters = dict(config.rope_parameters.copy()) + if self.is_sliding: + rope_parameters["rope_theta"] = getattr( + config, "rope_local_base_freq", 10000.0 + ) + + # KV sharing: layers in the last `num_kv_shared_layers` share KV + # cache with earlier layers of the same type. + kv_sharing_target_layer_name = None + self.is_kv_shared_layer = False + num_kv_shared_layers = getattr(config, "num_kv_shared_layers", 0) + if num_kv_shared_layers > 0: + first_kv_shared_layer_idx = config.num_hidden_layers - num_kv_shared_layers + if layer_idx >= first_kv_shared_layer_idx: + self.is_kv_shared_layer = True + # Find the last non-shared layer of the same attention type + prev_layers = config.layer_types[:first_kv_shared_layer_idx] + current_layer_type = config.layer_types[layer_idx] + kv_shared_layer_index = ( + len(prev_layers) - 1 - prev_layers[::-1].index(current_layer_type) + ) + if kv_shared_layer_index >= 0: + if ".layers." in prefix: + param_name_before_layers = prefix.split(".layers.")[0] + else: + raise ValueError( + "Unexpected prefix format for Gemma4Attention: " + f"'{prefix}'. Expected to contain '.layers.'." + ) + kv_sharing_target_layer_name = ( + f"{param_name_before_layers}.layers." + f"{kv_shared_layer_index}.self_attn.attn" + ) + + self.rotary_emb = get_rope( + self.head_dim, + max_position=max_position_embeddings, + rope_parameters=rope_parameters, + is_neox_style=True, + ) + + self.attn = Attention( + self.num_heads, + self.head_dim, + self.scaling, + num_kv_heads=self.num_kv_heads, + cache_config=cache_config, + quant_config=quant_config, + logits_soft_cap=attn_logits_soft_cap, + per_layer_sliding_window=sliding_window, + kv_sharing_target_layer_name=kv_sharing_target_layer_name, + prefix=f"{prefix}.attn", + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + **kwargs, + ) -> torch.Tensor: + # Unified QKV path (works for both k_eq_v and standard layers). + # For k_eq_v, K weights are loaded into both K and V slots of + # qkv_proj, so V == K automatically. + qkv, _ = self.qkv_proj(hidden_states) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + + # Q norm (always applied) + q = q.unflatten(-1, (self.num_heads, self.head_dim)) + q = self.q_norm(q) + q = q.flatten(-2, -1) + + if not self.is_kv_shared_layer: + # Non-shared: apply K norm + RoPE, V norm + k = k.unflatten(-1, (self.num_kv_heads, self.head_dim)) + k = self.k_norm(k) + k = k.flatten(-2, -1) + q, k = self.rotary_emb(positions, q, k) + + v = v.unflatten(-1, (self.num_kv_heads, self.head_dim)) + v = self.v_norm(v) + v = v.flatten(-2, -1) + else: + # Shared: only apply RoPE to Q + q = self.rotary_emb(positions, q, k)[0] + + attn_output = self.attn(q, k, v) + output, _ = self.o_proj(attn_output) + + return output + + +class Gemma4DecoderLayer(nn.Module): + def __init__( + self, + config, + cache_config: CacheConfig | None = None, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + self.hidden_size = config.hidden_size + self.hidden_size_per_layer_input = getattr( + config, "hidden_size_per_layer_input", 0 + ) + + layer_idx = extract_layer_index(prefix) + self.layer_idx = layer_idx + + # Gemma4 uses different head dimensions for sliding vs full attention + layer_type = config.layer_types[layer_idx] + self.is_full_attention = layer_type == "full_attention" + if self.is_full_attention: + head_dim = getattr(config, "global_head_dim", config.head_dim) + else: + head_dim = config.head_dim + + # Determine if this full-attention layer uses k_eq_v + # (laptop variant: no v_proj, K reused as V on full attention layers) + use_k_eq_v = self.is_full_attention and getattr( + config, "attention_k_eq_v", False + ) + + # For k_eq_v full-attention layers, use num_global_key_value_heads + # as the KV head count when k_eq_v is enabled. + if use_k_eq_v: + num_kv_heads = getattr( + config, "num_global_key_value_heads", config.num_key_value_heads + ) + else: + num_kv_heads = config.num_key_value_heads + + self.self_attn = Gemma4Attention( + config=config, + hidden_size=self.hidden_size, + num_heads=config.num_attention_heads, + num_kv_heads=num_kv_heads, + head_dim=head_dim, + max_position_embeddings=config.max_position_embeddings, + use_k_eq_v=use_k_eq_v, + cache_config=cache_config, + quant_config=quant_config, + attn_logits_soft_cap=getattr(config, "attn_logit_softcapping", None), + prefix=f"{prefix}.self_attn", + ) + + # Compute per-layer intermediate_size from config. + # When use_double_wide_mlp is set, intermediate_size doubles for + # KV-shared layers (layers >= first_kv_shared_layer_idx). + first_kv_shared_layer_idx = config.num_hidden_layers - getattr( + config, "num_kv_shared_layers", 0 + ) + is_kv_shared_layer = layer_idx >= first_kv_shared_layer_idx > 0 + use_double_wide_mlp = ( + getattr(config, "use_double_wide_mlp", False) and is_kv_shared_layer + ) + layer_intermediate_size = config.intermediate_size * ( + 2 if use_double_wide_mlp else 1 + ) + + self.mlp = Gemma4MLP( + hidden_size=self.hidden_size, + intermediate_size=layer_intermediate_size, + hidden_activation=config.hidden_activation, + quant_config=quant_config, + prefix=f"{prefix}.mlp", + ) + + # Layer norms: output = norm(x) * weight + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + self.pre_feedforward_layernorm = RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + self.post_feedforward_layernorm = RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + + # MoE (Mixture of Experts) — router + expert block parallel to MLP + self.enable_moe_block = getattr(config, "enable_moe_block", False) or getattr( + config, "use_second_mlp_block", False + ) + if self.enable_moe_block: + self.router = Gemma4Router( + config, + quant_config=quant_config, + prefix=f"{prefix}.router", + ) + self.moe = Gemma4MoE( + config, + quant_config=quant_config, + prefix=f"{prefix}.moe", + ) + self.post_feedforward_layernorm_1 = RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + self.post_feedforward_layernorm_2 = RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + self.pre_feedforward_layernorm_2 = RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + else: + self.router = None + self.moe = None + self.post_feedforward_layernorm_1 = None + self.post_feedforward_layernorm_2 = None + self.pre_feedforward_layernorm_2 = None + + # Per-Layer Embedding (PLE) components — present in each decoder layer + if ( + self.hidden_size_per_layer_input is not None + and self.hidden_size_per_layer_input > 0 + ): + # Gate: projects hidden_states → per-layer dim for gating + self.per_layer_input_gate = ReplicatedLinear( + self.hidden_size, + self.hidden_size_per_layer_input, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.per_layer_input_gate", + return_bias=False, + ) + # Projection: projects gated per-layer input back → hidden size + self.per_layer_projection = ReplicatedLinear( + self.hidden_size_per_layer_input, + self.hidden_size, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.per_layer_projection", + return_bias=False, + ) + # Post-PLE norm: output = norm(x) * weight + self.post_per_layer_input_norm = RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + else: + self.per_layer_input_gate = None + self.per_layer_projection = None + self.post_per_layer_input_norm = None + + # Layer scalar (loaded from checkpoint) — applies to ALL text layers + self.register_buffer("layer_scalar", torch.ones(1)) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + residual: torch.Tensor | None, + per_layer_input: torch.Tensor | None = None, + **kwargs, + ) -> tuple[torch.Tensor, torch.Tensor]: + # Gemma4 residual pattern: + # 1. input_norm(x) → attn → post_attn_norm → ADD residual + # 2. pre_ff_norm → mlp → post_ff_norm → ADD residual + residual = hidden_states + + hidden_states = self.input_layernorm(residual) + + hidden_states = self.self_attn( + positions=positions, + hidden_states=hidden_states, + **kwargs, + ) + + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = hidden_states + residual + residual = hidden_states + + # MLP runs unconditionally (same inputs for MoE and non-MoE) + hidden_states = self.pre_feedforward_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + + if self.enable_moe_block: + hidden_states_1 = self.post_feedforward_layernorm_1(hidden_states) + + # Router and MoE experts see the residual (pre-MLP state), + # matching the HF transformers forward path + router_logits = self.router(residual) + hidden_states_2 = self.pre_feedforward_layernorm_2(residual) + hidden_states_2 = self.moe(hidden_states_2, router_logits) + hidden_states_2 = self.post_feedforward_layernorm_2(hidden_states_2) + + # Combine MLP and MoE outputs + hidden_states = hidden_states_1 + hidden_states_2 + + hidden_states = self.post_feedforward_layernorm(hidden_states) + hidden_states = hidden_states + residual + + # Apply PLE (Per-Layer Embedding) if configured + if per_layer_input is not None and self.per_layer_input_gate is not None: + gate = self.per_layer_input_gate(hidden_states) + gate = torch.nn.functional.gelu(gate, approximate="tanh") + gated_per_layer = gate * per_layer_input + per_layer_contribution = self.per_layer_projection(gated_per_layer) + per_layer_contribution = self.post_per_layer_input_norm( + per_layer_contribution + ) + hidden_states = hidden_states + per_layer_contribution + + # Apply layer scalar for full-attention layers + # Apply per-layer scalar (all text layers) + hidden_states = hidden_states * self.layer_scalar + + return hidden_states, None + + +@support_torch_compile +class Gemma4Model(nn.Module): + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + config = _get_text_config(vllm_config.model_config.hf_config) + cache_config = vllm_config.cache_config + quant_config = vllm_config.quant_config + self.config = config + self.quant_config = quant_config + + # PLE config values (default to 0 if not present — disables PLE) + self.hidden_size_per_layer_input = getattr( + config, "hidden_size_per_layer_input", 0 + ) + self.vocab_size_per_layer_input = getattr( + config, "vocab_size_per_layer_input", config.vocab_size + ) + + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=f"{prefix}.embed_tokens", + ) + + # Per-Layer Embedding (PLE) components + if ( + self.hidden_size_per_layer_input is not None + and self.hidden_size_per_layer_input > 0 + ): + total_ple_dim = self.hidden_size_per_layer_input * config.num_hidden_layers + self.embed_tokens_per_layer = VocabParallelEmbedding( + self.vocab_size_per_layer_input, + total_ple_dim, + quant_config=quant_config, + prefix=f"{prefix}.embed_tokens_per_layer", + ) + # Scaled embedding factor (from config, not hardcoded) + # Register as buffer so it moves to GPU with the model + # and interacts correctly with torch.compile AOT caching. + self.register_buffer( + "embed_scale_per_layer", + torch.tensor(self.hidden_size_per_layer_input**0.5), + persistent=False, + ) + # Projection: hidden_size → total_ple_dim + # ColumnParallelLinear with gather_output=True + self.per_layer_model_projection = ColumnParallelLinear( + config.hidden_size, + total_ple_dim, + bias=False, + gather_output=True, + return_bias=False, + quant_config=quant_config, + prefix=f"{prefix}.per_layer_model_projection", + ) + # PLE projection norm: output = norm(x) * weight + self.per_layer_projection_norm = RMSNorm( + self.hidden_size_per_layer_input, + eps=config.rms_norm_eps, + ) + # Scale factor for combining projection + per_layer_inputs + # Register as buffer so it moves to GPU with the model + # and interacts correctly with torch.compile AOT caching. + self.register_buffer( + "per_layer_input_scale", + torch.rsqrt(torch.tensor(2.0)), + persistent=False, + ) + # Scaled projection: multiply output by hidden_size**-0.5. + # Register as buffer for GPU placement and torch.compile. + self.register_buffer( + "per_layer_projection_scale", + torch.tensor(config.hidden_size**-0.5), + persistent=False, + ) + else: + self.embed_tokens_per_layer = None + self.embed_scale_per_layer = None + self.per_layer_model_projection = None + self.per_layer_projection_norm = None + self.per_layer_input_scale = None + self.per_layer_projection_scale = None + + self.start_layer, self.end_layer, self.layers = make_layers( + config.num_hidden_layers, + lambda prefix: Gemma4DecoderLayer( + config, + cache_config=cache_config, + quant_config=quant_config, + prefix=prefix, + ), + prefix=f"{prefix}.layers", + ) + # Final norm: output = norm(x) * weight + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + # Embedding scale = sqrt(hidden_size) + # Downcast to model dtype (bfloat16 etc.) for numerical parity + self.register_buffer( + "normalizer", + torch.tensor(config.hidden_size**0.5), + persistent=False, + ) + # Custom factory that includes per_layer_inputs for PLE-enabled PP. + # per_layer_inputs has shape (batch, num_layers, per_layer_dim), + # which differs from the standard (batch, hidden_size) shape, + # so we can't use the default factory. + ple_dim = self.hidden_size_per_layer_input + num_layers = config.num_hidden_layers + hidden_size = config.hidden_size + + def _make_empty_intermediate_tensors( + batch_size: int, + dtype: torch.dtype, + device: torch.device, + ) -> IntermediateTensors: + tensors: dict[str, torch.Tensor] = { + "hidden_states": torch.zeros( + (batch_size, hidden_size), + dtype=dtype, + device=device, + ), + "residual": torch.zeros( + (batch_size, hidden_size), + dtype=dtype, + device=device, + ), + } + if ple_dim and ple_dim > 0: + tensors["per_layer_inputs"] = torch.zeros( + (batch_size, num_layers, ple_dim), + dtype=dtype, + device=device, + ) + return IntermediateTensors(tensors) + + self.make_empty_intermediate_tensors = _make_empty_intermediate_tensors + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.embed_tokens(input_ids) * self.normalizer + + def get_per_layer_inputs(self, input_ids: torch.Tensor) -> torch.Tensor: + """Get per-layer embeddings from embed_tokens_per_layer. + + Returns: + Per-layer embeddings (num_tokens, num_layers, + hidden_size_per_layer_input) + """ + if self.embed_tokens_per_layer is None: + return None + + # Handle out-of-vocab tokens for PLE (vocab_size_per_layer_input may + # be smaller than the main vocab_size). + per_layer_inputs_mask = torch.logical_and( + input_ids >= 0, + input_ids < self.vocab_size_per_layer_input, + ) + per_layer_inputs_tokens = torch.where( + per_layer_inputs_mask, input_ids, torch.zeros_like(input_ids) + ) + + # Get packed per-layer embeddings: (num_tokens, total_ple_dim) + per_layer_embeds = self.embed_tokens_per_layer(per_layer_inputs_tokens) + + # Apply embed_scale (sqrt of per-layer hidden dim) + per_layer_embeds = per_layer_embeds * self.embed_scale_per_layer + + # Reshape to (num_tokens, num_layers, hidden_size_per_layer_input) + per_layer_embeds = per_layer_embeds.reshape( + *input_ids.shape, + self.config.num_hidden_layers, + self.hidden_size_per_layer_input, + ) + return per_layer_embeds + + def project_per_layer_inputs( + self, + inputs_embeds: torch.Tensor, + per_layer_inputs: torch.Tensor | None, + ) -> torch.Tensor: + """Project inputs_embeds and combine with per_layer_inputs. + + Steps: + 1. Project inputs_embeds: hidden_size → total_ple_dim + 2. Scale by hidden_size^{-0.5} + 3. Reshape to (num_tokens, num_layers, per_layer_dim) + 4. Normalize with per_layer_projection_norm + 5. Combine: (projection + per_layer_inputs) * 1/sqrt(2) + """ + if self.per_layer_model_projection is None: + return None + + # Project from hidden_size to total_ple_dim + # Scaled projection: output = linear(input, weight) * scale + per_layer_projection = self.per_layer_model_projection(inputs_embeds) + per_layer_projection = per_layer_projection * self.per_layer_projection_scale + + # Reshape to (num_tokens, num_layers, hidden_size_per_layer_input) + per_layer_projection = per_layer_projection.reshape( + *inputs_embeds.shape[:-1], + self.config.num_hidden_layers, + self.hidden_size_per_layer_input, + ) + + # Normalize + per_layer_projection = self.per_layer_projection_norm(per_layer_projection) + + if per_layer_inputs is None: + return per_layer_projection + + # Combine: (projection + per_layer_inputs) * scale + return (per_layer_projection + per_layer_inputs) * self.per_layer_input_scale + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None, + inputs_embeds: torch.Tensor | None = None, + per_layer_inputs: torch.Tensor | None = None, + **kwargs, + ) -> torch.Tensor | IntermediateTensors: + if get_pp_group().is_first_rank: + if inputs_embeds is not None: + hidden_states = inputs_embeds + # When called from the multimodal wrapper, raw PLE + # embeddings are pre-computed and passed explicitly. + # Project them through per_layer_model_projection. + per_layer_inputs = self.project_per_layer_inputs( + hidden_states, per_layer_inputs + ) + else: + hidden_states = self.embed_input_ids(input_ids) + # Compute per-layer inputs for PLE + per_layer_embeds = self.get_per_layer_inputs(input_ids) + per_layer_inputs = self.project_per_layer_inputs( + hidden_states, per_layer_embeds + ) + residual = None + else: + assert intermediate_tensors is not None + hidden_states = intermediate_tensors["hidden_states"] + residual = intermediate_tensors["residual"] + per_layer_inputs = intermediate_tensors.get("per_layer_inputs") + + for layer_idx, layer in enumerate( + islice(self.layers, self.start_layer, self.end_layer) + ): + # Extract the per-layer embedding for this specific layer + if per_layer_inputs is not None: + actual_layer_idx = self.start_layer + layer_idx + layer_per_input = per_layer_inputs[ + :, actual_layer_idx, : + ] # (num_tokens, per_layer_dim) + else: + layer_per_input = None + hidden_states, residual = layer( + positions, + hidden_states, + residual, + per_layer_input=layer_per_input, + **kwargs, + ) + if not get_pp_group().is_last_rank: + return IntermediateTensors( + { + "hidden_states": hidden_states, + "residual": residual, + "per_layer_inputs": per_layer_inputs, + } + ) + # Gemma4 incorporates residual into hidden_states directly + # Apply norm without residual fusion when possible. + if residual is None: + hidden_states = self.norm(hidden_states) + else: + hidden_states, _ = self.norm(hidden_states, residual) + return hidden_states + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + ("qkv_proj", "q_proj", "q"), + ("qkv_proj", "k_proj", "k"), + ("qkv_proj", "v_proj", "v"), + ("gate_up_proj", "gate_proj", 0), + ("gate_up_proj", "up_proj", 1), + ] + + # MoE expert weight mapping: checkpoint 3D packed tensors are + # exploded in _weight_iterator to per-expert 2D weights like: + # moe.experts.{id}.gate_proj → FusedMoE w1 (shard of w13) + # moe.experts.{id}.up_proj → FusedMoE w3 (shard of w13) + # moe.experts.{id}.down_proj → FusedMoE w2 + # We build the mapping directly since Gemma4 uses bare param + # names (no .weight suffix) unlike standard MoE checkpoints. + num_experts = getattr(self.config, "num_experts", None) or 0 + expert_params_mapping = [ + # (param_name, weight_name, expert_id, shard_id) + ( + "experts.w13_weight" + if proj_name in ["gate_proj", "up_proj"] + else "experts.w2_weight", + f"experts.{expert_id}.{proj_name}", + expert_id, + shard_id, + ) + for expert_id in range(num_experts) + for shard_id, proj_name in [ + ("w1", "gate_proj"), + ("w2", "down_proj"), + ("w3", "up_proj"), + ] + ] + params_dict = dict(self.named_parameters()) + # Include buffers (e.g. layer_scalar) so they can be loaded too + params_dict.update(dict(self.named_buffers())) + loaded_params: set[str] = set() + for name, loaded_weight in weights: + if self.quant_config is not None and ( + scale_name := self.quant_config.get_cache_scale(name) + ): + param = params_dict[scale_name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + loaded_weight = loaded_weight[0] + weight_loader(param, loaded_weight) + loaded_params.add(scale_name) + continue + + if name.endswith((".k_scale", ".v_scale", ".q_scale", ".prob_scale")): + remapped_name = maybe_remap_kv_scale_name(name, params_dict) + if remapped_name is not None and remapped_name in params_dict: + param = params_dict[remapped_name] + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, loaded_weight) + loaded_params.add(remapped_name) + continue + + for param_name, shard_name, shard_id in stacked_params_mapping: + if shard_name not in name: + continue + stacked_name = name.replace(shard_name, param_name) + # k_eq_v layers use separate q_proj/k_proj instead of + # packed qkv_proj. If the stacked param doesn't exist, + # skip this mapping and fall through to direct load. + if stacked_name not in params_dict: + continue + if is_pp_missing_parameter(stacked_name, self): + continue + param = params_dict[stacked_name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + loaded_params.add(stacked_name) + break + else: + for ( + param_name, + weight_name, + expert_id, + shard_id, + ) in expert_params_mapping: + if weight_name not in name: + continue + moe_name = name.replace(weight_name, param_name) + if moe_name not in params_dict: + continue + if is_pp_missing_parameter(moe_name, self): + continue + param = params_dict[moe_name] + # Expert weights are already in the correct + # orientation for FusedMoE after _weight_iterator: + # gate/up: [I, H] → w1/w3 expects [I, H] + # down: [H, I] → w2 expects [H, I] + assert loaded_weight.dim() == 2, ( + f"Expected 2D expert weight for {weight_name}, " + f"got shape {loaded_weight.shape}" + ) + weight_loader = param.weight_loader + weight_loader( + param, + loaded_weight, + weight_name + ".weight", + shard_id=shard_id, + expert_id=expert_id, + ) + loaded_params.add(moe_name) + break + else: + if name.endswith(".bias") and name not in params_dict: + continue + name = maybe_remap_kv_scale_name(name, params_dict) + if name is None: + continue + if is_pp_missing_parameter(name, self): + continue + param = params_dict[name] + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, loaded_weight) + loaded_params.add(name) + + return loaded_params + + +class Gemma4ForCausalLM(nn.Module, SupportsLoRA, SupportsPP, MixtureOfExperts): + # Note: qkv_proj packing applies to non-k_eq_v layers (sliding + # attention and full attention without k_eq_v). k_eq_v layers use + # separate q_proj + k_proj without packing. + packed_modules_mapping = { + "qkv_proj": [ + "q_proj", + "k_proj", + "v_proj", + ], + "gate_up_proj": [ + "gate_proj", + "up_proj", + ], + } + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + config = _get_text_config(vllm_config.model_config.hf_config) + quant_config = vllm_config.quant_config + + super().__init__() + self.config = config + self.quant_config = quant_config + self.model = Gemma4Model( + vllm_config=vllm_config, + prefix=maybe_prefix(prefix, "model"), + ) + + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=maybe_prefix(prefix, "lm_head"), + ) + if config.tie_word_embeddings: + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens) + + self.logits_processor = LogitsProcessor( + config.vocab_size, + soft_cap=getattr(config, "final_logit_softcapping", None), + ) + self.make_empty_intermediate_tensors = ( + self.model.make_empty_intermediate_tensors + ) + + # --- MixtureOfExperts protocol --- + self.expert_weights: list[list[torch.Tensor]] = [] + self.moe_layers: list[nn.Module] = [] + example_moe: Gemma4MoE | None = None + + for layer in self.model.layers: + if hasattr(layer, "moe") and isinstance(layer.moe, Gemma4MoE): + example_moe = layer.moe + self.moe_layers.append(layer.moe.experts) + + self.num_moe_layers = len(self.moe_layers) + + if example_moe is not None: + self.num_logical_experts = example_moe.num_experts + self.num_physical_experts = example_moe.num_experts + self.num_local_physical_experts = example_moe.num_experts + self.num_routed_experts = example_moe.num_experts + else: + self.num_logical_experts = 0 + self.num_physical_experts = 0 + self.num_local_physical_experts = 0 + self.num_routed_experts = 0 + + self.num_expert_groups = 1 + self.num_shared_experts = 0 + self.num_redundant_experts = 0 + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.model.embed_input_ids(input_ids) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + **kwargs, + ) -> torch.Tensor | IntermediateTensors: + hidden_states = self.model( + input_ids, positions, intermediate_tensors, inputs_embeds, **kwargs + ) + return hidden_states + + def compute_logits( + self, + hidden_states: torch.Tensor, + ) -> torch.Tensor | None: + return self.logits_processor(self.lm_head, hidden_states) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + # Checkpoint weight names use "language_model." prefix (from the + # Gemma4ForConditionalGeneration wrapper). Strip it to map to our + # model tree which is just "model.*". + def _weight_iterator(): + use_k_eq_v = getattr(self.config, "attention_k_eq_v", False) + # Build set of k_eq_v layer indices (full_attention layers + # when attention_k_eq_v is enabled). These layers have k_proj + # but no v_proj in checkpoint — we duplicate k_proj as v_proj. + k_eq_v_layer_indices: set[int] = set() + if use_k_eq_v: + for idx, lt in enumerate(self.config.layer_types): + if lt == "full_attention": + k_eq_v_layer_indices.add(idx) + + for name, weight in weights: + # Remap "language_model." → "" to match our model tree. + # Checkpoint: model.language_model.layers.X.* + # Our model: model.layers.X.* + name = name.replace("language_model.", "") + + # Remap new HF checkpoint naming to internal vLLM + # naming: HF moved per_expert_scale to router and + # renamed moe → experts in the MoE block. + name = name.replace( + ".router.per_expert_scale", + ".moe.per_expert_scale", + ) + if ".experts.gate_up_proj" in name: + name = name.replace( + ".experts.gate_up_proj", + ".moe.gate_up_proj", + ) + elif ".experts.down_proj" in name: + name = name.replace( + ".experts.down_proj", + ".moe.down_proj", + ) + + # MoE expert weights: checkpoint stores as 3D packed + # tensors. Explode into per-expert 2D weights for + # FusedMoE weight_loader. + # + # Checkpoint format: + # moe.gate_up_proj: [E, 2*I, H] (fused gate + up) + # moe.down_proj: [E, H, I] + # + # FusedMoE expects per-expert: + # w1 (gate): [I, H] — first half of gate_up + # w3 (up): [I, H] — second half of gate_up + # w2 (down): [H, I] — as-is from checkpoint + # + # No transpose needed: checkpoint orientation already + # matches FusedMoE's expected layout. + if "moe.gate_up_proj" in name and weight.dim() == 3: + num_experts = weight.size(0) + intermediate_size = weight.size(1) // 2 + for expert_id in range(num_experts): + gate_weight = weight[expert_id, :intermediate_size, :] + up_weight = weight[expert_id, intermediate_size:, :] + base = name.replace("moe.", f"moe.experts.{expert_id}.") + yield base.replace("gate_up_proj", "gate_proj"), gate_weight + yield base.replace("gate_up_proj", "up_proj"), up_weight + continue + + if "moe.down_proj" in name and weight.dim() == 3: + num_experts = weight.size(0) + for expert_id in range(num_experts): + expert_name = name.replace("moe.", f"moe.experts.{expert_id}.") + yield expert_name, weight[expert_id] + continue + + # k_eq_v layers: checkpoint has k_proj but no v_proj. + # QKVParallelLinear expects both, so duplicate k_proj + # as v_proj so V gets identical weights to K. + # ONLY for full_attention layers — sliding layers have + # their own real v_proj weights. + if "self_attn.k_proj" in name and k_eq_v_layer_indices: + m = re.search(r"layers\.(\d+)\.", name) + if m and int(m.group(1)) in k_eq_v_layer_indices: + yield name, weight + yield name.replace("k_proj", "v_proj"), weight.clone() + continue + + yield name, weight + + # Skip multimodal weights — handled by the multimodal wrapper. + # Also skip lm_head when weights are tied. + skip = [ + "audio_tower.", + "vision_tower.", + "embed_audio.", + "embed_vision.", + ] + if self.config.tie_word_embeddings: + skip.append("lm_head.") + + loader = AutoWeightsLoader(self, skip_substrs=skip) + return loader.load_weights(_weight_iterator()) diff --git a/vllm/model_executor/models/gemma4_mm.py b/vllm/model_executor/models/gemma4_mm.py new file mode 100644 index 00000000000..fa597fe96a0 --- /dev/null +++ b/vllm/model_executor/models/gemma4_mm.py @@ -0,0 +1,1338 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Gemma 4 multimodal model (image + audio + video support). + +Adds vision tower, audio tower, and multimodal embedders on top of the +text-only Gemma4ForCausalLM. The vision/audio encoders are loaded via +AutoModel.from_config and run in eager mode while the language model uses +the vLLM-optimized path. + +Video support: Gemma4 does **not** have a native video tower. Videos are +decomposed into timestamped image frames (up to 32 frames at 70 soft tokens +each) and fed through the same vision tower as regular images. The +processor inserts ``mm:ss`` timestamps between frames so the model can +reason about temporal order. +""" + +import math +from collections.abc import Iterable, Mapping, Sequence +from typing import Annotated, Any, Literal + +import numpy as np +import torch +from PIL import Image as PILImage +from torch import nn +from transformers import AutoModel, BatchFeature +from transformers.models.gemma4 import ( + Gemma4Config, + Gemma4Processor, + Gemma4VisionConfig, +) +from transformers.models.gemma4.configuration_gemma4 import ( + Gemma4AudioConfig, + Gemma4TextConfig, +) + +from vllm.config import VllmConfig +from vllm.config.multimodal import BaseDummyOptions, VideoDummyOptions +from vllm.inputs import MultiModalDataDict +from vllm.logger import init_logger +from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.model_executor.layers.linear import ReplicatedLinear +from vllm.model_executor.models.gemma4 import Gemma4ForCausalLM +from vllm.model_executor.models.module_mapping import MultiModelKeys +from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.multimodal.inputs import ( + MultiModalFieldConfig, + MultiModalKwargsItems, + VideoItem, +) +from vllm.multimodal.parse import ( + AudioProcessorItems, + ImageProcessorItems, + MultiModalDataItems, + MultiModalDataParser, +) +from vllm.multimodal.processing import BaseDummyInputsBuilder +from vllm.multimodal.processing.processor import ( + BaseMultiModalProcessor, + BaseProcessingInfo, + PromptReplacement, + PromptUpdate, + PromptUpdateDetails, +) +from vllm.sequence import IntermediateTensors +from vllm.utils.tensor_schema import TensorSchema, TensorShape + +from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP +from .utils import ( + AutoWeightsLoader, + WeightsMapper, + init_vllm_registered_model, + maybe_prefix, +) + +logger = init_logger(__name__) + +# Video constants — match transformers Gemma4VideoProcessor defaults. +_VIDEO_MAX_SOFT_TOKENS = 70 # soft tokens per video frame (vs 280 for images) +_VIDEO_MAX_FRAMES = 32 # max sampled frames per video + + +# --------------------------------------------------------------------------- +# Input schema +# --------------------------------------------------------------------------- + + +class Gemma4ImagePixelInputs(TensorSchema): + """ + Pre-patchified image inputs from the Gemma4 image processor. + + Dimensions: + - bn: Batch size * number of images + - np: Number of patches (max_patches = max_soft_tokens * pooling_kernel_size²) + - pp: Patch pixels (patch_size² * 3) + + The HF Gemma4ImageProcessor outputs pixel_values as + (batch, max_patches, patch_pixels) — already patchified with + zero-padding for patches beyond the real image content. + pixel_position_ids provides (x, y) coordinates per patch, + with (-1, -1) for padding patches. + """ + + type: Literal["pixel_values"] = "pixel_values" + pixel_values: Annotated[ + torch.Tensor, + TensorShape("bn", "np", "pp"), + ] + pixel_position_ids: Annotated[ + torch.Tensor, + TensorShape("bn", "np", 2), + ] + + +class Gemma4AudioInputs(TensorSchema): + """ + Dimensions: + - bn: Batch size * number of audios + - s: Sequence length (MEL spectrogram frames) + - f: Number of features (MEL bins) + """ + + type: Literal["audio"] = "audio" + input_features_padded: Annotated[torch.Tensor, TensorShape("bn", "s", "f")] + input_features_mask: Annotated[torch.Tensor, TensorShape("bn", "s")] + + +Gemma4ImageInputs = Gemma4ImagePixelInputs + + +class Gemma4VideoInputs(TensorSchema): + """Video frame inputs — same tensor format as image inputs. + + Gemma4 has no separate video tower; video frames are processed + through the vision tower at lower resolution (max_soft_tokens=70). + """ + + type: Literal["pixel_values_videos"] = "pixel_values_videos" + pixel_values_videos: Annotated[ + torch.Tensor, + TensorShape("bn", "np", "pp"), + ] + pixel_position_ids_videos: Annotated[ + torch.Tensor, + TensorShape("bn", "np", 2), + ] + + +# --------------------------------------------------------------------------- +# Processing info +# --------------------------------------------------------------------------- + + +class Gemma4ProcessingInfo(BaseProcessingInfo): + def get_hf_config(self): + return self.ctx.get_hf_config(Gemma4Config) + + def get_default_tok_params(self): + """Gemma4's chat template already embeds a literal ```` token in + the rendered text. If ``add_special_tokens=True`` (the base-class + default), the tokenizer prepends *another* BOS, producing a + ``[2, 2, ...]`` double-BOS sequence that the model was not trained on. + + Setting ``add_special_tokens=False`` here prevents the duplicate and + ensures both ``llm.generate()`` and the chat/completions API behave + correctly. + """ + params = super().get_default_tok_params() + params = params.with_kwargs(add_special_tokens=False) + return params + + def get_hf_processor(self, **kwargs: object) -> Gemma4Processor: + return self.ctx.get_hf_processor( + Gemma4Processor, + **kwargs, + ) + + def validate_num_items(self, modality: str, num_items: int) -> None: + if ( + modality == "audio" + and num_items > 0 + and self.get_hf_config().audio_config is None + ): + model = self.ctx.model_config.model + raise ValueError( + f"Audio input was provided but the model " + f"'{model}' does not have an audio tower. " + f"Audio inference is only supported for Gemma4 " + f"models that include an audio_config " + f"(i.e., models that include an audio_config)." + ) + super().validate_num_items(modality, num_items) + + def get_supported_mm_limits(self) -> Mapping[str, int | None]: + limits: dict[str, int | None] = {"image": None} + if self.get_hf_config().audio_config is not None: + limits["audio"] = None + limits["video"] = None + return limits + + def get_mm_max_tokens_per_item( + self, seq_len: int, mm_counts: Mapping[str, int] + ) -> Mapping[str, int] | None: + config = self.get_hf_config() + # Upper bound: the pooler outputs default_output_length slots + # per image (280). After padding is stripped the actual count + # is ≤ this value, but vLLM needs the max for memory planning. + tokens_per_image = config.vision_config.default_output_length + tokens: dict[str, int] = {"image": tokens_per_image} + if config.audio_config is not None: + # Audio max tokens from the processor's audio_seq_length. + processor = self.get_hf_processor() + tokens["audio"] = processor.audio_seq_length + # Video: each frame ≤ 70 soft tokens + boi + eoi + ~6 ts tokens. + tokens["video"] = _VIDEO_MAX_FRAMES * (_VIDEO_MAX_SOFT_TOKENS + 2 + 6) + return tokens + + def get_data_parser(self) -> MultiModalDataParser: + config = self.get_hf_config() + kwargs: dict[str, Any] = {"video_needs_metadata": True} + if getattr(config, "audio_config", None) is not None: + processor = self.get_hf_processor() + kwargs["target_sr"] = processor.feature_extractor.sampling_rate + return MultiModalDataParser(**kwargs) + + def _compute_num_soft_tokens( + self, + image_width: int, + image_height: int, + max_soft_tokens: int | None = None, + ) -> int: + """Compute the number of soft tokens the vision tower produces + for an image of the given dimensions, after padding is stripped. + + Args: + max_soft_tokens: Override for the vision config's + ``default_output_length``. When *None*, the value from + the model config is used. + """ + vision_cfg = self.get_hf_config().vision_config + patch_size = vision_cfg.patch_size + pooling_kernel_size = vision_cfg.pooling_kernel_size + + if max_soft_tokens is None: + max_soft_tokens = vision_cfg.default_output_length + + unit = patch_size * pooling_kernel_size + max_patches = max_soft_tokens * pooling_kernel_size**2 + num_patches_orig = (image_height / patch_size) * (image_width / patch_size) + scale = math.sqrt(max_patches / num_patches_orig) + target_h = max(unit, int(math.floor(image_height * scale / unit)) * unit) + target_w = max(unit, int(math.floor(image_width * scale / unit)) * unit) + num_patches = (target_h // patch_size) * (target_w // patch_size) + return num_patches // (pooling_kernel_size**2) + + def get_image_repl( + self, + *, + image_width: int, + image_height: int, + processor: Gemma4Processor | None, + max_soft_tokens: int | None = None, + ) -> PromptUpdateDetails[list[int]]: + """Return the dynamic image token sequence for this image. + + Computes the exact number of soft tokens the vision tower will + produce after stripping padding. + + Args: + max_soft_tokens: Override for the default token budget. + When *None*, falls back to the model config value. + """ + if processor is None: + processor = self.get_hf_processor() + + num_soft = self._compute_num_soft_tokens( + image_width, + image_height, + max_soft_tokens=max_soft_tokens, + ) + config = self.get_hf_config() + token_ids = ( + [config.boi_token_id] + + [processor.image_token_id] * num_soft + + [config.eoi_token_id] + ) + return PromptUpdateDetails.select_token_id(token_ids, processor.image_token_id) + + def get_audio_repl( + self, + *, + audio_len: int, + processor: Gemma4Processor | None, + ) -> PromptUpdateDetails[list[int]]: + """Return the dynamic audio token sequence for this audio. + + Computes the number of soft tokens from the audio waveform + length using ``ceil(duration_ms / audio_ms_per_token)``. + """ + if processor is None: + processor = self.get_hf_processor() + + sampling_rate = processor.feature_extractor.sampling_rate + num_tokens = processor._compute_audio_num_tokens( + torch.zeros(audio_len), sampling_rate + ) + config = self.get_hf_config() + token_ids = ( + [config.boa_token_id] + + [processor.audio_token_id] * num_tokens + + [config.eoa_token_id] + ) + return PromptUpdateDetails.select_token_id(token_ids, processor.audio_token_id) + + def get_video_repl( + self, + *, + timestamps: list[float], + num_soft_tokens_per_frame: list[int], + processor: Gemma4Processor, + ) -> PromptUpdateDetails[list[int]]: + """Build the full token replacement for one video. + + Produces the same interleaved sequence as the HF Gemma4Processor: + mm:ss <|video|>*N mm:ss <|video|>*N ... + """ + tokenizer = self.ctx.get_tokenizer() + config = self.get_hf_config() + + boi_token_id = config.boi_token_id + eoi_token_id = config.eoi_token_id + video_token_id = processor.video_token_id + + all_token_ids: list[int] = [] + for i, (ts, n_tokens) in enumerate(zip(timestamps, num_soft_tokens_per_frame)): + # mm:ss timestamp — matches transformers: int-truncated, + # zero-padded. + minutes = int(ts // 60) + seconds = int(ts % 60) + ts_str = f"{minutes:02d}:{seconds:02d}" + + prefix = f" {ts_str} " if i > 0 else f"{ts_str} " + ts_token_ids = tokenizer.encode(prefix, add_special_tokens=False) + all_token_ids.extend(ts_token_ids) + + all_token_ids.append(boi_token_id) + all_token_ids.extend([video_token_id] * n_tokens) + all_token_ids.append(eoi_token_id) + + return PromptUpdateDetails.select_token_id(all_token_ids, video_token_id) + + +# --------------------------------------------------------------------------- +# Dummy inputs builder +# --------------------------------------------------------------------------- + + +class Gemma4DummyInputsBuilder(BaseDummyInputsBuilder[Gemma4ProcessingInfo]): + def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str: + num_images = mm_counts.get("image", 0) + num_audios = mm_counts.get("audio", 0) + num_videos = mm_counts.get("video", 0) + processor = self.info.get_hf_processor() + # Use image_token (<|image|>) with tab prefix — this is what the + # Gemma4 chat template inserts per image (\t<|image|>). + # _get_prompt_updates targets image_token and expands it to the + # full_image_sequence. + text = ("\t" + processor.image_token) * num_images + if num_audios > 0 and processor.audio_token: + text += processor.audio_token * num_audios + if num_videos > 0: + text += processor.video_token * num_videos + return text + + def get_dummy_mm_data( + self, + seq_len: int, + mm_counts: Mapping[str, int], + mm_options: Mapping[str, BaseDummyOptions] | None = None, + ) -> MultiModalDataDict: + num_images = mm_counts.get("image", 0) + num_audios = mm_counts.get("audio", 0) + num_videos = mm_counts.get("video", 0) + processor = self.info.get_hf_processor() + image_processor = processor.image_processor + # Use processor's configured image size for dummies. + # Gemma4ImageProcessor sets size=None (it uses patch_size / + # max_soft_tokens instead of the standard size dict), so we + # guard against None with `or {}`. + size = getattr(image_processor, "size", None) or {} + img_width = size.get("width", 224) + img_height = size.get("height", 224) + + image_overrides = mm_options.get("image") if mm_options else None + audio_overrides = mm_options.get("audio") if mm_options else None + video_overrides = mm_options.get("video") if mm_options else None + + data: MultiModalDataDict = { + "image": self._get_dummy_images( + width=img_width, + height=img_height, + num_images=num_images, + overrides=image_overrides, + ), + } + + if num_audios > 0: + audio_len = processor.feature_extractor.fft_length + data["audio"] = self._get_dummy_audios( + length=audio_len, + num_audios=num_audios, + overrides=audio_overrides, + ) + + if num_videos > 0: + data["video"] = self._get_dummy_videos( + width=img_width, + height=img_height, + num_frames=_VIDEO_MAX_FRAMES, + num_videos=num_videos, + overrides=video_overrides, + ) + + return data + + def _get_dummy_videos( + self, + *, + width: int, + height: int, + num_frames: int, + num_videos: int, + overrides: VideoDummyOptions | None = None, + ) -> list[VideoItem]: + num_frames = max(num_frames, 2) + videos = super()._get_dummy_videos( + width=width, + height=height, + num_frames=num_frames, + num_videos=num_videos, + overrides=overrides, + ) + videos = [v.copy() for v in videos] + + video_items: list[VideoItem] = [] + for video in videos: + video_num_frames = video.shape[0] + video_metadata = { + "fps": 2.0, + "duration": video_num_frames / 2.0, + "total_num_frames": video_num_frames, + "frames_indices": list(range(video_num_frames)), + "video_backend": "opencv", + "do_sample_frames": False, + } + video_items.append((video, video_metadata)) + + return video_items + + +# --------------------------------------------------------------------------- +# Multimodal processor +# --------------------------------------------------------------------------- + + +class Gemma4MultiModalProcessor(BaseMultiModalProcessor[Gemma4ProcessingInfo]): + def _call_hf_processor( + self, + prompt: str, + mm_data: Mapping[str, object], + mm_kwargs: Mapping[str, object], + tok_kwargs: Mapping[str, object], + ) -> BatchFeature: + # Validate max_soft_tokens early and exit cleanly on bad values. + _SUPPORTED_SOFT_TOKENS = (70, 140, 280, 560, 1120) + + merged_kwargs = self.info.ctx.get_merged_mm_kwargs(mm_kwargs) + val = merged_kwargs.get("max_soft_tokens") + if val is None: + val = merged_kwargs.get("images_kwargs", {}).get("max_soft_tokens") + + if val is not None and val not in _SUPPORTED_SOFT_TOKENS: + raise ValueError( + f"Unsupported max_soft_tokens value: {val}. " + f"Valid values are {_SUPPORTED_SOFT_TOKENS}." + ) + + mm_data = dict(mm_data) + + # ---- VIDEO HANDLING ---- + # Gemma4 decomposes video into timestamped image frames. + # Each frame is processed with max_soft_tokens=70 through the + # same vision tower, matching transformers processing_gemma4.py. + video_outputs: dict[str, Any] = {} + if videos := mm_data.pop("videos", []): + processor = self.info.get_hf_processor() + + all_video_pixel_values: list[torch.Tensor] = [] + all_video_position_ids: list[torch.Tensor] = [] + video_num_soft_tokens_per_video: list[list[int]] = [] + video_timestamps_per_video: list[list[float]] = [] + video_frame_counts: list[int] = [] + + for item in videos: + video_array, metadata = item + + # Convert frames to PIL images + if isinstance(video_array, np.ndarray): + frames = [ + PILImage.fromarray(video_array[i]) + for i in range(video_array.shape[0]) + ] + else: + frames = list(video_array) + + # Compute timestamps from metadata (same as transformers) + fps = metadata.get("fps") or 24 + frame_indices = metadata.get("frames_indices", list(range(len(frames)))) + timestamps = [idx / fps for idx in frame_indices] + + # Process frames as images with max_soft_tokens=70 + video_mm_kwargs = dict(mm_kwargs) + video_mm_kwargs["max_soft_tokens"] = _VIDEO_MAX_SOFT_TOKENS + + dummy_prompt = ("\t" + processor.image_token) * len(frames) + + frame_outputs = super()._call_hf_processor( + prompt=dummy_prompt, + mm_data={"images": frames}, + mm_kwargs=video_mm_kwargs, + tok_kwargs=tok_kwargs, + ) + + # Remap HF key name + if "image_position_ids" in frame_outputs: + frame_outputs["pixel_position_ids"] = frame_outputs.pop( + "image_position_ids" + ) + + all_video_pixel_values.append(frame_outputs["pixel_values"]) + all_video_position_ids.append(frame_outputs["pixel_position_ids"]) + + # Compute soft tokens per frame + num_soft_per_frame = [] + for img in frames: + w, h = img.size + n = self.info._compute_num_soft_tokens( + w, h, max_soft_tokens=_VIDEO_MAX_SOFT_TOKENS + ) + num_soft_per_frame.append(n) + + video_num_soft_tokens_per_video.append(num_soft_per_frame) + video_timestamps_per_video.append(timestamps) + video_frame_counts.append(len(frames)) + + # Build expanded replacement text and replace the + # <|video|> placeholder in the prompt. + # Use split(token, 1) to avoid collision — the + # replacement text itself contains <|video|> tokens. + ts_strs = [f"{int(s // 60):02d}:{int(s % 60):02d}" for s in timestamps] + replacement = " ".join( + f"{t} {processor.boi_token}" + f"{processor.video_token * n}" + f"{processor.eoi_token}" + for t, n in zip(ts_strs, num_soft_per_frame) + ) + parts = prompt.split(processor.video_token, 1) + if len(parts) == 2: + prompt = parts[0] + replacement + parts[1] + + video_outputs = { + "pixel_values_videos": torch.cat(all_video_pixel_values, dim=0), + "pixel_position_ids_videos": torch.cat(all_video_position_ids, dim=0), + "video_frame_counts": torch.tensor(video_frame_counts), + "video_num_soft_tokens": video_num_soft_tokens_per_video, + "video_timestamps": video_timestamps_per_video, + } + + # The processor accepts 'audio' not 'audios'. + if "audios" in mm_data: + mm_data["audio"] = mm_data.pop("audios") + + # Warn if any audio waveform exceeds the model's max duration. + if "audio" in mm_data: + processor = self.info.get_hf_processor() + sr = processor.feature_extractor.sampling_rate + max_tokens = processor.audio_seq_length + ms_per_tok = processor.audio_ms_per_token + max_duration_s = max_tokens * ms_per_tok / 1000.0 + audios = mm_data["audio"] + if not isinstance(audios, (list, tuple)): + audios = [audios] + for i, waveform in enumerate(audios): + duration_s = len(waveform) / sr + if duration_s > max_duration_s: + logger.warning( + "Audio duration exceeds max: %f > %f seconds", + duration_s, + max_duration_s, + ) + # vLLM's call_hf_processor (context.py) re-merges + # mm_processor_kwargs from the model config on every call via: + # config_kwargs | incoming_kwargs (right side wins) + # + # If we strip max_soft_tokens from incoming, the re-merge puts + # back the config's global default (e.g. 280), ignoring any + # per-prompt override. Instead, we keep it in the kwargs with + # the validated per-prompt value so it wins during the merge. + # + # NOTE: This requires a corresponding type annotation on the + # HF side (Gemma4ProcessorKwargs.images_kwargs) so that + # _merge_kwargs routes max_soft_tokens into images_kwargs. + patched_mm_kwargs = dict(mm_kwargs) + if val is not None: + patched_mm_kwargs["max_soft_tokens"] = val + + processed_outputs = super()._call_hf_processor( + prompt, + mm_data, + patched_mm_kwargs, + tok_kwargs, + ) + + # HF uses 'image_position_ids'; vLLM uses 'pixel_position_ids'. + # Remap here to keep a single translation point. + if "image_position_ids" in processed_outputs: + processed_outputs["pixel_position_ids"] = processed_outputs.pop( + "image_position_ids" + ) + + if "input_features" in processed_outputs: + # Keep padded features for batched audio tower execution. + processed_outputs["input_features_padded"] = processed_outputs[ + "input_features" + ] + # Unpad per-item so each item's cache entry is self-contained. + unpadded_features = [ + f[mask] + for f, mask in zip( + processed_outputs["input_features"], + processed_outputs["input_features_mask"], + ) + ] + processed_outputs["input_features"] = unpadded_features + + # Merge video outputs into the final result + combined_outputs = dict(processed_outputs, **video_outputs) + return BatchFeature(combined_outputs) + + def _get_mm_fields_config( + self, + hf_inputs: BatchFeature, + hf_processor_mm_kwargs: Mapping[str, object], + ) -> Mapping[str, MultiModalFieldConfig]: + fields = dict( + pixel_values=MultiModalFieldConfig.batched("image"), + pixel_position_ids=MultiModalFieldConfig.batched("image"), + input_features_padded=MultiModalFieldConfig.batched("audio"), + input_features_mask=MultiModalFieldConfig.batched("audio"), + ) + + # Video fields: frames stored flat, split per video by + # video_frame_counts. + video_frame_counts = hf_inputs.get("video_frame_counts") + if video_frame_counts is not None: + vfc = video_frame_counts + if not isinstance(vfc, torch.Tensor): + vfc = torch.tensor(vfc) + fields.update( + pixel_values_videos=( + MultiModalFieldConfig.flat_from_sizes("video", vfc) + ), + pixel_position_ids_videos=( + MultiModalFieldConfig.flat_from_sizes("video", vfc) + ), + video_frame_counts=MultiModalFieldConfig.batched( + "video", + ), + video_num_soft_tokens=MultiModalFieldConfig.batched( + "video", keep_on_cpu=True + ), + video_timestamps=MultiModalFieldConfig.batched( + "video", keep_on_cpu=True + ), + ) + + return fields + + def _get_prompt_updates( + self, + mm_items: MultiModalDataItems, + hf_processor_mm_kwargs: Mapping[str, Any], + out_mm_kwargs: MultiModalKwargsItems, + ) -> Sequence[PromptUpdate]: + hf_processor = self.info.get_hf_processor(**hf_processor_mm_kwargs) + + prompt_updates = [] + + if "image" in mm_items: + # Target image_token (<|image|>) — the single placeholder the + # Gemma4 chat template inserts once per image in the prompt. + # vLLM tokenizes the prompt without token expansion, so only + # one image_token exists per image in the token stream. + # The replacement expands it to the full image sequence + # (boi + N×image_token + eoi, where N = max_soft_tokens). + image_token = hf_processor.image_token + + def get_replacement_image(item_idx: int): + images = mm_items.get_items("image", ImageProcessorItems) + image_size = images.get_image_size(item_idx) + # Resolve the effective max_soft_tokens by merging + # per-prompt kwargs with the config-level defaults, + # consistent with how _call_hf_processor resolves it. + # Without this merge, a missing per-prompt override + # would fall back to vision_cfg.default_output_length + # instead of the config's mm_processor_kwargs default. + merged_kwargs = self.info.ctx.get_merged_mm_kwargs( + hf_processor_mm_kwargs, + ) + max_soft_tokens = merged_kwargs.get("max_soft_tokens") + return self.info.get_image_repl( + image_width=image_size.width, + image_height=image_size.height, + processor=hf_processor, + max_soft_tokens=max_soft_tokens, + ) + + prompt_updates.append( + PromptReplacement( + modality="image", + target=image_token, + replacement=get_replacement_image, + ) + ) + + if "video" in mm_items: + video_token = hf_processor.video_token + + def get_replacement_video(item_idx: int): + out_item = out_mm_kwargs["video"][item_idx] + timestamps = out_item["video_timestamps"].data + num_soft = out_item["video_num_soft_tokens"].data + return self.info.get_video_repl( + timestamps=timestamps, + num_soft_tokens_per_frame=num_soft, + processor=hf_processor, + ) + + prompt_updates.append( + PromptReplacement( + modality="video", + target=video_token, + replacement=get_replacement_video, + ) + ) + + if "audio" in mm_items: + audio_token = hf_processor.audio_token + + def get_replacement_audio(item_idx: int): + audios = mm_items.get_items("audio", AudioProcessorItems) + audio_len = audios.get_audio_length(item_idx) + return self.info.get_audio_repl( + audio_len=audio_len, + processor=hf_processor, + ) + + prompt_updates.append( + PromptReplacement( + modality="audio", + target=audio_token, + replacement=get_replacement_audio, + ) + ) + + return prompt_updates + + # NOTE: Gemma3/Gemma3n override _apply_token_matches and + # _find_mm_placeholders to merge adjacent newline tokens that arise + # when full_image_sequence contains "\n\n" wrappers. Gemma4's + # full_image_sequence has NO newlines (just BOI + 280×image_token + + # EOI), so the base class implementations work correctly as-is. + + +# --------------------------------------------------------------------------- +# Multimodal embedder +# --------------------------------------------------------------------------- + + +class Gemma4MultimodalEmbedder(nn.Module): + """Projects vision/audio soft tokens into LM embedding space. + + Architecture: + inputs_embeds → embedding_projection → embedding_post_projection_norm + + Unlike Gemma3n which has separate hard/soft embedding paths with + per-path normalization and a learned embedding table, Gemma4 uses a + simplified 2-layer design: a linear projection followed by RMSNorm + (without learnable scale). The checkpoint confirms this — only + ``embedding_projection.weight`` exists; there is no embedding table + or pre-projection norm weights. + """ + + def __init__( + self, + multimodal_config: Gemma4VisionConfig | Gemma4AudioConfig, + text_config: Gemma4TextConfig, + ): + super().__init__() + + self.eps = multimodal_config.rms_norm_eps + self.text_hidden_size = text_config.hidden_size + + # Audio tower uses output_proj_dims (1536) rather than hidden_size + # (1024); vision uses hidden_size (768) directly. + embedding_dim = ( + getattr(multimodal_config, "output_proj_dims", None) + or multimodal_config.hidden_size + ) + + self.embedding_projection = ReplicatedLinear( + embedding_dim, + self.text_hidden_size, + bias=False, + ) + + self.embedding_post_projection_norm = RMSNorm( + self.text_hidden_size, + eps=self.eps, + has_weight=False, + ) + + def forward(self, inputs_embeds: torch.Tensor) -> torch.Tensor: + """Project soft tokens from a multimodal tower into LM space.""" + embs_proj, _ = self.embedding_projection(inputs_embeds) + return self.embedding_post_projection_norm(embs_proj) + + +# --------------------------------------------------------------------------- +# Main model +# --------------------------------------------------------------------------- + + +@MULTIMODAL_REGISTRY.register_processor( + Gemma4MultiModalProcessor, + info=Gemma4ProcessingInfo, + dummy_inputs=Gemma4DummyInputsBuilder, +) +class Gemma4ForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP): + packed_modules_mapping = { + "qkv_proj": [ + "q_proj", + "k_proj", + "v_proj", + ], + "gate_up_proj": [ + "gate_proj", + "up_proj", + ], + } + + # Maps checkpoint prefixes to vLLM module paths. + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_prefix={ + "model.embed_audio.": "embed_audio.", + "model.embed_vision.": "embed_vision.", + "model.language_model.": "language_model.model.", + "model.vision_tower.": "vision_tower.", + "model.audio_tower.": "audio_tower.", + "lm_head.": "language_model.lm_head.", + "model": "language_model.model", + } + ) + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + config = vllm_config.model_config.hf_config + quant_config = vllm_config.quant_config + multimodal_config = vllm_config.model_config.multimodal_config + self.config = config + self.quant_config = quant_config + self.multimodal_config = multimodal_config + + # ---- Vision tower (shared by image and video) ---- + with self._mark_tower_model(vllm_config, {"image", "video"}): + self.vision_tower = AutoModel.from_config(config=config.vision_config) + self.embed_vision = Gemma4MultimodalEmbedder( + config.vision_config, config.text_config + ) + + # ---- Audio tower (variants with audio_config) ---- + if config.audio_config is not None: + with self._mark_tower_model(vllm_config, "audio"): + self.audio_tower = AutoModel.from_config(config=config.audio_config) + # AutoModel.from_config does NOT call post_init(), + # which is needed to initialize buffers that are absent + # from the checkpoint (e.g. inv_timescales for relative + # position embeddings, softcap, gradient_clipping). + self.audio_tower.post_init() + self.embed_audio = Gemma4MultimodalEmbedder( + config.audio_config, config.text_config + ) + else: + self.audio_tower = None + self.embed_audio = None + + # ---- Language model (vLLM optimised) ---- + with self._mark_language_model(vllm_config): + self.language_model: Gemma4ForCausalLM = init_vllm_registered_model( + vllm_config=vllm_config, + hf_config=config.text_config, + prefix=maybe_prefix(prefix, "language_model"), + architectures=["Gemma4ForCausalLM"], + ) + + # Pre-allocate PLE buffer for CUDA graph compatibility. + # Some variants have hidden_size_per_layer_input=None (no PLE). + ple_dim = config.text_config.hidden_size_per_layer_input + if ple_dim is not None: + self.per_layer_embeddings = torch.zeros( + vllm_config.scheduler_config.max_num_batched_tokens, + config.text_config.num_hidden_layers, + ple_dim, + device=(self.language_model.model.embed_tokens.weight.device), + dtype=(self.language_model.model.embed_tokens.weight.dtype), + ) + else: + self.per_layer_embeddings = None + + self.make_empty_intermediate_tensors = ( + self.language_model.make_empty_intermediate_tensors + ) + + # --- MixtureOfExperts delegation to language_model --- + self.expert_weights = self.language_model.expert_weights + self.moe_layers = self.language_model.moe_layers + self.num_moe_layers = self.language_model.num_moe_layers + self.num_logical_experts = self.language_model.num_logical_experts + self.num_physical_experts = self.language_model.num_physical_experts + self.num_local_physical_experts = self.language_model.num_local_physical_experts + self.num_routed_experts = self.language_model.num_routed_experts + self.num_expert_groups = self.language_model.num_expert_groups + self.num_shared_experts = self.language_model.num_shared_experts + self.num_redundant_experts = self.language_model.num_redundant_experts + + # ------------------------------------------------------------------ # + # Input parsing + # ------------------------------------------------------------------ # + + def _parse_and_validate_image_input( + self, **kwargs: object + ) -> Gemma4ImageInputs | None: + pixel_values = kwargs.pop("pixel_values", None) + pixel_position_ids = kwargs.pop("pixel_position_ids", None) + image_embeds = kwargs.pop("image_embeds", None) + assert image_embeds is None, "Gemma4 does not support image_embeds." + if pixel_values is None: + return None + return Gemma4ImagePixelInputs( + pixel_values=pixel_values, + pixel_position_ids=pixel_position_ids, + ) + + def _parse_and_validate_audio_input( + self, **kwargs: object + ) -> Gemma4AudioInputs | None: + input_features_padded = kwargs.pop("input_features_padded", None) + if input_features_padded is None: + return None + input_features_mask = kwargs.pop("input_features_mask", None) + if input_features_mask is None: + return None + return Gemma4AudioInputs( + input_features_padded=input_features_padded, + input_features_mask=input_features_mask, + ) + + def _parse_and_validate_video_input( + self, **kwargs: object + ) -> dict[str, torch.Tensor] | None: + pixel_values_videos = kwargs.pop("pixel_values_videos", None) + pixel_position_ids_videos = kwargs.pop("pixel_position_ids_videos", None) + video_frame_counts = kwargs.pop("video_frame_counts", None) + if pixel_values_videos is None: + return None + return { + "pixel_values_videos": pixel_values_videos, + "pixel_position_ids_videos": pixel_position_ids_videos, + "video_frame_counts": video_frame_counts, + } + + def _parse_and_validate_multimodal_inputs( + self, **kwargs: object + ) -> dict[str, Gemma4ImageInputs | Gemma4AudioInputs | Gemma4VideoInputs | None]: + mm_input_by_modality = {} + for input_key in list(kwargs): + if ( + input_key in ("pixel_values", "image_embeds") + and "image" not in mm_input_by_modality + ): + mm_input_by_modality["image"] = self._parse_and_validate_image_input( + **kwargs + ) + if ( + input_key == "pixel_values_videos" + and "video" not in mm_input_by_modality + ): + mm_input_by_modality["video"] = self._parse_and_validate_video_input( + **kwargs + ) + if ( + input_key == "input_features_padded" + and "audio" not in mm_input_by_modality + ): + mm_input_by_modality["audio"] = self._parse_and_validate_audio_input( + **kwargs + ) + return mm_input_by_modality + + # ------------------------------------------------------------------ # + # Image processing + # ------------------------------------------------------------------ # + + def _process_image_input( + self, + image_input: Gemma4ImageInputs, + ) -> list[torch.Tensor]: + pixel_values = image_input["pixel_values"] + pixel_position_ids = image_input["pixel_position_ids"] + + # The HF image processor now outputs pre-patchified data: + # pixel_values: (num_images, max_patches, patch_pixels) + # pixel_position_ids: (num_images, max_patches, 2) + # We call the vision tower's forward() directly, which handles + # patch embedding, encoding, pooling, padding removal, and + # optional standardization internally. + vt = self.vision_tower + pooling_k2 = self.config.vision_config.pooling_kernel_size**2 + + # TODO: Move this per-image loop into the input processor to + # reduce dynamism at the model runner / engine core. This + # requires spatially padding all images to uniform (H_max, + # W_max) in _call_hf_processor() so they arrive as a single + # stacked tensor, tracking padded regions via image_sizes + # metadata, and validating numerical equivalence with the + # current per-image path. + # + # Process each image individually through the vision tower. + # The vision tower's forward() strips padding and returns a + # flat tensor of valid tokens. We process per-image to get + # variable-length outputs matching the dynamic token count + # from get_image_repl. + per_image_features = [] + for i in range(pixel_values.shape[0]): + pv = pixel_values[i].unsqueeze(0) # (1, max_patches, patch_pixels) + pp = pixel_position_ids[i].unsqueeze(0) # (1, max_patches, 2) + + # Derive the pooler's output_length from the total patch + # count (including padding). The vision tower encoder + # processes ALL patches — padding patches get zero hidden + # states but still occupy sequence positions. The pooler's + # _avg_pool_by_positions requires: + # input_seq_len / output_length == k² + # where k == pooling_kernel_size. The image processor + # allocates max_patches = max_soft_tokens * k² total slots, + # so output_length = max_patches / k² == max_soft_tokens. + # Without this, the pooler falls back to + # config.image_seq_length (e.g. 280), which fails when a + # different max_soft_tokens was used at preprocessing time. + max_patches = pv.shape[1] + output_length = max_patches // pooling_k2 + + vt_output = vt(pv, pp, output_length=output_length) + # last_hidden_state: (num_valid_tokens, hidden_size) + # — already flat with padding stripped by the vision tower + per_image_features.append(vt_output.last_hidden_state) + + # Project each image's features into LM embedding space. + # Per-image loop is required because images have variable + # token counts after padding removal. + # Cast to match the projection layer's dtype (model may be + # bf16 while the vision tower outputs fp32). + target_dtype = self.embed_vision.embedding_projection.weight.dtype + return [ + self.embed_vision(inputs_embeds=img.unsqueeze(0).to(target_dtype)).squeeze( + 0 + ) + for img in per_image_features + ] + + # ------------------------------------------------------------------ # + # Video processing (frames through vision tower) + # ------------------------------------------------------------------ # + + def _process_video_input( + self, + video_input: dict[str, torch.Tensor], + ) -> list[torch.Tensor]: + """Process video frames through the vision tower. + + Reuses the image processing pipeline — Gemma4 has no separate + video tower; video frames are just images at lower resolution + (max_soft_tokens=70). + + Returns one concatenated embedding tensor per video (not per + frame), because vLLM treats one video as one multimodal item. + The flat_from_sizes field config groups all frames of a video + together, so embed_multimodal must return one tensor per video. + """ + pixel_values = video_input["pixel_values_videos"] + pixel_position_ids = video_input["pixel_position_ids_videos"] + frame_counts = video_input["video_frame_counts"] + + vt = self.vision_tower + pooling_k2 = self.config.vision_config.pooling_kernel_size**2 + target_dtype = self.embed_vision.embedding_projection.weight.dtype + + # Split flat tensors into per-video chunks + if isinstance(frame_counts, torch.Tensor): + fc_list = frame_counts.tolist() + else: + fc_list = list(frame_counts) + + pv_per_video = torch.split(pixel_values, fc_list, dim=0) + pp_per_video = torch.split(pixel_position_ids, fc_list, dim=0) + + per_video_embeddings = [] + for pv_chunk, pp_chunk in zip(pv_per_video, pp_per_video): + frame_embs = [] + for i in range(pv_chunk.shape[0]): + pv = pv_chunk[i].unsqueeze(0) + pp = pp_chunk[i].unsqueeze(0) + + max_patches = pv.shape[1] + output_length = max_patches // pooling_k2 + + vt_output = vt(pv, pp, output_length=output_length) + frame_emb = self.embed_vision( + inputs_embeds=( + vt_output.last_hidden_state.unsqueeze(0).to(target_dtype) + ) + ).squeeze(0) + frame_embs.append(frame_emb) + + # Concatenate all frames of this video into one tensor. + per_video_embeddings.append(torch.cat(frame_embs, dim=0)) + + return per_video_embeddings + + # ------------------------------------------------------------------ # + # Audio processing + # ------------------------------------------------------------------ # + + def _process_audio_input( + self, + audio_input: Gemma4AudioInputs, + ) -> list[torch.Tensor]: + input_features = audio_input["input_features_padded"].squeeze(1) + input_features_mask = audio_input["input_features_mask"].squeeze(1) + + # Run audio tower — mask uses standard HF convention + # (True=valid, False=padding). + audio_outputs = self.audio_tower(input_features, input_features_mask) + if isinstance(audio_outputs, tuple): + audio_encodings, audio_mask = audio_outputs + else: + audio_encodings = audio_outputs.last_hidden_state + audio_mask = audio_outputs.attention_mask + + # Project into LM embedding space. + audio_features = self.embed_audio(inputs_embeds=audio_encodings) + + # Strip padding per-batch element: only keep real (non-padding) + # tokens. audio_mask is True for valid positions (HF convention). + per_audio = [] + for enc, mask in zip(audio_features, audio_mask, strict=True): + per_audio.append(enc[mask]) # [num_real, hidden_size] + + return per_audio + + # ------------------------------------------------------------------ # + # MultiModalEmbeddings interface + # ------------------------------------------------------------------ # + + def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings: + mm_input_by_modality = self._parse_and_validate_multimodal_inputs(**kwargs) + multimodal_embeddings: list[torch.Tensor] = [] + + for modality, multimodal_input in mm_input_by_modality.items(): + if multimodal_input is None: + continue + if modality == "image": + multimodal_embeddings.extend( + self._process_image_input(multimodal_input) + ) + elif modality == "video": + multimodal_embeddings.extend( + self._process_video_input(multimodal_input) + ) + elif modality == "audio": + multimodal_embeddings.extend( + self._process_audio_input(multimodal_input) + ) + + return multimodal_embeddings + + def embed_input_ids( + self, + input_ids: torch.Tensor, + multimodal_embeddings: MultiModalEmbeddings | None = None, + *, + is_multimodal: torch.Tensor | None = None, + ) -> torch.Tensor: + # Cache per-layer embeddings (PLE) for the language model's + # forward pass. During profiling embed_input_ids is not called, + # so the pre-allocated zeros are used instead. + if self.per_layer_embeddings is not None: + # Mask multimodal tokens (image/audio) to 0 for PLE + # computation (using token_type_ids == 0 as text_mask). + # Replicate this: map image token positions to token 0. + if is_multimodal is not None: + is_multimodal = is_multimodal.to(input_ids.device) + ple_input_ids = torch.where( + is_multimodal, torch.zeros_like(input_ids), input_ids + ) + else: + ple_input_ids = input_ids + + per_layer_inputs = self.language_model.model.get_per_layer_inputs( + ple_input_ids + ) + if per_layer_inputs is not None: + per_layer_inputs = per_layer_inputs.reshape( + -1, + self.config.text_config.num_hidden_layers, + self.config.text_config.hidden_size_per_layer_input, + ) + self.per_layer_embeddings[: per_layer_inputs.shape[0]].copy_( + per_layer_inputs + ) + + if multimodal_embeddings is None or is_multimodal is None: + return super().embed_input_ids(input_ids) + + return super().embed_input_ids( + input_ids, + multimodal_embeddings=multimodal_embeddings, + is_multimodal=is_multimodal, + ) + + # ------------------------------------------------------------------ # + # Forward + # ------------------------------------------------------------------ # + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + **kwargs: object, + ) -> IntermediateTensors: + if intermediate_tensors is not None: + inputs_embeds = None + + # Select the pre-cached PLEs for this batch (None when PLE + # is disabled for variants without PLE). + per_layer_inputs = ( + self.per_layer_embeddings[: inputs_embeds.shape[0]] + if self.per_layer_embeddings is not None and inputs_embeds is not None + else None + ) + + hidden_states = self.language_model.model( + input_ids, + positions, + per_layer_inputs=per_layer_inputs, + intermediate_tensors=intermediate_tensors, + inputs_embeds=inputs_embeds, + **kwargs, + ) + + return hidden_states + + def compute_logits( + self, + hidden_states: torch.Tensor, + ) -> torch.Tensor | None: + return self.language_model.compute_logits(hidden_states) + + # ------------------------------------------------------------------ # + # Weight loading + # ------------------------------------------------------------------ # + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + # Some checkpoints have vestigial embed_vision.embedding and + # embed_audio.embedding weights from the Gemma3n architecture + # that are not used by Gemma4's MultimodalEmbedder (which only + # has embedding_projection + embedding_post_projection_norm). + ignore_prefixes = [ + "embed_vision.embedding.", + "embed_audio.embedding.", + ] + # Models without audio tower should skip + # audio weights entirely. + if self.audio_tower is None: + ignore_prefixes.extend( + [ + "audio_tower.", + "embed_audio.", + ] + ) + loader = AutoWeightsLoader( + self, + ignore_unexpected_prefixes=ignore_prefixes, + ) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) + + # ------------------------------------------------------------------ # + # LoRA / multimodal mapping + # ------------------------------------------------------------------ # + + def get_mm_mapping(self) -> MultiModelKeys: + """Get the module prefix mapping for multimodal models.""" + return MultiModelKeys.from_string_field( + language_model="language_model", + connector=["embed_vision", "embed_audio"], + tower_model=["vision_tower", "audio_tower"], + ) + + @classmethod + def get_placeholder_str(cls, modality: str, i: int) -> str | None: + if modality == "image": + return "" + if modality == "audio": + return "" + if modality == "video": + return "<|video|>" + raise ValueError(f"Unsupported modality: {modality}") diff --git a/vllm/model_executor/models/gpt_oss.py b/vllm/model_executor/models/gpt_oss.py index 482056250a1..a9ec8297422 100644 --- a/vllm/model_executor/models/gpt_oss.py +++ b/vllm/model_executor/models/gpt_oss.py @@ -20,11 +20,12 @@ from vllm.distributed import ( tensor_model_parallel_all_gather, ) from vllm.model_executor.layers.attention import Attention -from vllm.model_executor.layers.fused_moe import FusedMoE, GateLinear +from vllm.model_executor.layers.fused_moe import FusedMoE from vllm.model_executor.layers.fused_moe.config import FusedMoEParallelConfig from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ( QKVParallelLinear, + ReplicatedLinear, RowParallelLinear, ) from vllm.model_executor.layers.logits_processor import LogitsProcessor @@ -174,11 +175,13 @@ class MLPBlock(torch.nn.Module): self.hidden_size = config.hidden_size self.experts_per_token = config.num_experts_per_tok self.world_size = dist.get_world_size() if dist.is_initialized() else 1 - self.router = GateLinear( + self.router = ReplicatedLinear( config.hidden_size, config.num_local_experts, bias=True, + quant_config=None, prefix=f"{prefix}.router", + return_bias=False, ) assert config.intermediate_size % self.world_size == 0 self.experts = FusedMoE( @@ -206,7 +209,7 @@ class MLPBlock(torch.nn.Module): self, x[:, : self.hidden_size], self.router.weight, self.router.bias ) else: - g, _ = self.router(x) + g = self.router(x) x = self.experts(hidden_states=x, router_logits=g)[:, : self.hidden_size] if self.is_sequence_parallel: diff --git a/vllm/model_executor/models/nano_nemotron_vl.py b/vllm/model_executor/models/nano_nemotron_vl.py index 90425417a09..819bb4a3cb2 100644 --- a/vllm/model_executor/models/nano_nemotron_vl.py +++ b/vllm/model_executor/models/nano_nemotron_vl.py @@ -7,7 +7,6 @@ # LICENSE is in root directory. # -------------------------------------------------------- -import copy import math import warnings from collections.abc import Iterable, Mapping, Sequence @@ -17,7 +16,7 @@ from typing import Annotated, Literal, TypeAlias import torch import torch.nn as nn -from transformers import BatchFeature +from transformers import BatchFeature, PretrainedConfig from vllm.config import VllmConfig from vllm.config.multimodal import BaseDummyOptions, VideoDummyOptions @@ -210,11 +209,15 @@ class NanoNemotronVLProcessingInfo(BaseProcessingInfo): @cached_property def is_dynamic_tiler(self) -> bool: - return self.get_hf_processor().dynamic_tiler is not None + return BaseNanoNemotronVLProcessor.use_dynamic_resolution(self.get_hf_config()) - @cached_property + @property def supports_video(self): - return self.get_hf_processor().supports_video + return True + + @property + def supports_audio(self) -> bool: + return self.sound_config is not None def get_video_token(self) -> str | None: return IMG_CONTEXT @@ -223,8 +226,8 @@ class NanoNemotronVLProcessingInfo(BaseProcessingInfo): return self.ctx.get_mm_config().video_pruning_rate @property - def audio_extractor(self) -> ParakeetExtractor | None: - return self.get_hf_processor().audio_extractor + def sound_config(self) -> PretrainedConfig | None: + return getattr(self.get_hf_config(), "sound_config", None) def get_default_tok_params(self) -> TokenizeParams: return super().get_default_tok_params().with_kwargs(add_special_tokens=False) @@ -232,14 +235,14 @@ class NanoNemotronVLProcessingInfo(BaseProcessingInfo): def get_supported_mm_limits(self) -> Mapping[str, int | None]: image_limit = {"image": None} video_limit = {"video": None} if self.supports_video else {} - audio_limit = {"audio": None} if self.audio_extractor is not None else {} + audio_limit = {"audio": None} if self.supports_audio else {} return {**image_limit, **video_limit, **audio_limit} def get_data_parser(self): target_sr = None target_channels = None - if extractor := self.audio_extractor: - target_sr = extractor.sampling_rate + if self.sound_config: + target_sr = self.sound_config.sampling_rate target_channels = 1 return MultiModalDataParser( @@ -371,7 +374,7 @@ class NanoNemotronVLMultiModalProcessor( fields = self._get_image_fields_config(hf_inputs) if self.info.supports_video: fields |= self._get_video_fields_config(hf_inputs) - if self.info.audio_extractor: + if self.info.supports_audio: fields |= self._get_audio_fields_config(hf_inputs) return fields @@ -399,9 +402,8 @@ class NanoNemotronVLMultiModalProcessor( if isinstance(images, ImageEmbeddingItems): feature_size = images.get_feature_size(item_idx) - elif tiler := hf_processor.dynamic_tiler: - image = images.get(item_idx) - feature_size = tiler.get_cached_feature_size(image) + elif self.info.is_dynamic_tiler: + feature_size = out_mm_data["num_tokens_per_image"][item_idx] else: image_size = images.get_image_size(item_idx) max_num_tiles = hf_processor.max_num_tiles @@ -536,7 +538,7 @@ class NanoNemotronVLMultiModalProcessor( prompt_repls.append( self._get_prompt_repl_video(mm_items, hf_processor, out_mm_data) ) - if self.info.audio_extractor: + if self.info.supports_audio: prompt_repls.append( self._get_prompt_repl_audio(mm_items, hf_processor, out_mm_data) ) @@ -772,12 +774,14 @@ class NanoNemotronVLDummyInputsBuilder( else: dummy_video = {} - if extractor := self.info.audio_extractor: + if sound_config := self.info.sound_config: num_audios = mm_counts.get("audio", 0) audio_overrides = mm_options.get("audio") if mm_options else None tokens_per_audio = max(1, seq_len // max(num_audios, 1)) - max_audio_num_samples = MAX_AUDIO_LEN_S * extractor.sampling_rate - calculated_max_audio_num_samples = extractor.audio_length(tokens_per_audio) + max_audio_num_samples = MAX_AUDIO_LEN_S * sound_config.sampling_rate + calculated_max_audio_num_samples = ParakeetExtractor.audio_length( + sound_config, tokens_per_audio + ) audio_len = min(max_audio_num_samples, calculated_max_audio_num_samples) dummy_audio = { "audio": self._get_dummy_audios( @@ -1029,9 +1033,13 @@ class NemotronH_Nano_VL_V2( data=image_embeds, ) + pixel_values_flat = kwargs.pop("pixel_values_flat", None) + if pixel_values_flat is None: + return None + if self.dynamic_resolution: pixel_values_flat = DynamicResolutionImageTiler.stack( - kwargs.pop("pixel_values_flat"), self.patch_size + pixel_values_flat, self.patch_size ) return NanoNemotronVLImagePixelInputsDynamic( pixel_values_flat=pixel_values_flat, **kwargs @@ -1497,15 +1505,13 @@ class NemotronH_Nano_VL_V2( @classmethod def get_mamba_state_shape_from_config(cls, vllm_config: "VllmConfig"): text_config = vllm_config.model_config.hf_config.text_config - temp_vllm_config = copy.deepcopy(vllm_config) - temp_vllm_config.model_config.hf_config = text_config + temp_vllm_config = vllm_config.with_hf_config(text_config) return NemotronHForCausalLM.get_mamba_state_shape_from_config(temp_vllm_config) @classmethod def get_mamba_state_dtype_from_config(cls, vllm_config: "VllmConfig"): text_config = vllm_config.model_config.hf_config.text_config - temp_vllm_config = copy.deepcopy(vllm_config) - temp_vllm_config.model_config.hf_config = text_config + temp_vllm_config = vllm_config.with_hf_config(text_config) return NemotronHForCausalLM.get_mamba_state_dtype_from_config(temp_vllm_config) @classmethod diff --git a/vllm/model_executor/models/olmo_hybrid.py b/vllm/model_executor/models/olmo_hybrid.py index 97e56b3ff6f..d070132fc4e 100644 --- a/vllm/model_executor/models/olmo_hybrid.py +++ b/vllm/model_executor/models/olmo_hybrid.py @@ -68,6 +68,7 @@ from vllm.model_executor.layers.mamba.mamba_utils import ( MambaStateCopyFuncCalculator, MambaStateDtypeCalculator, MambaStateShapeCalculator, + is_conv_state_dim_first, ) from vllm.model_executor.layers.mamba.ops.causal_conv1d import ( causal_conv1d_fn, @@ -429,7 +430,13 @@ class OlmoHybridGatedDeltaNet(nn.Module, MambaBase): spec_state_indices_tensor = attn_metadata.spec_state_indices_tensor non_spec_state_indices_tensor = attn_metadata.non_spec_state_indices_tensor self_kv_cache = self.kv_cache - conv_state = self_kv_cache[0].transpose(-1, -2) + # conv_state must be (..., dim, width-1) for the conv kernels. + # DS layout stores it that way directly; SD layout needs a transpose. + conv_state = ( + self_kv_cache[0] + if is_conv_state_dim_first() + else self_kv_cache[0].transpose(-1, -2) + ) ssm_state = self_kv_cache[1] num_actual_tokens = attn_metadata.num_actual_tokens num_accepted_tokens = attn_metadata.num_accepted_tokens diff --git a/vllm/model_executor/models/parakeet.py b/vllm/model_executor/models/parakeet.py index 1a3fd5bad0c..67c61dd8064 100644 --- a/vllm/model_executor/models/parakeet.py +++ b/vllm/model_executor/models/parakeet.py @@ -159,5 +159,7 @@ class ParakeetExtractor(ParakeetFeatureExtractor): outputs["audio_num_clips"] = audio_num_clips return outputs - def audio_length(self, audio_tokens: int) -> int: - return int(audio_tokens * self.config.subsampling_factor * self.hop_length) + @staticmethod + def audio_length(raw_config: PretrainedConfig, audio_tokens: int) -> int: + config = ExtractorConfig.from_hf_config(raw_config) + return int(audio_tokens * config.subsampling_factor * config.hop_length) diff --git a/vllm/model_executor/models/phi4siglip.py b/vllm/model_executor/models/phi4siglip.py new file mode 100644 index 00000000000..d71a572f6ad --- /dev/null +++ b/vllm/model_executor/models/phi4siglip.py @@ -0,0 +1,429 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""vLLM support for microsoft/Phi-4-reasoning-vision-15B. + +Architecture: Siglip2 vision tower + MLP projector + Phi3 language model. +""" + +import math +from collections.abc import Iterable, Mapping, Sequence +from typing import Annotated, Any, Literal + +import torch +import torch.nn as nn +from transformers import BatchFeature, PretrainedConfig, Siglip2VisionConfig + +from vllm.config import VllmConfig +from vllm.config.multimodal import BaseDummyOptions +from vllm.inputs import MultiModalDataDict +from vllm.logger import init_logger +from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.multimodal.inputs import ( + MultiModalFieldConfig, + MultiModalKwargsItems, +) +from vllm.multimodal.parse import ( + ImageSize, + MultiModalDataItems, +) +from vllm.multimodal.processing import ( + BaseDummyInputsBuilder, + PromptReplacement, + PromptUpdate, +) +from vllm.multimodal.processing.processor import ( + BaseMultiModalProcessor, + BaseProcessingInfo, +) +from vllm.sequence import IntermediateTensors +from vllm.utils.tensor_schema import TensorSchema, TensorShape + +from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP +from .lfm2_siglip2 import Siglip2Model +from .llava import LlavaMultiModalProjector +from .utils import ( + AutoWeightsLoader, + WeightsMapper, + init_vllm_registered_model, + maybe_prefix, +) + +logger = init_logger(__name__) + +IMAGE_TOKEN_INDEX = -200 +DEFAULT_IMAGE_TOKEN = "" + +# The HF processor replaces "" with IMAGE_TOKEN_INDEX (-200) in input_ids. +# Negative token IDs cause OverflowError during decoding, so we remap to a real +# in-vocabulary token. The Phi-4-reasoning-vision tokenizer ships with reserved +# dummy tokens (<|dummy_0|> … <|dummy_83|>); we reuse the first one as the +# image placeholder. This mirrors how Phi-3-vision uses its dedicated <|image|> +# token (ID 32044). +_IMAGE_TOKEN_ID = 100256 # <|dummy_0|> in the Phi-4 tokenizer + + +# --------------------------------------------------------------------------- +# Processing +# --------------------------------------------------------------------------- + + +class Phi4SiglipProcessingInfo(BaseProcessingInfo): + def get_supported_mm_limits(self) -> Mapping[str, int | None]: + return {"image": None} + + def _get_vision_config(self) -> dict: + return self.get_hf_config().vision_config # type: ignore[attr-defined] + + def _get_patch_size(self) -> int: + vc = self._get_vision_config() + if isinstance(vc, dict): + return vc.get("patch_size", 16) + return getattr(vc, "patch_size", 16) + + def _get_max_num_patches(self) -> int: + return getattr(self.get_hf_config(), "max_num_patches", 3600) + + def _get_min_num_patches(self) -> int: + return getattr(self.get_hf_config(), "min_num_patches", 256) + + def get_num_image_tokens( + self, + *, + image_width: int, + image_height: int, + ) -> int: + patch_size = self._get_patch_size() + min_patches = self._get_min_num_patches() + max_patches = self._get_max_num_patches() + + num_patches_h = image_height // patch_size + num_patches_w = image_width // patch_size + num_patches = max(num_patches_h * num_patches_w, 1) + num_patches = max(min(num_patches, max_patches), min_patches) + return num_patches + + def get_image_size_with_most_features(self) -> ImageSize: + patch_size = self._get_patch_size() + max_patches = self._get_max_num_patches() + side = int(math.sqrt(max_patches)) * patch_size + return ImageSize(width=side, height=side) + + def get_mm_max_tokens_per_item( + self, seq_len: int, mm_counts: Mapping[str, int] + ) -> Mapping[str, int]: + return {"image": self._get_max_num_patches()} + + +class Phi4SiglipDummyInputsBuilder( + BaseDummyInputsBuilder[Phi4SiglipProcessingInfo], +): + def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str: + num_images = mm_counts.get("image", 0) + return DEFAULT_IMAGE_TOKEN * num_images + + def get_dummy_mm_data( + self, + seq_len: int, + mm_counts: Mapping[str, int], + mm_options: Mapping[str, BaseDummyOptions], + ) -> MultiModalDataDict: + num_images = mm_counts.get("image", 0) + size = self.info.get_image_size_with_most_features() + return { + "image": self._get_dummy_images( + width=size.width, + height=size.height, + num_images=num_images, + overrides=mm_options.get("image"), + ), + } + + +class Phi4SiglipMultiModalProcessor( + BaseMultiModalProcessor[Phi4SiglipProcessingInfo], +): + def _call_hf_processor( + self, + prompt: str, + mm_data: Mapping[str, object], + mm_kwargs: Mapping[str, object], + tok_kwargs: Mapping[str, object], + ) -> BatchFeature: + processed = super()._call_hf_processor( + prompt=prompt, + mm_data=mm_data, + mm_kwargs=mm_kwargs, + tok_kwargs=tok_kwargs, + ) + + # The HF processor's tokenizer_image_token() replaces the "" + # string with IMAGE_TOKEN_INDEX (-200) in input_ids. This breaks + # vLLM's prompt-replacement pipeline which needs to find "" + # as normal sub-tokens. Re-tokenize with the plain tokenizer so + # that "" stays as sub-tokens and can be located by + # PromptReplacement. + # NOTE: tokenizer.__call__() (not .encode()) must be used so that + # added/special tokens like <|user|>, <|end|> are kept as single IDs. + tokenizer = self.info.get_tokenizer() + new_ids = tokenizer(prompt).input_ids + processed["input_ids"] = torch.tensor([new_ids]) + + return processed + + def _hf_processor_applies_updates( + self, + prompt_text: str, + mm_items: MultiModalDataItems, + hf_processor_mm_kwargs: Mapping[str, object], + tokenization_kwargs: Mapping[str, object], + ) -> bool: + # The HF processor replaces "" with a single -200 placeholder + # but does NOT expand it into N vision-encoder tokens. Since we also + # re-tokenize the prompt (see _call_hf_processor), prompt updates are + # never applied by the HF processor — vLLM handles the expansion via + # _apply_prompt_updates. + return False + + def _get_mm_fields_config( + self, + hf_inputs: BatchFeature, + hf_processor_mm_kwargs: Mapping[str, object], + ) -> Mapping[str, MultiModalFieldConfig]: + return dict( + pixel_values=MultiModalFieldConfig.batched("image"), + pixel_attention_mask=MultiModalFieldConfig.batched("image"), + spatial_shapes=MultiModalFieldConfig.batched("image", keep_on_cpu=True), + ) + + def _get_prompt_updates( + self, + mm_items: MultiModalDataItems, + hf_processor_mm_kwargs: Mapping[str, Any], + out_mm_kwargs: MultiModalKwargsItems, + ) -> Sequence[PromptUpdate]: + def get_replacement(item_idx: int): + # Read the actual patch grid from the NaFlex processor's + # spatial_shapes output (same pattern as LFM2-VL). This avoids + # predicting from raw image dimensions, which can diverge from + # the NaFlex resize/tile logic. + out_item = out_mm_kwargs["image"][item_idx] + spatial_shapes = out_item["spatial_shapes"].data + assert isinstance(spatial_shapes, torch.Tensor) + num_tokens = int(spatial_shapes.prod().item()) + return [_IMAGE_TOKEN_ID] * num_tokens + + return [ + PromptReplacement( + modality="image", + target=DEFAULT_IMAGE_TOKEN, + replacement=get_replacement, + ), + ] + + +# --------------------------------------------------------------------------- +# Input schemas +# --------------------------------------------------------------------------- + + +class Phi4SiglipImagePixelInputs(TensorSchema): + """ + Dimensions: + - bn: Batch size * number of images + - d: Max number of patches (padded across images in the batch) + - fd: Features per patch (patch_size * patch_size * channels) + """ + + type: Literal["pixel_values"] = "pixel_values" + pixel_values: Annotated[torch.Tensor, TensorShape("bn", "d", "fd")] + pixel_attention_mask: Annotated[torch.Tensor, TensorShape("bn", "d")] + spatial_shapes: Annotated[torch.Tensor, TensorShape("bn", 2)] + + +# --------------------------------------------------------------------------- +# Model +# --------------------------------------------------------------------------- + + +@MULTIMODAL_REGISTRY.register_processor( + Phi4SiglipMultiModalProcessor, + info=Phi4SiglipProcessingInfo, + dummy_inputs=Phi4SiglipDummyInputsBuilder, +) +class Phi4ForCausalLMV(nn.Module, SupportsMultiModal, SupportsPP): + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_prefix={ + "model.vision_tower.vision_tower.vision_model.head.": None, + "model.vision_tower.vision_tower.": "vision_tower.", + "model.mm_projector.0.": "multi_modal_projector.linear_1.", + "model.mm_projector.2.": "multi_modal_projector.linear_2.", + "lm_head.": "language_model.lm_head.", + "model.": "language_model.model.", + }, + ) + + @classmethod + def get_placeholder_str(cls, modality: str, i: int) -> str | None: + if modality.startswith("image"): + return DEFAULT_IMAGE_TOKEN + raise ValueError("Only image modality is supported") + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__() + + config: PretrainedConfig = vllm_config.model_config.hf_config + quant_config = vllm_config.quant_config + self.config = config + + vision_config_dict: dict = getattr(config, "vision_config", {}) + if isinstance(vision_config_dict, dict): + if "patch_size" not in vision_config_dict: + vision_config_dict["patch_size"] = 16 + siglip2_config = Siglip2VisionConfig(**vision_config_dict) + else: + siglip2_config = vision_config_dict + + vision_hidden_size: int = config.mm_hidden_size # type: ignore[attr-defined] + text_hidden_size: int = config.hidden_size # type: ignore[attr-defined] + + with self._mark_tower_model(vllm_config, "image"): + layer_idx = -2 + num_hidden_layers = siglip2_config.num_hidden_layers + layer_idx + 1 + + self.vision_tower = Siglip2Model( + siglip2_config, + quant_config=quant_config, + num_hidden_layers_override=num_hidden_layers, + require_post_norm=False, + prefix=maybe_prefix(prefix, "vision_tower"), + ) + self.multi_modal_projector = LlavaMultiModalProjector( + vision_hidden_size=vision_hidden_size, + text_hidden_size=text_hidden_size, + projector_hidden_act="gelu", + multimodal_projector_bias=True, + quant_config=quant_config, + prefix=maybe_prefix(prefix, "multi_modal_projector"), + ) + + with self._mark_language_model(vllm_config): + self.language_model = init_vllm_registered_model( + vllm_config=vllm_config, + hf_config=config, + prefix=maybe_prefix(prefix, "language_model"), + architectures=["Phi3ForCausalLM"], + ) + + self.make_empty_intermediate_tensors = ( + self.language_model.make_empty_intermediate_tensors + ) + + self.configure_mm_token_handling( + vocab_size=config.vocab_size, # type: ignore[attr-defined] + mm_token_ids=[_IMAGE_TOKEN_ID], + ) + + def _packed_from_padded( + self, + pixel_values: torch.Tensor, + pixel_attention_mask: torch.Tensor, + spatial_shapes: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Convert padded NaFlex tensors to packed format for Siglip2Model.""" + valid_counts = pixel_attention_mask.sum(dim=1).to(torch.int32) + pixel_values_packed = pixel_values[pixel_attention_mask.bool()] + cu_seqlens = torch.zeros( + len(valid_counts) + 1, + dtype=torch.int32, + device=pixel_values.device, + ) + cu_seqlens[1:] = valid_counts.cumsum(0) + max_seqlen = valid_counts.max() + return ( + pixel_values_packed, + spatial_shapes, + cu_seqlens, + max_seqlen, + ) + + def _parse_and_validate_image_input( + self, **kwargs: object + ) -> Phi4SiglipImagePixelInputs | None: + pixel_values = kwargs.pop("pixel_values", None) + pixel_attention_mask = kwargs.pop("pixel_attention_mask", None) + spatial_shapes = kwargs.pop("spatial_shapes", None) + if pixel_values is None: + return None + + return Phi4SiglipImagePixelInputs( + type="pixel_values", + pixel_values=pixel_values, + pixel_attention_mask=pixel_attention_mask, + spatial_shapes=spatial_shapes, + ) + + def _process_image_input( + self, image_input: Phi4SiglipImagePixelInputs + ) -> MultiModalEmbeddings: + pixel_values = image_input["pixel_values"] + pixel_attention_mask = image_input["pixel_attention_mask"] + spatial_shapes = image_input["spatial_shapes"] + + ( + pixel_values_packed, + spatial_shapes_packed, + cu_seqlens, + max_seqlen, + ) = self._packed_from_padded(pixel_values, pixel_attention_mask, spatial_shapes) + + vision_features = self.vision_tower( + pixel_values_packed=pixel_values_packed, + spatial_shapes=spatial_shapes_packed, + cu_seqlens=cu_seqlens, + max_seqlen=max_seqlen, + select_layers=[-2], + ) + + if vision_features.dim() == 3: + vision_features = vision_features.squeeze(0) + + image_features = self.multi_modal_projector(vision_features) + + valid_counts = pixel_attention_mask.sum(dim=1).tolist() + return torch.split(image_features, [int(c) for c in valid_counts]) + + def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings: + image_input = self._parse_and_validate_image_input(**kwargs) + if image_input is None: + return [] + + return self._process_image_input(image_input) + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + **kwargs: object, + ) -> torch.Tensor | IntermediateTensors: + if intermediate_tensors is not None: + inputs_embeds = None + + hidden_states = self.language_model.model( + input_ids, + positions, + intermediate_tensors, + inputs_embeds=inputs_embeds, + ) + return hidden_states + + def compute_logits( + self, + hidden_states: torch.Tensor, + ) -> torch.Tensor | None: + return self.language_model.compute_logits(hidden_states) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + loader = AutoWeightsLoader(self) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/plamo2.py b/vllm/model_executor/models/plamo2.py index 44b1207745e..ce7acc1cb19 100644 --- a/vllm/model_executor/models/plamo2.py +++ b/vllm/model_executor/models/plamo2.py @@ -32,6 +32,7 @@ from vllm.model_executor.layers.mamba.mamba_utils import ( MambaStateCopyFuncCalculator, MambaStateDtypeCalculator, MambaStateShapeCalculator, + is_conv_state_dim_first, ) from vllm.model_executor.layers.mamba.ops.causal_conv1d import ( causal_conv1d_fn, @@ -266,7 +267,13 @@ class Plamo2MambaMixer(MambaBase, PluggableLayer): assert isinstance(attn_metadata, Mamba2AttentionMetadata) self_kv_cache = self.kv_cache # conv_state = (..., dim, width-1) yet contiguous along 'dim' - conv_state = self_kv_cache[0].transpose(-1, -2) + # conv_state must be (..., dim, width-1) for the conv kernels. + # DS layout stores it that way directly; SD layout needs a transpose. + conv_state = ( + self_kv_cache[0] + if is_conv_state_dim_first() + else self_kv_cache[0].transpose(-1, -2) + ) ssm_state = self_kv_cache[1] state_indices_tensor_p = attn_metadata.state_indices_tensor_p state_indices_tensor_d = attn_metadata.state_indices_tensor_d diff --git a/vllm/model_executor/models/qwen3_5_mtp.py b/vllm/model_executor/models/qwen3_5_mtp.py index 0eca47492c9..e49806365e3 100644 --- a/vllm/model_executor/models/qwen3_5_mtp.py +++ b/vllm/model_executor/models/qwen3_5_mtp.py @@ -75,13 +75,22 @@ class Qwen3_5MultiTokenPredictor(nn.Module): config.hidden_size, ) + # Workaround: mtp.fc is stored as BF16 in NVFP4 checkpoints but is + # missing from hf_quant_config.json exclude_modules. Force unquantized. + # Ref: https://github.com/vllm-project/vllm/pull/38650 + # Ref: https://github.com/NVIDIA/Model-Optimizer/pull/1124 + fc_quant = ( + None + if (quant_config and quant_config.get_name() == "modelopt_fp4") + else quant_config + ) self.fc = ColumnParallelLinear( self.config.hidden_size * 2, self.config.hidden_size, gather_output=True, bias=False, return_bias=False, - quant_config=quant_config, + quant_config=fc_quant, prefix=f"{prefix}.fc", ) diff --git a/vllm/model_executor/models/radio.py b/vllm/model_executor/models/radio.py index 9d1a070ca7d..7ec320c5348 100644 --- a/vllm/model_executor/models/radio.py +++ b/vllm/model_executor/models/radio.py @@ -176,7 +176,6 @@ class ViTPatchGenerator(nn.Module): temporal_patch_size=temporal_patch_size, **factory, ) - self._video_embedder_loaded = False if abs_pos: scale = embed_dim**-0.5 @@ -225,12 +224,7 @@ class ViTPatchGenerator(nn.Module): Returns: Embedded patches with temporal compression applied. """ - if not self._video_embedder_loaded: - raise ValueError( - "Temporal compression (video_temporal_patch_size > 1) requires " - "video_embedder weights, but they were never loaded. " - "Ensure the checkpoint was trained with temporal compression." - ) + assert self.temporal_patch_size > 1 T = self.temporal_patch_size input_size = x.shape[2:] @@ -794,9 +788,6 @@ class RadioModel(nn.Module): weight_loader(param, weight) loaded_params.add(vllm_key) - if "model.patch_generator.video_embedder.weight" in loaded_params: - self.model.patch_generator._video_embedder_loaded = True - return loaded_params def _extract_final( diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index 3e6d8cca0e9..1901381cbd3 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -111,6 +111,7 @@ _TEXT_GENERATION_MODELS = { "Gemma2ForCausalLM": ("gemma2", "Gemma2ForCausalLM"), "Gemma3ForCausalLM": ("gemma3", "Gemma3ForCausalLM"), "Gemma3nForCausalLM": ("gemma3n", "Gemma3nForCausalLM"), + "Gemma4ForCausalLM": ("gemma4", "Gemma4ForCausalLM"), "Qwen3NextForCausalLM": ("qwen3_next", "Qwen3NextForCausalLM"), "GlmForCausalLM": ("glm", "GlmForCausalLM"), "Glm4ForCausalLM": ("glm4", "Glm4ForCausalLM"), @@ -205,6 +206,7 @@ _TEXT_GENERATION_MODELS = { "SolarForCausalLM": ("solar", "SolarForCausalLM"), "TeleChatForCausalLM": ("telechat2", "TeleChat2ForCausalLM"), "TeleChat2ForCausalLM": ("telechat2", "TeleChat2ForCausalLM"), + "TeleChat3ForCausalLM": ("llama", "LlamaForCausalLM"), "TeleFLMForCausalLM": ("teleflm", "TeleFLMForCausalLM"), "XverseForCausalLM": ("llama", "LlamaForCausalLM"), "Zamba2ForCausalLM": ("zamba2", "Zamba2ForCausalLM"), @@ -350,6 +352,8 @@ _MULTIMODAL_MODELS = { "chameleon", "ChameleonForConditionalGeneration", ), + "Cheers": ("cheers", "CheersForConditionalGeneration"), + "CheersForConditionalGeneration": ("cheers", "CheersForConditionalGeneration"), "Cohere2VisionForConditionalGeneration": ( "cohere2_vision", "Cohere2VisionForConditionalGeneration", @@ -381,6 +385,7 @@ _MULTIMODAL_MODELS = { "gemma3n_mm", "Gemma3nForConditionalGeneration", ), + "Gemma4ForConditionalGeneration": ("gemma4_mm", "Gemma4ForConditionalGeneration"), "GlmAsrForConditionalGeneration": ("glmasr", "GlmAsrForConditionalGeneration"), "GLM4VForCausalLM": ("glm4v", "GLM4VForCausalLM"), "Glm4vForConditionalGeneration": ("glm4_1v", "Glm4vForConditionalGeneration"), @@ -476,6 +481,7 @@ _MULTIMODAL_MODELS = { "PaliGemmaForConditionalGeneration", ), "Phi3VForCausalLM": ("phi3v", "Phi3VForCausalLM"), + "Phi4ForCausalLMV": ("phi4siglip", "Phi4ForCausalLMV"), "Phi4MMForCausalLM": ("phi4mm", "Phi4MMForCausalLM"), "PixtralForConditionalGeneration": ("pixtral", "PixtralForConditionalGeneration"), "QwenVLForConditionalGeneration": ("qwen_vl", "QwenVLForConditionalGeneration"), diff --git a/vllm/model_executor/models/transformers/__init__.py b/vllm/model_executor/models/transformers/__init__.py index 93cd8ff5076..cb224e5cbc0 100644 --- a/vllm/model_executor/models/transformers/__init__.py +++ b/vllm/model_executor/models/transformers/__init__.py @@ -16,13 +16,11 @@ # limitations under the License. """Wrapper around `transformers` models""" -from vllm.compilation.decorators import support_torch_compile from vllm.model_executor.models.transformers.base import Base from vllm.model_executor.models.transformers.causal import CausalMixin from vllm.model_executor.models.transformers.legacy import LegacyMixin from vllm.model_executor.models.transformers.moe import MoEMixin from vllm.model_executor.models.transformers.multimodal import ( - DYNAMIC_ARG_DIMS, MultiModalDummyInputsBuilder, MultiModalMixin, MultiModalProcessingInfo, @@ -32,16 +30,13 @@ from vllm.model_executor.models.transformers.pooling import ( EmbeddingMixin, SequenceClassificationMixin, ) -from vllm.model_executor.models.transformers.utils import can_enable_torch_compile from vllm.multimodal import MULTIMODAL_REGISTRY # Text only models -@support_torch_compile(enable_if=can_enable_torch_compile) class TransformersForCausalLM(CausalMixin, Base): ... -@support_torch_compile(enable_if=can_enable_torch_compile) class TransformersMoEForCausalLM(MoEMixin, CausalMixin, Base): ... @@ -51,9 +46,6 @@ class TransformersMoEForCausalLM(MoEMixin, CausalMixin, Base): ... info=MultiModalProcessingInfo, dummy_inputs=MultiModalDummyInputsBuilder, ) -@support_torch_compile( - dynamic_arg_dims=DYNAMIC_ARG_DIMS, enable_if=can_enable_torch_compile -) class TransformersMultiModalForCausalLM(MultiModalMixin, CausalMixin, Base): ... @@ -62,20 +54,15 @@ class TransformersMultiModalForCausalLM(MultiModalMixin, CausalMixin, Base): ... info=MultiModalProcessingInfo, dummy_inputs=MultiModalDummyInputsBuilder, ) -@support_torch_compile( - dynamic_arg_dims=DYNAMIC_ARG_DIMS, enable_if=can_enable_torch_compile -) class TransformersMultiModalMoEForCausalLM( MoEMixin, MultiModalMixin, CausalMixin, Base ): ... # Embedding models -@support_torch_compile(enable_if=can_enable_torch_compile) class TransformersEmbeddingModel(EmbeddingMixin, LegacyMixin, Base): ... -@support_torch_compile(enable_if=can_enable_torch_compile) class TransformersMoEEmbeddingModel(EmbeddingMixin, MoEMixin, Base): ... @@ -84,20 +71,15 @@ class TransformersMoEEmbeddingModel(EmbeddingMixin, MoEMixin, Base): ... info=MultiModalProcessingInfo, dummy_inputs=MultiModalDummyInputsBuilder, ) -@support_torch_compile( - dynamic_arg_dims=DYNAMIC_ARG_DIMS, enable_if=can_enable_torch_compile -) class TransformersMultiModalEmbeddingModel(EmbeddingMixin, MultiModalMixin, Base): ... # Sequence classification models -@support_torch_compile(enable_if=can_enable_torch_compile) class TransformersForSequenceClassification( SequenceClassificationMixin, LegacyMixin, Base ): ... -@support_torch_compile(enable_if=can_enable_torch_compile) class TransformersMoEForSequenceClassification( SequenceClassificationMixin, MoEMixin, Base ): ... @@ -108,9 +90,6 @@ class TransformersMoEForSequenceClassification( info=MultiModalProcessingInfo, dummy_inputs=MultiModalDummyInputsBuilder, ) -@support_torch_compile( - dynamic_arg_dims=DYNAMIC_ARG_DIMS, enable_if=can_enable_torch_compile -) class TransformersMultiModalForSequenceClassification( SequenceClassificationMixin, MultiModalMixin, Base ): ... diff --git a/vllm/model_executor/models/transformers/base.py b/vllm/model_executor/models/transformers/base.py index d32bfe6cabb..8b3ef56c80a 100644 --- a/vllm/model_executor/models/transformers/base.py +++ b/vllm/model_executor/models/transformers/base.py @@ -16,6 +16,7 @@ # limitations under the License. """Transformers modeling backend base class.""" +import sys from collections.abc import Callable, Iterable from itertools import chain from operator import attrgetter @@ -29,6 +30,7 @@ from torch import nn from transformers import AutoModel from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS +from vllm.compilation.decorators import support_torch_compile from vllm.config.utils import getattr_iter from vllm.distributed import get_pp_group, get_tp_group from vllm.distributed.utils import get_pp_indices @@ -47,6 +49,7 @@ from vllm.model_executor.models.interfaces import ( ) from vllm.model_executor.models.interfaces_base import VllmModel from vllm.model_executor.models.transformers.utils import ( + can_enable_torch_compile, get_feature_request_tip, init_on_device_without_buffers, log_replacement, @@ -117,6 +120,7 @@ class Base( self.config = vllm_config.model_config.hf_config self.text_config = self.config.get_text_config() self.cache_config = vllm_config.cache_config + self.compilation_config = vllm_config.compilation_config self.device_config = vllm_config.device_config self.model_config = vllm_config.model_config self.parallel_config = vllm_config.parallel_config @@ -155,14 +159,16 @@ class Base( if "gptq" in quant_method_name: self.ignore_unexpected_suffixes.append(".bias") - # Patch config and init on "meta" to delay allocating GPU tensors self._patch_config() + from_config_kwargs = dict( + config=self.config, + dtype=self.model_config.dtype, + trust_remote_code=self.model_config.trust_remote_code, + ) + self._decorate_for_torch_compile(**from_config_kwargs) + # Init on "meta" to delay allocating GPU tensors with init_on_device_without_buffers("meta"): - self.model: PreTrainedModel = AutoModel.from_config( - self.config, - dtype=self.model_config.dtype, - trust_remote_code=self.model_config.trust_remote_code, - ) + self.model: PreTrainedModel = AutoModel.from_config(**from_config_kwargs) # Create weight name to module qualname mapper self._create_hf_to_vllm_mapper() @@ -218,6 +224,82 @@ class Base( if sub_config.dtype != (dtype := self.config.dtype): sub_config.dtype = dtype + def _get_decoder_cls(self, **kwargs: dict) -> type[PreTrainedModel]: + """ + Get the decoder class from the model. + + Args: + kwargs: The kwargs to create the model. + + Returns: + The decoder class. + """ + with torch.device("meta"): + model: PreTrainedModel = AutoModel.from_config(**kwargs) + decoder_cls = type(model.get_decoder()) + logger.debug("Identified decoder class as: %s", decoder_cls) + del model + return decoder_cls + + def _decorate_cls_for_torch_compile( + self, + cls: type[PreTrainedModel], + dynamic_arg_dims: dict[str, int] | None, + enable_if: Callable[["VllmConfig"], bool], + is_encoder: bool, + ): + """ + Decorate `cls` to indicate to vLLM that it supports torch compile. + + Args: + cls: The PreTrainedModel class to decorate. + dynamic_arg_dims: A mapping from argument name to the dynamic dimensions + of the argument. If None, default dynamic arg dims will be used. See + [`support_torch_compile`][vllm.compilation.decorators.support_torch_compile] + for more details. + enable_if: A function which takes in the vLLM config and returns whether + torch compile should be enabled for this class. + is_encoder: Whether the class being decorated is an encoder. + """ + logger.debug( + "Decorating `%s` as %s for torch compile with dynamic_arg_dims of %s", + cls.__name__, + "encoder" if is_encoder else "decoder", + dynamic_arg_dims, + ) + + @support_torch_compile( + dynamic_arg_dims=dynamic_arg_dims, + enable_if=enable_if, + is_encoder=is_encoder, + ) + class SupportTorchCompileWrapper(cls): ... + + # Patch the class in its module + module = sys.modules[cls.__module__] + setattr(module, cls.__name__, SupportTorchCompileWrapper) + + def _decorate_for_torch_compile(self, **kwargs: dict): + """ + Decorate the model's decoder class to indicate to vLLM that it supports torch + compile if `can_enable_torch_compile` is True. + + Args: + kwargs: The kwargs to create the model, which are needed to get the decoder + class. + """ + self._decorate_cls_for_torch_compile( + cls=self._get_decoder_cls(**kwargs), + # Applied to a PreTrainedModel so the batch dimension will exist + dynamic_arg_dims=dict[str, int]( + input_ids=1, # shape: [1, seq_len] + inputs_embeds=1, # shape: [1, seq_len, hidden_size] + position_ids=-1, # shape: [1, seq_len] or [3, 1, seq_len] for mrope + ), + enable_if=can_enable_torch_compile, + is_encoder=False, + ) + def _create_hf_to_vllm_mapper(self): """ Create a WeightsMapper to map checkpoint weight names to module qualnames. @@ -553,11 +635,6 @@ class Base( input_ids = None inputs_embeds = intermediate_tensors["hidden_states"] - if input_ids is not None: - input_ids = input_ids[None, ...] - if inputs_embeds is not None: - inputs_embeds = inputs_embeds[None, ...] - # If the model scales embeddings inside the input embedding layer we must # ensure they are scaled here since VocabParallelEmbedding will not do it if ( @@ -568,22 +645,29 @@ class Base( inputs_embeds = self.embed_input_ids(input_ids) input_ids = None - if self.model_config.uses_mrope: - position_ids = positions[:, None] - else: - position_ids = positions[None, ...] + # Add batch dimension before entering Transformers model + if input_ids is not None and input_ids.ndim == 1: + # [seq_len] -> [1, seq_len] + input_ids = input_ids[None, ...] + if inputs_embeds is not None and inputs_embeds.ndim == 2: + # [seq_len, hidden_size] -> [1, seq_len, hidden_size] + inputs_embeds = inputs_embeds[None, ...] + if positions.ndim == 1: + # [seq_len] -> [1, seq_len] + positions = positions[None, ...] outputs = self.model( input_ids=input_ids, inputs_embeds=inputs_embeds, use_cache=False, - position_ids=position_ids, + position_ids=positions, attention_instances=self.attention_instances, return_dict=False, **self._output_aux_hidden_states_kwargs, **kwargs, ) - # We must remove the batch dimension from these outputs + + # Remove batch dimension after exiting Transformers model hidden_states = outputs[0][0, ...] if self._output_aux_hidden_states_kwargs: aux_hidden_states = [x[0][0, ...] for x in outputs[1:]] diff --git a/vllm/model_executor/models/transformers/multimodal.py b/vllm/model_executor/models/transformers/multimodal.py index ddcd91f61e4..ab6ba91d243 100644 --- a/vllm/model_executor/models/transformers/multimodal.py +++ b/vllm/model_executor/models/transformers/multimodal.py @@ -20,7 +20,9 @@ from collections.abc import Mapping from typing import TYPE_CHECKING import torch +from transformers import AutoModel +from vllm.compilation.decorators import should_torch_compile_mm_encoder from vllm.config.utils import getattr_iter from vllm.inputs import MultiModalDataDict, MultiModalInput, mm_input from vllm.logger import init_logger @@ -46,19 +48,11 @@ from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors if TYPE_CHECKING: - from transformers import BatchFeature + from transformers import BatchFeature, PreTrainedModel from vllm.config import VllmConfig from vllm.config.multimodal import BaseDummyOptions -DYNAMIC_ARG_DIMS = { - "input_ids": 0, - # set `positions` to last dim to support Qwen-mrope - "positions": -1, - "intermediate_tensors": 0, - "inputs_embeds": 0, -} - logger = init_logger(__name__) @@ -274,6 +268,66 @@ class MultiModalMixin(SupportsMultiModal, SupportsMRoPE): # Skip SupportsMRoPE.__init__ and call the next class in MRO super(SupportsMRoPE, self).__init__(vllm_config=vllm_config, prefix=prefix) + def _get_encoder_cls( + self, modality: str = "image", **kwargs: dict + ) -> type["PreTrainedModel"]: + """ + Get the encoder class from the model. + + Args: + kwargs: The kwargs to create the model. + + Returns: + The encoder class. + """ + with torch.device("meta"): + model: PreTrainedModel = AutoModel.from_config(**kwargs) + encoder_cls = type(model.get_encoder(modality=modality)) + logger.debug("Identified encoder class as: %s", encoder_cls) + if type(model) is encoder_cls: + raise ValueError( + "Unable to infer vision encoder class from the model. " + "You must either: update the model so that " + "https://huggingface.co/docs/transformers/en/main_classes/model#transformers.PreTrainedModel.get_encoder" + " can detect the vision encoder correctly, or remove " + "'compile_mm_encoder'." + ) + del model + return encoder_cls + + def _decorate_for_torch_compile(self, **kwargs: dict): + """ + Decorate the model's decoder and encoder classes to indicate to vLLM that they + support torch compile if `can_enable_torch_compile` and + `should_torch_compile_mm_encoder` are True respectively. + + Args: + kwargs: The kwargs to create the model, which are needed to get the decoder + and encoder classes. + """ + super()._decorate_for_torch_compile(**kwargs) + # Decorate the vision encoder model class to support torch compile if needed + if self.compilation_config.compile_mm_encoder: + self.check_version("5.0.0", "multimodal encoder compilation support") + logger.warning_once( + "Multimodal encoder compilation with the Transformers modeling backend " + "is an experimental feature. It relies on:\n" + "- The vision encoder being torch compilable.\n" + "- All vision encoder tensor inputs must be type hinted as either " + "`torch.Tensor` or `torch.FloatTensor`.\n" + "- The 0-th dimension of all tensor inputs to the vision encoder being " + "the dynamic dimension (i.e., sequence length or number of patches).\n" + "Please report any issues you encounter to help us improve it." + ) + self._decorate_cls_for_torch_compile( + cls=self._get_encoder_cls(**kwargs), + # TODO: properly infer dynamic_arg_dims based on the encoder's forward + # method signature. Currently we assume dim 0 for all tensor inputs. + dynamic_arg_dims=None, + enable_if=should_torch_compile_mm_encoder, + is_encoder=True, + ) + def forward( self, input_ids: torch.Tensor | None, @@ -285,6 +339,10 @@ class MultiModalMixin(SupportsMultiModal, SupportsMRoPE): # Gemma3 and PaliGemma needs `token_type_ids` to work correctly # Other models will not have `token_type_ids` in kwargs kwargs = {k: v for k, v in kwargs.items() if k == "token_type_ids"} + # Positions shape handling for MRoPE models + if self.model_config.uses_mrope: + # [3, seq_len] -> [3, 1, seq_len] + positions = positions[:, None] model_output = super().forward( input_ids, positions, intermediate_tensors, inputs_embeds, **kwargs ) diff --git a/vllm/model_executor/models/utils.py b/vllm/model_executor/models/utils.py index 1fd18b5472b..a6c46339303 100644 --- a/vllm/model_executor/models/utils.py +++ b/vllm/model_executor/models/utils.py @@ -233,8 +233,15 @@ class AutoWeightsLoader: ): """ Add tensor names that are not in the model params that may be in the - safetensors, e.g., batch normalization stats. + safetensors, e.g., batch normalization stats and registered buffers. """ + # Add persistent registered buffers. + # Non-persistent buffers are excluded, matching PyTorch state_dict(). + non_persistent = getattr(module, "_non_persistent_buffers_set", set()) + for buf_name, buf in module.named_buffers(recurse=False): + if buf_name not in child_params and buf_name not in non_persistent: + child_params[buf_name] = buf + if isinstance( module, ( diff --git a/vllm/platforms/cpu.py b/vllm/platforms/cpu.py index ff6b22e55c9..27271448606 100644 --- a/vllm/platforms/cpu.py +++ b/vllm/platforms/cpu.py @@ -2,7 +2,6 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import glob -import json import os import platform import subprocess @@ -11,11 +10,11 @@ from dataclasses import dataclass from typing import TYPE_CHECKING import psutil -import regex as re import torch from vllm import envs from vllm.logger import init_logger +from vllm.utils.ompmultiprocessing import OMPProcessManager from vllm.utils.torch_utils import is_quantized_kv_cache from vllm.v1.attention.backends.registry import AttentionBackendEnum @@ -76,6 +75,10 @@ class CpuPlatform(Platform): dispatch_key: str = "CPU" dist_backend: str = "gloo" device_control_env_var = "CPU_VISIBLE_MEMORY_NODES" + omp_process_manager = None + smt = 1 # SMT level for OMP - 4 threads on PowerPC, 1 on others + global_cpu_mask = None + simulate_numa = int(os.environ.get("_SIM_MULTI_NUMA", 0)) @property def supported_dtypes(self) -> list[torch.dtype]: @@ -191,26 +194,10 @@ class CpuPlatform(Platform): cache_config.cpu_kvcache_space_bytes = CpuPlatform.get_device_total_memory() - # reserve at least one core for nixl_connector under p/d case - if vllm_config.kv_transfer_config and ( - envs.VLLM_CPU_NUM_OF_RESERVED_CPU == 0 - or envs.VLLM_CPU_NUM_OF_RESERVED_CPU is None - ): - os.environ["VLLM_CPU_NUM_OF_RESERVED_CPU"] = "1" - parallel_config = vllm_config.parallel_config - if ( - parallel_config.world_size > 1 - and parallel_config.distributed_executor_backend is not None - and parallel_config.distributed_executor_backend != "mp" - ): - logger.warning( - ( - "%s is not supported on CPU, fallback to mp " - "distributed executor backend." - ), - parallel_config.distributed_executor_backend, - ) + # OMP requires the MP executor to function correctly, UniProc is not + # supported as it is not possible to set the OMP environment correctly + if parallel_config.distributed_executor_backend == "uni": parallel_config.distributed_executor_backend = "mp" if parallel_config.worker_cls == "auto": parallel_config.worker_cls = "vllm.v1.worker.cpu_worker.CPUWorker" @@ -267,14 +254,6 @@ class CpuPlatform(Platform): # variable "NUMEXPR_MAX_THREADS" (64)'. os.environ["NUMEXPR_MAX_THREADS"] = str(get_max_threads()) - if envs.VLLM_CPU_OMP_THREADS_BIND != "nobind": - # Set default threads num for OpenMP parallel - os.environ["OMP_NUM_THREADS"] = str(torch.get_num_threads()) - else: - # In this case, setting the OpenMP configuration via - # OMP_NUM_THREADS is up to the user. - logger.info("Disabling binding processes to CPU cores...") - # Disable torch async compiling which won't work with daemonic processes os.environ["TORCHINDUCTOR_COMPILE_THREADS"] = "1" @@ -286,8 +265,8 @@ class CpuPlatform(Platform): ld_preload_str = os.getenv("LD_PRELOAD", "") - # Intel OpenMP setting - if "libiomp5.so" in ld_preload_str: + # Intel and CLANG OpenMP setting + if "libiomp5.so" in ld_preload_str or "libomp5" in ld_preload_str: # The time(milliseconds) that a thread should wait after # completing the execution of a parallel region, before sleeping. os.environ["KMP_BLOCKTIME"] = "1" @@ -324,37 +303,6 @@ class CpuPlatform(Platform): ld_preload_str = tcmalloc_so os.environ["LD_PRELOAD"] = ld_preload_str - if ( - platform.system() == "Linux" - and cpu_architecture in (CpuArchEnum.ARM, CpuArchEnum.POWERPC) - and not ("libomp" in ld_preload_str or "libgomp" in ld_preload_str) - ): - # We need to LD_PRELOAD PyTorch's libgomp, otherwise only - # one core will be properly utilized when we thread-bind - # See: https://github.com/vllm-project/vllm/issues/27369 - # TODO: Remove once: - # https://github.com/pytorch/pytorch/issues/166087 is fixed - - # We need to find the location of PyTorch's libgomp - torch_pkg = os.path.dirname(torch.__file__) - site_root = os.path.dirname(torch_pkg) - # Search both torch.libs and torch/lib - See: https://github.com/vllm-project/vllm/issues/30470 - torch_libs_paths = [ - os.path.join(site_root, "torch.libs"), - os.path.join(torch_pkg, "lib"), - ] - pytorch_libgomp_so_candidates = [] - for torch_libs in torch_libs_paths: - pytorch_libgomp_so_candidates.extend( - glob.glob(os.path.join(torch_libs, "libgomp*.so*")) - ) - if pytorch_libgomp_so_candidates: - pytorch_libgomp_so = pytorch_libgomp_so_candidates[0] - if ld_preload_str: - ld_preload_str += ":" - ld_preload_str += pytorch_libgomp_so - os.environ["LD_PRELOAD"] = ld_preload_str - os.environ["LOCAL_WORLD_SIZE"] = str( vllm_config.parallel_config.tensor_parallel_size ) @@ -369,6 +317,13 @@ class CpuPlatform(Platform): vllm_config.model_config.max_model_len, vllm_config.scheduler_config.DEFAULT_MAX_NUM_BATCHED_TOKENS, ) + # CI specific "quick" NUMA simulation - split all available CPUs + # into a fake NUMA topology + if os.environ.get("VLLM_CPU_SIM_MULTI_NUMA", None) is not None: + os.environ["_SIM_MULTI_NUMA"] = str( + vllm_config.parallel_config.world_size + * vllm_config.parallel_config._api_process_count + ) @classmethod def update_block_size_for_backend(cls, vllm_config: "VllmConfig") -> None: @@ -377,46 +332,71 @@ class CpuPlatform(Platform): pass @classmethod - def get_allowed_cpu_core_node_list(cls) -> tuple[list[int], list[LogicalCPUInfo]]: - assert platform.system() == "Linux" + def get_omp_manager(cls) -> OMPProcessManager: + # initialise the OMP resource management if need be and return the manager + if cls.omp_process_manager is None: + if cls.get_cpu_architecture() == CpuArchEnum.POWERPC: + cls.smt = 4 + cls.omp_process_manager = OMPProcessManager( + affinity=cls.get_global_cpu_mask(), smt=cls.smt + ) + # we need to fix up the topology returned by the OMP Manager for + # simulated NUMA environments in CI + if cls.simulate_numa > 0: + logger.info( + "Adjusting numa topology to resemble at least %d nodes", + int(cls.simulate_numa), + ) + om = cls.omp_process_manager + while len(om.omp_places) < cls.simulate_numa: + new_omp_places = [] + touched = False + for omp_place in om.omp_places: + if len(omp_place["mask"]) > 1: + touched = True + cpu_list = sorted(list(omp_place["mask"])) + new_omp_places.append( + { + "mask": set(cpu_list[0 : int(len(cpu_list) / 2)]), + "available": True, + } + ) + new_omp_places.append( + { + "mask": set(cpu_list[int(len(cpu_list) / 2) :]), + "available": True, + } + ) + if touched: + om.omp_places = new_omp_places + else: + raise ValueError( + "Cannot split the existing NUMA topology to match " + "simulation requirements" + ) - # Init LogicalCPUInfo from lscpu - lscpu_output = subprocess.check_output( - "lscpu -J -e=CPU,CORE,NODE", shell=True, text=True + return cls.omp_process_manager + + @classmethod + def get_global_cpu_mask(cls) -> set[int]: + # get global cpu mask + if cls.global_cpu_mask is None: + cls.global_cpu_mask = os.sched_getaffinity(0) + return cls.global_cpu_mask + + @classmethod + def reserve_cpus(cls, reserve: set[int]) -> bool: + # remove CPUs from global mask, for now there is no "release" mechanism + if cls.omp_process_manager is not None: + for place in cls.omp_process_manager.omp_places: + if not place["available"]: + return False + cls.global_cpu_mask = cls.get_global_cpu_mask() - reserve + # reinitialize OMP resource management + cls.omp_process_manager = OMPProcessManager( + affinity=cls.global_cpu_mask, smt=cls.smt ) - lscpu_output = re.sub(r'"node":\s*-\s*(,|\n)', r'"node": 0\1', lscpu_output) - logical_cpu_list: list[LogicalCPUInfo] = json.loads( - lscpu_output, object_hook=LogicalCPUInfo.json_decoder - )["cpus"] - - # Filter CPUs with invalid attributes - logical_cpu_list = [ - x - for x in logical_cpu_list - if -1 not in (x.id, x.physical_core, x.numa_node) - ] - - # Filter allowed CPUs - if hasattr(os, "sched_getaffinity"): - allowed_cpu_id_list = os.sched_getaffinity(0) - else: - raise NotImplementedError("Unsupported OS") - logical_cpu_list = [x for x in logical_cpu_list if x.id in allowed_cpu_id_list] - - # Get allowed NUMA nodes - allowed_numa_nodes = set() - for x in logical_cpu_list: - allowed_numa_nodes.add(x.numa_node) # type: ignore - allowed_numa_nodes_list = sorted(allowed_numa_nodes) - - env_key = CpuPlatform.device_control_env_var - if env_key in os.environ and os.environ[env_key] != "": - visible_nodes = [int(s) for s in os.environ[env_key].split(",")] - allowed_numa_nodes_list = [ - x for x in sorted(list(set(visible_nodes))) if x in allowed_numa_nodes - ] - - return allowed_numa_nodes_list, logical_cpu_list + return True @classmethod def discover_numa_topology(cls) -> list[list[int]]: diff --git a/vllm/platforms/interface.py b/vllm/platforms/interface.py index 7fba7a65fa4..1f54004f7b2 100644 --- a/vllm/platforms/interface.py +++ b/vllm/platforms/interface.py @@ -505,6 +505,7 @@ class Platform: FullAttentionSpec, MambaSpec, MLAAttentionSpec, + get_kv_quant_mode, ) cache_config = vllm_config.cache_config @@ -516,6 +517,8 @@ class Platform: else: kv_cache_dtype = STR_DTYPE_TO_TORCH_DTYPE[cache_config.cache_dtype] + kv_quant_mode = get_kv_quant_mode(cache_config.cache_dtype) + # Compute attention page size for 1 token if model_config.use_mla: attn_page_size_1_token = MLAAttentionSpec( @@ -523,6 +526,7 @@ class Platform: num_kv_heads=model_config.get_num_kv_heads(parallel_config), head_size=model_config.get_head_size(), dtype=kv_cache_dtype, + kv_quant_mode=kv_quant_mode, ).page_size_bytes else: attn_page_size_1_token = FullAttentionSpec( @@ -530,6 +534,7 @@ class Platform: num_kv_heads=model_config.get_num_kv_heads(parallel_config), head_size=model_config.get_head_size(), dtype=kv_cache_dtype, + kv_quant_mode=kv_quant_mode, ).page_size_bytes # Compute mamba page size diff --git a/vllm/platforms/rocm.py b/vllm/platforms/rocm.py index 26b081b47df..4f66d8ea2ba 100644 --- a/vllm/platforms/rocm.py +++ b/vllm/platforms/rocm.py @@ -64,6 +64,7 @@ _ROCM_DEVICE_ID_NAME_MAP: dict[str, str] = { "0x74a9": "AMD_Instinct_MI300X_HF", "0x74bd": "AMD_Instinct_MI300X_HF", "0x744c": "AMD_Radeon_RX7900XTX", + "0x7551": "AMD_Radeon_R9700", } @@ -178,6 +179,7 @@ def _get_gcn_arch() -> str: _GCN_ARCH = _get_gcn_arch() _ON_GFX1X = any(arch in _GCN_ARCH for arch in ["gfx11", "gfx12"]) +_ON_GFX12X = any(arch in _GCN_ARCH for arch in ["gfx12"]) _ON_MI3XX = any(arch in _GCN_ARCH for arch in ["gfx942", "gfx950"]) _ON_GFX9 = any(arch in _GCN_ARCH for arch in ["gfx90a", "gfx942", "gfx950"]) _ON_GFX942 = "gfx942" in _GCN_ARCH @@ -259,6 +261,10 @@ def on_gfx1x() -> bool: return _ON_GFX1X +def on_gfx12x() -> bool: + return _ON_GFX12X + + def on_mi3xx() -> bool: return _ON_MI3XX @@ -636,9 +642,9 @@ class RocmPlatform(Platform): physical_device_id = cls.device_id_to_physical_device_id(device_id) handle = amdsmi_get_processor_handles()[physical_device_id] asic_info = amdsmi_get_gpu_asic_info(handle) - device_name: str = asic_info["device_id"] - if device_name in _ROCM_DEVICE_ID_NAME_MAP: - return _ROCM_DEVICE_ID_NAME_MAP[device_name] + asic_info_device_id: str = asic_info["device_id"] + if asic_info_device_id in _ROCM_DEVICE_ID_NAME_MAP: + return _ROCM_DEVICE_ID_NAME_MAP[asic_info_device_id] return asic_info["market_name"] @classmethod diff --git a/vllm/platforms/xpu.py b/vllm/platforms/xpu.py index 2a56ff5c6e6..ffc765257ed 100644 --- a/vllm/platforms/xpu.py +++ b/vllm/platforms/xpu.py @@ -218,6 +218,57 @@ class XPUPlatform(Platform): # ref. https://openucx.readthedocs.io/en/master/faq.html os.environ["UCX_MEMTYPE_CACHE"] = "n" + @classmethod + def update_block_size_for_backend(cls, vllm_config: "VllmConfig") -> None: + super().update_block_size_for_backend(vllm_config) + from vllm.config.vllm import get_layers_from_vllm_config + from vllm.model_executor.layers.attention_layer_base import ( + AttentionLayerBase, + ) + from vllm.utils.math_utils import cdiv + + cache_config = vllm_config.cache_config + # special fix for GDN since kernel only supports block size dividable by 64 + attn_layers = get_layers_from_vllm_config( + vllm_config, + AttentionLayerBase, # type: ignore[type-abstract] + ) + + kernel_block_size = None + for layer in attn_layers.values(): + b = layer.get_attn_backend() + if b.get_name() == "GDN_ATTN": + kernel_block_size = 64 + break + + if kernel_block_size is None: + return + new_block_size = ( + cdiv(cache_config.block_size, kernel_block_size) * kernel_block_size + ) + if new_block_size == cache_config.block_size: + return + + if cache_config.mamba_cache_mode == "align": + cache_config.mamba_block_size = new_block_size + original_mamba_page_size_padded = cache_config.mamba_page_size_padded + if cache_config.mamba_page_size_padded is not None: + attn_page_size_1_token = ( + cache_config.mamba_page_size_padded // cache_config.block_size + ) + cache_config.mamba_page_size_padded = ( + new_block_size * attn_page_size_1_token + ) + cache_config.block_size = new_block_size + logger.info( + "[XPU]Setting attention block size to %d tokens to ensure multiple of %d, " + "set mamba_page_size_padded to %d bytes accordingly, before was %d bytes.", + new_block_size, + kernel_block_size, + cache_config.mamba_page_size_padded, + original_mamba_page_size_padded, + ) + @classmethod def support_hybrid_kv_cache(cls) -> bool: return True diff --git a/vllm/reasoning/__init__.py b/vllm/reasoning/__init__.py index 8c78db6f187..2d57b93369d 100644 --- a/vllm/reasoning/__init__.py +++ b/vllm/reasoning/__init__.py @@ -32,6 +32,10 @@ _REASONING_PARSERS_TO_REGISTER = { "ernie45_reasoning_parser", "Ernie45ReasoningParser", ), + "gemma4": ( + "gemma4_reasoning_parser", + "Gemma4ReasoningParser", + ), "glm45": ( "deepseek_v3_reasoning_parser", "DeepSeekV3ReasoningWithThinkingParser", diff --git a/vllm/reasoning/gemma4_reasoning_parser.py b/vllm/reasoning/gemma4_reasoning_parser.py new file mode 100644 index 00000000000..efcdcca237b --- /dev/null +++ b/vllm/reasoning/gemma4_reasoning_parser.py @@ -0,0 +1,193 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from collections.abc import Sequence +from typing import TYPE_CHECKING + +from vllm.entrypoints.openai.engine.protocol import DeltaMessage +from vllm.reasoning.basic_parsers import BaseThinkingReasoningParser +from vllm.tokenizers import TokenizerLike + +if TYPE_CHECKING: + from vllm.entrypoints.openai.chat_completion.protocol import ( + ChatCompletionRequest, + ) + from vllm.entrypoints.openai.responses.protocol import ResponsesRequest + +# Role label that Gemma4 emits at the start of the thinking channel. +# The model generates: <|channel>thought\n...reasoning... +# This prefix must be stripped to expose only the actual reasoning content. +_THOUGHT_PREFIX = "thought\n" + + +class Gemma4ReasoningParser(BaseThinkingReasoningParser): + """ + Reasoning parser for Google Gemma4 thinking models. + + Gemma4 uses <|channel>... tokens to delimit reasoning/thinking + content within its output. Thinking mode is activated by passing + ``enable_thinking=True`` in the chat template kwargs, which injects a + system turn containing <|think|> (token 98) to trigger chain-of-thought + reasoning. + + Output pattern when thinking is enabled:: + + <|channel>thought + ...chain of thought reasoning... + Final answer text here. + + The ``thought\\n`` role label inside the channel delimiters is a + structural artefact (analogous to ``user\\n`` in ``<|turn>user\\n...``). + This parser strips it so that downstream consumers see only the + actual reasoning text, consistent with the offline parser + (``vllm.reasoning.gemma4_utils._strip_thought_label``). + """ + + def __init__(self, tokenizer: TokenizerLike, *args, **kwargs): + super().__init__(tokenizer, *args, **kwargs) + # Instance state for streaming prefix stripping. + # Tracks only the reasoning text received from the base parser, + # independent of current_text (which may contain pre-reasoning + # content and lacks special token text due to + # skip_special_tokens=True). + self._reasoning_text: str = "" + self._prefix_stripped: bool = False + + @property + def start_token(self) -> str: + """The token that starts reasoning content.""" + return "<|channel>" + + @property + def end_token(self) -> str: + """The token that ends reasoning content.""" + return "" + + # ------------------------------------------------------------------ + # Non-streaming path + # ------------------------------------------------------------------ + + def extract_reasoning( + self, + model_output: str, + request: "ChatCompletionRequest | ResponsesRequest", + ) -> tuple[str | None, str | None]: + """Extract reasoning, stripping the ``thought\\n`` role label.""" + if self.start_token not in model_output and self.end_token not in model_output: + # Default to content history if no tags are present + # (or if they were stripped) + return None, model_output + + reasoning, content = super().extract_reasoning(model_output, request) + if reasoning is not None: + reasoning = _strip_thought_label(reasoning) + return reasoning, content + + # ------------------------------------------------------------------ + # Streaming path + # ------------------------------------------------------------------ + + def extract_reasoning_streaming( + self, + previous_text: str, + current_text: str, + delta_text: str, + previous_token_ids: Sequence[int], + current_token_ids: Sequence[int], + delta_token_ids: Sequence[int], + ) -> DeltaMessage | None: + """Extract streaming reasoning, stripping ``thought\\n`` from the + first reasoning delta(s). + + The ``thought\\n`` prefix may arrive as a single delta or split + across multiple deltas (e.g. ``"thought"`` then ``"\\n"``). We + buffer early reasoning tokens until we can determine whether the + prefix is present, then emit the buffered content minus the + prefix. + + Unlike the previous implementation which reconstructed accumulated + reasoning from ``current_text``, this uses instance state + (``_reasoning_text``) to track only the reasoning content returned + by the base parser. This is necessary because + ``skip_special_tokens=True`` (the vLLM default) causes the + ``<|channel>`` delimiter to be invisible in ``current_text``, + making it impossible to separate pre-reasoning content from + reasoning content via string matching. + """ + result = super().extract_reasoning_streaming( + previous_text, + current_text, + delta_text, + previous_token_ids, + current_token_ids, + delta_token_ids, + ) + if result is None: + return None + + if result.reasoning is None: + return result + + # Accumulate ONLY the reasoning text from base parser results. + # This is immune to pre-reasoning content pollution. + self._reasoning_text += result.reasoning + + # Once the prefix has been handled, all subsequent reasoning + # deltas pass through unchanged. + if self._prefix_stripped: + return result + + # ---- Prefix stripping logic ---- + + # Case 1: We've accumulated enough to confirm the prefix is + # present. Strip it and pass through the remainder. + if self._reasoning_text.startswith(_THOUGHT_PREFIX): + prefix_len = len(_THOUGHT_PREFIX) + # How much reasoning was accumulated before this delta? + prev_reasoning_len = len(self._reasoning_text) - len(result.reasoning) + if prev_reasoning_len >= prefix_len: + # Prefix was already consumed by prior deltas; this + # delta is entirely real content — pass through. + self._prefix_stripped = True + return result + else: + # Part or all of the prefix is in this delta. + chars_of_prefix_in_delta = prefix_len - prev_reasoning_len + stripped = result.reasoning[chars_of_prefix_in_delta:] + if stripped: + self._prefix_stripped = True + result.reasoning = stripped + return result + else: + # This entire delta was prefix — suppress it. + # Don't set _prefix_stripped yet; there may be more + # prefix chars to consume in the next delta. + if len(self._reasoning_text) >= prefix_len: + self._prefix_stripped = True + return None + + # Case 2: Accumulated text is a strict prefix of + # _THOUGHT_PREFIX (e.g. we've only seen "thou" so far). + # Buffer by suppressing — we can't yet tell if this will + # become the full prefix or diverge. + if _THOUGHT_PREFIX.startswith(self._reasoning_text): + return None + + # Case 3: Accumulated text doesn't match the thought prefix + # at all. This means prior deltas were buffered (suppressed + # by Case 2) but the text diverged. Re-emit the full + # accumulated text to avoid data loss. + self._prefix_stripped = True + result.reasoning = self._reasoning_text + return result + + +def _strip_thought_label(text: str) -> str: + """Remove the ``thought\\n`` role label from the beginning of text. + + Mirrors ``vllm.reasoning.gemma4_utils._strip_thought_label`` from the + offline parser. + """ + if text.startswith(_THOUGHT_PREFIX): + return text[len(_THOUGHT_PREFIX) :] + return text diff --git a/vllm/reasoning/gemma4_utils.py b/vllm/reasoning/gemma4_utils.py new file mode 100644 index 00000000000..9cdac72039e --- /dev/null +++ b/vllm/reasoning/gemma4_utils.py @@ -0,0 +1,130 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Copyright 2025 Google Inc. HuggingFace Inc. team. All rights reserved. + +"""Gemma4 thinking/reasoning output parsing utilities for offline inference. + +Standalone functions that parse decoded model text to extract structured +thinking content from Gemma4 models. These are pure-Python utilities with +zero heavy dependencies — they work on raw decoded strings from any +inference backend (vLLM, HuggingFace, TGI, etc.). + +For the OpenAI-compatible API reasoning parser (streaming + +non-streaming), see ``vllm.reasoning.gemma4_reasoning_parser``. +For tool call parsing, see ``vllm.tool_parsers.gemma4_utils``. + +Usage with vLLM offline inference:: + + from vllm import LLM, SamplingParams + from vllm.reasoning.gemma4_utils import parse_thinking_output + + llm = LLM(model="google/gemma-4-it") + outputs = llm.generate(prompt, SamplingParams(...)) + text = tokenizer.decode(outputs[0].outputs[0].token_ids, skip_special_tokens=False) + + # Extract thinking / answer (works with or without enable_thinking) + result = parse_thinking_output(text) + print(result["thinking"]) # chain-of-thought or None + print(result["answer"]) # final answer + +Ported from ``transformers.models.gemma4.utils_gemma4`` so that vLLM users +do not need a transformers dependency for output parsing. +""" + +# ---- Thinking Mode Utility ---- + +# Thinking delimiter tokens as they appear in decoded text. +# Gemma4 uses <|channel> (start) and (end) as thinking delimiters. +_THINKING_START_TAG = "<|channel>" +_THINKING_END_TAG = "" + +# Sentinel tokens that may appear in decoded output. +_TURN_END_TAG = "" + + +def parse_thinking_output(text: str) -> dict[str, str | None]: + """Parse decoded Gemma4 model output. + + Use this on **all** Gemma4 output regardless of whether thinking mode + was enabled. It handles three cases: + + 1. **Thinking enabled, tags present** — splits on ``<|channel>``/ + ```` to separate chain-of-thought from the answer and + strips the ``thought\\n`` role label. + 2. **Thinking disabled, spurious label** — strips the bare + ``thought\\n`` prefix that some Gemma4 models emit even + without thinking mode. + 3. **Clean output** — returns the text unchanged. + + The answer text is always cleaned of trailing sentinel tokens + (````, ````, etc.). + + Args: + text: Decoded model output text (from ``tokenizer.decode(...)``). + + Returns: + A dict with keys: + - ``"thinking"``: The chain-of-thought text, or ``None`` if no + thinking delimiters were found. + - ``"answer"``: The final answer text. + + Example:: + + >>> from vllm.reasoning.gemma4_utils import parse_thinking_output + >>> output_text = tokenizer.decode(outputs[0], skip_special_tokens=False) + >>> result = parse_thinking_output(output_text) + >>> print(result["thinking"]) # chain-of-thought reasoning or None + >>> print(result["answer"]) # final answer + """ + if _THINKING_END_TAG in text: + parts = text.split(_THINKING_END_TAG, 1) + thinking_block = parts[0] + answer = _clean_answer(parts[1]) + + # Extract thinking content: strip the start tag if present + if _THINKING_START_TAG in thinking_block: + thinking = thinking_block.split(_THINKING_START_TAG, 1)[1] + else: + thinking = thinking_block + + # Strip the "thought\n" channel role label the model emits inside + # <|channel>thought\n... (analogous to "user\n" in + # <|turn>user\n...). + thinking = _strip_thought_label(thinking.strip()) + thinking = thinking.strip() + + return {"thinking": thinking, "answer": answer} + + # No thinking delimiters found. + # Strip spurious "thought\n" role label that some Gemma4 models sometimes + # emit even without thinking mode enabled, then clean trailing tokens. + answer = _strip_thought_label(text) + answer = _clean_answer(answer) + return {"thinking": None, "answer": answer} + + +def _strip_thought_label(text: str) -> str: + """Strip the spurious ``thought\\n`` label from the start of text. + + Only strips when ``thought`` appears as the very first word followed by + a newline — preserving the word ``thought`` in any other context. + """ + if text.startswith("thought\n"): + return text[len("thought\n") :] + return text + + +def _clean_answer(text: str) -> str: + """Clean trailing sentinel tokens from the answer text. + + Strips ````, ````, and surrounding whitespace that the + model appends at the end of its response. + """ + text = text.strip() + # Strip trailing (Gemma4 turn-end marker) + if text.endswith(_TURN_END_TAG): + text = text[: -len(_TURN_END_TAG)].rstrip() + # Strip trailing if present + if text.endswith(""): + text = text[:-5].rstrip() + return text diff --git a/vllm/renderers/base.py b/vllm/renderers/base.py index 23f952eff4b..9947f8c9187 100644 --- a/vllm/renderers/base.py +++ b/vllm/renderers/base.py @@ -97,6 +97,7 @@ class BaseRenderer(ABC, Generic[_T]): self._async_tokenizer: AsyncMicrobatchTokenizer | None = None self.mm_processor: BaseMultiModalProcessor | None = None + self._readonly_mm_processor: BaseMultiModalProcessor | None = None self._mm_cache_stats: MultiModalCacheStats | None = None self._clear_mm_cache_async = make_async( self.clear_mm_cache, executor=self._executor @@ -124,6 +125,19 @@ class BaseRenderer(ABC, Generic[_T]): if mm_processor_cache: self._mm_cache_stats = MultiModalCacheStats() + # A second processor with its own processor-only cache. + # Used by the tokenize endpoint so that tokenize-only + # requests don't pollute the sender cache. + ro_cache = mm_registry.processor_only_cache_from_config(config) + if ro_cache is not None: + ro_tokenizer = copy.deepcopy(tokenizer) + with set_default_torch_num_threads(): + self._readonly_mm_processor = mm_registry.create_processor( + config.model_config, + tokenizer=ro_tokenizer, + cache=ro_cache, + ) + # This is used to generate internal request ID for MM processing # It has no relation to the request ID for engine core self._mm_req_counter = AtomicCounter() @@ -625,10 +639,15 @@ class BaseRenderer(ABC, Generic[_T]): mm_uuids: MultiModalUUIDDict | None, mm_processor_kwargs: Mapping[str, object] | None, tokenization_kwargs: dict[str, Any] | None, + *, + skip_mm_cache: bool = False, ) -> "MultiModalInput": mm_req_id = f"renderer{self.api_process_rank}-mm-{self._mm_req_counter.inc(1)}" - mm_processor = self.get_mm_processor() + if skip_mm_cache and self._readonly_mm_processor is not None: + mm_processor = self._readonly_mm_processor + else: + mm_processor = self.get_mm_processor() mm_data_items = mm_processor.info.parse_mm_data(mm_data) mm_uuid_items = parse_mm_uuids(mm_uuids) @@ -656,6 +675,8 @@ class BaseRenderer(ABC, Generic[_T]): def _process_tokens( self, prompt: TokensPrompt, + *, + skip_mm_cache: bool = False, ) -> TokensInput | MultiModalInput: """Process token inputs, with multimodal preprocessing offloaded to the shared thread pool in the async variant. @@ -670,6 +691,7 @@ class BaseRenderer(ABC, Generic[_T]): mm_processor_kwargs=prompt.get("mm_processor_kwargs"), tokenization_kwargs=None, # Tokenization already done in Step 2 mm_uuids=prompt.get("multi_modal_uuids"), + skip_mm_cache=skip_mm_cache, ) else: engine_input = tokens_input(prompt_token_ids) @@ -712,6 +734,8 @@ class BaseRenderer(ABC, Generic[_T]): async def _process_tokens_async( self, prompt: TokensPrompt, + *, + skip_mm_cache: bool = False, ) -> TokensInput | MultiModalInput: prompt_token_ids = prompt["prompt_token_ids"] @@ -723,6 +747,7 @@ class BaseRenderer(ABC, Generic[_T]): mm_processor_kwargs=prompt.get("mm_processor_kwargs"), tokenization_kwargs=None, mm_uuids=prompt.get("multi_modal_uuids"), + skip_mm_cache=skip_mm_cache, ) else: engine_input = tokens_input(prompt_token_ids) @@ -734,24 +759,33 @@ class BaseRenderer(ABC, Generic[_T]): return engine_input - def _process_singleton(self, prompt: SingletonTokPrompt) -> SingletonInput: - if "prompt_embeds" in prompt: - return self._process_embeds(prompt) # type: ignore[arg-type] - - return self._process_tokens(prompt) # type: ignore[arg-type] - - async def _process_singleton_async( + def _process_singleton( self, prompt: SingletonTokPrompt, + *, + skip_mm_cache: bool = False, ) -> SingletonInput: if "prompt_embeds" in prompt: return self._process_embeds(prompt) # type: ignore[arg-type] - return await self._process_tokens_async(prompt) # type: ignore[arg-type] + return self._process_tokens(prompt, skip_mm_cache=skip_mm_cache) # type: ignore[arg-type] + + async def _process_singleton_async( + self, + prompt: SingletonTokPrompt, + *, + skip_mm_cache: bool = False, + ) -> SingletonInput: + if "prompt_embeds" in prompt: + return self._process_embeds(prompt) # type: ignore[arg-type] + + return await self._process_tokens_async(prompt, skip_mm_cache=skip_mm_cache) # type: ignore[arg-type] def _process_enc_dec( self, prompt: EncoderDecoderTokPrompt, + *, + skip_mm_cache: bool = False, ) -> EncoderDecoderInput: enc_prompt = prompt["encoder_prompt"] dec_prompt = prompt["decoder_prompt"] @@ -764,9 +798,13 @@ class BaseRenderer(ABC, Generic[_T]): skip_decoder_start_token = self.mm_processor.skip_decoder_start_token return build_enc_dec_input( - encoder_input=self._process_singleton(enc_prompt), + encoder_input=self._process_singleton( + enc_prompt, skip_mm_cache=skip_mm_cache + ), decoder_input=( - None if dec_prompt is None else self._process_singleton(dec_prompt) + None + if dec_prompt is None + else self._process_singleton(dec_prompt, skip_mm_cache=skip_mm_cache) ), decoder_start_token_id=self.get_dec_start_token_id(), skip_decoder_start_token=skip_decoder_start_token, @@ -775,16 +813,20 @@ class BaseRenderer(ABC, Generic[_T]): async def _process_enc_dec_async( self, prompt: EncoderDecoderTokPrompt, + *, + skip_mm_cache: bool = False, ) -> EncoderDecoderInput: enc_prompt = prompt["encoder_prompt"] dec_prompt = prompt["decoder_prompt"] encoder_input, decoder_input = await asyncio.gather( - self._process_singleton_async(enc_prompt), + self._process_singleton_async(enc_prompt, skip_mm_cache=skip_mm_cache), ( asyncio.sleep(0) if dec_prompt is None - else self._process_singleton_async(dec_prompt) + else self._process_singleton_async( + dec_prompt, skip_mm_cache=skip_mm_cache + ) ), ) @@ -794,27 +836,40 @@ class BaseRenderer(ABC, Generic[_T]): decoder_start_token_id=self.get_dec_start_token_id(), ) - def process_for_engine(self, prompt: TokPrompt, arrival_time: float) -> EngineInput: + def process_for_engine( + self, + prompt: TokPrompt, + arrival_time: float, + *, + skip_mm_cache: bool = False, + ) -> EngineInput: engine_input: EngineInput if "encoder_prompt" in prompt: - engine_input = self._process_enc_dec(prompt) # type: ignore[arg-type] + engine_input = self._process_enc_dec(prompt, skip_mm_cache=skip_mm_cache) # type: ignore[arg-type] else: - engine_input = self._process_singleton(prompt) + engine_input = self._process_singleton(prompt, skip_mm_cache=skip_mm_cache) engine_input["arrival_time"] = arrival_time return engine_input async def process_for_engine_async( - self, prompt: TokPrompt, arrival_time: float + self, + prompt: TokPrompt, + arrival_time: float, + *, + skip_mm_cache: bool = False, ) -> EngineInput: engine_input: EngineInput if "encoder_prompt" in prompt: engine_input = await self._process_enc_dec_async( - prompt # type: ignore[arg-type] + prompt, # type: ignore[arg-type] + skip_mm_cache=skip_mm_cache, ) else: - engine_input = await self._process_singleton_async(prompt) + engine_input = await self._process_singleton_async( + prompt, skip_mm_cache=skip_mm_cache + ) engine_input["arrival_time"] = arrival_time @@ -827,6 +882,7 @@ class BaseRenderer(ABC, Generic[_T]): tok_params: TokenizeParams | None = None, *, prompt_extras: dict[str, Any] | None = None, + skip_mm_cache: bool = False, ): arrival_time = time.time() @@ -838,7 +894,10 @@ class BaseRenderer(ABC, Generic[_T]): self._apply_prompt_extras(tok_prompts, prompt_extras) - return [self.process_for_engine(prompt, arrival_time) for prompt in tok_prompts] + return [ + self.process_for_engine(prompt, arrival_time, skip_mm_cache=skip_mm_cache) + for prompt in tok_prompts + ] async def render_cmpl_async( self, @@ -846,6 +905,7 @@ class BaseRenderer(ABC, Generic[_T]): tok_params: TokenizeParams | None = None, *, prompt_extras: dict[str, Any] | None = None, + skip_mm_cache: bool = False, ): arrival_time = time.time() @@ -858,7 +918,12 @@ class BaseRenderer(ABC, Generic[_T]): self._apply_prompt_extras(tok_prompts, prompt_extras) return await asyncio.gather( - *(self.process_for_engine_async(p, arrival_time) for p in tok_prompts) + *( + self.process_for_engine_async( + p, arrival_time, skip_mm_cache=skip_mm_cache + ) + for p in tok_prompts + ) ) def render_chat( @@ -868,6 +933,7 @@ class BaseRenderer(ABC, Generic[_T]): tok_params: TokenizeParams | None = None, *, prompt_extras: dict[str, Any] | None = None, + skip_mm_cache: bool = False, ): arrival_time = time.time() @@ -890,7 +956,8 @@ class BaseRenderer(ABC, Generic[_T]): self._apply_prompt_extras(tok_prompts, prompt_extras) eng_prompts = [ - self.process_for_engine(prompt, arrival_time) for prompt in tok_prompts + self.process_for_engine(prompt, arrival_time, skip_mm_cache=skip_mm_cache) + for prompt in tok_prompts ] return out_conversations, eng_prompts @@ -902,6 +969,7 @@ class BaseRenderer(ABC, Generic[_T]): tok_params: TokenizeParams | None = None, *, prompt_extras: dict[str, Any] | None = None, + skip_mm_cache: bool = False, ): arrival_time = time.time() @@ -924,7 +992,12 @@ class BaseRenderer(ABC, Generic[_T]): self._apply_prompt_extras(tok_prompts, prompt_extras) eng_prompts = await asyncio.gather( - *(self.process_for_engine_async(p, arrival_time) for p in tok_prompts) + *( + self.process_for_engine_async( + p, arrival_time, skip_mm_cache=skip_mm_cache + ) + for p in tok_prompts + ) ) return out_conversations, eng_prompts diff --git a/vllm/tool_parsers/__init__.py b/vllm/tool_parsers/__init__.py index f480a635c6a..bffa00c4ef3 100644 --- a/vllm/tool_parsers/__init__.py +++ b/vllm/tool_parsers/__init__.py @@ -154,6 +154,10 @@ _TOOL_PARSERS_TO_REGISTER = { "functiongemma_tool_parser", "FunctionGemmaToolParser", ), + "gemma4": ( + "gemma4_tool_parser", + "Gemma4ToolParser", + ), } diff --git a/vllm/tool_parsers/gemma4_tool_parser.py b/vllm/tool_parsers/gemma4_tool_parser.py new file mode 100644 index 00000000000..3d0e4e7c4ab --- /dev/null +++ b/vllm/tool_parsers/gemma4_tool_parser.py @@ -0,0 +1,724 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Tool call parser for Google Gemma4 models. + +Gemma4 uses a custom serialization format (not JSON) for tool calls:: + + <|tool_call>call:func_name{key:<|"|>value<|"|>,num:42} + +Strings are delimited by ``<|"|>`` (token 52), keys are unquoted, and +multiple tool calls are concatenated without separators. + +Used when ``--enable-auto-tool-choice --tool-call-parser gemma4`` are set. + +For offline inference tool call parsing (direct ``tokenizer.decode()`` output), +see ``vllm.tool_parsers.gemma4_utils.parse_tool_calls``. +""" + +import json +from collections.abc import Sequence + +import regex as re + +from vllm.entrypoints.chat_utils import make_tool_call_id +from vllm.entrypoints.openai.chat_completion.protocol import ( + ChatCompletionRequest, +) +from vllm.entrypoints.openai.engine.protocol import ( + DeltaFunctionCall, + DeltaMessage, + DeltaToolCall, + ExtractedToolCallInformation, + FunctionCall, + ToolCall, +) +from vllm.entrypoints.openai.responses.protocol import ( + ResponsesRequest, +) +from vllm.logger import init_logger +from vllm.tokenizers import TokenizerLike +from vllm.tool_parsers.abstract_tool_parser import Tool, ToolParser +from vllm.tool_parsers.utils import find_common_prefix + +logger = init_logger(__name__) + +# Gemma4 special tokens for tool calls +TOOL_CALL_START = "<|tool_call>" +TOOL_CALL_END = "" +STRING_DELIM = '<|"|>' + + +# --------------------------------------------------------------------------- +# Gemma4 argument parser (used by both streaming and non-streaming paths) +# --------------------------------------------------------------------------- + + +def _parse_gemma4_value(value_str: str) -> object: + """Parse a single Gemma4 value (after key:) into a Python object.""" + value_str = value_str.strip() + if not value_str: + return value_str + + # Boolean + if value_str == "true": + return True + if value_str == "false": + return False + + # Number (int or float) + try: + if "." in value_str: + return float(value_str) + return int(value_str) + except ValueError: + pass + + # Bare string (no <|"|> delimiters — shouldn't happen but be safe) + return value_str + + +def _parse_gemma4_args(args_str: str) -> dict: + """Parse Gemma4's custom key:value format into a Python dict. + + Format examples:: + + location:<|"|>Tokyo<|"|> + location:<|"|>San Francisco<|"|>,unit:<|"|>celsius<|"|> + count:42,flag:true + nested:{inner_key:<|"|>val<|"|>} + items:[<|"|>a<|"|>,<|"|>b<|"|>] + + Returns a dict ready for ``json.dumps()``. + """ + if not args_str or not args_str.strip(): + return {} + + result: dict = {} + i = 0 + n = len(args_str) + + while i < n: + # Skip whitespace and commas + while i < n and args_str[i] in (" ", ",", "\n", "\t"): + i += 1 + if i >= n: + break + + # Parse key (unquoted, ends at ':') + key_start = i + while i < n and args_str[i] != ":": + i += 1 + if i >= n: + break + key = args_str[key_start:i].strip() + i += 1 # skip ':' + + # Parse value + if i >= n: + result[key] = "" + break + + # Skip whitespace after ':' + while i < n and args_str[i] in (" ", "\n", "\t"): + i += 1 + if i >= n: + result[key] = "" + break + + # String value: <|"|>...<|"|> + if args_str[i:].startswith(STRING_DELIM): + i += len(STRING_DELIM) + val_start = i + end_pos = args_str.find(STRING_DELIM, i) + if end_pos == -1: + # Unterminated string — take rest + result[key] = args_str[val_start:] + break + result[key] = args_str[val_start:end_pos] + i = end_pos + len(STRING_DELIM) + + # Nested object: {...} + elif args_str[i] == "{": + depth = 1 + obj_start = i + 1 + i += 1 + while i < n and depth > 0: + if args_str[i:].startswith(STRING_DELIM): + # Skip over string contents to avoid counting { inside strings + i += len(STRING_DELIM) + next_delim = args_str.find(STRING_DELIM, i) + i = n if next_delim == -1 else next_delim + len(STRING_DELIM) + continue + if args_str[i] == "{": + depth += 1 + elif args_str[i] == "}": + depth -= 1 + i += 1 + result[key] = _parse_gemma4_args(args_str[obj_start : i - 1]) + + # Array: [...] + elif args_str[i] == "[": + depth = 1 + arr_start = i + 1 + i += 1 + while i < n and depth > 0: + if args_str[i:].startswith(STRING_DELIM): + i += len(STRING_DELIM) + next_delim = args_str.find(STRING_DELIM, i) + i = n if next_delim == -1 else next_delim + len(STRING_DELIM) + continue + if args_str[i] == "[": + depth += 1 + elif args_str[i] == "]": + depth -= 1 + i += 1 + arr_content = args_str[arr_start : i - 1] + result[key] = _parse_gemma4_array(arr_content) + + # Bare value (number, boolean, etc.) + else: + val_start = i + while i < n and args_str[i] not in (",", "}", "]"): + i += 1 + result[key] = _parse_gemma4_value(args_str[val_start:i]) + + return result + + +def _parse_gemma4_array(arr_str: str) -> list: + """Parse a Gemma4 array content string into a Python list.""" + items: list = [] + i = 0 + n = len(arr_str) + + while i < n: + while i < n and arr_str[i] in (" ", ",", "\n", "\t"): + i += 1 + if i >= n: + break + + # String element + if arr_str[i:].startswith(STRING_DELIM): + i += len(STRING_DELIM) + end_pos = arr_str.find(STRING_DELIM, i) + if end_pos == -1: + items.append(arr_str[i:]) + break + items.append(arr_str[i:end_pos]) + i = end_pos + len(STRING_DELIM) + + # Nested object + elif arr_str[i] == "{": + depth = 1 + obj_start = i + 1 + i += 1 + while i < n and depth > 0: + if arr_str[i:].startswith(STRING_DELIM): + i += len(STRING_DELIM) + nd = arr_str.find(STRING_DELIM, i) + i = nd + len(STRING_DELIM) if nd != -1 else n + continue + if arr_str[i] == "{": + depth += 1 + elif arr_str[i] == "}": + depth -= 1 + i += 1 + items.append(_parse_gemma4_args(arr_str[obj_start : i - 1])) + + # Nested array + elif arr_str[i] == "[": + depth = 1 + sub_start = i + 1 + i += 1 + while i < n and depth > 0: + if arr_str[i] == "[": + depth += 1 + elif arr_str[i] == "]": + depth -= 1 + i += 1 + items.append(_parse_gemma4_array(arr_str[sub_start : i - 1])) + + # Bare value + else: + val_start = i + while i < n and arr_str[i] not in (",", "]"): + i += 1 + items.append(_parse_gemma4_value(arr_str[val_start:i])) + + return items + + +# --------------------------------------------------------------------------- +# Parser +# --------------------------------------------------------------------------- + + +class Gemma4ToolParser(ToolParser): + """ + Tool call parser for Google Gemma4 models. + + Handles the Gemma4 function call format:: + + <|tool_call>call:func_name{key:<|"|>value<|"|>} + + Used when ``--enable-auto-tool-choice --tool-call-parser gemma4`` + are set. + + Streaming strategy: **accumulate-then-parse-then-diff** + + Instead of trying to convert Gemma4's custom format to JSON + token-by-token (which fails because Gemma4 uses bare keys, custom + delimiters, and structural braces that differ from JSON), this parser: + + 1. Accumulates the raw Gemma4 argument string during streaming + 2. Parses it with ``_parse_gemma4_args()`` into a Python dict + 3. Converts to JSON with ``json.dumps()`` + 4. Diffs against the previously-streamed JSON string + 5. Emits only the new JSON fragment as the delta + + This follows the same pattern used by FunctionGemma, Hermes, and Llama + tool parsers. + """ + + def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None): + super().__init__(tokenizer, tools) + + if not self.model_tokenizer: + raise ValueError( + "The model tokenizer must be passed to the ToolParser " + "constructor during construction." + ) + + # Token strings + self.tool_call_start_token = TOOL_CALL_START + self.tool_call_end_token = TOOL_CALL_END + + # Token IDs + self.tool_call_start_token_id = self.vocab.get(TOOL_CALL_START) + self.tool_call_end_token_id = self.vocab.get(TOOL_CALL_END) + + if self.tool_call_start_token_id is None: + raise RuntimeError( + "Gemma4 ToolParser could not locate the tool call start " + f"token '{TOOL_CALL_START}' in the tokenizer!" + ) + + # Regex for non-streaming: extract complete tool calls. + # Supports function names with letters, digits, underscores, + # hyphens, and dots (e.g. "get-weather", "module.func"). + self.tool_call_regex = re.compile( + r"<\|tool_call>call:([\w\-\.]+)\{(.*?)\}", + re.DOTALL, + ) + + # Streaming state — reset per-request via _reset_streaming_state() + self._reset_streaming_state() + + # Delta buffer for handling multi-token special sequences + self.buffered_delta_text = "" + + def _reset_streaming_state(self) -> None: + """Reset all streaming state for a new request.""" + self.current_tool_id = -1 + self.current_tool_name_sent = False + self.prev_tool_call_arr: list[dict] = [] + self.streamed_args_for_tool: list[str] = [] + + def adjust_request( + self, request: ChatCompletionRequest | ResponsesRequest + ) -> ChatCompletionRequest | ResponsesRequest: + request = super().adjust_request(request) + if ( + isinstance(request, ChatCompletionRequest) + and request.tools + and request.tool_choice != "none" + ): + # Don't skip special tokens — <|tool_call> etc. are needed + request.skip_special_tokens = False + return request + + # ------------------------------------------------------------------ + # Delta buffering for multi-token special sequences + # ------------------------------------------------------------------ + + def _buffer_delta_text(self, delta_text: str) -> str: + """Buffer incoming delta text to handle multi-token special sequences. + + Accumulates partial tokens that could be the start of + ``<|tool_call>`` or ```` and only flushes them + when the complete sequence is recognized or the sequence breaks. + + This prevents partial special tokens (e.g., ``<|tool``) from being + emitted prematurely as content text. + """ + combined = self.buffered_delta_text + delta_text + + # Check if combined ends with a complete special token + if combined.endswith(TOOL_CALL_START) or combined.endswith(TOOL_CALL_END): + self.buffered_delta_text = "" + return combined + + # Check if combined ends with a partial prefix of a special token + for tag in [TOOL_CALL_START, TOOL_CALL_END]: + for i in range(1, len(tag)): + if combined.endswith(tag[:i]): + self.buffered_delta_text = combined[-i:] + return combined[:-i] + + # No partial match — flush everything + self.buffered_delta_text = "" + return combined + + # ------------------------------------------------------------------ + # Non-streaming extraction + # ------------------------------------------------------------------ + + def extract_tool_calls( + self, + model_output: str, + request: ChatCompletionRequest, + ) -> ExtractedToolCallInformation: + if self.tool_call_start_token not in model_output: + return ExtractedToolCallInformation( + tools_called=False, tool_calls=[], content=model_output + ) + + try: + matches = self.tool_call_regex.findall(model_output) + if not matches: + return ExtractedToolCallInformation( + tools_called=False, tool_calls=[], content=model_output + ) + + tool_calls: list[ToolCall] = [] + for func_name, args_str in matches: + arguments = _parse_gemma4_args(args_str) + tool_calls.append( + ToolCall( + type="function", + function=FunctionCall( + name=func_name, + arguments=json.dumps(arguments, ensure_ascii=False), + ), + ) + ) + + # Content = text before first tool call (if any) + content_end = model_output.find(self.tool_call_start_token) + content = model_output[:content_end].strip() if content_end > 0 else None + + return ExtractedToolCallInformation( + tools_called=True, + tool_calls=tool_calls, + content=content if content else None, + ) + + except Exception: + logger.exception("Error extracting tool calls from Gemma4 response") + return ExtractedToolCallInformation( + tools_called=False, tool_calls=[], content=model_output + ) + + # ------------------------------------------------------------------ + # Streaming extraction — accumulate-then-parse-then-diff + # ------------------------------------------------------------------ + + def extract_tool_calls_streaming( + self, + previous_text: str, + current_text: str, + delta_text: str, + previous_token_ids: Sequence[int], + current_token_ids: Sequence[int], + delta_token_ids: Sequence[int], + request: ChatCompletionRequest, + ) -> DeltaMessage | None: + # Buffer delta text to handle multi-token special sequences + delta_text = self._buffer_delta_text(delta_text) + # Reconstruct current_text after buffering to stay in sync + current_text = previous_text + delta_text + + # If no tool call token seen yet, emit as content + if self.tool_call_start_token not in current_text: + if delta_text: + return DeltaMessage(content=delta_text) + return None + + try: + return self._extract_streaming( + previous_text=previous_text, + current_text=current_text, + delta_text=delta_text, + ) + except Exception: + logger.exception("Error in Gemma4 streaming tool call extraction") + return None + + def _extract_streaming( + self, + previous_text: str, + current_text: str, + delta_text: str, + ) -> DeltaMessage | None: + """Tag-counting streaming parser. + + Uses the proven approach from FunctionGemma/Hermes: count start/end + tags in previous vs current text to determine phase, then + accumulate-parse-diff for arguments. + + Format: ``<|tool_call>call:name{args}`` + """ + start_count = current_text.count(self.tool_call_start_token) + end_count = current_text.count(self.tool_call_end_token) + prev_start_count = previous_text.count(self.tool_call_start_token) + prev_end_count = previous_text.count(self.tool_call_end_token) + + # Case 1: Not inside any tool call — emit as content + if ( + start_count == end_count + and prev_end_count == end_count + and self.tool_call_end_token not in delta_text + ): + if delta_text: + return DeltaMessage(content=delta_text) + return None + + # Case 2: Starting a new tool call + if start_count > prev_start_count and start_count > end_count: + self.current_tool_id += 1 + self.current_tool_name_sent = False + self.streamed_args_for_tool.append("") + self.prev_tool_call_arr.append({}) + logger.debug("Starting new tool call %d", self.current_tool_id) + # Don't return yet — fall through to try parsing if there's + # content after <|tool_call> in this same delta + # (but usually it's just the token itself, so return None) + if len(delta_text) <= len(self.tool_call_start_token): + return None + + # Case 3: Tool call just ended + if end_count > prev_end_count: + return self._handle_tool_call_end(current_text) + + # Case 4: In the middle of a tool call — parse partial content + if start_count > end_count: + return self._handle_tool_call_middle(current_text) + + # Default: generate text outside tool calls + if delta_text: + text = delta_text.replace(self.tool_call_start_token, "") + text = text.replace(self.tool_call_end_token, "") + if text: + return DeltaMessage(content=text) + return None + + def _extract_partial_call(self, current_text: str) -> tuple[str | None, str]: + """Extract function name and raw argument string from partial text. + + Returns (func_name, raw_args_str) or (None, "") if not parseable yet. + """ + # Get the text after the last <|tool_call> token + last_start = current_text.rfind(self.tool_call_start_token) + if last_start == -1: + return None, "" + + partial_call = current_text[last_start + len(self.tool_call_start_token) :] + + # Strip end token if present + if self.tool_call_end_token in partial_call: + partial_call = partial_call.split(self.tool_call_end_token)[0] + + # Expect "call:name{args...}" or "call:name{args...}" + if not partial_call.startswith("call:"): + return None, "" + + func_part = partial_call[5:] # skip "call:" + + if "{" not in func_part: + # Still accumulating function name, not ready yet + return None, "" + + func_name, _, args_part = func_part.partition("{") + func_name = func_name.strip() + + # Strip trailing '}' if present (Gemma4 structural brace) + if args_part.endswith("}"): + args_part = args_part[:-1] + + return func_name, args_part + + def _handle_tool_call_middle(self, current_text: str) -> DeltaMessage | None: + """Handle streaming when we're inside an active tool call. + + Accumulates the raw Gemma4 arguments, parses them into JSON, and + diffs against the previously-streamed JSON to emit only the new + fragment. + """ + func_name, args_part = self._extract_partial_call(current_text) + + if func_name is None: + return None + + # Step 1: Send function name (once) + if not self.current_tool_name_sent and func_name: + self.current_tool_name_sent = True + self.prev_tool_call_arr[self.current_tool_id] = { + "name": func_name, + "arguments": {}, + } + return DeltaMessage( + tool_calls=[ + DeltaToolCall( + index=self.current_tool_id, + type="function", + id=make_tool_call_id(), + function=DeltaFunctionCall( + name=func_name, + arguments="", + ).model_dump(exclude_none=True), + ) + ] + ) + + # Step 2: Parse and diff arguments + if self.current_tool_name_sent and args_part: + return self._emit_argument_diff(args_part) + + return None + + def _handle_tool_call_end(self, current_text: str) -> DeltaMessage | None: + """Handle streaming when a tool call has just completed. + + Performs a final parse of the complete tool call and flushes + any remaining un-streamed argument fragments. + """ + if self.current_tool_id < 0 or self.current_tool_id >= len( + self.prev_tool_call_arr + ): + logger.debug( + "Tool call end detected but no active tool call (current_tool_id=%d)", + self.current_tool_id, + ) + return None + + # Parse the complete tool call using regex for accuracy + all_matches = self.tool_call_regex.findall(current_text) + if self.current_tool_id < len(all_matches): + _, args_str = all_matches[self.current_tool_id] + final_args = _parse_gemma4_args(args_str) + final_args_json = json.dumps(final_args, ensure_ascii=False) + + prev_streamed = self.streamed_args_for_tool[self.current_tool_id] + if len(final_args_json) > len(prev_streamed): + diff = final_args_json[len(prev_streamed) :] + self.streamed_args_for_tool[self.current_tool_id] = final_args_json + self.prev_tool_call_arr[self.current_tool_id]["arguments"] = final_args + + return DeltaMessage( + tool_calls=[ + DeltaToolCall( + index=self.current_tool_id, + function=DeltaFunctionCall(arguments=diff).model_dump( + exclude_none=True + ), + ) + ] + ) + + return None + + def _emit_argument_diff(self, raw_args_str: str) -> DeltaMessage | None: + """Parse raw Gemma4 arguments, convert to JSON, diff, and emit. + + This is the core of the accumulate-then-parse-then-diff strategy: + 1. Parse ``raw_args_str`` with ``_parse_gemma4_args()`` + 2. Convert to JSON string with ``json.dumps()`` + 3. Withhold trailing closing characters (``"}``) that may move + as more tokens arrive + 4. Diff against previously streamed JSON and emit only new chars + + **Why withholding is necessary:** + + Gemma4's custom format produces *structurally incomplete* JSON + during streaming. For example, when ``<|"|>Paris`` arrives + without a closing delimiter, ``_parse_gemma4_args`` treats it + as a complete value and produces ``{"location": "Paris"}``. But + when ``, France<|"|>`` arrives next, the JSON becomes + ``{"location": "Paris, France"}``. If we had sent the closing + ``"}`` from the first parse, the concatenated client output + would be ``{"location": "Paris"}France"}``, which is garbage. + + The solution: **never send trailing closing chars during + streaming**. They get flushed by ``_handle_tool_call_end()`` + when the ```` end marker arrives. + + Args: + raw_args_str: The raw Gemma4 argument text accumulated so far + (without the surrounding ``{`` ``}``). + + Returns: + DeltaMessage with the argument diff, or None if no new content. + """ + try: + current_args = _parse_gemma4_args(raw_args_str) + except Exception: + logger.debug( + "Could not parse partial Gemma4 args yet: %s", + raw_args_str[:100], + ) + return None + + if not current_args: + return None + + current_args_json = json.dumps(current_args, ensure_ascii=False) + + # Withhold trailing closing characters that may shift as more + # tokens arrive. Strip trailing '}', '"', and ']' sequences + # to get the "safe prefix". + safe_json = current_args_json + while safe_json and safe_json[-1] in ("}", '"', "]"): + safe_json = safe_json[:-1] + + prev_streamed = self.streamed_args_for_tool[self.current_tool_id] + + if not safe_json or safe_json == prev_streamed: + return None + + # Use find_common_prefix to handle cases where the value changed + # structurally (e.g., a string grew). + if prev_streamed: + prefix = find_common_prefix(prev_streamed, safe_json) + sent_len = len(prev_streamed) + prefix_len = len(prefix) + + if prefix_len < sent_len: + # Structure changed — we sent too much. Truncate our + # tracking to the common prefix and wait for the final + # flush in _handle_tool_call_end. + self.streamed_args_for_tool[self.current_tool_id] = prefix + return None + + # Stream the new stable portion + diff = safe_json[sent_len:] + else: + # First emission + diff = safe_json + + if diff: + self.streamed_args_for_tool[self.current_tool_id] = safe_json + self.prev_tool_call_arr[self.current_tool_id]["arguments"] = current_args + + return DeltaMessage( + tool_calls=[ + DeltaToolCall( + index=self.current_tool_id, + function=DeltaFunctionCall(arguments=diff).model_dump( + exclude_none=True + ), + ) + ] + ) + + return None diff --git a/vllm/tool_parsers/gemma4_utils.py b/vllm/tool_parsers/gemma4_utils.py new file mode 100644 index 00000000000..439ad1125ce --- /dev/null +++ b/vllm/tool_parsers/gemma4_utils.py @@ -0,0 +1,183 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Copyright 2025 Google Inc. HuggingFace Inc. team. All rights reserved. + +"""Gemma4 tool call parsing utilities for offline inference. + +Standalone functions that parse decoded model text to extract tool calls +from Gemma4 models. These are pure-Python utilities with zero heavy +dependencies — they work on raw decoded strings from any inference +backend (vLLM, HuggingFace, TGI, etc.). + +For the OpenAI-compatible API server tool parser (streaming + +non-streaming), see ``vllm.tool_parsers.gemma4_tool_parser``. +For thinking/reasoning output parsing, see +``vllm.reasoning.gemma4_utils``. + +Usage with vLLM offline inference:: + + from vllm import LLM, SamplingParams + from vllm.tool_parsers.gemma4_utils import ( + parse_tool_calls, + has_tool_response_tag, + ) + + llm = LLM(model="google/gemma-4-it") + outputs = llm.generate(prompt, SamplingParams(...)) + text = tokenizer.decode(outputs[0].outputs[0].token_ids, skip_special_tokens=False) + + # Extract tool calls + tool_calls = parse_tool_calls(text) + for tc in tool_calls: + print(f"{tc['name']}({tc['arguments']})") + +Ported from ``transformers.models.gemma4.utils_gemma4`` so that vLLM users +do not need a transformers dependency for output parsing. +""" + +import json + +import regex as re + +# Tool call delimiter tokens as they appear in decoded text. +# Standard format: <|tool_call>call:name{args} +_TOOL_CALL_START_TAG = "<|tool_call>" +_TOOL_CALL_END_TAG = "" +_TOOL_RESPONSE_START_TAG = "<|tool_response>" + +# Gemma4 escape token as it appears in decoded text. +_ESCAPE_TOKEN = '<|"|>' + + +def _parse_tool_arguments(args_str: str) -> dict[str, str]: + """Parse tool call arguments from the Gemma4 compact format. + + Handles the ``key:<|"|>value<|"|>`` format used by Gemma4, with fallback + to heuristic key-value extraction. Also tolerates the slightly different + ``key: "value"`` format (space + plain quotes) that some chat templates + produce. + + Args: + args_str: Raw argument string from inside ``call:name{...}``. + + Returns: + Dictionary of argument name → value. + """ + if not args_str or not args_str.strip(): + return {} + + # Replace Gemma4 escape tokens with standard quotes. + cleaned = args_str.replace(_ESCAPE_TOKEN, '"') + + # Try JSON parsing first (handles nested values, arrays, etc.). + try: + parsed = json.loads("{" + cleaned + "}") + # Ensure all values are strings for consistency. + return {k: str(v) if not isinstance(v, str) else v for k, v in parsed.items()} + except (json.JSONDecodeError, ValueError): + pass + + # Fallback: extract key:"value" pairs (allow optional space after colon). + arguments = {} + for key, value in re.findall(r'(\w+):\s*"([^"]*)"', cleaned): + arguments[key] = value + + if not arguments: + # Last resort: extract key:value pairs (unquoted). + for key, value in re.findall(r"(\w+):\s*([^,}]+)", args_str): + arguments[key] = value.strip().strip('"').replace(_ESCAPE_TOKEN, "") + + return arguments + + +def parse_tool_calls(text: str, *, strict: bool = False) -> list[dict]: + """Parse tool calls from decoded Gemma4 model output. + + Uses a tiered parsing strategy to handle known output variations in + Gemma4 models, which may emit + non-standard tool call formats. + + Parsing tiers: + 1. **Standard**: ``<|tool_call>call:name{args}`` + (special token IDs 48/49 in decoded text) + 2. **Fallback** (when ``strict=False``): bare ``call:name{args}`` + patterns, including ``name{args}`` (fragmented tokens from + multimodal inputs) + + Args: + text: Decoded model output text (from ``tokenizer.decode(..., + skip_special_tokens=False)``). + strict: If ``True``, only match the standard ``<|tool_call>`` format. + If ``False`` (default), also try fallback patterns for + known Gemma4 output variations. + + Returns: + A list of dicts, each with keys: + - ``"name"``: The tool function name (e.g. ``"get_weather"``). + - ``"arguments"``: A dict of argument name → value. + + Example:: + + >>> from vllm.tool_parsers.gemma4_utils import parse_tool_calls + >>> output = tokenizer.decode(outputs[0], skip_special_tokens=False) + >>> tool_calls = parse_tool_calls(output) + >>> for tc in tool_calls: + ... print(f"Call: {tc['name']}({tc['arguments']})") + """ + results = [] + + # Tier 1: Standard format with special tokens. + # <|tool_call>call:name{args} + # Note: Some Gemma4 models emit instead of . + standard_pattern = r"<\|tool_call\>call:(\w+)\{(.*?)\}(?:|)" + for match in re.finditer(standard_pattern, text, re.DOTALL): + name, args_str = match.group(1), match.group(2) + results.append( + { + "name": name, + "arguments": _parse_tool_arguments(args_str), + } + ) + + if results or strict: + return results + + # Tier 2: Fallback for known Gemma4 output variations. + # Matches: name{args}, call:name{args}, or bare call:name{args} + fallback_pattern = r"(?:|(?:^|\s)call:)(\w+)\{(.*?)\}" + for match in re.finditer(fallback_pattern, text, re.DOTALL): + name, args_str = match.group(1), match.group(2) + results.append( + { + "name": name, + "arguments": _parse_tool_arguments(args_str), + } + ) + + return results + + +def has_tool_response_tag(text: str) -> bool: + """Check if model output properly ends with a tool response tag. + + Some Gemma4 models sometimes emit ```` instead of + ``<|tool_response>`` after a tool call. This helper detects + whether the model used the proper termination, so callers can + decide whether to inject ``<|tool_response>`` into the next prompt. + + Args: + text: Decoded model output text. + + Returns: + ``True`` if the output ends with ``<|tool_response>`` + (proper behavior), ``False`` otherwise. + + Example:: + + >>> from vllm.tool_parsers.gemma4_utils import has_tool_response_tag + >>> if not has_tool_response_tag(model_output): + ... # Model used instead — inject <|tool_response> manually + ... next_prompt = "<|tool_response>" + tool_result + """ + stripped = text.rstrip() + return stripped.endswith(_TOOL_RESPONSE_START_TAG) diff --git a/vllm/tool_parsers/qwen3coder_tool_parser.py b/vllm/tool_parsers/qwen3coder_tool_parser.py index 8a5b4f5a272..ea25ea2be92 100644 --- a/vllm/tool_parsers/qwen3coder_tool_parser.py +++ b/vllm/tool_parsers/qwen3coder_tool_parser.py @@ -131,58 +131,11 @@ class Qwen3CoderToolParser(ToolParser): logger.debug("Tool '%s' is not defined in the tools list.", func_name) return {} - @staticmethod - def _first_non_null_type(type_value: Any) -> str | None: - """Extract the first non-null type from a type value. - - Handles both scalar types ("integer") and type-as-array - (["integer", "null"]) per JSON Schema spec. - """ - if isinstance(type_value, list): - return next( - ( - str(t).strip().lower() - for t in type_value - if t is not None and str(t).lower() != "null" - ), - None, - ) - if type_value is not None and str(type_value).lower() != "null": - return str(type_value).strip().lower() - return None - - def _resolve_param_type(self, param_def: dict) -> str: - """Resolve the effective type string from a parameter definition. - - Handles direct "type" fields (including type-as-array), - anyOf/oneOf schemas emitted by Pydantic v2 for Optional[T], - and $ref schemas from Pydantic model inputs. - """ - if "type" in param_def: - resolved = self._first_non_null_type(param_def["type"]) - return resolved or "string" - - if "anyOf" in param_def or "oneOf" in param_def: - variants = param_def.get("anyOf") or param_def.get("oneOf", []) - for v in variants: - if not isinstance(v, dict): - continue - resolved = self._first_non_null_type(v.get("type")) - if resolved: - return resolved - - # $ref points to a schema definition (e.g. a Pydantic model). - # The referenced type is almost always an object, so treat it - # as such to route through json.loads. - if "$ref" in param_def: - return "object" - - return "string" - def _convert_param_value( self, param_value: str, param_name: str, param_config: dict, func_name: str ) -> Any: """Convert parameter value based on its type in the schema.""" + # Handle null value for any type if param_value.lower() == "null": return None @@ -197,10 +150,19 @@ class Qwen3CoderToolParser(ToolParser): ) return param_value - if not isinstance(param_config[param_name], dict): - return param_value - - param_type = self._resolve_param_type(param_config[param_name]) + if ( + isinstance(param_config[param_name], dict) + and "type" in param_config[param_name] + ): + param_type = str(param_config[param_name]["type"]).strip().lower() + elif ( + isinstance(param_config[param_name], dict) + and "anyOf" in param_config[param_name] + ): + # anyOf has no top-level "type"; treat as object to trigger json.loads. + param_type = "object" + else: + param_type = "string" if param_type in ["string", "str", "text", "varchar", "char", "enum"]: return param_value elif ( diff --git a/vllm/transformers_utils/config.py b/vllm/transformers_utils/config.py index d2713415729..be031a83af7 100644 --- a/vllm/transformers_utils/config.py +++ b/vllm/transformers_utils/config.py @@ -80,6 +80,7 @@ class LazyConfigDict(dict): _CONFIG_REGISTRY: dict[str, type[PretrainedConfig]] = LazyConfigDict( afmoe="AfmoeConfig", bagel="BagelConfig", + umm="CheersConfig", chatglm="ChatGLMConfig", colmodernvbert="ColModernVBertConfig", colpali="ColPaliConfig", diff --git a/vllm/transformers_utils/configs/__init__.py b/vllm/transformers_utils/configs/__init__.py index 75bfda3fbdf..49bb1772463 100644 --- a/vllm/transformers_utils/configs/__init__.py +++ b/vllm/transformers_utils/configs/__init__.py @@ -18,6 +18,7 @@ _CLASS_TO_MODULE: dict[str, str] = { "AfmoeConfig": "vllm.transformers_utils.configs.afmoe", "AXK1Config": "vllm.transformers_utils.configs.AXK1", "BagelConfig": "vllm.transformers_utils.configs.bagel", + "CheersConfig": "vllm.transformers_utils.configs.cheers", "ChatGLMConfig": "vllm.transformers_utils.configs.chatglm", "ColModernVBertConfig": "vllm.transformers_utils.configs.colmodernvbert", "ColPaliConfig": "vllm.transformers_utils.configs.colpali", @@ -75,6 +76,7 @@ __all__ = [ "AfmoeConfig", "AXK1Config", "BagelConfig", + "CheersConfig", "ChatGLMConfig", "ColModernVBertConfig", "ColPaliConfig", diff --git a/vllm/transformers_utils/configs/cheers.py b/vllm/transformers_utils/configs/cheers.py new file mode 100644 index 00000000000..e00d19761af --- /dev/null +++ b/vllm/transformers_utils/configs/cheers.py @@ -0,0 +1,109 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from transformers import PretrainedConfig, SiglipVisionConfig +from transformers.modeling_rope_utils import rope_config_validation + + +class CheersTextConfig(PretrainedConfig): + """Qwen2-based text config with Cheers-specific defaults.""" + + model_type = "umm" + base_config_key = "text_config" + + def __init__( + self, + vocab_size=152064, + hidden_size=3584, + intermediate_size=18944, + num_hidden_layers=28, + num_attention_heads=28, + num_key_value_heads=4, + hidden_act="silu", + max_position_embeddings=131072, + initializer_range=0.02, + rms_norm_eps=1e-6, + use_cache=True, + tie_word_embeddings=False, + rope_theta=1000000.0, + rope_scaling=None, + use_sliding_window=False, + sliding_window=131072, + max_window_layers=28, + layer_types=None, + attention_dropout=0.0, + **kwargs, + ): + self.vocab_size = vocab_size + self.max_position_embeddings = max_position_embeddings + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + self.use_sliding_window = use_sliding_window + self.sliding_window = sliding_window if self.use_sliding_window else None + self.max_window_layers = max_window_layers + if num_key_value_heads is None: + num_key_value_heads = num_attention_heads + self.num_key_value_heads = num_key_value_heads + self.hidden_act = hidden_act + self.initializer_range = initializer_range + self.rms_norm_eps = rms_norm_eps + self.use_cache = use_cache + self.rope_theta = rope_theta + self.rope_scaling = rope_scaling + self.attention_dropout = attention_dropout + if self.rope_scaling is not None and "type" in self.rope_scaling: + self.rope_scaling["rope_type"] = self.rope_scaling["type"] + rope_config_validation(self) + + self.layer_types = layer_types + if self.layer_types is None: + self.layer_types = [ + "sliding_attention" + if self.sliding_window is not None and i >= self.max_window_layers + else "full_attention" + for i in range(self.num_hidden_layers) + ] + + super().__init__( + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) + + +class CheersConfig(PretrainedConfig): + """Configuration class for Cheers (UMM) model.""" + + model_type = "umm" + + def __init__( + self, + text_config: dict | CheersTextConfig | None = None, + vision_representation_config: dict | SiglipVisionConfig | None = None, + vae_encoder_config: dict | None = None, + vae_decoder_config: dict | None = None, + **kwargs, + ): + super().__init__(**kwargs) + + if isinstance(text_config, dict): + self.text_config = CheersTextConfig(**text_config) + else: + self.text_config = text_config or CheersTextConfig() + + if isinstance(vision_representation_config, dict): + self.vision_representation_config = SiglipVisionConfig( + **vision_representation_config + ) + else: + self.vision_representation_config = ( + vision_representation_config or SiglipVisionConfig() + ) + + self.vae_encoder_config = vae_encoder_config or {"resolution": 512} + self.vae_decoder_config = vae_decoder_config or {"resolution": 512} + + @property + def hidden_size(self) -> int: + """Return the hidden size of the language model.""" + return self.text_config.hidden_size diff --git a/vllm/transformers_utils/configs/parakeet.py b/vllm/transformers_utils/configs/parakeet.py index 7c7a5ddd800..8309277b092 100644 --- a/vllm/transformers_utils/configs/parakeet.py +++ b/vllm/transformers_utils/configs/parakeet.py @@ -44,15 +44,19 @@ class ExtractorConfig: subsampling_factor: int subsampling_conv_kernel_size: int subsampling_conv_stride: int + hop_length: int = 160 + """Default `160`: Matches HF default""" clip_duration_s: int = 30 clip_min_duration_s: float = 0.1 @staticmethod def from_hf_config(config: PretrainedConfig) -> "ExtractorConfig": assert isinstance(config, PretrainedConfig) + hop_length = int(getattr(config, "hop_length", ExtractorConfig.hop_length)) return ExtractorConfig( feature_size=config.num_mel_bins, sampling_rate=config.sampling_rate, + hop_length=hop_length, subsampling_factor=config.subsampling_factor, subsampling_conv_kernel_size=config.subsampling_conv_kernel_size, subsampling_conv_stride=config.subsampling_conv_stride, diff --git a/vllm/transformers_utils/model_arch_config_convertor.py b/vllm/transformers_utils/model_arch_config_convertor.py index 3229539e313..ea7096ae8bd 100644 --- a/vllm/transformers_utils/model_arch_config_convertor.py +++ b/vllm/transformers_utils/model_arch_config_convertor.py @@ -448,6 +448,16 @@ class LongCatFlashMTPModelArchConfigConvertor(ModelArchConfigConvertorBase): return getattr(self.hf_text_config, "num_nextn_predict_layers", 1) +class Gemma4ModelArchConfigConvertor(ModelArchConfigConvertorBase): + def get_head_size(self) -> int: + # Gemma4 uses dual head dimensions: head_dim (sliding attention) + # and global_head_dim (full attention). Return the largest so + # that attention backends allocate buffers large enough for both. + head_dim = getattr(self.hf_text_config, "head_dim", 0) + global_head_dim = getattr(self.hf_text_config, "global_head_dim", 0) + return max(head_dim, global_head_dim) or super().get_head_size() + + # hf_config.model_type -> convertor class MODEL_ARCH_CONFIG_CONVERTORS = { "cohere_asr": CohereAsrModelArchConfigConvertor, @@ -459,6 +469,8 @@ MODEL_ARCH_CONFIG_CONVERTORS = { "mpt": MPTModelArchConfigConvertor, "dbrx": DbrxModelArchConfigConvertor, "falcon": FalconModelArchConfigConvertor, + "gemma4": Gemma4ModelArchConfigConvertor, + "gemma4_text": Gemma4ModelArchConfigConvertor, "RefinedWeb": FalconModelArchConfigConvertor, "RefinedWebModel": FalconModelArchConfigConvertor, "nemotron-nas": NemotronNasModelArchConfigConvertor, diff --git a/vllm/transformers_utils/processors/__init__.py b/vllm/transformers_utils/processors/__init__.py index d0994c25779..dc837674820 100644 --- a/vllm/transformers_utils/processors/__init__.py +++ b/vllm/transformers_utils/processors/__init__.py @@ -12,6 +12,7 @@ import importlib __all__ = [ "BagelProcessor", + "CheersProcessor", "CohereASRProcessor", "DeepseekVLV2Processor", "FireRedASR2Processor", @@ -39,6 +40,7 @@ __all__ = [ _CLASS_TO_MODULE: dict[str, str] = { "BagelProcessor": "vllm.transformers_utils.processors.bagel", + "CheersProcessor": "vllm.transformers_utils.processors.cheers", "CohereASRProcessor": "vllm.transformers_utils.processors.cohere_asr", "DeepseekVLV2Processor": "vllm.transformers_utils.processors.deepseek_vl2", "FireRedASR2Processor": "vllm.transformers_utils.processors.fireredasr2", diff --git a/vllm/transformers_utils/processors/cheers.py b/vllm/transformers_utils/processors/cheers.py new file mode 100644 index 00000000000..68eecbcffe7 --- /dev/null +++ b/vllm/transformers_utils/processors/cheers.py @@ -0,0 +1,89 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Cheers (UMM) processor for image and text inputs.""" + +from transformers import AutoProcessor +from transformers.feature_extraction_utils import BatchFeature +from transformers.image_utils import ImageInput +from transformers.processing_utils import ProcessingKwargs, ProcessorMixin, Unpack +from transformers.tokenization_utils_base import PreTokenizedInput, TextInput + + +class CheersProcessorKwargs(ProcessingKwargs, total=False): # type: ignore[call-arg] + _defaults = { + "images_kwargs": { + "return_tensors": "pt", + }, + } + + +class CheersProcessor(ProcessorMixin): + """ + Constructs a Cheers processor which wraps a + SigLIP image processor and a Qwen2 tokenizer. + """ + + attributes = ["image_processor", "tokenizer"] + image_processor_class = "AutoImageProcessor" + tokenizer_class = "AutoTokenizer" + + def __call__( + self, + text: TextInput + | PreTokenizedInput + | list[TextInput] + | list[PreTokenizedInput] = None, + images: ImageInput = None, + **kwargs: Unpack[CheersProcessorKwargs], + ): + output_kwargs = self._merge_kwargs( + CheersProcessorKwargs, + tokenizer_init_kwargs=self.tokenizer.init_kwargs, + **kwargs, + ) + + if images is not None: + import torch + + if isinstance(images, (list, tuple)): + all_pv = [] + all_ghw = [] + for img in images: + result = self.image_processor(img, **output_kwargs["images_kwargs"]) + all_pv.append(result["pixel_values"]) + if "grid_hws" in result: + all_ghw.append(result["grid_hws"]) + pixel_values = { + "pixel_values": torch.cat(all_pv, dim=0), + } + if all_ghw: + pixel_values["grid_hws"] = torch.cat(all_ghw, dim=0) + else: + pixel_values = self.image_processor( + images, **output_kwargs["images_kwargs"] + ) + else: + pixel_values = {} + + text_inputs = ( + self.tokenizer(text, **output_kwargs["text_kwargs"]) + if text is not None + else {} + ) + + return BatchFeature(data={**pixel_values, **text_inputs}) + + def batch_decode(self, *args, **kwargs): + return self.tokenizer.batch_decode(*args, **kwargs) + + def decode(self, *args, **kwargs): + return self.tokenizer.decode(*args, **kwargs) + + @property + def model_input_names(self): + tokenizer_input_names = self.tokenizer.model_input_names + image_processor_input_names = self.image_processor.model_input_names + return list(dict.fromkeys(tokenizer_input_names + image_processor_input_names)) + + +AutoProcessor.register("CheersProcessor", CheersProcessor) diff --git a/vllm/transformers_utils/processors/nano_nemotron_vl.py b/vllm/transformers_utils/processors/nano_nemotron_vl.py index 594290c1441..42659c8c143 100644 --- a/vllm/transformers_utils/processors/nano_nemotron_vl.py +++ b/vllm/transformers_utils/processors/nano_nemotron_vl.py @@ -356,15 +356,6 @@ class DynamicResolutionImageTiler: feature_sizes.append(param.num_embeddings) return images, feature_sizes - feature_size_cache: dict[Image.Image, int] = {} - - @classmethod - def get_cached_feature_size(cls, image: Image.Image) -> int: - feature_size = cls.feature_size_cache[id(image)] - # hard assert that we only use the feature size once - del cls.feature_size_cache[id(image)] - return feature_size - @dataclass class DynamicResolutionParams: media: Image.Image @@ -519,7 +510,6 @@ class DynamicResolutionImageTiler: param, token_count = self.process_media(media, tokens_for_media) params.append(param) token_counts.append(token_count) - self.feature_size_cache[id(param.media)] = param.num_embeddings # Step 2: Check if total tokens is within budget total_tokens = sum(token_counts) @@ -857,13 +847,12 @@ class NanoNemotronVLProcessor(BaseNanoNemotronVLProcessor): @property def supports_video(self) -> bool: - return self.video_token_id is not None + return True @property - def video_token_id(self) -> int | None: - if self.video_token is None: - return None - return self.tokenizer.get_vocab().get(self.video_token, None) + def video_token_id(self) -> int: + assert self.video_token is not None + return self.tokenizer.get_vocab()[self.video_token] @property def image_token_id(self) -> int: @@ -1055,6 +1044,13 @@ class NanoNemotronVLProcessor(BaseNanoNemotronVLProcessor): text_inputs = self.tokenizer(text, add_special_tokens=False) combined_inputs = {**text_inputs, **video_inputs, **audio_inputs} + frames_indices = combined_inputs.get("frames_indices") + ragged_frames_indices = ( + isinstance(frames_indices, list) + and len({len(frame_indices) for frame_indices in frames_indices}) > 1 + ) + if ragged_frames_indices: + combined_inputs.pop("frames_indices") if self.dynamic_tiler is None: batch = BatchFeature( @@ -1066,6 +1062,12 @@ class NanoNemotronVLProcessor(BaseNanoNemotronVLProcessor): # allow images to be exempt from the BatchFeature validation: # We will .stack() them in _parse_and_validate_image_input batch.update(image_inputs) + if ragged_frames_indices: + assert isinstance(frames_indices, list) + batch["frames_indices"] = [ + torch.as_tensor(frame_indices, dtype=torch.int64) + for frame_indices in frames_indices + ] return batch def get_image_repl( diff --git a/vllm/utils/flashinfer.py b/vllm/utils/flashinfer.py index 02cb330bf62..8ffac48cc87 100644 --- a/vllm/utils/flashinfer.py +++ b/vllm/utils/flashinfer.py @@ -290,10 +290,10 @@ def supports_trtllm_attention() -> bool: if envs.VLLM_BATCH_INVARIANT: return False - # Requires SM100 and NVIDIA artifactory to be accessible to download cubins - return ( - current_platform.is_device_capability_family(100) and has_nvidia_artifactory() - ) + # TRTLLM attention is currently only validated on SM100 (CC 10.0). + # SM103 (GB300) hangs with FlashInfer >= 0.6.7. + # See: https://github.com/flashinfer-ai/flashinfer/issues/2939 + return current_platform.is_device_capability(100) and has_nvidia_artifactory() def force_use_trtllm_attention() -> bool | None: diff --git a/vllm/utils/ompmultiprocessing.py b/vllm/utils/ompmultiprocessing.py new file mode 100644 index 00000000000..f2273e7e3f5 --- /dev/null +++ b/vllm/utils/ompmultiprocessing.py @@ -0,0 +1,174 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""OMP Aware Multiprocessing manager for running multiprocessing.Process() +Copyright (c) 2026 Red Hat Inc +Copyright (c) 2026 Cambridge Greys Ltd +""" + +import json +import os +import subprocess + + +def _int(arg): + """Relaxed parsing of ints which handles a - instead of a number. + The lscpu json may contain that for nodes in some cases. If that + is the case we parse it to zero + """ + try: + if int(arg) >= 0: + return int(arg) + except ValueError: + pass + return 0 + + +def parse_mask(mask): + """Expand a X-Y,Z list""" + result = [] + for token in mask.split(","): + try: + start, finish = token.split("-") + if int(start) > int(finish): + raise IndexError("Invalid Indexes for cpu ranges") + for cpu in range(int(start), int(finish) + 1): + result.append(cpu) + except ValueError: + result.append(int(token)) + return set(result) + + +def enumerate_resources(resource_map, mask=None, allowed=None): + """Enumerate system resources""" + if allowed is None: + allowed = os.sched_getaffinity(0) + if mask is not None: + allowed = allowed & mask + + try: + allowed_nodes = parse_mask(os.environ["CPU_VISIBLE_MEMORY_NODES"]) + except KeyError: + allowed_nodes = None + + lscpu: dict[str, dict] = {"cpus": {}, "cores": {}, "nodes": {}} + for cpu in resource_map["cpus"]: + cpunum = int(cpu["cpu"]) + if ( + cpunum in allowed + and cpunum >= 0 + and (allowed_nodes is None or _int(cpu["node"]) in allowed_nodes) + ): + lscpu["cpus"][cpunum] = [cpu] + core = _int(cpu["core"]) + if lscpu["cores"].get(core, None) is None: + lscpu["cores"][core] = [cpu] + else: + lscpu["cores"][core].append(cpu) + node = _int(cpu["node"]) + if lscpu["nodes"].get(node, None) is None: + lscpu["nodes"][node] = [cpu] + else: + lscpu["nodes"][node].append(cpu) + return lscpu + + +def produce_cpu_list(cpus, smt=1): + """Produce a CPU list with/without SMT pairs - main cpu list case""" + mask: list[int] = [] + for key, value in cpus.items(): + exists = 0 + for cpu in mask: + if cpu == value[0]["core"]: + exists += 1 + break + if exists < smt: + mask.append(int(key)) + return {"mask": set(mask), "available": True} + + +def produce_cpu_sublist(scpus, smt=1): + """Produce a CPU list with/without SMT pairs - resource leaf case""" + cpu_list: list[dict] = [] + for value in scpus: + exists = 0 + for cpu in cpu_list: + if int(cpu["core"]) == int(value["core"]): + exists += 1 + break + if exists < smt: + cpu_list.append(value) + mask = [] + for cpu in cpu_list: + mask.append(int(cpu["cpu"])) + + return {"mask": set(mask), "available": True} + + +def create_omp_places(resources, strategy, smt=True): + """Parse CPU topology and generate possible CPU masks""" + omp_places = [] + if strategy == "all": + omp_places.append(produce_cpu_list(resources["cpus"], smt)) + elif strategy == "cores": + for value in resources["cores"].values(): + omp_places.append(produce_cpu_sublist(value, smt)) + elif strategy == "nodes": + for value in resources["nodes"].values(): + omp_places.append(produce_cpu_sublist(value, smt)) + else: + raise NotImplementedError("Unknown strategy") + + return omp_places + + +# pylint: disable=too-few-public-methods +class OMPProcessManager: + """OMP aware wrapper to run mp Process()""" + + def __init__(self, strategy="nodes", smt=1, mock=None, affinity=None): + self.strategy = strategy + self.smt = smt + self.omp_places = [] + vllm_mask = os.environ.get("VLLM_CPU_OMP_THREADS_BIND", None) + self.setup_omp = vllm_mask != "nobind" + if self.setup_omp: + omp_places = [] + if vllm_mask is not None: + masks = [] + for spec in vllm_mask.split("|"): + masks.append(parse_mask(spec)) + else: + masks = [None] + if mock is None: + data = subprocess.run( + ["lscpu", "-Je"], check=True, capture_output=True + ).stdout + else: + with open(mock, mode="rb") as jf: + data = jf.read() + lscpu = json.loads(data) + for mask in masks: + resources = enumerate_resources(lscpu, mask, affinity) + omp_places.extend(create_omp_places(resources, strategy, smt)) + self.omp_places = sorted( + omp_places, + key=lambda p: "{:04d}-{:04d}".format(len(p["mask"]), max(p["mask"])), + reverse=True, + ) + + def run(self, what, *args, **kwargs): + """Run arg with correct OMP environment""" + if self.setup_omp: + for place in self.omp_places: + if place["available"]: + reserve = int(os.environ.get("VLLM_CPU_NUM_OF_RESERVED_CPU", 0)) + place["available"] = False + # pylint: disable=consider-using-f-string + os.environ["OMP_PLACES"] = "{}".format(place["mask"]) + os.environ["OMP_NUM_THREADS"] = "{}".format( + len(place["mask"]) - reserve + ) + os.environ["OMP_PROC_BIND"] = "TRUE" + return what(*args, **kwargs) + raise IndexError("Out of OMP places") + return what(*args, **kwargs) diff --git a/vllm/utils/torch_utils.py b/vllm/utils/torch_utils.py index 59c19a56e4f..94f8c096e31 100644 --- a/vllm/utils/torch_utils.py +++ b/vllm/utils/torch_utils.py @@ -37,6 +37,8 @@ STR_DTYPE_TO_TORCH_DTYPE = { "fp8_e4m3": torch.uint8, "fp8_e5m2": torch.uint8, "int8": torch.int8, + "int8_per_token_head": torch.int8, + "fp8_per_token_head": torch.uint8, "fp8_inc": torch.float8_e4m3fn, "fp8_ds_mla": torch.uint8, } @@ -62,7 +64,12 @@ T = TypeVar("T") def is_quantized_kv_cache(kv_cache_dtype: str) -> bool: - return kv_cache_dtype.startswith("fp8") + return kv_cache_dtype.startswith("fp8") or kv_cache_dtype.endswith("per_token_head") + + +def kv_cache_uses_per_token_head_scales(kv_cache_dtype: str) -> bool: + """Return True if *kv_cache_dtype* needs per-token-head scales.""" + return kv_cache_dtype.endswith("per_token_head") def is_strictly_contiguous(t: torch.Tensor) -> bool: diff --git a/vllm/v1/attention/backend.py b/vllm/v1/attention/backend.py index 32fac520c88..bb05b31bbb7 100644 --- a/vllm/v1/attention/backend.py +++ b/vllm/v1/attention/backend.py @@ -10,6 +10,11 @@ import numpy as np import torch from typing_extensions import deprecated +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + kFp8StaticTensorSym, + kNvfp4Dynamic, +) + if TYPE_CHECKING: from vllm.config import VllmConfig from vllm.config.cache import CacheDType @@ -17,7 +22,9 @@ if TYPE_CHECKING: from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey from vllm.platforms.interface import DeviceCapability from vllm.v1.attention.backends.utils import KVCacheLayoutType - from vllm.v1.kv_cache_interface import AttentionSpec + from vllm.v1.kv_cache_interface import AttentionSpec, KVQuantMode + +from vllm.v1.kv_cache_interface import get_kv_quant_mode class AttentionType(str, Enum): @@ -740,6 +747,13 @@ class AttentionImplBase(ABC, Generic[T]): class AttentionImpl(AttentionImplBase[T], Generic[T]): """Standard attention implementation with forward method.""" + kv_cache_dtype: str + + @property + def kv_quant_mode(self) -> "KVQuantMode": + """Return the KV cache quantization mode for this layer.""" + return get_kv_quant_mode(self.kv_cache_dtype) + @abstractmethod def __init__( self, @@ -864,6 +878,14 @@ class MLAAttentionImpl(AttentionImplBase[T], Generic[T]): """MQA-style decode forward pass.""" raise NotImplementedError + def fused_output_quant_supported(self, quant_key: "QuantKey"): + """ + Does this attention implementation support fused output quantization. + Since MLA quantization is done manually in forward_impl (common code), + all MLA backends support it by default. + """ + return quant_key in (kFp8StaticTensorSym, kNvfp4Dynamic) + def do_kv_cache_update( self, kv_c_normed: torch.Tensor, @@ -894,6 +916,14 @@ class SparseMLAAttentionImpl(AttentionImplBase[T], Generic[T]): They do not support prefill (MHA-style) attention. """ + def fused_output_quant_supported(self, quant_key: "QuantKey"): + """ + Does this attention implementation support fused output quantization. + Since MLA quantization is done manually in forward_impl (common code), + all MLA backends support it by default. + """ + return quant_key in (kFp8StaticTensorSym, kNvfp4Dynamic) + @abstractmethod def __init__( self, diff --git a/vllm/v1/attention/backends/cpu_attn.py b/vllm/v1/attention/backends/cpu_attn.py index 90151a25146..5216301ef64 100644 --- a/vllm/v1/attention/backends/cpu_attn.py +++ b/vllm/v1/attention/backends/cpu_attn.py @@ -6,6 +6,7 @@ from typing import ClassVar import torch from vllm import _custom_ops as ops +from vllm import envs from vllm.config import VllmConfig from vllm.logger import init_logger from vllm.platforms import CpuArchEnum, current_platform @@ -181,7 +182,7 @@ class CPUAttentionMetadataBuilder(AttentionMetadataBuilder[CPUAttentionMetadata] causal=causal, sliding_window_size=self.window_size, isa=self.isa, - enable_kv_split=True, + enable_kv_split=envs.VLLM_CPU_ATTN_SPLIT_KV, ) attn_metadata = CPUAttentionMetadata( diff --git a/vllm/v1/attention/backends/flex_attention.py b/vllm/v1/attention/backends/flex_attention.py index e832f6bdd82..5e202e00f8b 100644 --- a/vllm/v1/attention/backends/flex_attention.py +++ b/vllm/v1/attention/backends/flex_attention.py @@ -30,6 +30,7 @@ from vllm.utils.math_utils import cdiv from vllm.utils.torch_utils import is_quantized_kv_cache, is_torch_equal_or_newer from vllm.v1.attention.backend import ( AttentionBackend, + AttentionCGSupport, AttentionImpl, AttentionMetadataBuilder, AttentionType, @@ -315,6 +316,18 @@ class BlockSparsityHint(NamedTuple): hint_fn: _block_sparsity_hint_signature +def copy_to_persistent(dst, src): + try: + dst = dst.as_strided(src.shape, src.stride()) + except RuntimeError as e: + raise RuntimeError( + f"Fail to re-stride a persistent tensor of shape {dst.shape} " + f"for a tensor of shape {src.shape}" + ) from e + dst.copy_(src) + return dst + + @dataclass class FlexAttentionMetadata: causal: bool @@ -340,6 +353,9 @@ class FlexAttentionMetadata: physical_to_logical: torch.Tensor decode_offset: torch.Tensor num_blocks_per_seq: torch.Tensor + persistent_kv_indices: torch.Tensor + persistent_kv_num_blocks: torch.Tensor + persistent_doc_ids: torch.Tensor # For logging. num_input_tokens: int = 0 # Number of tokens including padding. @@ -656,8 +672,11 @@ class FlexAttentionMetadata: kv_indices = unique_static_unsorted( (used_pages_padded.long()), M=self.num_blocks ).to(torch.int32) + kv_indices = copy_to_persistent(self.persistent_kv_indices, kv_indices) kv_num_blocks = (kv_indices >= 0).sum(dim=-1).to(torch.int32) + kv_num_blocks = copy_to_persistent(self.persistent_kv_num_blocks, kv_num_blocks) + block_mask_kwargs = { "seq_lengths": (self.num_actual_tokens, self.total_cache_tokens), "kv_num_blocks": kv_num_blocks[None, None], @@ -694,6 +713,7 @@ class FlexAttentionMetadata: assert self.suffix_kv_lens is None, "Not implemented yet." # Create a lookup mapping from query indices -> request number self.doc_ids = _offsets_to_doc_ids_tensor(self.query_start_loc) + self.doc_ids = copy_to_persistent(self.persistent_doc_ids, self.doc_ids) self.num_blocks = self.total_cache_tokens // self.block_size self.mask_mod = self.get_mask_mod() @@ -701,6 +721,8 @@ class FlexAttentionMetadata: class FlexAttentionMetadataBuilder(AttentionMetadataBuilder[FlexAttentionMetadata]): + _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.ALWAYS + def __init__( self, kv_cache_spec: AttentionSpec, @@ -726,6 +748,38 @@ class FlexAttentionMetadataBuilder(AttentionMetadataBuilder[FlexAttentionMetadat self.q_block_size: int = 16 if supports_small_blocks else 128 self.kv_block_size: int = self.block_size if supports_small_blocks else 128 + self.max_model_len = self.model_config.max_model_len + max_num_seqs = vllm_config.scheduler_config.max_num_seqs + max_num_batched_tokens = vllm_config.scheduler_config.max_num_batched_tokens + self.max_num_q_block = ( + self.max_model_len + self.q_block_size - 1 + ) // self.q_block_size + self.persistent_kv_num_blocks = torch.empty( + self.max_num_q_block, dtype=torch.int32, device=device + ) + self.persistent_offset_tensor = torch.empty( + max_num_seqs, dtype=torch.int32, device=device + ) + self.persistent_doc_ids = torch.empty( + max_num_batched_tokens, dtype=torch.int32, device=device + ) + + # initialize later when we can access block_table + self.persistent_physical_to_logical = None + self.persistent_kv_indices = None + + def build_for_cudagraph_capture( + self, common_attn_metadata: CommonAttentionMetadata + ) -> FlexAttentionMetadata: + # Use actual max_seq_len instead of max_model_len to avoid + # torch.compile recompilation during CUDA graph capture. + common_attn_metadata.max_seq_len = ( + common_attn_metadata.seq_lens_cpu.max().item() + ) + return self.build( + common_prefix_len=0, common_attn_metadata=common_attn_metadata + ) + def build( self, common_prefix_len: int, @@ -765,8 +819,32 @@ class FlexAttentionMetadataBuilder(AttentionMetadataBuilder[FlexAttentionMetadat inverse_block_table = physical_to_logical_mapping( block_table_tensor, seq_lens, block_size, num_gpu_blocks ) + if self.persistent_physical_to_logical is None: + max_num_seqs = self.vllm_config.scheduler_config.max_num_seqs + self.persistent_physical_to_logical = torch.empty( + max_num_seqs, + num_gpu_blocks, + dtype=torch.long, + device=self.device, + ) + + if self.persistent_kv_indices is None: + max_num_kv_block = ( + self.max_model_len + self.kv_block_size - 1 + ) // self.kv_block_size + self.persistent_kv_indices = torch.empty( + self.max_model_len, + max_num_kv_block, + dtype=torch.int32, + device=self.device, + ) + + inverse_block_table = copy_to_persistent( + self.persistent_physical_to_logical, inverse_block_table + ) offset_tensor = common_attn_metadata.compute_num_computed_tokens() + offset_tensor = copy_to_persistent(self.persistent_offset_tensor, offset_tensor) out = FlexAttentionMetadata( causal=common_attn_metadata.causal, @@ -795,7 +873,20 @@ class FlexAttentionMetadataBuilder(AttentionMetadataBuilder[FlexAttentionMetadat direct_build=(self.direct_build and common_attn_metadata.causal), q_block_size=self.q_block_size, kv_block_size=self.kv_block_size, + persistent_kv_indices=self.persistent_kv_indices, + persistent_kv_num_blocks=self.persistent_kv_num_blocks, + persistent_doc_ids=self.persistent_doc_ids, ) + + # Pre-build block_mask so it is ready before CUDA graph capture. + # Without this, the lazy build in forward() would run non-graph-safe + # ops (e.g. torch.nonzero) inside capture. + if out.block_mask is None: + if out.direct_build: + out.block_mask = out._build_block_mask_direct() + else: + out.block_mask = out.build_block_mask() + return out def use_cascade_attention(self, *args, **kwargs) -> bool: diff --git a/vllm/v1/attention/backends/gdn_attn.py b/vllm/v1/attention/backends/gdn_attn.py index 41c69deb43a..5ebf040be7a 100644 --- a/vllm/v1/attention/backends/gdn_attn.py +++ b/vllm/v1/attention/backends/gdn_attn.py @@ -63,6 +63,10 @@ class GDNAttentionMetadata: num_accepted_tokens: torch.Tensor | None = None # shape: [batch,] + # Pre-computed FLA chunk metadata (avoids GPU->CPU sync in prepare_chunk_indices) + chunk_indices: torch.Tensor | None = None + chunk_offsets: torch.Tensor | None = None + # The following attributes are for triton implementation of causal_conv1d nums_dict: dict | None = None batch_ptr: torch.Tensor | None = None @@ -305,6 +309,26 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata] assert num_accepted_tokens is not None num_accepted_tokens = num_accepted_tokens[spec_sequence_masks] + chunk_indices: torch.Tensor | None = None + chunk_offsets: torch.Tensor | None = None + if num_prefills > 0: + # Only prefill batches use FLA chunk ops. + # Pre-compute on CPU and async-copy to GPU to avoid + # GPU→CPU sync (.tolist()) in prepare_chunk_indices. + from vllm.model_executor.layers.fla.ops.index import ( + prepare_chunk_indices, + prepare_chunk_offsets, + ) + from vllm.model_executor.layers.fla.ops.utils import FLA_CHUNK_SIZE + + gpu_device = query_start_loc.device + chunk_indices = prepare_chunk_indices( + non_spec_query_start_loc_cpu, FLA_CHUNK_SIZE + ).to(device=gpu_device, non_blocking=True) + chunk_offsets = prepare_chunk_offsets( + non_spec_query_start_loc_cpu, FLA_CHUNK_SIZE + ).to(device=gpu_device, non_blocking=True) + if num_prefills > 0: has_initial_state = context_lens_tensor > 0 if spec_sequence_masks is not None: @@ -405,6 +429,8 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata] num_spec_decode_tokens=num_spec_decode_tokens, num_actual_tokens=m.num_actual_tokens, has_initial_state=has_initial_state, + chunk_indices=chunk_indices, + chunk_offsets=chunk_offsets, spec_query_start_loc=spec_query_start_loc, non_spec_query_start_loc=non_spec_query_start_loc, spec_state_indices_tensor=spec_state_indices_tensor, diff --git a/vllm/v1/attention/backends/mla/rocm_aiter_mla.py b/vllm/v1/attention/backends/mla/rocm_aiter_mla.py index 6c1073b3aa7..8b764cd627a 100644 --- a/vllm/v1/attention/backends/mla/rocm_aiter_mla.py +++ b/vllm/v1/attention/backends/mla/rocm_aiter_mla.py @@ -129,9 +129,10 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]): from aiter import dtypes, get_mla_metadata_info_v1 - self._num_attention_heads = vllm_config.model_config.get_num_attention_heads( - vllm_config.parallel_config - ) + # For num_attention_heads < 16 (e.g. kimi-k2.5 head=8 with TP8), + # make sure get_mla_metadata_info_v1 / get_mla_metadata_v1 are consistent + # with the actual tensor shape passed to mla_decode_fwd. + self._num_attention_heads = max(16, self.num_heads) q_dtype = self.decode_attn_out_dtype kv_cache_dtype_str = getattr(vllm_config.cache_config, "cache_dtype", "auto") if kv_cache_dtype_str in ("fp8", "fp8_e4m3", "fp8_e5m2"): diff --git a/vllm/v1/attention/backends/mla/triton_mla.py b/vllm/v1/attention/backends/mla/triton_mla.py index 6fa1bbf2087..e8b09a43656 100644 --- a/vllm/v1/attention/backends/mla/triton_mla.py +++ b/vllm/v1/attention/backends/mla/triton_mla.py @@ -14,6 +14,7 @@ from vllm.model_executor.layers.attention.mla_attention import ( MLACommonMetadata, ) from vllm.platforms.interface import DeviceCapability +from vllm.triton_utils import triton from vllm.utils.torch_utils import is_quantized_kv_cache from vllm.v1.attention.backend import ( AttentionLayer, @@ -115,6 +116,8 @@ class TritonMLAImpl(MLACommonImpl[MLACommonMetadata]): if is_quantized_kv_cache(self.kv_cache_dtype): self.supports_quant_query_input = False + self._sm_count = torch.cuda.get_device_properties(0).multi_processor_count + def _flash_attn_varlen_diff_headdims( self, q, k, v, return_softmax_lse=False, softmax_scale=None, **kwargs ): @@ -149,7 +152,24 @@ class TritonMLAImpl(MLACommonImpl[MLACommonMetadata]): lse = torch.zeros(B, q_num_heads, dtype=q.dtype, device=q.device) # For batch invariance, use only 1 split to ensure deterministic reduction - num_kv_splits = 1 if envs.VLLM_BATCH_INVARIANT else 4 + if envs.VLLM_BATCH_INVARIANT: + num_kv_splits = 1 + else: + # Minimum work per split + # hardware dependent + min_work_per_split = 512 + + ideal_splits = max(1, attn_metadata.max_seq_len // min_work_per_split) + + # use power of 2 to avoid excessive kernel instantiations + ideal_splits = triton.next_power_of_2(ideal_splits) + + # Calculate SM-based maximum splits with occupancy multiplier + # 2-4x allows multiple blocks per SM for latency hiding + # hardware dependent + occupancy_multiplier = 2 + max_splits = self._sm_count * occupancy_multiplier + num_kv_splits = min(ideal_splits, max_splits) # TODO(lucas) Allocate ahead of time attn_logits = torch.empty( @@ -186,6 +206,7 @@ class TritonMLAImpl(MLACommonImpl[MLACommonMetadata]): PAGE_SIZE, k_scale=layer._k_scale, v_scale=layer._k_scale, + is_mla=True, ) return o, lse diff --git a/vllm/v1/attention/backends/triton_attn.py b/vllm/v1/attention/backends/triton_attn.py index 3dd081745fb..5b1eec3856f 100644 --- a/vllm/v1/attention/backends/triton_attn.py +++ b/vllm/v1/attention/backends/triton_attn.py @@ -33,9 +33,14 @@ from vllm.v1.attention.backends.utils import get_kv_cache_layout from vllm.v1.attention.ops.triton_prefill_attention import context_attention_fwd from vllm.v1.attention.ops.triton_reshape_and_cache_flash import ( triton_reshape_and_cache_flash, + triton_reshape_and_cache_flash_per_token_head_quant, ) from vllm.v1.attention.ops.triton_unified_attention import unified_attention -from vllm.v1.kv_cache_interface import AttentionSpec +from vllm.v1.kv_cache_interface import ( + AttentionSpec, + get_kv_quant_mode, + kv_cache_uses_per_token_head_scales, +) logger = init_logger(__name__) @@ -270,6 +275,8 @@ class TritonAttentionBackend(AttentionBackend): "fp8", "fp8_e4m3", "fp8_e5m2", + "int8_per_token_head", + "fp8_per_token_head", ] @staticmethod @@ -302,6 +309,18 @@ class TritonAttentionBackend(AttentionBackend): ) -> tuple[int, ...]: if block_size % 16 != 0: raise ValueError("Block size must be a multiple of 16.") + if kv_cache_uses_per_token_head_scales(cache_dtype_str): + # Pad head_size by sizeof(float32)/sizeof(cache_dtype) so + # the per-head scale fits inline. The backend extracts + # data[:head_size] and scale[head_size:] via typed views. + from vllm.utils.torch_utils import ( + STR_DTYPE_TO_TORCH_DTYPE, + get_dtype_size, + ) + + cache_dtype = STR_DTYPE_TO_TORCH_DTYPE[cache_dtype_str] + scale_pad = get_dtype_size(torch.float32) // get_dtype_size(cache_dtype) + return (num_blocks, 2, block_size, num_kv_heads, head_size + scale_pad) return (num_blocks, 2, block_size, num_kv_heads, head_size) @staticmethod @@ -365,6 +384,62 @@ class TritonAttentionBackend(AttentionBackend): class TritonAttentionImpl(AttentionImpl): + # Per-token-head quant: scale views carved from inline head padding. + _k_scale_cache: torch.Tensor | None = None + _v_scale_cache: torch.Tensor | None = None + + def _ensure_scale_caches(self, kv_cache: torch.Tensor) -> None: + """Extract per-head scale views from the padded head dimension. + + The KV cache shape is ``(num_blocks, 2, block_size, nkv, hs+pad)`` + where ``pad = sizeof(float32) / sizeof(cache_dtype)``. The last + ``pad`` elements of each head hold one float32 scale. We create + strided float32 views over those bytes. + + Scale shape: ``(num_blocks, block_size, num_kv_heads)`` + """ + if self._k_scale_cache is not None: + return + from vllm.utils.torch_utils import get_dtype_size + + num_blocks, _, block_size, nkv, padded_hs = kv_cache.shape + dtype_sz = kv_cache.element_size() + scale_pad = get_dtype_size(torch.float32) // dtype_sz # e.g. 4 + hs = padded_hs - scale_pad + + raw = kv_cache.untyped_storage() + base_f32 = torch.tensor([], dtype=torch.float32, device=kv_cache.device).set_( + raw + ) + + # In the raw bytes, each (block, kv_half, slot, head) occupies + # padded_hs * dtype_sz bytes. The scale float32 sits at byte + # offset hs * dtype_sz within that region. + kv_half_bytes = block_size * nkv * padded_hs * dtype_sz + full_block_f32 = 2 * kv_half_bytes // 4 # stride between blocks + slot_f32 = nkv * padded_hs * dtype_sz // 4 # stride between slots + head_f32 = padded_hs * dtype_sz // 4 # stride between heads + scale_off_f32 = hs * dtype_sz // 4 # offset to scale within head + + # K scales: kv_half=0 + self._k_scale_cache = torch.as_strided( + base_f32, + size=(num_blocks, block_size, nkv), + stride=(full_block_f32, slot_f32, head_f32), + storage_offset=scale_off_f32, + ) + self._k_scale_cache.fill_(1.0) + + # V scales: kv_half=1, offset by kv_half_bytes + v_base_f32 = kv_half_bytes // 4 + self._v_scale_cache = torch.as_strided( + base_f32, + size=(num_blocks, block_size, nkv), + stride=(full_block_f32, slot_f32, head_f32), + storage_offset=v_base_f32 + scale_off_f32, + ) + self._v_scale_cache.fill_(1.0) + def fused_output_quant_supported(self, quant_key: QuantKey): return quant_key == kFp8StaticTensorSym @@ -418,6 +493,9 @@ class TritonAttentionImpl(AttentionImpl): self.use_alibi_sqrt = use_alibi_sqrt self.supports_quant_query_input = current_platform.is_cuda() + self._kv_quant_mode = get_kv_quant_mode(kv_cache_dtype) + self._is_per_token_head_quant = self._kv_quant_mode.is_per_token_head + def forward( self, layer: torch.nn.Module, @@ -480,15 +558,35 @@ class TritonAttentionImpl(AttentionImpl): layer, ) - # For decoder and cross-attention, use KV cache as before - key_cache, value_cache = kv_cache.unbind(1) - if is_quantized_kv_cache(self.kv_cache_dtype): - if key_cache.dtype != self.fp8_dtype: + # Per-token-head quantized KV cache: use separate scale caches. + if self._is_per_token_head_quant: + self._ensure_scale_caches(kv_cache) + key_cache, value_cache = kv_cache.unbind(1) + if key_cache.dtype == torch.uint8: key_cache = key_cache.view(self.fp8_dtype) value_cache = value_cache.view(self.fp8_dtype) - assert layer._q_scale_float == 1.0, ( - "A non 1.0 q_scale is not currently supported." + k_descale = None + v_descale = None + k_scale_cache = self._k_scale_cache + v_scale_cache = self._v_scale_cache + # FP8 per-tensor / auto path (original flow). + else: + key_cache, value_cache = kv_cache.unbind(1) + if is_quantized_kv_cache(self.kv_cache_dtype): + if key_cache.dtype != self.fp8_dtype: + key_cache = key_cache.view(self.fp8_dtype) + value_cache = value_cache.view(self.fp8_dtype) + assert layer._q_scale_float == 1.0, ( + "A non 1.0 q_scale is not currently supported." + ) + descale_shape = ( + attn_metadata.query_start_loc.shape[0] - 1, + key_cache.shape[2], ) + k_descale = layer._k_scale.expand(descale_shape) + v_descale = layer._v_scale.expand(descale_shape) + k_scale_cache = None + v_scale_cache = None cu_seqlens_q = attn_metadata.query_start_loc seqused_k = attn_metadata.seq_lens @@ -502,7 +600,6 @@ class TritonAttentionImpl(AttentionImpl): softmax_segm_max = attn_metadata.softmax_segm_max softmax_segm_expsum = attn_metadata.softmax_segm_expsum - descale_shape = (cu_seqlens_q.shape[0] - 1, key_cache.shape[2]) mm_prefix_range_tensor = attn_metadata.mm_prefix_range_tensor unified_attention( @@ -522,8 +619,8 @@ class TritonAttentionImpl(AttentionImpl): block_table=block_table, softcap=self.logits_soft_cap, q_descale=None, # Not supported - k_descale=layer._k_scale.expand(descale_shape), - v_descale=layer._v_scale.expand(descale_shape), + k_descale=k_descale, + v_descale=v_descale, seq_threshold_3D=seq_threshold_3D, num_par_softmax_segments=num_par_softmax_segments, softmax_segm_output=softmax_segm_output, @@ -532,6 +629,9 @@ class TritonAttentionImpl(AttentionImpl): sinks=self.sinks, output_scale=output_scale, mm_prefix_range=mm_prefix_range_tensor, + kv_quant_mode=self._kv_quant_mode, + k_scale_cache=k_scale_cache, + v_scale_cache=v_scale_cache, ) return output @@ -555,10 +655,10 @@ class TritonAttentionImpl(AttentionImpl): attn_metadata: Encoder attention metadata layer: The attention layer """ - # For encoder attention, process FP8 quantization if needed + # Quantized KV cache is not supported for encoder attention. if is_quantized_kv_cache(self.kv_cache_dtype): raise NotImplementedError( - "quantization is not supported for encoder attention" + "quantized KV cache is not supported for encoder attention" ) # Use encoder-specific metadata for sequence information @@ -594,16 +694,28 @@ class TritonAttentionImpl(AttentionImpl): # For encoder attention, # we use direct Q, K, V tensors without caching return - # For decoder and cross-attention, use KV cache as before - key_cache, value_cache = kv_cache.unbind(1) - # Reshape the input keys and values and store them in the cache. + if self._is_per_token_head_quant: + self._ensure_scale_caches(kv_cache) + key_cache, value_cache = kv_cache.unbind(1) + if key_cache.dtype == torch.uint8: + key_cache = key_cache.view(self.fp8_dtype) + value_cache = value_cache.view(self.fp8_dtype) + triton_reshape_and_cache_flash_per_token_head_quant( + key, + value, + key_cache, + value_cache, + self._k_scale_cache, + self._v_scale_cache, + slot_mapping, + ) + return + # For decoder and cross-attention, use KV cache as before. + key_cache, value_cache = kv_cache.unbind(1) if is_quantized_kv_cache(self.kv_cache_dtype): key_cache = key_cache.view(self.fp8_dtype) value_cache = value_cache.view(self.fp8_dtype) - # triton kernel does not support uint8 kv_cache - # (because some explicit casts (e.g. float8_e4m3fnuz) - # are not supported) triton_reshape_and_cache_flash( key, value, @@ -616,6 +728,8 @@ class TritonAttentionImpl(AttentionImpl): ) def fused_rope_kvcache_supported(self): + if self._is_per_token_head_quant: + return False return rocm_aiter_ops.is_enabled() def do_rope_and_kv_cache_update( diff --git a/vllm/v1/attention/ops/merge_attn_states.py b/vllm/v1/attention/ops/merge_attn_states.py index 270f65d5efb..cf4338fb180 100644 --- a/vllm/v1/attention/ops/merge_attn_states.py +++ b/vllm/v1/attention/ops/merge_attn_states.py @@ -14,6 +14,7 @@ def merge_attn_states( suffix_lse: torch.Tensor, output_lse: torch.Tensor | None = None, prefill_tokens_with_context: int | None = None, + output_scale: torch.Tensor | None = None, ) -> None: """Merge partial attention outputs from prefix (KV cache) and suffix (new tokens) into a single output tensor using the log-sum-exp (LSE) @@ -41,27 +42,37 @@ def merge_attn_states( >= this value are decode or context-free prefill tokens whose output is taken directly from suffix_output. If None, all tokens are treated as having context. + output_scale: Optional scalar tensor for FP8 static quantization. + When provided, output must be FP8 dtype. """ # NOTE(DefTruth): Currently, custom merge_attn_states CUDA kernel - # does not support FP8 dtype, fallback to use Triton kernel. - def supported_dtypes(o: torch.Tensor) -> bool: - return o.dtype in [torch.float32, torch.half, torch.bfloat16] + # does not support FP8 dtype for inputs, fallback to use Triton kernel. + # However, when output_scale is provided, the inputs are still BF16/FP16 + # and the output is FP8 — both CUDA and Triton support this. + # FP8 output requires output_scale to be set. + if output.dtype not in (torch.float32, torch.half, torch.bfloat16): + assert output_scale is not None, ( + f"output_scale is required when output is {output.dtype}" + ) + + def supported_dtypes(prefix: torch.Tensor) -> bool: + return prefix.dtype in [torch.float32, torch.half, torch.bfloat16] # NOTE(DefTruth): Currently, custom merge_attn_states CUDA # kernel load/store 128b(16 bytes) per memory issue within # thread. Namely, the headsize(headdim) must be multiple of - # pack_size (float32 -> 4, half/bfloat16 -> 8). - def supported_headdim(o: torch.Tensor) -> bool: - headdim = o.shape[2] # [NUM_TOKENS, NUM_HEADS, HEAD_SIZE] - if o.dtype == torch.float32: + # pack_size based on input dtype (float32 -> 4, half/bfloat16 -> 8). + def supported_headdim(prefix: torch.Tensor) -> bool: + headdim = prefix.shape[2] # [NUM_TOKENS, NUM_HEADS, HEAD_SIZE] + if prefix.dtype == torch.float32: return headdim % 4 == 0 return headdim % 8 == 0 if ( current_platform.is_cuda() - and supported_dtypes(output) - and supported_headdim(output) + and supported_dtypes(prefix_output) + and supported_headdim(prefix_output) ): from vllm._custom_ops import merge_attn_states @@ -73,9 +84,12 @@ def merge_attn_states( suffix_lse, output_lse, prefill_tokens_with_context, + output_scale, ) else: - from vllm.v1.attention.ops.triton_merge_attn_states import merge_attn_states + from vllm.v1.attention.ops.triton_merge_attn_states import ( + merge_attn_states, + ) return merge_attn_states( output, @@ -85,4 +99,5 @@ def merge_attn_states( suffix_lse, output_lse, prefill_tokens_with_context, + output_scale, ) diff --git a/vllm/v1/attention/ops/triton_decode_attention.py b/vllm/v1/attention/ops/triton_decode_attention.py index 63263bc92e2..8118db0da8c 100644 --- a/vllm/v1/attention/ops/triton_decode_attention.py +++ b/vllm/v1/attention/ops/triton_decode_attention.py @@ -291,6 +291,7 @@ def _fwd_grouped_kernel_stage1( logit_cap: tl.constexpr, Lk: tl.constexpr, Lv: tl.constexpr, + IS_MLA: tl.constexpr = False, ): cur_batch = tl.program_id(0) cur_head_id = tl.program_id(1) @@ -310,7 +311,12 @@ def _fwd_grouped_kernel_stage1( cur_batch_req_idx = cur_batch offs_q = cur_batch * stride_qbs + cur_head[:, None] * stride_qh + offs_d[None, :] - q = tl.load(Q + offs_q, mask=(mask_h[:, None]) & (mask_d[None, :]), other=0.0) + q = tl.load( + Q + offs_q, + mask=(mask_h[:, None]) & (mask_d[None, :]), + other=0.0, + cache_modifier=".ca", + ) if BLOCK_DPE > 0: offs_dpe = BLOCK_DMODEL + tl.arange(0, BLOCK_DPE) @@ -319,7 +325,10 @@ def _fwd_grouped_kernel_stage1( cur_batch * stride_qbs + cur_head[:, None] * stride_qh + offs_dpe[None, :] ) qpe = tl.load( - Q + off_qpe, mask=(mask_h[:, None]) & (mask_dpe[None, :]), other=0.0 + Q + off_qpe, + mask=(mask_h[:, None]) & (mask_dpe[None, :]), + other=0.0, + cache_modifier=".ca", ) kv_len_per_split = tl.cdiv(cur_batch_seq_len, NUM_KV_SPLITS) @@ -331,9 +340,14 @@ def _fwd_grouped_kernel_stage1( acc = tl.zeros([BLOCK_H, BLOCK_DV], dtype=tl.float32) if split_kv_end > split_kv_start: + base_offs_k = cur_kv_head * stride_buf_kh + offs_d[:, None] + base_offs_v = cur_kv_head * stride_buf_vh + offs_dv[None, :] + if BLOCK_DPE > 0: + base_offs_kpe = cur_kv_head * stride_buf_kh + offs_dpe[:, None] + ks = tl.load(k_scale) vs = tl.load(v_scale) - for start_n in range(split_kv_start, split_kv_end, BLOCK_N): + for start_n in tl.range(split_kv_start, split_kv_end, BLOCK_N): offs_n = start_n + tl.arange(0, BLOCK_N) kv_page_number = tl.load( Req_to_tokens @@ -341,31 +355,29 @@ def _fwd_grouped_kernel_stage1( + offs_n // PAGE_SIZE, mask=offs_n < split_kv_end, other=0, + cache_modifier=".ca", ) kv_loc = kv_page_number * PAGE_SIZE + offs_n % PAGE_SIZE - offs_buf_k = ( - kv_loc[None, :] * stride_buf_kbs - + cur_kv_head * stride_buf_kh - + offs_d[:, None] - ) + + # explicitly facilitate overlapping load/compute + offs_buf_k = kv_loc[None, :] * stride_buf_kbs + base_offs_k k = tl.load( K_Buffer + offs_buf_k, mask=(offs_n[None, :] < split_kv_end) & (mask_d[:, None]), other=0.0, + cache_modifier=".cg", ) + if k.dtype.is_fp8(): k = (k.to(tl.float32) * ks).to(q.dtype) qk = tl.dot(q, k.to(q.dtype)) if BLOCK_DPE > 0: - offs_buf_kpe = ( - kv_loc[None, :] * stride_buf_kbs - + cur_kv_head * stride_buf_kh - + offs_dpe[:, None] - ) + offs_buf_kpe = kv_loc[None, :] * stride_buf_kbs + base_offs_kpe kpe = tl.load( K_Buffer + offs_buf_kpe, mask=(offs_n[None, :] < split_kv_end) & (mask_dpe[:, None]), other=0.0, + cache_modifier=".cg", ) if kpe.dtype.is_fp8(): kpe = (kpe.to(tl.float32) * ks).to(qpe.dtype) @@ -379,18 +391,20 @@ def _fwd_grouped_kernel_stage1( mask_h[:, None] & (offs_n[None, :] < split_kv_end), qk, float("-inf") ) - offs_buf_v = ( - kv_loc[:, None] * stride_buf_vbs - + cur_kv_head * stride_buf_vh - + offs_dv[None, :] - ) - v = tl.load( - V_Buffer + offs_buf_v, - mask=(offs_n[:, None] < split_kv_end) & (mask_dv[None, :]), - other=0.0, - ) - if v.dtype.is_fp8(): - v = (v.to(tl.float32) * vs).to(q.dtype) + if not IS_MLA: + offs_buf_v = kv_loc[:, None] * stride_buf_vbs + base_offs_v + v = tl.load( + V_Buffer + offs_buf_v, + mask=(offs_n[:, None] < split_kv_end) & (mask_dv[None, :]), + other=0.0, + ) + if v.dtype.is_fp8(): + v = (v.to(tl.float32) * vs).to(q.dtype) + else: + # MLA uses a single c_kv. + # loading the same c_kv to interpret it as v is not necessary. + # transpose the existing c_kv (aka k) for the dot product. + v = tl.trans(k) n_e_max = tl.maximum(tl.max(qk, 1), e_max) re_scale = tl.exp(e_max - n_e_max) @@ -441,7 +455,10 @@ def _decode_grouped_att_m_fwd( logit_cap, k_scale, v_scale, + is_mla=False, ): + # with is_mla there is only a single c_kv in smem. + # could increase BLOCK or num_stages. BLOCK = 32 Lk = k_buffer.shape[-1] Lv = v_buffer.shape[-1] @@ -514,6 +531,7 @@ def _decode_grouped_att_m_fwd( num_stages=num_stages, Lk=Lk, Lv=Lv, + IS_MLA=is_mla, **extra_kargs, ) @@ -673,6 +691,7 @@ def decode_attention_fwd_grouped( logit_cap=0.0, k_scale=None, v_scale=None, + is_mla=False, ): _decode_grouped_att_m_fwd( q, @@ -687,6 +706,7 @@ def decode_attention_fwd_grouped( logit_cap, k_scale, v_scale, + is_mla=is_mla, ) _decode_softmax_reducev_fwd( attn_logits, q, o, lse, v_buffer, b_seq_len, num_kv_splits @@ -708,6 +728,7 @@ def decode_attention_fwd( logit_cap=0.0, k_scale=None, v_scale=None, + is_mla=False, ): assert num_kv_splits == attn_logits.shape[2] @@ -753,4 +774,5 @@ def decode_attention_fwd( logit_cap, k_scale, v_scale, + is_mla=is_mla, ) diff --git a/vllm/v1/attention/ops/triton_merge_attn_states.py b/vllm/v1/attention/ops/triton_merge_attn_states.py index f5b4fbe0b2a..14a52ada97f 100644 --- a/vllm/v1/attention/ops/triton_merge_attn_states.py +++ b/vllm/v1/attention/ops/triton_merge_attn_states.py @@ -3,8 +3,11 @@ import torch +from vllm.platforms import current_platform from vllm.triton_utils import tl, triton +float8_info = torch.finfo(current_platform.fp8_dtype()) + # Implements section 2.2 of https://www.arxiv.org/pdf/2501.01005 # can be used to combine partial attention results (in the split-KV case) @@ -16,14 +19,15 @@ def merge_attn_states( suffix_lse: torch.Tensor, output_lse: torch.Tensor | None = None, prefill_tokens_with_context: int | None = None, + output_scale: torch.Tensor | None = None, ) -> None: num_tokens = output.shape[0] num_query_heads = output.shape[1] head_size = output.shape[2] padded_head_size = triton.next_power_of_2(head_size) # We assume the output stride on num_head is not always as same as the - # `suffix_output` and `prefix_output`, as them might be padded by the attention - # backend. + # `suffix_output` and `prefix_output`, as them might be padded by the + # attention backend. prefix_head_stride = prefix_output.stride(1) output_head_stride = output.stride(1) @@ -41,10 +45,12 @@ def merge_attn_states( suffix_lse, prefix_head_stride, output_head_stride, + output_scale, head_size, padded_head_size, output_lse is not None, prefill_tokens_with_context, + output_scale is not None, ) @@ -58,10 +64,14 @@ def merge_attn_states_kernel( suffix_lse, # [NUM_HEADS, NUM_TOKENS] prefix_head_stride, output_head_stride, + output_scale, # scale tensor or None HEAD_SIZE: tl.constexpr, PADDED_HEAD_SIZE: tl.constexpr, OUTPUT_LSE: tl.constexpr, prefill_tokens_with_context: tl.constexpr, + USE_FP8: tl.constexpr, + FP8_MIN: tl.constexpr = float8_info.min, + FP8_MAX: tl.constexpr = float8_info.max, ): token_idx = tl.program_id(0) num_tokens = tl.num_programs(0) @@ -87,6 +97,12 @@ def merge_attn_states_kernel( + head_arange, mask=head_mask, ) + + if USE_FP8: + s_out = s_out * (1.0 / tl.load(output_scale)) + s_out = tl.clamp(s_out, FP8_MIN, FP8_MAX) + s_out = s_out.to(output.dtype.element_ty) + tl.store( output + token_idx * num_heads * output_head_stride @@ -143,6 +159,12 @@ def merge_attn_states_kernel( p_scale = p_se / out_se s_scale = s_se / out_se out = p_out * p_scale + s_out * s_scale + + if USE_FP8: + out = out * (1.0 / tl.load(output_scale)) + out = tl.clamp(out, FP8_MIN, FP8_MAX) + out = out.to(output.dtype.element_ty) + tl.store( output + token_idx * num_heads * output_head_stride diff --git a/vllm/v1/attention/ops/triton_reshape_and_cache_flash.py b/vllm/v1/attention/ops/triton_reshape_and_cache_flash.py index eeec609626b..6e696fdb513 100644 --- a/vllm/v1/attention/ops/triton_reshape_and_cache_flash.py +++ b/vllm/v1/attention/ops/triton_reshape_and_cache_flash.py @@ -3,10 +3,16 @@ import torch +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + FP8_DTYPE, + get_fp8_min_max, +) from vllm.platforms import current_platform from vllm.triton_utils import tl, triton from vllm.utils.torch_utils import is_quantized_kv_cache +FP8_MIN, FP8_MAX = get_fp8_min_max() + @triton.jit def reshape_and_cache_kernel_flash( @@ -118,6 +124,198 @@ def reshape_and_cache_kernel_flash( return +# --------------------------------------------------------------------------- +# Per-token-head dynamic quantization kernel +# Grid: (num_tokens, NUM_KV_HEADS) +# Each program handles one (token, head) pair: +# 1. Loads K (or V) for that single head +# 2. Computes absmax across head_size → scale = absmax / QUANT_MAX +# 3. Quantizes and stores the data + per-head scale +# +# Parametrised by QUANT_MAX / QUANT_MIN so the same code path works +# for int8 (±127/128), fp8_e4m3 (±448), and other formats. +# --------------------------------------------------------------------------- +@triton.jit +def _reshape_cache_per_token_head( + key_ptr, # [num_tokens, num_kv_heads, head_size] + value_ptr, # [num_tokens, num_kv_heads, head_size_v] + key_cache_ptr, # [num_blocks, block_size, num_kv_heads, head_size] + value_cache_ptr, # [num_blocks, block_size, num_kv_heads, head_size_v] + k_scale_cache_ptr, # [num_blocks, block_size, num_kv_heads] float32 + v_scale_cache_ptr, # [num_blocks, block_size, num_kv_heads] float32 + slot_mapping_ptr, # [num_tokens] + stride_key_tok: tl.int64, + stride_key_head: tl.int64, + stride_val_tok: tl.int64, + stride_val_head: tl.int64, + stride_kc_blk: tl.int64, # key_cache stride over blocks + stride_kc_slot: tl.int64, # key_cache stride over slots + stride_kc_head: tl.int64, # key_cache stride over heads + stride_vc_blk: tl.int64, + stride_vc_slot: tl.int64, + stride_vc_head: tl.int64, + stride_ks_blk: tl.int64, # k_scale_cache stride[0] (blocks) + stride_ks_slot: tl.int64, # k_scale_cache stride[1] (slots) + stride_ks_head: tl.int64, # k_scale_cache stride[2] (heads) + stride_vs_blk: tl.int64, # v_scale_cache stride[0] (blocks) + stride_vs_slot: tl.int64, # v_scale_cache stride[1] (slots) + stride_vs_head: tl.int64, # v_scale_cache stride[2] (heads) + block_size: tl.constexpr, + head_size: tl.constexpr, + head_size_v: tl.constexpr, + HEAD_SIZE_PADDED: tl.constexpr, # next_power_of_2(max(head_size, head_size_v)) + QUANT_MAX: tl.constexpr = 127.0, + QUANT_MIN: tl.constexpr = -128.0, +): + tok = tl.program_id(0) + head = tl.program_id(1) + + slot = tl.load(slot_mapping_ptr + tok).to(tl.int64) + if slot < 0: + return + + blk = slot // block_size + slot_in_blk = slot % block_size + + dim_offs = tl.arange(0, HEAD_SIZE_PADDED) + + # ---- Key: load one head → absmax → quantize → store ------------------- + k_mask = dim_offs < head_size + k_h = tl.load( + key_ptr + tok * stride_key_tok + head * stride_key_head + dim_offs, + mask=k_mask, + other=0.0, + ).to(tl.float32) + + k_scale = tl.maximum(tl.max(tl.abs(k_h)) / QUANT_MAX, 1e-6) + tl.store( + k_scale_cache_ptr + + blk * stride_ks_blk + + slot_in_blk * stride_ks_slot + + head * stride_ks_head, + k_scale, + ) + + k_q = tl.clamp(k_h * (1.0 / k_scale), QUANT_MIN, QUANT_MAX) + tl.store( + key_cache_ptr + + blk * stride_kc_blk + + slot_in_blk * stride_kc_slot + + head * stride_kc_head + + dim_offs, + k_q, + mask=k_mask, + ) + + # ---- Value: same per-head approach ------------------------------------ + v_mask = dim_offs < head_size_v + v_h = tl.load( + value_ptr + tok * stride_val_tok + head * stride_val_head + dim_offs, + mask=v_mask, + other=0.0, + ).to(tl.float32) + + v_scale = tl.maximum(tl.max(tl.abs(v_h)) / QUANT_MAX, 1e-6) + tl.store( + v_scale_cache_ptr + + blk * stride_vs_blk + + slot_in_blk * stride_vs_slot + + head * stride_vs_head, + v_scale, + ) + + v_q = tl.clamp(v_h * (1.0 / v_scale), QUANT_MIN, QUANT_MAX) + tl.store( + value_cache_ptr + + blk * stride_vc_blk + + slot_in_blk * stride_vc_slot + + head * stride_vc_head + + dim_offs, + v_q, + mask=v_mask, + ) + + +# Mapping from cache torch dtype to (QUANT_MAX, QUANT_MIN) for the +# per-token-head quantization kernel. +_PER_TOKEN_HEAD_QUANT_PARAMS: dict[torch.dtype, tuple[float, float]] = { + torch.int8: (127.0, -128.0), + FP8_DTYPE: (FP8_MAX, FP8_MIN), +} + + +def triton_reshape_and_cache_flash_per_token_head_quant( + key: torch.Tensor, # [num_tokens, num_kv_heads, head_size] + value: torch.Tensor, # [num_tokens, num_kv_heads, head_size_v] + key_cache: torch.Tensor, # [num_blocks, block_size, num_kv_heads, head_size] + value_cache: torch.Tensor, # [num_blocks, block_size, num_kv_heads, head_size_v] + k_scale_cache: torch.Tensor, # [num_blocks, block_size, num_kv_heads] float32 + v_scale_cache: torch.Tensor, # [num_blocks, block_size, num_kv_heads] float32 + slot_mapping: torch.Tensor, # [num_tokens] +): + """Quantize key/value per (token, head) and write to paged cache. + + Computes one scale = absmax / QUANT_MAX per (token, head), stores + quantized data in key_cache/value_cache, and stores the float32 + scale in k_scale_cache/v_scale_cache. + + The quantization range (QUANT_MAX, QUANT_MIN) is derived from the + cache tensor dtype so the same code path works for int8 and fp8. + """ + cache_dtype = key_cache.dtype + quant_params = _PER_TOKEN_HEAD_QUANT_PARAMS.get(cache_dtype) + if quant_params is None: + raise ValueError( + f"Per-token-head quantization not supported for cache dtype " + f"{cache_dtype}. Supported: {list(_PER_TOKEN_HEAD_QUANT_PARAMS)}" + ) + quant_max, quant_min = quant_params + + num_tokens, num_kv_heads, head_size = key.shape + head_size_v = value.shape[2] + head_size_padded = triton.next_power_of_2(max(head_size, head_size_v)) + + block_size = key_cache.shape[1] + + if current_platform.is_rocm() or current_platform.is_xpu(): + num_warps = 4 + else: + num_warps = min(16, max(1, head_size_padded // 32)) + + _reshape_cache_per_token_head[(num_tokens, num_kv_heads)]( + key_ptr=key, + value_ptr=value, + key_cache_ptr=key_cache, + value_cache_ptr=value_cache, + k_scale_cache_ptr=k_scale_cache, + v_scale_cache_ptr=v_scale_cache, + slot_mapping_ptr=slot_mapping, + stride_key_tok=key.stride(0), + stride_key_head=key.stride(1), + stride_val_tok=value.stride(0), + stride_val_head=value.stride(1), + stride_kc_blk=key_cache.stride(0), + stride_kc_slot=key_cache.stride(1), + stride_kc_head=key_cache.stride(2), + stride_vc_blk=value_cache.stride(0), + stride_vc_slot=value_cache.stride(1), + stride_vc_head=value_cache.stride(2), + stride_ks_blk=k_scale_cache.stride(0), + stride_ks_slot=k_scale_cache.stride(1), + stride_ks_head=k_scale_cache.stride(2), + stride_vs_blk=v_scale_cache.stride(0), + stride_vs_slot=v_scale_cache.stride(1), + stride_vs_head=v_scale_cache.stride(2), + block_size=block_size, + head_size=head_size, + head_size_v=head_size_v, + HEAD_SIZE_PADDED=head_size_padded, + QUANT_MAX=quant_max, + QUANT_MIN=quant_min, + num_warps=num_warps, + ) + + def triton_reshape_and_cache_flash( key: torch.Tensor, # [num_tokens, num_heads, head_size] value: torch.Tensor, # [num_tokens, num_heads, head_size] @@ -224,7 +422,6 @@ def triton_reshape_and_cache_flash( block_size=block_size, x=x, USE_HEAD_MAJOR_LAYOUT=use_head_major_layout, - # FP8 flags FP8_KV_CACHE=FP8_KV_CACHE, # autotune parameters TILE_SIZE=TILE_SIZE, diff --git a/vllm/v1/attention/ops/triton_unified_attention.py b/vllm/v1/attention/ops/triton_unified_attention.py index ca5d0e33671..150f022f848 100644 --- a/vllm/v1/attention/ops/triton_unified_attention.py +++ b/vllm/v1/attention/ops/triton_unified_attention.py @@ -13,6 +13,7 @@ import vllm.envs as envs from vllm.logger import init_logger from vllm.platforms import current_platform from vllm.triton_utils import tl, triton +from vllm.v1.kv_cache_interface import KVQuantMode logger = init_logger(__name__) is_batch_invariant = envs.VLLM_BATCH_INVARIANT @@ -32,6 +33,63 @@ def apply_softcap(S, x): return x * (p1 - p2) / (p1 + p2) +@triton.jit +def _prepare_kv_tile( + data, + Q, + tensor_scale, + scale_cache_ptr, + physical_block_idx, + seq_offset, + kv_head_idx, + stride_s_blk, + stride_s_slot, + stride_s_head, + tile_mask, + BLOCK_SIZE: tl.constexpr, + KV_QUANT_MODE: tl.constexpr, +): + """Prepare a loaded KV tile for attention computation. + + Casts the raw KV data to Q's dtype and loads per-token-head scales + when applicable: + + - ``KV_QUANT_MODE == 0``: cast only (no-op for bf16/fp16). + - ``KV_QUANT_MODE == 1`` (FP8 per-tensor): dequantize inline + using the tensor-wide scale. + - ``KV_QUANT_MODE >= 2`` (per-token-head int8/fp8): cast to Q's + dtype and return per-head scales separately — the caller applies + them after the dot product for better numerical efficiency. + + Returns ``(data, token_head_scales)``. *token_head_scales* is only + meaningful when ``KV_QUANT_MODE >= 2``; callers gate its use on + the same constexpr so the compiler eliminates dead code. + """ + # KV_QUANT_MODE values: 0=none, 1=fp8 per-tensor, + # 2=int8 per-token-head, 3=fp8 per-token-head + + # Placeholder scales (float32) — never read when KV_QUANT_MODE < 2. + unused_scales = tile_mask.to(tl.float32) + + if KV_QUANT_MODE == 1: # FP8 per-tensor + if Q.dtype.is_fp8(): + return data.to(Q.dtype), unused_scales + return (data.to(tl.float32) * tl.load(tensor_scale)).to(Q.dtype), unused_scales + if KV_QUANT_MODE >= 2: # per-token-head (int8 or fp8) + scale_idx = ( + physical_block_idx * stride_s_blk + + (seq_offset % BLOCK_SIZE) * stride_s_slot + + kv_head_idx * stride_s_head + ) + token_head_scales = tl.load( + scale_cache_ptr + scale_idx, mask=tile_mask, other=1.0 + ) + return data.to(Q.dtype), token_head_scales + # .to(Q.dtype) is a no-op when data is already Q's type (bf16/fp16), + # but required so Triton sees consistent return types across branches. + return data.to(Q.dtype), unused_scales + + @triton.jit def find_seq_idx( query_start_len_ptr, @@ -105,8 +163,20 @@ def kernel_unified_attention_2d( num_seqs: tl.int32, BLOCK_M: tl.constexpr, # int USE_FP8: tl.constexpr, # bool + # KV cache quantization: 0=none, 1=fp8, 2=per-token-head + KV_QUANT_MODE: tl.constexpr = 0, FP8_MIN: tl.constexpr = float8_info.min, FP8_MAX: tl.constexpr = float8_info.max, + # Per-token-head scale caches (KV_QUANT_MODE >= 2) + # Shape: [num_blocks, block_size, num_kv_heads] + k_scale_cache_ptr=None, + v_scale_cache_ptr=None, + stride_ks_blk=0, + stride_ks_slot=0, + stride_ks_head=0, + stride_vs_blk=0, + stride_vs_slot=0, + stride_vs_head=0, ): q_block_global_idx = tl.program_id(0) kv_head_idx = tl.program_id(1) @@ -258,14 +328,21 @@ def kernel_unified_attention_2d( mask=dim_mask[:, None] & tile_mask[None, :], other=0.0, ) - - if K_load.dtype.is_fp8(): - if Q.dtype.is_fp8(): - K = K_load - else: - K = (K_load.to(tl.float32) * tl.load(k_scale)).to(Q.dtype) - else: - K = K_load + K, k_token_head_scales = _prepare_kv_tile( + K_load, + Q, + k_scale, + k_scale_cache_ptr, + physical_block_idx, + seq_offset, + kv_head_idx, + stride_ks_blk, + stride_ks_slot, + stride_ks_head, + tile_mask, + BLOCK_SIZE, + KV_QUANT_MODE, + ) # V : (TILE_SIZE, HEAD_SIZE) V_load = tl.load( @@ -273,14 +350,21 @@ def kernel_unified_attention_2d( mask=dim_mask[None, :] & tile_mask[:, None], other=0.0, ) - - if V_load.dtype.is_fp8(): - if Q.dtype.is_fp8(): - V = V_load - else: - V = (V_load.to(tl.float32) * tl.load(v_scale)).to(Q.dtype) - else: - V = V_load + V, v_token_head_scales = _prepare_kv_tile( + V_load, + Q, + v_scale, + v_scale_cache_ptr, + physical_block_idx, + seq_offset, + kv_head_idx, + stride_vs_blk, + stride_vs_slot, + stride_vs_head, + tile_mask, + BLOCK_SIZE, + KV_QUANT_MODE, + ) # Compute attention mask: causal by default (key <= query) query_abs_pos = context_len + query_pos[:, None] @@ -318,7 +402,12 @@ def kernel_unified_attention_2d( # S : (BLOCK_M, TILE_SIZE) S = tl.zeros(shape=(BLOCK_M, TILE_SIZE), dtype=tl.float32) - S += scale * tl.dot(Q, K) + # Per-token-head quant: fuse softmax_scale with per-head k_scale + # to avoid a separate BLOCK_M × TILE_SIZE multiply on S. + if KV_QUANT_MODE >= 2: + S += tl.dot(Q, K) * (scale * k_token_head_scales[None, :]) + else: + S += scale * tl.dot(Q, K) if USE_SOFTCAP: S = apply_softcap(S, softcap) @@ -382,7 +471,12 @@ def kernel_unified_attention_2d( ) # acc : (BLOCK_M, HEAD_SIZE_PADDED) - acc += tl.dot(P.to(V.dtype), V) + # Per-token-head quant: apply v_scale to P instead of V. + if KV_QUANT_MODE >= 2: + P_v = (P * v_token_head_scales[None, :]).to(V.dtype) + acc += tl.dot(P_v, V) + else: + acc += tl.dot(P.to(V.dtype), V) # epilogue acc = acc / L[:, None] @@ -453,6 +547,18 @@ def kernel_unified_attention_3d( USE_MM_PREFIX: tl.constexpr, # bool MAX_MM_RANGES: tl.constexpr, # int mm_prefix_range_ptr, # [num_seqs] - prefix length for each sequence + # KV cache quantization: 0=none, 1=fp8, 2=per-token-head + KV_QUANT_MODE: tl.constexpr = 0, + # Per-token-head scale caches (KV_QUANT_MODE >= 2) + # Shape: [num_blocks, block_size, num_kv_heads] + k_scale_cache_ptr=None, + v_scale_cache_ptr=None, + stride_ks_blk=0, + stride_ks_slot=0, + stride_ks_head=0, + stride_vs_blk=0, + stride_vs_slot=0, + stride_vs_head=0, ): q_block_global_idx = tl.program_id(0) kv_head_idx = tl.program_id(1) @@ -613,14 +719,21 @@ def kernel_unified_attention_3d( mask=dim_mask[:, None] & tile_mask[None, :], other=0.0, ) - - if K_load.dtype.is_fp8(): - if Q.dtype.is_fp8(): - K = K_load - else: - K = (K_load.to(tl.float32) * tl.load(k_scale)).to(Q.dtype) - else: - K = K_load + K, k_token_head_scales = _prepare_kv_tile( + K_load, + Q, + k_scale, + k_scale_cache_ptr, + physical_block_idx, + seq_offset, + kv_head_idx, + stride_ks_blk, + stride_ks_slot, + stride_ks_head, + tile_mask, + BLOCK_SIZE, + KV_QUANT_MODE, + ) # V : (TILE_SIZE, HEAD_SIZE) V_load = tl.load( @@ -628,14 +741,21 @@ def kernel_unified_attention_3d( mask=dim_mask[None, :] & tile_mask[:, None], other=0.0, ) - - if V_load.dtype.is_fp8(): - if Q.dtype.is_fp8(): - V = V_load - else: - V = (V_load.to(tl.float32) * tl.load(v_scale)).to(Q.dtype) - else: - V = V_load + V, v_token_head_scales = _prepare_kv_tile( + V_load, + Q, + v_scale, + v_scale_cache_ptr, + physical_block_idx, + seq_offset, + kv_head_idx, + stride_vs_blk, + stride_vs_slot, + stride_vs_head, + tile_mask, + BLOCK_SIZE, + KV_QUANT_MODE, + ) # Compute attention mask: causal by default (key <= query) query_abs_pos = context_len + query_pos[:, None] @@ -672,7 +792,13 @@ def kernel_unified_attention_3d( # S : (BLOCK_M, TILE_SIZE) S = tl.zeros(shape=(BLOCK_M, TILE_SIZE), dtype=tl.float32) - S += scale * tl.dot(Q, K) + + # Per-token-head quant: fuse softmax_scale with per-head k_scale + # to avoid a separate BLOCK_M × TILE_SIZE multiply on S. + if KV_QUANT_MODE >= 2: + S += tl.dot(Q, K) * (scale * k_token_head_scales[None, :]) + else: + S += scale * tl.dot(Q, K) if USE_SOFTCAP: S = apply_softcap(S, softcap) @@ -736,7 +862,12 @@ def kernel_unified_attention_3d( ) # acc : (BLOCK_M, HEAD_SIZE_PADDED) - acc += tl.dot(P.to(V.dtype), V) + # Per-token-head quant: apply v_scale to P instead of V. + if KV_QUANT_MODE >= 2: + P_v = (P * v_token_head_scales[None, :]).to(V.dtype) + acc += tl.dot(P_v, V) + else: + acc += tl.dot(P.to(V.dtype), V) segm_output_offset = ( query_offset_0[:, None].to(tl.int64) @@ -911,6 +1042,10 @@ def unified_attention( # Optional tensor for prefix lengths (PrefixLM support) mm_prefix_range=None, use_alibi_sqrt=False, + # KV cache quantization mode and per-token-head scale caches. + kv_quant_mode: KVQuantMode = KVQuantMode.NONE, + k_scale_cache=None, # [num_blocks, block_size, num_kv_heads] float32 + v_scale_cache=None, # [num_blocks, block_size, num_kv_heads] float32 ): assert causal, "Only causal attention is supported" assert q_descale is None, "Q scales not supported" @@ -1040,6 +1175,15 @@ def unified_attention( num_seqs=num_seqs, BLOCK_M=BLOCK_M, USE_FP8=output_scale is not None, + KV_QUANT_MODE=kv_quant_mode, + k_scale_cache_ptr=k_scale_cache, + v_scale_cache_ptr=v_scale_cache, + stride_ks_blk=k_scale_cache.stride(0) if k_scale_cache is not None else 0, + stride_ks_slot=k_scale_cache.stride(1) if k_scale_cache is not None else 0, + stride_ks_head=k_scale_cache.stride(2) if k_scale_cache is not None else 0, + stride_vs_blk=v_scale_cache.stride(0) if v_scale_cache is not None else 0, + stride_vs_slot=v_scale_cache.stride(1) if v_scale_cache is not None else 0, + stride_vs_head=v_scale_cache.stride(2) if v_scale_cache is not None else 0, ) else: kernel_unified_attention_3d[ @@ -1092,6 +1236,15 @@ def unified_attention( num_seqs=num_seqs, BLOCK_M=BLOCK_M, NUM_SEGMENTS_PER_SEQ=num_par_softmax_segments, + KV_QUANT_MODE=kv_quant_mode, + k_scale_cache_ptr=k_scale_cache, + v_scale_cache_ptr=v_scale_cache, + stride_ks_blk=k_scale_cache.stride(0) if k_scale_cache is not None else 0, + stride_ks_slot=k_scale_cache.stride(1) if k_scale_cache is not None else 0, + stride_ks_head=k_scale_cache.stride(2) if k_scale_cache is not None else 0, + stride_vs_blk=v_scale_cache.stride(0) if v_scale_cache is not None else 0, + stride_vs_slot=v_scale_cache.stride(1) if v_scale_cache is not None else 0, + stride_vs_head=v_scale_cache.stride(2) if v_scale_cache is not None else 0, ) reduce_segments[(q.shape[0], num_query_heads)]( output_ptr=out, diff --git a/vllm/v1/executor/abstract.py b/vllm/v1/executor/abstract.py index 2c3538d9ac2..ae5af0d6f78 100644 --- a/vllm/v1/executor/abstract.py +++ b/vllm/v1/executor/abstract.py @@ -7,6 +7,7 @@ from concurrent.futures import Future from functools import cached_property from typing import TYPE_CHECKING, Literal, TypeVar, overload +import vllm.envs as envs from vllm.config import VllmConfig from vllm.distributed.kv_transfer.kv_connector.utils import KVOutputAggregator from vllm.distributed.kv_transfer.kv_connector.v1.base import ( @@ -57,9 +58,14 @@ class Executor(ABC): ) executor_class = distributed_executor_backend elif distributed_executor_backend == "ray": - from vllm.v1.executor.ray_executor import RayDistributedExecutor + if envs.VLLM_USE_RAY_V2_EXECUTOR_BACKEND: + from vllm.v1.executor.ray_executor_v2 import RayExecutorV2 - executor_class = RayDistributedExecutor + executor_class = RayExecutorV2 + else: + from vllm.v1.executor.ray_executor import RayDistributedExecutor + + executor_class = RayDistributedExecutor elif distributed_executor_backend == "mp": from vllm.v1.executor.multiproc_executor import MultiprocExecutor diff --git a/vllm/v1/executor/multiproc_executor.py b/vllm/v1/executor/multiproc_executor.py index f9b77154067..ac61ded7987 100644 --- a/vllm/v1/executor/multiproc_executor.py +++ b/vllm/v1/executor/multiproc_executor.py @@ -67,25 +67,29 @@ logger = init_logger(__name__) class FutureWrapper(Future): def __init__( self, - futures_queue: deque[tuple["FutureWrapper", Callable]], + futures_queue: deque["FutureWrapper"], + get_response: Callable[[], Any], aggregate: Callable = lambda x: x, ): self.futures_queue = futures_queue + self.get_response = get_response self.aggregate = aggregate super().__init__() + self.futures_queue.appendleft(self) def result(self, timeout=None): if timeout is not None: raise RuntimeError("timeout not implemented") + # Drain any futures ahead of us in the queue. while not self.done(): - future, get_response = self.futures_queue.pop() - future.wait_for_response(get_response) + future = self.futures_queue.pop() + future._wait_for_response() return super().result() - def wait_for_response(self, get_response: Callable): + def _wait_for_response(self): try: - response = self.aggregate(get_response()) + response = self.aggregate(self.get_response()) with suppress(InvalidStateError): self.set_result(response) except Exception as e: @@ -115,7 +119,6 @@ class MultiprocExecutor(Executor): f"_parallel_size ({pcp_size}). " ) - # Set multiprocessing envs set_multiprocessing_worker_envs() # use the loopback address get_loopback_ip() for communication. @@ -168,16 +171,31 @@ class MultiprocExecutor(Executor): for local_rank in range(self.local_world_size): global_rank = global_start_rank + local_rank is_driver_worker = self._is_driver_worker(global_rank) - unready_worker_handle = WorkerProc.make_worker_process( - vllm_config=self.vllm_config, - local_rank=local_rank, - rank=global_rank, - distributed_init_method=distributed_init_method, - input_shm_handle=scheduler_output_handle, - shared_worker_lock=shared_worker_lock, - is_driver_worker=is_driver_worker, - inherited_fds=inherited_fds, - ) + if current_platform.is_cpu(): + om = current_platform.get_omp_manager() + logger.info("Configured OMP PLACES %s", str(om.omp_places)) + unready_worker_handle = om.run( + WorkerProc.make_worker_process, + vllm_config=self.vllm_config, + local_rank=local_rank, + rank=global_rank, + distributed_init_method=distributed_init_method, + input_shm_handle=scheduler_output_handle, + shared_worker_lock=shared_worker_lock, + is_driver_worker=is_driver_worker, + inherited_fds=inherited_fds, + ) + else: + unready_worker_handle = WorkerProc.make_worker_process( + vllm_config=self.vllm_config, + local_rank=local_rank, + rank=global_rank, + distributed_init_method=distributed_init_method, + input_shm_handle=scheduler_output_handle, + shared_worker_lock=shared_worker_lock, + is_driver_worker=is_driver_worker, + inherited_fds=inherited_fds, + ) unready_workers.append(unready_worker_handle) if inherited_fds is not None: inherited_fds.append(unready_worker_handle.death_writer.fileno()) @@ -218,7 +236,7 @@ class MultiprocExecutor(Executor): for response_mq in self.response_mqs: response_mq.wait_until_ready() - self.futures_queue = deque[tuple[FutureWrapper, Callable]]() + self.futures_queue = deque[FutureWrapper]() self._post_init_executor() @@ -384,17 +402,13 @@ class MultiprocExecutor(Executor): responses.append(result) return responses[0] if output_rank is not None else responses - if non_block: - future = FutureWrapper(self.futures_queue, aggregate=aggregate) - self.futures_queue.appendleft((future, get_response)) - return future + future = FutureWrapper( + self.futures_queue, + get_response=get_response, + aggregate=aggregate, + ) - # First drain any pending futures in the queue. - while self.futures_queue: - future, get_fut_response = self.futures_queue.pop() - future.wait_for_response(get_fut_response) - - return aggregate(get_response()) + return future if non_block else future.result() @staticmethod def _ensure_worker_termination(worker_procs: list[BaseProcess]): @@ -1000,24 +1014,26 @@ def set_multiprocessing_worker_envs(): _maybe_force_spawn() - # Configure thread parallelism if OMP_NUM_THREADS isn't set - # - # Helps to avoid CPU contention. The default of spawning a thread per - # core combined with multiprocessing for each GPU can have a negative - # impact on performance. The contention is amplified when running in a - # container where CPU limits can cause throttling. - default_omp_num_threads = 1 - if ( - "OMP_NUM_THREADS" not in os.environ - and (current_parallelism := torch.get_num_threads()) > default_omp_num_threads - ): - logger.warning_once( - "Reducing Torch parallelism from %d threads to %d to avoid " - "unnecessary CPU contention. Set OMP_NUM_THREADS in the " - "external environment to tune this value as needed.", - current_parallelism, - default_omp_num_threads, - scope="local", - ) - os.environ["OMP_NUM_THREADS"] = str(default_omp_num_threads) - torch.set_num_threads(default_omp_num_threads) + if not current_platform.is_cpu(): + # Configure thread parallelism if OMP_NUM_THREADS isn't set + # + # Helps to avoid CPU contention. The default of spawning a thread per + # core combined with multiprocessing for each GPU can have a negative + # impact on performance. The contention is amplified when running in a + # container where CPU limits can cause throttling. + default_omp_num_threads = 1 + if ( + "OMP_NUM_THREADS" not in os.environ + and (current_parallelism := torch.get_num_threads()) + > default_omp_num_threads + ): + logger.warning_once( + "Reducing Torch parallelism from %d threads to %d to avoid " + "unnecessary CPU contention. Set OMP_NUM_THREADS in the " + "external environment to tune this value as needed.", + current_parallelism, + default_omp_num_threads, + scope="local", + ) + os.environ["OMP_NUM_THREADS"] = str(default_omp_num_threads) + torch.set_num_threads(default_omp_num_threads) diff --git a/vllm/v1/executor/ray_env_utils.py b/vllm/v1/executor/ray_env_utils.py new file mode 100644 index 00000000000..6ce12b8ca91 --- /dev/null +++ b/vllm/v1/executor/ray_env_utils.py @@ -0,0 +1,18 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import os + +from vllm.ray.ray_env import RAY_NON_CARRY_OVER_ENV_VARS + + +def get_driver_env_vars( + worker_specific_vars: set[str], +) -> dict[str, str]: + """Return driver env vars to propagate to Ray workers. + + Returns everything from ``os.environ`` except ``worker_specific_vars`` + and user-configured exclusions (``RAY_NON_CARRY_OVER_ENV_VARS``). + """ + exclude_vars = worker_specific_vars | RAY_NON_CARRY_OVER_ENV_VARS + + return {key: value for key, value in os.environ.items() if key not in exclude_vars} diff --git a/vllm/v1/executor/ray_executor.py b/vllm/v1/executor/ray_executor.py index c4e5e7bc67e..1dda2c294e4 100644 --- a/vllm/v1/executor/ray_executor.py +++ b/vllm/v1/executor/ray_executor.py @@ -23,6 +23,7 @@ from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput from vllm.v1.engine import ReconfigureDistributedRequest, ReconfigureRankType from vllm.v1.executor.abstract import Executor from vllm.v1.executor.ray_utils import ( + WORKER_SPECIFIC_ENV_VARS, FutureWrapper, RayWorkerWrapper, initialize_ray_cluster, @@ -62,17 +63,6 @@ class RayWorkerMetaData: class RayDistributedExecutor(Executor): """Ray-based distributed executor""" - # These env vars are worker-specific, therefore are NOT copied - # from the driver to the workers - WORKER_SPECIFIC_ENV_VARS = { - "VLLM_HOST_IP", - "VLLM_HOST_PORT", - "LOCAL_RANK", - "CUDA_VISIBLE_DEVICES", - "HIP_VISIBLE_DEVICES", - "ROCR_VISIBLE_DEVICES", - } - uses_ray: bool = True supports_pp: bool = True @@ -335,7 +325,7 @@ class RayDistributedExecutor(Executor): # Environment variables to copy from driver to workers env_vars_to_copy = get_env_vars_to_copy( - exclude_vars=self.WORKER_SPECIFIC_ENV_VARS, + exclude_vars=WORKER_SPECIFIC_ENV_VARS, additional_vars=set(current_platform.additional_env_vars), destination="workers", ) diff --git a/vllm/v1/executor/ray_executor_v2.py b/vllm/v1/executor/ray_executor_v2.py new file mode 100644 index 00000000000..0665b5fc1b8 --- /dev/null +++ b/vllm/v1/executor/ray_executor_v2.py @@ -0,0 +1,524 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import copy +import os +import threading +import weakref +from collections import defaultdict, deque +from dataclasses import dataclass +from typing import Any + +import vllm.envs as envs +from vllm.config import VllmConfig +from vllm.distributed.device_communicators.shm_broadcast import ( + Handle, + MessageQueue, +) +from vllm.logger import init_logger +from vllm.platforms import current_platform +from vllm.utils.network_utils import ( + get_distributed_init_method, + get_open_port, +) +from vllm.v1.executor.multiproc_executor import ( + FutureWrapper, + MultiprocExecutor, + WorkerProc, +) +from vllm.v1.executor.ray_env_utils import get_driver_env_vars +from vllm.v1.executor.ray_utils import ( + WORKER_SPECIFIC_ENV_VARS, + build_actor_name, + get_bundles_for_indices, + get_bundles_sorted_by_node, + initialize_ray_cluster, + ray, +) + +if ray is not None: + from ray.actor import ActorHandle + from ray.types import ObjectRef + from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy +else: + ActorHandle = None + +logger = init_logger(__name__) + + +@dataclass +class RayWorkerHandle: + """Handle for a Ray worker actor, compatible with MultiprocExecutor.""" + + actor: ActorHandle + """Ray worker actor""" + + rank: int + """Rank of the worker""" + + local_rank: int + """Local rank of the worker""" + + node_id: str + """Node ID of the worker""" + + bundle_id_idx: int = -1 + """Placement group bundle index for the worker""" + + run_ref: ObjectRef | None = None + """run() ObjectRef used as a sentinel for health monitoring""" + + def run(self): + """Start the worker's busy loop""" + self.run_ref = self.actor.run.remote() + + +class RayWorkerProc(WorkerProc): + """Worker process that runs inside a Ray actor. + + Initialization is split into two phases: + 1. __init__: lightweight setup, stores init args (no device/model init) + 2. initialize_worker: called after GPU IDs are discovered, completes + the full WorkerProc initialization with the correct local_rank and + CUDA_VISIBLE_DEVICES. + + CUDA_VISIBLE_DEVICES setup flow: + + 1. RayExecutorV2 enables RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES so Ray does + not set CUDA_VISIBLE_DEVICES on RayWorkerProc actors at creation time. + 2. Each actor is scheduled with a placement group and bundle index; Ray resolves + the physical GPU ID for that bundle at placement time. + 3. After placement, the worker discovers that GPU ID and sets + CUDA_VISIBLE_DEVICES before finishing WorkerProc initialization. + + There is no workaround for this unset-and-reset sequence when the placement group + is externally managed: scheduling must complete before CUDA_VISIBLE_DEVICES can + match the GPU tied to the worker's bundle. + + This sequence allows multiple vLLM instances to coexist on the same node: + each instance is unaware which physical devices others hold, and the + externally managed placement group avoids CUDA_VISIBLE_DEVICES conflicts + by binding workers to specific placement group bundles. + """ + + def __init__( + self, + vllm_config: VllmConfig, + rank: int, + distributed_init_method: str, + input_shm_handle: Handle, + is_driver_worker: bool, + is_driver_node: bool = False, + ): + # Defer WorkerProc.__init__ until GPU IDs are known. + self._is_driver_node = is_driver_node + self._init_kwargs = dict( + vllm_config=vllm_config, + rank=rank, + distributed_init_method=distributed_init_method, + input_shm_handle=input_shm_handle, + shared_worker_lock=None, + is_driver_worker=is_driver_worker, + ) + + def get_node_and_gpu_ids(self) -> tuple[str, list[int]]: + """Return (node_id, gpu_ids) assigned to this actor by Ray.""" + node_id = ray.get_runtime_context().get_node_id() + device_key = current_platform.ray_device_key + if not device_key: + raise RuntimeError( + f"current platform {current_platform.device_name} does not support ray." + ) + gpu_ids = ray.get_runtime_context().get_accelerator_ids()[device_key] + return node_id, [int(x) for x in gpu_ids] + + def initialize_worker( + self, + local_rank: int, + env_vars: dict[str, str], + driver_env_vars: dict[str, str] | None = None, + ) -> None: + """Complete initialization after GPU assignment is known. + + *driver_env_vars* are applied with ``setdefault`` — they fill + in missing vars but never overwrite node-local values. + *env_vars* (e.g. CUDA_VISIBLE_DEVICES) always overwrite. + """ + if driver_env_vars: + for key, value in driver_env_vars.items(): + os.environ.setdefault(key, value) + for key, value in env_vars.items(): + os.environ[key] = value + + self.local_rank = local_rank + super().__init__( + local_rank=local_rank, + **self._init_kwargs, + ) + + def _init_message_queues( + self, input_shm_handle: Handle, vllm_config: VllmConfig + ) -> None: + """ + Workers on the same node as the executor use shared memory for + both the broadcast (input) MQ and the response MQ. Workers on + different nodes use TCP (n_local_reader=0). + """ + self.rpc_broadcast_mq = MessageQueue.create_from_handle( + input_shm_handle, self.worker.rank + ) + + n_local = 1 if self._is_driver_node else 0 + # Use ray.util.get_node_ip_address() to get Ray's internal IP. + # get_ip() returns host's external IP which is typically not + # routable between nodes within the cluster. + self.worker_response_mq = MessageQueue( + n_reader=1, + n_local_reader=n_local, + connect_ip=ray.util.get_node_ip_address(), + ) + self.peer_response_handles: list[dict] = [] + + def wait_for_init(self) -> dict: + """Respond to the driver's wait_until_ready() barrier.""" + assert self.worker_response_mq is not None + return { + "status": self.READY_STR, + "handle": self.worker_response_mq.export_handle(), + } + + def run(self) -> None: + """Main entry point called via actor.run.remote().""" + try: + assert self.rpc_broadcast_mq is not None + self.rpc_broadcast_mq.wait_until_ready() + assert self.worker_response_mq is not None + self.worker_response_mq.wait_until_ready() + + self.worker_busy_loop() + except Exception as e: + logger.exception("RayWorkerProc failed: %s", e) + raise + finally: + self.shutdown() + + +class RayExecutorV2(MultiprocExecutor): + """Ray-based distributed executor using MessageQueue communication. + + Inherits from MultiprocExecutor to reuse the MQ-based control plane + and NCCL data plane. Workers are Ray actors. + + Async scheduling is enabled, inherited from MultiprocExecutor. + This is cricitcal for RayExecutorV2 to be performant. + """ + + uses_ray: bool = True + supports_pp: bool = True + + def __init__(self, vllm_config: VllmConfig): + super().__init__(vllm_config) + + def _build_runtime_env(self) -> dict: + """Build a runtime_env dict for RayWorkerProc actors. + + Driver env vars are applied separately via initialize_worker + with setdefault semantics. + """ + base = self.parallel_config.ray_runtime_env + runtime_env: dict = copy.deepcopy(dict(base)) if base else {} + + env_vars = runtime_env.setdefault("env_vars", {}) + env_vars.update({v: "1" for v in current_platform.ray_noset_device_env_vars}) + if self.parallel_config.ray_workers_use_nsight: + runtime_env["nsight"] = { + "t": "cuda,cudnn,cublas", + "o": "'worker_process_%p'", + "cuda-graph-trace": "node", + } + return runtime_env + + @staticmethod + def _get_actor_resource_kwargs() -> dict[str, Any]: + """Return Ray actor resource kwargs for the current platform.""" + num_devices = envs.VLLM_RAY_PER_WORKER_GPUS + device_key = current_platform.ray_device_key + if device_key == "GPU": + return {"num_gpus": num_devices} + return {"num_gpus": 0, "resources": {device_key: num_devices}} + + def _init_executor(self) -> None: + """Initialize the RayExecutorV2 executor.""" + self._finalizer = weakref.finalize(self, self.shutdown) + self.is_failed = False + self.failure_callback = None + self.shutting_down = False + self.shutdown_lock = threading.Lock() + + # Step 1: Initialize Ray cluster and retrieve placement group + if ray is None: + raise ImportError("Using Ray backend requires installation of ray.") + initialize_ray_cluster(self.parallel_config, require_gpu_on_driver=False) + placement_group = self.parallel_config.placement_group + + tp_size, pp_size, pcp_size = self._get_parallel_sizes() + assert self.world_size == tp_size * pp_size * pcp_size, ( + f"world_size ({self.world_size}) must be equal to the " + f"tensor_parallel_size ({tp_size}) x pipeline" + f"_parallel_size ({pp_size}) x prefill_context" + f"_parallel_size ({pcp_size}). " + ) + + # Step 2: Build bundle assignments for worker rank placement + # while respecting VLLM_RAY_BUNDLE_INDICES. + if envs.VLLM_RAY_BUNDLE_INDICES: + bundle_to_node_id = get_bundles_for_indices( + placement_group, + list(map(int, envs.VLLM_RAY_BUNDLE_INDICES.split(","))), + self.world_size, + ) + else: + bundle_to_node_id = get_bundles_sorted_by_node(placement_group) + driver_node = ray.get_runtime_context().get_node_id() + + bundle_assignments: list[dict[str, Any]] = [] + for rank, (bundle_id_idx, node_id, node_ip) in enumerate(bundle_to_node_id): + bundle_assignments.append( + { + "rank": rank, + "bundle_id_idx": bundle_id_idx, + "node_id": node_id, + "node_ip": node_ip, + } + ) + + # Step 3: Resolve the IP for torch.distributed TCPStore. + # The TCPStore server runs on rank 0's node, so all workers + # must be able to reach this address. + dist_ip = bundle_assignments[0]["node_ip"] + distributed_init_method = get_distributed_init_method(dist_ip, get_open_port()) + + # Step 4: Create broadcast MessageQueue. + # Workers on the driver node use shared memory; the rest use TCP. + max_chunk_bytes = envs.VLLM_MQ_MAX_CHUNK_BYTES_MB * 1024 * 1024 + n_local = sum(1 for a in bundle_assignments if a["node_id"] == driver_node) + self.rpc_broadcast_mq = MessageQueue( + self.world_size, + n_local, + max_chunk_bytes=max_chunk_bytes, + connect_ip=ray.util.get_node_ip_address(), + ) + scheduler_output_handle = self.rpc_broadcast_mq.export_handle() + + # Step 5: Spawn RayWorkerProc actors into PG bundles (deferred init). + # Workers are created lightweight here; full initialization happens + # in Step 7 after GPU IDs are discovered. + self.ray_worker_handles: list[RayWorkerHandle] = [] + instance_id = self.vllm_config.instance_id + + # Collect driver env vars and apply but don't overwrite node-local values. + self.driver_env_vars = get_driver_env_vars( + worker_specific_vars=WORKER_SPECIFIC_ENV_VARS, + ) + + runtime_env = self._build_runtime_env() + resource_kwargs = self._get_actor_resource_kwargs() + + for bundle_idx in range(self.world_size): + bundle = bundle_assignments[bundle_idx] + is_driver_worker = self._is_driver_worker(bundle["rank"]) + is_driver_node = bundle["node_id"] == driver_node + + scheduling_strategy = PlacementGroupSchedulingStrategy( + placement_group=placement_group, + placement_group_bundle_index=bundle["bundle_id_idx"], + ) + + actor_name = build_actor_name( + instance_id, bundle["rank"], tp_size, pp_size, pcp_size + ) + + actor = ( + ray.remote(RayWorkerProc) + .options( + name=actor_name, + num_cpus=0, + **resource_kwargs, + scheduling_strategy=scheduling_strategy, + runtime_env=runtime_env, + ) + .remote( + vllm_config=self.vllm_config, + rank=bundle["rank"], + distributed_init_method=distributed_init_method, + input_shm_handle=scheduler_output_handle, + is_driver_worker=is_driver_worker, + is_driver_node=is_driver_node, + ) + ) + + handle = RayWorkerHandle( + actor=actor, + rank=bundle["rank"], + local_rank=-1, # Set in Step 7 after GPU ID discovery + node_id=bundle["node_id"], + bundle_id_idx=bundle["bundle_id_idx"], + ) + self.ray_worker_handles.append(handle) + + # Step 6: Discover GPU IDs assigned to each worker via Ray runtime context. + worker_node_and_gpu_ids = ray.get( + [h.actor.get_node_and_gpu_ids.remote() for h in self.ray_worker_handles] + ) + + node_workers: dict[str, list[int]] = defaultdict(list) + node_gpus: dict[str, list[int]] = defaultdict(list) + for i, (node_id, gpu_ids) in enumerate(worker_node_and_gpu_ids): + node_workers[node_id].append(i) + node_gpus[node_id].extend(gpu_ids) + for node_id, gpu_ids in node_gpus.items(): + node_gpus[node_id] = sorted(gpu_ids) + + # Step 7: Initialize workers with correct local_rank and + # CUDA_VISIBLE_DEVICES. Each worker sees all GPUs assigned to + # this executor on its node; local_rank indexes into that set. + init_worker_refs = [] + for i, (node_id, _) in enumerate(worker_node_and_gpu_ids): + local_rank = node_workers[node_id].index(i) + worker_env_vars = { + current_platform.device_control_env_var: ",".join( + map(str, node_gpus[node_id]) + ), + } + self.ray_worker_handles[i].local_rank = local_rank + init_worker_refs.append( + self.ray_worker_handles[i].actor.initialize_worker.remote( + local_rank, worker_env_vars, self.driver_env_vars + ) + ) + ray.get(init_worker_refs) + + # Step 8: Collect response MQ handles + init_results = ray.get( + [h.actor.wait_for_init.remote() for h in self.ray_worker_handles] + ) + + self.response_mqs: list[MessageQueue] = [] + for i, result in enumerate(init_results): + if result["status"] != RayWorkerProc.READY_STR: + raise RuntimeError(f"Worker {i} failed to initialize: {result}") + self.response_mqs.append( + MessageQueue.create_from_handle(result["handle"], 0) + ) + + # Step 9: Start run() before wait_until_ready() to avoid + # deadlock — workers send subscriptions inside run(). + for handle in self.ray_worker_handles: + handle.run() + + # Step 10: wait_until_ready() barrier + self.rpc_broadcast_mq.wait_until_ready() + for response_mq in self.response_mqs: + response_mq.wait_until_ready() + + self.futures_queue = deque[FutureWrapper]() + self._post_init_executor() + + self.start_worker_monitor() + self.output_rank = self._get_output_rank() + + def start_worker_monitor(self, inline=False) -> None: + """Monitor worker liveness via ray.wait() on run() ObjectRefs.""" + run_refs = [h.run_ref for h in self.ray_worker_handles if h.run_ref is not None] + if not run_refs: + raise RuntimeError("Ray workers have not started successfully.") + + self_ref = weakref.ref(self) + ref_to_rank = { + h.run_ref: h.rank for h in self.ray_worker_handles if h.run_ref is not None + } + + def _should_stop() -> bool: + executor = self_ref() + return not executor or executor.shutting_down + + def monitor_workers(): + # Poll with a timeout rather than blocking on ray.wait() + # because a blocking call would segfault if Ray is torn down + # while this thread is inside it. + while not _should_stop() and ray.is_initialized(): + try: + done, _ = ray.wait(run_refs, num_returns=1, timeout=5.0) + except Exception: + logger.exception( + "RayWorkerMonitor: unexpected error, exiting monitor thread" + ) + return + if not done or _should_stop(): + continue + + dead_ranks = [ref_to_rank[r] for r in done] + executor = self_ref() + if not executor: + return + executor.is_failed = True + logger.error( + "RayWorkerProc rank=%s died unexpectedly, shutting down executor.", + dead_ranks, + ) + executor.shutdown() + if executor.failure_callback is not None: + callback = executor.failure_callback + executor.failure_callback = None + callback() + return + + t = threading.Thread( + target=monitor_workers, daemon=True, name="RayWorkerMonitor" + ) + t.start() + self._monitor_thread = t + + def _join_monitor_thread(self) -> None: + """Wait for the monitor thread to exit. + + Must be called before tearing down Ray resources — the monitor + may be inside ray.wait() which would segfault if Ray is shut + down underneath it. When the monitor itself calls shutdown() + on worker death, we skip the join because the thread is about + to return anyway. + """ + monitor = getattr(self, "_monitor_thread", None) + if ( + monitor is not None + and monitor.is_alive() + and threading.current_thread() is not monitor + ): + monitor.join(timeout=10) + + def shutdown(self) -> None: + """Properly shut down the executor and its workers.""" + lock = getattr(self, "shutdown_lock", None) + if lock is None: + return + + with lock: + if getattr(self, "shutting_down", False): + return + self.shutting_down = True + + self._join_monitor_thread() + + for handle in getattr(self, "ray_worker_handles", []): + try: + ray.kill(handle.actor) + logger.debug("Killed actor rank=%d", handle.rank) + except Exception: + logger.exception("Failed to kill actor rank=%d", handle.rank) + + if rpc_broadcast_mq := getattr(self, "rpc_broadcast_mq", None): + rpc_broadcast_mq.shutdown() + self.rpc_broadcast_mq = None + + for mq in getattr(self, "response_mqs", []): + mq.shutdown() + self.response_mqs = [] diff --git a/vllm/v1/executor/ray_utils.py b/vllm/v1/executor/ray_utils.py index c363a837b7b..67b43bdc11d 100644 --- a/vllm/v1/executor/ray_utils.py +++ b/vllm/v1/executor/ray_utils.py @@ -26,6 +26,17 @@ if TYPE_CHECKING: logger = init_logger(__name__) PG_WAIT_TIMEOUT = 1800 +# Env vars that are worker-specific and must NOT be copied from the +# driver to Ray workers — they are set per-worker after GPU discovery. +WORKER_SPECIFIC_ENV_VARS: set[str] = { + "VLLM_HOST_IP", + "VLLM_HOST_PORT", + "LOCAL_RANK", + "CUDA_VISIBLE_DEVICES", + "HIP_VISIBLE_DEVICES", + "ROCR_VISIBLE_DEVICES", +} + try: import ray from ray.util import placement_group_table @@ -51,6 +62,8 @@ try: # that thread. self.compiled_dag_cuda_device_set = False + rpc_rank: int + def adjust_rank(self, rank_mapping: dict[int, int]) -> None: """ Adjust the rpc_rank based on the given mapping. @@ -214,13 +227,17 @@ def assert_ray_available(): def _verify_bundles( - placement_group: "PlacementGroup", parallel_config: ParallelConfig, device_str: str + placement_group: "PlacementGroup", + parallel_config: ParallelConfig, + device_str: str, + require_gpu_on_driver: bool = True, ): """Verify a given placement group has bundles located in the right place. There are 2 rules. - Warn if all tensor parallel workers cannot fit in a single node. - - Fail if driver node is not included in a placement group. + - Fail if driver node is not included in a placement group + (only when require_gpu_on_driver is True). """ assert ray.is_initialized(), ( "Ray is not initialized although distributed-executor-backend is ray." @@ -237,7 +254,7 @@ def _verify_bundles( node_id_to_bundle[node_id].append(bundles[bundle_idx]) driver_node_id = ray.get_runtime_context().get_node_id() - if driver_node_id not in node_id_to_bundle: + if require_gpu_on_driver and driver_node_id not in node_id_to_bundle: raise RuntimeError( f"driver node id {driver_node_id} is not included in a placement " f"group {placement_group.id}. Node id -> bundles " @@ -266,6 +283,115 @@ def _verify_bundles( ) +def build_actor_name( + instance_id: str, + rank: int, + tp_size: int, + pp_size: int, + pcp_size: int, +) -> str: + """Build a descriptive Ray actor name for dashboard visibility.""" + name = f"vllm_Worker_{instance_id}" + if tp_size > 1: + name += f"_TP{rank % tp_size}" + if pp_size > 1: + name += f"_PP{(rank // tp_size) % pp_size}" + if pcp_size > 1: + name += f"_PCP{rank // (tp_size * pp_size)}" + return name + + +def get_bundles_for_indices( + placement_group: "PlacementGroup", + bundle_indices: list[int], + world_size: int, +) -> list[tuple[int, str, str]]: + """ + Return GPU bundle indices paired with node IDs and node IPs for + explicit bundle indices specified via VLLM_RAY_BUNDLE_INDICES. + """ + assert len(bundle_indices) == world_size, ( + "VLLM_RAY_BUNDLE_INDICES must have the same size" + f" as the world size, but got {bundle_indices=} " + f"and {world_size=}" + ) + assert len(set(bundle_indices)) == len(bundle_indices), ( + "VLLM_RAY_BUNDLE_INDICES cannot have duplicate values," + f" but got {bundle_indices=}" + ) + + pg_data = placement_group_table(placement_group) + pg_bundle_to_node = pg_data["bundles_to_node_id"] + node_id_to_ip = { + n["NodeID"]: n["NodeManagerAddress"] for n in ray.nodes() if n["Alive"] + } + return [ + (bid, pg_bundle_to_node[bid], node_id_to_ip[pg_bundle_to_node[bid]]) + for bid in bundle_indices + ] + + +def get_bundles_sorted_by_node( + placement_group: "PlacementGroup", +) -> list[tuple[int, str, str]]: + """ + Return GPU bundle indices paired with node IDs and node IPs, + sorted driver-first. + + This utility has to be invoked from the driver node. + + Example: 3-node cluster, driver on node-A, PG bundles spread + across nodes: + + Input: [ + (0, node-C), + (1, node-A), + (2, node-B), + (3, node-C), + (4, node-A), + (5, node-B), + ] + Output: [ + (1, node-A), + (4, node-A), + (2, node-B), + (5, node-B), + (0, node-C), + (3, node-C), + ] + """ + pg_data = placement_group_table(placement_group) + bundle_to_node = pg_data["bundles_to_node_id"] + + ray_device_key = current_platform.ray_device_key + if not ray_device_key: + raise ValueError( + f"current platform {current_platform.device_name} does not support ray." + ) + + node_id_to_ip = { + n["NodeID"]: n["NodeManagerAddress"] for n in ray.nodes() if n["Alive"] + } + + bundle_specs = placement_group.bundle_specs + assert bundle_specs is not None + bundle_to_node_id: list[tuple[int, str, str]] = [] + for bundle_idx, bundle in enumerate(bundle_specs): + if bundle.get(ray_device_key): + node_id = bundle_to_node.get(bundle_idx) + bundle_to_node_id.append((bundle_idx, node_id, node_id_to_ip[node_id])) + + driver_node = ray.get_runtime_context().get_node_id() + + def _sort_key(item): + _, node_id, _ = item + return (0 if node_id == driver_node else 1, node_id) + + bundle_to_node_id.sort(key=_sort_key) + + return bundle_to_node_id + + def _wait_until_pg_ready(current_placement_group: "PlacementGroup"): """Wait until a placement group is ready. @@ -352,6 +478,7 @@ def _wait_until_pg_removed(current_placement_group: "PlacementGroup"): def initialize_ray_cluster( parallel_config: ParallelConfig, ray_address: str | None = None, + require_gpu_on_driver: bool = True, ): """Initialize the distributed cluster with Ray. @@ -363,10 +490,18 @@ def initialize_ray_cluster( parallel_config: The configurations for parallel execution. ray_address: The address of the Ray cluster. If None, uses the default Ray cluster address. + require_gpu_on_driver: If True (default), require at least one GPU + on the current (driver) node and pin the first PG bundle to it. + Set to False for executors like RayExecutorV2 where all GPU work + is delegated to remote Ray actors. """ assert_ray_available() from vllm.platforms import current_platform + # Disable Ray usage stats collection + if os.environ.get("RAY_USAGE_STATS_ENABLED", "0") != "1": + os.environ["RAY_USAGE_STATS_ENABLED"] = "0" + # Prevalidate GPU requirements before Ray processing if current_platform.is_cuda() and parallel_config.world_size > 1: available_gpus = current_platform.device_count() @@ -459,16 +594,20 @@ def initialize_ray_cluster( current_ip = get_ip() current_node_id = ray.get_runtime_context().get_node_id() current_node_resource = available_resources_per_node()[current_node_id] - if current_node_resource.get(device_str, 0) < 1: - raise ValueError( - f"Current node has no {device_str} available. " - f"{current_node_resource=}. vLLM engine cannot start without " - f"{device_str}. Make sure you have at least 1 {device_str} " - f"available in a node {current_node_id=} {current_ip=}." - ) - # This way, at least bundle is required to be created in a current - # node. - placement_group_specs[0][f"node:{current_ip}"] = 0.001 + # TODO (jeffreywang): require_gpu_on_driver should be always False + # after deprecating RayDistributedExecutor. + if require_gpu_on_driver: + if current_node_resource.get(device_str, 0) < 1: + raise ValueError( + f"Current node has no {device_str} available. " + f"{current_node_resource=}. vLLM engine cannot start " + f"without {device_str}. Make sure you have at least 1 " + f"{device_str} available in a node " + f"{current_node_id=} {current_ip=}." + ) + # This way, at least bundle is required to be created in a + # current node. + placement_group_specs[0][f"node:{current_ip}"] = 0.001 # By default, Ray packs resources as much as possible. current_placement_group = ray.util.placement_group( @@ -477,7 +616,9 @@ def initialize_ray_cluster( _wait_until_pg_ready(current_placement_group) assert current_placement_group is not None - _verify_bundles(current_placement_group, parallel_config, device_str) + _verify_bundles( + current_placement_group, parallel_config, device_str, require_gpu_on_driver + ) # Set the placement group in the parallel config parallel_config.placement_group = current_placement_group diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index 48ecf6b9dc8..6f8ad8e7d8e 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/vllm/v1/kv_cache_interface.py @@ -1,21 +1,70 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from __future__ import annotations + import copy from dataclasses import dataclass, fields, replace +from enum import IntEnum from math import prod +from typing import TYPE_CHECKING import torch from typing_extensions import Self -from vllm.config import VllmConfig from vllm.logger import init_logger + +if TYPE_CHECKING: + from vllm.config import VllmConfig from vllm.utils.math_utils import cdiv from vllm.utils.torch_utils import get_dtype_size logger = init_logger(__name__) +# --------------------------------------------------------------------------- +# KV cache quantization mode +# --------------------------------------------------------------------------- + + +class KVQuantMode(IntEnum): + """KV cache quantization mode. + + Used by attention backends and kernels to dispatch quantization logic + without string matching on ``kv_cache_dtype``. + """ + + NONE = 0 + FP8_PER_TENSOR = 1 # per-tensor scales (current fp8 path) + INT8_PER_TOKEN_HEAD = 2 # per-token-head dynamic scales for int8 + FP8_PER_TOKEN_HEAD = 3 # per-token-head dynamic scales for fp8 + + @property + def is_per_token_head(self) -> bool: + """True for any per-token-head quantization mode.""" + return self >= 2 + + +def get_kv_quant_mode(kv_cache_dtype: str) -> KVQuantMode: + """Map a ``kv_cache_dtype`` string to a :class:`KVQuantMode`.""" + if kv_cache_dtype == "int8_per_token_head": + return KVQuantMode.INT8_PER_TOKEN_HEAD + if kv_cache_dtype == "fp8_per_token_head": + return KVQuantMode.FP8_PER_TOKEN_HEAD + if kv_cache_dtype.startswith("fp8"): + return KVQuantMode.FP8_PER_TENSOR + return KVQuantMode.NONE + + +def is_quantized_kv_cache(kv_cache_dtype: str) -> bool: + return get_kv_quant_mode(kv_cache_dtype) != KVQuantMode.NONE + + +def kv_cache_uses_per_token_head_scales(kv_cache_dtype: str) -> bool: + """Return True if *kv_cache_dtype* needs per-token-head scales.""" + return get_kv_quant_mode(kv_cache_dtype).is_per_token_head + + @dataclass(frozen=True) class KVCacheSpec: """ @@ -66,11 +115,19 @@ class AttentionSpec(KVCacheSpec): num_kv_heads: int head_size: int dtype: torch.dtype + kv_quant_mode: KVQuantMode = KVQuantMode.NONE page_size_padded: int | None = None @property def page_size_bytes(self) -> int: real_page_size = self.real_page_size_bytes + # Per-token-head scales are stored in separate tensors managed + # by the attention backend, but the memory is carved from the + # raw KV cache allocation so it must be budgeted here. + if self.kv_quant_mode.is_per_token_head: + real_page_size += ( + 2 * self.block_size * self.num_kv_heads * get_dtype_size(torch.float32) + ) if self.page_size_padded is not None: assert self.page_size_padded >= real_page_size return self.page_size_padded @@ -159,6 +216,7 @@ class FullAttentionSpec(AttentionSpec): head_size=specs[0].head_size, head_size_v=specs[0].head_size_v, dtype=specs[0].dtype, + kv_quant_mode=specs[0].kv_quant_mode, page_size_padded=specs[0].page_size_padded, sliding_window=cls.merge_window_sizes(sliding_window), attention_chunk_size=cls.merge_window_sizes(attention_chunk_size), @@ -220,6 +278,7 @@ class MLAAttentionSpec(FullAttentionSpec): num_kv_heads=specs[0].num_kv_heads, head_size=specs[0].head_size, dtype=specs[0].dtype, + kv_quant_mode=specs[0].kv_quant_mode, page_size_padded=specs[0].page_size_padded, cache_dtype_str=cache_dtype_str_set.pop(), ) @@ -352,6 +411,7 @@ class SinkFullAttentionSpec(FullAttentionSpec): head_size_v=specs[0].head_size_v, sink_len=specs[0].sink_len, dtype=specs[0].dtype, + kv_quant_mode=specs[0].kv_quant_mode, page_size_padded=specs[0].page_size_padded, sliding_window=cls.merge_window_sizes(sliding_window), attention_chunk_size=cls.merge_window_sizes(attention_chunk_size), diff --git a/vllm/v1/kv_offload/worker/cpu_gpu.py b/vllm/v1/kv_offload/worker/cpu_gpu.py index eeabf0cdadd..cd0136d4882 100644 --- a/vllm/v1/kv_offload/worker/cpu_gpu.py +++ b/vllm/v1/kv_offload/worker/cpu_gpu.py @@ -149,6 +149,17 @@ class SingleDirectionOffloadingHandler(OffloadingHandler): # list of CUDA events available for re-use self._event_pool: list[torch.Event] = [] + # Pre-compute base pointers and block sizes for batch copies. + self._src_base_ptrs = np.array( + [t.data_ptr() for t in self.src_tensors], dtype=np.int64 + ) + self._dst_base_ptrs = np.array( + [t.data_ptr() for t in self.dst_tensors], dtype=np.int64 + ) + self._block_size_in_bytes_arr = np.array( + self.tensor_block_size_in_bytes, dtype=np.int64 + ) + def transfer_async(self, job_id: int, transfer_spec: TransferSpec) -> bool: src_spec, dst_spec = transfer_spec assert isinstance(src_spec, BlockIDsLoadStoreSpec) @@ -165,15 +176,35 @@ class SingleDirectionOffloadingHandler(OffloadingHandler): assert dst_sub_block_count == src_sub_block_count - src_sub_blocks_to_skip - src_to_dst = np.empty((dst_sub_block_count, 2), dtype=np.int64) + src_block_ids = np.empty(dst_sub_block_count, dtype=np.int64) + dst_block_ids = np.empty(dst_sub_block_count, dtype=np.int64) expand_block_ids( src_blocks, self.src_block_size_factor, - src_to_dst[:, 0], + src_block_ids, skip_count=src_sub_blocks_to_skip, ) - expand_block_ids(dst_blocks, self.dst_block_size_factor, src_to_dst[:, 1]) - src_to_dst_tensor = torch.from_numpy(src_to_dst) + expand_block_ids(dst_blocks, self.dst_block_size_factor, dst_block_ids) + + # Build flat pointer arrays for all tensors × all block pairs. + num_pairs = dst_sub_block_count + num_tensors = len(self.src_tensors) + total = num_pairs * num_tensors + + all_src = np.empty(total, dtype=np.int64) + all_dst = np.empty(total, dtype=np.int64) + all_sizes = np.empty(total, dtype=np.int64) + + for t_idx, bsz in enumerate(self._block_size_in_bytes_arr): + start = t_idx * num_pairs + end = start + num_pairs + all_src[start:end] = self._src_base_ptrs[t_idx] + src_block_ids * bsz + all_dst[start:end] = self._dst_base_ptrs[t_idx] + dst_block_ids * bsz + all_sizes[start:end] = bsz + + batch_src = torch.from_numpy(all_src) + batch_dst = torch.from_numpy(all_dst) + batch_sizes = torch.from_numpy(all_sizes) stream = self._stream_pool.pop() if self._stream_pool else torch.cuda.Stream() start_event = ( @@ -197,17 +228,8 @@ class SingleDirectionOffloadingHandler(OffloadingHandler): stream.wait_event(last_event) with torch.cuda.stream(stream): start_event.record(stream) - for src_tensor, dst_tensor, block_size_in_bytes in zip( - self.src_tensors, - self.dst_tensors, - self.tensor_block_size_in_bytes, - ): - ops.swap_blocks( - src_tensor, - dst_tensor, - block_size_in_bytes, - src_to_dst_tensor, - ) + if total > 0: + ops.swap_blocks_batch(batch_src, batch_dst, batch_sizes) end_event.record(stream) self._transfer_events[job_id] = end_event diff --git a/vllm/v1/pool/late_interaction.py b/vllm/v1/pool/late_interaction.py index 4a465bd2f7d..554c5947c61 100644 --- a/vllm/v1/pool/late_interaction.py +++ b/vllm/v1/pool/late_interaction.py @@ -56,16 +56,7 @@ def build_late_interaction_doc_params( ) -def compute_maxsim_score( - q_emb: torch.Tensor, - d_emb: torch.Tensor, -) -> torch.Tensor: - # compute in float32 for numerical stability - token_scores = torch.matmul(q_emb.float(), d_emb.float().T) - return token_scores.amax(dim=-1).sum() - - -def compute_maxsim_scores( +def compute_maxsim_score_batched( q_embs: Sequence[torch.Tensor], d_embs: Sequence[torch.Tensor], max_batch_size: int = 64, diff --git a/vllm/v1/worker/cpu_worker.py b/vllm/v1/worker/cpu_worker.py index 2547751c0d8..e759be30d68 100644 --- a/vllm/v1/worker/cpu_worker.py +++ b/vllm/v1/worker/cpu_worker.py @@ -1,18 +1,14 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import os -import platform import sys -from collections.abc import Callable from typing import Any import torch -from vllm import envs from vllm.config import VllmConfig from vllm.logger import init_logger from vllm.platforms import CpuArchEnum, current_platform -from vllm.platforms.cpu import CpuPlatform, LogicalCPUInfo from vllm.profiler.wrapper import TorchProfilerWrapper from vllm.utils.torch_utils import set_random_seed from vllm.v1.worker.cpu_model_runner import CPUModelRunner @@ -71,44 +67,6 @@ class CPUWorker(Worker): if current_platform.get_cpu_architecture() == CpuArchEnum.X86: check_preloaded_libs("libiomp") - # Setup OpenMP threads affinity. - omp_cpuids = envs.VLLM_CPU_OMP_THREADS_BIND - # Under numa binding some cores reserved for kv transfer in nixl_connector.py - if omp_cpuids == "auto" and platform.system() == "Linux": - cpu_arch = current_platform.get_cpu_architecture() - if cpu_arch in (CpuArchEnum.POWERPC, CpuArchEnum.S390X): - # For S390X/POWERPC SMT-8/4/2 - self.local_omp_cpuid = self._get_autobind_cpu_ids( - lambda cpus: [cpu for cpu in cpus if cpu.id % 8 < 4] - ) - elif cpu_arch == CpuArchEnum.X86: - # For x86 SMT-2, use 1 CPU per core - self.local_omp_cpuid = self._get_autobind_cpu_ids( - lambda cpus: cpus[-1:] - ) - elif cpu_arch == CpuArchEnum.ARM: - # For AArch64, no SMT - self.local_omp_cpuid = self._get_autobind_cpu_ids(lambda cpus: cpus) - else: - self.local_omp_cpuid = "nobind" - elif omp_cpuids == "nobind": - self.local_omp_cpuid = "nobind" - else: - local_dp_rank = self.parallel_config.data_parallel_rank_local - omp_cpuids_list = omp_cpuids.split("|") - if local_dp_rank is not None: - world_size = self.parallel_config.world_size - omp_cpuids_list = omp_cpuids_list[ - local_dp_rank * world_size : (local_dp_rank + 1) * world_size - ] - self.local_omp_cpuid = omp_cpuids_list[self.rank] - - if self.local_omp_cpuid != "nobind": - ret = torch.ops._C.init_cpu_threads_env(self.local_omp_cpuid) - if ret: - logger.info(ret) - - # After the thread binding, changing thread num is not allowed def skip_set_num_threads(x: int): logger.warning( "CPU backend doesn't allow to use " @@ -153,92 +111,6 @@ class CPUWorker(Worker): self.model_runner.warming_up_model() return self.compilation_config.compilation_time - def _get_autobind_cpu_ids( - self, cpu_selector: Callable[[list[LogicalCPUInfo]], list[LogicalCPUInfo]] - ) -> str: - """ - Return CPU ids to bind based on NUMA nodes. - Currently for rank N, only CPU ids on the N-th node in available NUMA - node list will be selected. - Args: - cpu_selector: a callable object to select CPUs from a CPU list - of a physical core. The input is a LogicalCPUInfo list, sorted by - the LogicalCPUInfo.id. A selected LogicalCPUInfo list should be - returned. - """ - # simulate multiple numa nodes, for testing - sim_multi_numa_nodes = os.environ.get("VLLM_CPU_SIM_MULTI_NUMA", "0") != "0" - - allowed_numa_nodes, logical_cpu_list = ( - CpuPlatform.get_allowed_cpu_core_node_list() - ) - local_world_size = self.parallel_config.local_world_size - assert len(allowed_numa_nodes) >= local_world_size or sim_multi_numa_nodes, ( - f"Not enough allowed NUMA nodes to bind threads of " - f"{local_world_size} local CPUWorkers. " - f"Allowed NUMA nodes are {allowed_numa_nodes}. " - "Please try to bind threads manually." - ) - - if not sim_multi_numa_nodes: - # Get CPUs on NUMA node `allowed_numa_nodes[local_rank]` - selected_numa_node = allowed_numa_nodes[self.local_rank] # type: ignore - logical_cpu_list = [ - x for x in logical_cpu_list if x.numa_node == selected_numa_node - ] - else: - # This is a bit tricky because the internal DP size - # is always 1 for non-MoE models - world_size_across_dp = ( - self.parallel_config.world_size - * self.parallel_config._api_process_count - ) - assert len(logical_cpu_list) >= world_size_across_dp - logical_cpu_list = sorted(logical_cpu_list, key=lambda x: x.numa_node) - sim_cpu_num_per_node = len(logical_cpu_list) // world_size_across_dp - assert self.parallel_config.data_parallel_rank_local is not None - start_idx = ( - self.local_rank - + self.parallel_config.world_size - * self.parallel_config.data_parallel_rank_local - ) * sim_cpu_num_per_node - logical_cpu_list = logical_cpu_list[ - start_idx : (start_idx + sim_cpu_num_per_node) - ] - - # Select CPUs from each physical core via cpu_selector - core_to_cpus: dict[int, list[LogicalCPUInfo]] = {} - for cpu_info in logical_cpu_list: - if cpu_info.physical_core not in core_to_cpus: - core_to_cpus[cpu_info.physical_core] = [] - core_to_cpus[cpu_info.physical_core].append(cpu_info) - logical_cpu_list = [] - for cpu_list in core_to_cpus.values(): - cpu_list = sorted(cpu_list, key=lambda x: x.id) - logical_cpu_list.extend(cpu_selector(cpu_list)) - logical_cpu_list = sorted(logical_cpu_list, key=lambda x: x.id) - - # Reserve CPUs for other processes - reserve_cpu_num = envs.VLLM_CPU_NUM_OF_RESERVED_CPU - if reserve_cpu_num is None: - need_reserve = ( - self.parallel_config.world_size > 1 - or self.parallel_config.data_parallel_size_local > 1 - ) - reserve_cpu_num = 1 if need_reserve else 0 - assert len(logical_cpu_list) > reserve_cpu_num, ( - f"VLLM_CPU_NUM_OF_RESERVED_CPU ({reserve_cpu_num}) " - f"should less than {len(logical_cpu_list)}." - ) - if reserve_cpu_num != 0: - logical_cpu_list = logical_cpu_list[:-reserve_cpu_num] - - logger.info( - "auto thread-binding list (id, physical core): %s", - [(x.id, x.physical_core) for x in logical_cpu_list], - ) - return ",".join([str(x.id) for x in logical_cpu_list]) - def profile(self, is_start: bool = True, profile_prefix: str | None = None): if self.profiler is None: raise RuntimeError("Profiler is not enabled.") diff --git a/vllm/v1/worker/gpu/kv_connector.py b/vllm/v1/worker/gpu/kv_connector.py index f3acb04a5a4..4349bcfe064 100644 --- a/vllm/v1/worker/gpu/kv_connector.py +++ b/vllm/v1/worker/gpu/kv_connector.py @@ -93,6 +93,10 @@ class ActiveKVConnector(KVConnector): output.invalid_block_ids = self.kv_connector.get_block_ids_with_load_errors() output.kv_connector_stats = self.kv_connector.get_kv_connector_stats() output.kv_cache_events = self.kv_connector.get_kv_connector_kv_cache_events() + output.kv_connector_worker_meta = ( + self.kv_connector.build_connector_worker_meta() + ) + if clear_metadata: self.kv_connector.clear_connector_metadata() return output diff --git a/vllm/v1/worker/gpu/pool/late_interaction_runner.py b/vllm/v1/worker/gpu/pool/late_interaction_runner.py index 221dee55869..da87c8f05d6 100644 --- a/vllm/v1/worker/gpu/pool/late_interaction_runner.py +++ b/vllm/v1/worker/gpu/pool/late_interaction_runner.py @@ -9,7 +9,7 @@ from vllm.v1.outputs import PoolerOutput from vllm.v1.pool.late_interaction import ( LATE_INTERACTION_MODE_CACHE_QUERY, LATE_INTERACTION_MODE_SCORE_DOC, - compute_maxsim_scores, + compute_maxsim_score_batched, ) @@ -116,7 +116,7 @@ class LateInteractionRunner: raise ValueError(f"Unsupported late-interaction mode: {mode!r}") if score_indices: - score_values = compute_maxsim_scores(score_queries, score_docs) + score_values = compute_maxsim_score_batched(score_queries, score_docs) for i, req_id, query_key, score in zip( score_indices, score_req_ids, score_query_keys, score_values ): diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 979ff8d33ca..bba707df04f 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -6077,6 +6077,7 @@ class GPUModelRunner( skip_eplb=True, remove_lora=False, num_active_loras=desc.num_active_loras, + profile_seq_lens=profile_seq_lens, ) self._dummy_run( desc.num_tokens, diff --git a/vllm/v1/worker/worker_base.py b/vllm/v1/worker/worker_base.py index 041fff637b8..6a80e3f3705 100644 --- a/vllm/v1/worker/worker_base.py +++ b/vllm/v1/worker/worker_base.py @@ -195,8 +195,8 @@ class WorkerWrapperBase: All workers have rpc_rank=0, but they have different ranks in the TP group. """ - self.rpc_rank = rpc_rank - self.global_rank = self.rpc_rank if global_rank is None else global_rank + self.rpc_rank: int = rpc_rank + self.global_rank: int = self.rpc_rank if global_rank is None else global_rank # Initialized after init_worker is called self.worker: WorkerBase