From f1d8d99717b6aebf19eac459e0c1fd04bdbe356c Mon Sep 17 00:00:00 2001 From: "Kai K." <59895482+KaletoAI@users.noreply.github.com> Date: Thu, 11 Jun 2026 17:14:21 +0200 Subject: [PATCH 01/52] [Bugfix] CohereModel.load_weights: skip modelopt _quantizer.* keys (#43495) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Kai Köhler --- vllm/model_executor/models/commandr.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/vllm/model_executor/models/commandr.py b/vllm/model_executor/models/commandr.py index 317269ec3b6..66adb9a3ca7 100644 --- a/vllm/model_executor/models/commandr.py +++ b/vllm/model_executor/models/commandr.py @@ -56,6 +56,7 @@ from vllm.sequence import IntermediateTensors from .interfaces import SupportsLoRA, SupportsPP, SupportsQuant from .utils import ( AutoWeightsLoader, + WeightsMapper, extract_layer_index, is_pp_missing_parameter, make_empty_intermediate_tensors_factory, @@ -397,6 +398,9 @@ class CohereForCausalLM(nn.Module, SupportsLoRA, SupportsPP, SupportsQuant): } # LoRA specific attributes embedding_modules = {"embed_tokens": "input_embeddings"} + # ModelOpt NVFP4 checkpoints carry raw quantizer-module state + # (e.g. "*.weight_quantizer._double_scale"); drop them before loading. See #41925. + hf_to_vllm_mapper = WeightsMapper(orig_to_new_substr={"_quantizer.": None}) def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() @@ -453,4 +457,4 @@ class CohereForCausalLM(nn.Module, SupportsLoRA, SupportsPP, SupportsQuant): loader = AutoWeightsLoader( self, skip_prefixes=["lm_head", "rotary_emb.inv_freq"] ) - return loader.load_weights(weights) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) From c2b4cd39acca972da59e53983cf3ddd3b3d32605 Mon Sep 17 00:00:00 2001 From: wineandchord Date: Thu, 11 Jun 2026 23:14:45 +0800 Subject: [PATCH 02/52] [Doc][Attention] Fix MLA top-of-file comments (#37047) Signed-off-by: wineandchord --- .../layers/attention/mla_attention.py | 20 +++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index b04edcc513c..b067cdd00e5 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -14,7 +14,7 @@ MLA has two possible ways of computing, a data-movement friendly approach and a compute friendly approach. We generally want to use the compute friendly approach for "prefill" (i.e. the ratio Sq / Skv is relatively large, often near 1) and the data-movement friendly approach for "decode" (i.e. the ratio -Sq / Skv is small). +Sq / Skv is small, often near 0). NOTE what we deem small and large is currently determined by if it is labelled prefill or decode by the scheduler, but this is something we should probably @@ -28,7 +28,7 @@ Deepseek's MLA attention works the following way: * For decode (i.e. the memory friendly approach) the attention "simulates" a multi-head attention, while the compute is similar to multi-query attention. -Below is example of both paths assuming batchsize = 1 +Below is an example of both paths assuming batch size = 1 ## More Extent Definitions: @@ -77,13 +77,13 @@ v = (kv_c @ W_UV.view(Lkv, N * V)).view(Skv, N, V) // MHA with QK headdim = P + R // V headdim = V -// spda_o shape [Sq, N, V] -spda_o = scaled_dot_product_attention( +// sdpa_o shape [Sq, N, V] +sdpa_o = scaled_dot_product_attention( torch.cat([q_nope, q_pe], dim=-1), torch.cat([k_nope, k_pe.unsqueeze(1).expand(-1, N, -1)], dim=-1), v ) -return spda_o @ W_O +return sdpa_o @ W_O NOTE: in the actual code, `kv_b_proj` is [W_UK; W_UV] concatenated per head @@ -105,16 +105,16 @@ k_pe = torch.cat([new_k_pe, cache_k_pe], dim=0) // MQA with QK headdim = Lkv + R // V headdim = Lkv -// spda_o shape [Sq, N, Lkv] +// sdpa_o shape [Sq, N, Lkv] // NOTE: this is less compute-friendly since Lkv > P // but is more data-movement friendly since its MQA vs MHA -spda_o = scaled_dot_product_attention( +sdpa_o = scaled_dot_product_attention( torch.cat([ql_nope, q_pe], dim=-1), torch.cat([kv_c, k_pe], dim=-1), kv_c ) -o = einsum("snl,lnv->snv", spda_o.reshape(-1, N, Lkv), W_UV) +o = einsum("snl,lnv->snv", sdpa_o.reshape(-1, N, Lkv), W_UV) return o.view(-1, N * V) @ W_O @@ -153,7 +153,7 @@ curr_o, curr_lse = scaled_dot_product_attention( torch.cat([q_nope, q_pe], dim=-1), torch.cat([new_k_nope, new_k_pe.unsqueeze(1).expand(-1, N, -1)], dim=-1), new_v, - casual=True, + causal=True, return_softmax_lse=True ) @@ -173,7 +173,7 @@ for chunk_idx in range(cdiv(C, MCC)): cache_k_pe_chunk.unsqueeze(1).expand(-1, N, -1)], dim=-1), cache_v_chunk, - casual=False, + causal=False, return_softmax_lse=True ) From 23eb7c8fbb7a07d69d10d340db226ee6042a2b02 Mon Sep 17 00:00:00 2001 From: fangyuchu Date: Thu, 11 Jun 2026 23:14:49 +0800 Subject: [PATCH 03/52] [Bugfix] Fix NixlEPAll2AllManager's dependency on --enable-elastic-ep to function (#44422) Signed-off-by: fangyuchu Co-authored-by: Tyler Michael Smith --- vllm/distributed/device_communicators/all2all.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/vllm/distributed/device_communicators/all2all.py b/vllm/distributed/device_communicators/all2all.py index fd1c826322c..967ce5d75c3 100644 --- a/vllm/distributed/device_communicators/all2all.py +++ b/vllm/distributed/device_communicators/all2all.py @@ -9,6 +9,7 @@ import torch.distributed as dist import vllm.envs as envs from vllm.distributed import get_dp_group, get_ep_group +from vllm.distributed.utils import StatelessProcessGroup from vllm.forward_context import get_forward_context from vllm.logger import init_logger from vllm.utils.flashinfer import ( @@ -342,7 +343,12 @@ class NixlEPAll2AllManager(All2AllManagerBase): _lock = threading.RLock() def __init__(self, cpu_group, tcp_store_group=None): - assert tcp_store_group is not None + if tcp_store_group is None: + tcp_store_group = StatelessProcessGroup( + rank=cpu_group.rank(), + world_size=cpu_group.size(), + store=dist.PrefixStore("nixl_ep", cpu_group.get_group_store()), + ) super().__init__(cpu_group, tcp_store_group) self.max_num_ep_ranks = envs.VLLM_NIXL_EP_MAX_NUM_RANKS From 4085ff7cb43d03bbfd05707238ea58a1561f87c2 Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Thu, 11 Jun 2026 08:27:31 -0700 Subject: [PATCH 04/52] [Core] Add kvcache watermark to reduce preemptions (#44594) Signed-off-by: Nick Hill Co-authored-by: Claude Opus 4.8 (1M context) --- benchmarks/kv_cache_watermark.sh | 248 +++++++++++++++++++++++++++++++ tests/v1/core/test_scheduler.py | 2 + tests/v1/core/utils.py | 2 + vllm/config/scheduler.py | 7 + vllm/engine/arg_utils.py | 4 + vllm/v1/core/kv_cache_manager.py | 28 +++- vllm/v1/core/sched/scheduler.py | 8 +- 7 files changed, 291 insertions(+), 8 deletions(-) create mode 100755 benchmarks/kv_cache_watermark.sh diff --git a/benchmarks/kv_cache_watermark.sh b/benchmarks/kv_cache_watermark.sh new file mode 100755 index 00000000000..258afa9fce1 --- /dev/null +++ b/benchmarks/kv_cache_watermark.sh @@ -0,0 +1,248 @@ +#!/bin/bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# +# Reproducible demonstration of the KV cache watermark (`--watermark`) for +# reducing preemption thrashing. +# +# The watermark is the fraction of total KV cache blocks the scheduler keeps +# free when admitting a waiting/preempted request into the running queue. +# +# Why this workload triggers thrashing: +# Requests are admitted based on the KV cache they need *at admission time*. +# With `--scheduler-reserve-full-isl` (default) the input length is reserved up +# front, but the *output* length is unknown and unreserved. A decode-heavy +# workload (output >> input) at high concurrency therefore over-admits while +# requests are short, then runs out of KV cache as they all grow during decode +# -> the scheduler preempts (recompute) recently-admitted requests, re-prefills +# them later, and repeats. The watermark keeps a block of KV cache free so +# running requests can grow into it instead of triggering this churn. +# +# This script launches `vllm serve` under a deliberately KV-constrained config +# and a decode-heavy workload, sweeping the watermark across several values, and +# reports the preemption count (scraped from /metrics), throughput, and latency +# percentiles for each. It then plots the results. +# +# Default workload: concurrency 200, input ~300 tokens, output ~4000 tokens +# (+/- 20% variance), sized to run each config for ~5 minutes. +# +# Usage: +# benchmarks/kv_cache_watermark.sh +# MODEL=Qwen/Qwen2.5-14B-Instruct TP=2 benchmarks/kv_cache_watermark.sh +# +# Run inside the vLLM virtualenv (so `vllm` and `python` resolve to it). +set -euo pipefail + +# ---- Config (override via environment) ------------------------------------- +MODEL=${MODEL:-Qwen/Qwen2.5-7B-Instruct} +TP=${TP:-1} +PORT=${PORT:-8000} +URL="http://127.0.0.1:${PORT}" +# Constrain the KV cache to a *near-critical* size: large enough that the engine +# can run stably, but small enough that greedy over-admission tips it into +# preemption thrashing. (Independent of GPU size, so the demo is reproducible.) +# At the default workload this fits ~1.5x the mean concurrent KV demand. +KV_CACHE_MEMORY_GB=${KV_CACHE_MEMORY_GB:-16} +MAX_MODEL_LEN=${MAX_MODEL_LEN:-8192} +MAX_NUM_SEQS=${MAX_NUM_SEQS:-256} +# Optional weight loader (e.g. fastsafetensors on the GCP cluster). +LOAD_FORMAT=${LOAD_FORMAT:-auto} +# Decode-heavy workload: moderate input, long output, with length variance. The +# long output means preempted requests have generated a lot before eviction, so +# resuming them re-prefills a long sequence (high recomputation cost). +INPUT_LEN=${INPUT_LEN:-1000} +OUTPUT_LEN=${OUTPUT_LEN:-5000} +RANGE_RATIO=${RANGE_RATIO:-0.2} +CONCURRENCY=${CONCURRENCY:-128} +# Enough prompts to keep each config saturated for ~5+ minutes. +NUM_PROMPTS=${NUM_PROMPTS:-450} +OUTDIR=${OUTDIR:-./watermark_bench_results} +# Watermark fractions compared. "label value" per line; value=0 disables it. +CONFIGS=${CONFIGS:-"off 0 +w0.02 0.02 +w0.05 0.05 +w0.10 0.10 +w0.15 0.15"} + +KV_CACHE_MEMORY_BYTES=$((KV_CACHE_MEMORY_GB * 1024 * 1024 * 1024)) +mkdir -p "$OUTDIR" + +SERVER_PID="" +cleanup() { [[ -n "$SERVER_PID" ]] && kill "$SERVER_PID" 2>/dev/null || true; } +trap cleanup EXIT + +scrape_preemptions() { + # Sum the vllm:num_preemptions_total counter across engines. + python - "${URL}/metrics" <<'PY' +import sys, urllib.request +total = 0.0 +try: + body = urllib.request.urlopen(sys.argv[1], timeout=10).read().decode("utf-8", "replace") + for line in body.splitlines(): + if line.startswith("vllm:num_preemptions_total"): + total += float(line.rsplit(" ", 1)[-1]) +except Exception as e: # noqa: BLE001 + print(f"scrape error: {e}", file=sys.stderr) +print(int(total)) +PY +} + +wait_for_server() { + for _ in $(seq 1 300); do + if curl -s "${URL}/health" >/dev/null 2>&1; then return 0; fi + if ! kill -0 "$SERVER_PID" 2>/dev/null; then + echo "ERROR: server process exited during startup" >&2; return 1 + fi + sleep 5 + done + echo "ERROR: server did not become ready" >&2; return 1 +} + +run_one() { + local label=$1 watermark=$2 + echo + echo "==================== watermark: ${label} (${watermark}) ====================" + vllm serve "$MODEL" \ + --tensor-parallel-size "$TP" \ + --load-format "$LOAD_FORMAT" \ + --kv-cache-memory-bytes "$KV_CACHE_MEMORY_BYTES" \ + --max-model-len "$MAX_MODEL_LEN" \ + --max-num-seqs "$MAX_NUM_SEQS" \ + --no-enable-prefix-caching \ + --watermark "$watermark" \ + --port "$PORT" >"${OUTDIR}/serve_${label}.log" 2>&1 & + SERVER_PID=$! + wait_for_server + sleep 5 + + local pre post + pre=$(scrape_preemptions) + vllm bench serve \ + --backend vllm \ + --base-url "$URL" \ + --model "$MODEL" \ + --dataset-name random \ + --random-input-len "$INPUT_LEN" \ + --random-output-len "$OUTPUT_LEN" \ + --random-range-ratio "$RANGE_RATIO" \ + --ignore-eos \ + --num-prompts "$NUM_PROMPTS" \ + --max-concurrency "$CONCURRENCY" \ + --percentile-metrics "ttft,tpot,itl,e2el" \ + --metric-percentiles "50,90,99" \ + --save-result \ + --result-dir "$OUTDIR" \ + --result-filename "bench_${label}.json" + post=$(scrape_preemptions) + echo "${label} ${watermark} $((post - pre))" >>"${OUTDIR}/preemptions.txt" + + kill "$SERVER_PID" 2>/dev/null || true + for _ in $(seq 1 60); do curl -s "${URL}/health" >/dev/null 2>&1 || break; sleep 2; done + SERVER_PID="" + sleep 10 +} + +: >"${OUTDIR}/preemptions.txt" +while read -r label watermark; do + [[ -z "${label:-}" ]] && continue + run_one "$label" "$watermark" +done <<<"$CONFIGS" + +echo +echo "==================== summary ====================" +python - "$OUTDIR" <<'PY' +import json, os, sys +outdir = sys.argv[1] +pre = {} +order = [] +for line in open(os.path.join(outdir, "preemptions.txt")): + label, watermark, n = line.split() + pre[label] = (float(watermark), int(n)) + order.append(label) + +def g(d, *names): + for n in names: + if d.get(n) is not None: + return d[n] + return float("nan") + +cols = ["watermark", "frac", "preempt", "out_tok/s", "req/s", + "TTFT_p50", "TTFT_p99", "ITL_p99", "E2EL_p50"] +print(" ".join(f"{c:>10}" for c in cols)) +rows = [] +for label in order: + watermark, n = pre[label] + d = json.load(open(os.path.join(outdir, f"bench_{label}.json"))) + rows.append(dict( + label=label, watermark=watermark, preempt=n, + out_tok_s=g(d, "output_throughput"), + req_s=g(d, "request_throughput"), + ttft_p50=g(d, "p50_ttft_ms", "median_ttft_ms"), + ttft_p99=g(d, "p99_ttft_ms"), + itl_p99=g(d, "p99_itl_ms"), + e2el_p50=g(d, "p50_e2el_ms", "median_e2el_ms"), + )) + print(" ".join(f"{str(v):>10}" for v in [ + label, watermark, n, + f"{rows[-1]['out_tok_s']:.0f}", + f"{rows[-1]['req_s']:.3f}", + f"{rows[-1]['ttft_p50']/1000:.2f}", + f"{rows[-1]['ttft_p99']/1000:.2f}", + f"{rows[-1]['itl_p99']:.2f}", + f"{rows[-1]['e2el_p50']/1000:.1f}", + ])) +print("\n(TTFT/E2EL in seconds; ITL in ms. Lower preempt is better.)") + +# ---- Plot ------------------------------------------------------------------- +try: + import matplotlib + matplotlib.use("Agg") + import matplotlib.pyplot as plt +except Exception as e: # noqa: BLE001 + print(f"\n(skip plot: matplotlib unavailable: {e})") + sys.exit(0) + +x = [r["watermark"] for r in rows] +xt = [f"{r['watermark']:g}\n({r['label']})" for r in rows] +idx = list(range(len(rows))) + +fig, axes = plt.subplots(2, 2, figsize=(12, 8)) +fig.suptitle( + f"KV cache watermark sweep — {os.path.basename(os.path.abspath(outdir))}", + fontsize=12, +) + +ax = axes[0][0] +ax.bar(idx, [r["preempt"] for r in rows], color="tab:red") +ax.set_title("Preemptions (lower is better)") +ax.set_ylabel("preemptions") +ax.set_xticks(idx); ax.set_xticklabels(xt) + +ax = axes[0][1] +ax.plot(idx, [r["out_tok_s"] for r in rows], "o-", color="tab:green") +ax.set_title("Output throughput (higher is better)") +ax.set_ylabel("tokens/s") +ax.set_xticks(idx); ax.set_xticklabels(xt) + +ax = axes[1][0] +ax.plot(idx, [r["itl_p99"] for r in rows], "o-", color="tab:blue") +ax.set_title("Inter-token latency p99 (lower is better)") +ax.set_ylabel("ITL p99 (ms)") +ax.set_xlabel("watermark fraction") +ax.set_xticks(idx); ax.set_xticklabels(xt) + +ax = axes[1][1] +ax.plot(idx, [r["ttft_p50"] / 1000 for r in rows], "o-", label="TTFT p50") +ax.plot(idx, [r["ttft_p99"] / 1000 for r in rows], "o-", label="TTFT p99") +ax.plot(idx, [r["e2el_p50"] / 1000 for r in rows], "o-", label="E2EL p50") +ax.set_title("Latency (lower is better)") +ax.set_ylabel("seconds") +ax.set_xlabel("watermark fraction") +ax.set_xticks(idx); ax.set_xticklabels(xt) +ax.legend() + +fig.tight_layout(rect=(0, 0, 1, 0.95)) +out_png = os.path.join(outdir, "watermark_results.png") +fig.savefig(out_png, dpi=120) +print(f"\nWrote plot: {out_png}") +PY diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py index 4d652beec81..1b789152e91 100644 --- a/tests/v1/core/test_scheduler.py +++ b/tests/v1/core/test_scheduler.py @@ -1849,6 +1849,8 @@ def create_scheduler_with_priority( enable_chunked_prefill=True, is_encoder_decoder=model_config.is_encoder_decoder, policy="priority", # Enable priority scheduling + # Ensure admission/preemption mechanics are deterministic + watermark=0.0, ) # Cache config, optionally force APC cache_config = CacheConfig( diff --git a/tests/v1/core/utils.py b/tests/v1/core/utils.py index 7213a669c53..7f34250cb21 100644 --- a/tests/v1/core/utils.py +++ b/tests/v1/core/utils.py @@ -90,6 +90,8 @@ def create_scheduler( enable_chunked_prefill=enable_chunked_prefill, async_scheduling=async_scheduling, is_encoder_decoder=model_config.is_encoder_decoder, + # Ensure admission/preemption mechanics are deterministic + watermark=0.0, ) # Cache config, optionally force APC cache_config = CacheConfig( diff --git a/vllm/config/scheduler.py b/vllm/config/scheduler.py index 9669bd1cc41..95f3ed48d47 100644 --- a/vllm/config/scheduler.py +++ b/vllm/config/scheduler.py @@ -143,6 +143,13 @@ class SchedulerConfig: checking the first chunk. Prevents over-admission and KV cache thrashing with chunked prefill.""" + watermark: float = Field(default=0.0, ge=0.0, lt=1.0) + """Fraction of total KV cache blocks to keep free (the watermark) when + admitting waiting or preempted requests into the running queue. This headroom + helps avoid frequent KV cache eviction and the resulting repeated preemption + of requests when GPU memory is scarce. Must be in the range [0.0, 1.0); 0.0 + (the default) disables the watermark.""" + async_scheduling: bool | None = None """If set to False, disable async scheduling. Async scheduling helps to avoid gaps in GPU utilization, leading to better latency and throughput. diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 0490cbc3e4b..f0dade83716 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -600,6 +600,8 @@ class EngineArgs: scheduler_reserve_full_isl: bool = SchedulerConfig.scheduler_reserve_full_isl + watermark: float = SchedulerConfig.watermark + disable_hybrid_kv_cache_manager: bool | None = ( SchedulerConfig.disable_hybrid_kv_cache_manager ) @@ -1408,6 +1410,7 @@ class EngineArgs: "--scheduler-reserve-full-isl", **scheduler_kwargs["scheduler_reserve_full_isl"], ) + scheduler_group.add_argument("--watermark", **scheduler_kwargs["watermark"]) scheduler_group.add_argument( "--disable-hybrid-kv-cache-manager", **scheduler_kwargs["disable_hybrid_kv_cache_manager"], @@ -2045,6 +2048,7 @@ class EngineArgs: max_long_partial_prefills=self.max_long_partial_prefills, long_prefill_token_threshold=self.long_prefill_token_threshold, scheduler_reserve_full_isl=self.scheduler_reserve_full_isl, + watermark=self.watermark, disable_hybrid_kv_cache_manager=self.disable_hybrid_kv_cache_manager, async_scheduling=self.async_scheduling, stream_interval=self.stream_interval, diff --git a/vllm/v1/core/kv_cache_manager.py b/vllm/v1/core/kv_cache_manager.py index 9f0bfc5880c..9af54e0a249 100644 --- a/vllm/v1/core/kv_cache_manager.py +++ b/vllm/v1/core/kv_cache_manager.py @@ -17,7 +17,7 @@ from vllm.v1.kv_cache_interface import ( get_kv_cache_spec_sliding_window, ) from vllm.v1.metrics.stats import PrefixCacheStats -from vllm.v1.request import Request +from vllm.v1.request import Request, RequestStatus logger = init_logger(__name__) @@ -122,6 +122,7 @@ class KVCacheManager: dcp_world_size: int = 1, pcp_world_size: int = 1, metrics_collector: KVCacheMetricsCollector | None = None, + watermark: float = 0.0, ) -> None: self.max_model_len = max_model_len # When unset, fall back to `max_model_len` so the recycling-aware cap @@ -155,6 +156,11 @@ class KVCacheManager: self.num_kv_cache_groups = len(kv_cache_config.kv_cache_groups) self.block_pool = self.coordinator.block_pool self.kv_cache_config = kv_cache_config + + # Watermark: minimum number of KV cache blocks to keep free when + # admitting waiting/preempted requests, to avoid frequent preemptions. + assert watermark >= 0.0, "watermark must be non-negative" + self.watermark_blocks = int(watermark * kv_cache_config.num_blocks) self.kv_cache_event_metadata = tuple( ( get_kv_cache_spec_kind(group.kv_cache_spec).value, @@ -247,6 +253,7 @@ class KVCacheManager: num_encoder_tokens: int = 0, full_sequence_must_fit: bool = False, reserved_blocks: int = 0, + has_scheduled_reqs: bool = True, ) -> KVCacheBlocks | None: """Add slots for a request with new tokens to append. @@ -277,6 +284,8 @@ class KVCacheManager: made if it fits within (free blocks - reserved_blocks). Used to gate async KV-connector loads so their initial allocation cannot consume blocks an already in-flight (prefilling) sequence is relying on. + has_scheduled_reqs: Whether any requests are already scheduled to run + this step, controls whether watermark is applied. Blocks layout: ``` @@ -351,6 +360,15 @@ class KVCacheManager: self.max_model_len, ) + watermark_blocks = 0 + # The watermark is applied to waiting/preempted requests only, and only + # when there's at least one request already scheduled. + if has_scheduled_reqs and request.status in ( + RequestStatus.WAITING, + RequestStatus.PREEMPTED, + ): + watermark_blocks = self.watermark_blocks + if full_sequence_must_fit: # First check and fail if the full request sequence won't fit. full_num_tokens = min(request.num_tokens, self.max_model_len) @@ -364,7 +382,8 @@ class KVCacheManager: num_tokens_main_model=full_num_tokens, apply_admission_cap=True, ) - if num_blocks_to_allocate > self.block_pool.get_num_free_blocks(): + required_blocks = num_blocks_to_allocate + watermark_blocks + if required_blocks > self.block_pool.get_num_free_blocks(): return None num_tokens_main_model = total_computed_tokens + num_new_tokens @@ -392,8 +411,11 @@ class KVCacheManager: num_tokens_main_model=num_tokens_main_model, ) + # Keep `reserved_blocks` free for other in-flight sequences, and an + # additional watermark of headroom for waiting/preempted admissions. available_blocks = self.block_pool.get_num_free_blocks() - reserved_blocks - if num_blocks_to_allocate > available_blocks: + required_blocks = num_blocks_to_allocate + watermark_blocks + if required_blocks > available_blocks: # Cannot allocate new blocks return None diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 160cdb74f57..9a3a9ffa7d6 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -242,6 +242,7 @@ class Scheduler(SchedulerInterface): scheduler_block_size=self.block_size, hash_block_size=hash_block_size, metrics_collector=self.kv_metrics_collector, + watermark=self.scheduler_config.watermark, ) # Bind GPU block pool to the KV connector. This must happen after # kv_cache_manager is constructed so block_pool is available. @@ -826,6 +827,7 @@ class Scheduler(SchedulerInterface): num_encoder_tokens=num_encoder_tokens, full_sequence_must_fit=self.scheduler_reserve_full_isl, reserved_blocks=reserved_blocks, + has_scheduled_reqs=bool(self.running), ) if new_blocks is None: @@ -2198,12 +2200,8 @@ class Scheduler(SchedulerInterface): ) def _inflight_prefill_reserved_blocks(self) -> int: - """Blocks in-flight prefills still need to finish (their reservation). + """Num blocks in-flight prefills still need to finish (their reservation).""" - Sums remaining full-ISL blocks over `self._inflight_prefills` (running - prefills + in-progress async loads). The candidate async load isn't yet - in the set, so it's naturally excluded. - """ return sum( self._request_remaining_blocks(req) for req in self._inflight_prefills ) From f81daf8880632eea46590a8222c082a1e27fd11f Mon Sep 17 00:00:00 2001 From: Jiangyun Zhu Date: Thu, 11 Jun 2026 23:36:31 +0800 Subject: [PATCH 05/52] [Attention] add triton diff-kv backend for mimo (#41797) Signed-off-by: zjy0516 --- .buildkite/test_areas/kernels.yaml | 13 + docs/design/attention_backends.md | 1 + .../test_triton_unified_attention_diffkv.py | 189 +++++++ vllm/model_executor/models/mimo_v2.py | 28 +- .../attention/backends/flash_attn_diffkv.py | 26 +- vllm/v1/attention/backends/registry.py | 3 + .../attention/backends/triton_attn_diffkv.py | 261 +++++++++ .../ops/triton_unified_attention_diffkv.py | 529 ++++++++++++++++++ 8 files changed, 1041 insertions(+), 9 deletions(-) create mode 100644 tests/kernels/attention/test_triton_unified_attention_diffkv.py create mode 100644 vllm/v1/attention/backends/triton_attn_diffkv.py create mode 100644 vllm/v1/attention/ops/triton_unified_attention_diffkv.py diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index 10b5b7527b8..9ec86845038 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -75,6 +75,19 @@ steps: - pytest -v -s kernels/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT parallelism: 2 +- label: Kernels Attention DiffKV Test (H100) + key: kernels-attention-diffkv-test-h100 + timeout_in_minutes: 20 + device: h100 + num_devices: 1 + source_file_dependencies: + - vllm/v1/attention/ops/triton_unified_attention_diffkv.py + - vllm/v1/attention/backends/triton_attn_diffkv.py + - vllm/v1/attention/backends/flash_attn_diffkv.py + - tests/kernels/attention/test_triton_unified_attention_diffkv.py + commands: + - pytest -v -s kernels/attention/test_triton_unified_attention_diffkv.py + - label: Kernels Quantization Test %N key: kernels-quantization-test timeout_in_minutes: 90 diff --git a/docs/design/attention_backends.md b/docs/design/attention_backends.md index 5d366253ef7..9ba7afcb9be 100644 --- a/docs/design/attention_backends.md +++ b/docs/design/attention_backends.md @@ -181,6 +181,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 | | `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 | +| `TRITON_ATTN_DIFFKV` | | fp16, bf16 | `auto`, `bfloat16` | Any | Any | ❌ | ❌ | ❌ | ❌ | Decoder | Any | | `TURBOQUANT` | | fp16, bf16 | `turboquant_k8v4`, `turboquant_4bit_nc`, `turboquant_k3v4_nc`, `turboquant_3bit_nc` | 16, 32, 64, 128 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | Any | > **†** FlashInfer uses TRTLLM attention on Blackwell (SM100), which supports sinks. Disable via `--attention-config.use_trtllm_attention=0`. diff --git a/tests/kernels/attention/test_triton_unified_attention_diffkv.py b/tests/kernels/attention/test_triton_unified_attention_diffkv.py new file mode 100644 index 00000000000..1a19cf34379 --- /dev/null +++ b/tests/kernels/attention/test_triton_unified_attention_diffkv.py @@ -0,0 +1,189 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Unit tests for the Triton DiffKV unified-attention kernel. +""" + +import pytest +import torch + +from vllm.platforms import current_platform +from vllm.utils.math_utils import next_power_of_2 +from vllm.utils.torch_utils import ( + canonicalize_singleton_dim_strides, + set_random_seed, +) +from vllm.v1.attention.backends.fa_utils import ( + get_flash_attn_version, + is_flash_attn_varlen_func_available, +) +from vllm.v1.attention.ops.triton_unified_attention_diffkv import ( + unified_attention_diffkv, +) + +DEVICE_TYPE = current_platform.device_type + +# (num_query_heads, num_kv_heads): MHA, GQA, and the num_kv_heads==1 +# (degenerate-stride) case. +NUM_HEADS = [(4, 4), (8, 2), (5, 1)] +# (head_size_qk, head_size_v). (192, 128) is the canonical asymmetric +# DiffKV shape; FA4 on Blackwell only supports head_size>128 when it is +# 192, and FA3 on Hopper supports it too -- so this pair is runnable on +# both. (128, 128) keeps the equal-dim path covered through the DiffKV +# kernel. +HEAD_SIZES = [(128, 128), (192, 128)] +BLOCK_SIZES = [16] +DTYPES = [torch.bfloat16] + +NUM_BLOCKS = 2048 + +# 0: 2D decode kernel; 8: 3D (split-KV) decode kernel. +SEQ_THRESHOLD_3D_VALUES = [0, 8] + +NUM_PAR_SOFTMAX_SEGMENTS = 16 + + +def _alloc_segm_buffers(seq_threshold_3D: int, num_query_heads: int, head_size_v: int): + """Allocate the split-KV softmax scratch (last dim == head_size_v).""" + head_size_v_padded = next_power_of_2(head_size_v) + segm_output = torch.empty( + ( + seq_threshold_3D, + num_query_heads, + NUM_PAR_SOFTMAX_SEGMENTS, + head_size_v_padded, + ), + dtype=torch.float32, + ) + segm_max = torch.empty( + (seq_threshold_3D, num_query_heads, NUM_PAR_SOFTMAX_SEGMENTS), + dtype=torch.float32, + ) + segm_expsum = torch.empty( + (seq_threshold_3D, num_query_heads, NUM_PAR_SOFTMAX_SEGMENTS), + dtype=torch.float32, + ) + return segm_output, segm_max, segm_expsum + + +@pytest.mark.parametrize( + "seq_lens", + [ + [(1, 1328), (5, 18), (129, 463)], # mixed prefill + decode + [(1, 523), (1, 37), (1, 2011)], # decode-only (exercises 3D path) + ], +) +@pytest.mark.parametrize("num_heads", NUM_HEADS) +@pytest.mark.parametrize("head_sizes", HEAD_SIZES) +@pytest.mark.parametrize("block_size", BLOCK_SIZES) +@pytest.mark.parametrize("sliding_window", [None, 128]) +@pytest.mark.parametrize("soft_cap", [None, 50.0]) +@pytest.mark.parametrize("dtype", DTYPES) +@pytest.mark.parametrize("seq_threshold_3D", SEQ_THRESHOLD_3D_VALUES) +@torch.inference_mode() +def test_triton_unified_attn_diffkv_vs_fa( + seq_lens: list[tuple[int, int]], + num_heads: tuple[int, int], + head_sizes: tuple[int, int], + sliding_window: int | None, + soft_cap: float | None, + dtype: torch.dtype, + block_size: int, + seq_threshold_3D: int, +) -> None: + head_size_qk, head_size_v = head_sizes + + # DiffKV requires FA3 (Hopper) / FA4 (Blackwell) as the reference. + fa_version = get_flash_attn_version(head_size=head_size_qk, head_size_v=head_size_v) + if not is_flash_attn_varlen_func_available() or fa_version not in (3, 4): + pytest.skip(f"FA DiffKV needs FA3/FA4 (got version {fa_version}).") + + from vllm.v1.attention.backends.fa_utils import flash_attn_varlen_func + + torch.set_default_device(DEVICE_TYPE) + set_random_seed(0) + + num_seqs = len(seq_lens) + query_lens = [x[0] for x in seq_lens] + kv_lens = [x[1] for x in seq_lens] + num_query_heads, num_kv_heads = num_heads + assert num_query_heads % num_kv_heads == 0 + max_query_len = max(query_lens) + max_kv_len = max(kv_lens) + window_size = (sliding_window - 1, 0) if sliding_window is not None else (-1, -1) + scale = head_size_qk**-0.5 + + query = torch.randn(sum(query_lens), num_query_heads, head_size_qk, dtype=dtype) + # Packed KV cache: [num_blocks, block_size, num_kv_heads, hqk + hv]. + kv_cache = torch.randn( + NUM_BLOCKS, + block_size, + num_kv_heads, + head_size_qk + head_size_v, + dtype=dtype, + ) + key_cache = kv_cache[..., :head_size_qk] + value_cache = kv_cache[..., head_size_qk:] + + 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 + ) + + # ---- FlashAttention DiffKV (ground truth) --------------------------- + # Mirror the backend: fix degenerate strides on size-1 dims so FA's + # TMA path sees ≥16-byte-aligned strides (matters for num_kv_heads==1). + fa_k = canonicalize_singleton_dim_strides(key_cache) + fa_v = canonicalize_singleton_dim_strides(value_cache) + fa_out = torch.empty(sum(query_lens), num_query_heads, head_size_v, dtype=dtype) + flash_attn_varlen_func( + q=query, + k=fa_k, + v=fa_v, + out=fa_out, + cu_seqlens_q=cu_query_lens, + max_seqlen_q=max_query_len, + seqused_k=kv_lens_t, + max_seqlen_k=max_kv_len, + softmax_scale=scale, + causal=True, + window_size=list(window_size), + block_table=block_tables, + softcap=soft_cap if soft_cap is not None else 0, + fa_version=fa_version, + ) + + # ---- Triton DiffKV -------------------------------------------------- + segm_output, segm_max, segm_expsum = _alloc_segm_buffers( + seq_threshold_3D, num_query_heads, head_size_v + ) + triton_out = torch.empty(sum(query_lens), num_query_heads, head_size_v, dtype=dtype) + unified_attention_diffkv( + q=query, + k=key_cache, + v=value_cache, + out=triton_out, + cu_seqlens_q=cu_query_lens, + seqused_k=kv_lens_t, + softmax_scale=scale, + causal=True, + window_size=window_size, + block_table=block_tables, + softcap=soft_cap if soft_cap is not None else 0, + max_seqlen_q=max_query_len, + seq_threshold_3D=seq_threshold_3D, + num_par_softmax_segments=NUM_PAR_SOFTMAX_SEGMENTS, + softmax_segm_output=segm_output, + softmax_segm_max=segm_max, + softmax_segm_expsum=segm_expsum, + ) + + ( + torch.testing.assert_close(triton_out, fa_out, atol=2e-2, rtol=2e-2), + f"triton vs FA max abs diff: {torch.max(torch.abs(triton_out - fa_out))}", + ) diff --git a/vllm/model_executor/models/mimo_v2.py b/vllm/model_executor/models/mimo_v2.py index 7c6d5363c0a..b5f618699cf 100644 --- a/vllm/model_executor/models/mimo_v2.py +++ b/vllm/model_executor/models/mimo_v2.py @@ -47,9 +47,7 @@ from vllm.model_executor.model_loader.weight_utils import ( from vllm.model_executor.models.utils import sequence_parallel_chunk from vllm.sequence import IntermediateTensors from vllm.v1.attention.backend import AttentionType -from vllm.v1.attention.backends.flash_attn_diffkv import ( - FlashAttentionDiffKVBackend, -) +from vllm.v1.attention.backends.registry import AttentionBackendEnum from .interfaces import MixtureOfExperts, SupportsPP from .utils import ( @@ -292,11 +290,27 @@ class MiMoV2Attention(nn.Module): sliding_window = sliding_window_size if sliding_window_size > -1 else None - # Use DiffKV backend when V has a different head dim than K + # Use DiffKV backend when V has a different head dim than K. + # Auto-pick FA-DiffKV when FA3/4 is usable on this device, else fall + # back to TRITON_ATTN_DIFFKV. Users can force a choice via + # `--attention-backend `. if self.v_head_dim != self.head_dim: - FlashAttentionDiffKVBackend.set_head_size_v(self.v_head_dim) - attn_backend = FlashAttentionDiffKVBackend - logger.info_once("Using FlashAttentionDiffKVBackend for attention.") + requested = get_current_vllm_config().attention_config.backend + if requested is not None and requested.name.endswith("_DIFFKV"): + backend_enum = requested + else: + fa_backend = AttentionBackendEnum.FLASH_ATTN_DIFFKV.get_class() + if fa_backend.is_supported_on_current_device( + head_size=self.head_dim, + head_size_v=self.v_head_dim, + has_sinks=self.attention_sink_bias is not None, + ): + backend_enum = AttentionBackendEnum.FLASH_ATTN_DIFFKV + else: + backend_enum = AttentionBackendEnum.TRITON_ATTN_DIFFKV + attn_backend = backend_enum.get_class() + attn_backend.set_head_size_v(self.v_head_dim) + logger.info_once("Using %s for attention.", attn_backend.get_name()) else: attn_backend = None diff --git a/vllm/v1/attention/backends/flash_attn_diffkv.py b/vllm/v1/attention/backends/flash_attn_diffkv.py index e788b0e3496..ff8fbfc022b 100644 --- a/vllm/v1/attention/backends/flash_attn_diffkv.py +++ b/vllm/v1/attention/backends/flash_attn_diffkv.py @@ -41,6 +41,30 @@ class FlashAttentionDiffKVBackend(FlashAttentionBackend): def set_head_size_v(cls, head_size_v: int) -> None: cls.head_size_v = head_size_v + @classmethod + def is_supported_on_current_device( + cls, + head_size: int, + head_size_v: int, + has_sinks: bool, + ) -> bool: + """Check whether FA3/4 with this DiffKV config is usable here. + + DiffKV (hdim_qk != hdim_v) requires FA3 or FA4 + """ + if not is_flash_attn_varlen_func_available(): + return False + try: + version = get_flash_attn_version( + requires_alibi=False, + head_size=head_size, + head_size_v=head_size_v, + has_sinks=has_sinks, + ) + except Exception: + return False + return version in (3, 4) + @staticmethod def get_name() -> str: return "FLASH_ATTN_DIFFKV" @@ -49,8 +73,6 @@ class FlashAttentionDiffKVBackend(FlashAttentionBackend): def get_impl_cls() -> type["FlashAttentionImpl"]: return FlashAttentionDiffKVImpl - # Do not modify the interface of get_kv_cache_shape, - # but consider head_size_v when returning result. @staticmethod def get_kv_cache_shape( num_blocks: int, diff --git a/vllm/v1/attention/backends/registry.py b/vllm/v1/attention/backends/registry.py index 24a59f03800..2cd2bb5b986 100644 --- a/vllm/v1/attention/backends/registry.py +++ b/vllm/v1/attention/backends/registry.py @@ -46,6 +46,9 @@ class AttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta): "vllm.v1.attention.backends.flash_attn_diffkv.FlashAttentionDiffKVBackend" ) TRITON_ATTN = "vllm.v1.attention.backends.triton_attn.TritonAttentionBackend" + TRITON_ATTN_DIFFKV = ( + "vllm.v1.attention.backends.triton_attn_diffkv.TritonAttentionDiffKVBackend" + ) ROCM_ATTN = "vllm.v1.attention.backends.rocm_attn.RocmAttentionBackend" ROCM_AITER_MLA = "vllm.v1.attention.backends.mla.rocm_aiter_mla.AiterMLABackend" ROCM_AITER_TRITON_MLA = ( diff --git a/vllm/v1/attention/backends/triton_attn_diffkv.py b/vllm/v1/attention/backends/triton_attn_diffkv.py new file mode 100644 index 00000000000..3420a0eba47 --- /dev/null +++ b/vllm/v1/attention/backends/triton_attn_diffkv.py @@ -0,0 +1,261 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Triton attention backend with different K/V head dimensions (DiffKV). + +The KV cache layout is identical to ``FlashAttentionDiffKVBackend`` — K +and V are packed along the last dim: + + [num_blocks, block_size, num_kv_heads, head_size_qk + head_size_v] + +so existing helpers (``triton_reshape_and_cache_flash_diffkv``) are reused. +""" + +from typing import ClassVar + +import torch + +from vllm.config import VllmConfig +from vllm.config.cache import CacheDType +from vllm.logger import init_logger +from vllm.utils.math_utils import next_power_of_2 +from vllm.utils.torch_utils import is_quantized_kv_cache +from vllm.v1.attention.backend import AttentionLayer, AttentionType +from vllm.v1.attention.backends.triton_attn import ( + TritonAttentionBackend, + TritonAttentionImpl, + TritonAttentionMetadata, + TritonAttentionMetadataBuilder, +) +from vllm.v1.attention.backends.utils import get_kv_cache_layout +from vllm.v1.attention.ops.triton_reshape_and_cache_flash import ( + triton_reshape_and_cache_flash_diffkv, +) +from vllm.v1.attention.ops.triton_unified_attention_diffkv import ( + unified_attention_diffkv, +) +from vllm.v1.kv_cache_interface import AttentionSpec + +logger = init_logger(__name__) + + +class TritonAttentionDiffKVMetadataBuilder(TritonAttentionMetadataBuilder): + """Override the parent's softmax buffer last-dim to head_size_v. + + The parent allocates ``softmax_segm_output`` with last-dim sized to + ``next_power_of_2(head_size)`` (== Q/K head size). For DiffKV the + accumulator and per-segment partial outputs are V-shaped, so we + re-allocate with ``next_power_of_2(head_size_v)`` instead. + """ + + def __init__( + self, + kv_cache_spec: AttentionSpec, + layer_names: list[str], + vllm_config: VllmConfig, + device: torch.device, + ): + super().__init__(kv_cache_spec, layer_names, vllm_config, device) + + head_size_v = TritonAttentionDiffKVBackend.head_size_v + head_size_v_padded = next_power_of_2(head_size_v) + self.softmax_segm_output = torch.empty( + ( + self.seq_threshold_3D, + self.num_heads_q, + self.num_par_softmax_segments, + head_size_v_padded, + ), + dtype=torch.float32, + device=device, + ) + + +class TritonAttentionDiffKVBackend(TritonAttentionBackend): + # V head dim — set per layer via ``set_head_size_v`` before instantiation. + head_size_v: int = 128 + + # No FP8 / int8 KV cache for the DiffKV path yet; require fp16/bf16/fp32. + supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [ + "auto", + "bfloat16", + ] + + @classmethod + def set_head_size_v(cls, head_size_v: int) -> None: + cls.head_size_v = head_size_v + + @staticmethod + def get_name() -> str: + return "TRITON_ATTN_DIFFKV" + + @staticmethod + def get_impl_cls() -> type["TritonAttentionDiffKVImpl"]: + return TritonAttentionDiffKVImpl + + @staticmethod + def get_builder_cls() -> type["TritonAttentionDiffKVMetadataBuilder"]: + return TritonAttentionDiffKVMetadataBuilder + + @staticmethod + def get_kv_cache_shape( + num_blocks: int, + block_size: int, + num_kv_heads: int, + head_size: int, + cache_dtype_str: str = "auto", + ) -> tuple[int, ...]: + if block_size % 16 != 0: + raise ValueError("Block size must be a multiple of 16.") + return ( + num_blocks, + block_size, + num_kv_heads, + head_size + TritonAttentionDiffKVBackend.head_size_v, + ) + + @staticmethod + def get_kv_cache_stride_order( + include_num_layers_dimension: bool = False, + ) -> tuple[int, ...]: + cache_layout = get_kv_cache_layout() + if cache_layout == "NHD" and include_num_layers_dimension: + # (num_blocks, num_layers, block_size, + # num_kv_heads, head_size + head_size_v) + return (1, 0, 2, 3, 4) + elif cache_layout == "NHD": + return (0, 1, 2, 3) + elif cache_layout == "HND" and include_num_layers_dimension: + # (num_blocks, num_kv_heads, num_layers, + # block_size, head_size + head_size_v) + return (1, 3, 0, 2, 4) + elif cache_layout == "HND": + return (0, 2, 1, 3) + else: + raise ValueError(f"Unknown cache layout format {cache_layout}.") + + @classmethod + def supports_head_size(cls, head_size: int) -> bool: + # DiffKV K head sizes (e.g. 192 for MiMo-V2.5) need to be allowed. + return head_size >= 32 + + @classmethod + def supports_attn_type(cls, attn_type: str) -> bool: + # DiffKV only implements decoder self-attention. Unlike the parent + # TritonAttentionBackend (which advertises all types), encoder + # attention is not supported, so gate it here at backend selection. + return attn_type == AttentionType.DECODER + + +class TritonAttentionDiffKVImpl(TritonAttentionImpl): + """Triton attention impl for the DiffKV packed KV cache layout.""" + + def __init__(self, *args, **kwargs) -> None: + super().__init__(*args, **kwargs) + if is_quantized_kv_cache(self.kv_cache_dtype): + raise NotImplementedError( + "TritonAttentionDiffKVBackend does not yet support quantized " + f"KV cache (got kv_cache_dtype={self.kv_cache_dtype!r})." + ) + if self._is_per_token_head_quant: + raise NotImplementedError( + "TritonAttentionDiffKVBackend does not support per-token-head " + "quantization." + ) + if self.chunk_lookback > -1: + raise NotImplementedError( + "TritonAttentionDiffKVBackend does not support chunked " + "attention with lookback." + ) + + def do_kv_cache_update( + self, + layer: AttentionLayer, + key: torch.Tensor, + value: torch.Tensor, + kv_cache: torch.Tensor, + slot_mapping: torch.Tensor, + ) -> None: + # Cache is packed [..., head_size_qk + head_size_v]; the diffkv + # reshape kernel writes K to [..., :head_size_qk] and V to + # [..., head_size_qk:hqk+hv]. + triton_reshape_and_cache_flash_diffkv( + key, + value, + kv_cache, + slot_mapping, + self.kv_cache_dtype, + layer._k_scale, + layer._v_scale, + ) + + def fused_rope_kvcache_supported(self): + # The fused rope+cache path assumes the standard 2-tensor layout. + return False + + def forward( + self, + layer: torch.nn.Module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + kv_cache: torch.Tensor, + attn_metadata: TritonAttentionMetadata, + output: torch.Tensor, + output_scale: torch.Tensor | None = None, + output_block_scale: torch.Tensor | None = None, + ) -> torch.Tensor: + """Forward pass. + + Shapes: + query: [num_tokens, num_heads, head_size_qk] + key: [num_tokens, num_kv_heads, head_size_qk] + value: [num_tokens, num_kv_heads, head_size_v] + kv_cache: [num_blocks, block_size, num_kv_heads, + head_size_qk + head_size_v] + output: [num_tokens, num_heads, head_size_v] + """ + if output_scale is not None or output_block_scale is not None: + raise NotImplementedError( + "fused output quantization is not supported for " + "TritonAttentionDiffKVImpl" + ) + + if attn_metadata is None: + return output.fill_(0) + + assert attn_metadata.use_cascade is False, ( + "Cascade attention not supported for TritonAttentionDiffKVImpl" + ) + + num_actual_tokens = attn_metadata.num_actual_tokens + head_size_qk = self.head_size + head_size_v = TritonAttentionDiffKVBackend.head_size_v + + # Slice the packed cache into K / V views. Strides on dims 0/1/2 + # match the original cache; dim 3 stays contiguous (stride 1). + key_cache = kv_cache[..., :head_size_qk] + value_cache = kv_cache[..., head_size_qk : head_size_qk + head_size_v] + + unified_attention_diffkv( + q=query[:num_actual_tokens], + k=key_cache, + v=value_cache, + out=output[:num_actual_tokens], + cu_seqlens_q=attn_metadata.query_start_loc, + seqused_k=attn_metadata.seq_lens, + softmax_scale=self.scale, + causal=True, + alibi_slopes=self.alibi_slopes, + use_alibi_sqrt=self.use_alibi_sqrt, + window_size=self.sliding_window, + block_table=attn_metadata.block_table, + softcap=self.logits_soft_cap, + sinks=self.sinks, + max_seqlen_q=attn_metadata.max_query_len, + seq_threshold_3D=attn_metadata.seq_threshold_3D, + num_par_softmax_segments=attn_metadata.num_par_softmax_segments, + softmax_segm_output=attn_metadata.softmax_segm_output, + softmax_segm_max=attn_metadata.softmax_segm_max, + softmax_segm_expsum=attn_metadata.softmax_segm_expsum, + ) + return output diff --git a/vllm/v1/attention/ops/triton_unified_attention_diffkv.py b/vllm/v1/attention/ops/triton_unified_attention_diffkv.py new file mode 100644 index 00000000000..ef4f2835b5c --- /dev/null +++ b/vllm/v1/attention/ops/triton_unified_attention_diffkv.py @@ -0,0 +1,529 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Triton unified attention with different K/V head dimensions (DiffKV). + +This is a slimmed fork of ``triton_unified_attention.py`` for models like +MiMo-V2.5 where the V tensor's head dimension differs from K's. The KV cache +is the same packed layout used by ``FlashAttentionDiffKVBackend``: + + kv_cache: [num_blocks, block_size, num_kv_heads, head_size_qk + head_size_v] + +We slice ``key_cache = kv_cache[..., :head_size_qk]`` and +``value_cache = kv_cache[..., head_size_qk:]`` on the host, so the kernel +takes two cache pointers but with two distinct head sizes. + +Both 2D and 3D launches are supported: + - 2D: one program per (q-block, kv-head); tile-loop walks the full KV + sequence; final output written directly. Used for prefill and large + decode batches. + - 3D: one program per (q-block, kv-head, segm); each program covers a + KV slice and writes per-segment partials (max/expsum/output). A + follow-up ``kernel_reduce_segments_diffkv`` combines them. Selected + for decode-only batches whose 2D grid would under-fill the GPU. +""" + +from typing import Any + +import torch + +import vllm.envs as envs +from vllm.logger import init_logger +from vllm.triton_utils import tl, triton +from vllm.v1.attention.ops.triton_attention_helpers import ( + apply_alibi_to_score, + apply_softcap, + cdiv_fn, + compute_kv_seq_mask, + compute_tile_loop_bounds, + find_seq_idx, + init_softmax_M, + resolve_seq_and_query_len, + softmax_step, + store_segm_reduce_scalars, +) + +logger = init_logger(__name__) + +is_batch_invariant = envs.VLLM_BATCH_INVARIANT + + +@triton.jit +def kernel_unified_attention_diffkv( + # Output destinations. In 2D mode we write the final result into + # ``output_ptr``; in 3D mode we write per-segment partials into + # ``segm_*`` and ``output_ptr`` is unused (callers may pass any + # non-null pointer). + output_ptr, + segm_output_ptr, + segm_max_ptr, + segm_expsum_ptr, + query_ptr, + key_cache_ptr, # view of packed cache: [..., :head_size_qk] + value_cache_ptr, # view of packed cache: [..., head_size_qk:hqk+hv] + sink_ptr, + block_tables_ptr, + seq_lens_ptr, + alibi_slopes_ptr, + scale, + softcap, + num_query_heads: tl.constexpr, + num_queries_per_kv: tl.constexpr, + block_table_stride: tl.int64, + query_stride_0: tl.int64, + query_stride_1: tl.int64, # == HEAD_SIZE_QK + output_stride_0: tl.int64, + output_stride_1: tl.int64, # == HEAD_SIZE_V + BLOCK_SIZE: tl.constexpr, + TILE_SIZE: tl.constexpr, + HEAD_SIZE_QK: tl.constexpr, + HEAD_SIZE_QK_PADDED: tl.constexpr, + HEAD_SIZE_V: tl.constexpr, + HEAD_SIZE_V_PADDED: tl.constexpr, + USE_ALIBI_SLOPES: tl.constexpr, + USE_ALIBI_SQRT: tl.constexpr, + USE_SOFTCAP: tl.constexpr, + USE_SINKS: tl.constexpr, + SLIDING_WINDOW: tl.constexpr, + # Strides for both cache views (they share the same packed buffer, so + # dims 0/1/2 strides match; only the per-head extent differs). + stride_k_cache_0: tl.int64, + stride_k_cache_1: tl.int64, + stride_k_cache_2: tl.int64, + stride_k_cache_3: tl.constexpr, + stride_v_cache_0: tl.int64, + stride_v_cache_1: tl.int64, + stride_v_cache_2: tl.int64, + stride_v_cache_3: tl.constexpr, + query_start_len_ptr, + BLOCK_Q: tl.constexpr, + num_seqs: tl.int32, + BLOCK_M: tl.constexpr, + NUM_SEGMENTS_PER_SEQ: tl.constexpr, + # ``IS_3D`` toggles between 2D layout (one program walks the full KV + # sequence) and 3D layout (split-KV / FlashDecoding-style: per-segm + # programs write partials, finalized by ``kernel_reduce_segments_diffkv``). + IS_3D: tl.constexpr, +): + q_block_global_idx = tl.program_id(0) + kv_head_idx = tl.program_id(1) + segm_idx = tl.program_id(2) if IS_3D else 0 + + ( + seq_idx, + q_block_local_idx, + cur_batch_in_all_start_index, + cur_batch_query_len, + seq_len, + ) = resolve_seq_and_query_len( + query_start_len_ptr, seq_lens_ptr, q_block_global_idx, num_seqs, BLOCK_Q + ) + + if q_block_local_idx * BLOCK_Q >= cur_batch_query_len: + return + + if IS_3D: + tiles_per_segment = cdiv_fn(seq_len, NUM_SEGMENTS_PER_SEQ * TILE_SIZE) + if segm_idx * tiles_per_segment * TILE_SIZE >= seq_len: + return + else: + tiles_per_segment = 0 + + offs_m = tl.arange(0, BLOCK_M) + offs_d_qk = tl.arange(0, HEAD_SIZE_QK_PADDED) + offs_d_v = tl.arange(0, HEAD_SIZE_V_PADDED) + offs_t = tl.arange(0, TILE_SIZE) + query_pos = q_block_local_idx * BLOCK_Q + offs_m // num_queries_per_kv + + query_offset_0 = cur_batch_in_all_start_index + query_pos + query_offset_1 = kv_head_idx * num_queries_per_kv + offs_m % num_queries_per_kv + query_offset = ( + query_offset_0[:, None] * query_stride_0 + + query_offset_1[:, None] * query_stride_1 + + offs_d_qk[None, :] + ) + + dim_mask_qk = tl.where(offs_d_qk < HEAD_SIZE_QK, 1, 0).to(tl.int1) + dim_mask_v = tl.where(offs_d_v < HEAD_SIZE_V, 1, 0).to(tl.int1) + query_mask_0 = tl.where(query_pos < cur_batch_query_len, 1, 0).to(tl.int1) + query_mask_1 = tl.where(query_offset_1 < num_query_heads, 1, 0).to(tl.int1) + + # Q : (BLOCK_M, HEAD_SIZE_QK_PADDED) + Q = tl.load( + query_ptr + query_offset, + mask=dim_mask_qk[None, :] & query_mask_0[:, None] & query_mask_1[:, None], + other=0.0, + ) + + block_table_offset = seq_idx * block_table_stride + + M = init_softmax_M( + sink_ptr, query_offset_1, query_mask_1, segm_idx, BLOCK_M, USE_SINKS, IS_3D + ) + L = tl.full([BLOCK_M], 1.0, dtype=tl.float32) + # acc : (BLOCK_M, HEAD_SIZE_V_PADDED) + acc = tl.zeros([BLOCK_M, HEAD_SIZE_V_PADDED], dtype=tl.float32) + + context_len = seq_len - cur_batch_query_len + + if USE_ALIBI_SLOPES: + alibi_slope = tl.load( + alibi_slopes_ptr + query_offset_1, mask=query_mask_1, other=0.0 + ) + + loop_lo, loop_hi, max_seq_prefix_len = compute_tile_loop_bounds( + context_len, + seq_len, + cur_batch_query_len, + q_block_local_idx, + segm_idx, + tiles_per_segment, + TILE_SIZE, + BLOCK_M, + BLOCK_Q, + num_queries_per_kv, + SLIDING_WINDOW, + False, # USE_MM_PREFIX + IS_3D, + ) + + for j in range(loop_lo, loop_hi): + seq_offset = j * TILE_SIZE + offs_t + tile_mask = seq_offset < max_seq_prefix_len + + physical_block_idx = tl.load( + block_tables_ptr + block_table_offset + seq_offset // BLOCK_SIZE + ).to(tl.int64) + + v_offset = ( + physical_block_idx[:, None] * stride_v_cache_0 + + kv_head_idx * stride_v_cache_2 + + offs_d_v[None, :] * stride_v_cache_3 + + (seq_offset % BLOCK_SIZE)[:, None] * stride_v_cache_1 + ) + k_offset = ( + physical_block_idx[None, :] * stride_k_cache_0 + + kv_head_idx * stride_k_cache_2 + + offs_d_qk[:, None] * stride_k_cache_3 + + (seq_offset % BLOCK_SIZE)[None, :] * stride_k_cache_1 + ) + # K : (HEAD_SIZE_QK_PADDED, TILE_SIZE) + K_load = tl.load( + key_cache_ptr + k_offset, + mask=dim_mask_qk[:, None] & tile_mask[None, :], + other=0.0, + ) + K = K_load.to(Q.dtype) + # V : (TILE_SIZE, HEAD_SIZE_V_PADDED) + V_load = tl.load( + value_cache_ptr + v_offset, + mask=dim_mask_v[None, :] & tile_mask[:, None], + other=0.0, + ) + V = V_load.to(Q.dtype) + + query_abs_pos = context_len + query_pos[:, None] + seq_mask = compute_kv_seq_mask( + query_abs_pos, + seq_offset, + seq_idx, + None, # mm_prefix_range_ptr + SLIDING_WINDOW, + False, # USE_MM_PREFIX + 0, # MAX_MM_RANGES + ) + + # S : (BLOCK_M, TILE_SIZE) + S = tl.zeros(shape=(BLOCK_M, TILE_SIZE), dtype=tl.float32) + S += scale * tl.dot(Q, K) + + if USE_SOFTCAP: + S = apply_softcap(S, softcap) + + S = tl.where( + query_mask_1[:, None] & query_mask_0[:, None] & seq_mask, S, float("-inf") + ) + + if USE_ALIBI_SLOPES: + S = apply_alibi_to_score( + S, alibi_slope, seq_offset, context_len, query_pos, USE_ALIBI_SQRT + ) + + M, L, P, alpha = softmax_step(S, M, L) + acc = acc * alpha[:, None] + + if SLIDING_WINDOW: + qpos_lo = q_block_local_idx * BLOCK_Q + V = tl.where( + (context_len + qpos_lo - seq_offset[:, None]) < SLIDING_WINDOW, + V, + 0.0, + ) + acc += tl.dot(P.to(V.dtype), V) + + # ---- Epilogue -------------------------------------------------------- + if IS_3D: + # Store per-segment partials; finalized by reduce_segments_diffkv. + segm_output_offset = ( + query_offset_0[:, None].to(tl.int64) + * (num_query_heads * NUM_SEGMENTS_PER_SEQ * HEAD_SIZE_V_PADDED) + + query_offset_1[:, None] * (NUM_SEGMENTS_PER_SEQ * HEAD_SIZE_V_PADDED) + + segm_idx * HEAD_SIZE_V_PADDED + + tl.arange(0, HEAD_SIZE_V_PADDED)[None, :] + ) + tl.store( + segm_output_ptr + segm_output_offset, + acc, + mask=dim_mask_v[None, :] & query_mask_0[:, None] & query_mask_1[:, None], + ) + store_segm_reduce_scalars( + segm_max_ptr, + segm_expsum_ptr, + query_offset_0, + query_offset_1, + segm_idx, + M, + L, + query_mask_0, + query_mask_1, + num_query_heads, + NUM_SEGMENTS_PER_SEQ, + ) + else: + acc = acc / L[:, None] + output_offset = ( + query_offset_0[:, None] * output_stride_0 + + query_offset_1[:, None] * output_stride_1 + + offs_d_v[None, :] + ) + tl.store( + output_ptr + output_offset, + acc, + mask=dim_mask_v[None, :] & query_mask_0[:, None] & query_mask_1[:, None], + ) + + +@triton.jit +def kernel_reduce_segments_diffkv( + output_ptr, # [num_tokens, num_query_heads, head_size_v] + segm_output_ptr, + # [num_tokens, num_query_heads, max_num_segments, head_size_v] + segm_max_ptr, # [num_tokens, num_query_heads, max_num_segments] + segm_expsum_ptr, # [num_tokens, num_query_heads, max_num_segments] + seq_lens_ptr, # [num_seqs] + num_seqs, + num_query_heads: tl.constexpr, + output_stride_0: tl.int64, + output_stride_1: tl.int64, # == HEAD_SIZE_V + TILE_SIZE: tl.constexpr, + HEAD_SIZE_V: tl.constexpr, + HEAD_SIZE_V_PADDED: tl.constexpr, + query_start_len_ptr, # [num_seqs+1] + BLOCK_Q: tl.constexpr, + NUM_SEGMENTS_PER_SEQ: tl.constexpr, +): + """Combine per-segment partials into the final softmax output. + + Mirrors ``reduce_segments`` from triton_unified_attention.py but + indexes V's head size (``HEAD_SIZE_V``) instead of the shared one. + """ + query_token_idx = tl.program_id(0) + query_head_idx = tl.program_id(1) + + seq_idx = find_seq_idx( + query_start_len_ptr, query_token_idx, num_seqs, BLOCK_Q, False + ) + seq_len = tl.load(seq_lens_ptr + seq_idx) + + tiles_per_segment = cdiv_fn(seq_len, NUM_SEGMENTS_PER_SEQ * TILE_SIZE) + act_num_segments = cdiv_fn(seq_len, tiles_per_segment * TILE_SIZE) + segm_mask = tl.arange(0, NUM_SEGMENTS_PER_SEQ) < tl.full( + [NUM_SEGMENTS_PER_SEQ], act_num_segments, dtype=tl.int32 + ) + dim_mask = tl.where(tl.arange(0, HEAD_SIZE_V_PADDED) < HEAD_SIZE_V, 1, 0).to( + tl.int1 + ) + + segm_offset = ( + query_token_idx.to(tl.int64) * (num_query_heads * NUM_SEGMENTS_PER_SEQ) + + query_head_idx * NUM_SEGMENTS_PER_SEQ + + tl.arange(0, NUM_SEGMENTS_PER_SEQ) + ) + segm_max = tl.load(segm_max_ptr + segm_offset, mask=segm_mask, other=float("-inf")) + overall_max = tl.max(segm_max) + + segm_expsum = tl.load(segm_expsum_ptr + segm_offset, mask=segm_mask, other=0.0) + segm_expsum = segm_expsum * tl.exp(segm_max - overall_max) + overall_expsum = tl.sum(segm_expsum) + + segm_output_offset = ( + query_token_idx.to(tl.int64) + * (num_query_heads * NUM_SEGMENTS_PER_SEQ * HEAD_SIZE_V_PADDED) + + query_head_idx * (NUM_SEGMENTS_PER_SEQ * HEAD_SIZE_V_PADDED) + + tl.arange(0, NUM_SEGMENTS_PER_SEQ)[:, None] * HEAD_SIZE_V_PADDED + + tl.arange(0, HEAD_SIZE_V_PADDED)[None, :] + ) + segm_output = tl.load( + segm_output_ptr + segm_output_offset, + mask=segm_mask[:, None] & dim_mask[None, :], + other=0.0, + ) + segm_output *= tl.exp(segm_max - overall_max)[:, None] + acc_sum = tl.sum(segm_output, axis=0) + acc = tl.where(overall_expsum == 0.0, 0.0, acc_sum / overall_expsum) + + output_offset = ( + query_token_idx * output_stride_0 + + query_head_idx * output_stride_1 + + tl.arange(0, HEAD_SIZE_V_PADDED) + ) + tl.store(output_ptr + output_offset, acc, mask=dim_mask) + + +def unified_attention_diffkv( + q, # [num_tokens, num_query_heads, head_size_qk] + k, # view: [num_blocks, block_size, num_kv_heads, head_size_qk] + v, # view: [num_blocks, block_size, num_kv_heads, head_size_v] + out, # [num_tokens, num_query_heads, head_size_v] + cu_seqlens_q, + seqused_k, + softmax_scale, + causal, + window_size, + block_table, + softcap, + max_seqlen_q: int = 1, + alibi_slopes=None, + sinks=None, + use_alibi_sqrt=False, + # 3D / split-KV softmax buffers. When all four are provided and the + # batch is decode-only with few sequences, the 3D path is taken. + seq_threshold_3D: int | None = None, + num_par_softmax_segments: int | None = None, + softmax_segm_output: torch.Tensor | None = None, + softmax_segm_max: torch.Tensor | None = None, + softmax_segm_expsum: torch.Tensor | None = None, +): + assert causal, "Only causal attention is supported" + + if sinks is not None: + assert sinks.shape[0] == q.shape[1], "Sinks must be num_query_heads size" + + use_alibi_slopes = alibi_slopes is not None + + block_size = v.shape[1] + num_seqs = len(seqused_k) + num_query_heads = q.shape[1] + num_kv_heads = k.shape[2] + num_queries_per_kv = num_query_heads // num_kv_heads + head_size_qk = q.shape[2] + head_size_v = v.shape[3] + + BLOCK_M = ( + 16 if num_queries_per_kv <= 16 else triton.next_power_of_2(num_queries_per_kv) + ) + BLOCK_Q = BLOCK_M // num_queries_per_kv + + total_num_q_blocks = q.shape[0] // BLOCK_Q + num_seqs + + sliding_window_val = 1 + window_size[0] if window_size[0] >= 0 else 0 + + # Decide between 2D and 3D launch. Mirrors the standard launcher: + # 3D requires preallocated softmax buffers, decode-only batches, and + # a small number of sequences (otherwise 2D already saturates the SM). + use_3d = not ( + seq_threshold_3D is None + or num_par_softmax_segments is None + or softmax_segm_output is None + or softmax_segm_max is None + or softmax_segm_expsum is None + or max_seqlen_q > 1 + or num_seqs > seq_threshold_3D + or is_batch_invariant + ) + + # Tile size: 32 for prefill-class kernels. Decode (small Q) prefers + # smaller tiles to expose more parallelism along the KV dim. + tile_size = 32 if not use_3d else (16 if q.element_size() >= 2 else 32) + + grid: tuple[Any, ...] + if use_3d: + grid = (total_num_q_blocks, num_kv_heads, num_par_softmax_segments) + segm_output_ptr = softmax_segm_output + segm_max_ptr = softmax_segm_max + segm_expsum_ptr = softmax_segm_expsum + num_segments = num_par_softmax_segments + else: + grid = (total_num_q_blocks, num_kv_heads) + # 2D never touches the segm tensors but Triton wants a non-null + # pointer; reuse ``out``. + segm_output_ptr = out + segm_max_ptr = out + segm_expsum_ptr = out + num_segments = 1 + + kernel_unified_attention_diffkv[grid]( + output_ptr=out, + segm_output_ptr=segm_output_ptr, + segm_max_ptr=segm_max_ptr, + segm_expsum_ptr=segm_expsum_ptr, + query_ptr=q, + key_cache_ptr=k, + value_cache_ptr=v, + sink_ptr=sinks, + block_tables_ptr=block_table, + seq_lens_ptr=seqused_k, + alibi_slopes_ptr=alibi_slopes, + scale=softmax_scale, + softcap=softcap, + num_query_heads=num_query_heads, + num_queries_per_kv=num_queries_per_kv, + block_table_stride=block_table.stride(0), + query_stride_0=q.stride(0), + query_stride_1=q.stride(1), + output_stride_0=out.stride(0), + output_stride_1=out.stride(1), + BLOCK_SIZE=block_size, + TILE_SIZE=tile_size, + HEAD_SIZE_QK=head_size_qk, + HEAD_SIZE_QK_PADDED=triton.next_power_of_2(head_size_qk), + HEAD_SIZE_V=head_size_v, + HEAD_SIZE_V_PADDED=triton.next_power_of_2(head_size_v), + USE_ALIBI_SLOPES=use_alibi_slopes, + USE_ALIBI_SQRT=use_alibi_sqrt, + USE_SOFTCAP=(softcap > 0), + USE_SINKS=(sinks is not None), + SLIDING_WINDOW=sliding_window_val, + stride_k_cache_0=k.stride(0), + stride_k_cache_1=k.stride(1), + stride_k_cache_2=k.stride(2), + stride_k_cache_3=k.stride(3), + stride_v_cache_0=v.stride(0), + stride_v_cache_1=v.stride(1), + stride_v_cache_2=v.stride(2), + stride_v_cache_3=v.stride(3), + query_start_len_ptr=cu_seqlens_q, + BLOCK_Q=BLOCK_Q, + num_seqs=num_seqs, + BLOCK_M=BLOCK_M, + NUM_SEGMENTS_PER_SEQ=num_segments, + IS_3D=use_3d, + ) + + if use_3d: + kernel_reduce_segments_diffkv[(q.shape[0], num_query_heads)]( + output_ptr=out, + segm_output_ptr=softmax_segm_output, + segm_max_ptr=softmax_segm_max, + segm_expsum_ptr=softmax_segm_expsum, + seq_lens_ptr=seqused_k, + num_seqs=num_seqs, + num_query_heads=num_query_heads, + output_stride_0=out.stride(0), + output_stride_1=out.stride(1), + TILE_SIZE=tile_size, + HEAD_SIZE_V=head_size_v, + HEAD_SIZE_V_PADDED=triton.next_power_of_2(head_size_v), + query_start_len_ptr=cu_seqlens_q, + BLOCK_Q=BLOCK_Q, + NUM_SEGMENTS_PER_SEQ=num_par_softmax_segments, + ) From 79f8c5bd8c8a6be1519e6c569653e502a06cd46b Mon Sep 17 00:00:00 2001 From: vraiti Date: Thu, 11 Jun 2026 11:43:14 -0400 Subject: [PATCH 06/52] [Metrics] Scope unregister_vllm_metrics() to strictly "vllm:" metrics (#42331) `unregister_vllm_metrics()` currently uses "vllm" in `collector._name` to decide which collectors to remove from the Prometheus registry, removing every even metrics registered by other subsystems or downstream extensions like "vllm_omni:" Signed-off-by: vraiti Signed-off-by: Mark McLoughlin --- vllm/v1/metrics/prometheus.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/v1/metrics/prometheus.py b/vllm/v1/metrics/prometheus.py index 1eacb785aa8..c8740276713 100644 --- a/vllm/v1/metrics/prometheus.py +++ b/vllm/v1/metrics/prometheus.py @@ -64,7 +64,7 @@ def unregister_vllm_metrics(): registry = REGISTRY # Unregister any existing vLLM collectors for collector in list(registry._collector_to_names): - if hasattr(collector, "_name") and "vllm" in collector._name: + if hasattr(collector, "_name") and collector._name.startswith("vllm:"): registry.unregister(collector) From 2ec6594db9c2397cc3c315ff3ce3b38e0d40e176 Mon Sep 17 00:00:00 2001 From: "Xiaohong (Sean) Chen" Date: Thu, 11 Jun 2026 11:59:08 -0400 Subject: [PATCH 07/52] [Kernel][Helion][1/N] Add Helion kernel for per_token_group_fp8_quant (#36902) Signed-off-by: Sean Chen Co-authored-by: Yanan Cao Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .buildkite/test-amd.yaml | 2 +- .buildkite/test_areas/kernels.yaml | 2 +- setup.py | 2 +- .../helion/test_per_token_group_fp8_quant.py | 243 +++ tests/kernels/helion/test_register.py | 3 + tests/kernels/helion/utils.py | 30 + .../nvidia_b200.json | 1938 +++++++++++++++++ .../nvidia_h100.json | 1893 ++++++++++++++++ .../helion/ops/per_token_group_fp8_quant.py | 232 ++ vllm/kernels/helion/register.py | 6 +- 10 files changed, 4347 insertions(+), 4 deletions(-) create mode 100644 tests/kernels/helion/test_per_token_group_fp8_quant.py create mode 100644 tests/kernels/helion/utils.py create mode 100644 vllm/kernels/helion/configs/per_token_group_fp8_quant/nvidia_b200.json create mode 100644 vllm/kernels/helion/configs/per_token_group_fp8_quant/nvidia_h100.json create mode 100644 vllm/kernels/helion/ops/per_token_group_fp8_quant.py diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index 186f7222539..148aea73c7f 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -398,7 +398,7 @@ steps: - tests/kernels/helion/ - vllm/platforms/rocm.py commands: - - pip install helion==1.0.0 + - pip install helion==1.1.0 - pytest -v -s kernels/helion/ - label: Kernels Mamba Test # TBD diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index 9ec86845038..159f940530e 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -237,7 +237,7 @@ steps: - vllm/utils/import_utils.py - tests/kernels/helion/ commands: - - pip install helion==1.0.0 + - pip install helion==1.1.0 - pytest -v -s kernels/helion/ diff --git a/setup.py b/setup.py index 0a820587958..657a65161e7 100644 --- a/setup.py +++ b/setup.py @@ -1229,7 +1229,7 @@ setup( # NOTE: When updating helion version, also update CI files: # - .buildkite/test_areas/kernels.yaml # - .buildkite/test-amd.yaml - "helion": ["helion==1.0.0"], + "helion": ["helion==1.1.0"], # Optional deps for gRPC server (vllm serve --grpc) "grpc": ["smg-grpc-servicer[vllm] >= 0.5.2"], # Optional deps for OpenTelemetry tracing diff --git a/tests/kernels/helion/test_per_token_group_fp8_quant.py b/tests/kernels/helion/test_per_token_group_fp8_quant.py new file mode 100644 index 00000000000..304734c77e5 --- /dev/null +++ b/tests/kernels/helion/test_per_token_group_fp8_quant.py @@ -0,0 +1,243 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for the per_token_group_fp8_quant helion kernel + +Run `pytest tests/kernels/helion/test_per_token_group_fp8_quant.py`. +""" + +from typing import Any + +import pytest +import torch +from torch._subclasses.fake_tensor import FakeTensorMode + +from tests.kernels.helion.utils import skip_if_platform_unsupported +from tests.kernels.quant_utils import FP8_DTYPE +from vllm.kernels.helion.case_key import CaseKey +from vllm.kernels.helion.config_manager import ConfigManager +from vllm.kernels.helion.ops.per_token_group_fp8_quant import ( + _pick_cache, + baseline, + per_token_group_fp8_quant, + pick_config, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + get_fp8_min_max, +) +from vllm.utils.import_utils import has_helion + +if not has_helion(): + pytest.skip( + "Helion is not installed. Install with: pip install vllm[helion]", + allow_module_level=True, + ) + + +def _generate_fake_input( + num_tokens: int, hidden_size: int, group_size: int +) -> tuple[Any, ...]: + with FakeTensorMode(): + input = torch.randn( + (num_tokens, hidden_size), device="cuda", dtype=torch.bfloat16 + ) + output_q = torch.empty(input.shape, device=input.device, dtype=FP8_DTYPE) + output_s = torch.empty( + (num_tokens, hidden_size // group_size), + device=input.device, + dtype=torch.float32, + ) + use_ue8m0 = False + column_major = False + fp8_min, fp8_max = get_fp8_min_max() + eps = 1e-10 + args = ( + input, + output_q, + output_s, + group_size, + eps, + fp8_min, + fp8_max, + use_ue8m0, + column_major, + ) + return args + + +@pytest.fixture(autouse=True) +def reset_config_manager_singleton(): + ConfigManager.reset_instance() + ConfigManager() + yield + ConfigManager.reset_instance() + + +class TestPerTokenGroupFp8QuantConfigPicker: + def setup_method(self): + _pick_cache.clear() + + def test_config_picker_exact_match(self): + config_keys = [ + CaseKey({"hidden_size": 2048, "group_size": 64, "num_tokens": 16}), + CaseKey({"hidden_size": 4096, "group_size": 128, "num_tokens": 16}), + ] + + args = _generate_fake_input(16, 4096, 128) + selected_key = pick_config(args, config_keys) + assert selected_key == CaseKey( + {"hidden_size": 4096, "group_size": 128, "num_tokens": 16} + ) + + def test_config_picker_closest_match(self): + config_keys = [ + CaseKey({"hidden_size": 2048, "group_size": 64, "num_tokens": 16}), + CaseKey({"hidden_size": 2048, "group_size": 64, "num_tokens": 32}), + CaseKey({"hidden_size": 2048, "group_size": 128, "num_tokens": 16}), + CaseKey({"hidden_size": 2048, "group_size": 128, "num_tokens": 32}), + CaseKey({"hidden_size": 4096, "group_size": 64, "num_tokens": 16}), + CaseKey({"hidden_size": 4096, "group_size": 64, "num_tokens": 32}), + CaseKey({"hidden_size": 4096, "group_size": 128, "num_tokens": 16}), + CaseKey({"hidden_size": 4096, "group_size": 128, "num_tokens": 32}), + ] + + args = _generate_fake_input(20, 3000, 70) + selected_key = pick_config(args, config_keys) + assert selected_key == CaseKey( + {"hidden_size": 2048, "group_size": 64, "num_tokens": 32} + ) + + def test_config_picker_no_configs(self): + config_keys: list[dict] = [] + + args = _generate_fake_input(16, 4096, 128) + selected_key = pick_config(args, config_keys) + assert selected_key is None + + def test_config_picker_fallback_to_largest(self): + config_keys = [ + CaseKey({"hidden_size": 2048, "group_size": 64, "num_tokens": 16}), + CaseKey({"hidden_size": 2048, "group_size": 64, "num_tokens": 32}), + CaseKey({"hidden_size": 2048, "group_size": 128, "num_tokens": 16}), + CaseKey({"hidden_size": 2048, "group_size": 128, "num_tokens": 32}), + CaseKey({"hidden_size": 4096, "group_size": 64, "num_tokens": 16}), + CaseKey({"hidden_size": 4096, "group_size": 64, "num_tokens": 32}), + CaseKey({"hidden_size": 4096, "group_size": 128, "num_tokens": 16}), + CaseKey({"hidden_size": 4096, "group_size": 128, "num_tokens": 32}), + ] + + args = _generate_fake_input(64, 8192, 256) + selected_key = pick_config(args, config_keys) + assert selected_key == CaseKey( + {"hidden_size": 4096, "group_size": 128, "num_tokens": 32} + ) + + +class TestPerTokenGroupFp8QuantCorrectness: + @pytest.mark.parametrize( + "shape", [(31, 128), (32, 128), (63, 256), (64, 256), (16, 512), (2048, 5120)] + ) + @pytest.mark.parametrize("column_major", [False, True]) + @pytest.mark.parametrize("tma_aligned", [False, True]) + @pytest.mark.parametrize("scale_ue8m0", [False, True]) + @pytest.mark.parametrize("group_size", [64, 128]) + def test_per_token_group_fp8_quant( + self, + shape, + column_major: bool, + tma_aligned: bool, + scale_ue8m0: bool, + group_size: int, + ): + skip_if_platform_unsupported("per_token_group_fp8_quant") + + torch.manual_seed(42) + num_tokens, hidden_size = shape + fp8_min, fp8_max = get_fp8_min_max() + eps = 1e-10 + input = ( + torch.randn((num_tokens, hidden_size), device="cuda", dtype=torch.bfloat16) + * 8 + ) + ref_q = torch.empty(input.shape, device=input.device, dtype=FP8_DTYPE) + ops_q = ref_q.clone() + + groups_per_row = hidden_size // group_size + if column_major: + if tma_aligned: + tma_alignment = 4 + tma_aligned_m = ( + (num_tokens + tma_alignment - 1) // tma_alignment * tma_alignment + ) + shape = (num_tokens, groups_per_row) + stride = (1, tma_aligned_m) + ref_s = torch.empty_strided( + shape, stride, device=input.device, dtype=torch.float32 + ) + else: + ref_s = torch.empty( + (groups_per_row, num_tokens), + device=input.device, + dtype=torch.float32, + ).transpose(0, 1) + else: + ref_s = torch.empty( + (num_tokens, groups_per_row), device=input.device, dtype=torch.float32 + ) + + ops_s = ref_s.clone() + + baseline( + input, + ref_q, + ref_s, + group_size, + eps, + fp8_min, + fp8_max, + scale_ue8m0, + column_major, + tma_aligned, + ) + per_token_group_fp8_quant( + input, + ops_q, + ops_s, + group_size, + eps, + fp8_min, + fp8_max, + scale_ue8m0, + column_major, + tma_aligned, + ) + + assert torch.allclose(ref_s, ops_s) + # allow 1 ULP difference + assert ( + ref_q.view(torch.uint8).to(torch.int16) + - ops_q.view(torch.uint8).to(torch.int16) + ).abs().max() <= 1 + + +class TestPerTokenGroupFp8QuantIntegration: + def test_kernel_registration_integration(self): + from vllm.kernels.helion.register import get_registered_kernels + + registered_kernels = get_registered_kernels() + assert "per_token_group_fp8_quant" in registered_kernels + + kernel_wrapper = registered_kernels["per_token_group_fp8_quant"] + assert kernel_wrapper.op_name == "per_token_group_fp8_quant" + assert kernel_wrapper._config_picker is not None + assert kernel_wrapper._mutates_args == ["output_q", "output_s"] + + def test_fake_impl_functionality(self): + skip_if_platform_unsupported("per_token_group_fp8_quant") + from vllm.kernels.helion.register import get_registered_kernels + + registered_kernels = get_registered_kernels() + kernel_wrapper = registered_kernels["per_token_group_fp8_quant"] + fake_impl = kernel_wrapper._fake_impl + + args = _generate_fake_input(16, 4096, 128) + assert fake_impl(*args) is None diff --git a/tests/kernels/helion/test_register.py b/tests/kernels/helion/test_register.py index c82c3c8358e..9876135056b 100644 --- a/tests/kernels/helion/test_register.py +++ b/tests/kernels/helion/test_register.py @@ -713,6 +713,7 @@ class TestHelionKernelWrapper: new_op = Mock() registered_ops: dict[str, Mock] = {} + mutates_args = ["y"] class MockNamespace: def __getattr__(self, name): @@ -748,6 +749,7 @@ class TestHelionKernelWrapper: raw_kernel_func=sample_kernel, op_name="test_kernel", fake_impl=fake_impl, + mutates_args=mutates_args, config_picker=default_picker, ) result = wrapper._get_or_register_custom_op() @@ -755,6 +757,7 @@ class TestHelionKernelWrapper: mock_register.assert_called_once() assert result is new_op assert mock_register.call_args[1]["op_func"] is mock_decorated + assert mock_register.call_args[1]["mutates_args"] is mutates_args class TestKernelRegistry: diff --git a/tests/kernels/helion/utils.py b/tests/kernels/helion/utils.py new file mode 100644 index 00000000000..38893fc8fec --- /dev/null +++ b/tests/kernels/helion/utils.py @@ -0,0 +1,30 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Helion Kernel test utils""" + +import pytest +import torch + +from vllm.kernels.helion.config_manager import ConfigManager + + +def skip_if_platform_unsupported(op_name: str): + try: + from vllm.kernels.helion.utils import get_canonical_gpu_name + + if not torch.cuda.is_available(): + pytest.skip("CUDA not available") + + platform = get_canonical_gpu_name() + + try: + config_manager = ConfigManager.get_instance() + except RuntimeError: + config_manager = ConfigManager() + + configs = config_manager.get_platform_configs(op_name, platform) + if len(configs) == 0: + pytest.skip(f"Current GPU platform not supported for {op_name} kernel") + + except (ImportError, RuntimeError, KeyError): + pytest.skip(f"Error detecting platform support for {op_name} kernel") diff --git a/vllm/kernels/helion/configs/per_token_group_fp8_quant/nvidia_b200.json b/vllm/kernels/helion/configs/per_token_group_fp8_quant/nvidia_b200.json new file mode 100644 index 00000000000..23f68e88c6e --- /dev/null +++ b/vllm/kernels/helion/configs/per_token_group_fp8_quant/nvidia_b200.json @@ -0,0 +1,1938 @@ +[ + { + "key": { + "hidden_size": 2048, + "group_size": 128, + "num_tokens": 1 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 4 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first" + ], + "num_warps": 2, + "num_stages": 6, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 128, + "num_tokens": 2 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 4 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 4, + "num_stages": 5, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 128, + "num_tokens": 4 + }, + "config": { + "block_sizes": [ + 16 + ], + "loop_orders": [ + [ + 0, + 1, + 2 + ] + ], + "l2_groupings": [ + 1 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 4, + "num_stages": 1, + "indexing": [ + "pointer", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 128, + "num_tokens": 8 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 4 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 4, + "num_stages": 5, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 128, + "num_tokens": 16 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 4 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first" + ], + "num_warps": 2, + "num_stages": 6, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 128, + "num_tokens": 32 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 4 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 4, + "num_stages": 5, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 128, + "num_tokens": 64 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 4 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 4, + "num_stages": 5, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 128, + "num_tokens": 128 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 4 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 4, + "num_stages": 5, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 128, + "num_tokens": 256 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 1 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first" + ], + "num_warps": 4, + "num_stages": 1, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 128, + "num_tokens": 512 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 2 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first" + ], + "num_warps": 4, + "num_stages": 7, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 128, + "num_tokens": 1024 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 4 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 2, + "num_stages": 8, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 128, + "num_tokens": 2048 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 4 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 2, + "num_stages": 8, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 128, + "num_tokens": 4096 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 64 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 2, + "num_stages": 5, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 128, + "num_tokens": 8192 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 4 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 2, + "num_stages": 8, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 128, + "num_tokens": 1 + }, + "config": { + "block_sizes": [ + 16 + ], + "loop_orders": [ + [ + 0, + 1, + 2 + ] + ], + "l2_groupings": [ + 1 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 4, + "num_stages": 1, + "indexing": [ + "pointer", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 128, + "num_tokens": 2 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 4 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 4, + "num_stages": 5, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 128, + "num_tokens": 4 + }, + "config": { + "block_sizes": [ + 1 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 8 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last" + ], + "num_warps": 2, + "num_stages": 3, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 128, + "num_tokens": 8 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 4 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 4, + "num_stages": 5, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 128, + "num_tokens": 16 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 4 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 4, + "num_stages": 5, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 128, + "num_tokens": 32 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 4 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 4, + "num_stages": 5, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 128, + "num_tokens": 64 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 4 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 4, + "num_stages": 5, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 128, + "num_tokens": 128 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 2 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first" + ], + "num_warps": 4, + "num_stages": 7, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 128, + "num_tokens": 256 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 4 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 2, + "num_stages": 8, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 128, + "num_tokens": 512 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 16 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 2, + "num_stages": 2, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 128, + "num_tokens": 1024 + }, + "config": { + "block_sizes": [ + 16 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 32 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first" + ], + "num_warps": 2, + "num_stages": 2, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 128, + "num_tokens": 2048 + }, + "config": { + "block_sizes": [ + 16 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 16 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 4, + "num_stages": 6, + "indexing": [ + "pointer", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 128, + "num_tokens": 4096 + }, + "config": { + "block_sizes": [ + 16 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 1 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 2, + "num_stages": 1, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 128, + "num_tokens": 8192 + }, + "config": { + "block_sizes": [ + 16 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 1 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 2, + "num_stages": 1, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "group_size": 128, + "num_tokens": 1 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 0, + 1, + 2 + ] + ], + "l2_groupings": [ + 1 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 8, + "num_stages": 1, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "group_size": 128, + "num_tokens": 2 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 4 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 4, + "num_stages": 5, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "group_size": 128, + "num_tokens": 4 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 4 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 4, + "num_stages": 5, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "group_size": 128, + "num_tokens": 8 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 4 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first" + ], + "num_warps": 2, + "num_stages": 6, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "group_size": 128, + "num_tokens": 16 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 1 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 8, + "num_stages": 1, + "indexing": [ + "pointer", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "group_size": 128, + "num_tokens": 32 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 4 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 4, + "num_stages": 5, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "group_size": 128, + "num_tokens": 64 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 4 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 4, + "num_stages": 5, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "group_size": 128, + "num_tokens": 128 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 1 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first" + ], + "num_warps": 4, + "num_stages": 1, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "group_size": 128, + "num_tokens": 256 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 2 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 4, + "num_stages": 7, + "indexing": [ + "pointer", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "group_size": 128, + "num_tokens": 512 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 32 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 1, + "num_stages": 1, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "group_size": 128, + "num_tokens": 1024 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 16 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 2, + "num_stages": 2, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "group_size": 128, + "num_tokens": 2048 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 4 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 2, + "num_stages": 4, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "group_size": 128, + "num_tokens": 4096 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 2 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 1, + "num_stages": 1, + "indexing": [ + "pointer", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "group_size": 128, + "num_tokens": 8192 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 2 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 1, + "num_stages": 1, + "indexing": [ + "pointer", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + } +] \ No newline at end of file diff --git a/vllm/kernels/helion/configs/per_token_group_fp8_quant/nvidia_h100.json b/vllm/kernels/helion/configs/per_token_group_fp8_quant/nvidia_h100.json new file mode 100644 index 00000000000..08a0d97ccf2 --- /dev/null +++ b/vllm/kernels/helion/configs/per_token_group_fp8_quant/nvidia_h100.json @@ -0,0 +1,1893 @@ +[ + { + "key": { + "hidden_size": 5120, + "group_size": 128, + "num_tokens": 16 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 64 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 2, + "num_stages": 8, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 128, + "num_tokens": 1 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 4 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 4, + "num_stages": 5, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 128, + "num_tokens": 2 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 2 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 4, + "num_stages": 5, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 128, + "num_tokens": 4 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 2 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 4, + "num_stages": 5, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 128, + "num_tokens": 8 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 2 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 4, + "num_stages": 5, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 128, + "num_tokens": 16 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 2 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 4, + "num_stages": 5, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 128, + "num_tokens": 32 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 2 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 4, + "num_stages": 5, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 128, + "num_tokens": 64 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 4 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 4, + "num_stages": 1, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 128, + "num_tokens": 128 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 1 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first" + ], + "num_warps": 4, + "num_stages": 1, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 128, + "num_tokens": 256 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 4 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first" + ], + "num_warps": 2, + "num_stages": 6, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 128, + "num_tokens": 512 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 4 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 2, + "num_stages": 3, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 128, + "num_tokens": 1024 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 4 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first" + ], + "num_warps": 2, + "num_stages": 6, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 128, + "num_tokens": 2048 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 8 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last" + ], + "num_warps": 2, + "num_stages": 4, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 128, + "num_tokens": 4096 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 2 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first" + ], + "num_warps": 2, + "num_stages": 4, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 128, + "num_tokens": 8192 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 64 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first" + ], + "num_warps": 2, + "num_stages": 8, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 128, + "num_tokens": 1 + }, + "config": { + "block_sizes": [ + 1 + ], + "loop_orders": [ + [ + 1, + 0, + 2 + ] + ], + "l2_groupings": [ + 8 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first" + ], + "num_warps": 1, + "num_stages": 4, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 128, + "num_tokens": 2 + }, + "config": { + "block_sizes": [ + 2 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 8 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last" + ], + "num_warps": 2, + "num_stages": 8, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 128, + "num_tokens": 4 + }, + "config": { + "block_sizes": [ + 1 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 8 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first" + ], + "num_warps": 1, + "num_stages": 8, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 128, + "num_tokens": 8 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 2 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last" + ], + "num_warps": 4, + "num_stages": 3, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 128, + "num_tokens": 16 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 2 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 4, + "num_stages": 6, + "indexing": [ + "pointer", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 128, + "num_tokens": 32 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 32 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last" + ], + "num_warps": 4, + "num_stages": 3, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 128, + "num_tokens": 64 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 2 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 4, + "num_stages": 6, + "indexing": [ + "pointer", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 128, + "num_tokens": 128 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 8 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first" + ], + "num_warps": 2, + "num_stages": 3, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 128, + "num_tokens": 256 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 64 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 2, + "num_stages": 5, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 128, + "num_tokens": 512 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 16 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last" + ], + "num_warps": 2, + "num_stages": 1, + "indexing": [ + "pointer", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 128, + "num_tokens": 1024 + }, + "config": { + "block_sizes": [ + 16 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 16 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 4, + "num_stages": 6, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 128, + "num_tokens": 2048 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 8 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first" + ], + "num_warps": 2, + "num_stages": 2, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 128, + "num_tokens": 4096 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 8 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first" + ], + "num_warps": 2, + "num_stages": 3, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 128, + "num_tokens": 8192 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 4 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first" + ], + "num_warps": 2, + "num_stages": 6, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "group_size": 128, + "num_tokens": 1 + }, + "config": { + "block_sizes": [ + 1 + ], + "loop_orders": [ + [ + 1, + 0, + 2 + ] + ], + "l2_groupings": [ + 8 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first" + ], + "num_warps": 1, + "num_stages": 4, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "group_size": 128, + "num_tokens": 2 + }, + "config": { + "block_sizes": [ + 2 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 64 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first" + ], + "num_warps": 1, + "num_stages": 6, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "group_size": 128, + "num_tokens": 4 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 8 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 4, + "num_stages": 4, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "group_size": 128, + "num_tokens": 8 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 32 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last" + ], + "num_warps": 4, + "num_stages": 3, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "group_size": 128, + "num_tokens": 32 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 32 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 4, + "num_stages": 1, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "group_size": 128, + "num_tokens": 64 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 2 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 4, + "num_stages": 5, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "group_size": 128, + "num_tokens": 128 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 16 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 4, + "num_stages": 7, + "indexing": [ + "pointer", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "group_size": 128, + "num_tokens": 256 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 8 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last" + ], + "num_warps": 2, + "num_stages": 4, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "group_size": 128, + "num_tokens": 512 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 8 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first" + ], + "num_warps": 2, + "num_stages": 3, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "group_size": 128, + "num_tokens": 1024 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 32 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 2, + "num_stages": 6, + "indexing": [ + "pointer", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "group_size": 128, + "num_tokens": 2048 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 32 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "" + ], + "num_warps": 2, + "num_stages": 6, + "indexing": [ + "pointer", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "group_size": 128, + "num_tokens": 4096 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 32 + ], + "range_unroll_factors": [ + 2 + ], + "range_warp_specializes": [], + "range_multi_buffers": [ + true + ], + "range_flattens": [ + true + ], + "load_eviction_policies": [ + "first" + ], + "num_warps": 2, + "num_stages": 6, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "persistent_interleaved", + "num_sm_multiplier": 128, + "maxnreg": 256 + } + }, + { + "key": { + "hidden_size": 5120, + "group_size": 128, + "num_tokens": 8192 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 64 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first" + ], + "num_warps": 2, + "num_stages": 7, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + } +] \ No newline at end of file diff --git a/vllm/kernels/helion/ops/per_token_group_fp8_quant.py b/vllm/kernels/helion/ops/per_token_group_fp8_quant.py new file mode 100644 index 00000000000..8b73fac4b8e --- /dev/null +++ b/vllm/kernels/helion/ops/per_token_group_fp8_quant.py @@ -0,0 +1,232 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from itertools import product +from typing import Any + +import torch + +from vllm.kernels.helion.case_key import CaseKey +from vllm.logger import init_logger +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + get_fp8_min_max, +) +from vllm.platforms import current_platform +from vllm.utils.import_utils import has_helion + +if not has_helion(): + raise ImportError( + "Helion kernel requires helion to be installed. " + "Install it with: pip install helion" + ) + +import helion +import helion.language as hl + +from vllm.kernels.helion.register import register_kernel + +logger = init_logger(__name__) + + +def generate_inputs() -> dict[CaseKey, tuple[Any, ...]]: + # TODO(xiaohongchen1991): it is difficult for kernel author to cover all + # input property combination. Currently, dtypes are fixed. We need + # optimization to bucket/skip some combinations + num_tokens_list = [1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192] + hidden_size_list = [2048, 4096, 5120] + group_size_list = [128] + in_dtype: torch.dtype = torch.bfloat16 + out_dtype: torch.dtype = current_platform.fp8_dtype() + scale_dtype: torch.dtype = torch.float32 + + use_ue8m0 = False + column_major = False + fp8_min, fp8_max = get_fp8_min_max() + eps = 1e-10 + + inputs = {} + + for hidden_size, group_size, num_tokens in product( + hidden_size_list, group_size_list, num_tokens_list + ): + input = torch.randn(num_tokens, hidden_size, device="cuda", dtype=in_dtype) + output_q = torch.empty(input.shape, device=input.device, dtype=out_dtype) + output_s = torch.empty( + (num_tokens, hidden_size // group_size), + device=input.device, + dtype=scale_dtype, + ) + config_key = CaseKey( + { + "hidden_size": hidden_size, + "group_size": group_size, + "num_tokens": num_tokens, + } + ) + inputs[config_key] = ( + input, + output_q, + output_s, + group_size, + eps, + fp8_min, + fp8_max, + use_ue8m0, + column_major, + False, + ) + + return inputs + + +_pick_cache: dict[tuple[int, int, int], CaseKey | None] = {} + + +def pick_config(args: tuple[Any, ...], config_keys: list[CaseKey]) -> CaseKey | None: + """Pick the best pre-tuned config for the given input shape. + + Selection strategy: + 1. Find the closest hidden_size among available configs + (exact match preferred). + 2. Find the closest group_size among available configs + (exact match preferred). + 3. Among the num_tokens values tuned for that hidden_size and group_size, pick + the smallest num_tokens >= the input's num_tokens. If the input is + larger than all available num_tokens, fall back to the largest. + """ + + if not config_keys: + return None + + input, _, _, group_size, *_ = args + num_tokens, hidden_size = input.shape + + cache_key = (num_tokens, group_size, hidden_size) + cached = _pick_cache.get(cache_key) + if cached is not None: + return cached + + configs: dict[int, dict[int, list[int]]] = {} + for key in config_keys: + if key.is_default(): + continue + configs.setdefault(key["hidden_size"], {}).setdefault( + key["group_size"], [] + ).append(key["num_tokens"]) + + if not configs: + return None + + best_hidden_size = min(configs, key=lambda s: abs(s - hidden_size)) + best_group_size = min(configs[best_hidden_size], key=lambda s: abs(s - group_size)) + available_num_tokens = sorted(configs[best_hidden_size][best_group_size]) + best_num_tokens = next( + (n for n in available_num_tokens if n >= num_tokens), available_num_tokens[-1] + ) + + result = CaseKey( + { + "hidden_size": best_hidden_size, + "group_size": best_group_size, + "num_tokens": best_num_tokens, + } + ) + _pick_cache[cache_key] = result + return result + + +def fake_impl( + input: torch.Tensor, # [num_tokens, hidden_size] + output_q: torch.Tensor, # [num_tokens, hidden_size] + output_s: torch.Tensor, # [num_tokens, groups_per_row] + group_size: int, + eps: float, + fp8_min: float, + fp8_max: float, + scale_ue8m0: bool, + # Unused dummy args + # Kept for consistency with existing kernel interface + dummy_is_scale_transposed: bool = False, + dummy_is_tma_aligned: bool = False, +) -> None: + return + + +def baseline( + input: torch.Tensor, # [num_tokens, hidden_size] + output_q: torch.Tensor, # [num_tokens, hidden_size] + output_s: torch.Tensor, # [num_tokens, groups_per_row] + group_size: int, + eps: float, + fp8_min: float, + fp8_max: float, + scale_ue8m0: bool, + dummy_is_scale_transposed: bool = False, + dummy_is_tma_aligned: bool = False, +) -> None: + torch.ops._C.per_token_group_fp8_quant( + input, + output_q, + output_s, + group_size, + eps, + fp8_min, + fp8_max, + scale_ue8m0, + dummy_is_scale_transposed, + dummy_is_tma_aligned, + ) + + +@register_kernel( + mutates_args=["output_q", "output_s"], + config_picker=pick_config, + input_generator=generate_inputs, + fake_impl=fake_impl, + helion_settings=helion.Settings( + autotune_baseline_fn=baseline, + ), +) # type: ignore[misc] +def per_token_group_fp8_quant( + input: torch.Tensor, # [num_tokens, hidden_size] + output_q: torch.Tensor, # [num_tokens, hidden_size] + output_s: torch.Tensor, # [num_tokens, groups_per_row] + group_size: int, + eps: float, + fp8_min: float, + fp8_max: float, + scale_ue8m0: bool, + # Unused dummy args + # Kept for consistency with existing kernel interface + dummy_is_scale_transposed: bool = False, + dummy_is_tma_aligned: bool = False, +) -> None: + # This code assumes batch_dim and num_tokens are flattened + assert input.ndim == 2 + num_tokens, hidden_size = input.shape + hl.specialize(hidden_size) + hl.specialize(group_size) + + groups_per_row = output_s.shape[1] + hl.specialize(groups_per_row) + assert hidden_size % group_size == 0 and hidden_size // group_size == groups_per_row + assert output_s.ndim == 2 and output_s.dtype == torch.float32 + + input = input.view(num_tokens, -1, group_size) + output_q = output_q.view(num_tokens, -1, group_size) + for tile_m, tile_gn, tile_n in hl.tile( + [num_tokens, groups_per_row, group_size], block_size=[1, None, group_size] + ): + x_blk = input[tile_m, tile_gn, tile_n] + y_s_blk = torch.clamp(torch.amax(torch.abs(x_blk), dim=-1), min=eps) + y_s_blk = y_s_blk / fp8_max + + if scale_ue8m0: + y_s_blk = torch.exp2(torch.ceil(torch.log2(y_s_blk))) + + y_q_blk = torch.clamp(x_blk / y_s_blk[:, :, None], fp8_min, fp8_max).to( + output_q.dtype + ) + + output_s[tile_m, tile_gn] = y_s_blk + output_q[tile_m, tile_gn, tile_n] = y_q_blk diff --git a/vllm/kernels/helion/register.py b/vllm/kernels/helion/register.py index f18120da45f..764022de77d 100644 --- a/vllm/kernels/helion/register.py +++ b/vllm/kernels/helion/register.py @@ -260,6 +260,7 @@ class HelionKernelWrapper: op_name: str, fake_impl: Callable, config_picker: ConfigPicker, + mutates_args: list[str] | None = None, helion_settings: helion.Settings | None = None, input_generator: (Callable[[], dict[CaseKey, tuple[Any, ...]]] | None) = None, ): @@ -272,6 +273,7 @@ class HelionKernelWrapper: self.helion_settings = helion_settings self._config_picker = config_picker self._input_generator = input_generator + self._mutates_args = mutates_args self._configured_kernel: ConfiguredHelionKernel | None = None # TODO(@gmagogsfm): Remove this disable flag once integrated with vLLM IR, # which handles op enablement/disablement. @@ -357,7 +359,7 @@ class HelionKernelWrapper: direct_register_custom_op( op_name=self.op_name, op_func=configured_kernel._decorated_kernel, - mutates_args=None, + mutates_args=self._mutates_args, fake_impl=self._fake_impl, target_lib=vllm_helion_lib, ) @@ -402,6 +404,7 @@ def register_kernel( *, config_picker: ConfigPicker, fake_impl: Callable | None = None, + mutates_args: list[str] | None = None, helion_settings: helion.Settings | None = None, input_generator: (Callable[[], dict[CaseKey, tuple[Any, ...]]] | None) = None, ) -> Callable[[Callable], HelionKernelWrapper]: @@ -455,6 +458,7 @@ def register_kernel( op_name=final_op_name, fake_impl=final_fake_impl, config_picker=config_picker, + mutates_args=mutates_args, helion_settings=helion_settings, input_generator=input_generator, ) From b8142294b7e757f3a39729c4f400bafaed534681 Mon Sep 17 00:00:00 2001 From: wentian-byte <3400259131@qq.com> Date: Fri, 12 Jun 2026 00:39:24 +0800 Subject: [PATCH 08/52] [Bugfix] Restrict FlashInfer cuDNN FP8 ViT attention gate to Blackwell (SM 100) (#45251) Signed-off-by: Wentian Byte <3400259131@qq.com> --- .../layers/attention/mm_encoder_attention.py | 5 +++-- vllm/utils/flashinfer.py | 15 +++++++++++---- 2 files changed, 14 insertions(+), 6 deletions(-) diff --git a/vllm/model_executor/layers/attention/mm_encoder_attention.py b/vllm/model_executor/layers/attention/mm_encoder_attention.py index 1731cc26bc3..2ca051ad9e4 100644 --- a/vllm/model_executor/layers/attention/mm_encoder_attention.py +++ b/vllm/model_executor/layers/attention/mm_encoder_attention.py @@ -396,8 +396,9 @@ class MMEncoderAttention(CustomOp): if not is_flashinfer_cudnn_fp8_prefill_attn_supported(): raise ValueError( "mm_encoder_attn_dtype='fp8' requires the FlashInfer " - "cuDNN backend with cuDNN >= 9.17.1 on a GPU with native " - "FP8 support." + "cuDNN backend with cuDNN >= 9.17.1 on Blackwell (SM 100) " + "or newer. cuDNN's FP8 SDPA path with bf16/fp16 output is " + "not available on Hopper (H100/H200) or earlier." ) self.fp8_enabled = True diff --git a/vllm/utils/flashinfer.py b/vllm/utils/flashinfer.py index 95f8b4b7ec0..e0518277865 100644 --- a/vllm/utils/flashinfer.py +++ b/vllm/utils/flashinfer.py @@ -934,20 +934,27 @@ def should_use_flashinfer_for_blockscale_fp8_gemm( return should_use_flashinfer -_MIN_CUDNN_FP8 = 91701 # cuDNN >= 9.17.1 required for FP8 attention +_MIN_CUDNN_FP8 = 91701 # cuDNN >= 9.17.1 required for FP8 ViT attention @functools.cache def is_flashinfer_cudnn_fp8_prefill_attn_supported() -> bool: """Check if FP8 ViT attention is supported on this platform. - Requires native FP8 hardware support, the FlashInfer cuDNN backend, + Requires Blackwell (SM 100) or newer, the FlashInfer cuDNN backend, and cuDNN >= 9.17.1. + + cuDNN's FP8 SDPA forward path with bf16/fp16 output (used by + ``MMEncoderAttention._forward_flashinfer``) gates internally on + ``prop.major >= 10``; on Hopper it raises a misleading + ``cudnnGraphNotSupportedError: ... cuDNN version 9.13.0 and newer`` + even when the installed cuDNN is new enough. See PR #38065 for the + original Blackwell-only design intent. """ from vllm.v1.attention.backends.registry import AttentionBackendEnum - # cuDNN SDPA FP8 requires Hopper (SM 90) or newer. - if not current_platform.has_device_capability(90): + # cuDNN SDPA FP8 with bf16/fp16 output requires Blackwell (SM 100) or newer. + if not current_platform.has_device_capability(100): return False try: From 3b03a2cf4772838da622d81315941bb41bcc03ff Mon Sep 17 00:00:00 2001 From: Chao-Ju Chen Date: Fri, 12 Jun 2026 01:50:59 +0800 Subject: [PATCH 09/52] [Rust Frontend] Support continuous_usage_stats stream option (#43965) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: Bugen Zhao Signed-off-by: RickyChen / 陳昭儒 Signed-off-by: Bugen Zhao --- .../src/routes/openai/chat_completions.rs | 71 +++++++---- .../routes/openai/chat_completions/convert.rs | 53 +++++++- .../openai/chat_completions/validate.rs | 9 -- .../server/src/routes/openai/completions.rs | 35 +++++- .../src/routes/openai/completions/convert.rs | 58 +++++++++ .../src/routes/openai/completions/validate.rs | 9 -- .../src/server/src/routes/openai/utils/mod.rs | 1 + .../server/src/routes/openai/utils/usage.rs | 35 ++++++ rust/src/server/src/routes/tests.rs | 115 ++++++++++++++++++ 9 files changed, 335 insertions(+), 51 deletions(-) create mode 100644 rust/src/server/src/routes/openai/utils/usage.rs diff --git a/rust/src/server/src/routes/openai/chat_completions.rs b/rust/src/server/src/routes/openai/chat_completions.rs index e93c049b2d1..6274a4e98ac 100644 --- a/rust/src/server/src/routes/openai/chat_completions.rs +++ b/rust/src/server/src/routes/openai/chat_completions.rs @@ -37,6 +37,7 @@ use crate::routes::openai::utils::logprobs::{ use crate::routes::openai::utils::types::{ ChatLogProbs, FunctionCallDelta, FunctionCallResponse, ToolCall, ToolCallDelta, Usage, }; +use crate::routes::openai::utils::usage::ContinuousUsage; use crate::routes::openai::utils::validated_json::ValidatedJson; use crate::state::AppState; use crate::utils::{resolve_request_context, unix_timestamp}; @@ -129,6 +130,8 @@ async fn collect_chat_completion( ResponseOptions { // Ignored: non-streaming responses always include usage. include_usage: _, + // Ignored: non-streaming responses are collected before usage is attached. + include_continuous_usage: _, requested_logprobs, include_prompt_logprobs, include_reasoning, @@ -249,6 +252,7 @@ async fn chat_completion_chunk_stream( }: ApiServerOptions, ResponseOptions { include_usage, + include_continuous_usage, requested_logprobs, // Ignored: chat streaming prompt logprobs are rejected for Python parity. include_prompt_logprobs: _, @@ -265,33 +269,47 @@ async fn chat_completion_chunk_stream( // starts or ends, omit its token metadata as well as its visible delta. let mut inside_hidden_reasoning = false; let mut suppress_current_update_metadata = false; + let mut continuous_usage = ContinuousUsage::default(); + + /// Yield a chunk with optional continuous usage attached. + macro_rules! yield_chunk { + ($chunk:expr) => {{ + let mut chunk = $chunk; + if include_continuous_usage { + chunk.usage = Some(continuous_usage.to_usage()); + } + y.yield_ok(chunk).await; + }}; + } // If the client requested logprobs or token_ids, we need to buffer chunks until // we receive the separate `LogprobsDelta` event, so that we can emit one // combined chunk with both the semantic delta and its per-update metadata. - let mut pending_chunk = - (requested_logprobs || return_token_ids).then(PendingChatChunk::default); + // Continuous usage also buffers so the token count from `LogprobsDelta` can + // be attached to the matching semantic chunk. + let mut pending_chunk = (requested_logprobs || return_token_ids || include_continuous_usage) + .then(PendingChatChunk::default); while let Some(next) = stream.next().await { match next { Ok(ChatEvent::Start { prompt_token_ids, .. }) => { + continuous_usage.set_prompt_tokens(prompt_token_ids.len()); let mut chunk = start_chunk(&request_id, &response_model, created); if return_token_ids { chunk.prompt_token_ids = Some(prompt_token_ids.to_vec()); } - y.yield_ok(chunk).await; + yield_chunk!(chunk); // When echo=true, emit the last assistant message content as a delta chunk. if let Some(echo_text) = &echo { - y.yield_ok(block_delta_chunk( + yield_chunk!(block_delta_chunk( &request_id, &response_model, created, AssistantBlockKind::Text, echo_text.clone(), - )) - .await; + )); } } Ok(ChatEvent::BlockDelta { kind, delta, .. }) => { @@ -301,14 +319,13 @@ async fn chat_completion_chunk_stream( if let Some(pending_chunk) = pending_chunk.as_mut() { pending_chunk.push_block_delta(kind, delta); } else { - y.yield_ok(block_delta_chunk( + yield_chunk!(block_delta_chunk( &request_id, &response_model, created, kind, delta, - )) - .await; + )); } } else { suppress_current_update_metadata = true; @@ -318,6 +335,8 @@ async fn chat_completion_chunk_stream( logprobs, token_ids, }) => { + let delta_token_count = token_ids.len(); + continuous_usage.add_output_tokens(delta_token_count); let include_metadata = !suppress_current_update_metadata && !inside_hidden_reasoning; suppress_current_update_metadata = false; @@ -339,16 +358,15 @@ async fn chat_completion_chunk_stream( if let Some(chunk) = pending_chunk.take_chunk(&request_id, &response_model, created) { - y.yield_ok(chunk).await; + yield_chunk!(chunk); } } else if let Some(logprobs) = openai_logprobs { - y.yield_ok(logprobs_only_chunk( + yield_chunk!(logprobs_only_chunk( &request_id, &response_model, created, logprobs, - )) - .await; + )); } } Ok(ChatEvent::BlockStart { kind, .. }) => { @@ -376,15 +394,14 @@ async fn chat_completion_chunk_stream( if let Some(pending_chunk) = pending_chunk.as_mut() { pending_chunk.push_tool_call_start(tool_index, id, name); } else { - y.yield_ok(tool_call_start_chunk( + yield_chunk!(tool_call_start_chunk( &request_id, &response_model, created, tool_index, id, name, - )) - .await; + )); } } Ok(ChatEvent::ToolCallArgumentsDelta { index, delta }) => { @@ -392,21 +409,20 @@ async fn chat_completion_chunk_stream( if let Some(pending_chunk) = pending_chunk.as_mut() { pending_chunk.push_tool_call_arguments(tool_index, delta); } else { - y.yield_ok(tool_call_arguments_chunk( + yield_chunk!(tool_call_arguments_chunk( &request_id, &response_model, created, tool_index, delta, - )) - .await; + )); } } Ok(ChatEvent::ToolCallEnd { .. }) => { debug!("ending current tool call"); } Ok(ChatEvent::Done { - usage, + usage: final_usage, finish_reason, .. }) => { @@ -414,18 +430,23 @@ async fn chat_completion_chunk_stream( info!( stream = true, model = %response_model, - prompt_tokens = usage.prompt_token_count, - output_tokens = usage.output_token_count, + prompt_tokens = final_usage.prompt_token_count, + output_tokens = final_usage.output_token_count, finish_reason = finish_reason.as_str(), "chat completion finished" ); } + continuous_usage.set_final_counts( + final_usage.prompt_token_count, + final_usage.output_token_count, + ); + if let Some(pending_chunk) = pending_chunk.as_mut() && let Some(chunk) = pending_chunk.take_chunk(&request_id, &response_model, created) { - y.yield_ok(chunk).await; + yield_chunk!(chunk); } match final_chunk( @@ -435,7 +456,7 @@ async fn chat_completion_chunk_stream( finish_reason, saw_tool_calls, ) { - Ok(chunk) => y.yield_ok(chunk).await, + Ok(chunk) => yield_chunk!(chunk), Err(error) => { error!( error = %error.to_error_response().error.message, @@ -450,7 +471,7 @@ async fn chat_completion_chunk_stream( &request_id, &response_model, created, - Usage::from_token_usage(usage, enable_prompt_tokens_details), + Usage::from_token_usage(final_usage, enable_prompt_tokens_details), )) .await; } diff --git a/rust/src/server/src/routes/openai/chat_completions/convert.rs b/rust/src/server/src/routes/openai/chat_completions/convert.rs index 2b3e3ddb360..aa430db76cc 100644 --- a/rust/src/server/src/routes/openai/chat_completions/convert.rs +++ b/rust/src/server/src/routes/openai/chat_completions/convert.rs @@ -33,6 +33,8 @@ pub(super) struct PreparedRequest { pub(super) struct ResponseOptions { /// Whether the caller asked for the final streamed usage chunk. pub include_usage: bool, + /// Whether every streamed chunk should carry cumulative usage. + pub include_continuous_usage: bool, /// Whether the caller requested output logprobs on chat choices. pub requested_logprobs: bool, /// Whether the caller requested top-level prompt logprobs. @@ -82,6 +84,12 @@ pub(super) fn prepare_chat_request( let include_usage = (request.stream_options.as_ref()) .and_then(|options| options.include_usage) .unwrap_or(false); + let include_continuous_usage = include_usage + && request + .stream_options + .as_ref() + .and_then(|options| options.continuous_usage_stats) + .unwrap_or(false); let requested_logprobs = request.logprobs; // Auto-enable prompt logprobs for non-streaming echo, matching Python vLLM's @@ -154,6 +162,7 @@ pub(super) fn prepare_chat_request( response_model, options: ResponseOptions { include_usage, + include_continuous_usage, requested_logprobs, include_prompt_logprobs, include_reasoning, @@ -375,8 +384,8 @@ mod tests { AssistantRole, ChatCompletionMessage, ChatCompletionRequest, }; use crate::routes::openai::utils::types::{ - ChatMessage, ContentPart, Function, FunctionCallResponse, ImageUrl, MessageContent, Tool, - ToolCall, ToolChoice, ToolChoiceValue, VideoUrl, + ChatMessage, ContentPart, Function, FunctionCallResponse, ImageUrl, MessageContent, + StreamOptions, Tool, ToolCall, ToolChoice, ToolChoiceValue, VideoUrl, }; use crate::utils::{ResolvedRequestContext, resolve_request_context}; @@ -456,6 +465,46 @@ mod tests { assert_eq!(prepared.chat_request.tool_choice, ChatToolChoice::Auto); } + #[test] + fn prepare_chat_request_maps_stream_usage_and_token_format_options() { + let mut request = base_request(); + request.return_tokens_as_token_ids = Some(true); + request.stream_options = Some(StreamOptions { + include_usage: Some(true), + continuous_usage_stats: Some(true), + }); + + let prepared = prepare_chat_request( + request, + &served(&["Qwen/Qwen1.5-0.5B-Chat"]), + ResolvedRequestContext::default(), + ) + .expect("request is valid"); + + assert!(prepared.options.include_usage); + assert!(prepared.options.include_continuous_usage); + assert!(prepared.options.return_tokens_as_token_ids); + } + + #[test] + fn prepare_chat_request_gates_continuous_usage_on_include_usage() { + let mut request = base_request(); + request.stream_options = Some(StreamOptions { + include_usage: None, + continuous_usage_stats: Some(true), + }); + + let prepared = prepare_chat_request( + request, + &served(&["Qwen/Qwen1.5-0.5B-Chat"]), + ResolvedRequestContext::default(), + ) + .expect("request is valid"); + + assert!(!prepared.options.include_usage); + assert!(!prepared.options.include_continuous_usage); + } + #[test] fn prepare_chat_request_keeps_optional_sampling_fields_unset() { let prepared = prepare_chat_request( diff --git a/rust/src/server/src/routes/openai/chat_completions/validate.rs b/rust/src/server/src/routes/openai/chat_completions/validate.rs index fb64428e4b2..a623925e649 100644 --- a/rust/src/server/src/routes/openai/chat_completions/validate.rs +++ b/rust/src/server/src/routes/openai/chat_completions/validate.rs @@ -137,15 +137,6 @@ pub(super) fn validate_request_compat( "repetition_detection is not supported.", )?; - if let Some(options) = &request.stream_options - && options.continuous_usage_stats.is_some() - { - bail_invalid_request!( - param = "stream_options", - "continuous_usage_stats is not supported." - ); - } - Ok(()) } diff --git a/rust/src/server/src/routes/openai/completions.rs b/rust/src/server/src/routes/openai/completions.rs index b6e4383c7d1..9dc2e19154f 100644 --- a/rust/src/server/src/routes/openai/completions.rs +++ b/rust/src/server/src/routes/openai/completions.rs @@ -31,6 +31,7 @@ use crate::routes::openai::completions::types::{ CompletionStreamChoice, CompletionStreamResponse, }; use crate::routes::openai::utils::types::LogProbs; +use crate::routes::openai::utils::usage::ContinuousUsage; use crate::routes::openai::utils::validated_json::ValidatedJson; use crate::state::AppState; use crate::utils::{resolve_request_context, unix_timestamp}; @@ -127,6 +128,8 @@ async fn collect_completion( ResponseOptions { // Ignored: non-streaming responses always include usage. include_usage: _, + // Ignored: non-streaming responses are collected before usage is attached. + include_continuous_usage: _, echo, requested_logprobs, include_prompt_logprobs, @@ -218,6 +221,7 @@ async fn completion_chunk_stream( }: ApiServerOptions, ResponseOptions { include_usage, + include_continuous_usage, echo, requested_logprobs, // Ignored: streaming prompt logprobs are rejected for Python parity. @@ -230,6 +234,18 @@ async fn completion_chunk_stream( pin_mut!(stream); let mut visible_text_len = 0_u32; let mut first_chunk = true; + let mut continuous_usage = ContinuousUsage::default(); + + /// Yield a chunk with optional continuous usage attached. + macro_rules! yield_chunk { + ($chunk:expr) => {{ + let mut chunk = $chunk; + if include_continuous_usage { + chunk.usage = Some(continuous_usage.to_usage()); + } + y.yield_ok(CompletionSseChunk::Chunk(chunk)).await; + }}; + } while let Some(next) = stream.next().await { match next { @@ -237,6 +253,7 @@ async fn completion_chunk_stream( prompt_token_ids, .. }) => { debug!("completion stream started"); + continuous_usage.set_prompt_tokens(prompt_token_ids.len()); if let Some(prompt) = echo.as_ref() { visible_text_len = text_len(prompt); let mut chunk = @@ -247,7 +264,7 @@ async fn completion_chunk_stream( } first_chunk = false; } - y.yield_ok(CompletionSseChunk::Chunk(chunk)).await; + yield_chunk!(chunk); } else if return_token_ids { // Emit a chunk with prompt_token_ids in the first streaming response let mut chunk = @@ -256,7 +273,7 @@ async fn completion_chunk_stream( choice.prompt_token_ids = Some(prompt_token_ids.to_vec()); } first_chunk = false; - y.yield_ok(CompletionSseChunk::Chunk(chunk)).await; + yield_chunk!(chunk); } } Ok(DecodedTextEvent::TextDelta { @@ -281,10 +298,12 @@ async fn completion_chunk_stream( None }; let mut chunk = delta_chunk(&request_id, &response_model, created, delta, logprobs); + let delta_token_count = token_ids.len(); + continuous_usage.add_output_tokens(delta_token_count); if return_token_ids && let Some(choice) = chunk.choices.first_mut() { choice.token_ids = Some(token_ids); } - y.yield_ok(CompletionSseChunk::Chunk(chunk)).await; + yield_chunk!(chunk); visible_text_len = visible_text_len.saturating_add(delta_text_len); if let Some(finished) = finished { @@ -298,13 +317,17 @@ async fn completion_chunk_stream( "completion finished" ); } - y.yield_ok(CompletionSseChunk::Chunk(final_chunk( + continuous_usage.set_final_counts( + finished.usage.prompt_token_count, + finished.usage.output_token_count, + ); + let final_chunk = final_chunk( &request_id, &response_model, created, finished.finish_reason, - )?)) - .await; + )?; + yield_chunk!(final_chunk); if include_usage { y.yield_ok(CompletionSseChunk::Usage(usage_chunk( diff --git a/rust/src/server/src/routes/openai/completions/convert.rs b/rust/src/server/src/routes/openai/completions/convert.rs index 1dd73a4f530..2f6c760a990 100644 --- a/rust/src/server/src/routes/openai/completions/convert.rs +++ b/rust/src/server/src/routes/openai/completions/convert.rs @@ -25,6 +25,8 @@ pub(super) struct PreparedRequest { pub(super) struct ResponseOptions { /// Whether the caller asked for the final streamed usage chunk. pub include_usage: bool, + /// Whether every streamed chunk should carry cumulative usage. + pub include_continuous_usage: bool, /// Original text prompt that should be echoed back northbound when /// `echo=true`. pub echo: Option, @@ -74,6 +76,12 @@ pub(super) fn prepare_completion_request( let include_usage = (request.stream_options.as_ref()) .and_then(|options| options.include_usage) .unwrap_or(false); + let include_continuous_usage = include_usage + && request + .stream_options + .as_ref() + .and_then(|options| options.continuous_usage_stats) + .unwrap_or(false); let include_prompt_logprobs = prompt_logprobs.is_some(); let echo = request.echo.then(|| request.prompt.as_text().cloned()).flatten(); @@ -129,6 +137,7 @@ pub(super) fn prepare_completion_request( response_model, options: ResponseOptions { include_usage, + include_continuous_usage, echo, requested_logprobs: request.logprobs, include_prompt_logprobs, @@ -247,6 +256,55 @@ mod tests { assert!(!prepared.text_request.decode_options.skip_special_tokens); } + #[test] + fn prepare_completion_request_maps_stream_usage_and_token_format_options() { + let request: CompletionRequest = serde_json::from_value(json!({ + "model": "Qwen/Qwen1.5-0.5B-Chat", + "prompt": "hello", + "stream": true, + "stream_options": { + "include_usage": true, + "continuous_usage_stats": true + }, + "return_tokens_as_token_ids": true + })) + .expect("parse request"); + + let prepared = prepare_completion_request( + request, + &served(&["Qwen/Qwen1.5-0.5B-Chat"]), + ResolvedRequestContext::default(), + ) + .expect("prepare"); + + assert!(prepared.options.include_usage); + assert!(prepared.options.include_continuous_usage); + assert!(prepared.options.return_tokens_as_token_ids); + } + + #[test] + fn prepare_completion_request_gates_continuous_usage_on_include_usage() { + let request: CompletionRequest = serde_json::from_value(json!({ + "model": "Qwen/Qwen1.5-0.5B-Chat", + "prompt": "hello", + "stream": true, + "stream_options": { + "continuous_usage_stats": true + } + })) + .expect("parse request"); + + let prepared = prepare_completion_request( + request, + &served(&["Qwen/Qwen1.5-0.5B-Chat"]), + ResolvedRequestContext::default(), + ) + .expect("prepare"); + + assert!(!prepared.options.include_usage); + assert!(!prepared.options.include_continuous_usage); + } + #[test] fn prepare_completion_request_accepts_text_echo() { let request: CompletionRequest = serde_json::from_value(json!({ diff --git a/rust/src/server/src/routes/openai/completions/validate.rs b/rust/src/server/src/routes/openai/completions/validate.rs index 2af8c8add11..2af41877bfd 100644 --- a/rust/src/server/src/routes/openai/completions/validate.rs +++ b/rust/src/server/src/routes/openai/completions/validate.rs @@ -95,15 +95,6 @@ pub(super) fn validate_request_compat( ); } - if let Some(options) = &request.stream_options - && options.continuous_usage_stats.is_some() - { - bail_invalid_request!( - param = "stream_options", - "continuous_usage_stats is not supported." - ); - } - Ok(()) } diff --git a/rust/src/server/src/routes/openai/utils/mod.rs b/rust/src/server/src/routes/openai/utils/mod.rs index 039df87f9dd..7ec1251ddf3 100644 --- a/rust/src/server/src/routes/openai/utils/mod.rs +++ b/rust/src/server/src/routes/openai/utils/mod.rs @@ -2,4 +2,5 @@ pub mod logprobs; pub mod structured_outputs; pub mod token_ids; pub mod types; +pub mod usage; pub mod validated_json; diff --git a/rust/src/server/src/routes/openai/utils/usage.rs b/rust/src/server/src/routes/openai/utils/usage.rs new file mode 100644 index 00000000000..c8c9d1e7262 --- /dev/null +++ b/rust/src/server/src/routes/openai/utils/usage.rs @@ -0,0 +1,35 @@ +use super::types::Usage; + +/// Tracks cumulative token counts for OpenAI streaming chunks. +/// +/// This helper is intentionally only a counter. Callers decide whether to +/// attach `counts()` to each streamed data chunk, while final usage-only chunks +/// should still be built from the authoritative terminal `TokenUsage`. +#[derive(Debug, Clone, Default)] +pub(crate) struct ContinuousUsage { + prompt_tokens: usize, + output_tokens: usize, +} + +impl ContinuousUsage { + /// Record the prompt-token count reported when a stream starts. + pub(crate) fn set_prompt_tokens(&mut self, prompt_tokens: usize) { + self.prompt_tokens = prompt_tokens; + } + + /// Add newly decoded output tokens to the running completion count. + pub(crate) fn add_output_tokens(&mut self, output_tokens: usize) { + self.output_tokens = self.output_tokens.saturating_add(output_tokens); + } + + /// Replace the running counts with the final counts reported by generation. + pub(crate) fn set_final_counts(&mut self, prompt_tokens: usize, output_tokens: usize) { + self.prompt_tokens = prompt_tokens; + self.output_tokens = output_tokens; + } + + /// Build a streaming usage snapshot without prompt cache details. + pub(crate) fn to_usage(&self) -> Usage { + Usage::from_counts(self.prompt_tokens, self.output_tokens, None) + } +} diff --git a/rust/src/server/src/routes/tests.rs b/rust/src/server/src/routes/tests.rs index 68ffe04a3b7..c6de4034026 100644 --- a/rust/src/server/src/routes/tests.rs +++ b/rust/src/server/src/routes/tests.rs @@ -151,6 +151,14 @@ fn sse_data_payloads(text: &str) -> Vec<&str> { text.lines().filter_map(|line| line.strip_prefix("data: ")).collect() } +fn sse_json_payloads(text: &str) -> Vec { + sse_data_payloads(text) + .into_iter() + .filter(|payload| *payload != "[DONE]") + .map(|payload| serde_json::from_str(payload).expect("sse json payload")) + .collect() +} + type TestFuture<'a> = Pin + Send + 'a>>; fn boxed_test_future<'a>(future: impl Future + Send + 'a) -> TestFuture<'a> { @@ -2341,6 +2349,60 @@ async fn include_usage_adds_final_usage_chunk_before_done() { assert_eq!(usage_chunk["usage"]["total_tokens"], 25); } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn stream_continuous_usage_stats_adds_usage_to_chat_chunks() { + let (app, engine_task) = test_app_with_stream_output_specs(default_stream_output_specs()).await; + let response = app + .clone() + .call( + Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "Qwen/Qwen1.5-0.5B-Chat", + "stream": true, + "stream_options": { + "include_usage": true, + "continuous_usage_stats": true + }, + "messages": [{"role": "user", "content": "hello"}] + }) + .to_string(), + )) + .expect("build request"), + ) + .await + .expect("call app"); + + assert_eq!(response.status(), StatusCode::OK); + + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body"); + engine_task.await.expect("mock engine task"); + let text = String::from_utf8(body.to_vec()).expect("utf8 body"); + let payloads = sse_json_payloads(&text); + + assert!( + payloads.iter().all(|payload| payload.get("usage").is_some()), + "{text}" + ); + assert!( + payloads.iter().any(|payload| { + payload["choices"].as_array().is_some_and(|choices| !choices.is_empty()) + && payload["usage"]["completion_tokens"] == json!(1) + }), + "{text}" + ); + let usage_chunk = payloads + .iter() + .find(|payload| payload["choices"] == json!([])) + .expect("final usage chunk"); + assert_eq!(usage_chunk["usage"]["prompt_tokens"], 22); + assert_eq!(usage_chunk["usage"]["completion_tokens"], 3); +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[serial] async fn stream_without_include_usage_keeps_existing_shape() { @@ -3434,6 +3496,59 @@ async fn completions_happy_path_returns_sse_stream() { assert_eq!(usage_chunk["usage"]["completion_tokens"], 3); } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn completions_stream_continuous_usage_stats_adds_usage_to_chunks() { + let (app, engine_task) = test_app_with_stream_output_specs(default_stream_output_specs()).await; + let response = app + .clone() + .call( + Request::builder() + .method("POST") + .uri("/v1/completions") + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "Qwen/Qwen1.5-0.5B-Chat", + "prompt": "hello", + "stream": true, + "stream_options": { + "include_usage": true, + "continuous_usage_stats": true + } + }) + .to_string(), + )) + .expect("build request"), + ) + .await + .expect("call app"); + + assert_eq!(response.status(), StatusCode::OK); + + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body"); + engine_task.await.expect("mock engine task"); + let text = String::from_utf8(body.to_vec()).expect("utf8 body"); + let payloads = sse_json_payloads(&text); + + assert!( + payloads.iter().all(|payload| payload.get("usage").is_some()), + "{text}" + ); + assert!( + payloads.iter().any(|payload| { + payload["choices"].as_array().is_some_and(|choices| !choices.is_empty()) + && payload["usage"]["completion_tokens"] == json!(1) + }), + "{text}" + ); + let usage_chunk = payloads + .iter() + .find(|payload| payload["choices"] == json!([])) + .expect("final usage chunk"); + assert_eq!(usage_chunk["usage"]["completion_tokens"], 3); +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[serial] async fn completions_echo_stream_emits_separate_prompt_chunk() { From 235b63c0046d2fbf4ab1bf810a1eb729f1f3fc27 Mon Sep 17 00:00:00 2001 From: Ben Browning Date: Thu, 11 Jun 2026 16:01:29 -0400 Subject: [PATCH 10/52] [Bugfix] Fix Anthropic tool_use content handling dropping args (#45287) Signed-off-by: Ben Browning --- .../test_anthropic_messages_conversion.py | 223 +++++++++++++++++- vllm/entrypoints/anthropic/serving.py | 34 ++- 2 files changed, 252 insertions(+), 5 deletions(-) diff --git a/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py b/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py index ad9fed1d355..21d5154c675 100644 --- a/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py +++ b/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py @@ -6,13 +6,29 @@ Tests the image source handling and tool_result content parsing in AnthropicServingMessages._convert_anthropic_to_openai_request(). Also covers extended-thinking edge cases such as ``redacted_thinking`` -blocks echoed back by Anthropic clients. +blocks echoed back by Anthropic clients, and streaming conversion in +``message_stream_converter``. """ +import json +from unittest.mock import MagicMock + +import pytest + from vllm.entrypoints.anthropic.protocol import ( AnthropicMessagesRequest, ) from vllm.entrypoints.anthropic.serving import AnthropicServingMessages +from vllm.entrypoints.openai.chat_completion.protocol import ( + ChatCompletionResponseStreamChoice, + ChatCompletionStreamResponse, +) +from vllm.entrypoints.openai.engine.protocol import ( + DeltaFunctionCall, + DeltaMessage, + DeltaToolCall, + UsageInfo, +) _convert = AnthropicServingMessages._convert_anthropic_to_openai_request _img_url = AnthropicServingMessages._convert_image_source_to_url @@ -775,3 +791,208 @@ class TestInlineSystemMessageInMessagesArray: assert result.messages[0]["role"] == "system" assert result.messages[0]["content"] == "Top-level prompt.Inline hint." assert result.messages[1]["role"] == "user" + + +# ====================================================================== +# Streaming conversion: message_stream_converter +# ====================================================================== + + +def _make_stream_converter(): + obj = MagicMock(spec=AnthropicServingMessages) + obj.stop_reason_map = { + "stop": "end_turn", + "length": "max_tokens", + "tool_calls": "tool_use", + } + obj.message_stream_converter = ( + AnthropicServingMessages.message_stream_converter.__get__(obj) + ) + return obj + + +def _parse_sse_events(raw_events: list[str]) -> list[tuple[str, dict]]: + results = [] + for raw in raw_events: + headers = dict( + line.split(": ", 1) for line in raw.strip().split("\n") if ": " in line + ) + if "event" in headers and "data" in headers: + results.append((headers["event"], json.loads(headers["data"]))) + return results + + +def _make_stream_chunk( + *, + delta: DeltaMessage | None = None, + finish_reason: str | None = None, + choices: list[ChatCompletionResponseStreamChoice] | None = None, + usage: UsageInfo | None = None, +) -> str: + if choices is None: + choices = [ + ChatCompletionResponseStreamChoice( + index=0, + delta=delta or DeltaMessage(), + finish_reason=finish_reason, + ) + ] + chunk = ChatCompletionStreamResponse( + id="chatcmpl-test", + created=0, + model="test-model", + choices=choices, + usage=usage, + ) + return f"data: {chunk.model_dump_json()}" + + +def _tc(*, args, id=None, name=None): + return DeltaToolCall( + index=0, + id=id, + function=DeltaFunctionCall(name=name, arguments=args), + ) + + +class TestMessageStreamConverterToolUseContentBuffering: + """Regression test for tool_use arguments being silently dropped. + + With speculative decoding or multi-token prediction, a single delta + can carry both the final tool_call argument fragment and trailing + content. + """ + + @pytest.mark.asyncio + async def test_tool_use_args_not_dropped_when_content_in_same_chunk( + self, + ): + async def sse_input(): + yield _make_stream_chunk( + delta=DeltaMessage(role="assistant"), + usage=UsageInfo(prompt_tokens=10, total_tokens=10), + ) + yield _make_stream_chunk( + delta=DeltaMessage( + tool_calls=[ + _tc(id="call_abc123", name="read_file", args=""), + ] + ) + ) + yield _make_stream_chunk( + delta=DeltaMessage( + tool_calls=[ + _tc(args='{"path":"/tmp/f"'), + ] + ) + ) + # BUG TRIGGER: final tool_call args and trailing content in + # one delta, as happens with spec decoding / multi-token + # prediction where multiple tokens land in a single chunk. + yield _make_stream_chunk( + delta=DeltaMessage( + content="\nOkay", + tool_calls=[_tc(args="}")], + ) + ) + yield _make_stream_chunk(finish_reason="tool_calls") + yield _make_stream_chunk( + choices=[], + usage=UsageInfo( + prompt_tokens=10, + total_tokens=30, + completion_tokens=20, + ), + ) + yield "data: [DONE]" + + converter = _make_stream_converter() + output = [] + async for event in converter.message_stream_converter(sse_input()): + output.append(event) + + events = _parse_sse_events(output) + + assert events[0][0] == "message_start" + + arg_fragments = [ + data["delta"]["partial_json"] + for _, data in events + if data.get("delta", {}).get("type") == "input_json_delta" + ] + full_args = "".join(arg_fragments) + assert full_args == '{"path":"/tmp/f"}' + + text_deltas = [ + data["delta"]["text"] + for _, data in events + if data.get("delta", {}).get("type") == "text_delta" + ] + assert text_deltas == ["\nOkay"] + + block_starts = [ + (data["content_block"]["type"], data.get("index")) + for ev_type, data in events + if ev_type == "content_block_start" + ] + assert block_starts[0] == ("tool_use", 0) + assert block_starts[1] == ("text", 1) + + msg_deltas = [data for ev_type, data in events if ev_type == "message_delta"] + assert msg_deltas[0]["delta"]["stop_reason"] == "tool_use" + + assert events[-1][0] == "message_stop" + + @pytest.mark.asyncio + async def test_buffered_content_flushed_on_done_without_usage_chunk(self): + """Content buffered during tool_use must be emitted even if the + stream jumps straight from finish_reason to [DONE], skipping the + empty-choices usage chunk.""" + + async def sse_input(): + yield _make_stream_chunk( + delta=DeltaMessage(role="assistant"), + usage=UsageInfo(prompt_tokens=10, total_tokens=10), + ) + yield _make_stream_chunk( + delta=DeltaMessage( + tool_calls=[ + _tc(id="call_xyz", name="get_weather", args=""), + ] + ) + ) + yield _make_stream_chunk( + delta=DeltaMessage( + tool_calls=[_tc(args='{"city":"NYC"}')], + ) + ) + yield _make_stream_chunk( + delta=DeltaMessage(content="\nDone"), + finish_reason="tool_calls", + ) + # No empty-choices usage chunk — go straight to [DONE]. + yield "data: [DONE]" + + converter = _make_stream_converter() + output = [] + async for event in converter.message_stream_converter(sse_input()): + output.append(event) + + events = _parse_sse_events(output) + + text_deltas = [ + data["delta"]["text"] + for _, data in events + if data.get("delta", {}).get("type") == "text_delta" + ] + assert text_deltas == ["\nDone"] + + block_starts = [ + data["content_block"]["type"] + for ev_type, data in events + if ev_type == "content_block_start" + ] + assert "tool_use" in block_starts + assert "text" in block_starts + + assert events[-1][0] == "message_stop" diff --git a/vllm/entrypoints/anthropic/serving.py b/vllm/entrypoints/anthropic/serving.py index 8f6cccdb0fc..266a3154212 100644 --- a/vllm/entrypoints/anthropic/serving.py +++ b/vllm/entrypoints/anthropic/serving.py @@ -564,6 +564,7 @@ class AnthropicServingMessages(OpenAIServingChat): self.block_signature: str | None = None self.signature_emitted: bool = False self.tool_use_id: str | None = None + self.pending_content: list[str] = [] def reset(self) -> None: self.block_type = None @@ -571,6 +572,7 @@ class AnthropicServingMessages(OpenAIServingChat): self.block_signature = None self.signature_emitted = False self.tool_use_id = None + self.pending_content.clear() def start(self, block: AnthropicContentBlock) -> None: self.block_type = block.type @@ -635,10 +637,30 @@ class AnthropicServingMessages(OpenAIServingChat): state.start(block) return event + def stop_and_flush() -> list[str]: + buffered = list(state.pending_content) + state.pending_content.clear() + events = stop_active_block() + if not buffered: + return events + text = "".join(buffered) + events.append(start_block(AnthropicContentBlock(type="text", text=""))) + pc_chunk = AnthropicStreamEvent( + index=state.block_index, + type="content_block_delta", + delta=AnthropicDelta(type="text_delta", text=text), + ) + pc_data = pc_chunk.model_dump_json(exclude_unset=True) + events.append(wrap_data_with_event(pc_data, "content_block_delta")) + events.extend(stop_active_block()) + return events + async for item in generator: if item.startswith("data:"): data_str = item[5:].strip().rstrip("\n") if data_str == "[DONE]": + for event in stop_and_flush(): + yield event stop_message = AnthropicStreamEvent( type="message_stop", ) @@ -675,7 +697,7 @@ class AnthropicServingMessages(OpenAIServingChat): # last chunk including usage info if len(origin_chunk.choices) == 0: - for event in stop_active_block(): + for event in stop_and_flush(): yield event stop_reason = self.stop_reason_map.get( finish_reason or "stop" @@ -707,7 +729,7 @@ class AnthropicServingMessages(OpenAIServingChat): pass else: if state.block_type != "thinking": - for event in stop_active_block(): + for event in stop_and_flush(): yield event start_event = start_block( AnthropicContentBlock( @@ -733,9 +755,13 @@ class AnthropicServingMessages(OpenAIServingChat): if origin_chunk.choices[0].delta.content is not None: if origin_chunk.choices[0].delta.content == "": pass + elif state.block_type == "tool_use": + state.pending_content.append( + origin_chunk.choices[0].delta.content + ) else: if state.block_type != "text": - for event in stop_active_block(): + for event in stop_and_flush(): yield event start_event = start_block( AnthropicContentBlock(type="text", text="") @@ -773,7 +799,7 @@ class AnthropicServingMessages(OpenAIServingChat): state.tool_use_id != tool_call.id and tool_name is not None ): - for event in stop_active_block(): + for event in stop_and_flush(): yield event start_event = start_block( AnthropicContentBlock( From c9340e6f350a009cf835878abad2a0e379b9e6a4 Mon Sep 17 00:00:00 2001 From: Tiezhen WANG <38108242+xianbaoqian@users.noreply.github.com> Date: Fri, 12 Jun 2026 04:02:51 +0800 Subject: [PATCH 11/52] [Model] Remove InternLMForCausalLM registry alias (#45128) Signed-off-by: Xianbao QIAN Co-authored-by: Claude --- docs/models/supported_models.md | 1 - tests/distributed/test_pipeline_parallel.py | 2 -- tests/models/registry.py | 3 --- vllm/model_executor/models/apertus.py | 1 - vllm/model_executor/models/exaone.py | 1 - vllm/model_executor/models/exaone4.py | 1 - vllm/model_executor/models/exaone_moe.py | 1 - vllm/model_executor/models/granite.py | 1 - vllm/model_executor/models/jais2.py | 1 - vllm/model_executor/models/llama.py | 1 - vllm/model_executor/models/nemotron.py | 1 - vllm/model_executor/models/nemotron_nas.py | 1 - vllm/model_executor/models/registry.py | 2 +- vllm/model_executor/models/solar.py | 1 - 14 files changed, 1 insertion(+), 17 deletions(-) diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 6f7cc6dab4b..1823ddcecc6 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -423,7 +423,6 @@ th { | `HunYuanMoEV1ForCausalLM` | Hunyuan-A13B | `tencent/Hunyuan-A13B-Instruct`, `tencent/Hunyuan-A13B-Pretrain`, `tencent/Hunyuan-A13B-Instruct-FP8`, etc. | ✅︎ | ✅︎ | | `HYV3ForCausalLM` | HY3 | `tencent/Hy3-preview-Base`, `tencent/Hy3-preview` | ✅︎ | ✅︎ | | `HyperCLOVAXForCausalLM` | HyperCLOVAX-SEED-Think-14B | `naver-hyperclovax/HyperCLOVAX-SEED-Think-14B` | ✅︎ | ✅︎ | -| `InternLMForCausalLM` | InternLM | `internlm/internlm-7b`, `internlm/internlm-chat-7b`, etc. | ✅︎ | ✅︎ | | `InternLM2ForCausalLM` | InternLM2 | `internlm/internlm2-7b`, `internlm/internlm2-chat-7b`, etc. | ✅︎ | ✅︎ | | `InternLM3ForCausalLM` | InternLM3 | `internlm/internlm3-8b-instruct`, etc. | ✅︎ | ✅︎ | | `IQuestCoderForCausalLM` | IQuestCoderV1 | `IQuestLab/IQuest-Coder-V1-40B-Instruct`, etc. | | | diff --git a/tests/distributed/test_pipeline_parallel.py b/tests/distributed/test_pipeline_parallel.py index 93f3abfc088..85307403200 100644 --- a/tests/distributed/test_pipeline_parallel.py +++ b/tests/distributed/test_pipeline_parallel.py @@ -124,8 +124,6 @@ TEXT_GENERATION_MODELS = { "EleutherAI/pythia-1.4b": PPTestSettings.fast(), "ibm/PowerLM-3b": PPTestSettings.fast(), "ibm/PowerMoE-3b": PPTestSettings.fast(), - # Uses Llama - # "internlm/internlm-chat-7b": PPTestSettings.fast(), "internlm/internlm2-chat-7b": PPTestSettings.fast(), "ai21labs/Jamba-tiny-dev": PPTestSettings.fast(), "pfnet/plamo-2-1b": PPTestSettings.fast(), diff --git a/tests/models/registry.py b/tests/models/registry.py index d2d2794962f..120a0ca8b85 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -338,9 +338,6 @@ _TEXT_GENERATION_EXAMPLE_MODELS = { "naver-hyperclovax/HyperCLOVAX-SEED-Think-14B", min_transformers_version="5.9.0", ), - "InternLMForCausalLM": _HfExamplesInfo( - "internlm/internlm-chat-7b", trust_remote_code=True - ), "InternLM2ForCausalLM": _HfExamplesInfo( "internlm/internlm2-chat-7b", trust_remote_code=True ), diff --git a/vllm/model_executor/models/apertus.py b/vllm/model_executor/models/apertus.py index 0711fb03f84..a857769cbe1 100644 --- a/vllm/model_executor/models/apertus.py +++ b/vllm/model_executor/models/apertus.py @@ -252,7 +252,6 @@ class ApertusDecoderLayer(nn.Module): self.hidden_size = config.hidden_size max_position_embeddings = getattr(config, "max_position_embeddings", 8192) # Support abacusai/Smaug-72B-v0.1 with attention_bias - # Support internlm/internlm-7b with bias attention_bias = getattr(config, "attention_bias", False) or getattr( config, "bias", False ) diff --git a/vllm/model_executor/models/exaone.py b/vllm/model_executor/models/exaone.py index dca05f72c69..be45d7dfb2b 100644 --- a/vllm/model_executor/models/exaone.py +++ b/vllm/model_executor/models/exaone.py @@ -243,7 +243,6 @@ class ExaoneDecoderLayer(nn.Module): self.hidden_size = config.hidden_size max_position_embeddings = getattr(config, "max_position_embeddings", 8192) # Support abacusai/Smaug-72B-v0.1 with attention_bias - # Support internlm/internlm-7b with bias attention_bias = getattr(config, "attention_bias", False) or getattr( config, "bias", False ) diff --git a/vllm/model_executor/models/exaone4.py b/vllm/model_executor/models/exaone4.py index e38dbb5ee29..a36b8e0e922 100644 --- a/vllm/model_executor/models/exaone4.py +++ b/vllm/model_executor/models/exaone4.py @@ -230,7 +230,6 @@ class Exaone4DecoderLayer(nn.Module): self.hidden_size = config.hidden_size max_position_embeddings = getattr(config, "max_position_embeddings", 8192) # Support abacusai/Smaug-72B-v0.1 with attention_bias - # Support internlm/internlm-7b with bias attention_bias = getattr(config, "attention_bias", False) or getattr( config, "bias", False ) diff --git a/vllm/model_executor/models/exaone_moe.py b/vllm/model_executor/models/exaone_moe.py index 3373983f5c9..18900557f61 100644 --- a/vllm/model_executor/models/exaone_moe.py +++ b/vllm/model_executor/models/exaone_moe.py @@ -179,7 +179,6 @@ class ExaoneMoeDecoderLayer(nn.Module): self.hidden_size = config.hidden_size max_position_embeddings = getattr(config, "max_position_embeddings", 8192) # Support abacusai/Smaug-72B-v0.1 with attention_bias - # Support internlm/internlm-7b with bias attention_bias = getattr(config, "attention_bias", False) or getattr( config, "bias", False ) diff --git a/vllm/model_executor/models/granite.py b/vllm/model_executor/models/granite.py index 2adc29f8d25..7470e7e7381 100644 --- a/vllm/model_executor/models/granite.py +++ b/vllm/model_executor/models/granite.py @@ -199,7 +199,6 @@ class GraniteDecoderLayer(nn.Module): self.residual_multiplier = config.residual_multiplier max_position_embeddings = getattr(config, "max_position_embeddings", 8192) # Support abacusai/Smaug-72B-v0.1 with attention_bias - # Support internlm/internlm-7b with bias attention_bias = getattr(config, "attention_bias", False) or getattr( config, "bias", False ) diff --git a/vllm/model_executor/models/jais2.py b/vllm/model_executor/models/jais2.py index dafa0f03ae9..67b0ac5033f 100644 --- a/vllm/model_executor/models/jais2.py +++ b/vllm/model_executor/models/jais2.py @@ -225,7 +225,6 @@ class Jais2DecoderLayer(nn.Module): self.hidden_size = config.hidden_size max_position_embeddings = getattr(config, "max_position_embeddings", 8192) # Support abacusai/Smaug-72B-v0.1 with attention_bias - # Support internlm/internlm-7b with bias attention_bias = getattr(config, "attention_bias", False) or getattr( config, "bias", False ) diff --git a/vllm/model_executor/models/llama.py b/vllm/model_executor/models/llama.py index 39044f5e8b4..c35896264a9 100644 --- a/vllm/model_executor/models/llama.py +++ b/vllm/model_executor/models/llama.py @@ -268,7 +268,6 @@ class LlamaDecoderLayer(nn.Module): self.hidden_size = config.hidden_size max_position_embeddings = getattr(config, "max_position_embeddings", 8192) # Support abacusai/Smaug-72B-v0.1 with attention_bias - # Support internlm/internlm-7b with bias attention_bias = getattr(config, "attention_bias", False) or getattr( config, "bias", False ) diff --git a/vllm/model_executor/models/nemotron.py b/vllm/model_executor/models/nemotron.py index 7b2e6b93b27..f5c526e33ed 100644 --- a/vllm/model_executor/models/nemotron.py +++ b/vllm/model_executor/models/nemotron.py @@ -237,7 +237,6 @@ class NemotronDecoderLayer(nn.Module): self.hidden_size = config.hidden_size max_position_embeddings = getattr(config, "max_position_embeddings", 8192) # Support abacusai/Smaug-72B-v0.1 with attention_bias - # Support internlm/internlm-7b with bias attention_bias = getattr(config, "attention_bias", False) or getattr( config, "bias", False ) diff --git a/vllm/model_executor/models/nemotron_nas.py b/vllm/model_executor/models/nemotron_nas.py index b974a3eb085..06a2096ec69 100644 --- a/vllm/model_executor/models/nemotron_nas.py +++ b/vllm/model_executor/models/nemotron_nas.py @@ -141,7 +141,6 @@ class DeciLMDecoderLayer(nn.Module): self.hidden_size = config.hidden_size max_position_embeddings = getattr(config, "max_position_embeddings", 8192) # Support abacusai/Smaug-72B-v0.1 with attention_bias - # Support internlm/internlm-7b with bias attention_bias = getattr(config, "attention_bias", False) or getattr( config, "bias", False ) diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index e1ce0efae2f..175f0f2dab2 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -140,7 +140,6 @@ _TEXT_GENERATION_MODELS = { "HCXVisionForCausalLM": ("hyperclovax_vision", "HCXVisionForCausalLM"), "HCXVisionV2ForCausalLM": ("hyperclovax_vision_v2", "HCXVisionV2ForCausalLM"), "HyperCLOVAXForCausalLM": ("hyperclovax", "HyperCLOVAXForCausalLM"), - "InternLMForCausalLM": ("llama", "LlamaForCausalLM"), "InternLM2ForCausalLM": ("internlm2", "InternLM2ForCausalLM"), "InternLM2VEForCausalLM": ("internlm2_ve", "InternLM2VEForCausalLM"), "InternLM3ForCausalLM": ("llama", "LlamaForCausalLM"), @@ -715,6 +714,7 @@ _PREVIOUSLY_SUPPORTED_MODELS = { "ErnieForTokenClassification": "0.23.0", "QWenLMHeadModel": "0.23.0", "QwenVLForConditionalGeneration": "0.23.0", + "InternLMForCausalLM": "0.23.0", # encoder-decoder models except whisper # have been removed for V0 deprecation. "DonutForConditionalGeneration": "0.10.2", diff --git a/vllm/model_executor/models/solar.py b/vllm/model_executor/models/solar.py index 454a0e97112..fcb2ae429cb 100644 --- a/vllm/model_executor/models/solar.py +++ b/vllm/model_executor/models/solar.py @@ -198,7 +198,6 @@ class SolarDecoderLayer(nn.Module): self.hidden_size = config.hidden_size max_position_embeddings = getattr(config, "max_position_embeddings", 8192) # Support abacusai/Smaug-72B-v0.1 with attention_bias - # Support internlm/internlm-7b with bias attention_bias = getattr(config, "attention_bias", False) or getattr( config, "bias", False ) From 5a6c7b7ab569f49491b5428a7983be5b17b85378 Mon Sep 17 00:00:00 2001 From: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Date: Thu, 11 Jun 2026 16:22:26 -0400 Subject: [PATCH 12/52] [Bug] Fix test flashmla for DSv4 (#45052) Signed-off-by: yewentao256 --- tests/kernels/attention/test_flashmla_sparse.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/tests/kernels/attention/test_flashmla_sparse.py b/tests/kernels/attention/test_flashmla_sparse.py index 9e4e7c2ec9a..d92dabe9d3e 100644 --- a/tests/kernels/attention/test_flashmla_sparse.py +++ b/tests/kernels/attention/test_flashmla_sparse.py @@ -29,8 +29,10 @@ def test_sparse_flashmla_metadata_smoke(): topk=topk, is_fp8_kvcache=True, ) - assert tile_md.dtype == torch.int32 - assert num_splits.dtype == torch.int32 + assert isinstance(tile_md, fm.FlashMLASchedMeta) + assert tile_md.tile_scheduler_metadata is None + assert tile_md.num_splits is None + assert num_splits is None def test_sparse_flashmla_decode_smoke(): @@ -116,7 +118,7 @@ def test_sparse_flashmla_prefill_smoke(): kv = torch.zeros((s_kv, h_kv, d_qk), dtype=torch.bfloat16, device=device) indices = torch.zeros((s_q, h_kv, topk), dtype=torch.int32, device=device) - out, max_logits, lse = fm.flash_mla_sparse_prefill(q, kv, indices, 1.0, d_v) + out, max_logits, lse = fm.flash_mla_sparse_fwd(q, kv, indices, 1.0, d_v) assert out.shape == (s_q, h_q, d_v) assert max_logits.shape == (s_q, h_q) assert lse.shape == (s_q, h_q) From f712fd0d7db6e0b2c7fbdb6e77cae155c81fd8c5 Mon Sep 17 00:00:00 2001 From: yzong-rh Date: Thu, 11 Jun 2026 17:18:30 -0400 Subject: [PATCH 13/52] [Refactor] Chat Completions Harmony Refactor, non-streaming path. (#45171) Signed-off-by: Yifan Zong --- .../chat_completion/test_serving_chat.py | 48 +- .../openai/parser/test_harmony_utils.py | 106 ---- tests/parser/test_harmony.py | 452 ++++++++++++++++++ tests/tool_parsers/test_openai_tool_parser.py | 415 ---------------- .../openai/chat_completion/serving.py | 84 +--- .../openai/parser/harmony_utils.py | 64 +-- vllm/entrypoints/openai/responses/serving.py | 1 + vllm/parser/__init__.py | 2 + vllm/parser/abstract_parser.py | 3 + vllm/parser/harmony.py | 240 ++++++++++ vllm/parser/mistral.py | 7 +- vllm/parser/parser_manager.py | 10 + vllm/reasoning/gptoss_reasoning_parser.py | 35 +- vllm/tool_parsers/__init__.py | 4 +- vllm/tool_parsers/gptoss_tool_parser.py | 47 ++ vllm/tool_parsers/openai_tool_parser.py | 120 ----- 16 files changed, 822 insertions(+), 816 deletions(-) create mode 100644 tests/parser/test_harmony.py delete mode 100644 tests/tool_parsers/test_openai_tool_parser.py create mode 100644 vllm/parser/harmony.py create mode 100644 vllm/tool_parsers/gptoss_tool_parser.py delete mode 100644 vllm/tool_parsers/openai_tool_parser.py diff --git a/tests/entrypoints/openai/chat_completion/test_serving_chat.py b/tests/entrypoints/openai/chat_completion/test_serving_chat.py index 22077bd4a31..e523cc2d4a3 100644 --- a/tests/entrypoints/openai/chat_completion/test_serving_chat.py +++ b/tests/entrypoints/openai/chat_completion/test_serving_chat.py @@ -38,12 +38,12 @@ from vllm.entrypoints.serve.render.serving import OpenAIServingRender from vllm.exceptions import VLLMValidationError from vllm.inputs import TokensPrompt from vllm.outputs import CompletionOutput, RequestOutput +from vllm.parser import HarmonyParser from vllm.renderers.hf import HfRenderer from vllm.renderers.mistral import MistralRenderer from vllm.tokenizers import get_tokenizer from vllm.tokenizers.mistral import MistralTokenizer from vllm.tokenizers.registry import cached_tokenizer_from_config -from vllm.tool_parsers import ToolParserManager from vllm.v1.engine.async_llm import AsyncLLM GPT_OSS_MODEL_NAME = "openai/gpt-oss-20b" @@ -575,7 +575,13 @@ def _build_serving_render( ) -def _build_serving_chat(engine: AsyncLLM) -> OpenAIServingChat: +def _build_serving_chat( + engine: AsyncLLM, + *, + reasoning_parser: str = "", + tool_parser: str | None = None, + enable_auto_tools: bool = False, +) -> OpenAIServingChat: models = OpenAIServingModels( engine_client=engine, base_model_paths=BASE_MODEL_PATHS, @@ -590,6 +596,9 @@ def _build_serving_chat(engine: AsyncLLM) -> OpenAIServingChat: chat_template=CHAT_TEMPLATE, chat_template_content_format="auto", request_logger=None, + reasoning_parser=reasoning_parser, + tool_parser=tool_parser, + enable_auto_tools=enable_auto_tools, ) return serving_chat @@ -637,7 +646,7 @@ async def test_serving_chat_returns_correct_model_name(): serving_chat = _build_serving_chat(mock_engine) messages = [{"role": "user", "content": "what is 1+1?"}] - async def return_model_name(*args): + async def return_model_name(*args, **kwargs): return args[3] serving_chat.chat_completion_full_generator = return_model_name @@ -1210,15 +1219,21 @@ class TestServingChatWithHarmony: mock_engine = MagicMock(spec=AsyncLLM) mock_engine.errored = False mock_engine.model_config = MockModelConfig() + mock_engine.model_config.hf_config = MockHFConfig(model_type="gpt_oss") + mock_engine.model_config.hf_text_config = MockHFConfig(model_type="gpt_oss") mock_engine.input_processor = MagicMock() mock_engine.renderer = _build_renderer(mock_engine.model_config) return mock_engine @pytest.fixture() def serving_chat(self, mock_engine) -> OpenAIServingChat: - chat = _build_serving_chat(mock_engine) - chat.use_harmony = True - chat.tool_parser = ToolParserManager.get_tool_parser("openai") + chat = _build_serving_chat( + mock_engine, + reasoning_parser="openai_gptoss", + tool_parser="openai", + enable_auto_tools=True, + ) + assert chat.parser_cls is HarmonyParser return chat def mock_request_output_from_req_and_token_ids( @@ -1277,6 +1292,7 @@ class TestServingChatWithHarmony: stream: bool = False, ) -> ChatCompletionResponse: harmony_token_ids = get_encoding().encode(harmony_str, allowed_special="all") + tokenizer = get_tokenizer(GPT_OSS_MODEL_NAME) async def result_generator(): if stream: @@ -1304,11 +1320,12 @@ class TestServingChatWithHarmony: request_id=req.request_id, model_name=req.model, conversation=[], - tokenizer=get_tokenizer(req.model), + tokenizer=tokenizer, request_metadata=RequestResponseMetadata( request_id=req.request_id, model_name=req.model, ), + chat_template_kwargs=serving_chat._effective_chat_template_kwargs(req), ) if stream: @@ -1316,11 +1333,18 @@ class TestServingChatWithHarmony: return await result @pytest.mark.asyncio - async def test_simple_chat(self, serving_chat, stream): + @pytest.mark.parametrize( + "include_reasoning", [True, False], ids=["with_reasoning", "no_reasoning"] + ) + async def test_simple_chat(self, serving_chat, stream, include_reasoning): messages = [{"role": "user", "content": "what is 1+1?"}] # Test the Harmony messages for the first turn's input - req = ChatCompletionRequest(model=MODEL_NAME, messages=messages) + req = ChatCompletionRequest( + model=MODEL_NAME, + messages=messages, + include_reasoning=include_reasoning, + ) input_messages, _ = ( serving_chat.openai_serving_render._make_request_with_harmony(req) ) @@ -1342,7 +1366,11 @@ class TestServingChatWithHarmony: response = await self.generate_response_from_harmony_str( serving_chat, req, response_str, stream=stream ) - verify_chat_response(response, content=final_str, reasoning=reasoning_str) + verify_chat_response( + response, + content=final_str, + reasoning=reasoning_str if include_reasoning else None, + ) # Add the output messages from the first turn as input to the second turn for choice in response.choices: diff --git a/tests/entrypoints/openai/parser/test_harmony_utils.py b/tests/entrypoints/openai/parser/test_harmony_utils.py index d2985264e0c..0027c2763fa 100644 --- a/tests/entrypoints/openai/parser/test_harmony_utils.py +++ b/tests/entrypoints/openai/parser/test_harmony_utils.py @@ -11,12 +11,10 @@ from vllm.entrypoints.openai.parser.harmony_utils import ( auto_drop_analysis_messages, create_tool_definition, extract_function_from_recipient, - get_encoding, get_system_message, has_custom_tools, is_function_recipient, parse_chat_input_to_harmony_message, - parse_chat_output, ) from vllm.entrypoints.openai.responses.harmony import ( response_input_to_harmony, @@ -941,110 +939,6 @@ class TestAutoDropAnalysisMessages: assert cleaned_messages == messages[1:] -class TestParseChatOutput: - def test_parse_chat_output_interrupted_first_message(self) -> None: - harmony_str = "<|channel|>final<|message|>I'm in the middle of answering" - token_ids = get_encoding().encode(harmony_str, allowed_special="all") - reasoning, final_content, _ = parse_chat_output(token_ids) - assert reasoning is None - assert final_content == "I'm in the middle of answering" - - def test_parse_chat_output_interrupted_reasoning_first_message(self) -> None: - harmony_str = "<|channel|>analysis<|message|>I'm in the middle of thinking" - token_ids = get_encoding().encode(harmony_str, allowed_special="all") - reasoning, final_content, _ = parse_chat_output(token_ids) - assert reasoning == "I'm in the middle of thinking" - assert final_content is None - - def test_parse_chat_output_complete_reasoning_interrupted_content(self) -> None: - harmony_str = ( - "<|channel|>analysis<|message|>I'm thinking.<|end|>" - "<|start|>assistant<|channel|>final" - "<|message|>I'm in the middle of answering" - ) - token_ids = get_encoding().encode(harmony_str, allowed_special="all") - reasoning, final_content, _ = parse_chat_output(token_ids) - assert reasoning == "I'm thinking." - assert final_content == "I'm in the middle of answering" - - def test_parse_chat_output_complete_content(self) -> None: - harmony_str = "<|channel|>final<|message|>The answer is 4.<|end|>" - token_ids = get_encoding().encode(harmony_str, allowed_special="all") - reasoning, final_content, _ = parse_chat_output(token_ids) - assert reasoning is None - assert final_content == "The answer is 4." - - def test_parse_chat_output_complete_commentary(self) -> None: - harmony_str = ( - "<|channel|>commentary<|message|>I need to call some tools.<|end|>" - ) - token_ids = get_encoding().encode(harmony_str, allowed_special="all") - reasoning, final_content, _ = parse_chat_output(token_ids) - assert reasoning is None - assert final_content == "I need to call some tools." - - def test_parse_chat_output_complete_reasoning(self) -> None: - harmony_str = ( - "<|channel|>analysis<|message|>I've thought hard about this.<|end|>" - ) - token_ids = get_encoding().encode(harmony_str, allowed_special="all") - reasoning, final_content, _ = parse_chat_output(token_ids) - assert reasoning == "I've thought hard about this." - assert final_content is None - - def test_parse_chat_output_complete_reasoning_and_content(self) -> None: - harmony_str = ( - "<|channel|>analysis<|message|>I've thought hard about this.<|end|>" - "<|start|>assistant<|channel|>final<|message|>The answer is 4.<|end|>" - ) - token_ids = get_encoding().encode(harmony_str, allowed_special="all") - reasoning, final_content, _ = parse_chat_output(token_ids) - assert reasoning == "I've thought hard about this." - assert final_content == "The answer is 4." - - def test_parse_chat_output_commentary_with_recipient_excluded(self) -> None: - """Commentary with a recipient (tool call) should not appear in - final_content — those are handled separately by the tool parser. - - The first message is a preamble (visible), the second is a tool - call (excluded). Only the preamble should appear in final_content. - """ - harmony_str = ( - "<|channel|>commentary" - "<|message|>Let me check the weather.<|end|>" - "<|start|>assistant to=functions.get_weather" - "<|channel|>commentary" - '<|message|>{"location": "SF"}<|end|>' - ) - token_ids = get_encoding().encode(harmony_str, allowed_special="all") - reasoning, final_content, _ = parse_chat_output(token_ids) - assert reasoning is None - assert final_content == "Let me check the weather." - - def test_parse_chat_output_interrupted_preamble(self) -> None: - """Partial/interrupted preamble (commentary without recipient) should - appear in final_content, not reasoning.""" - harmony_str = "<|channel|>commentary<|message|>I'll search for that" - token_ids = get_encoding().encode(harmony_str, allowed_special="all") - reasoning, final_content, _ = parse_chat_output(token_ids) - assert reasoning is None - assert final_content == "I'll search for that" - - def test_parse_chat_output_preamble_then_final(self) -> None: - """Preamble followed by a final message should both appear in - final_content, joined by newline.""" - harmony_str = ( - "<|channel|>commentary" - "<|message|>Let me look that up.<|end|>" - "<|start|>assistant<|channel|>final" - "<|message|>The answer is 42.<|end|>" - ) - token_ids = get_encoding().encode(harmony_str, allowed_special="all") - reasoning, final_content, _ = parse_chat_output(token_ids) - assert reasoning is None - assert final_content == "Let me look that up.\nThe answer is 42." - - def test_has_custom_tools() -> None: assert not has_custom_tools(set()) assert not has_custom_tools({"web_search_preview", "code_interpreter", "container"}) diff --git a/tests/parser/test_harmony.py b/tests/parser/test_harmony.py new file mode 100644 index 00000000000..98687b08edd --- /dev/null +++ b/tests/parser/test_harmony.py @@ -0,0 +1,452 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import json +from collections.abc import Sequence + +import pytest +from openai_harmony import ( + Conversation, + Message, + RenderConversationConfig, + Role, +) +from transformers import AutoTokenizer + +from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest +from vllm.entrypoints.openai.engine.protocol import FunctionCall +from vllm.entrypoints.openai.parser.harmony_utils import ( + get_encoding, +) +from vllm.parser.harmony import HarmonyParser +from vllm.parser.parser_manager import ParserManager + +REASONING_MODEL_NAME = "openai/gpt-oss-20b" + + +@pytest.fixture(scope="module") +def gpt_oss_tokenizer(): + return AutoTokenizer.from_pretrained(REASONING_MODEL_NAME) + + +@pytest.fixture +def harmony_parser(gpt_oss_tokenizer): + parser_cls = ParserManager.get_parser( + tool_parser_name="openai", + reasoning_parser_name="openai_gptoss", + enable_auto_tools=True, + model_name=REASONING_MODEL_NAME, + is_harmony=True, + ) + assert parser_cls is HarmonyParser + return parser_cls(gpt_oss_tokenizer) + + +@pytest.fixture +def chat_request(): + return ChatCompletionRequest( + model="openai/gpt-oss-20b", + messages=[{"role": "user", "content": "Hello"}], + ) + + +def encode_output(harmony_str: str) -> list[int]: + return get_encoding().encode(harmony_str, allowed_special="all") + + +def assistant(content: str, channel: str) -> Message: + return Message.from_role_and_content(Role.ASSISTANT, content).with_channel(channel) + + +def tool_call( + recipient: str, + content: str, + channel: str = "commentary", + content_type: str | None = "json", +) -> Message: + message = assistant(content, channel).with_recipient(recipient) + return message if content_type is None else message.with_content_type(content_type) + + +def get_model_output_tokens( + prompt_messages: Sequence[Message], + response_messages: Sequence[Message], +) -> list[int]: + enc = get_encoding() + # Keep analysis messages when synthesizing model-output-only token sequences + # for parser tests; the default render path drops them after a later final turn. + config = RenderConversationConfig(auto_drop_analysis=False) + prompt_ids = enc.render_conversation_for_completion( + Conversation.from_messages(list(prompt_messages)), + Role.ASSISTANT, + config=config, + ) + full_ids = enc.render_conversation_for_completion( + Conversation.from_messages([*prompt_messages, *response_messages]), + Role.ASSISTANT, + config=config, + ) + assert full_ids[: len(prompt_ids)] == prompt_ids + return full_ids[len(prompt_ids) :] + + +def get_text(msg: Message) -> str: + return msg.content[0].text if msg.content else "" + + +def visible_segments(result) -> list[tuple[str | None, str | None, str]]: + return [ + (segment.channel, segment.recipient, segment.delta) + for segment in result.segments + if not segment.is_boundary and segment.delta + ] + + +def tool_call_tuples(tool_calls: list[FunctionCall] | None) -> list[tuple[str, str]]: + return [] if tool_calls is None else [(tc.name, tc.arguments) for tc in tool_calls] + + +class TestParse: + # Rendered conversation outputs. + + def test_reasoning_only(self, harmony_parser, chat_request): + prompt = [Message.from_role_and_content(Role.USER, "Why?")] + response = [assistant("This is reasoning", "analysis")] + + reasoning, content, tool_calls = harmony_parser.parse( + "", + chat_request, + model_output_token_ids=get_model_output_tokens(prompt, response), + ) + + assert reasoning == "This is reasoning" + assert content is None + assert tool_calls is None + + def test_content_only(self, harmony_parser, chat_request): + prompt = [Message.from_role_and_content(Role.USER, "Hello")] + response = [assistant("This is a test", "final")] + + reasoning, content, tool_calls = harmony_parser.parse( + "", + chat_request, + model_output_token_ids=get_model_output_tokens(prompt, response), + ) + + assert reasoning is None + assert content == "This is a test" + assert tool_calls is None + + def test_reasoning_and_content(self, harmony_parser, chat_request): + prompt = [Message.from_role_and_content(Role.USER, "What is 2+2?")] + response = [ + assistant("I should think first.", "analysis"), + assistant("The answer is 4.", "final"), + ] + + reasoning, content, tool_calls = harmony_parser.parse( + "", + chat_request, + model_output_token_ids=get_model_output_tokens(prompt, response), + ) + + assert reasoning == "I should think first." + assert content == "The answer is 4." + assert tool_calls is None + + @pytest.mark.parametrize( + "tool_args", + [ + '{"location": "Tokyo"}', + '{\n"location": "Tokyo"\n}', + ], + ) + @pytest.mark.parametrize("tool_channel", ["commentary", "analysis"]) + def test_single_tool_call( + self, harmony_parser, chat_request, tool_args, tool_channel + ): + prompt = [ + Message.from_role_and_content(Role.USER, "What is the weather in Tokyo?") + ] + response = [tool_call("functions.get_current_weather", tool_args, tool_channel)] + + reasoning, content, tool_calls = harmony_parser.parse( + "", + chat_request, + model_output_token_ids=get_model_output_tokens(prompt, response), + ) + + assert reasoning is None + assert content is None + assert tool_call_tuples(tool_calls) == [ + ("get_current_weather", json.dumps({"location": "Tokyo"})) + ] + + def test_multiple_tool_calls_varied_formats(self, harmony_parser, chat_request): + prompt = [ + Message.from_role_and_content( + Role.USER, "What is the weather in Tokyo based on where I'm at?" + ) + ] + response = [ + tool_call("functions.get_current_weather", '{"location": "Tokyo"}'), + tool_call("functions.get_user_location", '{"location": "Tokyo"}'), + tool_call( + "functions.no_content_type", + '{"location": "Tokyo"}', + content_type=None, + ), + tool_call("functions.not_json_no_content_type", "foo", content_type=None), + tool_call("functions.empty_args", "{}"), + tool_call("functions.no_args", ""), + ] + + _, content, tool_calls = harmony_parser.parse( + "", + chat_request, + model_output_token_ids=get_model_output_tokens(prompt, response), + ) + + assert content is None + assert tool_call_tuples(tool_calls) == [ + ("get_current_weather", json.dumps({"location": "Tokyo"})), + ("get_user_location", json.dumps({"location": "Tokyo"})), + ("no_content_type", json.dumps({"location": "Tokyo"})), + ("not_json_no_content_type", "foo"), + ("empty_args", json.dumps({})), + ("no_args", ""), + ] + + def test_tool_call_bare_recipient(self, harmony_parser, chat_request): + prompt = [Message.from_role_and_content(Role.USER, "Weather?")] + response = [tool_call("get_current_weather", '{"location": "Tokyo"}')] + + _, _, tool_calls = harmony_parser.parse( + "", + chat_request, + model_output_token_ids=get_model_output_tokens(prompt, response), + ) + + assert tool_call_tuples(tool_calls) == [ + ("get_current_weather", json.dumps({"location": "Tokyo"})) + ] + + def test_multiple_tool_calls_bare_recipients(self, harmony_parser, chat_request): + prompt = [Message.from_role_and_content(Role.USER, "Use both tools.")] + response = [ + tool_call("get_current_weather", '{"location": "Tokyo"}'), + tool_call("get_user_location", "{}"), + ] + + _, _, tool_calls = harmony_parser.parse( + "", + chat_request, + model_output_token_ids=get_model_output_tokens(prompt, response), + ) + + assert tool_call_tuples(tool_calls) == [ + ("get_current_weather", json.dumps({"location": "Tokyo"})), + ("get_user_location", json.dumps({})), + ] + + def test_assistant_recipient_not_tool(self, harmony_parser, chat_request): + prompt = [Message.from_role_and_content(Role.USER, "Hello")] + response = [ + tool_call("assistant", "Some tool response", content_type=None), + assistant("Here is the answer", "final"), + ] + + reasoning, content, tool_calls = harmony_parser.parse( + "", + chat_request, + model_output_token_ids=get_model_output_tokens(prompt, response), + ) + + assert reasoning is None + assert content == "Here is the answer" + assert tool_calls is None + + def test_tool_call_dotted_name(self, harmony_parser, chat_request): + prompt = [Message.from_role_and_content(Role.USER, "Compute 2+3")] + response = [tool_call("math.sum", '{"a": 2, "b": 3}')] + + _, _, tool_calls = harmony_parser.parse( + "", + chat_request, + model_output_token_ids=get_model_output_tokens(prompt, response), + ) + + assert tool_call_tuples(tool_calls) == [ + ("math.sum", json.dumps({"a": 2, "b": 3})) + ] + + def test_tool_calls_with_final_content(self, harmony_parser, chat_request): + prompt = [Message.from_role_and_content(Role.USER, "What is the weather?")] + response = [ + assistant("User asked about the weather.", "analysis"), + tool_call("functions.get_current_weather", '{"location": "Tokyo"}'), + assistant("This tool call will get the weather.", "final"), + ] + + reasoning, content, tool_calls = harmony_parser.parse( + "", + chat_request, + model_output_token_ids=get_model_output_tokens(prompt, response), + ) + + assert reasoning == "User asked about the weather." + assert content == "This tool call will get the weather." + assert tool_call_tuples(tool_calls) == [ + ("get_current_weather", json.dumps({"location": "Tokyo"})) + ] + + # Raw/truncated Harmony output streams. + + def test_interrupted_first_message(self, harmony_parser, chat_request): + reasoning, content, tool_calls = harmony_parser.parse( + "", + chat_request, + model_output_token_ids=encode_output( + "<|channel|>final<|message|>I'm in the middle of answering" + ), + ) + + assert reasoning is None + assert content == "I'm in the middle of answering" + assert tool_calls is None + + def test_interrupted_reasoning_first_message(self, harmony_parser, chat_request): + reasoning, content, tool_calls = harmony_parser.parse( + "", + chat_request, + model_output_token_ids=encode_output( + "<|channel|>analysis<|message|>I'm in the middle of thinking" + ), + ) + + assert reasoning == "I'm in the middle of thinking" + assert content is None + assert tool_calls is None + + def test_truncated_output(self, harmony_parser, chat_request): + reasoning, content, tool_calls = harmony_parser.parse( + "", + chat_request, + model_output_token_ids=encode_output( + "<|channel|>analysis<|message|>I'm thinking.<|end|>" + "<|start|>assistant<|channel|>final<|message|>" + "I'm in the middle of answering" + ), + ) + + assert reasoning == "I'm thinking." + assert content == "I'm in the middle of answering" + assert tool_calls is None + + @pytest.mark.parametrize( + ("harmony_str", "expected_content"), + [ + ( + "<|channel|>commentary<|message|>I'll search for that", + "I'll search for that", + ), + ( + "<|channel|>commentary<|message|>Let me look that up.<|end|>" + "<|start|>assistant<|channel|>final<|message|>The answer is 42.<|end|>", + "Let me look that up.\nThe answer is 42.", + ), + ], + ) + def test_commentary_preambles( + self, + harmony_parser, + chat_request, + harmony_str, + expected_content, + ): + reasoning, content, tool_calls = harmony_parser.parse( + "", + chat_request, + model_output_token_ids=encode_output(harmony_str), + ) + + assert reasoning is None + assert content == expected_content + assert tool_calls is None + + def test_commentary_with_recipient_excluded(self, harmony_parser, chat_request): + reasoning, content, tool_calls = harmony_parser.parse( + "", + chat_request, + model_output_token_ids=encode_output( + "<|channel|>commentary" + "<|message|>Let me check the weather.<|end|>" + "<|start|>assistant to=functions.get_weather" + "<|channel|>commentary" + '<|message|>{"location": "SF"}<|end|>' + ), + ) + + assert reasoning is None + assert content == "Let me check the weather." + assert tool_call_tuples(tool_calls) == [ + ("get_weather", json.dumps({"location": "SF"})) + ] + + +class TestProcessChunk: + def test_empty(self, harmony_parser): + result = harmony_parser.process_chunk([]) + assert result.segments == [] + assert result.reasoning_token_count == 0 + + def test_single_channel(self, harmony_parser): + result = harmony_parser.process_chunk( + encode_output("<|channel|>final<|message|>Hello") + ) + + assert visible_segments(result) == [("final", None, "Hello")] + + def test_cross_channel(self, harmony_parser): + result = harmony_parser.process_chunk( + encode_output( + "<|channel|>analysis<|message|>Think<|end|>" + "<|start|>assistant<|channel|>final<|message|>Answer" + ) + ) + + assert visible_segments(result) == [ + ("analysis", None, "Think"), + ("final", None, "Answer"), + ] + + def test_boundary_detection(self, harmony_parser): + result = harmony_parser.process_chunk( + encode_output("<|channel|>final<|message|>Done<|end|>") + ) + + boundary_segments = [ + segment for segment in result.segments if segment.is_boundary + ] + assert len(boundary_segments) == 1 + assert boundary_segments[0].completed_message is not None + assert boundary_segments[0].completed_message.channel == "final" + assert get_text(boundary_segments[0].completed_message) == "Done" + + def test_multi_boundary(self, harmony_parser): + result = harmony_parser.process_chunk( + encode_output( + "<|channel|>analysis<|message|>One<|end|>" + "<|start|>assistant<|channel|>final<|message|>Two<|end|>" + ) + ) + + boundary_segments = [ + segment for segment in result.segments if segment.is_boundary + ] + assert [ + get_text(segment.completed_message) for segment in boundary_segments + ] == [ + "One", + "Two", + ] diff --git a/tests/tool_parsers/test_openai_tool_parser.py b/tests/tool_parsers/test_openai_tool_parser.py deleted file mode 100644 index 843fbca621f..00000000000 --- a/tests/tool_parsers/test_openai_tool_parser.py +++ /dev/null @@ -1,415 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import json - -import pytest -from openai_harmony import ( - Conversation, - DeveloperContent, - HarmonyEncodingName, - Message, - Role, - SystemContent, - load_harmony_encoding, -) - -from vllm.entrypoints.openai.engine.protocol import FunctionCall, ToolCall -from vllm.tokenizers import get_tokenizer -from vllm.tool_parsers.openai_tool_parser import OpenAIToolParser - -MODEL = "gpt2" - - -@pytest.fixture(scope="module") -def openai_tokenizer(): - # The parser does not use the tokenizer, but the constructor requires it. - return get_tokenizer(MODEL) - - -@pytest.fixture -def openai_tool_parser(openai_tokenizer): - return OpenAIToolParser(openai_tokenizer) - - -@pytest.fixture(scope="module") -def harmony_encoding(): - return load_harmony_encoding(HarmonyEncodingName.HARMONY_GPT_OSS) - - -def assert_tool_calls( - actual_tool_calls: list[ToolCall], - expected_tool_calls: list[ToolCall], -): - assert len(actual_tool_calls) == len(expected_tool_calls) - - for actual_tool_call, expected_tool_call in zip( - actual_tool_calls, expected_tool_calls - ): - assert isinstance(actual_tool_call.id, str) - assert len(actual_tool_call.id) > 16 # Default from protocol.py - assert actual_tool_call.type == "function" - assert actual_tool_call.function == expected_tool_call.function - - -def test_extract_tool_calls_no_tools(openai_tool_parser, harmony_encoding): - convo = Conversation.from_messages( - [ - Message.from_role_and_content( - Role.SYSTEM, - SystemContent.new(), - ), - Message.from_role_and_content( - Role.DEVELOPER, - DeveloperContent.new().with_instructions("Talk like a pirate!"), - ), - Message.from_role_and_content(Role.USER, "Arrr, how be you?"), - Message.from_role_and_content( - Role.ASSISTANT, "This is a test" - ).with_channel("final"), - ] - ) - token_ids = harmony_encoding.render_conversation_for_completion( - convo, Role.ASSISTANT - ) - extracted_info = openai_tool_parser.extract_tool_calls( - "", - request=None, - token_ids=token_ids, - ) - assert not extracted_info.tools_called - assert extracted_info.tool_calls == [] - assert extracted_info.content == "This is a test" - - -@pytest.mark.parametrize( - "tool_args", - [ - '{"location": "Tokyo"}', - '{\n"location": "Tokyo"\n}', - ], -) -def test_extract_tool_calls_single_tool( - openai_tool_parser, harmony_encoding, tool_args -): - convo = Conversation.from_messages( - [ - Message.from_role_and_content(Role.USER, "What is the weather in Tokyo?"), - Message.from_role_and_content( - Role.ASSISTANT, - 'User asks: "What is the weather in Tokyo?" We need to use get_current_weather tool.', # noqa: E501 - ).with_channel("analysis"), - Message.from_role_and_content(Role.ASSISTANT, tool_args) - .with_channel("commentary") - .with_recipient("functions.get_current_weather") - .with_content_type("json"), - ] - ) - token_ids = harmony_encoding.render_conversation_for_completion( - convo, Role.ASSISTANT - ) - - extracted_info = openai_tool_parser.extract_tool_calls( - "", - request=None, - token_ids=token_ids, - ) - assert extracted_info.tools_called - expected_tool_calls = [ - ToolCall( - function=FunctionCall( - name="get_current_weather", - arguments=json.dumps({"location": "Tokyo"}), - ) - ) - ] - assert_tool_calls(extracted_info.tool_calls, expected_tool_calls) - assert extracted_info.content is None - - -def test_extract_tool_calls_multiple_tools( - openai_tool_parser, - harmony_encoding, -): - convo = Conversation.from_messages( - [ - Message.from_role_and_content( - Role.USER, "What is the weather in Tokyo based on where I'm at?" - ), - Message.from_role_and_content( - Role.ASSISTANT, - 'User asks: "What is the weather in Tokyo?" based on their location. We need to use get_current_weather tool and get_user_location tool.', # noqa: E501 - ).with_channel("analysis"), - Message.from_role_and_content(Role.ASSISTANT, '{"location": "Tokyo"}') - .with_channel("commentary") - .with_recipient("functions.get_current_weather") - .with_content_type("json"), - Message.from_role_and_content(Role.ASSISTANT, '{"location": "Tokyo"}') - .with_channel("commentary") - .with_recipient("functions.get_user_location") - .with_content_type("json"), - Message.from_role_and_content(Role.ASSISTANT, '{"location": "Tokyo"}') - .with_channel("commentary") - .with_recipient("functions.no_content_type"), - Message.from_role_and_content(Role.ASSISTANT, "foo") - .with_channel("commentary") - .with_recipient("functions.not_json_no_content_type"), - Message.from_role_and_content(Role.ASSISTANT, "{}") - .with_channel("commentary") - .with_recipient("functions.empty_args") - .with_content_type("json"), - Message.from_role_and_content(Role.ASSISTANT, "") - .with_channel("commentary") - .with_recipient("functions.no_args") - .with_content_type("json"), - ] - ) - token_ids = harmony_encoding.render_conversation_for_completion( - convo, - Role.ASSISTANT, - ) - - extracted_info = openai_tool_parser.extract_tool_calls( - "", - request=None, - token_ids=token_ids, - ) - assert extracted_info.tools_called - expected_tool_calls = [ - ToolCall( - function=FunctionCall( - name="get_current_weather", - arguments=json.dumps({"location": "Tokyo"}), - ) - ), - ToolCall( - function=FunctionCall( - name="get_user_location", - arguments=json.dumps({"location": "Tokyo"}), - ) - ), - ToolCall( - function=FunctionCall( - name="no_content_type", - arguments=json.dumps({"location": "Tokyo"}), - ) - ), - ToolCall( - function=FunctionCall( - name="not_json_no_content_type", - arguments="foo", - ) - ), - ToolCall( - function=FunctionCall( - name="empty_args", - arguments=json.dumps({}), - ) - ), - ToolCall( - function=FunctionCall( - name="no_args", - arguments="", - ) - ), - ] - assert_tool_calls(extracted_info.tool_calls, expected_tool_calls) - assert extracted_info.content is None - - -def test_extract_tool_calls_bare_function_name( - openai_tool_parser, - harmony_encoding, -): - convo = Conversation.from_messages( - [ - Message.from_role_and_content(Role.USER, "What is the weather in Tokyo?"), - Message.from_role_and_content( - Role.ASSISTANT, - "We need to use get_current_weather tool.", - ).with_channel("analysis"), - Message.from_role_and_content(Role.ASSISTANT, '{"location": "Tokyo"}') - .with_channel("commentary") - .with_recipient("get_current_weather") - .with_content_type("json"), - ] - ) - token_ids = harmony_encoding.render_conversation_for_completion( - convo, Role.ASSISTANT - ) - - extracted_info = openai_tool_parser.extract_tool_calls( - "", - request=None, - token_ids=token_ids, - ) - assert extracted_info.tools_called - expected_tool_calls = [ - ToolCall( - function=FunctionCall( - name="get_current_weather", - arguments=json.dumps({"location": "Tokyo"}), - ) - ) - ] - assert_tool_calls(extracted_info.tool_calls, expected_tool_calls) - assert extracted_info.content is None - - -def test_extract_tool_calls_bare_function_name_multiple( - openai_tool_parser, - harmony_encoding, -): - convo = Conversation.from_messages( - [ - Message.from_role_and_content( - Role.USER, "What is the weather in Tokyo based on where I'm at?" - ), - Message.from_role_and_content( - Role.ASSISTANT, - "We need to use both tools.", - ).with_channel("analysis"), - Message.from_role_and_content(Role.ASSISTANT, '{"location": "Tokyo"}') - .with_channel("commentary") - .with_recipient("get_current_weather") - .with_content_type("json"), - Message.from_role_and_content(Role.ASSISTANT, "{}") - .with_channel("commentary") - .with_recipient("get_user_location") - .with_content_type("json"), - ] - ) - token_ids = harmony_encoding.render_conversation_for_completion( - convo, Role.ASSISTANT - ) - - extracted_info = openai_tool_parser.extract_tool_calls( - "", - request=None, - token_ids=token_ids, - ) - assert extracted_info.tools_called - expected_tool_calls = [ - ToolCall( - function=FunctionCall( - name="get_current_weather", - arguments=json.dumps({"location": "Tokyo"}), - ) - ), - ToolCall( - function=FunctionCall( - name="get_user_location", - arguments=json.dumps({}), - ) - ), - ] - assert_tool_calls(extracted_info.tool_calls, expected_tool_calls) - - -def test_extract_tool_calls_assistant_recipient_ignored( - openai_tool_parser, - harmony_encoding, -): - convo = Conversation.from_messages( - [ - Message.from_role_and_content(Role.USER, "Hello"), - Message.from_role_and_content(Role.ASSISTANT, "Some tool response") - .with_channel("commentary") - .with_recipient("assistant"), - Message.from_role_and_content( - Role.ASSISTANT, "Here is the answer" - ).with_channel("final"), - ] - ) - token_ids = harmony_encoding.render_conversation_for_completion( - convo, Role.ASSISTANT - ) - - extracted_info = openai_tool_parser.extract_tool_calls( - "", - request=None, - token_ids=token_ids, - ) - assert not extracted_info.tools_called - assert extracted_info.tool_calls == [] - assert extracted_info.content == "Here is the answer" - - -def test_extract_tool_calls_dotted_function_name( - openai_tool_parser, - harmony_encoding, -): - convo = Conversation.from_messages( - [ - Message.from_role_and_content(Role.USER, "Compute 2+3"), - Message.from_role_and_content(Role.ASSISTANT, '{"a": 2, "b": 3}') - .with_channel("commentary") - .with_recipient("math.sum") - .with_content_type("json"), - ] - ) - token_ids = harmony_encoding.render_conversation_for_completion( - convo, Role.ASSISTANT - ) - - extracted_info = openai_tool_parser.extract_tool_calls( - "", - request=None, - token_ids=token_ids, - ) - assert extracted_info.tools_called - expected_tool_calls = [ - ToolCall( - function=FunctionCall( - name="math.sum", - arguments=json.dumps({"a": 2, "b": 3}), - ) - ) - ] - assert_tool_calls(extracted_info.tool_calls, expected_tool_calls) - - -def test_extract_tool_calls_with_content( - openai_tool_parser, - harmony_encoding, -): - final_content = "This tool call will get the weather." - convo = Conversation.from_messages( - [ - Message.from_role_and_content( - Role.USER, "What is the weather in Tokyo based on where I'm at?" - ), - Message.from_role_and_content( - Role.ASSISTANT, - 'User asks: "What is the weather in Tokyo?" based on their location. We need to use get_current_weather tool and get_user_location tool.', # noqa: E501 - ).with_channel("analysis"), - Message.from_role_and_content(Role.ASSISTANT, '{"location": "Tokyo"}') - .with_channel("commentary") - .with_recipient("functions.get_current_weather") - .with_content_type("json"), - Message.from_role_and_content(Role.ASSISTANT, final_content).with_channel( - "final" - ), - ] - ) - token_ids = harmony_encoding.render_conversation_for_completion( - convo, - Role.ASSISTANT, - ) - - extracted_info = openai_tool_parser.extract_tool_calls( - "", - request=None, - token_ids=token_ids, - ) - assert extracted_info.tools_called - expected_tool_calls = [ - ToolCall( - function=FunctionCall( - name="get_current_weather", - arguments=json.dumps({"location": "Tokyo"}), - ) - ), - ] - assert_tool_calls(extracted_info.tool_calls, expected_tool_calls) - assert extracted_info.content == final_content diff --git a/vllm/entrypoints/openai/chat_completion/serving.py b/vllm/entrypoints/openai/chat_completion/serving.py index 2da89917a8d..4924ceb8b5d 100644 --- a/vllm/entrypoints/openai/chat_completion/serving.py +++ b/vllm/entrypoints/openai/chat_completion/serving.py @@ -54,7 +54,6 @@ from vllm.entrypoints.openai.engine.serving import ( from vllm.entrypoints.openai.models.serving import OpenAIServingModels from vllm.entrypoints.openai.parser.harmony_utils import ( get_streamable_parser_for_assistant, - parse_chat_output, ) from vllm.entrypoints.serve.utils.api_utils import get_max_tokens, should_include_usage from vllm.entrypoints.serve.utils.request_logger import RequestLogger @@ -135,6 +134,7 @@ class OpenAIServingChat(OpenAIServing): reasoning_parser_name=reasoning_parser, enable_auto_tools=enable_auto_tools, model_name=self.model_config.model, + is_harmony=self.model_config.hf_config.model_type == "gpt_oss", ) if ( is_mistral_tool_parser(self.tool_parser) @@ -359,14 +359,6 @@ class OpenAIServingChat(OpenAIServing): assert len(generators) == 1 (result_generator,) = generators - parser: Parser | None = None - if self.parser_cls is not None: - parser = self.parser_cls( - tokenizer, - request.tools, - chat_template_kwargs=chat_template_kwargs, - ) - if request.stream: return self.chat_completion_stream_generator( request, @@ -387,7 +379,7 @@ class OpenAIServingChat(OpenAIServing): conversation, tokenizer, request_metadata, - parser, + chat_template_kwargs=chat_template_kwargs, ) def get_chat_request_role(self, request: ChatCompletionRequest) -> str: @@ -840,7 +832,7 @@ class OpenAIServingChat(OpenAIServing): conversation: list[ConversationMessage], tokenizer: TokenizerLike, request_metadata: RequestResponseMetadata, - parser: Parser | None = None, + chat_template_kwargs: dict[str, Any] | None = None, ) -> ErrorResponse | ChatCompletionResponse: created_time = int(time.time()) final_res: RequestOutput | None = None @@ -871,7 +863,6 @@ class OpenAIServingChat(OpenAIServing): self._raise_if_error(output.finish_reason, request_id) token_ids = output.token_ids out_logprobs = output.logprobs - tool_call_info = None if request.logprobs and request.top_logprobs is not None: assert out_logprobs is not None, "Did not output logprobs" @@ -885,75 +876,20 @@ class OpenAIServingChat(OpenAIServing): else: logprobs = None - if self.use_harmony: - reasoning, content, _ = parse_chat_output(token_ids) - if not request.include_reasoning: - reasoning = None - - if self.tool_parser is not None: - if tokenizer is None: - raise ValueError( - "Tokenizer not available when `skip_tokenizer_init=True`" - ) - - tool_parser = self.tool_parser(tokenizer, request.tools) - # NOTE: We use token_ids for openai tool parser - tool_call_info = tool_parser.extract_tool_calls( - "", - request=request, - token_ids=token_ids, # type: ignore - ) - content = tool_call_info.content - message = ChatMessage( - role=role, - reasoning=reasoning, - content=content, - tool_calls=tool_call_info.tool_calls, - ) - else: - message = ChatMessage( - role=role, - reasoning=reasoning, - content=content, - ) - - # Encode routed_experts for transport. JSON can't carry raw - # bytes, so we write the ndarray as a ``.npy`` byte stream - # and base64-encode it. ``pybase64`` is ~3x faster than the - # stdlib ``base64`` on large payloads thanks to SIMD. - routed_experts_b64 = None - if output.routed_experts is not None: - buf = io.BytesIO() - np.save(buf, output.routed_experts) - routed_experts_b64 = base64.b64encode(buf.getvalue()).decode( - "ascii" - ) - - choice_data = ChatCompletionResponseChoice( - index=output.index, - message=message, - logprobs=logprobs, - finish_reason=( - "tool_calls" - if (tool_call_info is not None and tool_call_info.tools_called) - else output.finish_reason - if output.finish_reason - else "stop" - ), - stop_reason=output.stop_reason, - token_ids=( - as_list(output.token_ids) if request.return_token_ids else None - ), - routed_experts=routed_experts_b64, + parser: Parser | None = None + if self.parser_cls is not None: + parser = self.parser_cls( + tokenizer, + request.tools, + chat_template_kwargs=chat_template_kwargs, ) - choices.append(choice_data) - continue if parser is not None: reasoning, content, tool_calls = parser.parse( output.text, request, enable_auto_tools=self.enable_auto_tools, + model_output_token_ids=token_ids, ) if not request.include_reasoning: reasoning = None diff --git a/vllm/entrypoints/openai/parser/harmony_utils.py b/vllm/entrypoints/openai/parser/harmony_utils.py index 771faabe609..82316efb86d 100644 --- a/vllm/entrypoints/openai/parser/harmony_utils.py +++ b/vllm/entrypoints/openai/parser/harmony_utils.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import datetime -from collections.abc import Iterable, Sequence +from collections.abc import Sequence from typing import Any from openai.types.responses.tool import Tool @@ -456,65 +456,3 @@ def render_for_completion(messages: list[Message]) -> list[int]: def get_streamable_parser_for_assistant() -> StreamableParser: return StreamableParser(get_encoding(), role=Role.ASSISTANT) - - -def parse_output_into_messages(token_ids: Iterable[int]) -> StreamableParser: - parser = get_streamable_parser_for_assistant() - for token_id in token_ids: - parser.process(token_id) - return parser - - -def parse_chat_output( - token_ids: Sequence[int], -) -> tuple[str | None, str | None, bool]: - """ - Parse the output of a Harmony chat completion into reasoning and final content. - Note that when the `openai` tool parser is used, serving_chat only uses this - for the reasoning content and gets the final content from the tool call parser. - - When the `openai` tool parser is not enabled, or when `GptOssReasoningParser` is - in use,this needs to return the final content without any tool calls parsed. - - Empty reasoning or final content is returned as None instead of an empty string. - """ - parser = parse_output_into_messages(token_ids) - output_msgs = parser.messages - is_tool_call = False # TODO: update this when tool call is supported - - # Get completed messages from the parser - # - analysis channel: hidden reasoning - # - commentary channel without recipient (preambles): visible to user - # - final channel: visible to user - # - commentary with recipient (tool calls): handled separately by tool parser - reasoning_texts = [ - msg.content[0].text for msg in output_msgs if msg.channel == "analysis" - ] - final_texts = [ - msg.content[0].text - for msg in output_msgs - if msg.channel == "final" or (msg.channel == "commentary" and not msg.recipient) - ] - - # Extract partial messages from the parser - if parser.current_channel == "analysis" and parser.current_content: - reasoning_texts.append(parser.current_content) - elif parser.current_channel == "final" and parser.current_content: - final_texts.append(parser.current_content) - elif ( - parser.current_channel == "commentary" - and not parser.current_recipient - and parser.current_content - ): - # Preambles (commentary without recipient) are visible to user - final_texts.append(parser.current_content) - - # Flatten multiple messages into a single string - reasoning: str | None = "\n".join(reasoning_texts) - final_content: str | None = "\n".join(final_texts) - - # Return None instead of empty string since existing callers check for None - reasoning = reasoning or None - final_content = final_content or None - - return reasoning, final_content, is_tool_call diff --git a/vllm/entrypoints/openai/responses/serving.py b/vllm/entrypoints/openai/responses/serving.py index 51831f60835..69fbcce818f 100644 --- a/vllm/entrypoints/openai/responses/serving.py +++ b/vllm/entrypoints/openai/responses/serving.py @@ -191,6 +191,7 @@ class OpenAIServingResponses(OpenAIServing): reasoning_parser_name=reasoning_parser, enable_auto_tools=enable_auto_tools, model_name=self.model_config.model, + is_harmony=self.model_config.hf_config.model_type == "gpt_oss", ) self.enable_prompt_tokens_details = enable_prompt_tokens_details self.enable_force_include_usage = enable_force_include_usage diff --git a/vllm/parser/__init__.py b/vllm/parser/__init__.py index de815b2e1fd..e13c2ece9f0 100644 --- a/vllm/parser/__init__.py +++ b/vllm/parser/__init__.py @@ -5,10 +5,12 @@ from vllm.parser.abstract_parser import ( DelegatingParser, Parser, ) +from vllm.parser.harmony import HarmonyParser from vllm.parser.parser_manager import ParserManager __all__ = [ "Parser", "DelegatingParser", + "HarmonyParser", "ParserManager", ] diff --git a/vllm/parser/abstract_parser.py b/vllm/parser/abstract_parser.py index 48db01c14e0..4fe7b7ec4d5 100644 --- a/vllm/parser/abstract_parser.py +++ b/vllm/parser/abstract_parser.py @@ -282,6 +282,7 @@ class Parser: model_output: str, request: ChatCompletionRequest | ResponsesRequest, enable_auto_tools: bool = False, + model_output_token_ids: Sequence[int] = (), ) -> tuple[str | None, str | None, list[FunctionCall] | None]: """Parse a complete model output, extracting reasoning and tool calls. @@ -289,6 +290,7 @@ class Parser: model_output: The complete model-generated string. request: The request object used to generate the output. enable_auto_tools: Whether to enable automatic tool call parsing. + model_output_token_ids: The generated raw output token IDs. Returns: A tuple of (reasoning, content, tool_calls). @@ -642,6 +644,7 @@ class DelegatingParser(Parser): model_output: str, request: ChatCompletionRequest | ResponsesRequest, enable_auto_tools: bool = False, + model_output_token_ids: Sequence[int] = (), ) -> tuple[str | None, str | None, list[FunctionCall] | None]: reasoning, content = self.extract_reasoning(model_output, request) tool_calls, content = self._extract_tool_calls( diff --git a/vllm/parser/harmony.py b/vllm/parser/harmony.py new file mode 100644 index 00000000000..c1eb7ea042e --- /dev/null +++ b/vllm/parser/harmony.py @@ -0,0 +1,240 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import json +from collections.abc import Sequence +from dataclasses import dataclass +from enum import Enum, auto +from typing import TYPE_CHECKING, NamedTuple + +from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest +from vllm.entrypoints.openai.engine.protocol import ( + DeltaMessage, + FunctionCall, +) +from vllm.entrypoints.openai.parser.harmony_utils import ( + extract_function_from_recipient, + get_streamable_parser_for_assistant, + is_function_recipient, +) +from vllm.entrypoints.openai.responses.protocol import ResponsesRequest +from vllm.parser.abstract_parser import DelegatingParser +from vllm.reasoning.gptoss_reasoning_parser import GptOssReasoningParser +from vllm.tool_parsers.gptoss_tool_parser import GptOssToolParser + +if TYPE_CHECKING: + from openai_harmony import Message, Role + from openai_harmony import StreamState as HarmonyStreamState + + +class _SegmentType(Enum): + TOOL = auto() + REASONING = auto() + CONTENT = auto() + IGNORE = auto() + + @staticmethod + def from_channel_and_recipient( + channel: str | None, recipient: str | None + ) -> _SegmentType: + if recipient and is_function_recipient(recipient): + return _SegmentType.TOOL + if channel == "analysis": + return _SegmentType.REASONING + if channel == "final" or (channel == "commentary" and recipient is None): + return _SegmentType.CONTENT + return _SegmentType.IGNORE + + +class Segment(NamedTuple): + channel: str | None + recipient: str | None + delta: str + is_boundary: bool = False + completed_message: Message | None = None + + +@dataclass +class ChunkResult: + segments: list[Segment] + reasoning_token_count: int + + +class HarmonyParser(DelegatingParser): + def __init__(self, tokenizer, tools=None, *args, **kwargs): + super().__init__(tokenizer, tools, *args, **kwargs) + + if self._reasoning_parser and not isinstance( + self._reasoning_parser, GptOssReasoningParser + ): + raise ValueError( + "Harmony requires GptOssReasoningParser, " + f"got {self._reasoning_parser.__class__.__name__}." + ) + + if self._tool_parser and not isinstance(self._tool_parser, GptOssToolParser): + raise ValueError( + "Harmony requires GptOssToolParser, " + f"got {self._tool_parser.__class__.__name__}." + ) + + self._harmony_parser = get_streamable_parser_for_assistant() + + @property + def messages(self) -> list[Message]: + return self._harmony_parser.messages + + @property + def state(self) -> HarmonyStreamState: + return self._harmony_parser.state + + @property + def current_role(self) -> Role | None: + return self._harmony_parser.current_role + + @property + def current_channel(self) -> str | None: + return self._harmony_parser.current_channel + + @property + def current_recipient(self) -> str | None: + return self._harmony_parser.current_recipient + + @property + def current_content(self) -> str: + return self._harmony_parser.current_content + + @property + def current_content_type(self) -> str | None: + return self._harmony_parser.current_content_type + + def parse( + self, + model_output: str, + request: ChatCompletionRequest | ResponsesRequest, + enable_auto_tools: bool = False, + model_output_token_ids: Sequence[int] = (), + ) -> tuple[str | None, str | None, list[FunctionCall] | None]: + """Parse Harmony output from token IDs. + + Tool calls are always extracted regardless of ``enable_auto_tools``. + Callers must decide whether to surface them. + """ + result = self.process_chunk(model_output_token_ids) + + reasoning_parts: list[str] = [] + content_parts: list[str] = [] + tool_calls: list[FunctionCall] = [] + + def _append_parsed_message( + channel: str | None, + recipient: str | None, + text: str, + content_type: str | None = None, + ) -> None: + segment_type = _SegmentType.from_channel_and_recipient(channel, recipient) + match segment_type: + case _SegmentType.REASONING if self.reasoning_parser and text: + reasoning_parts.append(text) + case _SegmentType.CONTENT if text: + content_parts.append(text) + case _SegmentType.TOOL if self.tool_parser: + assert recipient is not None + if content_type is not None and "json" not in content_type: + arguments = text + else: + try: + arguments = json.dumps(json.loads(text)) + except json.JSONDecodeError: + arguments = text + tool_calls.append( + FunctionCall( + name=extract_function_from_recipient(recipient), + arguments=arguments, + ) + ) + + for segment in result.segments: + msg = segment.completed_message + if msg is None: + continue + if msg.author.role != "assistant" or not msg.content: + continue + _append_parsed_message( + channel=msg.channel, + recipient=msg.recipient, + text=msg.content[0].text, + content_type=msg.content_type, + ) + + if ( + self.current_channel is not None + or self.current_recipient is not None + or self.current_content + ): + _append_parsed_message( + channel=self.current_channel, + recipient=self.current_recipient, + text=self.current_content, + content_type=self.current_content_type, + ) + + reasoning = "\n".join(reasoning_parts) or None + content = "\n".join(content_parts) or None + return reasoning, content, tool_calls or None + + def parse_delta( + self, + delta_text: str, + delta_token_ids: list[int], + request: ChatCompletionRequest | ResponsesRequest, + prompt_token_ids: list[int] | None = None, + *, + finished: bool, + ) -> DeltaMessage | None: + raise NotImplementedError( + "HarmonyParser streaming parsing is deferred. " + "Use the existing harmony streaming path." + ) + + def process_chunk(self, token_ids: Sequence[int]) -> ChunkResult: + if not token_ids: + return ChunkResult(segments=[], reasoning_token_count=0) + + from openai_harmony import StreamState + + segments: list[Segment] = [] + reasoning_token_count = 0 + for token_id in token_ids: + self._harmony_parser.process(token_id) + channel = self.current_channel + recipient = self.current_recipient + delta = self._harmony_parser.last_content_delta or "" + completed_message = None + is_boundary = self.state == StreamState.EXPECT_START + if is_boundary and self.messages: + completed_message = self.messages[-1] + + if channel == "analysis" or ( + channel == "commentary" and recipient is not None + ): + reasoning_token_count += 1 + + segments.append( + Segment( + channel=channel, + recipient=recipient, + delta=delta, + is_boundary=is_boundary, + completed_message=completed_message, + ) + ) + + # TODO: Optionally merge and suppress empty Segments + + return ChunkResult( + segments=segments, + reasoning_token_count=reasoning_token_count, + ) diff --git a/vllm/parser/mistral.py b/vllm/parser/mistral.py index c7f557a5a95..52f16136ee3 100644 --- a/vllm/parser/mistral.py +++ b/vllm/parser/mistral.py @@ -3,6 +3,7 @@ from __future__ import annotations +from collections.abc import Sequence from typing import TYPE_CHECKING from vllm.entrypoints.openai.engine.protocol import DeltaMessage, FunctionCall @@ -43,10 +44,14 @@ class MistralParser(DelegatingParser): model_output: str, request: ChatCompletionRequest | ResponsesRequest, enable_auto_tools: bool = False, + model_output_token_ids: Sequence[int] = (), ) -> tuple[str | None, str | None, list[FunctionCall] | None]: self._maybe_force_auto_tool_parsing(request) reasoning, content, tool_calls = super().parse( - model_output, request, enable_auto_tools + model_output, + request, + enable_auto_tools, + model_output_token_ids, ) if tool_calls: from vllm.tool_parsers.mistral_tool_parser import MistralToolCall diff --git a/vllm/parser/parser_manager.py b/vllm/parser/parser_manager.py index 6c2fdf52dd3..1b5133f5a8f 100644 --- a/vllm/parser/parser_manager.py +++ b/vllm/parser/parser_manager.py @@ -79,6 +79,7 @@ class ParserManager: reasoning_parser_name: str | None = None, enable_auto_tools: bool = False, model_name: str | None = None, + is_harmony: bool = False, ) -> type[Parser] | None: """ Get a Parser that handles both reasoning and tool parsing. @@ -91,6 +92,8 @@ class ParserManager: reasoning_parser_name: The name of the reasoning parser. enable_auto_tools: Whether auto tool choice is enabled. model_name: The model name for parser-specific warnings. + is_harmony: Whether the selected model uses the Harmony format. + If True, HarmonyParser is always returned. Returns: A Parser class, or None if neither parser is specified. @@ -108,6 +111,13 @@ class ParserManager: from vllm.utils.mistral import is_mistral_tool_parser + if is_harmony: + from vllm.parser.harmony import HarmonyParser + + HarmonyParser.reasoning_parser_cls = reasoning_parser_cls + HarmonyParser.tool_parser_cls = tool_parser_cls + return HarmonyParser + if is_mistral_tool_parser(tool_parser_cls): from vllm.parser.mistral import MistralParser diff --git a/vllm/reasoning/gptoss_reasoning_parser.py b/vllm/reasoning/gptoss_reasoning_parser.py index 1ba933cca31..d7bdca82912 100644 --- a/vllm/reasoning/gptoss_reasoning_parser.py +++ b/vllm/reasoning/gptoss_reasoning_parser.py @@ -8,7 +8,6 @@ from transformers import PreTrainedTokenizerBase from vllm.entrypoints.mcp.tool_server import ToolServer from vllm.entrypoints.openai.engine.protocol import DeltaMessage -from vllm.entrypoints.openai.parser.harmony_utils import parse_chat_output from vllm.logger import init_logger from vllm.reasoning import ReasoningParser @@ -132,10 +131,10 @@ class GptOssReasoningParser(ReasoningParser): return self.is_reasoning_end(input_ids[n - window :]) def extract_content_ids(self, input_ids: list[int]) -> list[int]: - _, content, _ = parse_chat_output(input_ids) - if content is None: - return [] - return self.model_tokenizer.encode(content) + raise NotImplementedError( + "GptOssReasoningParser only provides boundary detection. " + "Use HarmonyParser for output parsing." + ) def extract_reasoning_streaming( self, @@ -146,25 +145,10 @@ class GptOssReasoningParser(ReasoningParser): current_token_ids: Sequence[int], delta_token_ids: Sequence[int], ) -> DeltaMessage | None: - prev_reasoning, prev_content, _ = parse_chat_output(list(previous_token_ids)) - cur_reasoning, cur_content, _ = parse_chat_output(list(current_token_ids)) - reasoning_delta = None - content_delta = None - if cur_reasoning is not None: - prev_r = prev_reasoning or "" - if cur_reasoning.startswith(prev_r): - reasoning_delta = cur_reasoning[len(prev_r) :] or None - else: - reasoning_delta = cur_reasoning - if cur_content is not None: - prev_c = prev_content or "" - if cur_content.startswith(prev_c): - content_delta = cur_content[len(prev_c) :] or None - else: - content_delta = cur_content - if reasoning_delta is None and content_delta is None: - return None - return DeltaMessage(reasoning=reasoning_delta, content=content_delta) + raise NotImplementedError( + "GptOssReasoningParser only provides boundary detection. " + "Use HarmonyParser for output parsing." + ) def extract_reasoning( self, @@ -172,7 +156,8 @@ class GptOssReasoningParser(ReasoningParser): request: "ChatCompletionRequest | ResponsesRequest", ) -> tuple[str | None, str | None]: raise NotImplementedError( - "gpt-oss has a special branch for parsing reasoning in non-streaming mode. This method shouldn't be used." # noqa: E501 + "GptOssReasoningParser only provides boundary detection. " + "Use HarmonyParser for output parsing." ) # This function prepares the structural tag to format reasoning output diff --git a/vllm/tool_parsers/__init__.py b/vllm/tool_parsers/__init__.py index bf832f178be..9c534e77f66 100644 --- a/vllm/tool_parsers/__init__.py +++ b/vllm/tool_parsers/__init__.py @@ -143,8 +143,8 @@ _TOOL_PARSERS_TO_REGISTER = { "Olmo3PythonicToolParser", ), "openai": ( - "openai_tool_parser", - "OpenAIToolParser", + "gptoss_tool_parser", + "GptOssToolParser", ), "phi4_mini_json": ( "phi4mini_tool_parser", diff --git a/vllm/tool_parsers/gptoss_tool_parser.py b/vllm/tool_parsers/gptoss_tool_parser.py new file mode 100644 index 00000000000..6857e6bbe72 --- /dev/null +++ b/vllm/tool_parsers/gptoss_tool_parser.py @@ -0,0 +1,47 @@ +# 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, + ExtractedToolCallInformation, +) +from vllm.tool_parsers.abstract_tool_parser import Tool, ToolParser + +if TYPE_CHECKING: + from vllm.tokenizers import TokenizerLike + + +class GptOssToolParser(ToolParser): + """ + Stub tool parser for gpt-oss/harmony models. + + All output parsing is handled by HarmonyParser. This stub exists as a + capability declaration via HarmonyParser.tool_parser_cls. + """ + + def __init__(self, tokenizer: "TokenizerLike", tools: list[Tool] | None = None): + super().__init__(tokenizer, tools) + + def extract_tool_calls( + self, model_output, request, **kwargs + ) -> ExtractedToolCallInformation: + raise NotImplementedError( + "GptOssToolParser is a stub. Use HarmonyParser for tool parsing." + ) + + 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, + ) -> DeltaMessage | None: + raise NotImplementedError( + "GptOssToolParser is a stub. Use HarmonyParser for tool parsing." + ) diff --git a/vllm/tool_parsers/openai_tool_parser.py b/vllm/tool_parsers/openai_tool_parser.py deleted file mode 100644 index e5c37fbd3df..00000000000 --- a/vllm/tool_parsers/openai_tool_parser.py +++ /dev/null @@ -1,120 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import json -from collections.abc import Sequence -from typing import TYPE_CHECKING - -from vllm.entrypoints.openai.chat_completion.protocol import ( - ChatCompletionRequest, -) -from vllm.entrypoints.openai.engine.protocol import ( - DeltaMessage, - ExtractedToolCallInformation, - FunctionCall, - ToolCall, -) -from vllm.entrypoints.openai.parser.harmony_utils import ( - extract_function_from_recipient, - is_function_recipient, - parse_output_into_messages, -) -from vllm.logger import init_logger -from vllm.tool_parsers.abstract_tool_parser import ( - Tool, - ToolParser, -) - -if TYPE_CHECKING: - from vllm.tokenizers import TokenizerLike -else: - TokenizerLike = object - -logger = init_logger(__name__) - - -class OpenAIToolParser(ToolParser): - def __init__(self, tokenizer: "TokenizerLike", tools: list[Tool] | None = None): - super().__init__(tokenizer, tools) - - def extract_tool_calls( - self, - model_output: str, - request: ChatCompletionRequest, - token_ids: Sequence[int] | None = None, - ) -> ExtractedToolCallInformation: - if token_ids is None: - raise NotImplementedError( - "OpenAIToolParser requires token IDs and does not support text-based extraction." # noqa: E501 - ) - - parser = parse_output_into_messages(token_ids) - tool_calls = [] - final_content = None - commentary_content = None - - if len(parser.messages) > 0: - for msg in parser.messages: - if msg.author.role != "assistant": - continue - if len(msg.content) < 1: - continue - msg_text = msg.content[0].text - if msg.recipient and is_function_recipient(msg.recipient): - # If no content-type is given assume JSON, as that's the - # most common case with gpt-oss models. - if not msg.content_type or "json" in msg.content_type: - # load and dump the JSON text to check validity and - # remove any extra newlines or other odd formatting - try: - tool_args = json.dumps(json.loads(msg_text)) - except json.JSONDecodeError: - logger.exception( - "Error decoding JSON tool call from response." - ) - tool_args = msg_text - else: - tool_args = msg_text - tool_calls.append( - ToolCall( - type="function", - function=FunctionCall( - name=extract_function_from_recipient(msg.recipient), - arguments=tool_args, - ), - ) - ) - elif msg.channel == "final": - final_content = msg_text - elif msg.channel == "commentary" and not msg.recipient: - commentary_content = msg_text - - # Extract partial content from the parser state if the generation was truncated - if parser.current_content: - if parser.current_channel == "final": - final_content = parser.current_content - elif ( - parser.current_channel == "commentary" and not parser.current_recipient - ): - commentary_content = parser.current_content - - return ExtractedToolCallInformation( - tools_called=len(tool_calls) > 0, - tool_calls=tool_calls, - # prefer final content over commentary content if both are present - # commentary content is tool call preambles meant to be shown to the user - content=final_content or commentary_content, - ) - - 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: - raise NotImplementedError( - "Not being used, manual parsing in serving_chat.py" # noqa: E501 - ) From 8a91228dbe363d1d113deb2a82e289429130dd01 Mon Sep 17 00:00:00 2001 From: Dao007forever Date: Thu, 11 Jun 2026 14:33:48 -0700 Subject: [PATCH 14/52] [Bugfix][KVConnector][Mooncake] Close MooncakeDistributedStore on connector teardown (#45206) Signed-off-by: Dao Le Co-authored-by: Claude --- .../unit/test_mooncake_store_connector.py | 63 +++++++++++++++++++ .../unit/test_mooncake_store_worker.py | 31 +++++++++ .../v1/mooncake/store/connector.py | 15 +++++ .../kv_connector/v1/mooncake/store/worker.py | 16 +++++ 4 files changed, 125 insertions(+) diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_connector.py b/tests/v1/kv_connector/unit/test_mooncake_store_connector.py index 69593011db9..d3992b02b68 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_connector.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_connector.py @@ -614,3 +614,66 @@ def test_lookup_key_server_reset_skips_drain_when_no_send_thread(): assert call_order == ["remove_all"] assert sent == [protocol.RESP_OK] + + +def test_shutdown_closes_worker_store(): + vllm_config = _make_vllm_config() + kv_cache_config = _make_kv_cache_config() + + with ( + set_current_vllm_config(vllm_config), + patch( + "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store." + "connector.MooncakeStoreWorker" + ) as mock_worker_cls, + ): + connector = mooncake_store_connector.MooncakeStoreConnector( + vllm_config, KVConnectorRole.WORKER, kv_cache_config + ) + + worker = mock_worker_cls.return_value + connector.shutdown() + + worker.close.assert_called_once_with() + + +def test_del_invokes_shutdown_and_closes_store(): + vllm_config = _make_vllm_config() + kv_cache_config = _make_kv_cache_config() + + with ( + set_current_vllm_config(vllm_config), + patch( + "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store." + "connector.MooncakeStoreWorker" + ) as mock_worker_cls, + ): + connector = mooncake_store_connector.MooncakeStoreConnector( + vllm_config, KVConnectorRole.WORKER, kv_cache_config + ) + + worker = mock_worker_cls.return_value + # __del__ is the GC backstop; it must route through shutdown() -> close(). + connector.__del__() + + worker.close.assert_called_once_with() + + +def test_shutdown_scheduler_role_is_noop(): + vllm_config = _make_vllm_config() + kv_cache_config = _make_kv_cache_config() + + with ( + set_current_vllm_config(vllm_config), + patch( + "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store." + "connector.MooncakeStoreScheduler" + ), + ): + connector = mooncake_store_connector.MooncakeStoreConnector( + vllm_config, KVConnectorRole.SCHEDULER, kv_cache_config + ) + + # Scheduler role holds no store handle, so shutdown must be a safe no-op. + assert connector.connector_worker is None + connector.shutdown() diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py b/tests/v1/kv_connector/unit/test_mooncake_store_worker.py index 8cd5e6e5358..1130a7d6a78 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_worker.py @@ -1558,3 +1558,34 @@ def test_lookup_records_mooncake_metrics(): assert isinstance(stats, MooncakeStoreConnectorStats) assert len(stats.data["lookup_exists"]) == 1 assert stats.data["lookup_exists"][0]["num_keys"] == 2 + + +def test_store_worker_close_releases_store(): + worker = _make_bare_worker() + store = worker.store + + worker.close() + + store.close.assert_called_once_with() + assert worker.store is None + + +def test_store_worker_close_is_idempotent(): + worker = _make_bare_worker() + store = worker.store + + worker.close() + worker.close() + + # Second call short-circuits because store was already released. + store.close.assert_called_once_with() + + +def test_store_worker_close_swallows_store_errors(): + worker = _make_bare_worker() + worker.store.close.side_effect = RuntimeError("boom") + + # A failure tearing down the store must not propagate out of close(). + worker.close() + + assert worker.store is None diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py index 14d4b381a3c..d53cd13c2e4 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py @@ -153,6 +153,21 @@ class MooncakeStoreConnector(KVConnectorBase_V1, SupportsHMA): else: self.connector_worker = MooncakeStoreWorker(vllm_config, kv_cache_config) + def shutdown(self): + """Release connector resources on teardown. + + Closes the worker's MooncakeDistributedStore handle so its + TransferEngine and RDMA registrations are released. Invoked from the + engine's explicit shutdown path and as a backstop from ``__del__``; + a no-op on the scheduler role, which holds no store handle. + """ + worker = getattr(self, "connector_worker", None) + if worker is not None: + worker.close() + + def __del__(self): + self.shutdown() + # ============================================================ # Scheduler-side methods # ============================================================ diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py index 9c3ac83e06a..105762ccfcf 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py @@ -1426,6 +1426,22 @@ class MooncakeStoreWorker: return self.kv_send_thread.get_kv_events() return [] + def close(self) -> None: + """Release the MooncakeDistributedStore handle on teardown. + + Closing the store frees its TransferEngine, the registered RDMA + buffers, and the connection to the master server. Idempotent so it is + safe to call from both the explicit shutdown path and ``__del__``. + """ + store = getattr(self, "store", None) + if store is None: + return + self.store = None + try: + store.close() + except Exception as e: + logger.warning("Error closing MooncakeDistributedStore: %s", e) + # ============================================================ # Lookup Key Server From 9bbf42be266f88a4fabc65a0c3336edc442821cf Mon Sep 17 00:00:00 2001 From: Neil Schemenauer Date: Thu, 11 Jun 2026 15:59:11 -0700 Subject: [PATCH 15/52] Make mistral_common optional by deferring MistralToolCall import (#45305) Signed-off-by: Neil Schemenauer --- vllm/tool_parsers/streaming.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/vllm/tool_parsers/streaming.py b/vllm/tool_parsers/streaming.py index 7f6638dcb94..53b3f06bb8c 100644 --- a/vllm/tool_parsers/streaming.py +++ b/vllm/tool_parsers/streaming.py @@ -14,7 +14,6 @@ from vllm.entrypoints.openai.engine.protocol import ( DeltaMessage, DeltaToolCall, ) -from vllm.tool_parsers.mistral_tool_parser import MistralToolCall from vllm.tool_parsers.utils import partial_json_loads from vllm.utils.mistral import is_mistral_tokenizer @@ -77,6 +76,9 @@ def extract_named_tool_call_streaming( ) else: if is_mistral_tokenizer(tokenizer): + # Import mistral_common only if we need it. + from vllm.tool_parsers.mistral_tool_parser import MistralToolCall + tool_call_id = MistralToolCall.generate_random_id() else: tool_call_id = make_tool_call_id( From 6f573f486bc659adf51a8d4639e225097f2d8d39 Mon Sep 17 00:00:00 2001 From: jpwang Date: Fri, 12 Jun 2026 08:21:01 +0800 Subject: [PATCH 16/52] [Bugfix] Initialize missing attributes in mistral eagle (#45217) Signed-off-by: jpwang Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .../test_mistral_large_3_eagle.py | 146 ++++++++++++++++++ .../models/mistral_large_3_eagle.py | 10 ++ 2 files changed, 156 insertions(+) create mode 100644 tests/model_executor/test_mistral_large_3_eagle.py diff --git a/tests/model_executor/test_mistral_large_3_eagle.py b/tests/model_executor/test_mistral_large_3_eagle.py new file mode 100644 index 00000000000..d8ef109af98 --- /dev/null +++ b/tests/model_executor/test_mistral_large_3_eagle.py @@ -0,0 +1,146 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest +import torch +import torch.nn as nn + +from vllm.config.compilation import CompilationMode +from vllm.model_executor.models import deepseek_v2 as deepseek_mod +from vllm.model_executor.models import mistral_large_3_eagle as eagle_mod + + +class DummyPPGroup: + world_size = 1 + is_first_rank = True + is_last_rank = True + + +class DummyEmbedding(nn.Module): + def __init__(self, vocab_size, hidden_size, *args, **kwargs): + super().__init__() + self.hidden_size = hidden_size + + def forward(self, input_ids): + return torch.zeros( + (*input_ids.shape, self.hidden_size), + dtype=torch.float32, + device=input_ids.device, + ) + + +class DummyLinear(nn.Module): + def __init__(self, in_features, out_features, *args, **kwargs): + super().__init__() + self.out_features = out_features + + def forward(self, x): + return torch.zeros( + (*x.shape[:-1], self.out_features), + dtype=x.dtype, + device=x.device, + ) + + +class DummyNorm(nn.Module): + def __init__(self, *args, **kwargs): + super().__init__() + + def forward(self, hidden_states, residual=None): + return hidden_states, residual + + +class DummyDecoderLayer(nn.Module): + def __init__(self, *args, **kwargs): + super().__init__() + + def forward(self, positions, hidden_states, residual, llama_4_scaling=None): + return hidden_states, residual + + +def make_vllm_config( + *, model_type="mistral3", qk_nope_head_dim=128, qk_rope_head_dim=64 +): + hf_config = SimpleNamespace( + model_type=model_type, + first_k_dense_replace=0, + vocab_size=32000, + hidden_size=16, + num_hidden_layers=1, + rms_norm_eps=1e-5, + qk_nope_head_dim=qk_nope_head_dim, + qk_rope_head_dim=qk_rope_head_dim, + ) + + return SimpleNamespace( + model_config=SimpleNamespace(hf_config=hf_config), + quant_config=None, + parallel_config=SimpleNamespace( + eplb_config=SimpleNamespace(num_redundant_experts=0), + ), + scheduler_config=SimpleNamespace(max_num_batched_tokens=8), + cache_config=None, + compilation_config=SimpleNamespace(mode=CompilationMode.NONE), + ) + + +@pytest.fixture(autouse=True) +def patch_heavy_modules(monkeypatch): + monkeypatch.setattr(eagle_mod, "get_pp_group", lambda: DummyPPGroup()) + monkeypatch.setattr(deepseek_mod, "get_pp_group", lambda: DummyPPGroup()) + + monkeypatch.setattr(eagle_mod, "VocabParallelEmbedding", DummyEmbedding) + monkeypatch.setattr(eagle_mod, "RowParallelLinear", DummyLinear) + monkeypatch.setattr(eagle_mod, "RMSNorm", DummyNorm) + monkeypatch.setattr(eagle_mod, "DeepseekV2DecoderLayer", DummyDecoderLayer) + + +@pytest.mark.cpu_test +@pytest.mark.parametrize( + ("model_type", "qk_nope_head_dim", "qk_rope_head_dim", "expected_use_mha"), + [ + # MLA-style config: should not use MHA. + ("mistral3", 128, 64, False), + # No MLA dims: should use MHA, matching DeepseekV2Model.__init__ logic. + ("mistral3", 0, 0, True), + # DeepSeek model type always uses MHA by the parent logic. + ("deepseek", 128, 64, True), + ], +) +def test_eagle_mistral_large3_initializes_deepseek_runtime_attrs( + model_type, + qk_nope_head_dim, + qk_rope_head_dim, + expected_use_mha, +): + vllm_config = make_vllm_config( + model_type=model_type, + qk_nope_head_dim=qk_nope_head_dim, + qk_rope_head_dim=qk_rope_head_dim, + ) + + model = eagle_mod.EagleMistralLarge3Model(vllm_config=vllm_config) + + assert model.aux_hidden_state_layers == () + assert model.use_mha is expected_use_mha + + # Add this if your fix also copies num_redundant_experts from + # DeepseekV2Model.__init__. + assert model.num_redundant_experts == 0 + + +@pytest.mark.cpu_test +def test_eagle_mistral_large3_forward_reuses_deepseek_parent_forward(): + vllm_config = make_vllm_config() + model = eagle_mod.EagleMistralLarge3Model(vllm_config=vllm_config) + + input_ids = torch.tensor([[1, 2, 3]]) + positions = torch.tensor([[0, 1, 2]]) + hidden_states = torch.zeros((1, 3, 16)) + + output = model(input_ids, positions, hidden_states) + + assert isinstance(output, torch.Tensor) + assert output.shape == hidden_states.shape diff --git a/vllm/model_executor/models/mistral_large_3_eagle.py b/vllm/model_executor/models/mistral_large_3_eagle.py index 3fcc048f9fa..bde5bc9451f 100644 --- a/vllm/model_executor/models/mistral_large_3_eagle.py +++ b/vllm/model_executor/models/mistral_large_3_eagle.py @@ -75,6 +75,16 @@ class EagleMistralLarge3Model(DeepseekV2Model): ) self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.aux_hidden_state_layers: tuple[int, ...] = () + + # Needed by load_weights + qk_nope_head_dim = getattr(config, "qk_nope_head_dim", 0) + qk_rope_head_dim = getattr(config, "qk_rope_head_dim", 0) + self.use_mha = config.model_type == "deepseek" or all( + dim == 0 for dim in (qk_nope_head_dim, qk_rope_head_dim) + ) + self.num_redundant_experts = ( + vllm_config.parallel_config.eplb_config.num_redundant_experts + ) self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( ["hidden_states", "residual"], config.hidden_size ) From e0871ad2259768add6dc43e2972bd364d0d13086 Mon Sep 17 00:00:00 2001 From: yzong-rh Date: Thu, 11 Jun 2026 21:09:47 -0400 Subject: [PATCH 17/52] [Refactor] Chat Completions Streaming Harmony Refactor and Bugfixes (#45104) Signed-off-by: Yifan Zong --- .../test_serving_chat_stream_harmony.py | 471 ------------------ tests/parser/test_harmony.py | 334 ++++++++++++- .../openai/chat_completion/serving.py | 73 +-- .../openai/chat_completion/stream_harmony.py | 167 ------- vllm/parser/harmony.py | 82 ++- 5 files changed, 394 insertions(+), 733 deletions(-) delete mode 100644 tests/entrypoints/openai/chat_completion/test_serving_chat_stream_harmony.py delete mode 100644 vllm/entrypoints/openai/chat_completion/stream_harmony.py diff --git a/tests/entrypoints/openai/chat_completion/test_serving_chat_stream_harmony.py b/tests/entrypoints/openai/chat_completion/test_serving_chat_stream_harmony.py deleted file mode 100644 index 1c058adaf0a..00000000000 --- a/tests/entrypoints/openai/chat_completion/test_serving_chat_stream_harmony.py +++ /dev/null @@ -1,471 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Unit tests for harmony streaming delta extraction. -""" - -from dataclasses import dataclass, field -from unittest.mock import patch - -import pytest - -from vllm.entrypoints.openai.chat_completion.stream_harmony import ( - TokenState, - extract_harmony_streaming_delta, -) - - -@dataclass -class MockMessage: - """Mock message object for testing.""" - - channel: str | None = None - recipient: str | None = None - - -@dataclass -class MockStreamableParser: - """Mock StreamableParser for testing without openai_harmony dependency.""" - - messages: list[MockMessage] = field(default_factory=list) - - -class TestExtractHarmonyStreamingDelta: - """Tests for extract_harmony_streaming_delta function.""" - - @pytest.mark.parametrize( - "delta_text,expected_content", - [ - ("Hello, world!", "Hello, world!"), - ("", ""), - ], - ) - def test_final_channel_returns_content_delta(self, delta_text, expected_content): - """Test that final channel returns a DeltaMessage with content.""" - parser = MockStreamableParser() - - # Updated to use TokenState list - token_states = [TokenState(channel="final", recipient=None, text=delta_text)] - - delta_message, tools_streamed = extract_harmony_streaming_delta( - harmony_parser=parser, - token_states=token_states, - prev_recipient=None, - include_reasoning=False, - ) - - assert delta_message is not None - assert delta_message.content == expected_content - assert tools_streamed is False - - @pytest.mark.parametrize( - "include_reasoning,expected_has_message", - [ - (True, True), - (False, False), - ], - ) - def test_analysis_channel_reasoning(self, include_reasoning, expected_has_message): - """Test analysis channel respects include_reasoning flag.""" - parser = MockStreamableParser() - text = "Let me think..." - token_states = [TokenState(channel="analysis", recipient=None, text=text)] - - delta_message, tools_streamed = extract_harmony_streaming_delta( - harmony_parser=parser, - token_states=token_states, - prev_recipient=None, - include_reasoning=include_reasoning, - ) - - if expected_has_message: - assert delta_message is not None - assert delta_message.reasoning == text - else: - assert delta_message is None - assert tools_streamed is False - - @pytest.mark.parametrize("channel", ["commentary", "analysis"]) - @patch("vllm.entrypoints.openai.chat_completion.stream_harmony.make_tool_call_id") - def test_new_tool_call(self, mock_make_tool_call_id, channel): - """Test new tool call creation when recipient changes.""" - mock_make_tool_call_id.return_value = "call_test123" - parser = MockStreamableParser() - - token_states = [ - TokenState(channel=channel, recipient="functions.get_weather", text="") - ] - - delta_message, tools_streamed = extract_harmony_streaming_delta( - harmony_parser=parser, - token_states=token_states, - prev_recipient=None, - include_reasoning=False, - ) - - assert delta_message is not None - assert len(delta_message.tool_calls) == 1 - tool_call = delta_message.tool_calls[0] - assert tool_call.id == "call_test123" - assert tool_call.type == "function" - assert tool_call.function.name == "get_weather" - assert tool_call.function.arguments == "" - assert tool_call.index == 0 - assert tools_streamed is True - - @pytest.mark.parametrize("channel", ["commentary", "analysis"]) - def test_tool_call_argument_streaming(self, channel): - """Test streaming tool call arguments (same recipient).""" - parser = MockStreamableParser() - args_text = '{"location": "Paris"}' - - token_states = [ - TokenState( - channel=channel, - recipient="functions.get_weather", - text=args_text, - ) - ] - - delta_message, tools_streamed = extract_harmony_streaming_delta( - harmony_parser=parser, - token_states=token_states, - prev_recipient="functions.get_weather", - include_reasoning=False, - ) - - assert delta_message is not None - tool_call = delta_message.tool_calls[0] - assert tool_call.id is None - assert tool_call.function.arguments == args_text - assert tool_call.index == 0 - assert tools_streamed is True - - @pytest.mark.parametrize("channel", ["commentary", "analysis"]) - def test_tool_call_empty_arguments_returns_none(self, channel): - """Test empty delta_text with same recipient returns None.""" - parser = MockStreamableParser() - - token_states = [ - TokenState(channel=channel, recipient="functions.get_weather", text="") - ] - - delta_message, tools_streamed = extract_harmony_streaming_delta( - harmony_parser=parser, - token_states=token_states, - prev_recipient="functions.get_weather", - include_reasoning=False, - ) - - assert delta_message is None - assert tools_streamed is False - - def test_tool_call_index_from_previous_messages(self): - """Test tool call index accounts for previous function messages.""" - messages = [ - MockMessage(channel="analysis", recipient=None), # Not counted - MockMessage(channel="commentary", recipient="functions.tool1"), # Counted - MockMessage(channel="final", recipient=None), # Not counted - ] - parser = MockStreamableParser(messages=messages) - - token_states = [ - TokenState( - channel="commentary", - recipient="functions.tool2", - text="args", - ) - ] - - delta_message, _ = extract_harmony_streaming_delta( - harmony_parser=parser, - token_states=token_states, - prev_recipient="functions.tool2", - include_reasoning=False, - ) - - assert delta_message.tool_calls[0].index == 1 - - def test_returns_preambles_as_content(self): - """Test that commentary with no recipient (preamble) is user content.""" - parser = MockStreamableParser() - delta_text = "some text" - - token_states = [ - TokenState(channel="commentary", recipient=None, text=delta_text) - ] - - delta_message, tools_streamed = extract_harmony_streaming_delta( - harmony_parser=parser, - token_states=token_states, - prev_recipient=None, - include_reasoning=True, - ) - - assert delta_message.content == delta_text - assert tools_streamed is False - - @pytest.mark.parametrize("channel", ["commentary", "analysis"]) - @patch("vllm.entrypoints.openai.chat_completion.stream_harmony.make_tool_call_id") - def test_new_tool_call_dotted_function_name(self, mock_make_tool_call_id, channel): - mock_make_tool_call_id.return_value = "call_dotted123" - parser = MockStreamableParser() - - token_states = [TokenState(channel=channel, recipient="math.sum", text="")] - - delta_message, tools_streamed = extract_harmony_streaming_delta( - harmony_parser=parser, - token_states=token_states, - prev_recipient=None, - include_reasoning=False, - ) - - assert delta_message is not None - assert len(delta_message.tool_calls) == 1 - tool_call = delta_message.tool_calls[0] - assert tool_call.id == "call_dotted123" - assert tool_call.type == "function" - assert tool_call.function.name == "math.sum" - assert tool_call.function.arguments == "" - assert tool_call.index == 0 - assert tools_streamed is True - - @pytest.mark.parametrize( - "channel,recipient", - [ - (None, None), - ("unknown_channel", None), - ("commentary", "browser.search"), - ("commentary", "assistant"), - ], - ) - def test_returns_none_for_invalid_inputs(self, channel, recipient): - """Test that invalid channel/recipient combinations return None.""" - parser = MockStreamableParser() - - token_states = [ - TokenState(channel=channel, recipient=recipient, text="some text") - ] - - delta_message, tools_streamed = extract_harmony_streaming_delta( - harmony_parser=parser, - token_states=token_states, - prev_recipient=None, - include_reasoning=True, - ) - - assert delta_message is None - assert tools_streamed is False - - def test_consecutive_token_grouping(self): - """ - Test that consecutive tokens with the same channel/recipient - are merged into a single processing group. - """ - parser = MockStreamableParser() - token_states = [ - TokenState("final", None, "H"), - TokenState("final", None, "el"), - TokenState("final", None, "lo"), - TokenState("final", None, ","), - TokenState("final", None, " World"), - ] - - delta_message, _ = extract_harmony_streaming_delta( - harmony_parser=parser, - token_states=token_states, - prev_recipient=None, - include_reasoning=False, - ) - - assert delta_message is not None - assert delta_message.content == "Hello, World" - - @patch("vllm.entrypoints.openai.chat_completion.stream_harmony.make_tool_call_id") - def test_complex_batch_permutation(self, mock_make_id): - """ - Test a complex permutation: Reasoning -> Tool Call -> Content. - This verifies that multiple distinct actions in one batch - are all captured in the single DeltaMessage. - """ - mock_make_id.return_value = "call_batch_test" - parser = MockStreamableParser() - - token_states = [ - # 1. Reasoning - TokenState("analysis", None, "Reasoning about query..."), - # 2. Tool Calling - TokenState("commentary", "functions.search", '{"query":'), - TokenState("commentary", "functions.search", ' "vllm"}'), - # 3. Final Content - TokenState("final", None, "."), - ] - - delta_message, tools_streamed = extract_harmony_streaming_delta( - harmony_parser=parser, - token_states=token_states, - prev_recipient=None, - include_reasoning=True, - ) - - assert delta_message is not None - - assert delta_message.reasoning == "Reasoning about query..." - - # We expect 2 objects for 1 logical tool call: - # 1. The definition (id, name, type) - # 2. The arguments payload - assert len(delta_message.tool_calls) == 2 - - header = delta_message.tool_calls[0] - payload = delta_message.tool_calls[1] - - assert header.function.name == "search" - assert header.id == "call_batch_test" - assert header.index == 0 - - assert payload.index == 0 - assert payload.function.arguments == '{"query": "vllm"}' - - assert delta_message.content == "." - assert tools_streamed is True - - @patch("vllm.entrypoints.openai.chat_completion.stream_harmony.make_tool_call_id") - def test_tool_call_index_consistency_with_ongoing_call(self, mock_make_id): - """ - Test that an ongoing tool call continuation and subsequent new calls - maintain correct indexing when interleaved with content. - """ - mock_make_id.side_effect = ["id_b", "id_c"] - - messages = [ - MockMessage(channel="commentary", recipient="functions.previous_tool") - ] - parser = MockStreamableParser(messages=messages) - - token_states = [ - TokenState("commentary", "functions.tool_a", '{"key_a": "val_a"}'), - TokenState("final", None, "Thinking..."), - TokenState("commentary", "functions.tool_b", '{"key_b": "val_b"}'), - TokenState("final", None, " Thinking again..."), - TokenState("commentary", "functions.tool_c", '{"key_c": "val_c"}'), - ] - - delta_message, _ = extract_harmony_streaming_delta( - harmony_parser=parser, - token_states=token_states, - prev_recipient="functions.tool_a", - include_reasoning=False, - ) - - assert delta_message is not None - - tool_a_deltas = [t for t in delta_message.tool_calls if t.index == 1] - assert len(tool_a_deltas) > 0 - assert tool_a_deltas[0].id is None - assert tool_a_deltas[0].function.arguments == '{"key_a": "val_a"}' - - tool_b_header = next(t for t in delta_message.tool_calls if t.id == "id_b") - assert tool_b_header.index == 2 - tool_b_args = next( - t for t in delta_message.tool_calls if t.index == 2 and t.id is None - ) - assert tool_b_args.function.arguments == '{"key_b": "val_b"}' - - tool_c_start = next(t for t in delta_message.tool_calls if t.id == "id_c") - assert tool_c_start.index == 3 - tool_c_args = next( - t for t in delta_message.tool_calls if t.index == 3 and t.id is None - ) - assert tool_c_args.function.arguments == '{"key_c": "val_c"}' - - assert delta_message.content == "Thinking... Thinking again..." - - -class TestToolCallsOnNonStandardChannels: - """Tool calls are detected by recipient, not channel. - - Models sometimes emit tool calls on unexpected channels (e.g. ``comment`` - instead of ``commentary``). These tests verify that the streaming delta - extraction is channel-agnostic for tool call detection. - """ - - @patch("vllm.entrypoints.openai.chat_completion.stream_harmony.make_tool_call_id") - def test_prefixed_tool_call_on_comment_channel(self, mock_make_tool_call_id): - mock_make_tool_call_id.return_value = "call_comment_chan" - parser = MockStreamableParser() - - token_states = [ - TokenState(channel="comment", recipient="functions.get_weather", text="") - ] - - delta_message, tools_streamed = extract_harmony_streaming_delta( - harmony_parser=parser, - token_states=token_states, - prev_recipient=None, - include_reasoning=False, - ) - - assert delta_message is not None - assert len(delta_message.tool_calls) == 1 - assert delta_message.tool_calls[0].function.name == "get_weather" - assert tools_streamed is True - - @patch("vllm.entrypoints.openai.chat_completion.stream_harmony.make_tool_call_id") - def test_bare_tool_call_on_comment_channel(self, mock_make_tool_call_id): - mock_make_tool_call_id.return_value = "call_bare_comment" - parser = MockStreamableParser() - - token_states = [TokenState(channel="comment", recipient="get_weather", text="")] - - delta_message, tools_streamed = extract_harmony_streaming_delta( - harmony_parser=parser, - token_states=token_states, - prev_recipient=None, - include_reasoning=False, - ) - - assert delta_message is not None - assert len(delta_message.tool_calls) == 1 - assert delta_message.tool_calls[0].function.name == "get_weather" - assert tools_streamed is True - - def test_tool_call_arguments_on_comment_channel(self): - parser = MockStreamableParser() - args_text = '{"location": "Paris"}' - - token_states = [ - TokenState( - channel="comment", recipient="functions.get_weather", text=args_text - ) - ] - - delta_message, tools_streamed = extract_harmony_streaming_delta( - harmony_parser=parser, - token_states=token_states, - prev_recipient="functions.get_weather", - include_reasoning=False, - ) - - assert delta_message is not None - assert delta_message.tool_calls[0].function.arguments == args_text - assert tools_streamed is True - - def test_base_index_counts_tool_calls_on_comment_channel(self): - messages = [ - MockMessage(channel="comment", recipient="functions.tool1"), - ] - parser = MockStreamableParser(messages=messages) - - token_states = [ - TokenState(channel="commentary", recipient="functions.tool2", text="args") - ] - - delta_message, _ = extract_harmony_streaming_delta( - harmony_parser=parser, - token_states=token_states, - prev_recipient="functions.tool2", - include_reasoning=False, - ) - - assert delta_message.tool_calls[0].index == 1 diff --git a/tests/parser/test_harmony.py b/tests/parser/test_harmony.py index 98687b08edd..2740ccbca04 100644 --- a/tests/parser/test_harmony.py +++ b/tests/parser/test_harmony.py @@ -94,16 +94,36 @@ def get_text(msg: Message) -> str: return msg.content[0].text if msg.content else "" -def visible_segments(result) -> list[tuple[str | None, str | None, str]]: +def tool_call_tuples(tool_calls: list[FunctionCall] | None) -> list[tuple[str, str]]: + return [] if tool_calls is None else [(tc.name, tc.arguments) for tc in tool_calls] + + +def tool_call_headers(delta_message) -> list: + if delta_message is None or not delta_message.tool_calls: + return [] return [ - (segment.channel, segment.recipient, segment.delta) - for segment in result.segments - if not segment.is_boundary and segment.delta + tool_call + for tool_call in delta_message.tool_calls + if tool_call.function and tool_call.function.name ] -def tool_call_tuples(tool_calls: list[FunctionCall] | None) -> list[tuple[str, str]]: - return [] if tool_calls is None else [(tc.name, tc.arguments) for tc in tool_calls] +def tool_call_payloads(delta_message) -> list: + if delta_message is None or not delta_message.tool_calls: + return [] + return [ + tool_call + for tool_call in delta_message.tool_calls + if tool_call.function and tool_call.function.arguments + ] + + +def combined_tool_arguments(delta_message) -> dict[int, str]: + combined: dict[int, str] = {} + for tool_call in tool_call_payloads(delta_message): + combined.setdefault(tool_call.index, "") + combined[tool_call.index] += tool_call.function.arguments + return combined class TestParse: @@ -394,6 +414,276 @@ class TestParse: ] +class TestParseDelta: + def test_basic(self, gpt_oss_tokenizer, chat_request): + parser = HarmonyParser(gpt_oss_tokenizer) + + first_delta = parser.parse_delta( + delta_text="", + delta_token_ids=encode_output("<|channel|>analysis<|message|>Thinking"), + request=chat_request, + finished=False, + ) + second_delta = parser.parse_delta( + delta_text="", + delta_token_ids=encode_output( + "<|end|><|start|>assistant<|channel|>final<|message|>Answer" + ), + request=chat_request, + finished=False, + ) + + assert first_delta is not None + assert first_delta.reasoning == "Thinking" + assert first_delta.content is None + assert second_delta is not None + assert second_delta.content == "Answer" + assert second_delta.reasoning is None + + def test_multi_token(self, gpt_oss_tokenizer, chat_request): + parser = HarmonyParser(gpt_oss_tokenizer) + + delta = parser.parse_delta( + delta_text="", + delta_token_ids=encode_output("<|channel|>final<|message|>Hello, world!"), + request=chat_request, + finished=False, + ) + + assert delta is not None + assert delta.content == "Hello, world!" + assert delta.reasoning is None + assert not delta.tool_calls + + @pytest.mark.parametrize("tool_channel", ["commentary", "analysis"]) + def test_tool_call_split_across_deltas( + self, gpt_oss_tokenizer, chat_request, tool_channel + ): + parser = HarmonyParser(gpt_oss_tokenizer) + + first_delta = parser.parse_delta( + delta_text="", + delta_token_ids=encode_output( + "<|channel|>analysis<|message|>Thinking<|end|>" + f"<|start|>assistant to=functions.get_weather<|channel|>{tool_channel}" + '<|constrain|>json<|message|>{"location": ' + ), + request=chat_request, + finished=False, + ) + second_delta = parser.parse_delta( + delta_text="", + delta_token_ids=encode_output('"Paris"}<|call|>'), + request=chat_request, + finished=False, + ) + + assert first_delta is not None + assert first_delta.reasoning == "Thinking" + assert first_delta.content is None + assert [tool.function.name for tool in tool_call_headers(first_delta)] == [ + "get_weather" + ] + assert combined_tool_arguments(first_delta) == {0: '{"location": '} + assert {tool.index for tool in first_delta.tool_calls} == {0} + + assert second_delta is not None + assert second_delta.reasoning is None + assert second_delta.content is None + assert not tool_call_headers(second_delta) + assert combined_tool_arguments(second_delta) == {0: '"Paris"}'} + assert {tool.index for tool in second_delta.tool_calls} == {0} + + def test_commentary_preamble_streaming(self, gpt_oss_tokenizer, chat_request): + parser = HarmonyParser(gpt_oss_tokenizer) + + delta = parser.parse_delta( + delta_text="", + delta_token_ids=encode_output( + "<|channel|>commentary<|message|>I'll search for that" + ), + request=chat_request, + finished=False, + ) + + assert delta is not None + assert delta.content == "I'll search for that" + assert delta.reasoning is None + assert not delta.tool_calls + + def test_multiple_choices(self, gpt_oss_tokenizer, chat_request): + parser_a = HarmonyParser(gpt_oss_tokenizer) + parser_b = HarmonyParser(gpt_oss_tokenizer) + + delta_a = parser_a.parse_delta( + delta_text="", + delta_token_ids=encode_output( + "<|channel|>analysis<|message|>Check weather<|end|>" + "<|start|>assistant to=functions.get_weather<|channel|>commentary" + '<|constrain|>json<|message|>{"location": "Paris"}' + ), + request=chat_request, + finished=False, + ) + delta_b = parser_b.parse_delta( + delta_text="", + delta_token_ids=encode_output( + "<|channel|>analysis<|message|>Check time<|end|>" + "<|start|>assistant to=functions.get_time<|channel|>commentary" + '<|constrain|>json<|message|>{"timezone": "UTC"}' + ), + request=chat_request, + finished=False, + ) + + assert [tool.function.name for tool in tool_call_headers(delta_a)] == [ + "get_weather" + ] + assert [tool.function.name for tool in tool_call_headers(delta_b)] == [ + "get_time" + ] + assert {tool.index for tool in delta_a.tool_calls} == {0} + assert {tool.index for tool in delta_b.tool_calls} == {0} + + def test_dotted_function_name(self, gpt_oss_tokenizer, chat_request): + parser = HarmonyParser(gpt_oss_tokenizer) + + delta = parser.parse_delta( + delta_text="", + delta_token_ids=encode_output( + "<|channel|>analysis<|message|>Compute this<|end|>" + "<|start|>assistant to=math.sum<|channel|>commentary" + '<|constrain|>json<|message|>{"a": 2, "b": 3}' + ), + request=chat_request, + finished=False, + ) + + assert delta is not None + assert [tool.function.name for tool in tool_call_headers(delta)] == ["math.sum"] + assert {tool.index for tool in delta.tool_calls} == {0} + + @pytest.mark.parametrize("recipient", ["assistant", "browser"]) + def test_builtin_recipient_skipped( + self, + gpt_oss_tokenizer, + chat_request, + recipient, + ): + parser = HarmonyParser(gpt_oss_tokenizer) + prompt = [Message.from_role_and_content(Role.USER, "Hello")] + response = [tool_call(recipient, "Ignore this", content_type=None)] + + delta = parser.parse_delta( + delta_text="", + delta_token_ids=get_model_output_tokens(prompt, response), + request=chat_request, + finished=False, + ) + + assert delta is None + + def test_cross_channel_with_tool(self, gpt_oss_tokenizer, chat_request): + parser = HarmonyParser(gpt_oss_tokenizer) + + delta = parser.parse_delta( + delta_text="", + delta_token_ids=encode_output( + "<|channel|>analysis<|message|>Reasoning about query...<|end|>" + "<|start|>assistant to=functions.search<|channel|>commentary" + '<|constrain|>json<|message|>{"query": "vllm"}<|call|>' + "<|start|>assistant<|channel|>final<|message|>Done" + ), + request=chat_request, + finished=False, + ) + + assert delta is not None + assert delta.reasoning == "Reasoning about query..." + assert delta.content == "Done" + assert [tool.function.name for tool in tool_call_headers(delta)] == ["search"] + assert combined_tool_arguments(delta) == {0: '{"query": "vllm"}'} + + def test_tool_index_across_calls(self, gpt_oss_tokenizer, chat_request): + parser = HarmonyParser(gpt_oss_tokenizer) + + first_delta = parser.parse_delta( + delta_text="", + delta_token_ids=encode_output( + "<|channel|>analysis<|message|>Thinking<|end|>" + "<|start|>assistant to=functions.get_weather<|channel|>commentary" + '<|constrain|>json<|message|>{"location": "Paris"}<|call|>' + ), + request=chat_request, + finished=False, + ) + second_delta = parser.parse_delta( + delta_text="", + delta_token_ids=encode_output( + "<|start|>assistant to=functions.get_time<|channel|>commentary" + '<|constrain|>json<|message|>{"timezone": "UTC"}<|call|>' + ), + request=chat_request, + finished=False, + ) + + assert [tool.index for tool in tool_call_headers(first_delta)] == [0] + assert [tool.index for tool in tool_call_headers(second_delta)] == [1] + assert [tool.function.name for tool in tool_call_headers(second_delta)] == [ + "get_time" + ] + + def test_multi_tool_interleaved(self, gpt_oss_tokenizer, chat_request): + parser = HarmonyParser(gpt_oss_tokenizer) + + first_delta = parser.parse_delta( + delta_text="", + delta_token_ids=encode_output( + "<|channel|>analysis<|message|>Plan<|end|>" + "<|start|>assistant to=functions.tool_a<|channel|>commentary" + '<|constrain|>json<|message|>{"a": 1}<|call|>' + "<|start|>assistant to=functions.tool_b<|channel|>commentary" + '<|constrain|>json<|message|>{"b": ' + ), + request=chat_request, + finished=False, + ) + second_delta = parser.parse_delta( + delta_text="", + delta_token_ids=encode_output("2"), + request=chat_request, + finished=False, + ) + third_delta = parser.parse_delta( + delta_text="", + delta_token_ids=encode_output( + "}<|call|><|start|>assistant<|channel|>final<|message|>Done<|end|>" + "<|start|>assistant to=functions.tool_c<|channel|>commentary" + '<|constrain|>json<|message|>{"c": 3}' + ), + request=chat_request, + finished=False, + ) + + assert [tool.index for tool in tool_call_headers(first_delta)] == [0, 1] + assert combined_tool_arguments(first_delta) == { + 0: '{"a": 1}', + 1: '{"b": ', + } + + assert second_delta is not None + assert [tool.index for tool in tool_call_payloads(second_delta)] == [1] + assert combined_tool_arguments(second_delta) == {1: "2"} + + assert third_delta is not None + assert third_delta.content == "Done" + assert combined_tool_arguments(third_delta) == { + 1: "}", + 2: '{"c": 3}', + } + assert [tool.index for tool in tool_call_headers(third_delta)] == [2] + + class TestProcessChunk: def test_empty(self, harmony_parser): result = harmony_parser.process_chunk([]) @@ -405,7 +695,9 @@ class TestProcessChunk: encode_output("<|channel|>final<|message|>Hello") ) - assert visible_segments(result) == [("final", None, "Hello")] + assert [ + (s.channel, s.recipient, s.delta) for s in result.segments if s.delta + ] == [("final", None, "Hello")] def test_cross_channel(self, harmony_parser): result = harmony_parser.process_chunk( @@ -415,24 +707,13 @@ class TestProcessChunk: ) ) - assert visible_segments(result) == [ + assert [ + (s.channel, s.recipient, s.delta) for s in result.segments if s.delta + ] == [ ("analysis", None, "Think"), ("final", None, "Answer"), ] - def test_boundary_detection(self, harmony_parser): - result = harmony_parser.process_chunk( - encode_output("<|channel|>final<|message|>Done<|end|>") - ) - - boundary_segments = [ - segment for segment in result.segments if segment.is_boundary - ] - assert len(boundary_segments) == 1 - assert boundary_segments[0].completed_message is not None - assert boundary_segments[0].completed_message.channel == "final" - assert get_text(boundary_segments[0].completed_message) == "Done" - def test_multi_boundary(self, harmony_parser): result = harmony_parser.process_chunk( encode_output( @@ -442,11 +723,14 @@ class TestProcessChunk: ) boundary_segments = [ - segment for segment in result.segments if segment.is_boundary + segment + for segment in result.segments + if segment.completed_message is not None ] assert [ - get_text(segment.completed_message) for segment in boundary_segments + (segment.completed_message.channel, get_text(segment.completed_message)) + for segment in boundary_segments ] == [ - "One", - "Two", + ("analysis", "One"), + ("final", "Two"), ] diff --git a/vllm/entrypoints/openai/chat_completion/serving.py b/vllm/entrypoints/openai/chat_completion/serving.py index 4924ceb8b5d..52d18519eff 100644 --- a/vllm/entrypoints/openai/chat_completion/serving.py +++ b/vllm/entrypoints/openai/chat_completion/serving.py @@ -33,10 +33,6 @@ from vllm.entrypoints.openai.chat_completion.protocol import ( ChatCompletionStreamResponse, ChatMessage, ) -from vllm.entrypoints.openai.chat_completion.stream_harmony import ( - TokenState, - extract_harmony_streaming_delta, -) from vllm.entrypoints.openai.engine.protocol import ( DeltaMessage, ErrorResponse, @@ -52,9 +48,6 @@ from vllm.entrypoints.openai.engine.serving import ( clamp_prompt_logprobs, ) from vllm.entrypoints.openai.models.serving import OpenAIServingModels -from vllm.entrypoints.openai.parser.harmony_utils import ( - get_streamable_parser_for_assistant, -) from vllm.entrypoints.serve.utils.api_utils import get_max_tokens, should_include_usage from vllm.entrypoints.serve.utils.request_logger import RequestLogger from vllm.entrypoints.serve.utils.tool_calls_utils import ( @@ -155,7 +148,6 @@ class OpenAIServingChat(OpenAIServing): if mc.generation_config not in ("auto", "vllm") else getattr(mc, "override_generation_config", {}).get("max_new_tokens") ) - self.use_harmony = self.model_config.hf_config.model_type == "gpt_oss" self.tool_call_id_type = get_tool_call_id_type(self.model_config) # NOTE(woosuk): While OpenAI's chat completion API supports browsing @@ -408,11 +400,6 @@ class OpenAIServingChat(OpenAIServing): finish_reason_sent = [False] * num_choices num_prompt_tokens = 0 num_cached_tokens = None - if self.use_harmony: - harmony_parsers = [ - get_streamable_parser_for_assistant() for _ in range(num_choices) - ] - harmony_tools_streamed = [False] * num_choices tools_streamed = [False] * num_choices if isinstance(request.tool_choice, ChatCompletionNamedToolChoiceParam): @@ -443,6 +430,7 @@ class OpenAIServingChat(OpenAIServing): ] for p in parsers: if p is not None: + # NOTE: HarmonyParser ignores _stream_state (uses its own FSM). p._stream_state.tool_call_id_type = self.tool_call_id_type p._stream_state.history_tool_call_cnt = history_tool_call_cnt else: @@ -572,32 +560,7 @@ class OpenAIServingChat(OpenAIServing): else: logprobs = None - if self.use_harmony: - harmony_parser = harmony_parsers[i] - prev_recipient = harmony_parser.current_recipient - - # Track accumulated content per token with their state - token_states: list[TokenState] = [] - for token_id in output.token_ids: - harmony_parser.process(token_id) - token_delta = harmony_parser.last_content_delta or "" - token_states.append( - TokenState( - harmony_parser.current_channel, - harmony_parser.current_recipient, - token_delta, - ) - ) - delta_text = "".join(delta for _, _, delta in token_states) - cur_channel = harmony_parser.current_channel - - # handle the case where several tokens where generated at once - # including the final token, leading to a delta in the text - # but the current channel to be empty (start state) - if not cur_channel and delta_text: - cur_channel = "final" - else: - delta_text = output.text + delta_text = output.text if ( not delta_text @@ -609,17 +572,7 @@ class OpenAIServingChat(OpenAIServing): delta_message: DeltaMessage | None - if self.use_harmony: - delta_message, tools_streamed_flag = ( - extract_harmony_streaming_delta( - harmony_parser=harmony_parser, - token_states=token_states, - prev_recipient=prev_recipient, - include_reasoning=request.include_reasoning, - ) - ) - harmony_tools_streamed[i] |= tools_streamed_flag - elif parser is not None: + if parser is not None: delta_message = parser.parse_delta( delta_text=delta_text, delta_token_ids=as_list(output.token_ids), @@ -627,8 +580,20 @@ class OpenAIServingChat(OpenAIServing): prompt_token_ids=res.prompt_token_ids, finished=output.finish_reason is not None, ) - if delta_message and delta_message.tool_calls: - tools_streamed[i] = True + if delta_message is not None: + if delta_message.tool_calls: + tools_streamed[i] = True + + if ( + delta_message.reasoning + and not request.include_reasoning + ): + delta_message.reasoning = None + if not ( + delta_message.content or delta_message.tool_calls + ): + delta_message = None + # handle streaming just a content delta (no parsers) else: delta_message = DeltaMessage(content=delta_text) @@ -706,9 +671,7 @@ class OpenAIServingChat(OpenAIServing): # finish_reason is: # "tool_calls" for "auto" or "required" tool calls, # and "stop" for named tool calls. - if (tools_streamed[i] and not tool_choice_function_name) or ( - self.use_harmony and harmony_tools_streamed[i] - ): + if tools_streamed[i] and not tool_choice_function_name: finish_reason_ = "tool_calls" else: finish_reason_ = ( diff --git a/vllm/entrypoints/openai/chat_completion/stream_harmony.py b/vllm/entrypoints/openai/chat_completion/stream_harmony.py deleted file mode 100644 index 271f8e8c85a..00000000000 --- a/vllm/entrypoints/openai/chat_completion/stream_harmony.py +++ /dev/null @@ -1,167 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Harmony-specific streaming delta extraction for chat completions. - -This module handles the extraction of DeltaMessage objects from -harmony parser state during streaming chat completions. -""" - -from typing import NamedTuple - -from openai_harmony import StreamableParser - -from vllm.entrypoints.chat_utils import make_tool_call_id -from vllm.entrypoints.openai.engine.protocol import ( - DeltaFunctionCall, - DeltaMessage, - DeltaToolCall, -) -from vllm.entrypoints.openai.parser.harmony_utils import ( - extract_function_from_recipient, - is_function_recipient, -) - - -class TokenState(NamedTuple): - channel: str | None - recipient: str | None - text: str - - -def extract_harmony_streaming_delta( - harmony_parser: StreamableParser, - token_states: list[TokenState], - prev_recipient: str | None, - include_reasoning: bool, -) -> tuple[DeltaMessage | None, bool]: - """ - Extract a DeltaMessage from harmony parser state during streaming. - - Args: - harmony_parser: The StreamableParser instance tracking parse state - token_states: List of TokenState tuples for each token - prev_recipient: Previous recipient for detecting tool call transitions - include_reasoning: Whether to include reasoning content - - Returns: - A tuple of (DeltaMessage or None, tools_streamed_flag) - """ - - if not token_states: - return None, False - - tools_streamed = False - - # Group consecutive tokens with same channel/recipient - groups: list[TokenState] = [] - - current_channel = token_states[0].channel - current_recipient = token_states[0].recipient - current_text = token_states[0].text - - for i in range(1, len(token_states)): - state = token_states[i] - if state.channel == current_channel and state.recipient == current_recipient: - current_text += state.text - else: - groups.append(TokenState(current_channel, current_recipient, current_text)) - current_channel = state.channel - current_recipient = state.recipient - current_text = state.text - - groups.append(TokenState(current_channel, current_recipient, current_text)) - - # Process each group and create delta messages - delta_message = None - combined_content = "" - combined_reasoning = "" - tool_messages = [] - content_encountered = False - - # Calculate base_index once before the loop - # This counts completed tool calls in messages - base_index = 0 - for msg in harmony_parser.messages: - if msg.recipient and is_function_recipient(msg.recipient): - base_index += 1 - - # If there's an ongoing tool call from previous chunk, - # the next new tool call starts at base_index + 1 - if prev_recipient and is_function_recipient(prev_recipient): - next_tool_index = base_index + 1 - # Ongoing call is at base_index - ongoing_tool_index = base_index - else: - # No ongoing call, next new call is at base_index - next_tool_index = base_index - ongoing_tool_index = None - - for group in groups: - if group.channel == "final": - combined_content += group.text - content_encountered = True - elif group.recipient and is_function_recipient(group.recipient): - opened_new_call = False - if prev_recipient != group.recipient: - # New tool call - emit the opening message - tool_name = extract_function_from_recipient(group.recipient) - tool_messages.append( - DeltaToolCall( - id=make_tool_call_id(), - type="function", - function=DeltaFunctionCall( - name=tool_name, - arguments="", - ), - index=next_tool_index, - ) - ) - opened_new_call = True - prev_recipient = group.recipient - # Increment for subsequent new tool calls - next_tool_index += 1 - - if group.text: - # Stream arguments for the ongoing tool call - if opened_new_call: - # Just opened in this group - tool_call_index = next_tool_index - 1 - else: - # Continuing from previous chunk - # If ongoing_tool_index is None here, it means - # we're continuing a call but prev_recipient - # wasn't a function. Use base_index. - tool_call_index = ( - ongoing_tool_index - if ongoing_tool_index is not None - else base_index - ) - tool_messages.append( - DeltaToolCall( - index=tool_call_index, - function=DeltaFunctionCall(arguments=group.text), - ) - ) - elif group.channel == "commentary" and group.recipient is None: - # Tool call preambles meant to be shown to the user - combined_content += group.text - content_encountered = True - elif group.channel == "analysis" and include_reasoning: - combined_reasoning += group.text - - # Combine all non-empty fields into a single message - if content_encountered or combined_reasoning or tool_messages: - delta_kwargs: dict[str, str | list[DeltaToolCall]] = {} - if content_encountered: - delta_kwargs["content"] = combined_content - if combined_reasoning: - delta_kwargs["reasoning"] = combined_reasoning - if tool_messages: - delta_kwargs["tool_calls"] = tool_messages - tools_streamed = True - delta_message = DeltaMessage(**delta_kwargs) - else: - delta_message = None - - return delta_message, tools_streamed diff --git a/vllm/parser/harmony.py b/vllm/parser/harmony.py index c1eb7ea042e..f19d3675dab 100644 --- a/vllm/parser/harmony.py +++ b/vllm/parser/harmony.py @@ -9,9 +9,12 @@ from dataclasses import dataclass from enum import Enum, auto from typing import TYPE_CHECKING, NamedTuple +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, FunctionCall, ) from vllm.entrypoints.openai.parser.harmony_utils import ( @@ -52,7 +55,6 @@ class Segment(NamedTuple): channel: str | None recipient: str | None delta: str - is_boundary: bool = False completed_message: Message | None = None @@ -81,10 +83,8 @@ class HarmonyParser(DelegatingParser): ) self._harmony_parser = get_streamable_parser_for_assistant() - - @property - def messages(self) -> list[Message]: - return self._harmony_parser.messages + self._next_tool_call_index = 0 + self._num_processed_messages = 0 @property def state(self) -> HarmonyStreamState: @@ -194,17 +194,69 @@ class HarmonyParser(DelegatingParser): *, finished: bool, ) -> DeltaMessage | None: - raise NotImplementedError( - "HarmonyParser streaming parsing is deferred. " - "Use the existing harmony streaming path." - ) + prev_recipient = self.current_recipient + result = self.process_chunk(delta_token_ids) + combined_content = "" + combined_reasoning = "" + tool_messages: list[DeltaToolCall] = [] + + for segment in result.segments: + if segment.completed_message is not None: + prev_recipient = None + continue + + segment_type = _SegmentType.from_channel_and_recipient( + segment.channel, segment.recipient + ) + match segment_type: + case _SegmentType.REASONING: + combined_reasoning += segment.delta + case _SegmentType.CONTENT: + combined_content += segment.delta + case _SegmentType.TOOL: + assert segment.recipient is not None + if prev_recipient != segment.recipient: + tool_name = extract_function_from_recipient(segment.recipient) + tool_messages.append( + DeltaToolCall( + # HarmonyParser does not use _stream_state; + # "random" tool_call_id_type is always used + id=make_tool_call_id(), + type="function", + function=DeltaFunctionCall( + name=tool_name, + arguments=segment.delta, + ), + index=self._next_tool_call_index, + ) + ) + self._next_tool_call_index += 1 + prev_recipient = segment.recipient + elif segment.delta: + tool_call_index = self._next_tool_call_index - 1 + tool_messages.append( + DeltaToolCall( + index=tool_call_index, + function=DeltaFunctionCall(arguments=segment.delta), + ) + ) + + if not combined_content and not combined_reasoning and not tool_messages: + return None + + delta_message = DeltaMessage() + if combined_content: + delta_message.content = combined_content + if combined_reasoning: + delta_message.reasoning = combined_reasoning + if tool_messages: + delta_message.tool_calls = tool_messages + return delta_message def process_chunk(self, token_ids: Sequence[int]) -> ChunkResult: if not token_ids: return ChunkResult(segments=[], reasoning_token_count=0) - from openai_harmony import StreamState - segments: list[Segment] = [] reasoning_token_count = 0 for token_id in token_ids: @@ -213,9 +265,10 @@ class HarmonyParser(DelegatingParser): recipient = self.current_recipient delta = self._harmony_parser.last_content_delta or "" completed_message = None - is_boundary = self.state == StreamState.EXPECT_START - if is_boundary and self.messages: - completed_message = self.messages[-1] + _messages = self._harmony_parser.messages + if len(_messages) > self._num_processed_messages: + completed_message = _messages[self._num_processed_messages] + self._num_processed_messages += 1 if channel == "analysis" or ( channel == "commentary" and recipient is not None @@ -227,7 +280,6 @@ class HarmonyParser(DelegatingParser): channel=channel, recipient=recipient, delta=delta, - is_boundary=is_boundary, completed_message=completed_message, ) ) From 4bc83323f2ea8e85c87ae5fb5ff2d792a8f61f9d Mon Sep 17 00:00:00 2001 From: littlecircle0730 <43994952+littlecircle0730@users.noreply.github.com> Date: Thu, 11 Jun 2026 19:20:39 -0700 Subject: [PATCH 18/52] [Bugfix] OffloadingConnector: respect skip_reading_prefix_cache flag (#44592) Signed-off-by: Hsiao-Yuan Chen Signed-off-by: littlecircle0730 Signed-off-by: littlecircle0730 <43994952+littlecircle0730@users.noreply.github.com> Co-authored-by: Hsiao-Yuan Chen Co-authored-by: Or Ozeri --- .../offloading_connector/test_scheduler.py | 51 +++++++++++++++++++ .../unit/offloading_connector/utils.py | 6 ++- .../kv_connector/v1/offloading/scheduler.py | 6 ++- 3 files changed, 61 insertions(+), 2 deletions(-) diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py index 20c230a4c2a..11da73b3152 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py @@ -1381,3 +1381,54 @@ def test_stale_sliding_window_block_after_prepare_store_failure( expected_stored=(2, 3), expected_flushed=(2, 3) if not async_scheduling else (), ) + + +@pytest.mark.parametrize("async_scheduling", [True, False]) +def test_skip_reading_prefix_cache(request_runner, async_scheduling: bool): + """When skip_reading_prefix_cache=True, the offloading connector must not + load any blocks from CPU even if a matching prefix is cached there.""" + block_size = 4 + block_size_factor = 3 + offloaded_block_size = block_size * block_size_factor + num_gpu_blocks = 100 + + runner = request_runner( + block_size=block_size, + num_gpu_blocks=num_gpu_blocks, + async_scheduling=async_scheduling, + block_size_factor=block_size_factor, + ) + + # Populate the CPU offload cache with one block. + runner.new_request(token_ids=[0] * offloaded_block_size) + runner.manager.prepare_store.side_effect = lambda keys, req_context: ( + generate_store_output(keys) + ) + runner.run( + decoded_tokens=[EOS_TOKEN_ID], + expected_stored=(0, 1, 2), + expected_flushed=(0, 1, 2) if not async_scheduling else (), + ) + + # Reset GPU prefix cache so the next request cannot hit locally. + runner.scheduler.reset_prefix_cache() + + # New request with identical tokens but skip_reading_prefix_cache=True. + # The offloading connector must not load anything from CPU, but must + # still offload the freshly computed blocks (state management intact). + runner.new_request( + token_ids=[0] * offloaded_block_size, + skip_reading_prefix_cache=True, + ) + runner.manager.prepare_store.side_effect = lambda keys, req_context: ( + generate_store_output(keys) + ) + runner.run( + decoded_tokens=[EOS_TOKEN_ID], + expected_loaded=(), # no CPU loads must happen + expected_stored=(0, 1, 2), # tokens still offloaded to CPU + expected_flushed=(0, 1, 2) if not async_scheduling else (), + ) + + # The external lookup must have been completely skipped. + runner.manager.lookup.assert_not_called() diff --git a/tests/v1/kv_connector/unit/offloading_connector/utils.py b/tests/v1/kv_connector/unit/offloading_connector/utils.py index 22d00b0c834..f6a354ebd43 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/utils.py +++ b/tests/v1/kv_connector/unit/offloading_connector/utils.py @@ -324,10 +324,14 @@ class RequestRunner: self, token_ids: list[int], kv_transfer_params: dict | None = None, + skip_reading_prefix_cache: bool = False, ): self.req_id += 1 - sampling_params = SamplingParams(max_tokens=1000) + sampling_params = SamplingParams( + max_tokens=1000, + skip_reading_prefix_cache=skip_reading_prefix_cache or None, + ) sampling_params.update_from_generation_config({}, EOS_TOKEN_ID) req = Request( diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py index 24e7143e630..94d68972822 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py @@ -571,7 +571,11 @@ class OffloadingConnectorScheduler: req_status.update_offload_keys() req_status.num_locally_computed_tokens = num_computed_tokens - num_hit_tokens = self._lookup(req_status) + num_hit_tokens: int | None + if request.skip_reading_prefix_cache: + num_hit_tokens = 0 + else: + num_hit_tokens = self._lookup(req_status) req_status.update_num_hit_blocks(num_computed_tokens + (num_hit_tokens or 0)) self._touch(req_status) From fcf5115c45b9acfe3a77052ddbb7dfb0f4d5ef18 Mon Sep 17 00:00:00 2001 From: Fangzhou Ai <31551580+Fangzhou-Ai@users.noreply.github.com> Date: Thu, 11 Jun 2026 23:17:52 -0400 Subject: [PATCH 19/52] [ROCm][DSv4][Perf] Flash-decode split-K decode attention kernel (#44899) Co-authored-by: vLLM Contributor --- .../attention/test_rocm_triton_attn_dsv4.py | 140 +++++ .../v1/attention/ops/rocm_aiter_mla_sparse.py | 545 +++++++++++++++++- 2 files changed, 675 insertions(+), 10 deletions(-) diff --git a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py index d4fa9697cb7..f328f339332 100644 --- a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py +++ b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py @@ -10,6 +10,25 @@ pytestmark = pytest.mark.skipif( not current_platform.is_rocm(), reason="Only used by ROCm" ) + +def _on_gfx950() -> bool: + if not current_platform.is_rocm(): + return False + try: + from vllm.platforms.rocm import _ON_GFX950 + + return bool(_ON_GFX950) + except Exception: + return False + + +# The flash-decode split-K decode path is only tuned for AMD gfx950; other +# architectures take the fallback decode kernel, so its tests are skipped there. +requires_gfx950 = pytest.mark.skipif( + not _on_gfx950(), + reason="split-K decode kernel is only tuned for AMD gfx950", +) + NOPE_HEAD_DIM = 448 ROPE_HEAD_DIM = 64 HEAD_DIM = NOPE_HEAD_DIM + ROPE_HEAD_DIM @@ -156,6 +175,20 @@ def _ref_sparse_decode_ragged( return out.to(torch.bfloat16) +def _ragged_from_rows( + rows: list[list[int]], device: torch.device +) -> tuple[torch.Tensor, torch.Tensor]: + """Flatten per-query slot lists into ragged (indices, indptr) tensors.""" + flat = [slot for row in rows for slot in row] + indptr = [0] + for row in rows: + indptr.append(indptr[-1] + len(row)) + return ( + torch.tensor(flat, dtype=torch.int32, device=device), + torch.tensor(indptr, dtype=torch.int32, device=device), + ) + + def _ref_combine_topk_swa_ragged( device: torch.device, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: @@ -375,3 +408,110 @@ def test_combine_topk_swa_indices_ragged() -> None: ) torch.testing.assert_close(actual_indptr, expected_indptr) torch.testing.assert_close(actual_lens, expected_lens) + + +@requires_gfx950 +@torch.inference_mode() +def test_decode_num_splits_heuristic(monkeypatch) -> None: + """Split-count heuristic added with the flash-decode split-K decode path.""" + from vllm.v1.attention.ops import rocm_aiter_mla_sparse as mod + + # Pin the CU count so the heuristic is deterministic off-device. + monkeypatch.setattr(mod, "_decode_cu_count", lambda: 256) + + # A batch that already fills the device should not be split. + assert mod._decode_num_splits(256, 1, avg_main_len=128.0, avg_extra_len=0.0) == 1 + # A tiny batch on a large device should split to add parallelism. + assert mod._decode_num_splits(2, 1, avg_main_len=256.0, avg_extra_len=0.0) > 1 + + # The chosen count always stays within the searched [1, 16] range, and a + # zero-length workload never splits (no work to parallelize). + for num_queries in (1, 4, 24, 224, 1024): + splits = mod._decode_num_splits( + num_queries, 1, avg_main_len=512.0, avg_extra_len=128.0 + ) + assert 1 <= splits <= 16 + assert mod._decode_num_splits(2, 1, avg_main_len=0.0, avg_extra_len=0.0) >= 1 + + +@requires_gfx950 +@pytest.mark.parametrize("num_splits", [1, 2, 3, 4, 8]) +@pytest.mark.parametrize("with_extra", [True, False]) +@pytest.mark.parametrize("with_sink", [True, False]) +@torch.inference_mode() +def test_sparse_attn_decode_split_k_kernel( + monkeypatch, num_splits: int, with_extra: bool, with_sink: bool +) -> None: + """Flash-decode split-K decode path (partial + reduce kernels). + + This path is the gfx950 production path (``_ON_GFX950``), so the test only + runs on gfx950. The split count is pinned so the partial/reduce kernels are + exercised across split counts. ``num_splits=8`` drives splits past the + shortest segment length, covering the empty-split edge case handled by the + reduce kernel. + """ + from vllm.v1.attention.ops import rocm_aiter_mla_sparse as mod + + device = torch.device("cuda") + torch.manual_seed(7) + block_size = 4 + num_heads = 3 + + main_rows = [[0, 2, 4, 6, 1, 3, 7, 5], [4, 1, 6, 0, 2]] + num_queries = len(main_rows) + q = ( + torch.randn( + num_queries, num_heads, HEAD_DIM, dtype=torch.bfloat16, device=device + ) + * 0.125 + ) + main_kv = torch.randn(8, HEAD_DIM, dtype=torch.bfloat16, device=device) * 0.125 + main_cache = _pack_fp8_ds_mla_cache(main_kv, block_size) + main_indices, main_indptr = _ragged_from_rows(main_rows, device) + + extra_rows: list[list[int]] | None = None + extra_cache: torch.Tensor | None = None + extra_indices: torch.Tensor | None = None + extra_indptr: torch.Tensor | None = None + if with_extra: + rows = [[1, 3, 0, 5, 2, 4], [3, 0, 6]] + extra_kv = torch.randn(7, HEAD_DIM, dtype=torch.bfloat16, device=device) * 0.125 + extra_rows = rows + extra_cache = _pack_fp8_ds_mla_cache(extra_kv, block_size) + extra_indices, extra_indptr = _ragged_from_rows(rows, device) + + attn_sink = ( + torch.tensor([-0.1, 0.0, 0.1], dtype=torch.float32, device=device) + if with_sink + else None + ) + scale = HEAD_DIM**-0.5 + + # Pin the split count so each parametrized value is exercised deterministically. + monkeypatch.setattr(mod, "_decode_num_splits", lambda *args, **kwargs: num_splits) + + actual = mod._rocm_sparse_attn_decode_ragged_triton( + q=q, + main_cache=main_cache, + main_indices=main_indices, + main_indptr=main_indptr, + scale=scale, + attn_sink=attn_sink, + nope_head_dim=NOPE_HEAD_DIM, + rope_head_dim=ROPE_HEAD_DIM, + extra_cache=extra_cache, + extra_indices=extra_indices, + extra_indptr=extra_indptr, + ) + expected = _ref_sparse_decode_ragged( + q=q, + main_cache=main_cache, + main_rows=main_rows, + scale=scale, + attn_sink=attn_sink, + block_size=block_size, + extra_cache=extra_cache, + extra_rows=extra_rows, + ) + + torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2) diff --git a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py index 12fd3a17421..8104e808f67 100644 --- a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py @@ -1406,6 +1406,348 @@ def _sparse_attn_decode_ragged_kernel( ) +@triton.jit +def _sparse_attn_decode_partial_kernel( + q_ptr, + main_cache_ptr, + main_indices_ptr, + main_indptr_ptr, + extra_cache_ptr, + extra_indices_ptr, + extra_indptr_ptr, + part_m_ptr, + part_l_ptr, + part_acc_ptr, + q_stride0, + q_stride1, + main_cache_stride0, + extra_cache_stride0, + pm_stride0, + pm_stride_s, + pa_stride0, + pa_stride_s, + pa_stride_h, + main_num_rows, + extra_num_rows, + main_block_size, + extra_block_size, + scale, + num_heads, + HAS_EXTRA: tl.constexpr, + NOPE_DIM: tl.constexpr, + NOPE_BLOCK: tl.constexpr, + ROPE_DIM: tl.constexpr, + IS_FNUZ: tl.constexpr, + BLOCK_H: tl.constexpr, + BLOCK_K: tl.constexpr, + NUM_SPLITS: tl.constexpr, + NUM_STAGES: tl.constexpr, +): + query_idx = tl.program_id(0) + split_id = tl.program_id(1) + pid_h = tl.program_id(2) + + head_offsets = pid_h * BLOCK_H + tl.arange(0, BLOCK_H) + head_mask = head_offsets < num_heads + nope_offsets = tl.arange(0, NOPE_BLOCK) + nope_mask = nope_offsets < NOPE_DIM + rope_offsets = tl.arange(0, ROPE_DIM) + + q_row_ptr = q_ptr + query_idx * q_stride0 + head_offsets[:, None] * q_stride1 + q_nope = tl.load( + q_row_ptr + nope_offsets[None, :], + mask=head_mask[:, None] & nope_mask[None, :], + other=0.0, + ) + q_rope = tl.load( + q_row_ptr + NOPE_DIM + rope_offsets[None, :], + mask=head_mask[:, None], + other=0.0, + ) + + neg_large = -3.4028234663852886e38 + m_i = tl.full((BLOCK_H,), neg_large, dtype=tl.float32) + l_i = tl.zeros((BLOCK_H,), dtype=tl.float32) + acc_nope = tl.zeros((BLOCK_H, NOPE_BLOCK), dtype=tl.float32) + acc_rope = tl.zeros((BLOCK_H, ROPE_DIM), dtype=tl.float32) + k_offsets = tl.arange(0, BLOCK_K) + + zero_nope = tl.zeros((BLOCK_K, NOPE_BLOCK), dtype=tl.bfloat16) + zero_rope = tl.zeros((BLOCK_K, ROPE_DIM), dtype=tl.bfloat16) + + # Each split processes a contiguous slice of this query's main (SWA) and + # extra (topk) segments. Slices are handled independently so a block never + # straddles the main/extra boundary. + main_start = tl.load(main_indptr_ptr + query_idx) + main_end = tl.load(main_indptr_ptr + query_idx + 1) + main_len = main_end - main_start + main_chunk = (main_len + NUM_SPLITS - 1) // NUM_SPLITS + main_lo = split_id * main_chunk + main_hi = tl.minimum(main_lo + main_chunk, main_len) + + for k_start in tl.range(main_lo, main_hi, BLOCK_K, num_stages=NUM_STAGES): + k_pos = k_start + k_offsets + in_range = k_pos < main_hi + slot = tl.load(main_indices_ptr + main_start + k_pos, mask=in_range, other=-1) + valid = in_range & (slot >= 0) & (slot < main_num_rows) + safe_slot = tl.where(valid, slot, 0) + + block_idx = safe_slot // main_block_size + pos_in_block = safe_slot % main_block_size + cache_block_ptr = main_cache_ptr + block_idx.to(tl.int64) * main_cache_stride0 + token_data_ptr = cache_block_ptr + pos_in_block * 576 + token_scale_ptr = cache_block_ptr + main_block_size * 576 + pos_in_block * 8 + + x_uint8 = tl.load( + token_data_ptr[:, None] + nope_offsets[None, :], + mask=valid[:, None] & nope_mask[None, :], + other=0, + ) + if IS_FNUZ: + x_fp8 = x_uint8.to(tl.float8e4b15, bitcast=True) + else: + x_fp8 = x_uint8.to(tl.float8e4nv, bitcast=True) + encoded_scales = tl.load( + token_scale_ptr[:, None] + nope_offsets[None, :] // 64, + mask=valid[:, None] & nope_mask[None, :], + other=127, + ) + scales = tl.exp2(encoded_scales.to(tl.float32) - 127.0) + k_nope = x_fp8.to(tl.bfloat16) * scales.to(tl.bfloat16) + k_nope = tl.where(valid[:, None] & nope_mask[None, :], k_nope, zero_nope) + k_nope = tl.where(k_nope == k_nope, k_nope, zero_nope) + + rope_ptr = (token_data_ptr + NOPE_DIM).to(tl.pointer_type(tl.bfloat16)) + k_rope = tl.load( + rope_ptr[:, None] + rope_offsets[None, :], + mask=valid[:, None], + other=0.0, + ) + k_rope = tl.where(valid[:, None], k_rope, zero_rope) + k_rope = tl.where(k_rope == k_rope, k_rope, zero_rope) + + scores = tl.dot(q_nope, tl.trans(k_nope)) + tl.dot(q_rope, tl.trans(k_rope)) + scores *= scale + scores = tl.where(head_mask[:, None] & valid[None, :], scores, neg_large) + + m_block = tl.max(scores, axis=1) + m_new = tl.maximum(m_i, m_block) + alpha = tl.exp(m_i - m_new) + p = tl.exp(scores - m_new[:, None]) + p = tl.where(head_mask[:, None] & valid[None, :], p, 0.0) + l_new = l_i * alpha + tl.sum(p, axis=1) + + acc_nope = acc_nope * alpha[:, None] + tl.dot(p.to(k_nope.dtype), k_nope) + acc_rope = acc_rope * alpha[:, None] + tl.dot(p.to(k_rope.dtype), k_rope) + m_i = m_new + l_i = l_new + + if HAS_EXTRA: + extra_start = tl.load(extra_indptr_ptr + query_idx) + extra_end = tl.load(extra_indptr_ptr + query_idx + 1) + extra_len = extra_end - extra_start + extra_chunk = (extra_len + NUM_SPLITS - 1) // NUM_SPLITS + extra_lo = split_id * extra_chunk + extra_hi = tl.minimum(extra_lo + extra_chunk, extra_len) + + for k_start in tl.range(extra_lo, extra_hi, BLOCK_K, num_stages=NUM_STAGES): + k_pos = k_start + k_offsets + in_range = k_pos < extra_hi + slot = tl.load( + extra_indices_ptr + extra_start + k_pos, mask=in_range, other=-1 + ) + valid = in_range & (slot >= 0) & (slot < extra_num_rows) + safe_slot = tl.where(valid, slot, 0) + + block_idx = safe_slot // extra_block_size + pos_in_block = safe_slot % extra_block_size + cache_block_ptr = ( + extra_cache_ptr + block_idx.to(tl.int64) * extra_cache_stride0 + ) + token_data_ptr = cache_block_ptr + pos_in_block * 576 + token_scale_ptr = ( + cache_block_ptr + extra_block_size * 576 + pos_in_block * 8 + ) + + x_uint8 = tl.load( + token_data_ptr[:, None] + nope_offsets[None, :], + mask=valid[:, None] & nope_mask[None, :], + other=0, + ) + if IS_FNUZ: + x_fp8 = x_uint8.to(tl.float8e4b15, bitcast=True) + else: + x_fp8 = x_uint8.to(tl.float8e4nv, bitcast=True) + encoded_scales = tl.load( + token_scale_ptr[:, None] + nope_offsets[None, :] // 64, + mask=valid[:, None] & nope_mask[None, :], + other=127, + ) + scales = tl.exp2(encoded_scales.to(tl.float32) - 127.0) + k_nope = x_fp8.to(tl.bfloat16) * scales.to(tl.bfloat16) + k_nope = tl.where(valid[:, None] & nope_mask[None, :], k_nope, zero_nope) + k_nope = tl.where(k_nope == k_nope, k_nope, zero_nope) + + rope_ptr = (token_data_ptr + NOPE_DIM).to(tl.pointer_type(tl.bfloat16)) + k_rope = tl.load( + rope_ptr[:, None] + rope_offsets[None, :], + mask=valid[:, None], + other=0.0, + ) + k_rope = tl.where(valid[:, None], k_rope, zero_rope) + k_rope = tl.where(k_rope == k_rope, k_rope, zero_rope) + + scores = tl.dot(q_nope, tl.trans(k_nope)) + tl.dot( + q_rope, + tl.trans(k_rope), + ) + scores *= scale + scores = tl.where(head_mask[:, None] & valid[None, :], scores, neg_large) + + m_block = tl.max(scores, axis=1) + m_new = tl.maximum(m_i, m_block) + alpha = tl.exp(m_i - m_new) + p = tl.exp(scores - m_new[:, None]) + p = tl.where(head_mask[:, None] & valid[None, :], p, 0.0) + l_new = l_i * alpha + tl.sum(p, axis=1) + + acc_nope = acc_nope * alpha[:, None] + tl.dot(p.to(k_nope.dtype), k_nope) + acc_rope = acc_rope * alpha[:, None] + tl.dot(p.to(k_rope.dtype), k_rope) + m_i = m_new + l_i = l_new + + # Store raw (un-normalized) partial state for this split. Softmax sink and + # final normalization happen in the reduce kernel. + pm_base = query_idx * pm_stride0 + split_id * pm_stride_s + head_offsets + tl.store(part_m_ptr + pm_base, m_i, mask=head_mask) + tl.store(part_l_ptr + pm_base, l_i, mask=head_mask) + acc_base = ( + part_acc_ptr + + query_idx * pa_stride0 + + split_id * pa_stride_s + + head_offsets[:, None] * pa_stride_h + ) + tl.store( + acc_base + nope_offsets[None, :], + acc_nope, + mask=head_mask[:, None] & nope_mask[None, :], + ) + tl.store( + acc_base + NOPE_DIM + rope_offsets[None, :], + acc_rope, + mask=head_mask[:, None], + ) + + +@triton.jit +def _sparse_attn_decode_reduce_kernel( + part_m_ptr, + part_l_ptr, + part_acc_ptr, + attn_sink_ptr, + out_ptr, + out_stride0, + out_stride1, + pm_stride0, + pm_stride_s, + pa_stride0, + pa_stride_s, + pa_stride_h, + num_heads, + HAS_ATTN_SINK: tl.constexpr, + COMB_DIM: tl.constexpr, + BLOCK_H: tl.constexpr, + NUM_SPLITS: tl.constexpr, + SPLITS_PAD: tl.constexpr, +): + query_idx = tl.program_id(0) + pid_h = tl.program_id(1) + + head_offsets = pid_h * BLOCK_H + tl.arange(0, BLOCK_H) + head_mask = head_offsets < num_heads + comb_offsets = tl.arange(0, COMB_DIM) + # SPLITS_PAD is NUM_SPLITS rounded up to a power of two so the parallel + # split-axis load is a legal arange for any split count; padding lanes are + # masked off. + split_offsets = tl.arange(0, SPLITS_PAD) + split_mask = split_offsets < NUM_SPLITS + + neg_large = -3.4028234663852886e38 + + # Phase 1: load every split's running max/sum at once and reduce the max + # in parallel (tl.max over the split axis) instead of walking the splits + # serially. This breaks the long online-softmax dependency chain that made + # the reduce latency-bound. + load_mask = split_mask[:, None] & head_mask[None, :] + pm_split = ( + part_m_ptr + + query_idx * pm_stride0 + + split_offsets[:, None] * pm_stride_s + + head_offsets[None, :] + ) + m_all = tl.load(pm_split, mask=load_mask, other=neg_large) # [S, H] + l_all = tl.load( + part_l_ptr + + query_idx * pm_stride0 + + split_offsets[:, None] * pm_stride_s + + head_offsets[None, :], + mask=load_mask, + other=0.0, + ) + + m_comb = tl.max(m_all, axis=0) # [H] + if HAS_ATTN_SINK: + sink = tl.load( + attn_sink_ptr + head_offsets, mask=head_mask, other=neg_large + ).to(tl.float32) + m_final = tl.maximum(m_comb, sink) + else: + m_final = m_comb + + w_all = tl.exp(m_all - m_final[None, :]) # [S, H] + w_all = tl.where(load_mask, w_all, 0.0) + l_final = tl.sum(w_all * l_all, axis=0) # [H] + if HAS_ATTN_SINK: + l_final = l_final + tl.exp(sink - m_final) + denom = tl.maximum(l_final, 1.0e-30) + + # Phase 2: weighted sum of the per-split accumulators. The combine weight + # for each split only depends on the (already known) global max, so the + # acc loads carry no cross-split dependency and the compiler can pipeline + # them; only the cheap FMA into `acc` is loop-carried. + acc = tl.zeros((BLOCK_H, COMB_DIM), dtype=tl.float32) + for s in tl.static_range(NUM_SPLITS): + m_s = tl.load( + part_m_ptr + query_idx * pm_stride0 + s * pm_stride_s + head_offsets, + mask=head_mask, + other=neg_large, + ) + w_s = tl.exp(m_s - m_final) + acc_base = ( + part_acc_ptr + + query_idx * pa_stride0 + + s * pa_stride_s + + head_offsets[:, None] * pa_stride_h + ) + acc_s = tl.load( + acc_base + comb_offsets[None, :], + mask=head_mask[:, None], + other=0.0, + ) + acc += w_s[:, None] * acc_s + + out = tl.where(l_final[:, None] > 0.0, acc / denom[:, None], 0.0) + + out_row_ptr = ( + out_ptr + query_idx * out_stride0 + head_offsets[:, None] * out_stride1 + ) + tl.store( + out_row_ptr + comb_offsets[None, :], + out, + mask=head_mask[:, None], + ) + + def _rocm_sparse_attn_prefill_ragged_triton( q: torch.Tensor, kv: torch.Tensor, @@ -1502,6 +1844,101 @@ def _rocm_sparse_attn_prefill_triton( ) +@functools.lru_cache +def _decode_cu_count() -> int: + try: + return torch.cuda.get_device_properties(0).multi_processor_count + except Exception: + return 256 # For gfx950 arch, gated behind a fallback path for other archs. + + +def _decode_partial_iters( + avg_main_len: float, avg_extra_len: float, splits: int, block_k: int +) -> int: + """BLOCK_K iterations one partial workgroup walks for ``splits`` splits. + + Each split processes ``ceil(seg_len / splits)`` tokens of a segment, walked + ``BLOCK_K`` at a time, and the main/extra segments are handled separately. + """ + main_iters = ( + math.ceil(math.ceil(avg_main_len / splits) / block_k) if avg_main_len > 0 else 0 + ) + extra_iters = ( + math.ceil(math.ceil(avg_extra_len / splits) / block_k) + if avg_extra_len > 0 + else 0 + ) + return main_iters + extra_iters + + +def _decode_num_splits( + num_queries: int, + heads_blocks: int, + avg_main_len: float = 0.0, + avg_extra_len: float = 0.0, + block_k: int = 32, +) -> int: + """Pick a flash-decode split count to keep the GPU busy across batch sizes. + + Decode launches only ``num_queries * heads_blocks`` workgroups otherwise, + which severely under-fills the device for the low-concurrency regime that + dominates latency. Splitting the KV sequence adds parallelism. + + We model the relative partial-kernel latency for a given split count ``s`` + as ``waves * (1/s + mu)`` where ``waves = ceil(base * s / CU)`` and ``mu`` + is a small per-wave overhead penalty: + + - ``waves / s`` captures the partial compute: each wave walks roughly + ``total_tokens / s`` tokens and there are ``waves`` of them, so dividing + by ``s`` makes more splits cheaper *until* they spill into extra waves. + - ``mu * waves`` charges per-wave launch/tail overhead so we do not + over-split into many mostly-idle waves (e.g. batch 224 on 256 CUs is + best left at 1 split rather than 8 splits across 7 waves). + + The minimiser naturally prefers split counts that pack the device into full + waves (``base * s`` near a multiple of ``CU``) and falls back to 1 split + once the batch already fills the device. Ties favour the smaller split + count (less reduce work). + + Finally we "snap down" the chosen split count to the smallest value that + yields the same wave count *and* the same per-workgroup BLOCK_K iteration + count. Because latency tracks iteration count (not raw token count), extra + splits that do not lower the iteration count add only reduce/HBM overhead + for no parallelism gain (e.g. batch 24: s8 and s10 both walk 4 extra iters + in one wave, so s8 is strictly better). Snapping needs the average segment + lengths, which the caller derives sync-free from the ragged index sizes. + """ + base = max(1, num_queries * heads_blocks) + # Target ~1 workgroup per CU: enough to fill the device while keeping the + # reduce cost (which grows with split count) small. Tuned on gfx950. + cu = max(1, _decode_cu_count()) + # Per-wave overhead penalty: higher values discourage split counts that + # spill into extra GPU waves. Tuned on gfx950. + mu = 0.04 + best_splits = 1 + best_cost = None + # Search up to 16 splits; beyond that the reduce/HBM overhead dominates. + for splits in range(1, 17): + waves = (base * splits + cu - 1) // cu + cost = waves * (1.0 / splits + mu) + if best_cost is None or cost < best_cost - 1e-9: + best_splits = splits + best_cost = cost + + if best_splits > 1 and (avg_main_len > 0 or avg_extra_len > 0): + target_waves = (base * best_splits + cu - 1) // cu + target_iters = _decode_partial_iters( + avg_main_len, avg_extra_len, best_splits, block_k + ) + for splits in range(1, best_splits): + waves = (base * splits + cu - 1) // cu + iters = _decode_partial_iters(avg_main_len, avg_extra_len, splits, block_k) + if waves == target_waves and iters == target_iters: + best_splits = splits + break + return best_splits + + def _rocm_sparse_attn_decode_ragged_triton( q: torch.Tensor, main_cache: torch.Tensor, @@ -1575,9 +2012,70 @@ def _rocm_sparse_attn_decode_ragged_triton( extra_indptr = torch.zeros(num_queries + 1, device=q.device, dtype=torch.int32) block_h = 16 - block_k = 16 if head_dim >= 256 else 32 out = torch.empty_like(q, dtype=torch.bfloat16) - _sparse_attn_decode_ragged_kernel[(num_queries, triton.cdiv(num_heads, block_h))]( + heads_blocks = triton.cdiv(num_heads, block_h) + nope_block = triton.next_power_of_2(nope_head_dim) + comb_dim = nope_head_dim + rope_head_dim + is_fnuz = current_platform.is_fp8_fnuz() + + if not _ON_GFX950: # Fallback path for un-tuned architectures. + block_k = 16 if head_dim >= 256 else 32 + _sparse_attn_decode_ragged_kernel[(num_queries, heads_blocks)]( + q, + main_cache, + main_indices, + main_indptr, + extra_cache, + extra_indices, + extra_indptr, + attn_sink, + out, + q.stride(0), + q.stride(1), + out.stride(0), + out.stride(1), + main_cache.stride(0), + extra_cache.stride(0), + main_cache.shape[0] * main_cache.shape[1], + extra_cache.shape[0] * extra_cache.shape[1], + main_cache.shape[1], + extra_cache.shape[1], + scale, + num_heads, + HAS_ATTN_SINK=has_attn_sink, + HAS_EXTRA=has_extra, + NOPE_DIM=nope_head_dim, + NOPE_BLOCK=nope_block, + ROPE_DIM=rope_head_dim, + IS_FNUZ=is_fnuz, + BLOCK_H=block_h, + BLOCK_K=block_k, + num_warps=8, + ) + return out + + block_k = 32 # KV tokens walked per split-K iteration. Tuned on gfx950. + # Average per-query segment lengths, read sync-free from the ragged index + # sizes, let the split heuristic avoid over-splitting + # main_indices/extra_indices are flat [nnz] int32. + inv_q = 1.0 / max(1, num_queries) + avg_main_len = main_indices.numel() * inv_q + avg_extra_len = (extra_indices.numel() * inv_q) if has_extra else 0.0 + num_splits = _decode_num_splits( + num_queries, heads_blocks, avg_main_len, avg_extra_len, block_k + ) + + part_m = torch.empty( + (num_queries, num_splits, num_heads), dtype=torch.float32, device=q.device + ) + part_l = torch.empty_like(part_m) + part_acc = torch.empty( + (num_queries, num_splits, num_heads, comb_dim), + dtype=torch.float32, + device=q.device, + ) + + _sparse_attn_decode_partial_kernel[(num_queries, num_splits, heads_blocks)]( q, main_cache, main_indices, @@ -1585,29 +2083,56 @@ def _rocm_sparse_attn_decode_ragged_triton( extra_cache, extra_indices, extra_indptr, - attn_sink, - out, + part_m, + part_l, + part_acc, q.stride(0), q.stride(1), - out.stride(0), - out.stride(1), main_cache.stride(0), extra_cache.stride(0), + part_m.stride(0), + part_m.stride(1), + part_acc.stride(0), + part_acc.stride(1), + part_acc.stride(2), main_cache.shape[0] * main_cache.shape[1], extra_cache.shape[0] * extra_cache.shape[1], main_cache.shape[1], extra_cache.shape[1], scale, num_heads, - HAS_ATTN_SINK=has_attn_sink, HAS_EXTRA=has_extra, NOPE_DIM=nope_head_dim, - NOPE_BLOCK=triton.next_power_of_2(nope_head_dim), + NOPE_BLOCK=nope_block, ROPE_DIM=rope_head_dim, - IS_FNUZ=current_platform.is_fp8_fnuz(), + IS_FNUZ=is_fnuz, BLOCK_H=block_h, BLOCK_K=block_k, - num_warps=8, + NUM_SPLITS=num_splits, + NUM_STAGES=1, + num_warps=4, + ) + + _sparse_attn_decode_reduce_kernel[(num_queries, heads_blocks)]( + part_m, + part_l, + part_acc, + attn_sink, + out, + out.stride(0), + out.stride(1), + part_m.stride(0), + part_m.stride(1), + part_acc.stride(0), + part_acc.stride(1), + part_acc.stride(2), + num_heads, + HAS_ATTN_SINK=has_attn_sink, + COMB_DIM=comb_dim, + BLOCK_H=block_h, + NUM_SPLITS=num_splits, + SPLITS_PAD=triton.next_power_of_2(num_splits), + num_warps=4, ) return out From c1076839c9f14a51c0eb963ca8ae12c2de3c0f63 Mon Sep 17 00:00:00 2001 From: Ting SUN Date: Fri, 12 Jun 2026 11:21:46 +0800 Subject: [PATCH 20/52] [Bugfix][Model] Pass revision by name in Run:ai and bitsandbytes index downloads (#45308) Signed-off-by: Ting Sun --- .../test_runai_model_streamer_loader.py | 25 +++++++++++++++++ .../models/quantization/test_bitsandbytes.py | 28 +++++++++++++++++++ .../model_loader/bitsandbytes_loader.py | 4 +-- .../model_loader/runai_streamer_loader.py | 5 +++- 4 files changed, 59 insertions(+), 3 deletions(-) diff --git a/tests/model_executor/model_loader/runai_streamer_loader/test_runai_model_streamer_loader.py b/tests/model_executor/model_loader/runai_streamer_loader/test_runai_model_streamer_loader.py index c7158dae537..82c0f8813e2 100644 --- a/tests/model_executor/model_loader/runai_streamer_loader/test_runai_model_streamer_loader.py +++ b/tests/model_executor/model_loader/runai_streamer_loader/test_runai_model_streamer_loader.py @@ -1,11 +1,15 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import types +from unittest.mock import patch + import pytest from vllm import SamplingParams from vllm.config.load import LoadConfig from vllm.model_executor.model_loader import get_model_loader +from vllm.model_executor.model_loader import runai_streamer_loader as rsl load_format = "runai_streamer" test_model = "openai-community/gpt2" @@ -53,3 +57,24 @@ def test_runai_model_loader_download_files_gcs( with vllm_runner(test_gcs_model, load_format=load_format) as llm: deserialized_outputs = llm.generate(prompts, sampling_params) assert deserialized_outputs + + +def test_runai_passes_revision_by_name(): + # revision must reach download_safetensors_index_file_from_hf as the + # ``revision`` keyword, not the positional ``subfolder`` slot. + fake_self = types.SimpleNamespace( + load_config=types.SimpleNamespace(download_dir="/cache", ignore_patterns=[]) + ) + with ( + patch.object(rsl, "is_runai_obj_uri", return_value=False), + patch.object(rsl, "download_weights_from_hf", return_value="/folder"), + patch.object( + rsl, "list_safetensors", return_value=["/folder/model.safetensors"] + ), + patch.object(rsl, "download_safetensors_index_file_from_hf") as mock_idx, + ): + rsl.RunaiModelStreamerLoader._prepare_weights(fake_self, "org/model", "myrev") + + mock_idx.assert_called_once() + assert mock_idx.call_args.kwargs.get("revision") == "myrev" + assert "myrev" not in mock_idx.call_args.args diff --git a/tests/models/quantization/test_bitsandbytes.py b/tests/models/quantization/test_bitsandbytes.py index d6f2b86c7af..03c19b0bf62 100644 --- a/tests/models/quantization/test_bitsandbytes.py +++ b/tests/models/quantization/test_bitsandbytes.py @@ -5,12 +5,16 @@ Run `pytest tests/quantization/test_bitsandbytes.py`. """ +import types +from unittest.mock import MagicMock, patch + import pytest from packaging.version import Version from transformers import BitsAndBytesConfig from transformers import __version__ as TRANSFORMERS_VERSION from tests.quantization.utils import is_quant_method_supported +from vllm.model_executor.model_loader import bitsandbytes_loader as bnb from vllm.platforms import current_platform from ...utils import compare_two_settings, multi_gpu_test @@ -300,3 +304,27 @@ def validate_generated_texts( f"HF Output: '{hf_str}'\n" f"vLLM Output: '{vllm_str}'" ) + + +def test_bitsandbytes_passes_revision_by_name(): + # revision must reach download_safetensors_index_file_from_hf as the + # ``revision`` keyword, not a positional slot. + fake_self = types.SimpleNamespace( + load_config=types.SimpleNamespace(download_dir="/cache"), + _get_weight_files=MagicMock( + return_value=("/folder", ["/folder/model.safetensors"], "*.safetensors") + ), + ) + with ( + patch.object(bnb, "download_safetensors_index_file_from_hf") as mock_idx, + patch.object( + bnb, + "filter_duplicate_safetensors_files", + return_value=["/folder/model.safetensors"], + ), + ): + bnb.BitsAndBytesModelLoader._prepare_weights(fake_self, "org/model", "myrev") + + mock_idx.assert_called_once() + assert mock_idx.call_args.kwargs.get("revision") == "myrev" + assert "myrev" not in mock_idx.call_args.args diff --git a/vllm/model_executor/model_loader/bitsandbytes_loader.py b/vllm/model_executor/model_loader/bitsandbytes_loader.py index d10f3bfcbe9..064a74023a2 100644 --- a/vllm/model_executor/model_loader/bitsandbytes_loader.py +++ b/vllm/model_executor/model_loader/bitsandbytes_loader.py @@ -140,8 +140,8 @@ class BitsAndBytesModelLoader(BaseModelLoader): download_safetensors_index_file_from_hf( model_name_or_path, index_file, - self.load_config.download_dir, - revision, + cache_dir=self.load_config.download_dir, + revision=revision, ) hf_weights_files = filter_duplicate_safetensors_files( hf_weights_files, hf_folder, index_file diff --git a/vllm/model_executor/model_loader/runai_streamer_loader.py b/vllm/model_executor/model_loader/runai_streamer_loader.py index 47c3c99b19a..0df14227919 100644 --- a/vllm/model_executor/model_loader/runai_streamer_loader.py +++ b/vllm/model_executor/model_loader/runai_streamer_loader.py @@ -70,7 +70,10 @@ class RunaiModelStreamerLoader(BaseModelLoader): if not is_local and not is_object_storage_path: download_safetensors_index_file_from_hf( - model_name_or_path, index_file, self.load_config.download_dir, revision + model_name_or_path, + index_file, + cache_dir=self.load_config.download_dir, + revision=revision, ) if not hf_weights_files: From 2263f8a3de64f4cf16488fb43369714c736612a0 Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Thu, 11 Jun 2026 20:26:17 -0700 Subject: [PATCH 21/52] [CI][BugFix] Fix broken `test_mamba_prefix_cache.py` due to stale mock (#45345) Signed-off-by: Nick Hill --- tests/v1/e2e/general/test_mamba_prefix_cache.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/v1/e2e/general/test_mamba_prefix_cache.py b/tests/v1/e2e/general/test_mamba_prefix_cache.py index ceae041c6f9..e857b127285 100644 --- a/tests/v1/e2e/general/test_mamba_prefix_cache.py +++ b/tests/v1/e2e/general/test_mamba_prefix_cache.py @@ -181,6 +181,7 @@ def get_fake_allocate_slots_fn(original_allocate_slots_fn: Callable): num_encoder_tokens: int = 0, full_sequence_must_fit: bool = False, reserved_blocks: int = 0, + has_scheduled_reqs: bool = True, ): ret = original_allocate_slots_fn( self, @@ -194,6 +195,7 @@ def get_fake_allocate_slots_fn(original_allocate_slots_fn: Callable): num_encoder_tokens, full_sequence_must_fit, reserved_blocks, + has_scheduled_reqs, ) if cur_step_action is not None: cur_block_ids = self.coordinator.single_type_managers[0].req_to_blocks[ From 42ae5e7ac61910815bf368da22f67a721179ee45 Mon Sep 17 00:00:00 2001 From: sasindharan <117493393+sasindharan@users.noreply.github.com> Date: Fri, 12 Jun 2026 09:07:42 +0530 Subject: [PATCH 22/52] [Bugfix] Fix --enable-prompt-tokens-details omitting zero cached tokens (#44383) Signed-off-by: Sasindharan Sankar Co-authored-by: Sasindharan Sankar Co-authored-by: Chauncey --- .../openai/completion/test_completion.py | 9 ++-- .../serve/disagg/test_generate_stream.py | 43 +++++++++++++++++++ .../openai/chat_completion/serving.py | 7 ++- vllm/entrypoints/openai/completion/serving.py | 4 +- vllm/entrypoints/serve/disagg/serving.py | 7 ++- 5 files changed, 61 insertions(+), 9 deletions(-) diff --git a/tests/entrypoints/openai/completion/test_completion.py b/tests/entrypoints/openai/completion/test_completion.py index 8ca0d1604b1..a16fa83fe32 100644 --- a/tests/entrypoints/openai/completion/test_completion.py +++ b/tests/entrypoints/openai/completion/test_completion.py @@ -58,9 +58,12 @@ async def test_single_completion(client: openai.AsyncOpenAI, model_name: str) -> choice = completion.choices[0] assert len(choice.text) >= 5 assert choice.finish_reason == "length" - assert completion.usage == openai.types.CompletionUsage( - completion_tokens=5, prompt_tokens=6, total_tokens=11 - ) + assert completion.usage is not None + assert completion.usage.completion_tokens == 5 + assert completion.usage.prompt_tokens == 6 + assert completion.usage.total_tokens == 11 + assert completion.usage.prompt_tokens_details is not None + assert completion.usage.prompt_tokens_details.cached_tokens == 0 # test using token IDs completion = await client.completions.create( diff --git a/tests/entrypoints/serve/disagg/test_generate_stream.py b/tests/entrypoints/serve/disagg/test_generate_stream.py index ac5b8bcd915..bd52863342d 100644 --- a/tests/entrypoints/serve/disagg/test_generate_stream.py +++ b/tests/entrypoints/serve/disagg/test_generate_stream.py @@ -512,3 +512,46 @@ async def test_stream_prompt_tokens_details(): usage_chunk = parsed[-2] assert usage_chunk["choices"] == [] assert usage_chunk["usage"]["prompt_tokens_details"]["cached_tokens"] == 2 + + +@pytest.mark.asyncio +async def test_stream_prompt_tokens_details_zero_cached(): + """enable_prompt_tokens_details includes cached_tokens=0 in final usage. + + Regression test for https://github.com/vllm-project/vllm/issues/44377: + zero cached tokens must not be treated as falsy and omitted. + """ + 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=0, + ) + + 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"] == [] + # Zero cached tokens must be present, not omitted + assert usage_chunk["usage"]["prompt_tokens_details"] is not None + assert usage_chunk["usage"]["prompt_tokens_details"]["cached_tokens"] == 0 diff --git a/vllm/entrypoints/openai/chat_completion/serving.py b/vllm/entrypoints/openai/chat_completion/serving.py index 52d18519eff..45b79c6a7ef 100644 --- a/vllm/entrypoints/openai/chat_completion/serving.py +++ b/vllm/entrypoints/openai/chat_completion/serving.py @@ -732,7 +732,7 @@ class OpenAIServingChat(OpenAIServing): completion_tokens=completion_tokens, total_tokens=num_prompt_tokens + completion_tokens, ) - if self.enable_prompt_tokens_details and num_cached_tokens: + if self.enable_prompt_tokens_details and num_cached_tokens is not None: final_usage.prompt_tokens_details = PromptTokenUsageInfo( cached_tokens=num_cached_tokens ) @@ -1023,7 +1023,10 @@ class OpenAIServingChat(OpenAIServing): completion_tokens=num_generated_tokens, total_tokens=num_prompt_tokens + num_generated_tokens, ) - if self.enable_prompt_tokens_details and final_res.num_cached_tokens: + if ( + self.enable_prompt_tokens_details + and final_res.num_cached_tokens is not None + ): usage.prompt_tokens_details = PromptTokenUsageInfo( cached_tokens=final_res.num_cached_tokens ) diff --git a/vllm/entrypoints/openai/completion/serving.py b/vllm/entrypoints/openai/completion/serving.py index ed85323d806..bd7e26b2b16 100644 --- a/vllm/entrypoints/openai/completion/serving.py +++ b/vllm/entrypoints/openai/completion/serving.py @@ -443,7 +443,7 @@ class OpenAIServingCompletion(OpenAIServing): total_tokens=total_prompt_tokens + total_completion_tokens, ) - if self.enable_prompt_tokens_details and num_cached_tokens: + if self.enable_prompt_tokens_details and num_cached_tokens is not None: final_usage_info.prompt_tokens_details = PromptTokenUsageInfo( cached_tokens=num_cached_tokens ) @@ -583,7 +583,7 @@ class OpenAIServingCompletion(OpenAIServing): if ( self.enable_prompt_tokens_details and last_final_res - and last_final_res.num_cached_tokens + and last_final_res.num_cached_tokens is not None ): usage.prompt_tokens_details = PromptTokenUsageInfo( cached_tokens=last_final_res.num_cached_tokens diff --git a/vllm/entrypoints/serve/disagg/serving.py b/vllm/entrypoints/serve/disagg/serving.py index 72aeb843773..0bb29c68d01 100644 --- a/vllm/entrypoints/serve/disagg/serving.py +++ b/vllm/entrypoints/serve/disagg/serving.py @@ -307,7 +307,10 @@ class ServingTokens(OpenAIServing): completion_tokens=num_generated_tokens, total_tokens=num_prompt_tokens + num_generated_tokens, ) - if self.enable_prompt_tokens_details and final_res.num_cached_tokens: + if ( + self.enable_prompt_tokens_details + and final_res.num_cached_tokens is not None + ): # This info is not available at the /coordinator level usage.prompt_tokens_details = PromptTokenUsageInfo( cached_tokens=final_res.num_cached_tokens @@ -424,7 +427,7 @@ class ServingTokens(OpenAIServing): total_tokens=num_prompt_tokens + total_completion_tokens, ) - if self.enable_prompt_tokens_details and num_cached_tokens: + if self.enable_prompt_tokens_details and num_cached_tokens is not None: final_usage_info.prompt_tokens_details = PromptTokenUsageInfo( cached_tokens=num_cached_tokens ) From e0b9fb12902b0bed54d2f1b866a7ae00b30aa814 Mon Sep 17 00:00:00 2001 From: Ekagra Ranjan <3116519+ekagra-ranjan@users.noreply.github.com> Date: Fri, 12 Jun 2026 00:05:11 -0400 Subject: [PATCH 23/52] [ASR] Optimize CPU preproc to get 2.5x RTFx via multi-threading (#44612) Signed-off-by: Ekagra Ranjan <3116519+ekagra-ranjan@users.noreply.github.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- vllm/entrypoints/serve/utils/server_utils.py | 7 ++ .../speech_to_text/base/serving.py | 95 ++++++++++++------- vllm/envs.py | 11 +++ vllm/utils/async_utils.py | 26 +++++ 4 files changed, 105 insertions(+), 34 deletions(-) diff --git a/vllm/entrypoints/serve/utils/server_utils.py b/vllm/entrypoints/serve/utils/server_utils.py index 3b6dfde447e..d24d492b61e 100644 --- a/vllm/entrypoints/serve/utils/server_utils.py +++ b/vllm/entrypoints/serve/utils/server_utils.py @@ -474,6 +474,13 @@ async def lifespan(app: FastAPI): finally: if task is not None: task.cancel() + for attr_name in ( + "openai_serving_transcription", + "openai_serving_translation", + ): + serving = getattr(app.state, attr_name, None) + if serving is not None and hasattr(serving, "shutdown"): + serving.shutdown() finally: # Ensure app state including engine ref is gc'd del app.state diff --git a/vllm/entrypoints/speech_to_text/base/serving.py b/vllm/entrypoints/speech_to_text/base/serving.py index 1c6a0d77fe2..9c0ecac41c1 100644 --- a/vllm/entrypoints/speech_to_text/base/serving.py +++ b/vllm/entrypoints/speech_to_text/base/serving.py @@ -6,6 +6,7 @@ import math import time import zlib from collections.abc import AsyncGenerator, Callable, Set +from concurrent.futures import ThreadPoolExecutor from functools import cached_property from typing import Final, Literal, TypeAlias, TypeVar, cast @@ -37,7 +38,7 @@ from vllm.renderers.inputs import DictPrompt, EncoderDecoderDictPrompt from vllm.renderers.inputs.preprocess import parse_enc_dec_prompt, parse_model_prompt from vllm.sampling_params import BeamSearchParams, SamplingParams from vllm.tokenizers import get_tokenizer -from vllm.utils.async_utils import merge_async_iterators +from vllm.utils.async_utils import make_async_with_semaphore, merge_async_iterators from ..transcription.protocol import ( TranscriptionResponse, @@ -63,6 +64,7 @@ T = TypeVar("T", bound=SpeechToTextResponse) V = TypeVar("V", bound=SpeechToTextResponseVerbose) S = TypeVar("S", bound=SpeechToTextSegment) + ResponseType: TypeAlias = ( TranscriptionResponse | TranslationResponse @@ -131,6 +133,19 @@ class OpenAISpeechToText(OpenAIServing): self.default_sampling_params, ) + # setup preprocess resources + # we keep separate thread pool for frontend preprocessing instead + # of reusing the one from Renderer which showed lower throughput + # https://github.com/vllm-project/vllm/pull/44612#issuecomment-4662757781 + num_audio_preprocess_workers = envs.VLLM_MAX_AUDIO_PREPROCESS_WORKERS + self._preprocess_executor = ThreadPoolExecutor( + max_workers=num_audio_preprocess_workers, + thread_name_prefix="stt-preprocess", + ) + self._decode_and_chunk_speech_async = make_async_with_semaphore( + self._decode_and_chunk_speech, executor=self._preprocess_executor + ) + @cached_property def model_cls(self) -> type[SupportsTranscription]: from vllm.model_executor.model_loader import get_model_cls @@ -138,6 +153,49 @@ class OpenAISpeechToText(OpenAIServing): model_cls = get_model_cls(self.model_config) return cast(type[SupportsTranscription], model_cls) + def shutdown(self) -> None: + self._preprocess_executor.shutdown(wait=False) + + def _decode_and_chunk_speech( + self, + audio_data: bytes, + ) -> tuple[list[np.ndarray], float]: + # Decode audio bytes. For container formats (MP4, M4A, WebM) that + # soundfile cannot detect from a BytesIO stream, _load_audio_bytes + # transparently falls back to ffmpeg via an in-memory fd. + # NOTE resample to model SR here for efficiency. This is also a + # pre-requisite for chunking, as it assumes Whisper SR. + try: + with io.BytesIO(audio_data) as buf: + y, sr = load_audio( + buf, + sr=self.asr_config.sample_rate, + max_duration_s=self.max_audio_decode_duration_s, + ) + except Exception as exc: + raise ValueError("Invalid or unsupported audio file.") from exc + + duration = get_audio_duration(y=y, sr=sr) + do_split_audio = self.asr_config.allow_audio_chunking and ( + self.asr_config.max_audio_clip_s is not None + and duration > self.asr_config.max_audio_clip_s + ) + + if not do_split_audio: + chunks = [y] + else: + assert self.asr_config.max_audio_clip_s is not None + assert self.asr_config.min_energy_split_window_size is not None + chunks = split_audio( + audio_data=y, + sample_rate=int(sr), + max_clip_duration_s=self.asr_config.max_audio_clip_s, + overlap_duration_s=self.asr_config.overlap_chunk_second, + min_energy_window_size=self.asr_config.min_energy_split_window_size, + ) + + return chunks, duration + async def _detect_language( self, audio_chunk: np.ndarray, @@ -210,39 +268,8 @@ class OpenAISpeechToText(OpenAIServing): value=len(audio_data) / 1024**2, ) - # Decode audio bytes. For container formats (MP4, M4A, WebM) that - # soundfile cannot detect from a BytesIO stream, _load_audio_bytes - # transparently falls back to ffmpeg via an in-memory fd. - # NOTE resample to model SR here for efficiency. This is also a - # pre-requisite for chunking, as it assumes Whisper SR. - try: - with io.BytesIO(audio_data) as buf: - y, sr = load_audio( - buf, - sr=self.asr_config.sample_rate, - max_duration_s=self.max_audio_decode_duration_s, - ) - except Exception as exc: - raise ValueError("Invalid or unsupported audio file.") from exc - - duration = get_audio_duration(y=y, sr=sr) - do_split_audio = self.asr_config.allow_audio_chunking and ( - self.asr_config.max_audio_clip_s is not None - and duration > self.asr_config.max_audio_clip_s - ) - - if not do_split_audio: - chunks = [y] - else: - assert self.asr_config.max_audio_clip_s is not None - assert self.asr_config.min_energy_split_window_size is not None - chunks = split_audio( - audio_data=y, - sample_rate=int(sr), - max_clip_duration_s=self.asr_config.max_audio_clip_s, - overlap_duration_s=self.asr_config.overlap_chunk_second, - min_energy_window_size=self.asr_config.min_energy_split_window_size, - ) + # Run cpu intensive preprocess step in a separate thread pool executor. + chunks, duration = await self._decode_and_chunk_speech_async(audio_data) if request.language is None and getattr( self.model_cls, "supports_explicit_language_detection", False diff --git a/vllm/envs.py b/vllm/envs.py index d0133638f16..479aab2323c 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -78,6 +78,7 @@ if TYPE_CHECKING: VLLM_MEDIA_LOADING_THREAD_COUNT: int = 8 VLLM_MAX_AUDIO_CLIP_FILESIZE_MB: int = 25 VLLM_MAX_AUDIO_DECODE_DURATION_S: int = 600 + VLLM_MAX_AUDIO_PREPROCESS_WORKERS: int = max(1, min(os.cpu_count() or 1, 2)) VLLM_VIDEO_LOADER_BACKEND: str = "opencv" VLLM_MEDIA_CONNECTOR: str = "http" VLLM_MM_HASHER_ALGORITHM: str = "blake3" @@ -928,6 +929,15 @@ environment_variables: dict[str, Callable[[], Any]] = { "VLLM_MAX_AUDIO_DECODE_DURATION_S": lambda: int( os.getenv("VLLM_MAX_AUDIO_DECODE_DURATION_S", "600") ), + # Maximum number of worker threads used for STT preprocessing. The default + # intentionally caps at 2 because that performed best in profiling. + # https://github.com/vllm-project/vllm/pull/44612#issuecomment-4662757781 + "VLLM_MAX_AUDIO_PREPROCESS_WORKERS": lambda: int( + os.getenv( + "VLLM_MAX_AUDIO_PREPROCESS_WORKERS", + str(max(1, min(os.cpu_count() or 1, 2))), + ) + ), # Backend for Video IO — selects the frame-sampling algorithm. # - "opencv": uniform sampling. # - "opencv_dynamic": duration-aware dynamic sampling. @@ -1997,6 +2007,7 @@ def compile_factors() -> dict[str, object]: "VLLM_MEDIA_LOADING_THREAD_COUNT", "VLLM_MAX_AUDIO_CLIP_FILESIZE_MB", "VLLM_MAX_AUDIO_DECODE_DURATION_S", + "VLLM_MAX_AUDIO_PREPROCESS_WORKERS", "VLLM_VIDEO_LOADER_BACKEND", "VLLM_MEDIA_CONNECTOR", "VLLM_OBJECT_STORAGE_SHM_BUFFER_NAME", diff --git a/vllm/utils/async_utils.py b/vllm/utils/async_utils.py index 725868c39a3..9f368be7b2d 100644 --- a/vllm/utils/async_utils.py +++ b/vllm/utils/async_utils.py @@ -248,6 +248,32 @@ def make_async( return _async_wrapper +def make_async_with_semaphore( + func: Callable[P, T], + executor: ThreadPoolExecutor, +) -> Callable[P, Awaitable[T]]: + """ + Take a blocking function, and run it on in an executor thread. + + This function prevents the blocking function from blocking the + asyncio event loop. + The code in this function needs to be thread safe. + + The function is wrapped in a semaphore to limit the number of + concurrent executions making it easier to cancel tasks before they start. + """ + + semaphore = asyncio.Semaphore(executor._max_workers) + + async def _async_wrapper(*args: P.args, **kwargs: P.kwargs) -> T: + loop = asyncio.get_event_loop() + p_func = partial(func, *args, **kwargs) + async with semaphore: + return await loop.run_in_executor(executor, p_func) + + return _async_wrapper + + def run_in_loop(loop: AbstractEventLoop, function: Callable, *args): if in_loop(loop): function(*args) From b927004c44e20c8cb86918d500adb431b1661607 Mon Sep 17 00:00:00 2001 From: Varun Sundar Rabindranath Date: Fri, 12 Jun 2026 00:07:35 -0400 Subject: [PATCH 24/52] [Bugfix] Mamba CPU Offloading (#44599) Signed-off-by: varun sundar rabindranath Co-authored-by: varun sundar rabindranath --- .../unit/test_offloading_connector.py | 88 +++++++++++++++++++ .../kv_connector/v1/offloading/scheduler.py | 27 +++++- 2 files changed, 114 insertions(+), 1 deletion(-) diff --git a/tests/v1/kv_connector/unit/test_offloading_connector.py b/tests/v1/kv_connector/unit/test_offloading_connector.py index c432b1b20ed..34a8ec57281 100644 --- a/tests/v1/kv_connector/unit/test_offloading_connector.py +++ b/tests/v1/kv_connector/unit/test_offloading_connector.py @@ -554,3 +554,91 @@ def test_fs_tiering_offloading(tmp_path) -> None: finally: subscriber.close() del llm + + +@pytest.mark.skipif( + not current_platform.is_cuda(), + reason="HMA mamba-align CPU offload test is CUDA-only", +) +@pytest.mark.parametrize( + "model,block_size,tp_size", + [ + # ("Qwen/Qwen3.6-35B-A3B", 1056, 2), + # ("tiiuae/falcon-mamba-7b", 16, 1), + ("state-spaces/mamba-1.4b-hf", 16, 1) + ], +) +def test_mamba_align_cpu_offload(model: str, block_size: int, tp_size: int): + kv_transfer_config = KVTransferConfig( + kv_connector="OffloadingConnector", + kv_role="kv_both", + kv_connector_extra_config={ + "cpu_bytes_to_use": 4 << 30, + "block_size": block_size, + }, + ) + llm = LLM( + model=model, + max_model_len=block_size * 10, + gpu_memory_utilization=0.85, + tensor_parallel_size=tp_size, + kv_transfer_config=kv_transfer_config, + language_model_only=True, + enable_prefix_caching=True, + mamba_cache_mode="align", + disable_hybrid_kv_cache_manager=False, + ) + + _PROMPT_SIZE: int = block_size * 2 + _PROMPT_TEXT = "Hi. Give me a set of trivia questions and their answers " + + # build prompt ids to match prompt_size + tokenizer = llm.get_tokenizer() + raw_ids: list[int] = tokenizer.encode(_PROMPT_TEXT) + while len(raw_ids) < _PROMPT_SIZE: + raw_ids = tokenizer.encode("....") + raw_ids + initial_ids: list[int] = raw_ids[:_PROMPT_SIZE] + + sampling_params = SamplingParams(max_tokens=128, temperature=0, ignore_eos=True) + + failures: list[str] = [] + + def _get_output_str(outputs): + return outputs[0].outputs[0].text + + def _verify(llm, prompt, label: str): + cold_outputs = llm.generate([prompt], sampling_params, use_tqdm=False) + _wait_for_prefix_cache_reset(llm) + cpu_outputs = llm.generate([prompt], sampling_params, use_tqdm=False) + + cold_text = _get_output_str(cold_outputs) + cpu_text = _get_output_str(cpu_outputs) + print(f"{label} : cold outputs\n{cold_text}") + print(f"{label} : cpu outputs\n{cpu_text}") + + if cold_text != cpu_text: + failures.append( + f"{label}: mismatch\n cold: {cold_text!r}\n cpu: {cpu_text!r}" + ) + + try: + # Mamba has only a single state. The CPU cache stores are triggered + # at offload block boundaries. When the prompt is exactly at the boundary, + # The CPU offload should not load the cached block. + # This is because we'd use that state to recompute the last token. This + # does not work for mamba as there is only one KV value and that is for + # for the token at the boundary. + # This is fine for other attention types as we have all the necessary + # token KV values in the hit blocks. + prompt = TokensPrompt(prompt_token_ids=initial_ids) + _verify(llm, prompt, "block-boundary-prompt") + + # Test for prompt token ids at non-block boundaries. + # Reuse is okay for this case. + prompt = TokensPrompt(prompt_token_ids=[0] + initial_ids) + _verify(llm, prompt, "block-mid-prompt") + + assert not failures, "\n\n".join(failures) + + finally: + del llm diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py index 94d68972822..1d3d83709be 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py @@ -19,7 +19,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import ( _TransferMetricName, ) from vllm.logger import init_logger -from vllm.utils.math_utils import cdiv +from vllm.utils.math_utils import cdiv, round_down from vllm.v1.core.kv_cache_manager import KVCacheBlocks from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.kv_cache_interface import ( @@ -94,6 +94,24 @@ def get_sliding_window_size_in_blocks( return None +def resolve_mamba_align_size(spec: "OffloadingSpec") -> int | None: + """Scan all KV cache groups in *spec* and return the single mamba alignment + size, or None if no group requires mamba alignment. + + For MambaSpec groups in "align" cache mode the hit window must be rounded + down to a multiple of the offloaded block size. Asserts that all such + groups agree on the same value. + """ + mamba_align_size: int | None = None + for idx, gpu_block_size in enumerate(spec.gpu_block_size): + kv_spec = spec.kv_cache_config.kv_cache_groups[idx].kv_cache_spec + if isinstance(kv_spec, MambaSpec) and kv_spec.mamba_cache_mode == "align": + offload_block_size = gpu_block_size * spec.block_size_factor + assert mamba_align_size is None or mamba_align_size == offload_block_size + mamba_align_size = offload_block_size + return mamba_align_size + + class SchedulerOffloadConfig(NamedTuple): kv_group_configs: tuple[GroupOffloadConfig, ...] block_size_factor: int @@ -290,6 +308,7 @@ class OffloadingConnectorScheduler: # used by _lookup self._sliding_window_groups: tuple[int, ...] = tuple(sliding_window_groups) self._lookup_groups = tuple(full_attention_groups) + self._sliding_window_groups + self._mamba_align_size: int | None = resolve_mamba_align_size(spec) self._req_status: dict[ReqId, RequestOffloadState] = {} self._current_batch_load_jobs: dict[int, TransferJob] = {} @@ -408,6 +427,12 @@ class OffloadingConnectorScheduler: # for sliding window attention, we must reduce by 1 to make sure # we still have a hit after reduction max_hit_size_tokens -= 1 + if self._mamba_align_size is not None: + # Constrain hit-window to the mamba block size. + max_hit_size_tokens = round_down( + max_hit_size_tokens, self._mamba_align_size + ) + num_hit_tokens: int = 0 defer_lookup = False lookup_groups = self._lookup_groups From 226ba9fc9e285556e7269e4efe102530a4d9fedb Mon Sep 17 00:00:00 2001 From: Ekagra Ranjan <3116519+ekagra-ranjan@users.noreply.github.com> Date: Fri, 12 Jun 2026 00:11:16 -0400 Subject: [PATCH 25/52] [ASR] Add Long Audio benchmark and correctness test (#44587) Signed-off-by: Ekagra Ranjan <3116519+ekagra-ranjan@users.noreply.github.com> --- docs/benchmarking/cli.md | 4 +- tests/benchmarks/test_audio_dataset.py | 200 ++++++++++++++++ .../test_transcription_api_correctness.py | 220 ++++++++++++++++-- vllm/benchmarks/datasets/datasets.py | 97 ++++++-- vllm/benchmarks/lib/endpoint_request_func.py | 55 ++++- 5 files changed, 530 insertions(+), 46 deletions(-) create mode 100644 tests/benchmarks/test_audio_dataset.py diff --git a/docs/benchmarking/cli.md b/docs/benchmarking/cli.md index 3d8fda95a34..22406f2eaa2 100644 --- a/docs/benchmarking/cli.md +++ b/docs/benchmarking/cli.md @@ -37,7 +37,7 @@ th { | HuggingFace-HumanEval | ✅ | ✅ | `openai/openai_humaneval` | | HuggingFace-GSM8K | ✅ | ✅ | `openai/gsm8k` | | HuggingFace-Blazedit | ✅ | ✅ | `vdaita/edit_5k_char`, `vdaita/edit_10k_char` | -| HuggingFace-ASR | ✅ | ✅ | `openslr/librispeech_asr`, `facebook/voxpopuli`, `LIUM/tedlium`, `edinburghcstr/ami`, `speechcolab/gigaspeech`, `kensho/spgispeech` | +| HuggingFace-ASR | ✅ | ✅ | `openslr/librispeech_asr`, `facebook/voxpopuli`, `LIUM/tedlium`, `edinburghcstr/ami`, `speechcolab/gigaspeech`, `kensho/spgispeech`, `ArtificialAnalysis/Earnings22-Cleaned-AA`, `D4nt3/esb-datasets-earnings22-validation-tiny-filtered` | | Spec Bench | ✅ | ✅ | `wget https://raw.githubusercontent.com/hemingkx/Spec-Bench/refs/heads/main/data/spec_bench/question.jsonl` | | SPEED-Bench | ✅ | ✅ | `curl -LsSf https://raw.githubusercontent.com/NVIDIA-NeMo/Skills/refs/heads/main/nemo_skills/dataset/speed-bench/prepare.py \| python3 -` | | Custom | ✅ | ✅ | Local file: `data.jsonl` | @@ -532,7 +532,7 @@ vllm bench serve \ --blazedit-max-distance 0.99 ``` -`openslr/librispeech_asr`, `facebook/voxpopuli`, `LIUM/tedlium`, `edinburghcstr/ami`, `speechcolab/gigaspeech`, `kensho/spgispeech` +`openslr/librispeech_asr`, `facebook/voxpopuli`, `LIUM/tedlium`, `edinburghcstr/ami`, `speechcolab/gigaspeech`, `kensho/spgispeech`, `ArtificialAnalysis/Earnings22-Cleaned-AA`, `D4nt3/esb-datasets-earnings22-validation-tiny-filtered` ```bash vllm bench serve \ diff --git a/tests/benchmarks/test_audio_dataset.py b/tests/benchmarks/test_audio_dataset.py new file mode 100644 index 00000000000..5957011c484 --- /dev/null +++ b/tests/benchmarks/test_audio_dataset.py @@ -0,0 +1,200 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import asyncio +from pathlib import Path +from typing import Protocol, cast + +import numpy as np +import pytest +import soundfile as sf + +import vllm.benchmarks.datasets.datasets as datasets_module +import vllm.benchmarks.lib.endpoint_request_func as request_func_module +from vllm.benchmarks.lib.endpoint_request_func import RequestFuncInput + +pytestmark = pytest.mark.skip_global_cleanup + + +class _ReadableBinary(Protocol): + def read(self, size: int = -1) -> bytes: ... + + +class _TokenizedPrompt: + def __init__(self, prompt: str) -> None: + self.input_ids = prompt.split() + + +class _Tokenizer: + def __init__(self, name_or_path: str = "openai/whisper-large-v3") -> None: + self.name_or_path = name_or_path + + def __call__(self, prompt: str) -> _TokenizedPrompt: + return _TokenizedPrompt(prompt) + + +def _write_wav(path: Path, duration_s: float = 0.1, sample_rate: int = 16_000) -> None: + num_samples = int(duration_s * sample_rate) + sf.write(path, np.zeros(num_samples, dtype=np.float32), sample_rate) + + +class _FakeFormData: + def __init__(self) -> None: + self.fields: list[tuple[str, object, dict[str, str]]] = [] + + def add_field(self, name: str, value: object, **kwargs: str) -> None: + self.fields.append((name, value, kwargs)) + + +class _FakeContent: + async def iter_any(self): + yield b'data: {"choices":[{"delta":{"content":"hello"}}]}\n\n' + yield b'data: {"usage":{"completion_tokens":1}}\n\n' + yield b"data: [DONE]\n\n" + + +class _FakeResponse: + def __init__(self) -> None: + self.status = 200 + self.reason = "OK" + self.content = _FakeContent() + + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc, tb): + return False + + +class _FakeSession: + def __init__(self) -> None: + self.uploaded_bytes: bytes | None = None + self.upload_filename: str | None = None + self.fields: list[tuple[str, object, dict[str, str]]] | None = None + + def post(self, *, url: str, data: _FakeFormData, headers: dict[str, str]): + del url, headers + self.fields = list(data.fields) + _, file_obj, file_kwargs = self.fields[0] + file_obj = cast(_ReadableBinary, file_obj) + self.uploaded_bytes = file_obj.read() + self.upload_filename = file_kwargs.get("filename") + return _FakeResponse() + + +def test_asr_dataset_sample_handles_local_audio_paths(tmp_path: Path) -> None: + audio_path = tmp_path / "earnings.wav" + _write_wav(audio_path, duration_s=0.1) + + dataset = object.__new__(datasets_module.ASRDataset) + dataset.data = [ + { + "audio": { + "path": str(audio_path), + "bytes": None, + }, + "text": "quarterly earnings call", + } + ] + + samples = dataset.sample( + tokenizer=_Tokenizer(), + num_requests=1, + output_len=32, + asr_min_audio_len_sec=0.0, + asr_max_audio_len_sec=1.0, + ) + + assert len(samples) == 1 + assert samples[0].multi_modal_data == {"audio_path": str(audio_path)} + assert ( + samples[0].prompt == "<|startoftranscript|><|en|><|transcribe|><|notimestamps|>" + ) + + +def test_asr_dataset_sample_handles_embedded_audio_bytes(tmp_path: Path) -> None: + audio_path = tmp_path / "earnings.wav" + _write_wav(audio_path, duration_s=0.1) + + dataset = object.__new__(datasets_module.ASRDataset) + dataset.data = [ + { + "audio": { + "path": None, + "bytes": audio_path.read_bytes(), + }, + "text": "quarterly earnings call", + } + ] + + samples = dataset.sample( + tokenizer=_Tokenizer(), + num_requests=1, + output_len=32, + asr_min_audio_len_sec=0.0, + asr_max_audio_len_sec=1.0, + ) + + assert len(samples) == 1 + assert isinstance(samples[0].multi_modal_data, dict) + audio, sample_rate = samples[0].multi_modal_data["audio"] + assert sample_rate == 16_000 + assert isinstance(audio, np.ndarray) + assert audio.size > 0 + + +def test_async_request_openai_audio_handles_local_audio_paths( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + audio_path = tmp_path / "earnings.wav" + _write_wav(audio_path, duration_s=0.25) + + monkeypatch.setattr(request_func_module.aiohttp, "FormData", _FakeFormData) + session = _FakeSession() + request_input = RequestFuncInput( + prompt="", + api_url="http://localhost:8000/v1/audio/transcriptions", + prompt_len=1, + output_len=32, + model="openai/whisper-large-v3", + multi_modal_content={"audio_path": str(audio_path)}, + ) + + output = asyncio.run( + request_func_module.async_request_openai_audio(request_input, session) + ) + + assert session.upload_filename == audio_path.name + assert session.uploaded_bytes == audio_path.read_bytes() + assert output.success is True + assert output.generated_text == "hello" + assert output.output_tokens == 1 + assert output.input_audio_duration == pytest.approx(0.25, abs=1e-2) + + +def test_async_request_openai_audio_handles_decoded_audio_arrays( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(request_func_module.aiohttp, "FormData", _FakeFormData) + session = _FakeSession() + request_input = RequestFuncInput( + prompt="", + api_url="http://localhost:8000/v1/audio/transcriptions", + prompt_len=1, + output_len=32, + model="openai/whisper-large-v3", + multi_modal_content={ + "audio": (np.zeros(1_600, dtype=np.float32), 16_000), + }, + ) + + output = asyncio.run( + request_func_module.async_request_openai_audio(request_input, session) + ) + + assert session.upload_filename == "audio.wav" + assert session.uploaded_bytes is not None + assert output.success is True + assert output.generated_text == "hello" diff --git a/tests/entrypoints/speech_to_text/correctness/test_transcription_api_correctness.py b/tests/entrypoints/speech_to_text/correctness/test_transcription_api_correctness.py index fedbd74795b..af61ebc5264 100644 --- a/tests/entrypoints/speech_to_text/correctness/test_transcription_api_correctness.py +++ b/tests/entrypoints/speech_to_text/correctness/test_transcription_api_correctness.py @@ -16,10 +16,11 @@ from statistics import mean, median import pytest import soundfile import torch -from datasets import load_dataset +from datasets import Audio, load_dataset from evaluate import load from transformers.models.whisper.english_normalizer import EnglishTextNormalizer +from vllm.benchmarks.datasets.datasets import ASRDataset from vllm.multimodal.audio import get_audio_duration from vllm.tokenizers import get_tokenizer @@ -38,6 +39,20 @@ def to_bytes(y, sr): return buffer +def load_audio_sample(audio): + # Avoid torchcodec in CI by decoding dataset audio with soundfile. + if "array" in audio and "sampling_rate" in audio: + return audio["array"], audio["sampling_rate"] + + if audio.get("path"): + return soundfile.read(audio["path"], dtype="float32") + + if audio.get("bytes") is not None: + return soundfile.read(io.BytesIO(audio["bytes"]), dtype="float32") + + raise ValueError("Audio sample did not contain array, path, or bytes data") + + # not all models have a normalizer so use the one from whisper as a standard option normalizer_model_info = HF_EXAMPLE_MODELS.find_hf_info("openai/whisper-large-v3") normalizer_tokenizer = get_tokenizer( @@ -48,7 +63,7 @@ normalizer_tokenizer = get_tokenizer( normalizer = EnglishTextNormalizer(normalizer_tokenizer.english_spelling_normalizer) -async def transcribe_audio(client, tokenizer, y, sr): +async def transcribe_audio(client, tokenizer, y, sr, extra_body=None): # Send loaded audio directly instead of loading from disk, # don't account for that time though with to_bytes(y, sr) as f: @@ -58,6 +73,7 @@ async def transcribe_audio(client, tokenizer, y, sr): model=tokenizer.name_or_path, language="en", temperature=0.0, + extra_body=extra_body, ) end_time = time.perf_counter() # NOTE there's no streaming in transcriptions, can't measure ttft @@ -68,17 +84,21 @@ async def transcribe_audio(client, tokenizer, y, sr): return latency, num_output_tokens, transcription.text -async def bound_transcribe(sem, client, tokenizer, audio, reference): +async def bound_transcribe( + sem, client, tokenizer, audio, sr, reference, extra_body=None +): # Use semaphore to limit concurrent requests. async with sem: - result = await transcribe_audio(client, tokenizer, *audio) + result = await transcribe_audio( + client, tokenizer, audio, sr, extra_body=extra_body + ) # Normalize *english* output/reference for evaluation. out = normalizer(result[2]) ref = normalizer(reference) return result[:2] + (out, ref) -async def process_dataset(model, client, data, concurrent_request): +async def process_dataset(model, client, data, concurrent_request, extra_body=None): sem = asyncio.Semaphore(concurrent_request) model_info = HF_EXAMPLE_MODELS.find_hf_info(model) @@ -89,14 +109,16 @@ async def process_dataset(model, client, data, concurrent_request): ) # Warmup call as the first `load_audio` server-side is quite slow. - audio, sr = data[0]["audio"]["array"], data[0]["audio"]["sampling_rate"] - _ = await bound_transcribe(sem, client, tokenizer, (audio, sr), "") + audio, sr = load_audio_sample(data[0]["audio"]) + _ = await bound_transcribe(sem, client, tokenizer, audio, sr, "", extra_body) tasks: list[asyncio.Task] = [] for sample in data: - audio, sr = sample["audio"]["array"], sample["audio"]["sampling_rate"] + audio, sr = load_audio_sample(sample["audio"]) task = asyncio.create_task( - bound_transcribe(sem, client, tokenizer, (audio, sr), sample["text"]) + bound_transcribe( + sem, client, tokenizer, audio, sr, sample["text"], extra_body + ) ) tasks.append(task) return await asyncio.gather(*tasks) @@ -121,19 +143,36 @@ def print_performance_metrics(results, total_time): def add_duration(sample): - y, sr = sample["audio"]["array"], sample["audio"]["sampling_rate"] + y, sr = load_audio_sample(sample["audio"]) sample["duration_ms"] = get_audio_duration(y=y, sr=sr) * 1000 return sample -def load_hf_dataset(dataset_repo: str, split="validation", **hf_kwargs): - ## Load and filter the dataset - dataset = load_dataset(dataset_repo, split=split, **hf_kwargs) - if "duration_ms" not in dataset[0]: - # compute duration to filter +def load_asr_dataset_rows(dataset_repo: str, split="validation", **hf_kwargs): + if dataset_repo in ASRDataset.SUPPORTED_DATASET_PATHS: + asr_dataset_kwargs = { + "dataset_path": dataset_repo, + "dataset_split": split, + "disable_shuffle": True, + "no_stream": True, + } + for key in ("dataset_subset", "hf_name", "trust_remote_code"): + if key in hf_kwargs: + asr_dataset_kwargs[key] = hf_kwargs[key] + return ASRDataset(**asr_dataset_kwargs).data + + return load_dataset(dataset_repo, split=split, **hf_kwargs) + + +def load_shortform_eval_dataset(dataset_repo: str, split="validation", **hf_kwargs): + ## Load and filter the dataset. + dataset = load_asr_dataset_rows(dataset_repo, split=split, **hf_kwargs) + dataset = dataset.cast_column("audio", Audio(decode=False)) + if "duration_ms" not in dataset.column_names: + # Compute duration to filter. dataset = dataset.map(add_duration) - # Whisper max supported duration + # Whisper max supported duration. dataset = dataset.filter(lambda example: example["duration_ms"] < 30000) return dataset @@ -145,11 +184,16 @@ def run_evaluation( max_concurrent_reqs: int, n_examples: int = -1, print_metrics: bool = True, + extra_body=None, ): if n_examples > 0: dataset = dataset.select(range(n_examples)) start = time.perf_counter() - results = asyncio.run(process_dataset(model, client, dataset, max_concurrent_reqs)) + results = asyncio.run( + process_dataset( + model, client, dataset, max_concurrent_reqs, extra_body=extra_body + ) + ) end = time.perf_counter() total_time = end - start print(f"Total Test Time: {total_time:.4f} seconds") @@ -164,6 +208,106 @@ def run_evaluation( return wer_score +LONGFORM_DATASET_REPO = ASRDataset.EARNINGS22_CLEANED_DATASET +LONGFORM_DATASET_SPLIT = "test" +LONGFORM_NUM_SAMPLES = 6 + + +def load_longform_dataset(): + dataset = load_asr_dataset_rows( + LONGFORM_DATASET_REPO, + split=LONGFORM_DATASET_SPLIT, + ) + assert len(dataset) >= LONGFORM_NUM_SAMPLES + return dataset.select(range(LONGFORM_NUM_SAMPLES)) + + +async def transcribe_audio_path(client, tokenizer, audio_path: str, extra_body=None): + with open(audio_path, "rb") as f: + start_time = time.perf_counter() + transcription = await client.audio.transcriptions.create( + file=f, + model=tokenizer.name_or_path, + language="en", + temperature=0.0, + extra_body=extra_body, + ) + end_time = time.perf_counter() + + latency = end_time - start_time + num_output_tokens = len( + tokenizer(transcription.text, add_special_tokens=False).input_ids + ) + return latency, num_output_tokens, transcription.text + + +async def bound_transcribe_path( + sem, client, tokenizer, audio_path, reference, extra_body=None +): + async with sem: + result = await transcribe_audio_path( + client, tokenizer, audio_path, extra_body=extra_body + ) + out = normalizer(result[2]) + ref = normalizer(reference) + return result[:2] + (out, ref) + + +async def process_longform_dataset( + model, client, data, concurrent_request, extra_body=None +): + sem = asyncio.Semaphore(concurrent_request) + + model_info = HF_EXAMPLE_MODELS.find_hf_info(model) + tokenizer = get_tokenizer( + model, + tokenizer_mode=model_info.tokenizer_mode, + trust_remote_code=model_info.trust_remote_code, + ) + + warmup_path = data[0]["audio"]["path"] + _ = await bound_transcribe_path(sem, client, tokenizer, warmup_path, "", extra_body) + + tasks: list[asyncio.Task] = [] + for sample in data: + audio_path = sample["audio"]["path"] + task = asyncio.create_task( + bound_transcribe_path( + sem, client, tokenizer, audio_path, sample["text"], extra_body + ) + ) + tasks.append(task) + return await asyncio.gather(*tasks) + + +def run_longform_evaluation( + model: str, + client, + dataset, + max_concurrent_reqs: int, + print_metrics: bool = True, + extra_body=None, +): + start = time.perf_counter() + results = asyncio.run( + process_longform_dataset( + model, client, dataset, max_concurrent_reqs, extra_body=extra_body + ) + ) + end = time.perf_counter() + total_time = end - start + print(f"Total Test Time: {total_time:.4f} seconds") + if print_metrics: + print_performance_metrics(results, total_time) + + predictions = [res[2] for res in results] + references = [res[3] for res in results] + wer = load("wer") + wer_score = 100 * wer.compute(references=references, predictions=predictions) + print("WER:", wer_score) + return wer_score + + # alternatives "openai/whisper-large-v2", "openai/whisper-large-v3-turbo".. # NOTE: Expected WER measured with equivalent hf.transformers args: # whisper-large-v3 + esb-datasets-earnings22-validation-tiny-filtered. @@ -184,7 +328,6 @@ def test_wer_correctness( ): model_name, expected_wer = model_config model_info = HF_EXAMPLE_MODELS.find_hf_info(model_name) - # TODO refactor to use `ASRDataset` server_args = [ "--enforce-eager", f"--tokenizer_mode={model_info.tokenizer_mode}", @@ -197,7 +340,7 @@ def test_wer_correctness( model_name, server_args, ) as remote_server: - dataset = load_hf_dataset(dataset_repo) + dataset = load_shortform_eval_dataset(dataset_repo) if not max_concurrent_request: # No max concurrency @@ -216,3 +359,42 @@ def test_wer_correctness( if expected_wer: torch.testing.assert_close(wer, expected_wer, atol=1e-1, rtol=1e-2) + + +# 14-22mins of 6 audio samples of total ~115 mins and just 37MB. +# checks for long audio transcription correctness and RMS split. +@pytest.mark.parametrize( + "model_config", + [("openai/whisper-large-v3", 9.5)], +) +def test_long_audio_wer_correctness(model_config): + model_name, expected_wer = model_config + model_info = HF_EXAMPLE_MODELS.find_hf_info(model_name) + server_args = [ + f"--tokenizer_mode={model_info.tokenizer_mode}", + ] + + if model_info.trust_remote_code: + server_args.append("--trust-remote-code") + + # 1800 seconds is 30 minutes + env_dict = { + "VLLM_MAX_AUDIO_DECODE_DURATION_S": "1800", + } + + with RemoteOpenAIServer( + model_name, + server_args, + env_dict=env_dict, + ) as remote_server: + dataset = load_longform_dataset() + client = remote_server.get_async_client() + wer = run_longform_evaluation( + model=model_name, + client=client, + dataset=dataset, + max_concurrent_reqs=LONGFORM_NUM_SAMPLES, + ) + + print(f"Expected WER: {expected_wer}, Actual WER: {wer}") + torch.testing.assert_close(wer, expected_wer, atol=1e-1, rtol=1e-2) diff --git a/vllm/benchmarks/datasets/datasets.py b/vllm/benchmarks/datasets/datasets.py index abdcedd12be..25ceadc41a1 100644 --- a/vllm/benchmarks/datasets/datasets.py +++ b/vllm/benchmarks/datasets/datasets.py @@ -4001,20 +4001,27 @@ class ASRDataset(HuggingFaceDataset): Dataset class for processing a ASR dataset for transcription. Tested on the following set: - +----------------+----------------------------------------+--------------------------+-----------------------------+ - | Dataset | Domain | Speaking Style | hf-subset | - +----------------+----------------------------------------+--------------------------+-----------------------------+ - | TED-LIUM | TED talks | Oratory | release1, release2, release3| - | | | | release3-speaker-adaptation | - | VoxPopuli | European Parliament | Oratory | en, de, it, fr, ... | - | LibriSpeech | Audiobook | Narrated | "LIUM/tedlium" | - | GigaSpeech | Audiobook, podcast, YouTube | Narrated, spontaneous | xs, s, m, l, xl, dev, test | - | SPGISpeech | Financial meetings | Oratory, spontaneous | S, M, L, dev, test | - | AMI | Meetings | Spontaneous | ihm, sdm | - +----------------+----------------------------------------+--------------------------+-----------------------------+ + +---------------------------+----------------------------------------+--------------------------+-----------------------------+ + | Dataset | Domain | Speaking Style | hf-subset | + +---------------------------+----------------------------------------+--------------------------+-----------------------------+ + | TED-LIUM | TED talks | Oratory | release1, release2, release3| + | | | | release3-speaker-adaptation | + | VoxPopuli | European Parliament | Oratory | en, de, it, fr, ... | + | LibriSpeech | Audiobook | Narrated | "LIUM/tedlium" | + | GigaSpeech | Audiobook, podcast, YouTube | Narrated, spontaneous | xs, s, m, l, xl, dev, test | + | SPGISpeech | Financial meetings | Oratory, spontaneous | S, M, L, dev, test | + | Earnings22-Cleaned-AA | Long form earnings calls | Prepared remarks, Q&A | test | + | Earnings22-Tiny-Filtered | Earnings calls | Prepared remarks, Q&A | validation | + | AMI | Meetings | Spontaneous | ihm, sdm | + +---------------------------+----------------------------------------+--------------------------+-----------------------------+ """ # noqa: E501 + EARNINGS22_CLEANED_DATASET = "ArtificialAnalysis/Earnings22-Cleaned-AA" + EARNINGS22_TINY_FILTERED_DATASET = ( + "D4nt3/esb-datasets-earnings22-validation-tiny-filtered" + ) + SUPPORTED_DATASET_PATHS = { "openslr/librispeech_asr", "facebook/voxpopuli", @@ -4022,11 +4029,52 @@ class ASRDataset(HuggingFaceDataset): "edinburghcstr/ami", "speechcolab/gigaspeech", "kensho/spgispeech", + EARNINGS22_CLEANED_DATASET, + EARNINGS22_TINY_FILTERED_DATASET, } DEFAULT_OUTPUT_LEN = 1024 IS_MULTIMODAL = True + def load_data(self) -> None: + if self.hf_name == self.EARNINGS22_CLEANED_DATASET: + # This subset stores repo-local MP3 paths instead of a HF `Audio` + # column, so eagerly materialize it back into the common schema. + self.data = load_dataset( + self.dataset_path, + name=self.dataset_subset, + split=self.dataset_split, + streaming=False, + trust_remote_code=self.trust_remote_code, + ) + if not getattr(self, "disable_shuffle", False): + self.data = self.data.shuffle(seed=self.random_seed) + self._materialize_local_audio_column() + return + if self.hf_name == self.EARNINGS22_TINY_FILTERED_DATASET: + super().load_data() + self._disable_audio_decode() + return + + super().load_data() + + def _disable_audio_decode(self) -> None: + from datasets import Audio + + self.data = self.data.cast_column("audio", Audio(decode=False)) + + def _materialize_local_audio_column(self) -> None: + local_path_root = Path( + hf_api().snapshot_download(self.hf_name, repo_type="dataset") + ) + self.data = self.data.map( + lambda item: { + "audio": str(local_path_root / item["url"]), + "text": item["transcript"], + } + ) + self._disable_audio_decode() + def sample( self, tokenizer: TokenizerLike, @@ -4052,14 +4100,35 @@ class ASRDataset(HuggingFaceDataset): if len(sampled_requests) >= num_requests: break audio = item["audio"] - y, sr = audio["array"], audio["sampling_rate"] - duration_s = get_audio_duration(y=y, sr=sr) + if ( + isinstance(audio, dict) + and "array" in audio + and "sampling_rate" in audio + ): + y, sr = audio["array"], audio["sampling_rate"] + duration_s = get_audio_duration(y=y, sr=sr) + mm_content = {"audio": (y, sr)} + elif isinstance(audio, str): + duration_s = sf.info(audio).duration + mm_content = {"audio_path": audio} + elif isinstance(audio, dict) and audio.get("path"): + duration_s = sf.info(audio["path"]).duration + mm_content = {"audio_path": audio["path"]} + elif isinstance(audio, dict) and audio.get("bytes") is not None: + with BytesIO(audio["bytes"]) as audio_buffer: + y, sr = sf.read(audio_buffer, dtype="float32") + duration_s = get_audio_duration(y=y, sr=sr) + mm_content = {"audio": (y, sr)} + else: + raise ValueError( + "ASR samples must provide decoded audio arrays, " + "embedded audio bytes, or a local audio path." + ) if duration_s < asr_min_audio_len_sec or duration_s > asr_max_audio_len_sec: skipped += 1 continue durations.append(duration_s) - mm_content = {"audio": (y, sr)} sampled_requests.append( SampleRequest( prompt=prompt, diff --git a/vllm/benchmarks/lib/endpoint_request_func.py b/vllm/benchmarks/lib/endpoint_request_func.py index d282033ba1f..db58f422b80 100644 --- a/vllm/benchmarks/lib/endpoint_request_func.py +++ b/vllm/benchmarks/lib/endpoint_request_func.py @@ -445,7 +445,6 @@ async def async_request_openai_audio( api_url = request_func_input.api_url _validate_api_url(api_url, "OpenAI Audio API", {"transcriptions", "translations"}) - content = [{"type": "text", "text": request_func_input.prompt}] payload = { "model": request_func_input.model_name if request_func_input.model_name @@ -469,19 +468,26 @@ async def async_request_openai_audio( buffer.seek(0) return buffer - mm_audio = request_func_input.multi_modal_content - if not isinstance(mm_audio, dict) or "audio" not in mm_audio: - raise TypeError("multi_modal_content must be a dict containing 'audio'") - with to_bytes(*mm_audio["audio"]) as f: + async def send_audio_file( + audio_file: io.BytesIO | Any, + *, + input_audio_duration: float, + filename: str | None = None, + content_type: str | None = None, + ) -> RequestFuncOutput: form = aiohttp.FormData() - form.add_field("file", f, content_type="audio/wav") + add_field_kwargs: dict[str, str] = {} + if filename is not None: + add_field_kwargs["filename"] = filename + if content_type is not None: + add_field_kwargs["content_type"] = content_type + form.add_field("file", audio_file, **add_field_kwargs) for key, value in payload.items(): form.add_field(key, str(value)) output = RequestFuncOutput() output.prompt_len = request_func_input.prompt_len - output.input_audio_duration = soundfile.info(f).duration - f.seek(0) + output.input_audio_duration = input_audio_duration generated_text = "" ttft = 0.0 @@ -541,9 +547,36 @@ async def async_request_openai_audio( exc_info = sys.exc_info() output.error = "".join(traceback.format_exception(*exc_info)) - if pbar: - pbar.update(1) - return output + if pbar: + pbar.update(1) + return output + + mm_audio = request_func_input.multi_modal_content + if not isinstance(mm_audio, dict): + raise TypeError( + "multi_modal_content must be a dict containing 'audio' or 'audio_path'" + ) + if "audio" in mm_audio: + with to_bytes(*mm_audio["audio"]) as f: + input_audio_duration = soundfile.info(f).duration + f.seek(0) + return await send_audio_file( + f, + input_audio_duration=input_audio_duration, + filename="audio.wav", + content_type="audio/wav", + ) + if "audio_path" in mm_audio: + audio_path = mm_audio["audio_path"] + with open(audio_path, "rb") as f: + return await send_audio_file( + f, + input_audio_duration=soundfile.info(audio_path).duration, + filename=os.path.basename(audio_path), + ) + raise TypeError( + "multi_modal_content must be a dict containing 'audio' or 'audio_path'" + ) async def _run_pooling_request( From 7021be66e8c351fa819fd07b4053e946f11c1147 Mon Sep 17 00:00:00 2001 From: Chris Leonard Date: Fri, 12 Jun 2026 00:22:37 -0400 Subject: [PATCH 26/52] [11a/n] Migrate Marlin kernels to torch stable ABI (#45176) Signed-off-by: Chris Leonard --- CMakeLists.txt | 278 +++++------ .../moe/marlin_moe_wna16/kernel.h | 4 +- .../moe/marlin_moe_wna16/marlin_template.h | 8 +- .../gptq_allspark/allspark_utils.cuh | 2 +- .../quantization/marlin/.gitignore | 0 .../quantization/marlin/awq_marlin_repack.cu | 80 ++-- .../quantization/marlin/dequant.h | 0 .../quantization/marlin/generate_kernels.py | 2 +- .../quantization/marlin/gptq_marlin_repack.cu | 91 ++-- .../quantization/marlin/kernel.h | 0 .../quantization/marlin/marlin.cu | 443 ++++++++++-------- .../quantization/marlin/marlin.cuh | 8 - .../quantization/marlin/marlin_dtypes.cuh | 0 .../marlin/marlin_int4_fp8_preprocess.cu | 118 +++++ .../quantization/marlin/marlin_mma.h | 0 .../quantization/marlin/marlin_template.h | 0 csrc/libtorch_stable/torch_bindings.cpp | 29 ++ .../marlin/marlin_int4_fp8_preprocess.cu | 106 ----- csrc/torch_bindings.cpp | 29 -- 19 files changed, 626 insertions(+), 572 deletions(-) rename csrc/{ => libtorch_stable}/quantization/marlin/.gitignore (100%) rename csrc/{ => libtorch_stable}/quantization/marlin/awq_marlin_repack.cu (77%) rename csrc/{ => libtorch_stable}/quantization/marlin/dequant.h (100%) rename csrc/{ => libtorch_stable}/quantization/marlin/generate_kernels.py (99%) rename csrc/{ => libtorch_stable}/quantization/marlin/gptq_marlin_repack.cu (77%) rename csrc/{ => libtorch_stable}/quantization/marlin/kernel.h (100%) rename csrc/{ => libtorch_stable}/quantization/marlin/marlin.cu (61%) rename csrc/{ => libtorch_stable}/quantization/marlin/marlin.cuh (93%) rename csrc/{ => libtorch_stable}/quantization/marlin/marlin_dtypes.cuh (100%) create mode 100644 csrc/libtorch_stable/quantization/marlin/marlin_int4_fp8_preprocess.cu rename csrc/{ => libtorch_stable}/quantization/marlin/marlin_mma.h (100%) rename csrc/{ => libtorch_stable}/quantization/marlin/marlin_template.h (100%) delete mode 100644 csrc/quantization/marlin/marlin_int4_fp8_preprocess.cu diff --git a/CMakeLists.txt b/CMakeLists.txt index 0a48ddca68a..c03360a5d4e 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -358,145 +358,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA") SRCS "${VLLM_EXT_SRC}" CUDA_ARCHS "${CUDA_ARCHS}") - # Only build Marlin kernels if we are building for at least some compatible archs. - # Keep building Marlin for 9.0 as there are some group sizes and shapes that - # are not supported by Machete yet. - - # marlin arches for fp16 output - # Family-conditional 12.0f (one cubin for SM12x family) requires CUDA >= 13.0; - # fall back to architecture-specific 12.0a;12.1a on CUDA < 13.0 (e.g. 12.8). - if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0) - cuda_archs_loose_intersection(MARLIN_ARCHS "8.0+PTX;12.0f" "${CUDA_ARCHS}") - else() - cuda_archs_loose_intersection(MARLIN_ARCHS "8.0+PTX;12.0a;12.1a" "${CUDA_ARCHS}") - endif() - # marlin has limited support for turing - cuda_archs_loose_intersection(MARLIN_SM75_ARCHS "7.5" "${CUDA_ARCHS}") - # marlin arches for bf16 output (we need 9.0 for bf16 atomicAdd PTX) - if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0) - cuda_archs_loose_intersection(MARLIN_BF16_ARCHS "8.0+PTX;9.0+PTX;12.0f" "${CUDA_ARCHS}") - else() - cuda_archs_loose_intersection(MARLIN_BF16_ARCHS "8.0+PTX;9.0+PTX;12.0a;12.1a" "${CUDA_ARCHS}") - endif() - # marlin arches for fp8 input - # - sm80 doesn't support fp8 computation - # - sm90 and sm100 don't support QMMA.16832.F32.E4M3.E4M3 SAAS instruction - # so we only enable fp8 computation for SM89 (e.g. RTX 40x0) and 12.0 (e.g. RTX 50x0) - if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0) - cuda_archs_loose_intersection(MARLIN_FP8_ARCHS "8.9;12.0f" "${CUDA_ARCHS}") - else() - cuda_archs_loose_intersection(MARLIN_FP8_ARCHS "8.9;12.0a;12.1a" "${CUDA_ARCHS}") - endif() - # marlin arches for other files - cuda_archs_loose_intersection(MARLIN_OTHER_ARCHS "7.5;8.0+PTX" "${CUDA_ARCHS}") - - if (MARLIN_OTHER_ARCHS) - - # - # For the Marlin kernels we automatically generate sources for various - # preselected input type pairs and schedules. - # Generate sources: - set(MARLIN_GEN_SCRIPT - ${CMAKE_CURRENT_SOURCE_DIR}/csrc/quantization/marlin/generate_kernels.py) - file(MD5 ${MARLIN_GEN_SCRIPT} MARLIN_GEN_SCRIPT_HASH) - list(JOIN CUDA_ARCHS "," CUDA_ARCHS_STR) - set(MARLIN_GEN_SCRIPT_HASH_AND_ARCH "${MARLIN_GEN_SCRIPT_HASH}(ARCH:${CUDA_ARCHS_STR})") - - message(STATUS "Marlin generation script hash: ${MARLIN_GEN_SCRIPT_HASH_AND_ARCH}") - message(STATUS "Last run Marlin generate script hash: $CACHE{MARLIN_GEN_SCRIPT_HASH_AND_ARCH}") - - if (NOT DEFINED CACHE{MARLIN_GEN_SCRIPT_HASH_AND_ARCH} - OR NOT $CACHE{MARLIN_GEN_SCRIPT_HASH_AND_ARCH} STREQUAL ${MARLIN_GEN_SCRIPT_HASH_AND_ARCH}) - execute_process( - COMMAND ${CMAKE_COMMAND} -E env - PYTHONPATH=$ENV{PYTHONPATH} - ${Python_EXECUTABLE} ${MARLIN_GEN_SCRIPT} ${CUDA_ARCHS_STR} - RESULT_VARIABLE marlin_generation_result - OUTPUT_VARIABLE marlin_generation_result - OUTPUT_FILE ${CMAKE_CURRENT_BINARY_DIR}/marlin_generation.log - ERROR_FILE ${CMAKE_CURRENT_BINARY_DIR}/marlin_generation.log - ) - - if (NOT marlin_generation_result EQUAL 0) - message(FATAL_ERROR "Marlin generation failed." - " Result: \"${marlin_generation_result}\"" - "\nCheck the log for details: " - "${CMAKE_CURRENT_BINARY_DIR}/marlin_generation.log") - else() - set(MARLIN_GEN_SCRIPT_HASH_AND_ARCH ${MARLIN_GEN_SCRIPT_HASH_AND_ARCH} - CACHE STRING "Last run Marlin generate script hash and arch" FORCE) - message(STATUS "Marlin generation completed successfully.") - endif() - else() - message(STATUS "Marlin generation script has not changed, skipping generation.") - endif() - - if (MARLIN_ARCHS) - file(GLOB MARLIN_TEMPLATE_KERNEL_SRC "csrc/quantization/marlin/sm80_kernel_*_float16.cu") - set_gencode_flags_for_srcs( - SRCS "${MARLIN_TEMPLATE_KERNEL_SRC}" - CUDA_ARCHS "${MARLIN_ARCHS}") - if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8) - set_source_files_properties(${MARLIN_TEMPLATE_KERNEL_SRC} - PROPERTIES COMPILE_FLAGS "-static-global-template-stub=false") - endif() - list(APPEND VLLM_EXT_SRC ${MARLIN_TEMPLATE_KERNEL_SRC}) - - file(GLOB MARLIN_TEMPLATE_BF16_KERNEL_SRC "csrc/quantization/marlin/sm80_kernel_*_bfloat16.cu") - set_gencode_flags_for_srcs( - SRCS "${MARLIN_TEMPLATE_BF16_KERNEL_SRC}" - CUDA_ARCHS "${MARLIN_BF16_ARCHS}") - if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8) - set_source_files_properties(${MARLIN_TEMPLATE_BF16_KERNEL_SRC} - PROPERTIES COMPILE_FLAGS "-static-global-template-stub=false") - endif() - list(APPEND VLLM_EXT_SRC ${MARLIN_TEMPLATE_BF16_KERNEL_SRC}) - endif() - - if (MARLIN_SM75_ARCHS) - file(GLOB MARLIN_TEMPLATE_SM75_KERNEL_SRC "csrc/quantization/marlin/sm75_kernel_*.cu") - set_gencode_flags_for_srcs( - SRCS "${MARLIN_TEMPLATE_SM75_KERNEL_SRC}" - CUDA_ARCHS "${MARLIN_SM75_ARCHS}") - if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8) - set_source_files_properties(${MARLIN_TEMPLATE_SM75_KERNEL_SRC} - PROPERTIES COMPILE_FLAGS "-static-global-template-stub=false") - endif() - list(APPEND VLLM_EXT_SRC ${MARLIN_TEMPLATE_SM75_KERNEL_SRC}) - endif() - - if (MARLIN_FP8_ARCHS) - file(GLOB MARLIN_TEMPLATE_FP8_KERNEL_SRC "csrc/quantization/marlin/sm89_kernel_*.cu") - set_gencode_flags_for_srcs( - SRCS "${MARLIN_TEMPLATE_FP8_KERNEL_SRC}" - CUDA_ARCHS "${MARLIN_FP8_ARCHS}") - if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8) - set_source_files_properties(${MARLIN_TEMPLATE_FP8_KERNEL_SRC} - PROPERTIES COMPILE_FLAGS "-static-global-template-stub=false") - endif() - list(APPEND VLLM_EXT_SRC ${MARLIN_TEMPLATE_FP8_KERNEL_SRC}) - endif() - - set(MARLIN_SRCS - "csrc/quantization/marlin/marlin.cu" - "csrc/quantization/marlin/marlin_int4_fp8_preprocess.cu" - "csrc/quantization/marlin/gptq_marlin_repack.cu" - "csrc/quantization/marlin/awq_marlin_repack.cu") - set_gencode_flags_for_srcs( - SRCS "${MARLIN_SRCS}" - CUDA_ARCHS "${MARLIN_OTHER_ARCHS}") - if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8) - set_source_files_properties(${MARLIN_SRCS} - PROPERTIES COMPILE_FLAGS "-static-global-template-stub=false") - endif() - list(APPEND VLLM_EXT_SRC "${MARLIN_SRCS}") - - message(STATUS "Building Marlin kernels for archs: ${MARLIN_OTHER_ARCHS}") - else() - message(STATUS "Not building Marlin kernels as no compatible archs found" - " in CUDA target architectures") - endif() - # Expert-specialization MXFP8 blockscaled grouped kernels (SM100+). if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0) cuda_archs_loose_intersection(ES_MXFP8_GROUPED_MM_ARCHS "10.0f;11.0f" "${CUDA_ARCHS}") @@ -676,6 +537,145 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") SRCS "${VLLM_STABLE_EXT_SRC}" CUDA_ARCHS "${CUDA_ARCHS}") + # Only build Marlin kernels if we are building for at least some compatible archs. + # Keep building Marlin for 9.0 as there are some group sizes and shapes that + # are not supported by Machete yet. + + # marlin arches for fp16 output + # Family-conditional 12.0f (one cubin for SM12x family) requires CUDA >= 13.0; + # fall back to architecture-specific 12.0a;12.1a on CUDA < 13.0 (e.g. 12.8). + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0) + cuda_archs_loose_intersection(MARLIN_ARCHS "8.0+PTX;12.0f" "${CUDA_ARCHS}") + else() + cuda_archs_loose_intersection(MARLIN_ARCHS "8.0+PTX;12.0a;12.1a" "${CUDA_ARCHS}") + endif() + # marlin has limited support for turing + cuda_archs_loose_intersection(MARLIN_SM75_ARCHS "7.5" "${CUDA_ARCHS}") + # marlin arches for bf16 output (we need 9.0 for bf16 atomicAdd PTX) + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0) + cuda_archs_loose_intersection(MARLIN_BF16_ARCHS "8.0+PTX;9.0+PTX;12.0f" "${CUDA_ARCHS}") + else() + cuda_archs_loose_intersection(MARLIN_BF16_ARCHS "8.0+PTX;9.0+PTX;12.0a;12.1a" "${CUDA_ARCHS}") + endif() + # marlin arches for fp8 input + # - sm80 doesn't support fp8 computation + # - sm90 and sm100 don't support QMMA.16832.F32.E4M3.E4M3 SAAS instruction + # so we only enable fp8 computation for SM89 (e.g. RTX 40x0) and 12.0 (e.g. RTX 50x0) + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0) + cuda_archs_loose_intersection(MARLIN_FP8_ARCHS "8.9;12.0f" "${CUDA_ARCHS}") + else() + cuda_archs_loose_intersection(MARLIN_FP8_ARCHS "8.9;12.0a;12.1a" "${CUDA_ARCHS}") + endif() + # marlin arches for other files + cuda_archs_loose_intersection(MARLIN_OTHER_ARCHS "7.5;8.0+PTX" "${CUDA_ARCHS}") + + if (MARLIN_OTHER_ARCHS) + + # + # For the Marlin kernels we automatically generate sources for various + # preselected input type pairs and schedules. + # Generate sources: + set(MARLIN_GEN_SCRIPT + ${CMAKE_CURRENT_SOURCE_DIR}/csrc/libtorch_stable/quantization/marlin/generate_kernels.py) + file(MD5 ${MARLIN_GEN_SCRIPT} MARLIN_GEN_SCRIPT_HASH) + list(JOIN CUDA_ARCHS "," CUDA_ARCHS_STR) + set(MARLIN_GEN_SCRIPT_HASH_AND_ARCH "${MARLIN_GEN_SCRIPT_HASH}(ARCH:${CUDA_ARCHS_STR})") + + message(STATUS "Marlin generation script hash: ${MARLIN_GEN_SCRIPT_HASH_AND_ARCH}") + message(STATUS "Last run Marlin generate script hash: $CACHE{MARLIN_GEN_SCRIPT_HASH_AND_ARCH}") + + if (NOT DEFINED CACHE{MARLIN_GEN_SCRIPT_HASH_AND_ARCH} + OR NOT $CACHE{MARLIN_GEN_SCRIPT_HASH_AND_ARCH} STREQUAL ${MARLIN_GEN_SCRIPT_HASH_AND_ARCH}) + execute_process( + COMMAND ${CMAKE_COMMAND} -E env + PYTHONPATH=$ENV{PYTHONPATH} + ${Python_EXECUTABLE} ${MARLIN_GEN_SCRIPT} ${CUDA_ARCHS_STR} + RESULT_VARIABLE marlin_generation_result + OUTPUT_VARIABLE marlin_generation_result + OUTPUT_FILE ${CMAKE_CURRENT_BINARY_DIR}/marlin_generation.log + ERROR_FILE ${CMAKE_CURRENT_BINARY_DIR}/marlin_generation.log + ) + + if (NOT marlin_generation_result EQUAL 0) + message(FATAL_ERROR "Marlin generation failed." + " Result: \"${marlin_generation_result}\"" + "\nCheck the log for details: " + "${CMAKE_CURRENT_BINARY_DIR}/marlin_generation.log") + else() + set(MARLIN_GEN_SCRIPT_HASH_AND_ARCH ${MARLIN_GEN_SCRIPT_HASH_AND_ARCH} + CACHE STRING "Last run Marlin generate script hash and arch" FORCE) + message(STATUS "Marlin generation completed successfully.") + endif() + else() + message(STATUS "Marlin generation script has not changed, skipping generation.") + endif() + + if (MARLIN_ARCHS) + file(GLOB MARLIN_TEMPLATE_KERNEL_SRC "csrc/libtorch_stable/quantization/marlin/sm80_kernel_*_float16.cu") + set_gencode_flags_for_srcs( + SRCS "${MARLIN_TEMPLATE_KERNEL_SRC}" + CUDA_ARCHS "${MARLIN_ARCHS}") + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8) + set_source_files_properties(${MARLIN_TEMPLATE_KERNEL_SRC} + PROPERTIES COMPILE_FLAGS "-static-global-template-stub=false") + endif() + list(APPEND VLLM_STABLE_EXT_SRC ${MARLIN_TEMPLATE_KERNEL_SRC}) + + file(GLOB MARLIN_TEMPLATE_BF16_KERNEL_SRC "csrc/libtorch_stable/quantization/marlin/sm80_kernel_*_bfloat16.cu") + set_gencode_flags_for_srcs( + SRCS "${MARLIN_TEMPLATE_BF16_KERNEL_SRC}" + CUDA_ARCHS "${MARLIN_BF16_ARCHS}") + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8) + set_source_files_properties(${MARLIN_TEMPLATE_BF16_KERNEL_SRC} + PROPERTIES COMPILE_FLAGS "-static-global-template-stub=false") + endif() + list(APPEND VLLM_STABLE_EXT_SRC ${MARLIN_TEMPLATE_BF16_KERNEL_SRC}) + endif() + + if (MARLIN_SM75_ARCHS) + file(GLOB MARLIN_TEMPLATE_SM75_KERNEL_SRC "csrc/libtorch_stable/quantization/marlin/sm75_kernel_*.cu") + set_gencode_flags_for_srcs( + SRCS "${MARLIN_TEMPLATE_SM75_KERNEL_SRC}" + CUDA_ARCHS "${MARLIN_SM75_ARCHS}") + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8) + set_source_files_properties(${MARLIN_TEMPLATE_SM75_KERNEL_SRC} + PROPERTIES COMPILE_FLAGS "-static-global-template-stub=false") + endif() + list(APPEND VLLM_STABLE_EXT_SRC ${MARLIN_TEMPLATE_SM75_KERNEL_SRC}) + endif() + + if (MARLIN_FP8_ARCHS) + file(GLOB MARLIN_TEMPLATE_FP8_KERNEL_SRC "csrc/libtorch_stable/quantization/marlin/sm89_kernel_*.cu") + set_gencode_flags_for_srcs( + SRCS "${MARLIN_TEMPLATE_FP8_KERNEL_SRC}" + CUDA_ARCHS "${MARLIN_FP8_ARCHS}") + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8) + set_source_files_properties(${MARLIN_TEMPLATE_FP8_KERNEL_SRC} + PROPERTIES COMPILE_FLAGS "-static-global-template-stub=false") + endif() + list(APPEND VLLM_STABLE_EXT_SRC ${MARLIN_TEMPLATE_FP8_KERNEL_SRC}) + endif() + + set(MARLIN_SRCS + "csrc/libtorch_stable/quantization/marlin/marlin.cu" + "csrc/libtorch_stable/quantization/marlin/marlin_int4_fp8_preprocess.cu" + "csrc/libtorch_stable/quantization/marlin/gptq_marlin_repack.cu" + "csrc/libtorch_stable/quantization/marlin/awq_marlin_repack.cu") + set_gencode_flags_for_srcs( + SRCS "${MARLIN_SRCS}" + CUDA_ARCHS "${MARLIN_OTHER_ARCHS}") + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8) + set_source_files_properties(${MARLIN_SRCS} + PROPERTIES COMPILE_FLAGS "-static-global-template-stub=false") + endif() + list(APPEND VLLM_STABLE_EXT_SRC "${MARLIN_SRCS}") + + message(STATUS "Building Marlin kernels for archs: ${MARLIN_OTHER_ARCHS}") + else() + message(STATUS "Not building Marlin kernels as no compatible archs found" + " in CUDA target architectures") + endif() + # DeepSeek V3 fused A GEMM kernel (requires SM 9.0+, Hopper and later) if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0) cuda_archs_loose_intersection(DSV3_FUSED_A_GEMM_ARCHS "9.0a;10.0f;11.0f" "${CUDA_ARCHS}") diff --git a/csrc/libtorch_stable/moe/marlin_moe_wna16/kernel.h b/csrc/libtorch_stable/moe/marlin_moe_wna16/kernel.h index 09ed1a470bd..783736ab509 100644 --- a/csrc/libtorch_stable/moe/marlin_moe_wna16/kernel.h +++ b/csrc/libtorch_stable/moe/marlin_moe_wna16/kernel.h @@ -3,8 +3,8 @@ #define MARLIN_NAMESPACE_NAME marlin_moe_wna16 #endif -#include "quantization/marlin/marlin.cuh" -#include "quantization/marlin/marlin_dtypes.cuh" +#include "libtorch_stable/quantization/marlin/marlin.cuh" +#include "libtorch_stable/quantization/marlin/marlin_dtypes.cuh" #include "core/scalar_type.hpp" #define MARLIN_KERNEL_PARAMS \ diff --git a/csrc/libtorch_stable/moe/marlin_moe_wna16/marlin_template.h b/csrc/libtorch_stable/moe/marlin_moe_wna16/marlin_template.h index 9858df94573..04f90101be4 100644 --- a/csrc/libtorch_stable/moe/marlin_moe_wna16/marlin_template.h +++ b/csrc/libtorch_stable/moe/marlin_moe_wna16/marlin_template.h @@ -23,10 +23,10 @@ #define MARLIN_NAMESPACE_NAME marlin_moe_wna16 #endif -#include "quantization/marlin/marlin.cuh" -#include "quantization/marlin/marlin_dtypes.cuh" -#include "quantization/marlin/dequant.h" -#include "quantization/marlin/marlin_mma.h" +#include "libtorch_stable/quantization/marlin/marlin.cuh" +#include "libtorch_stable/quantization/marlin/marlin_dtypes.cuh" +#include "libtorch_stable/quantization/marlin/dequant.h" +#include "libtorch_stable/quantization/marlin/marlin_mma.h" #include "core/scalar_type.hpp" #define STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t) \ diff --git a/csrc/libtorch_stable/quantization/gptq_allspark/allspark_utils.cuh b/csrc/libtorch_stable/quantization/gptq_allspark/allspark_utils.cuh index ce96c2d11fe..ac33d5f2ce6 100644 --- a/csrc/libtorch_stable/quantization/gptq_allspark/allspark_utils.cuh +++ b/csrc/libtorch_stable/quantization/gptq_allspark/allspark_utils.cuh @@ -6,7 +6,7 @@ #include -#include "quantization/marlin/marlin_dtypes.cuh" +#include "libtorch_stable/quantization/marlin/marlin_dtypes.cuh" using marlin::MarlinScalarType2; namespace allspark { diff --git a/csrc/quantization/marlin/.gitignore b/csrc/libtorch_stable/quantization/marlin/.gitignore similarity index 100% rename from csrc/quantization/marlin/.gitignore rename to csrc/libtorch_stable/quantization/marlin/.gitignore diff --git a/csrc/quantization/marlin/awq_marlin_repack.cu b/csrc/libtorch_stable/quantization/marlin/awq_marlin_repack.cu similarity index 77% rename from csrc/quantization/marlin/awq_marlin_repack.cu rename to csrc/libtorch_stable/quantization/marlin/awq_marlin_repack.cu index 307bae6738e..55ce5b4e732 100644 --- a/csrc/quantization/marlin/awq_marlin_repack.cu +++ b/csrc/libtorch_stable/quantization/marlin/awq_marlin_repack.cu @@ -1,6 +1,13 @@ #include "marlin.cuh" -#include "core/registration.h" +#include +#include +#include +#include +#include +#include + +#include "libtorch_stable/torch_utils.h" namespace marlin { @@ -218,56 +225,55 @@ __global__ void awq_marlin_repack_kernel( b_q_weight_ptr, out_ptr, size_k, size_n); \ } -torch::Tensor awq_marlin_repack(torch::Tensor& b_q_weight, int64_t size_k, - int64_t size_n, int64_t num_bits, - bool is_a_8bit) { +torch::stable::Tensor awq_marlin_repack(torch::stable::Tensor& b_q_weight, + int64_t size_k, int64_t size_n, + int64_t num_bits, bool is_a_8bit) { // Verify compatibility with marlin tile of 16x64 - TORCH_CHECK(size_k % marlin::tile_k_size == 0, "size_k = ", size_k, - " is not divisible by tile_k_size = ", marlin::tile_k_size); - TORCH_CHECK(size_n % marlin::tile_n_size == 0, "size_n = ", size_n, - " is not divisible by tile_n_size = ", marlin::tile_n_size); + STD_TORCH_CHECK(size_k % marlin::tile_k_size == 0, "size_k = ", size_k, + " is not divisible by tile_k_size = ", marlin::tile_k_size); + STD_TORCH_CHECK(size_n % marlin::tile_n_size == 0, "size_n = ", size_n, + " is not divisible by tile_n_size = ", marlin::tile_n_size); - TORCH_CHECK(num_bits == 4 || num_bits == 8, - "num_bits must be 4 or 8. Got = ", num_bits); + STD_TORCH_CHECK(num_bits == 4 || num_bits == 8, + "num_bits must be 4 or 8. Got = ", num_bits); int const pack_factor = 32 / num_bits; // Verify B - TORCH_CHECK(b_q_weight.size(0) == size_k, - "b_q_weight.size(0) = ", b_q_weight.size(0), - " is not size_k = ", size_k); - TORCH_CHECK((size_n / pack_factor) == b_q_weight.size(1), - "Shape mismatch: b_q_weight.size(1) = ", b_q_weight.size(1), - ", size_n = ", size_n, ", pack_factor = ", pack_factor); + STD_TORCH_CHECK(b_q_weight.size(0) == size_k, + "b_q_weight.size(0) = ", b_q_weight.size(0), + " is not size_k = ", size_k); + STD_TORCH_CHECK((size_n / pack_factor) == b_q_weight.size(1), + "Shape mismatch: b_q_weight.size(1) = ", b_q_weight.size(1), + ", size_n = ", size_n, ", pack_factor = ", pack_factor); // Verify device and strides - TORCH_CHECK(b_q_weight.device().is_cuda(), "b_q_weight is not on GPU"); - TORCH_CHECK(b_q_weight.is_contiguous(), "b_q_weight is not contiguous"); - TORCH_CHECK(b_q_weight.dtype() == at::kInt, "b_q_weight type is not kInt"); + STD_TORCH_CHECK(b_q_weight.is_cuda(), "b_q_weight is not on GPU"); + STD_TORCH_CHECK(b_q_weight.is_contiguous(), "b_q_weight is not contiguous"); + STD_TORCH_CHECK( + b_q_weight.scalar_type() == torch::headeronly::ScalarType::Int, + "b_q_weight type is not kInt"); + + const int32_t device_index = b_q_weight.get_device_index(); + torch::stable::accelerator::DeviceGuard device_guard(device_index); + const cudaStream_t stream = get_current_cuda_stream(device_index); // Alloc buffers - const at::cuda::OptionalCUDAGuard device_guard(device_of(b_q_weight)); - auto options = torch::TensorOptions() - .dtype(b_q_weight.dtype()) - .device(b_q_weight.device()); - torch::Tensor out = torch::empty( + torch::stable::Tensor out = torch::stable::empty( {size_k / marlin::tile_size, size_n * marlin::tile_size / pack_factor}, - options); + b_q_weight.scalar_type(), std::nullopt, b_q_weight.device()); // Get ptrs uint32_t const* b_q_weight_ptr = - reinterpret_cast(b_q_weight.data_ptr()); - uint32_t* out_ptr = reinterpret_cast(out.data_ptr()); + reinterpret_cast(b_q_weight.const_data_ptr()); + uint32_t* out_ptr = reinterpret_cast(out.mutable_data_ptr()); - // Get dev info - int dev = b_q_weight.get_device(); - cudaStream_t stream = at::cuda::getCurrentCUDAStream(dev); int blocks; - cudaDeviceGetAttribute(&blocks, cudaDevAttrMultiProcessorCount, dev); + cudaDeviceGetAttribute(&blocks, cudaDevAttrMultiProcessorCount, device_index); int max_shared_mem = 0; cudaDeviceGetAttribute(&max_shared_mem, - cudaDevAttrMaxSharedMemoryPerBlockOptin, dev); - TORCH_CHECK(max_shared_mem > 0); + cudaDevAttrMaxSharedMemoryPerBlockOptin, device_index); + STD_TORCH_CHECK(max_shared_mem > 0); if (false) { } @@ -276,13 +282,13 @@ torch::Tensor awq_marlin_repack(torch::Tensor& b_q_weight, int64_t size_k, CALL_IF(4, true) CALL_IF(8, true) else { - TORCH_CHECK(false, "Unsupported repack config: num_bits = ", num_bits, - ", is_a_8bit = ", is_a_8bit); + STD_TORCH_CHECK(false, "Unsupported repack config: num_bits = ", num_bits, + ", is_a_8bit = ", is_a_8bit); } return out; } -TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) { - m.impl("awq_marlin_repack", &awq_marlin_repack); +STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) { + m.impl("awq_marlin_repack", TORCH_BOX(&awq_marlin_repack)); } diff --git a/csrc/quantization/marlin/dequant.h b/csrc/libtorch_stable/quantization/marlin/dequant.h similarity index 100% rename from csrc/quantization/marlin/dequant.h rename to csrc/libtorch_stable/quantization/marlin/dequant.h diff --git a/csrc/quantization/marlin/generate_kernels.py b/csrc/libtorch_stable/quantization/marlin/generate_kernels.py similarity index 99% rename from csrc/quantization/marlin/generate_kernels.py rename to csrc/libtorch_stable/quantization/marlin/generate_kernels.py index 7b316037ec6..2a038479893 100644 --- a/csrc/quantization/marlin/generate_kernels.py +++ b/csrc/libtorch_stable/quantization/marlin/generate_kernels.py @@ -303,7 +303,7 @@ def generate_new_kernels(): if not SUPPORT_FP8 and kernel_selector_str != FILE_HEAD_COMMENT: kernel_selector_str += ( "else if (a_type == vllm::kFE4M3fn)\n" - " TORCH_CHECK(false, " + " STD_TORCH_CHECK(false, " '"marlin kernel with fp8 activation is not built.");' ) diff --git a/csrc/quantization/marlin/gptq_marlin_repack.cu b/csrc/libtorch_stable/quantization/marlin/gptq_marlin_repack.cu similarity index 77% rename from csrc/quantization/marlin/gptq_marlin_repack.cu rename to csrc/libtorch_stable/quantization/marlin/gptq_marlin_repack.cu index 796e6c5359d..cafa212bccb 100644 --- a/csrc/quantization/marlin/gptq_marlin_repack.cu +++ b/csrc/libtorch_stable/quantization/marlin/gptq_marlin_repack.cu @@ -1,6 +1,13 @@ #include "marlin.cuh" -#include "core/registration.h" +#include +#include +#include +#include +#include +#include + +#include "libtorch_stable/torch_utils.h" namespace marlin { @@ -275,64 +282,66 @@ __global__ void gptq_marlin_repack_kernel( b_q_weight_ptr, perm_ptr, out_ptr, size_k, size_n); \ } -torch::Tensor gptq_marlin_repack(torch::Tensor& b_q_weight, torch::Tensor& perm, - int64_t size_k, int64_t size_n, - int64_t num_bits, bool is_a_8bit) { +torch::stable::Tensor gptq_marlin_repack(torch::stable::Tensor& b_q_weight, + torch::stable::Tensor& perm, + int64_t size_k, int64_t size_n, + int64_t num_bits, bool is_a_8bit) { // Verify compatibility with marlin tile of 16x64 - TORCH_CHECK(size_k % marlin::tile_k_size == 0, "size_k = ", size_k, - " is not divisible by tile_k_size = ", marlin::tile_k_size); - TORCH_CHECK(size_n % marlin::tile_n_size == 0, "size_n = ", size_n, - " is not divisible by tile_n_size = ", marlin::tile_n_size); + STD_TORCH_CHECK(size_k % marlin::tile_k_size == 0, "size_k = ", size_k, + " is not divisible by tile_k_size = ", marlin::tile_k_size); + STD_TORCH_CHECK(size_n % marlin::tile_n_size == 0, "size_n = ", size_n, + " is not divisible by tile_n_size = ", marlin::tile_n_size); - TORCH_CHECK(num_bits == 4 || num_bits == 8, - "num_bits must be 4 or 8. Got = ", num_bits); + STD_TORCH_CHECK(num_bits == 4 || num_bits == 8, + "num_bits must be 4 or 8. Got = ", num_bits); int const pack_factor = 32 / num_bits; // Verify B - TORCH_CHECK((size_k / pack_factor) == b_q_weight.size(0), - "Shape mismatch: b_q_weight.size(0) = ", b_q_weight.size(0), - ", size_k = ", size_k, ", pack_factor = ", pack_factor); - TORCH_CHECK(b_q_weight.size(1) == size_n, - "b_q_weight.size(1) = ", b_q_weight.size(1), - " is not size_n = ", size_n); + STD_TORCH_CHECK((size_k / pack_factor) == b_q_weight.size(0), + "Shape mismatch: b_q_weight.size(0) = ", b_q_weight.size(0), + ", size_k = ", size_k, ", pack_factor = ", pack_factor); + STD_TORCH_CHECK(b_q_weight.size(1) == size_n, + "b_q_weight.size(1) = ", b_q_weight.size(1), + " is not size_n = ", size_n); // Verify device and strides - TORCH_CHECK(b_q_weight.device().is_cuda(), "b_q_weight is not on GPU"); - TORCH_CHECK(b_q_weight.is_contiguous(), "b_q_weight is not contiguous"); - TORCH_CHECK(b_q_weight.dtype() == at::kInt, "b_q_weight type is not kInt"); + STD_TORCH_CHECK(b_q_weight.is_cuda(), "b_q_weight is not on GPU"); + STD_TORCH_CHECK(b_q_weight.is_contiguous(), "b_q_weight is not contiguous"); + STD_TORCH_CHECK( + b_q_weight.scalar_type() == torch::headeronly::ScalarType::Int, + "b_q_weight type is not kInt"); - TORCH_CHECK(perm.device().is_cuda(), "perm is not on GPU"); - TORCH_CHECK(perm.is_contiguous(), "perm is not contiguous"); - TORCH_CHECK(perm.dtype() == at::kInt, "perm type is not at::kInt"); + STD_TORCH_CHECK(perm.is_cuda(), "perm is not on GPU"); + STD_TORCH_CHECK(perm.is_contiguous(), "perm is not contiguous"); + STD_TORCH_CHECK(perm.scalar_type() == torch::headeronly::ScalarType::Int, + "perm type is not at::kInt"); + + const int32_t device_index = b_q_weight.get_device_index(); + torch::stable::accelerator::DeviceGuard device_guard(device_index); + const cudaStream_t stream = get_current_cuda_stream(device_index); // Alloc buffers - const at::cuda::OptionalCUDAGuard device_guard(device_of(b_q_weight)); - auto options = torch::TensorOptions() - .dtype(b_q_weight.dtype()) - .device(b_q_weight.device()); - torch::Tensor out = torch::empty( + torch::stable::Tensor out = torch::stable::empty( {size_k / marlin::tile_size, size_n * marlin::tile_size / pack_factor}, - options); + b_q_weight.scalar_type(), std::nullopt, b_q_weight.device()); // Detect if there is act_order bool has_perm = perm.size(0) != 0; // Get ptrs uint32_t const* b_q_weight_ptr = - reinterpret_cast(b_q_weight.data_ptr()); - uint32_t const* perm_ptr = reinterpret_cast(perm.data_ptr()); - uint32_t* out_ptr = reinterpret_cast(out.data_ptr()); + reinterpret_cast(b_q_weight.const_data_ptr()); + uint32_t const* perm_ptr = + reinterpret_cast(perm.const_data_ptr()); + uint32_t* out_ptr = reinterpret_cast(out.mutable_data_ptr()); - // Get dev info - int dev = b_q_weight.get_device(); - cudaStream_t stream = at::cuda::getCurrentCUDAStream(dev); int blocks; - cudaDeviceGetAttribute(&blocks, cudaDevAttrMultiProcessorCount, dev); + cudaDeviceGetAttribute(&blocks, cudaDevAttrMultiProcessorCount, device_index); int max_shared_mem = 0; cudaDeviceGetAttribute(&max_shared_mem, - cudaDevAttrMaxSharedMemoryPerBlockOptin, dev); - TORCH_CHECK(max_shared_mem > 0); + cudaDevAttrMaxSharedMemoryPerBlockOptin, device_index); + STD_TORCH_CHECK(max_shared_mem > 0); if (false) { } @@ -345,13 +354,13 @@ torch::Tensor gptq_marlin_repack(torch::Tensor& b_q_weight, torch::Tensor& perm, CALL_IF(8, false, true) else { - TORCH_CHECK(false, "Unsupported repack config: num_bits = ", num_bits, - ", has_perm = ", has_perm, ", is_a_8bit = ", is_a_8bit); + STD_TORCH_CHECK(false, "Unsupported repack config: num_bits = ", num_bits, + ", has_perm = ", has_perm, ", is_a_8bit = ", is_a_8bit); } return out; } -TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) { - m.impl("gptq_marlin_repack", &gptq_marlin_repack); +STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) { + m.impl("gptq_marlin_repack", TORCH_BOX(&gptq_marlin_repack)); } diff --git a/csrc/quantization/marlin/kernel.h b/csrc/libtorch_stable/quantization/marlin/kernel.h similarity index 100% rename from csrc/quantization/marlin/kernel.h rename to csrc/libtorch_stable/quantization/marlin/kernel.h diff --git a/csrc/quantization/marlin/marlin.cu b/csrc/libtorch_stable/quantization/marlin/marlin.cu similarity index 61% rename from csrc/quantization/marlin/marlin.cu rename to csrc/libtorch_stable/quantization/marlin/marlin.cu index 721c206c33f..63fea239e4a 100644 --- a/csrc/quantization/marlin/marlin.cu +++ b/csrc/libtorch_stable/quantization/marlin/marlin.cu @@ -24,7 +24,15 @@ #endif #include "kernel.h" -#include "core/registration.h" + +#include +#include +#include +#include +#include +#include + +#include "libtorch_stable/torch_utils.h" #define STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t) \ static_assert(std::is_same::value || \ @@ -46,19 +54,22 @@ __global__ void permute_cols_kernel(int4 const* __restrict__ a_int4_ptr, } // namespace marlin -torch::Tensor marlin_gemm( - torch::Tensor& a, std::optional c_or_none, - torch::Tensor& b_q_weight, - std::optional const& b_bias_or_none, torch::Tensor& b_scales, - std::optional const& b_zeros_or_none, - std::optional const& g_idx_or_none, - std::optional const& perm_or_none, torch::Tensor& workspace, - vllm::ScalarTypeId const& b_type_id, int64_t size_m, int64_t size_n, - int64_t size_k, bool is_k_full, bool use_atomic_add, bool use_fp32_reduce, - bool is_zp_float) { - TORCH_CHECK_NOT_IMPLEMENTED(false, - "marlin_gemm(..) requires CUDA_ARCH >= 7.5"); - return torch::empty({1, 1}); +torch::stable::Tensor marlin_gemm( + torch::stable::Tensor& a, std::optional c_or_none, + torch::stable::Tensor& b_q_weight, + std::optional const& b_bias_or_none, + torch::stable::Tensor& b_scales, + std::optional const& a_scales_or_none, + std::optional const& global_scale_or_none, + std::optional const& b_zeros_or_none, + std::optional const& g_idx_or_none, + std::optional const& perm_or_none, + torch::stable::Tensor& workspace, vllm::ScalarTypeId const& b_type_id, + int64_t size_m, int64_t size_n, int64_t size_k, bool is_k_full, + bool use_atomic_add, bool use_fp32_reduce, bool is_zp_float) { + STD_TORCH_CHECK_NOT_IMPLEMENTED(false, + "marlin_gemm(..) requires CUDA_ARCH >= 7.5"); + return torch::stable::empty({1, 1}); } #else @@ -323,18 +334,18 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias, int thread_n_init, int sms, bool use_atomic_add, bool use_fp32_reduce, bool is_zp_float) { bool is_a_8bit = a_type.size_bits() == 8; - TORCH_CHECK(prob_m > 0 && prob_n > 0 && prob_k > 0, "Invalid MNK = [", prob_m, - ", ", prob_n, ", ", prob_k, "]"); + STD_TORCH_CHECK(prob_m > 0 && prob_n > 0 && prob_k > 0, "Invalid MNK = [", + prob_m, ", ", prob_n, ", ", prob_k, "]"); int group_blocks = 0; if (has_act_order) { if (is_k_full) { - TORCH_CHECK(group_size != -1); + STD_TORCH_CHECK(group_size != -1); group_blocks = group_size / 16; - TORCH_CHECK(prob_k % group_blocks == 0, "prob_k = ", prob_k, - " is not divisible by group_blocks = ", group_blocks); + STD_TORCH_CHECK(prob_k % group_blocks == 0, "prob_k = ", prob_k, + " is not divisible by group_blocks = ", group_blocks); } else { - TORCH_CHECK(group_size == 0); + STD_TORCH_CHECK(group_size == 0); group_blocks = 0; } } else { @@ -342,8 +353,8 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias, group_blocks = -1; } else { group_blocks = group_size / 16; - TORCH_CHECK(prob_k % group_blocks == 0, "prob_k = ", prob_k, - " is not divisible by group_blocks = ", group_blocks); + STD_TORCH_CHECK(prob_k % group_blocks == 0, "prob_k = ", prob_k, + " is not divisible by group_blocks = ", group_blocks); } } @@ -384,25 +395,25 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias, int max_shared_mem = 0; cudaDeviceGetAttribute(&max_shared_mem, cudaDevAttrMaxSharedMemoryPerBlockOptin, dev); - TORCH_CHECK(max_shared_mem > 0); + STD_TORCH_CHECK(max_shared_mem > 0); int major_capability, minor_capability; cudaDeviceGetAttribute(&major_capability, cudaDevAttrComputeCapabilityMajor, dev); cudaDeviceGetAttribute(&minor_capability, cudaDevAttrComputeCapabilityMinor, dev); - TORCH_CHECK(major_capability * 10 + minor_capability >= 75, - "marlin kernel only support Turing or newer GPUs."); + STD_TORCH_CHECK(major_capability * 10 + minor_capability >= 75, + "marlin kernel only support Turing or newer GPUs."); int stages = 4; if (major_capability == 7 && minor_capability == 5) { stages = 2; - TORCH_CHECK(a_type == vllm::kFloat16 || a_type == vllm::kS8, - "Turing only support FP16 or INT8 activation."); + STD_TORCH_CHECK(a_type == vllm::kFloat16 || a_type == vllm::kS8, + "Turing only support FP16 or INT8 activation."); } if (a_type == vllm::kFE4M3fn) { - TORCH_CHECK(major_capability * 10 + minor_capability >= 89, - "FP8 only support Ada Lovelace or newer GPUs."); - TORCH_CHECK( + STD_TORCH_CHECK(major_capability * 10 + minor_capability >= 89, + "FP8 only support Ada Lovelace or newer GPUs."); + STD_TORCH_CHECK( major_capability * 10 + minor_capability == 89 || major_capability == 12, "Marlin W4A8-FP8 only support SM89 or SM12x device (It is slower than " @@ -432,10 +443,10 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias, if (thread_k != -1 && thread_n != -1) { thread_tfg = thread_config_t{thread_k, thread_n, default_threads}; exec_cfg = exec_config_t{1, thread_tfg}; - TORCH_CHECK(prob_n % thread_n == 0, "prob_n = ", prob_n, - " is not divisible by thread_n = ", thread_n); - TORCH_CHECK(prob_k % thread_k == 0, "prob_k = ", prob_k, - " is not divisible by thread_k = ", thread_k); + STD_TORCH_CHECK(prob_n % thread_n == 0, "prob_n = ", prob_n, + " is not divisible by thread_n = ", thread_n); + STD_TORCH_CHECK(prob_k % thread_k == 0, "prob_k = ", prob_k, + " is not divisible by thread_k = ", thread_k); } else { // Auto config exec_cfg = determine_exec_config( @@ -474,7 +485,7 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias, int thread_k_blocks = thread_k / 16; int thread_n_blocks = thread_n / 16; - TORCH_CHECK( + STD_TORCH_CHECK( is_valid_config(thread_tfg, thread_m_blocks, prob_m_split, prob_n, prob_k, num_bits, group_size, has_act_order, is_k_full, has_zp, is_zp_float, is_a_8bit, stages, @@ -495,14 +506,15 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias, num_threads, is_zp_float, stages); if (kernel == MarlinDefault) { - TORCH_CHECK(false, "Unsupported shapes: MNK = [", prob_m, ", ", prob_n, - ", ", prob_k, "]", ", has_act_order = ", has_act_order, - ", num_groups = ", num_groups, ", group_size = ", group_size, - ", prob_m_split = ", prob_m_split, - ", thread_m_blocks = ", thread_m_blocks, - ", thread_n_blocks = ", thread_n_blocks, - ", thread_k_blocks = ", thread_k_blocks, - ", num_threads = ", num_threads, ", num_bits = ", num_bits); + STD_TORCH_CHECK( + false, "Unsupported shapes: MNK = [", prob_m, ", ", prob_n, ", ", + prob_k, "]", ", has_act_order = ", has_act_order, + ", num_groups = ", num_groups, ", group_size = ", group_size, + ", prob_m_split = ", prob_m_split, + ", thread_m_blocks = ", thread_m_blocks, + ", thread_n_blocks = ", thread_n_blocks, + ", thread_k_blocks = ", thread_k_blocks, + ", num_threads = ", num_threads, ", num_bits = ", num_bits); } cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, @@ -530,71 +542,76 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias, } // namespace marlin -torch::Tensor marlin_gemm( - torch::Tensor& a, std::optional c_or_none, - torch::Tensor& b_q_weight, - std::optional const& b_bias_or_none, torch::Tensor& b_scales, - std::optional const& a_scales_or_none, - std::optional const& global_scale_or_none, - std::optional const& b_zeros_or_none, - std::optional const& g_idx_or_none, - std::optional const& perm_or_none, torch::Tensor& workspace, - vllm::ScalarTypeId const& b_type_id, int64_t size_m, int64_t size_n, - int64_t size_k, bool is_k_full, bool use_atomic_add, bool use_fp32_reduce, - bool is_zp_float) { +torch::stable::Tensor marlin_gemm( + torch::stable::Tensor& a, std::optional c_or_none, + torch::stable::Tensor& b_q_weight, + std::optional const& b_bias_or_none, + torch::stable::Tensor& b_scales, + std::optional const& a_scales_or_none, + std::optional const& global_scale_or_none, + std::optional const& b_zeros_or_none, + std::optional const& g_idx_or_none, + std::optional const& perm_or_none, + torch::stable::Tensor& workspace, vllm::ScalarTypeId const& b_type_id, + int64_t size_m, int64_t size_n, int64_t size_k, bool is_k_full, + bool use_atomic_add, bool use_fp32_reduce, bool is_zp_float) { vllm::ScalarTypeId a_type_id, c_type_id, s_type_id; - auto c_dtype = a.dtype(); - if (a.scalar_type() == at::ScalarType::Half) { + auto c_scalar_type = a.scalar_type(); + if (a.scalar_type() == torch::headeronly::ScalarType::Half) { a_type_id = vllm::kFloat16.id(); c_type_id = vllm::kFloat16.id(); - } else if (a.scalar_type() == at::ScalarType::BFloat16) { + } else if (a.scalar_type() == torch::headeronly::ScalarType::BFloat16) { a_type_id = vllm::kBFloat16.id(); c_type_id = vllm::kBFloat16.id(); } else { - c_dtype = b_scales.dtype(); - if (b_scales.scalar_type() == at::ScalarType::Half) { + c_scalar_type = b_scales.scalar_type(); + if (b_scales.scalar_type() == torch::headeronly::ScalarType::Half) { c_type_id = vllm::kFloat16.id(); - } else if (b_scales.scalar_type() == at::ScalarType::BFloat16) { + } else if (b_scales.scalar_type() == + torch::headeronly::ScalarType::BFloat16) { c_type_id = vllm::kBFloat16.id(); } else { c_type_id = vllm::kBFloat16.id(); - TORCH_CHECK(c_or_none.has_value(), "c must be passed for W4A8-FP4"); - torch::Tensor c = c_or_none.value(); - c_dtype = c.dtype(); + STD_TORCH_CHECK(c_or_none.has_value(), "c must be passed for W4A8-FP4"); + torch::stable::Tensor c = c_or_none.value(); + c_scalar_type = c.scalar_type(); - if (c.scalar_type() == at::ScalarType::Half) { + if (c.scalar_type() == torch::headeronly::ScalarType::Half) { c_type_id = vllm::kFloat16.id(); - } else if (c.scalar_type() == at::ScalarType::BFloat16) { + } else if (c.scalar_type() == torch::headeronly::ScalarType::BFloat16) { c_type_id = vllm::kBFloat16.id(); } else { - TORCH_CHECK(false, "unsupported c dtype"); + STD_TORCH_CHECK(false, "unsupported c dtype"); } } - if (a.scalar_type() == at::ScalarType::Float8_e4m3fn) { + if (a.scalar_type() == torch::headeronly::ScalarType::Float8_e4m3fn) { a_type_id = vllm::kFE4M3fn.id(); - } else if (a.scalar_type() == at::ScalarType::Char) { + } else if (a.scalar_type() == torch::headeronly::ScalarType::Char) { a_type_id = vllm::kS8.id(); } else { - TORCH_CHECK(false, "unsupported `a` scalar_type"); + STD_TORCH_CHECK(false, "unsupported `a` scalar_type"); } } s_type_id = c_type_id; if (b_type_id == vllm::kFE2M1f.id()) { - if (b_scales.scalar_type() == at::ScalarType::Float8_e4m3fn) { + if (b_scales.scalar_type() == + torch::headeronly::ScalarType::Float8_e4m3fn) { s_type_id = vllm::kFE4M3fn.id(); - } else if (b_scales.scalar_type() == at::ScalarType::Float8_e8m0fnu) { + } else if (b_scales.scalar_type() == + torch::headeronly::ScalarType::Float8_e8m0fnu) { s_type_id = vllm::kFE8M0fnu.id(); } else { - TORCH_CHECK(false, - "When b_type = float4_e2m1f, b_scale scalar type must be", - "float8_e4m3fn (for NVFP4) or float8_e8m0fnu (for MXFP4)."); + STD_TORCH_CHECK( + false, "When b_type = float4_e2m1f, b_scale scalar type must be", + "float8_e4m3fn (for NVFP4) or float8_e8m0fnu (for MXFP4)."); } } else if (b_type_id == vllm::kFE4M3fn.id() && - b_scales.scalar_type() == at::ScalarType::Float8_e8m0fnu) { + b_scales.scalar_type() == + torch::headeronly::ScalarType::Float8_e8m0fnu) { s_type_id = vllm::kFE8M0fnu.id(); } @@ -606,54 +623,58 @@ torch::Tensor marlin_gemm( int pack_factor = 32 / b_type.size_bits(); // Verify A - TORCH_CHECK(a.size(0) == size_m, "Shape mismatch: a.size(0) = ", a.size(0), - ", size_m = ", size_m); - TORCH_CHECK(a.size(1) == size_k, "Shape mismatch: a.size(1) = ", a.size(1), - ", size_k = ", size_k); + STD_TORCH_CHECK(a.size(0) == size_m, + "Shape mismatch: a.size(0) = ", a.size(0), + ", size_m = ", size_m); + STD_TORCH_CHECK(a.size(1) == size_k, + "Shape mismatch: a.size(1) = ", a.size(1), + ", size_k = ", size_k); // Verify B - TORCH_CHECK( + STD_TORCH_CHECK( size_k % MARLIN_NAMESPACE_NAME::tile_size == 0, "size_k = ", size_k, " is not divisible by tile_size = ", MARLIN_NAMESPACE_NAME::tile_size); - TORCH_CHECK((size_k / MARLIN_NAMESPACE_NAME::tile_size) == b_q_weight.size(0), - "Shape mismatch: b_q_weight.size(0) = ", b_q_weight.size(0), - ", size_k = ", size_k, - ", tile_size = ", MARLIN_NAMESPACE_NAME::tile_size); - TORCH_CHECK( + STD_TORCH_CHECK( + (size_k / MARLIN_NAMESPACE_NAME::tile_size) == b_q_weight.size(0), + "Shape mismatch: b_q_weight.size(0) = ", b_q_weight.size(0), + ", size_k = ", size_k, + ", tile_size = ", MARLIN_NAMESPACE_NAME::tile_size); + STD_TORCH_CHECK( b_q_weight.size(1) % MARLIN_NAMESPACE_NAME::tile_size == 0, "b_q_weight.size(1) = ", b_q_weight.size(1), " is not divisible by tile_size = ", MARLIN_NAMESPACE_NAME::tile_size); int actual_size_n = (b_q_weight.size(1) / MARLIN_NAMESPACE_NAME::tile_size) * pack_factor; - TORCH_CHECK(size_n == actual_size_n, "size_n = ", size_n, - ", actual_size_n = ", actual_size_n); + STD_TORCH_CHECK(size_n == actual_size_n, "size_n = ", size_n, + ", actual_size_n = ", actual_size_n); // Verify device and strides - TORCH_CHECK(a.device().is_cuda(), "A is not on GPU"); - TORCH_CHECK(a.stride(1) == 1, "A.stride(1) is not 1"); + STD_TORCH_CHECK(a.device().is_cuda(), "A is not on GPU"); + STD_TORCH_CHECK(a.stride(1) == 1, "A.stride(1) is not 1"); // We use int4 (16 bytes) to load A, so A must aligned to 16 bytes - TORCH_CHECK(a.stride(0) % 8 == 0, "A.stride(0) must divisible by 8"); - TORCH_CHECK(((uint64_t)a.data_ptr()) % 16 == 0, "A must aligned to 16 bytes"); + STD_TORCH_CHECK(a.stride(0) % 8 == 0, "A.stride(0) must divisible by 8"); + STD_TORCH_CHECK(((uint64_t)a.const_data_ptr()) % 16 == 0, + "A must aligned to 16 bytes"); - TORCH_CHECK(b_q_weight.device().is_cuda(), "b_q_weight is not on GPU"); - TORCH_CHECK(b_q_weight.is_contiguous(), "b_q_weight is not contiguous"); + STD_TORCH_CHECK(b_q_weight.device().is_cuda(), "b_q_weight is not on GPU"); + STD_TORCH_CHECK(b_q_weight.is_contiguous(), "b_q_weight is not contiguous"); - TORCH_CHECK(b_scales.device().is_cuda(), "b_scales is not on GPU"); - TORCH_CHECK(b_scales.is_contiguous(), "b_scales is not contiguous"); + STD_TORCH_CHECK(b_scales.device().is_cuda(), "b_scales is not on GPU"); + STD_TORCH_CHECK(b_scales.is_contiguous(), "b_scales is not contiguous"); - torch::Tensor a_scales; - auto options = torch::TensorOptions().dtype(c_dtype).device(a.device()); - auto options_fp32 = - torch::TensorOptions().dtype(at::kFloat).device(a.device()); + torch::stable::Tensor a_scales; + const auto device = a.device(); if (a_scales_or_none.has_value()) { a_scales = a_scales_or_none.value(); - TORCH_CHECK(a_type.size_bits() == 8, - "a_scales can only be used for 8bit activation."); + STD_TORCH_CHECK(a_type.size_bits() == 8, + "a_scales can only be used for 8bit activation."); } else { - a_scales = torch::empty({0}, options_fp32); - TORCH_CHECK(a_type.size_bits() != 8, - "the a_scales parameter must be passed for 8bit activation."); + a_scales = torch::stable::empty({0}, torch::headeronly::ScalarType::Float, + std::nullopt, device); + STD_TORCH_CHECK( + a_type.size_bits() != 8, + "the a_scales parameter must be passed for 8bit activation."); } // thread_k: `k` size of a thread_tile in `weights` (can usually be left as @@ -664,84 +685,93 @@ torch::Tensor marlin_gemm( int thread_n = -1; // sms: number of SMs to use for the kernel int sms = -1; - cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, a.get_device()); + const int32_t device_index = a.get_device_index(); + cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, device_index); // Alloc buffers - const at::cuda::OptionalCUDAGuard device_guard(device_of(a)); - torch::Tensor c; + torch::stable::accelerator::DeviceGuard device_guard(device_index); + torch::stable::Tensor c; if (c_or_none.has_value()) { c = c_or_none.value(); - TORCH_CHECK(c.device().is_cuda(), "c is not on GPU"); - TORCH_CHECK(c.is_contiguous(), "c is not contiguous"); - TORCH_CHECK(c.size(0) == size_m, "Shape mismatch: c.size(0) = ", c.size(0), - ", size_m = ", size_m); - TORCH_CHECK(c.size(1) == size_n, "Shape mismatch: c.size(1) = ", c.size(1), - ", size_n = ", size_n); + STD_TORCH_CHECK(c.device().is_cuda(), "c is not on GPU"); + STD_TORCH_CHECK(c.is_contiguous(), "c is not contiguous"); + STD_TORCH_CHECK(c.size(0) == size_m, + "Shape mismatch: c.size(0) = ", c.size(0), + ", size_m = ", size_m); + STD_TORCH_CHECK(c.size(1) == size_n, + "Shape mismatch: c.size(1) = ", c.size(1), + ", size_n = ", size_n); } else { - c = torch::empty({size_m, size_n}, options); + c = torch::stable::empty({size_m, size_n}, c_scalar_type, std::nullopt, + device); } if (size_m == 0) return c; // Alloc C tmp buffer that is going to be used for the global reduce - torch::Tensor c_tmp; + torch::stable::Tensor c_tmp; if (use_fp32_reduce) { int max_m_block_size = (size_m + 16 - 1) / 16 * 16; max_m_block_size = min(max_m_block_size, 64); int max_c_tmp_size = sms * max_m_block_size * MARLIN_NAMESPACE_NAME::max_thread_n; - c_tmp = torch::empty({max_c_tmp_size}, options_fp32); + c_tmp = torch::stable::empty({max_c_tmp_size}, + torch::headeronly::ScalarType::Float, + std::nullopt, device); } else { - c_tmp = torch::empty({0}, options_fp32); + c_tmp = torch::stable::empty({0}, torch::headeronly::ScalarType::Float, + std::nullopt, device); } // Detect groupsize and act_order int num_groups = -1; int group_size = -1; - int rank = b_scales.sizes().size(); - TORCH_CHECK(rank == 2, "b_scales rank = ", rank, " is not 2"); - TORCH_CHECK(b_scales.size(1) == size_n, "b_scales dim 1 = ", b_scales.size(1), - " is not size_n = ", size_n); + int rank = b_scales.dim(); + STD_TORCH_CHECK(rank == 2, "b_scales rank = ", rank, " is not 2"); + STD_TORCH_CHECK(b_scales.size(1) == size_n, + "b_scales dim 1 = ", b_scales.size(1), + " is not size_n = ", size_n); num_groups = b_scales.size(0); - torch::Tensor g_idx, perm, a_tmp; + torch::stable::Tensor g_idx, perm, a_tmp; if (g_idx_or_none.has_value() && perm_or_none.has_value()) { g_idx = g_idx_or_none.value(); perm = perm_or_none.value(); - TORCH_CHECK(g_idx.device().is_cuda(), "g_idx is not on GPU"); - TORCH_CHECK(g_idx.is_contiguous(), "g_idx is not contiguous"); - TORCH_CHECK(perm.device().is_cuda(), "perm is not on GPU"); - TORCH_CHECK(perm.is_contiguous(), "perm is not contiguous"); + STD_TORCH_CHECK(g_idx.device().is_cuda(), "g_idx is not on GPU"); + STD_TORCH_CHECK(g_idx.is_contiguous(), "g_idx is not contiguous"); + STD_TORCH_CHECK(perm.device().is_cuda(), "perm is not on GPU"); + STD_TORCH_CHECK(perm.is_contiguous(), "perm is not contiguous"); // Verify g_idx and perm - TORCH_CHECK((g_idx.size(-1) == 0 && perm.size(-1) == 0) || - (g_idx.size(-1) == size_k && perm.size(-1) == size_k), - "Unexpected g_idx.size(-1) = ", g_idx.size(-1), - " and perm.size(-1) = ", perm.size(-1), - ", where size_k = ", size_k); + STD_TORCH_CHECK((g_idx.size(-1) == 0 && perm.size(-1) == 0) || + (g_idx.size(-1) == size_k && perm.size(-1) == size_k), + "Unexpected g_idx.size(-1) = ", g_idx.size(-1), + " and perm.size(-1) = ", perm.size(-1), + ", where size_k = ", size_k); } else { - g_idx = torch::empty({0}, options); - perm = torch::empty({0}, options); - a_tmp = torch::empty({0}, options); + g_idx = torch::stable::empty({0}, c_scalar_type, std::nullopt, device); + perm = torch::stable::empty({0}, c_scalar_type, std::nullopt, device); + a_tmp = torch::stable::empty({0}, c_scalar_type, std::nullopt, device); } bool has_act_order = g_idx.size(-1) > 0 && perm.size(-1) > 0; if (has_act_order) { - a_tmp = torch::empty({size_m, size_k}, options); + a_tmp = torch::stable::empty({size_m, size_k}, c_scalar_type, std::nullopt, + device); if (is_k_full) { - TORCH_CHECK(num_groups > 1, "For act_order, num_groups must be > 1"); - TORCH_CHECK(size_k % num_groups == 0, "size_k = ", size_k, - ", is not divisible by num_groups = ", num_groups); + STD_TORCH_CHECK(num_groups > 1, "For act_order, num_groups must be > 1"); + STD_TORCH_CHECK(size_k % num_groups == 0, "size_k = ", size_k, + ", is not divisible by num_groups = ", num_groups); group_size = size_k / num_groups; } else { group_size = 0; } } else { - a_tmp = torch::empty({0}, options); + a_tmp = torch::stable::empty({0}, c_scalar_type, std::nullopt, device); if (num_groups > 1) { - TORCH_CHECK( + STD_TORCH_CHECK( size_k % num_groups == 0, "size_k = ", size_k, ", is not divisible by b_scales.size(0) = ", b_scales.size(0)); group_size = size_k / num_groups; @@ -750,109 +780,114 @@ torch::Tensor marlin_gemm( } } - torch::Tensor global_scale; + torch::stable::Tensor global_scale; if (global_scale_or_none.has_value()) { global_scale = global_scale_or_none.value(); - TORCH_CHECK(b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn, - "global_scale can only be used for nvfp4 format."); + STD_TORCH_CHECK(b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn, + "global_scale can only be used for nvfp4 format."); } else { - global_scale = torch::empty({0}, options_fp32); - TORCH_CHECK(!(b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn), - "the global_scale parameter must be passed for nvfp4 format."); + global_scale = torch::stable::empty( + {0}, torch::headeronly::ScalarType::Float, std::nullopt, device); + STD_TORCH_CHECK( + !(b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn), + "the global_scale parameter must be passed for nvfp4 format."); } bool has_bias = b_bias_or_none.has_value(); - torch::Tensor b_bias; + torch::stable::Tensor b_bias; if (has_bias) { b_bias = b_bias_or_none.value(); - TORCH_CHECK(b_bias.device().is_cuda(), "b_bias is not on GPU"); - TORCH_CHECK(b_bias.is_contiguous(), "b_bias is not contiguous"); - TORCH_CHECK(b_bias.size(0) == size_n, "b_bias.size(0) != size_n"); - TORCH_CHECK(b_bias.stride(0) == 1, "b_bias.stride(0) != 1"); + STD_TORCH_CHECK(b_bias.device().is_cuda(), "b_bias is not on GPU"); + STD_TORCH_CHECK(b_bias.is_contiguous(), "b_bias is not contiguous"); + STD_TORCH_CHECK(b_bias.size(0) == size_n, "b_bias.size(0) != size_n"); + STD_TORCH_CHECK(b_bias.stride(0) == 1, "b_bias.stride(0) != 1"); } else { - b_bias = torch::empty({0}, options); + b_bias = torch::stable::empty({0}, c_scalar_type, std::nullopt, device); } - torch::Tensor b_zeros; + torch::stable::Tensor b_zeros; if (b_zeros_or_none.has_value()) { b_zeros = b_zeros_or_none.value(); - TORCH_CHECK(b_zeros.device().is_cuda(), "b_zeros is not on GPU"); - TORCH_CHECK(b_zeros.is_contiguous(), "b_zeros is not contiguous"); + STD_TORCH_CHECK(b_zeros.device().is_cuda(), "b_zeros is not on GPU"); + STD_TORCH_CHECK(b_zeros.is_contiguous(), "b_zeros is not contiguous"); } else { - b_zeros = torch::empty({0}, options); + b_zeros = torch::stable::empty({0}, c_scalar_type, std::nullopt, device); } bool has_zp = b_zeros.size(-1) > 0; if (has_zp) { - TORCH_CHECK( + STD_TORCH_CHECK( b_type == vllm::kU4 || b_type == vllm::kU8, "b_type must be u4 or u8 when has_zp = True. Got = ", b_type.str()); } else { - TORCH_CHECK(b_type == vllm::kU4B8 || b_type == vllm::kU8B128 || - b_type == vllm::kS4 || b_type == vllm::kS8 || - b_type == vllm::kFE4M3fn || b_type == vllm::kFE2M1f, - "b_type must be uint4b8, uint8b128, int4, int8, " - "float8_e4m3fn or float4_e2m1f when has_zp = False. Got = ", - b_type.str()); + STD_TORCH_CHECK(b_type == vllm::kU4B8 || b_type == vllm::kU8B128 || + b_type == vllm::kS4 || b_type == vllm::kS8 || + b_type == vllm::kFE4M3fn || b_type == vllm::kFE2M1f, + "b_type must be uint4b8, uint8b128, int4, int8, " + "float8_e4m3fn or float4_e2m1f when has_zp = False. Got = ", + b_type.str()); } if (has_zp && is_zp_float) { - TORCH_CHECK(a.scalar_type() == at::ScalarType::Half, - "Computation type must be float16 (half) when using float zero " - "points."); + STD_TORCH_CHECK( + a.scalar_type() == torch::headeronly::ScalarType::Half, + "Computation type must be float16 (half) when using float zero " + "points."); } // Verify b_zeros if (has_zp) { - int rank = b_zeros.sizes().size(); - TORCH_CHECK(rank == 2, "b_zeros rank = ", rank, " is not 2"); + int rank = b_zeros.dim(); + STD_TORCH_CHECK(rank == 2, "b_zeros rank = ", rank, " is not 2"); if (is_zp_float) { - TORCH_CHECK(b_zeros.size(1) == size_n, - "b_zeros dim 1 = ", b_zeros.size(1), - " is not size_n = ", size_n); - TORCH_CHECK(num_groups == b_zeros.size(0), - "b_zeros dim 0 = ", b_zeros.size(0), - " is not num_groups = ", num_groups); - TORCH_CHECK(num_groups != -1, "num_groups must be != -1"); + STD_TORCH_CHECK(b_zeros.size(1) == size_n, + "b_zeros dim 1 = ", b_zeros.size(1), + " is not size_n = ", size_n); + STD_TORCH_CHECK(num_groups == b_zeros.size(0), + "b_zeros dim 0 = ", b_zeros.size(0), + " is not num_groups = ", num_groups); + STD_TORCH_CHECK(num_groups != -1, "num_groups must be != -1"); } else { - TORCH_CHECK(b_zeros.size(0) == num_groups, - "b_zeros dim 0 = ", b_zeros.size(0), - " is not num_groups = ", num_groups); - TORCH_CHECK(b_zeros.size(1) == size_n / pack_factor, - "b_zeros dim 1 = ", b_zeros.size(1), - " is not size_n / pack_factor = ", size_n / pack_factor); + STD_TORCH_CHECK(b_zeros.size(0) == num_groups, + "b_zeros dim 0 = ", b_zeros.size(0), + " is not num_groups = ", num_groups); + STD_TORCH_CHECK(b_zeros.size(1) == size_n / pack_factor, + "b_zeros dim 1 = ", b_zeros.size(1), + " is not size_n / pack_factor = ", size_n / pack_factor); } } // Verify workspace size - TORCH_CHECK(size_n % MARLIN_NAMESPACE_NAME::min_thread_n == 0, - "size_n = ", size_n, ", is not divisible by min_thread_n = ", - MARLIN_NAMESPACE_NAME::min_thread_n); + STD_TORCH_CHECK(size_n % MARLIN_NAMESPACE_NAME::min_thread_n == 0, + "size_n = ", size_n, ", is not divisible by min_thread_n = ", + MARLIN_NAMESPACE_NAME::min_thread_n); int min_workspace_size = sms; - TORCH_CHECK(workspace.numel() >= min_workspace_size, - "workspace.numel = ", workspace.numel(), - " is below min_workspace_size = ", min_workspace_size); + STD_TORCH_CHECK(workspace.numel() >= min_workspace_size, + "workspace.numel = ", workspace.numel(), + " is below min_workspace_size = ", min_workspace_size); - int dev = a.get_device(); - - TORCH_CHECK(a_scales.scalar_type() == at::ScalarType::Float, - "scalar type of a_scales must be float"); - TORCH_CHECK(global_scale.scalar_type() == at::ScalarType::Float, - "scalar type of global_scale must be float"); + STD_TORCH_CHECK( + a_scales.scalar_type() == torch::headeronly::ScalarType::Float, + "scalar type of a_scales must be float"); + STD_TORCH_CHECK( + global_scale.scalar_type() == torch::headeronly::ScalarType::Float, + "scalar type of global_scale must be float"); if (a_type.size_bits() == 16) { - TORCH_CHECK( + STD_TORCH_CHECK( a.scalar_type() == c.scalar_type(), "scalar type of a must be the same with c for 16 bit activation"); } marlin::marlin_mm( - a.data_ptr(), b_q_weight.data_ptr(), c.data_ptr(), c_tmp.data_ptr(), - b_bias.data_ptr(), a_scales.data_ptr(), b_scales.data_ptr(), - global_scale.data_ptr(), b_zeros.data_ptr(), g_idx.data_ptr(), - perm.data_ptr(), a_tmp.data_ptr(), size_m, size_n, size_k, a.stride(0), - workspace.data_ptr(), a_type, b_type, c_type, s_type, has_bias, - has_act_order, is_k_full, has_zp, num_groups, group_size, dev, - at::cuda::getCurrentCUDAStream(dev), thread_k, thread_n, sms, + a.const_data_ptr(), b_q_weight.const_data_ptr(), c.mutable_data_ptr(), + c_tmp.mutable_data_ptr(), b_bias.mutable_data_ptr(), + a_scales.mutable_data_ptr(), b_scales.mutable_data_ptr(), + global_scale.mutable_data_ptr(), b_zeros.mutable_data_ptr(), + g_idx.mutable_data_ptr(), perm.mutable_data_ptr(), + a_tmp.mutable_data_ptr(), size_m, size_n, size_k, a.stride(0), + workspace.mutable_data_ptr(), a_type, b_type, c_type, s_type, has_bias, + has_act_order, is_k_full, has_zp, num_groups, group_size, device_index, + get_current_cuda_stream(device_index), thread_k, thread_n, sms, use_atomic_add, use_fp32_reduce, is_zp_float); return c; @@ -860,6 +895,6 @@ torch::Tensor marlin_gemm( #endif -TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) { - m.impl("marlin_gemm", &marlin_gemm); +STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) { + m.impl("marlin_gemm", TORCH_BOX(&marlin_gemm)); } diff --git a/csrc/quantization/marlin/marlin.cuh b/csrc/libtorch_stable/quantization/marlin/marlin.cuh similarity index 93% rename from csrc/quantization/marlin/marlin.cuh rename to csrc/libtorch_stable/quantization/marlin/marlin.cuh index d3a91568349..bfb65e874b3 100644 --- a/csrc/quantization/marlin/marlin.cuh +++ b/csrc/libtorch_stable/quantization/marlin/marlin.cuh @@ -2,14 +2,6 @@ #ifndef _marlin_cuh #define _marlin_cuh - // These torch headers are only needed by non-stable callers (e.g. ops.cu). - // Guard them so that stable ABI targets can still include marlin.cuh - // for Vec, constants, and cp_async helpers without pulling in torch/all.h. - #ifndef TORCH_TARGET_VERSION - #include - #include - #include - #endif #include #include #include diff --git a/csrc/quantization/marlin/marlin_dtypes.cuh b/csrc/libtorch_stable/quantization/marlin/marlin_dtypes.cuh similarity index 100% rename from csrc/quantization/marlin/marlin_dtypes.cuh rename to csrc/libtorch_stable/quantization/marlin/marlin_dtypes.cuh diff --git a/csrc/libtorch_stable/quantization/marlin/marlin_int4_fp8_preprocess.cu b/csrc/libtorch_stable/quantization/marlin/marlin_int4_fp8_preprocess.cu new file mode 100644 index 00000000000..f8ef6b12a01 --- /dev/null +++ b/csrc/libtorch_stable/quantization/marlin/marlin_int4_fp8_preprocess.cu @@ -0,0 +1,118 @@ + +#include "marlin.cuh" + +#include +#include +#include +#include +#include +#include + +#include "libtorch_stable/torch_utils.h" + +// for only non-zp format (like gptq) +__global__ void marlin_int4_fp8_preprocess_kernel_without_zp( + // qweight: (size_k * size_n // 8,) + const int32_t* __restrict__ qweight, + // output: same shape with qweight + int32_t* __restrict__ output) { + int32_t val = qweight[blockIdx.x * 32 + threadIdx.x]; + int32_t new_val = 0; + +#pragma unroll + for (int32_t i = 0; i < 8; i++) { + int32_t single_val = val & 0xF; + single_val = single_val >= 8 ? single_val - 8 : 15 - single_val; + new_val |= single_val << (i * 4); + val >>= 4; + } + + output[blockIdx.x * 32 + threadIdx.x] = new_val; +} + +// for awq format only (with zp and with awq weight layout) +__global__ void marlin_int4_fp8_preprocess_kernel_awq( + // AWQ qweight: (size_k, size_n // 8) + const int32_t* __restrict__ qweight, + // output: same shape with qweight + int32_t* __restrict__ output, + // AWQ zeros: (size_k // group_size, size_n // 8) + const int32_t* __restrict__ qzeros, int32_t size_n, int32_t size_k, + int32_t group_size) { + int32_t val = + qweight[(blockIdx.x * 32 + threadIdx.x) * size_n / 8 + blockIdx.y]; + int32_t zero = + qzeros[(blockIdx.x * 32 + threadIdx.x) / group_size * size_n / 8 + + blockIdx.y]; + int32_t new_val = 0; + +#pragma unroll + for (int32_t i = 0; i < 8; i++) { + int32_t single_val = val & 0xF; + int32_t single_zero = zero & 0xF; + + single_val = + single_val >= single_zero ? single_val - single_zero : 15 - single_val; + new_val |= single_val << (i * 4); + val >>= 4; + zero >>= 4; + } + + output[(blockIdx.x * 32 + threadIdx.x) * size_n / 8 + blockIdx.y] = new_val; +} + +torch::stable::Tensor marlin_int4_fp8_preprocess( + torch::stable::Tensor& qweight, + std::optional qzeros_or_none, bool inplace) { + STD_TORCH_CHECK(qweight.is_cuda(), "qweight is not on GPU"); + STD_TORCH_CHECK(qweight.scalar_type() == torch::headeronly::ScalarType::Int, + "qweight.dtype != torch.int32"); + + const int32_t device_index = qweight.get_device_index(); + torch::stable::accelerator::DeviceGuard device_guard(device_index); + const cudaStream_t stream = get_current_cuda_stream(device_index); + + torch::stable::Tensor output = + inplace ? qweight : torch::stable::empty_like(qweight); + + if (!qzeros_or_none.has_value()) { + STD_TORCH_CHECK(qweight.numel() * 8 % 256 == 0, + "qweight.numel() * 8 % 256 != 0"); + + int blocks = qweight.numel() * 8 / 256; + marlin_int4_fp8_preprocess_kernel_without_zp<<>>( + reinterpret_cast(qweight.const_data_ptr()), + reinterpret_cast(output.mutable_data_ptr())); + } else { + int32_t size_k = qweight.size(0); + int32_t size_n = qweight.size(1) * 8; + torch::stable::Tensor qzeros = qzeros_or_none.value(); + + STD_TORCH_CHECK(size_k % 32 == 0, "size_k % 32 != 0"); + STD_TORCH_CHECK(qzeros.is_cuda(), "qzeros is not on GPU"); + STD_TORCH_CHECK(qzeros.scalar_type() == torch::headeronly::ScalarType::Int, + "qweight.dtype != torch.int32"); + STD_TORCH_CHECK(qzeros.get_device_index() == device_index, + "qzeros is not on the same device with qweight"); + + int32_t group_size = qweight.size(0) / qzeros.size(0); + STD_TORCH_CHECK(qweight.size(1) == qzeros.size(1), + "qweight.size(1) != qzeros.size(1)"); + STD_TORCH_CHECK(qweight.size(0) % qzeros.size(0) == 0, + "qweight.size(0) % qzeros.size(0) != 0"); + STD_TORCH_CHECK(group_size % 8 == 0, "group_size % 8 != 0"); + + dim3 blocks(size_k / 32, size_n / 8); + marlin_int4_fp8_preprocess_kernel_awq<<>>( + reinterpret_cast(qweight.const_data_ptr()), + reinterpret_cast(output.mutable_data_ptr()), + reinterpret_cast(qzeros.const_data_ptr()), size_n, + size_k, group_size); + } + + return output; +} + +STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) { + m.impl("marlin_int4_fp8_preprocess", TORCH_BOX(&marlin_int4_fp8_preprocess)); +} diff --git a/csrc/quantization/marlin/marlin_mma.h b/csrc/libtorch_stable/quantization/marlin/marlin_mma.h similarity index 100% rename from csrc/quantization/marlin/marlin_mma.h rename to csrc/libtorch_stable/quantization/marlin/marlin_mma.h diff --git a/csrc/quantization/marlin/marlin_template.h b/csrc/libtorch_stable/quantization/marlin/marlin_template.h similarity index 100% rename from csrc/quantization/marlin/marlin_template.h rename to csrc/libtorch_stable/quantization/marlin/marlin_template.h diff --git a/csrc/libtorch_stable/torch_bindings.cpp b/csrc/libtorch_stable/torch_bindings.cpp index 816f2665048..204feed4a25 100644 --- a/csrc/libtorch_stable/torch_bindings.cpp +++ b/csrc/libtorch_stable/torch_bindings.cpp @@ -33,6 +33,35 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { // TODO: Remove this once ROCm upgrade to torch 2.11. ops.def("get_cuda_view_from_cpu_tensor(Tensor cpu_tensor) -> Tensor"); + + // Marlin GEMM + ops.def( + "marlin_gemm(Tensor a, Tensor? c_or_none, Tensor b_q_weight, " + "Tensor? b_bias_or_none,Tensor b_scales, " + "Tensor? a_scales, Tensor? global_scale, Tensor? b_zeros_or_none, " + "Tensor? " + "g_idx_or_none, Tensor? perm_or_none, Tensor workspace, int b_type_id, " + "SymInt size_m, SymInt size_n, SymInt size_k, bool is_k_full, " + "bool use_atomic_add, bool use_fp32_reduce, bool is_zp_float) -> Tensor"); + // conditionally compiled so impl registrations are in source file + + // gptq_marlin repack from GPTQ. + ops.def( + "gptq_marlin_repack(Tensor b_q_weight, Tensor perm, " + "SymInt size_k, SymInt size_n, int num_bits, bool is_a_8bit) -> Tensor"); + // conditionally compiled so impl registrations are in source file + + // awq_marlin repack from AWQ. + ops.def( + "awq_marlin_repack(Tensor b_q_weight, SymInt size_k, " + "SymInt size_n, int num_bits, bool is_a_8bit) -> Tensor"); + // conditionally compiled so impl registrations are in source file + + // preprocess W-int4A-fp8 weight for marlin kernel + ops.def( + "marlin_int4_fp8_preprocess(Tensor qweight, " + "Tensor? qzeros_or_none, bool inplace) -> Tensor"); + // conditionally compiled so impl registrations are in source file #endif #ifndef USE_ROCM diff --git a/csrc/quantization/marlin/marlin_int4_fp8_preprocess.cu b/csrc/quantization/marlin/marlin_int4_fp8_preprocess.cu deleted file mode 100644 index 7d4c97fb57e..00000000000 --- a/csrc/quantization/marlin/marlin_int4_fp8_preprocess.cu +++ /dev/null @@ -1,106 +0,0 @@ - - -#include "marlin.cuh" - -#include "core/registration.h" - -// for only non-zp format (like gptq) -__global__ void marlin_int4_fp8_preprocess_kernel_without_zp( - // qweight: (size_k * size_n // 8,) - const int32_t* __restrict__ qweight, - // output: same shape with qweight - int32_t* __restrict__ output) { - int32_t val = qweight[blockIdx.x * 32 + threadIdx.x]; - int32_t new_val = 0; - -#pragma unroll - for (int32_t i = 0; i < 8; i++) { - int32_t single_val = val & 0xF; - single_val = single_val >= 8 ? single_val - 8 : 15 - single_val; - new_val |= single_val << (i * 4); - val >>= 4; - } - - output[blockIdx.x * 32 + threadIdx.x] = new_val; -} - -// for awq format only (with zp and with awq weight layout) -__global__ void marlin_int4_fp8_preprocess_kernel_awq( - // AWQ qweight: (size_k, size_n // 8) - const int32_t* __restrict__ qweight, - // output: same shape with qweight - int32_t* __restrict__ output, - // AWQ zeros: (size_k // group_size, size_n // 8) - const int32_t* __restrict__ qzeros, int32_t size_n, int32_t size_k, - int32_t group_size) { - int32_t val = - qweight[(blockIdx.x * 32 + threadIdx.x) * size_n / 8 + blockIdx.y]; - int32_t zero = - qzeros[(blockIdx.x * 32 + threadIdx.x) / group_size * size_n / 8 + - blockIdx.y]; - int32_t new_val = 0; - -#pragma unroll - for (int32_t i = 0; i < 8; i++) { - int32_t single_val = val & 0xF; - int32_t single_zero = zero & 0xF; - - single_val = - single_val >= single_zero ? single_val - single_zero : 15 - single_val; - new_val |= single_val << (i * 4); - val >>= 4; - zero >>= 4; - } - - output[(blockIdx.x * 32 + threadIdx.x) * size_n / 8 + blockIdx.y] = new_val; -} - -torch::Tensor marlin_int4_fp8_preprocess( - torch::Tensor& qweight, std::optional qzeros_or_none, - bool inplace) { - TORCH_CHECK(qweight.device().is_cuda(), "qweight is not on GPU"); - TORCH_CHECK(qweight.scalar_type() == at::ScalarType::Int, - "qweight.dtype != torch.int32"); - - const at::cuda::OptionalCUDAGuard device_guard(device_of(qweight)); - - torch::Tensor output = inplace ? qweight : torch::empty_like(qweight); - - if (!qzeros_or_none.has_value()) { - TORCH_CHECK(qweight.numel() * 8 % 256 == 0, - "qweight.numel() * 8 % 256 != 0"); - - int blocks = qweight.numel() * 8 / 256; - marlin_int4_fp8_preprocess_kernel_without_zp<<>>( - (const int32_t*)qweight.data_ptr(), (int32_t*)output.data_ptr()); - } else { - int32_t size_k = qweight.size(0); - int32_t size_n = qweight.size(1) * 8; - torch::Tensor qzeros = qzeros_or_none.value(); - - TORCH_CHECK(size_k % 32 == 0, "size_k % 32 != 0"); - TORCH_CHECK(qzeros.device().is_cuda(), "qzeros is not on GPU"); - TORCH_CHECK(qzeros.scalar_type() == at::ScalarType::Int, - "qweight.dtype != torch.int32"); - TORCH_CHECK(device_of(qweight) == device_of(qzeros), - "qzeros is not on the same device with qweight"); - - int32_t group_size = qweight.size(0) / qzeros.size(0); - TORCH_CHECK(qweight.size(1) == qzeros.size(1), - "qweight.size(1) != qzeros.size(1)"); - TORCH_CHECK(qweight.size(0) % qzeros.size(0) == 0, - "qweight.size(0) % qzeros.size(0) != 0"); - TORCH_CHECK(group_size % 8 == 0, "group_size % 8 != 0"); - - dim3 blocks(size_k / 32, size_n / 8); - marlin_int4_fp8_preprocess_kernel_awq<<>>( - (const int32_t*)qweight.data_ptr(), (int32_t*)output.data_ptr(), - (const int32_t*)qzeros.data_ptr(), size_n, size_k, group_size); - } - - return output; -} - -TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) { - m.impl("marlin_int4_fp8_preprocess", &marlin_int4_fp8_preprocess); -} diff --git a/csrc/torch_bindings.cpp b/csrc/torch_bindings.cpp index 58524c4c5db..941e4a61c1a 100644 --- a/csrc/torch_bindings.cpp +++ b/csrc/torch_bindings.cpp @@ -101,35 +101,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { ") -> Tensor"); // conditionally compiled so impl registration is in source file - // Marlin Optimized Quantized GEMM (supports GPTQ, AWQ, FP8, NVFP4, MXFP4). - ops.def( - "marlin_gemm(Tensor a, Tensor? c_or_none, Tensor b_q_weight, " - "Tensor? b_bias_or_none,Tensor b_scales, " - "Tensor? a_scales, Tensor? global_scale, Tensor? b_zeros_or_none, " - "Tensor? " - "g_idx_or_none, Tensor? perm_or_none, Tensor workspace, int b_type_id, " - "SymInt size_m, SymInt size_n, SymInt size_k, bool is_k_full, " - "bool use_atomic_add, bool use_fp32_reduce, bool is_zp_float) -> Tensor"); - // conditionally compiled so impl registration is in source file - - // gptq_marlin repack from GPTQ. - ops.def( - "gptq_marlin_repack(Tensor b_q_weight, Tensor perm, " - "SymInt size_k, SymInt size_n, int num_bits, bool is_a_8bit) -> Tensor"); - // conditionally compiled so impl registrations are in source file - - // awq_marlin repack from AWQ. - ops.def( - "awq_marlin_repack(Tensor b_q_weight, SymInt size_k, " - "SymInt size_n, int num_bits, bool is_a_8bit) -> Tensor"); - // conditionally compiled so impl registrations are in source file - - // preprocess W-int4A-fp8 weight for marlin kernel - ops.def( - "marlin_int4_fp8_preprocess(Tensor qweight, " - "Tensor? qzeros_or_none, bool inplace) -> Tensor"); - // conditionally compiled so impl registrations are in source file - #endif } From 6fbfdd183145443274df49c09c46cd13ea27af5f Mon Sep 17 00:00:00 2001 From: Dao007forever Date: Thu, 11 Jun 2026 21:42:41 -0700 Subject: [PATCH 27/52] [NIXL] Per-region KV transfer classification for mixed full-attn + MLA groups (#44583) --- .../kv_connector/unit/test_nixl_connector.py | 81 +++++++ tests/v1/kv_connector/unit/test_tp_mapping.py | 12 +- .../kv_connector/v1/nixl/worker.py | 209 ++++++++++++------ 3 files changed, 234 insertions(+), 68 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_nixl_connector.py b/tests/v1/kv_connector/unit/test_nixl_connector.py index a2a46684bb7..c5784d1c200 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector.py @@ -1063,6 +1063,87 @@ class TestNixlHandshake: # whole block is moved. worker.add_remote_agent(meta, remote_tp_size=1) + @patch( + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper", + FakeNixlWrapper, + ) + def test_handshake_mixed_fa_mla_hetero_tp(self, default_vllm_config, dist_init): + """Mixed full-attn (SPLIT) + MLA (REPLICATE) single KV group under + heterogeneous TP must NOT raise (previously a NotImplementedError), + and the per-region gate must still reject a wrong block_len. + """ + vllm_config = create_vllm_config() + with patch( + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.get_tensor_model_parallel_world_size", # noqa: E501 + return_value=2, + ): + connector = NixlConnector( + vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16) + ) + connector.connector_worker = FakeNixlConnectorWorker( + vllm_config, connector.engine_id, hand_shake_latency=0 + ) + worker = connector.connector_worker + + # Region 0: full-attn (SPLIT). Region 1: MLA (REPLICATE). + fa_len = 4096 * worker.block_size + idx_len = 512 * worker.block_size + worker.slot_size_per_layer = [4096, 512] + worker.block_len_per_layer = [fa_len, idx_len] + worker._region_is_mla = [False, True] + worker.num_blocks = 1 + worker.dst_num_blocks[worker.engine_id] = worker.num_blocks + worker.src_blocks_data = [ + (0, fa_len, worker.tp_rank), + (0, idx_len, worker.tp_rank), + ] + worker.num_descs = len(worker.src_blocks_data) + + # D_TP=2, P_TP=1 -> tp_ratio=2. SPLIT region scales by tp_ratio; + # REPLICATE region is unchanged. + tp_ratio = 2 + meta = NixlAgentMetadata( + engine_id=FakeNixlConnectorWorker.REMOTE_ENGINE_ID, + agent_metadata=FakeNixlWrapper.AGENT_METADATA, + kv_caches_base_addr=[0, 0], + device_id=0, + num_blocks=1, + block_lens=[fa_len * tp_ratio, idx_len], + kv_cache_layout=worker.kv_cache_layout, + block_size=worker.block_size, + ssm_sizes=(0, 0), + attn_backend_name=worker.backend_name, + physical_blocks_per_logical_kv_block=1, + ) + worker.add_remote_agent(meta, remote_tp_size=1) + assert ( + FakeNixlConnectorWorker.REMOTE_ENGINE_ID in worker.dst_xfer_side_handles + ) + # Gate rejects an MLA region wrongly scaled by tp_ratio. + worker2 = FakeNixlConnectorWorker( + vllm_config, connector.engine_id, hand_shake_latency=0 + ) + worker2.block_len_per_layer = [fa_len, idx_len] + worker2._region_is_mla = [False, True] + worker2.num_blocks = 1 + worker2.dst_num_blocks[worker2.engine_id] = worker2.num_blocks + bad_meta = NixlAgentMetadata( + engine_id=FakeNixlConnectorWorker.REMOTE_ENGINE_ID, + agent_metadata=FakeNixlWrapper.AGENT_METADATA, + kv_caches_base_addr=[0, 0], + device_id=0, + num_blocks=1, + # WRONG: MLA region scaled by tp_ratio (it should be replicated). + block_lens=[fa_len * tp_ratio, idx_len * tp_ratio], + kv_cache_layout=worker2.kv_cache_layout, + block_size=worker2.block_size, + ssm_sizes=(0, 0), + attn_backend_name=worker2.backend_name, + physical_blocks_per_logical_kv_block=1, + ) + with pytest.raises(AssertionError): + worker2.add_remote_agent(bad_meta, remote_tp_size=1) + # NOTE: resource cleanup in mp backend is a bit finicky, so the order in which # we put here is important. First run ray, it will clean up the resources, then diff --git a/tests/v1/kv_connector/unit/test_tp_mapping.py b/tests/v1/kv_connector/unit/test_tp_mapping.py index 95d49faf042..5ab6b68400c 100644 --- a/tests/v1/kv_connector/unit/test_tp_mapping.py +++ b/tests/v1/kv_connector/unit/test_tp_mapping.py @@ -73,9 +73,19 @@ class TestTPMappingStructure: def _make_mock_worker_for_splits(group_spec_types): - """Build a mock NixlConnectorWorker with _group_spec_types for split tests.""" + """Build a mock NixlConnectorWorker with _group_spec_types for split tests. + + No per-region replicate flags are configured (``block_len_per_layer`` empty + and ``num_regions == 0``), so ``_fa_desc_replicated`` takes its early-return + path and treats every FA descriptor as SPLIT, matching the legacy behavior + these tests assert. + """ worker = object.__new__(NixlConnectorWorker) worker._group_spec_types = group_spec_types + worker.transfer_topo = SimpleNamespace(virtually_split_kv_in_blocks=False) + worker.block_len_per_layer = [] + worker.num_regions = 0 + worker._region_is_mla = [] return worker diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/worker.py index e4b20c01f4d..213a3b03144 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/worker.py @@ -71,6 +71,7 @@ from vllm.v1.attention.backends.utils import get_kv_cache_layout from vllm.v1.kv_cache_interface import ( FullAttentionSpec, MambaSpec, + MLAAttentionSpec, UniformTypeKVCacheSpecs, ) from vllm.v1.worker.block_table import BlockTable @@ -178,19 +179,63 @@ class NixlConnectorWorker: else 0 ) + # Per-FA-descriptor replicate flag, in _build_fa_local emission order. + fa_desc_replicated = self._fa_desc_replicated(num_fa_descs) + for p_idx, p_rank in enumerate(plan.all_source_ranks): fa_slot = plan.rank_to_attention_slot.get(p_rank, 0) handle: list[tuple[int, int, int]] = [] for j, (addr, local_len, dev) in enumerate(src_blocks_data): if j < num_fa_descs: - chunk = local_len // fa_num_splits - handle.append((addr + fa_slot * chunk, chunk, dev)) + if fa_desc_replicated[j]: + # REPLICATE (MLA): whole block written on every rank. + handle.append((addr, local_len, dev)) + else: + # SPLIT (full-attn): this rank's head slice. + chunk = local_len // fa_num_splits + handle.append((addr + fa_slot * chunk, chunk, dev)) else: chunk = local_len // ssm_num_splits handle.append((addr + p_idx * chunk, chunk, dev)) yield handle + def _fa_desc_replicated(self, num_fa_descs: int) -> list[bool]: + """Per-FA-descriptor replicate flag, in _build_fa_local emission order + (region-major; K then optional V per region). Length ``num_fa_descs``. + """ + assert self.transfer_topo is not None + n_regions = len(self.block_len_per_layer) + # Unset only when the worker is built directly in unit tests; a real + # model always registers regions (no-KV-cache crashes long before here). + # Fall back to all-SPLIT to preserve the pre-per-region behavior. + if n_regions == 0 or self.num_regions == 0: + return [False] * num_fa_descs + # Descriptors (blocks) per stream; all streams share the same count. + nblk = num_fa_descs // self.num_regions + virtually_split = self.transfer_topo.virtually_split_kv_in_blocks + flags: list[bool] = [] + for i in range(n_regions): + replicated = self._is_region_replicated(i) + # REPLICATE (MLA) is key-only -> 1 stream; SPLIT emits K and V + # (2 streams) under the virtually-split layout. + num_streams = 1 if replicated or not virtually_split else 2 + flags.extend([replicated] * (num_streams * nblk)) + assert len(flags) == num_fa_descs, ( + f"FA desc flags {len(flags)} != num_fa_descs {num_fa_descs}" + ) + return flags + + def _is_region_replicated(self, region_idx: int) -> bool: + """Whether region ``region_idx`` is transferred REPLICATE vs SPLIT. + + REPLICATE (MLA): identical on every rank, whole block read from one + rank at offset 0, key-only. SPLIT (full-attn): head-sharded across TP. + Defaults to SPLIT when the per-region map is unset (e.g. tests that set + block_len_per_layer without register_kv_caches). + """ + return region_idx < len(self._region_is_mla) and self._region_is_mla[region_idx] + def __init__( self, vllm_config: "VllmConfig", @@ -450,6 +495,15 @@ class NixlConnectorWorker: for g in self.kv_cache_config.kv_cache_groups ) + # Per-region MLA flag, 1:1 with block_len_per_layer. True -> REPLICATE + # (MLA), False -> SPLIT (head-sharded full-attn). Mixed only for models + # combining both (e.g. GQA main + MLA Eagle-3 draft). + self._region_is_mla = list[bool]() + + # Enable different block lengths for different layers *only* when MLA is used. + # This is not used for SSM layers, which use the counterpart `mamba_ssm_size`. + self.block_len_per_layer = list[int]() + # Per-engine TP mappings. Generated during handshake. self.tp_mappings: dict[EngineId, TPMapping] = {} @@ -849,9 +903,6 @@ class NixlConnectorWorker: # to better exploit the memory layout (ie num_blocks is the first dim). tensor_size_bytes = None - # Enable different block lengths for different layers *only* when MLA is used. - # This is not used for SSM layers, which use the counterpart `mamba_ssm_size`. - self.block_len_per_layer = list[int]() for layer_name, cache_or_caches in xfer_buffers.items(): # NOTE (NickLucche) Hybrid SSM models assume a layout that is similar to # that of FI, with block laid out as in `get_backend_aware_kv_block_len`. @@ -895,8 +946,6 @@ class NixlConnectorWorker: # `page_size` accounts for physical blocks, st KVCache is always # [`num_blocks` * `page_size`] curr_tensor_size_bytes = num_blocks * physical_page_size - if tensor_size_bytes is None: - tensor_size_bytes = curr_tensor_size_bytes # TODO (NickLucche) we could eventually unify how we handle FA/FI regions, # registering a single tensor for both K/V and splitting logically like FI. @@ -920,6 +969,20 @@ class NixlConnectorWorker: ) else: self.block_len_per_layer.append(physical_page_size) + is_mla_region = isinstance(layer_spec, MLAAttentionSpec) + self._region_is_mla.append(is_mla_region) + + # HeteroTP cannot transfer differently-sized regions, so every + # non-MLA region in a group must share one tensor size (this also + # holds for Mamba-like models). The sole exception is the DeepSeek + # MLA indexer, which sits in a UniformTypeKVCacheSpecs group at a + # different size; MLA regions are therefore exempt. + if not is_mla_region: + if tensor_size_bytes is None: + tensor_size_bytes = curr_tensor_size_bytes + assert tensor_size_bytes == curr_tensor_size_bytes, ( + "All non-MLA kv cache tensors must have the same size" + ) if cache.shape[0] != num_blocks: raise AssertionError( @@ -937,12 +1000,6 @@ class NixlConnectorWorker: f"{self.transfer_topo.is_kv_layout_blocks_first}" ) - if not self.use_mla: - # Different kv cache shape is not supported by HeteroTP. - # This must also hold true for Mamba-like models. - assert tensor_size_bytes == curr_tensor_size_bytes, ( - "All kv cache tensors must have the same size" - ) # Need to make sure the device ID is non-negative for NIXL, # Torch uses -1 to indicate CPU tensors. self.device_id = max(cache.get_device(), 0) @@ -953,7 +1010,11 @@ class NixlConnectorWorker: logger.debug( "Different block lengths collected: %s", set(self.block_len_per_layer) ) - assert len(self.block_len_per_layer) == len(seen_base_addresses) + assert ( + len(self.block_len_per_layer) + == len(seen_base_addresses) + == len(self._region_is_mla) + ) self.kv_caches_base_addr[self.engine_id][self.tp_rank] = seen_base_addresses self.num_regions = len(caches_data) @@ -967,7 +1028,12 @@ class NixlConnectorWorker: # of 'virtual' regions here and halve `block_len` below. # Similarly for Mamba layers, we register SSM+Conv as a single region and # then duplicate it logically to be able to index SSM/Conv separately. - self.num_regions *= 2 + # Exception: key-only REPLICATE regions (MLA) have no V half, so + # they contribute a single desc stream and are not doubled. + self.num_regions = sum( + 1 if self._is_region_replicated(i) else 2 + for i in range(len(self._region_is_mla)) + ) # Total local FA descriptors (boundary between FA and mamba descs). self.num_descs = self.num_regions * self.num_blocks @@ -1133,10 +1199,13 @@ class NixlConnectorWorker: addr = base_addr + block_offset result.append((addr, kv_block_len, self.device_id)) - if self.transfer_topo.virtually_split_kv_in_blocks: + if ( + self.transfer_topo.virtually_split_kv_in_blocks + and not self._is_region_replicated(i) + ): # Separate and interleave K/V regions to maintain the same # descs ordering. This is needed for selecting contiguous heads - # when split across TP ranks. + # when split across TP ranks. (Skipped for key-only REPLICATE.) second_split = self.get_backend_aware_kv_block_len( layer_idx=i, first_split=False, mamba_view=False ) @@ -1158,10 +1227,13 @@ class NixlConnectorWorker: fa_group_idx = next( i for i, t in enumerate(self._group_spec_types) if _is_attention_spec(t) ) - num_attn_reads = len(plan.source_ranks_per_group[fa_group_idx]) + # SPLIT regions read their head slice from this many remote ranks at a + # per-rank offset; REPLICATE regions read the whole block once. + split_reads = len(plan.source_ranks_per_group[fa_group_idx]) num_blocks = nixl_agent_meta.num_blocks result: list[tuple[int, int, int]] = [] for i, base_addr in enumerate(nixl_agent_meta.kv_caches_base_addr): + replicated = self._is_region_replicated(i) # Read our whole local region size from remote.. local_block_len = self.get_backend_aware_kv_block_len( layer_idx=i, first_split=True, mamba_view=False @@ -1171,8 +1243,13 @@ class NixlConnectorWorker: # ..using remote kv_block_len as transfer unit local_block_len = remote_kv_block_len - local_block_len = local_block_len // num_attn_reads - rank_offset = plan.rank_offset_factor * remote_kv_block_len + # REPLICATE reads the whole block once at offset 0; SPLIT gathers + # its head slice from `split_reads` remote ranks at a per-rank offset. + num_reads = 1 if replicated else split_reads + rank_offset = ( + 0 if replicated else plan.rank_offset_factor * remote_kv_block_len + ) + local_block_len = local_block_len // num_reads page_size = nixl_agent_meta.block_lens[i] for block_id in range(num_blocks): @@ -1182,12 +1259,13 @@ class NixlConnectorWorker: addr = base_addr + block_offset + rank_offset result.append((addr, local_block_len, nixl_agent_meta.device_id)) - if self.transfer_topo.virtually_split_kv_in_blocks: + emits_v = self.transfer_topo.virtually_split_kv_in_blocks and not replicated + if emits_v: # With FlashInfer index V separately to allow head splitting. second_split = self.get_backend_aware_kv_block_len( layer_idx=i, first_split=False, mamba_view=False ) - second_split = second_split // num_attn_reads + second_split = second_split // num_reads for block_id in range(num_blocks): block_offset = block_id * page_size addr = base_addr + block_offset + rank_offset @@ -1527,49 +1605,43 @@ class NixlConnectorWorker: "Use HND layout on the prefill side." ) - # Block len can only vary across layers when using MLA. - remote_block_len = nixl_agent_meta.block_lens[0] - if self.use_mla or self.transfer_topo.is_kv_replicated(remote_engine_id): - # With replicated KV cache, only the number of blocks can differ. - # TODO (ZhanqiuHu): For mamba models, validate FA and mamba - # block_lens separately. - if not self._has_mamba: - for i in range(len(self.block_len_per_layer)): - assert ( - self.block_len_per_layer[i] // block_size_ratio - == nixl_agent_meta.block_lens[i] - ), "KV cache sizes must match between P and D when replicated" - else: - # When MLA is not used, this is a list of the same block length - for block_len in nixl_agent_meta.block_lens: - assert block_len == remote_block_len, ( - "All remote layers must have the same block size" - ) - - # HMA hybrid models (mamba+attention) pad block_len to - # max(attn_page, mamba_page), so the linear tp_ratio scaling - # assumption only holds for pure-attention models. - if not self._has_mamba: - if tp_ratio > 0: - assert ( - remote_block_len - == (self.block_len_per_layer[0] * tp_ratio) // block_size_ratio - ), ( - "Remote P worker KV layer cache must be of shape [2, N," - " local_kv_heads*tp_ratio, page_size, head_dim] and " - "same dtype." + # Per-region block_len validation enforcing the P/D invariant. + # REPLICATE regions (MLA, or a whole-model MLA / replicated-KV transfer) + # only allow the number of blocks to differ; SPLIT regions scale with + # tp_ratio. Mamba uses the ssm_sizes counterpart, so skip block_len here. + if not self._has_mamba: + assert len(self.block_len_per_layer) == len(nixl_agent_meta.block_lens), ( + "Number of KV layers must match between prefill and decode" + ) + model_replicated = self.use_mla or self.transfer_topo.is_kv_replicated( + remote_engine_id + ) + for i, local_len in enumerate(self.block_len_per_layer): + replicated = model_replicated or self._is_region_replicated(i) + remote_len = nixl_agent_meta.block_lens[i] + if replicated: + # Whole block copied; only the number of blocks may differ. + assert local_len // block_size_ratio == remote_len, ( + "KV cache sizes must match between P and D when " + f"replicated (region {i}: local={local_len}, " + f"remote={remote_len}, bsr={block_size_ratio})." + ) + elif tp_ratio > 0: + # D_TP >= P_TP: remote holds tp_ratio x local heads. + assert remote_len == (local_len * tp_ratio) // block_size_ratio, ( + f"SPLIT region {i}: remote P KV block_len {remote_len} " + f"must equal local {local_len} * tp_ratio {tp_ratio} " + f"// block_size_ratio {block_size_ratio}." ) else: + # P_TP > D_TP: local holds |tp_ratio| x remote heads. assert block_size_ratio == 1, ( - "Different local/remote block sizes are not supported" - " when P TP > D TP." + "Different local/remote block sizes are not supported " + "when P TP > D TP." ) - assert remote_block_len == self.block_len_per_layer[0] // ( - -tp_ratio - ), ( - "Remote P worker KV layer cache must be of shape [2, N," - " local_kv_heads/tp_ratio, page_size, head_dim] and " - "same dtype." + assert remote_len == local_len // (-tp_ratio), ( + f"SPLIT region {i}: remote P KV block_len {remote_len} " + f"must equal local {local_len} // |tp_ratio| {-tp_ratio}." ) # TP workers that handhshake with same remote have same #blocks. @@ -2450,13 +2522,16 @@ class NixlConnectorWorker: |1st_split-2nd_split| |1st_split-2nd_split | """ assert self.transfer_topo is not None - if self.transfer_topo.virtually_split_kv_in_blocks: - if mamba_view: - block_len = self._mamba_ssm_size[not first_split] - else: - block_len = self.block_len_per_layer[layer_idx] // 2 + virtually_split = self.transfer_topo.virtually_split_kv_in_blocks + if virtually_split and mamba_view: + block_len = self._mamba_ssm_size[not first_split] else: - block_len = self.block_len_per_layer[layer_idx] + # Per-descriptor block length: a SPLIT region (full-attn under the + # virtually-split layout) emits separate K and V and uses + # block_len//2; REPLICATE (MLA, key-only) and non-split layouts use + # the whole block. + half_block = virtually_split and not self._is_region_replicated(layer_idx) + block_len = self.block_len_per_layer[layer_idx] // (2 if half_block else 1) return block_len def get_kv_connector_stats(self) -> KVConnectorStats | None: From 1ce3cdc5c14f656f81f91e264d557e0b6e6fea54 Mon Sep 17 00:00:00 2001 From: Divakar Verma <137818590+divakar-amd@users.noreply.github.com> Date: Fri, 12 Jun 2026 01:16:14 -0400 Subject: [PATCH 28/52] [ROCm][CI] fix fp8 support for test_deepep_moe (#45302) Signed-off-by: Divakar Verma --- tests/kernels/moe/test_deepep_moe.py | 21 ++++++++++++--------- 1 file changed, 12 insertions(+), 9 deletions(-) diff --git a/tests/kernels/moe/test_deepep_moe.py b/tests/kernels/moe/test_deepep_moe.py index 83cd2f09d1e..4080ca18459 100644 --- a/tests/kernels/moe/test_deepep_moe.py +++ b/tests/kernels/moe/test_deepep_moe.py @@ -27,6 +27,7 @@ from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernel from vllm.model_executor.layers.quantization.utils.fp8_utils import ( per_token_group_quant_fp8, ) +from vllm.platforms import current_platform from vllm.utils.import_utils import has_deep_ep from vllm.utils.torch_utils import set_random_seed from vllm.v1.worker.workspace import init_workspace_manager @@ -64,7 +65,7 @@ def make_weights( return w1, w2, None, None # per-out-channel weight quantization - assert dtype == torch.float8_e4m3fn + assert dtype == current_platform.fp8_dtype() w1 = torch.empty((e, 2 * n, k), device="cuda", dtype=torch.float16) w2 = torch.empty((e, k, n), device="cuda", dtype=torch.float16) @@ -105,9 +106,11 @@ class TestTensors: @staticmethod def make(config: TestConfig, low_latency_mode: bool) -> "TestTensors": # TODO (varun) - check that float16 works ? - assert config.dtype in [torch.bfloat16, torch.float8_e4m3fn] + assert config.dtype in [torch.bfloat16, current_platform.fp8_dtype()] token_dtype = ( - torch.bfloat16 if config.dtype == torch.float8_e4m3fn else config.dtype + torch.bfloat16 + if config.dtype == current_platform.fp8_dtype() + else config.dtype ) rank_tokens = ( torch.randn((config.m, config.k), device="cuda", dtype=token_dtype) / 10 @@ -216,10 +219,10 @@ def deep_ep_moe_impl( return expert_map.to(device=device, dtype=torch.int32) hidden_size = test_tensors.rank_tokens.size(1) - is_quantized = w1.dtype == torch.float8_e4m3fn + is_quantized = w1.dtype == current_platform.fp8_dtype() q_dtype = None if is_quantized: - q_dtype = torch.float8_e4m3fn + q_dtype = current_platform.fp8_dtype() out_hidden_states = torch.empty_like(test_tensors.rank_tokens) total_num_tokens = test_tensors.rank_tokens.size(0) @@ -318,7 +321,7 @@ def torch_moe_impl( .to(a.dtype) ) - is_quantized = w1.dtype == torch.float8_e4m3fn + is_quantized = w1.dtype == current_platform.fp8_dtype() a_dtype = a.dtype if is_quantized: w1 = w1.to(dtype=torch.float32) * w1_scale @@ -367,7 +370,7 @@ def _deep_ep_moe( "FP8 dispatch interface is available only in low-latency mode" ) - is_quantized = w1.dtype == torch.float8_e4m3fn + is_quantized = w1.dtype == current_platform.fp8_dtype() device_idx = torch.accelerator.current_device_index() w1 = w1.to(device=device_idx) w2 = w2.to(device=device_idx) @@ -441,7 +444,7 @@ MNKs = [ (222, 1024, 2048), ] -DTYPES = [torch.bfloat16, torch.float8_e4m3fn] +DTYPES = [torch.bfloat16, current_platform.fp8_dtype()] @pytest.mark.parametrize("dtype", DTYPES) @@ -496,7 +499,7 @@ MNKs = [ (64, 1024, 2560), (222, 1024, 2560), ] -DTYPES = [torch.float8_e4m3fn, torch.bfloat16] +DTYPES = [current_platform.fp8_dtype(), torch.bfloat16] USE_FP8_DISPATCH = [True, False] From eb28452b10a1376d143b2847a78b31726db346dd Mon Sep 17 00:00:00 2001 From: Lucas Wilkinson Date: Fri, 12 Jun 2026 01:17:35 -0400 Subject: [PATCH 29/52] [Model] Add DiffusionGemma Support (#45163) Signed-off-by: Lucas Wilkinson Signed-off-by: Matthew Bonanni Co-authored-by: Martin Kukla Co-authored-by: Matthew Bonanni Co-authored-by: Dipika Sikka Co-authored-by: NickLucche Co-authored-by: jiahanc <173873397+jiahanc@users.noreply.github.com> Co-authored-by: Alec Kohlhoff <134344302+aleckohlhoff@users.noreply.github.com> Co-authored-by: Porras Huang <20535584+porrashuang@users.noreply.github.com> Co-authored-by: Claude Co-authored-by: scoootscooob <167050519+scoootscooob@users.noreply.github.com> --- benchmarks/kernels/benchmark_moe.py | 6 + cmake/external_projects/vllm_flash_attn.cmake | 2 +- docs/design/attention_backends.md | 2 +- .../attention/test_mixed_causal_attn.py | 318 ++++ tests/models/registry.py | 4 + tests/models/utils.py | 4 +- tests/tool_parsers/test_gemma4_tool_parser.py | 82 + tests/v1/cudagraph/test_cudagraph_dispatch.py | 1 + .../unit/test_handshake_pp_aggregation.py | 2 +- .../worker/test_gpu_model_runner_v2_eplb.py | 19 +- vllm/benchmarks/serve.py | 118 +- vllm/config/__init__.py | 3 + vllm/config/diffusion.py | 26 + vllm/config/model.py | 5 + vllm/config/vllm.py | 19 +- vllm/engine/arg_utils.py | 16 + .../experts/flashinfer_cutlass_moe.py | 2 + .../fused_moe/experts/trtllm_nvfp4_moe.py | 1 + .../quantization/utils/flashinfer_utils.py | 1 + vllm/model_executor/models/config.py | 55 + vllm/model_executor/models/diffusion_gemma.py | 1363 +++++++++++++++++ vllm/model_executor/models/gemma4.py | 4 +- vllm/model_executor/models/registry.py | 4 + vllm/tool_parsers/gemma4_tool_parser.py | 180 ++- vllm/transformers_utils/config.py | 1 + vllm/transformers_utils/configs/__init__.py | 4 + .../configs/diffusion_gemma.py | 44 + .../model_arch_config_convertor.py | 1 + vllm/v1/attention/backend.py | 6 +- vllm/v1/attention/backends/fa_utils.py | 6 + vllm/v1/attention/backends/flash_attn.py | 38 +- vllm/v1/attention/backends/triton_attn.py | 9 +- .../attention/ops/triton_attention_helpers.py | 57 +- .../attention/ops/triton_unified_attention.py | 70 +- .../ops/triton_unified_attention_diffkv.py | 1 + vllm/v1/core/sched/async_scheduler.py | 10 +- vllm/v1/core/sched/scheduler.py | 22 +- vllm/v1/cudagraph_dispatcher.py | 6 +- vllm/v1/engine/core.py | 27 +- vllm/v1/metrics/loggers.py | 9 +- vllm/v1/spec_decode/metrics.py | 150 +- vllm/v1/structured_output/__init__.py | 15 +- vllm/v1/worker/gpu/input_batch.py | 23 +- vllm/v1/worker/gpu/model_runner.py | 107 +- vllm/v1/worker/gpu/model_states/__init__.py | 5 + vllm/v1/worker/gpu/model_states/interface.py | 16 + vllm/v1/worker/gpu/sample/output.py | 1 + vllm/v1/worker/gpu/sample/sampler.py | 17 +- .../gpu/spec_decode/rejection_sampler.py | 14 +- vllm/v1/worker/gpu/spec_decode/utils.py | 4 + vllm/v1/worker/gpu/warmup.py | 25 +- vllm/vllm_flash_attn/flash_attn_interface.py | 2 + 52 files changed, 2695 insertions(+), 232 deletions(-) create mode 100644 tests/kernels/attention/test_mixed_causal_attn.py create mode 100644 vllm/config/diffusion.py create mode 100644 vllm/model_executor/models/diffusion_gemma.py create mode 100644 vllm/transformers_utils/configs/diffusion_gemma.py diff --git a/benchmarks/kernels/benchmark_moe.py b/benchmarks/kernels/benchmark_moe.py index f885b1e0952..5d0876f9125 100644 --- a/benchmarks/kernels/benchmark_moe.py +++ b/benchmarks/kernels/benchmark_moe.py @@ -792,6 +792,12 @@ def get_model_params(config): topk = text_config.num_experts_per_tok intermediate_size = text_config.moe_intermediate_size hidden_size = text_config.hidden_size + elif architecture == "DiffusionGemmaForBlockDiffusion": + text_config = config.get_text_config() + E = text_config.num_experts + topk = text_config.top_k_experts + intermediate_size = text_config.moe_intermediate_size + hidden_size = text_config.hidden_size elif architecture == "HunYuanMoEV1ForCausalLM": E = config.num_experts topk = config.moe_topk[0] diff --git a/cmake/external_projects/vllm_flash_attn.cmake b/cmake/external_projects/vllm_flash_attn.cmake index 1e4feb0ff9e..ea7ac544b9d 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 dd62dac706b1cf7895bd99b18c6cb7e7e117ee25 + GIT_TAG 803020a8fa15407871341d41eba4919ade2ee1ee 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/docs/design/attention_backends.md b/docs/design/attention_backends.md index 9ba7afcb9be..a585cd77ffb 100644 --- a/docs/design/attention_backends.md +++ b/docs/design/attention_backends.md @@ -180,7 +180,7 @@ Priority is **1 = highest** (tried first). | `ROCM_AITER_FA` | | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32 | 64, 128, 256 | ✅ | ✅ | ❌ | ❌ | Decoder | N/A | | `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 | -| `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 | +| `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 | | `TRITON_ATTN_DIFFKV` | | fp16, bf16 | `auto`, `bfloat16` | Any | Any | ❌ | ❌ | ❌ | ❌ | Decoder | Any | | `TURBOQUANT` | | fp16, bf16 | `turboquant_k8v4`, `turboquant_4bit_nc`, `turboquant_k3v4_nc`, `turboquant_3bit_nc` | 16, 32, 64, 128 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | Any | diff --git a/tests/kernels/attention/test_mixed_causal_attn.py b/tests/kernels/attention/test_mixed_causal_attn.py new file mode 100644 index 00000000000..5343f701f28 --- /dev/null +++ b/tests/kernels/attention/test_mixed_causal_attn.py @@ -0,0 +1,318 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for per-request causal/non-causal attention (mixed batches). + +Validates that both triton and flash-attention backends correctly handle +batches where some sequences use causal masking and others use non-causal +(bidirectional) masking — needed by DiffusionGemma. +""" + +import pytest +import torch + +from vllm.platforms import current_platform +from vllm.utils.torch_utils import set_random_seed + +# Mixed causal/non-causal attention is only validated on a subset of GPUs: +# the Triton path on Hopper (SM90) and B200 (SM100); the FA4 path on Hopper +# (SM90) only. +_device_capability = current_platform.get_device_capability() +_major = _device_capability.major if _device_capability is not None else None + +NUM_HEADS = [(4, 4), (8, 2)] +HEAD_SIZES = [128] +BLOCK_SIZES = [16] +DTYPES = [torch.bfloat16] + + +def ref_paged_attn( + query: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + query_lens: list[int], + kv_lens: list[int], + block_tables: torch.Tensor, + scale: float, + per_seq_causal: list[bool], + sliding_window: int | None = None, +) -> torch.Tensor: + num_seqs = len(query_lens) + block_tables_np = block_tables.cpu().numpy() + _, block_size, num_kv_heads, head_size = key_cache.shape + + outputs: list[torch.Tensor] = [] + start_idx = 0 + for i in range(num_seqs): + query_len = query_lens[i] + kv_len = kv_lens[i] + q = query[start_idx : start_idx + query_len] + q = q * scale + + num_kv_blocks = (kv_len + block_size - 1) // block_size + block_indices = block_tables_np[i, :num_kv_blocks] + k = key_cache[block_indices].view(-1, num_kv_heads, head_size)[:kv_len] + v = value_cache[block_indices].view(-1, num_kv_heads, head_size)[:kv_len] + + if q.shape[1] != k.shape[1]: + k = torch.repeat_interleave(k, q.shape[1] // k.shape[1], dim=1) + v = torch.repeat_interleave(v, q.shape[1] // v.shape[1], dim=1) + + attn = torch.einsum("qhd,khd->hqk", q, k).float() + + if per_seq_causal[i]: + mask = torch.triu( + torch.ones(query_len, kv_len, device=attn.device), + diagonal=kv_len - query_len + 1, + ).bool() + else: + mask = torch.zeros(query_len, kv_len, device=attn.device).bool() + + if sliding_window is not None: + sw_mask = ( + torch.triu( + torch.ones(query_len, kv_len, device=attn.device), + diagonal=kv_len - (query_len + sliding_window) + 1, + ) + .bool() + .logical_not() + ) + mask |= sw_mask + + attn.masked_fill_(mask, float("-inf")) + attn = torch.softmax(attn, dim=-1).to(v.dtype) + out = torch.einsum("hqk,khd->qhd", attn, v) + outputs.append(out) + start_idx += query_len + + return torch.cat(outputs, dim=0) + + +# ---- Triton backend test ---- + + +@pytest.mark.skipif( + _major not in (9, 10), + reason="Triton mixed causal attention requires Hopper (SM90) or B200 (SM100).", +) +@pytest.mark.parametrize( + "seq_lens", + [[(1, 128), (5, 64), (1, 256)]], +) +@pytest.mark.parametrize( + "per_seq_causal", + [[True, False, True], [False, True, False], [True, True, False]], +) +@pytest.mark.parametrize("num_heads", NUM_HEADS) +@pytest.mark.parametrize("head_size", HEAD_SIZES) +@pytest.mark.parametrize("block_size", BLOCK_SIZES) +@pytest.mark.parametrize("dtype", DTYPES) +@torch.inference_mode() +def test_triton_mixed_causal( + seq_lens: list[tuple[int, int]], + per_seq_causal: list[bool], + num_heads: tuple[int, int], + head_size: int, + dtype: torch.dtype, + block_size: int, +): + if not current_platform.is_cuda(): + pytest.skip("Triton attention requires CUDA") + + from vllm.v1.attention.ops.triton_unified_attention import unified_attention + + set_random_seed(42) + device = "cuda" + + num_query_heads, num_kv_heads = num_heads + assert len(seq_lens) == len(per_seq_causal) + + query_lens = [s[0] for s in seq_lens] + kv_lens = [s[1] for s in seq_lens] + num_seqs = len(seq_lens) + + num_query_tokens = sum(query_lens) + max_kv_len = max(kv_lens) + max_num_blocks = (max_kv_len + block_size - 1) // block_size + num_blocks = max_num_blocks * num_seqs + 10 + + scale = head_size**-0.5 + query = torch.randn( + num_query_tokens, num_query_heads, head_size, dtype=dtype, device=device + ) + key_cache = torch.randn( + num_blocks, block_size, num_kv_heads, head_size, dtype=dtype, device=device + ) + value_cache = torch.randn( + num_blocks, block_size, num_kv_heads, head_size, dtype=dtype, device=device + ) + + block_tables_list = [] + for i in range(num_seqs): + n_blocks = (kv_lens[i] + block_size - 1) // block_size + blocks = list(range(i * max_num_blocks, i * max_num_blocks + n_blocks)) + blocks += [0] * (max_num_blocks - n_blocks) + block_tables_list.append(blocks) + block_tables = torch.tensor(block_tables_list, dtype=torch.int32, device=device) + + cu_seqlens_q = torch.zeros(num_seqs + 1, dtype=torch.int32, device=device) + for i, ql in enumerate(query_lens): + cu_seqlens_q[i + 1] = cu_seqlens_q[i] + ql + + seqused_k = torch.tensor(kv_lens, dtype=torch.int32, device=device) + max_seqlen_q = max(query_lens) + max_seqlen_k = max(kv_lens) + + causal_tensor = torch.tensor(per_seq_causal, dtype=torch.bool, device=device) + + output = torch.empty_like(query) + unified_attention( + q=query, + k=key_cache, + v=value_cache, + out=output, + cu_seqlens_q=cu_seqlens_q, + max_seqlen_q=max_seqlen_q, + seqused_k=seqused_k, + max_seqlen_k=max_seqlen_k, + softmax_scale=scale, + causal=causal_tensor, + window_size=(-1, -1), + block_table=block_tables, + softcap=0.0, + q_descale=None, + k_descale=1.0, + v_descale=1.0, + ) + + ref_output = ref_paged_attn( + query, + key_cache, + value_cache, + query_lens, + kv_lens, + block_tables, + scale, + per_seq_causal, + ) + + torch.testing.assert_close(output, ref_output, atol=1e-2, rtol=1e-2) + + +# ---- Flash Attention 4 backend test (native per_seq_causal) ---- + + +@pytest.mark.skipif( + _major != 9, + reason="FA4 mixed causal attention requires Hopper (SM90).", +) +@pytest.mark.parametrize( + "seq_lens", + [[(1, 128), (5, 64), (1, 256)]], +) +@pytest.mark.parametrize( + "per_seq_causal", + [[True, False, True], [False, True, False]], +) +@pytest.mark.parametrize("num_heads", NUM_HEADS) +@pytest.mark.parametrize("head_size", HEAD_SIZES) +@pytest.mark.parametrize("block_size", BLOCK_SIZES) +@pytest.mark.parametrize("dtype", DTYPES) +@torch.inference_mode() +def test_flash_attn4_mixed_causal( + seq_lens: list[tuple[int, int]], + per_seq_causal: list[bool], + num_heads: tuple[int, int], + head_size: int, + dtype: torch.dtype, + block_size: int, +): + if not current_platform.is_cuda(): + pytest.skip("Flash attention requires CUDA") + + try: + from vllm.vllm_flash_attn import ( + fa_version_unsupported_reason, + flash_attn_varlen_func, + is_fa_version_supported, + ) + except ImportError: + pytest.skip("vllm_flash_attn not available") + + if not is_fa_version_supported(4): + reason = fa_version_unsupported_reason(4) + pytest.skip(f"FA4 not supported: {reason}") + + set_random_seed(42) + device = "cuda" + + num_query_heads, num_kv_heads = num_heads + assert len(seq_lens) == len(per_seq_causal) + + query_lens = [s[0] for s in seq_lens] + kv_lens = [s[1] for s in seq_lens] + num_seqs = len(seq_lens) + + num_query_tokens = sum(query_lens) + max_kv_len = max(kv_lens) + max_num_blocks = (max_kv_len + block_size - 1) // block_size + num_blocks = max_num_blocks * num_seqs + 10 + + scale = head_size**-0.5 + query = torch.randn( + num_query_tokens, num_query_heads, head_size, dtype=dtype, device=device + ) + key_cache = torch.randn( + num_blocks, block_size, num_kv_heads, head_size, dtype=dtype, device=device + ) + value_cache = torch.randn( + num_blocks, block_size, num_kv_heads, head_size, dtype=dtype, device=device + ) + + block_tables_list = [] + for i in range(num_seqs): + n_blocks = (kv_lens[i] + block_size - 1) // block_size + blocks = list(range(i * max_num_blocks, i * max_num_blocks + n_blocks)) + blocks += [0] * (max_num_blocks - n_blocks) + block_tables_list.append(blocks) + block_tables = torch.tensor(block_tables_list, dtype=torch.int32, device=device) + + cu_seqlens_q = torch.zeros(num_seqs + 1, dtype=torch.int32, device=device) + for i, ql in enumerate(query_lens): + cu_seqlens_q[i + 1] = cu_seqlens_q[i] + ql + + seqused_k = torch.tensor(kv_lens, dtype=torch.int32, device=device) + per_seq_causal_tensor = torch.tensor( + per_seq_causal, dtype=torch.int32, device=device + ) + + ref_output = ref_paged_attn( + query, + key_cache, + value_cache, + query_lens, + kv_lens, + block_tables, + scale, + per_seq_causal, + ) + + output = torch.empty_like(query) + flash_attn_varlen_func( + q=query, + k=key_cache, + v=value_cache, + out=output, + cu_seqlens_q=cu_seqlens_q, + max_seqlen_q=max(query_lens), + seqused_k=seqused_k, + max_seqlen_k=max(kv_lens), + softmax_scale=scale, + # The kernel must be compiled causal for `dynamic_causal` to take effect. + causal=True, + block_table=block_tables, + softcap=0.0, + dynamic_causal=per_seq_causal_tensor, + fa_version=4, + ) + + torch.testing.assert_close(output, ref_output, atol=1e-2, rtol=1e-2) diff --git a/tests/models/registry.py b/tests/models/registry.py index 120a0ca8b85..ed15ac5f46f 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -898,6 +898,10 @@ _MULTIMODAL_EXAMPLE_MODELS = { ), "FuyuForCausalLM": _HfExamplesInfo("adept/fuyu-8b"), "Gemma3ForConditionalGeneration": _HfExamplesInfo("google/gemma-3-4b-it"), + "DiffusionGemmaForBlockDiffusion": _HfExamplesInfo( + "google/diffusiongemma-26B-A4B-it", + trust_remote_code=True, + ), "Gemma4ForConditionalGeneration": _HfExamplesInfo( "google/gemma-4-E2B-it", min_transformers_version="5.5.0", diff --git a/tests/models/utils.py b/tests/models/utils.py index a5d1844a307..259cdac13c0 100644 --- a/tests/models/utils.py +++ b/tests/models/utils.py @@ -486,6 +486,7 @@ def dummy_hf_overrides( "Gemma3nForConditionalGeneration", "Gemma4ForCausalLM", "Gemma4ForConditionalGeneration", + "DiffusionGemmaForBlockDiffusion", ) else 1 ) @@ -558,7 +559,8 @@ def dummy_hf_overrides( ) # e.g.: Qwen/Qwen2-Audio-7B-Instruct - if hasattr(hf_config, "audio_config"): + # audio_config may exist but be None (e.g. audio-less Gemma4 variants). + if getattr(hf_config, "audio_config", None) is not None: hf_config.audio_config.update( { "num_layers": 1, diff --git a/tests/tool_parsers/test_gemma4_tool_parser.py b/tests/tool_parsers/test_gemma4_tool_parser.py index 6f3709e19a4..eea084a2bb4 100644 --- a/tests/tool_parsers/test_gemma4_tool_parser.py +++ b/tests/tool_parsers/test_gemma4_tool_parser.py @@ -702,6 +702,88 @@ class TestStreamingExtraction: ' \n' ) + def _collect_tool_calls_by_index(self, results): + """Group streamed tool-call fragments by their ``index``. + + Returns ``{index: {"name": str | None, "arguments": str}}`` where + ``arguments`` is the concatenation of every streamed argument + fragment for that index (which should form valid JSON once complete). + """ + by_index: dict[int, dict[str, Any]] = {} + for delta, _ in results: + if not (delta and delta.tool_calls): + continue + for tc in delta.tool_calls: + entry = by_index.setdefault(tc.index, {"name": None, "arguments": ""}) + func = tc.function + if isinstance(func, dict): + name = func.get("name") + arg = func.get("arguments", "") + else: + name = getattr(func, "name", None) + arg = getattr(func, "arguments", "") or "" + if name: + entry["name"] = name + if arg: + entry["arguments"] += arg + return by_index + + def test_streaming_single_chunk_complete_tool_call(self, parser, mock_request): + """A backend may deliver a whole tool call in one streaming delta. + + The start token, ``call:name{...}`` payload and the end token all + arrive in a single chunk. The parser must still emit one + ``DeltaToolCall`` with the correct name + complete arguments JSON + (rather than swallowing it and finishing with finish_reason="stop"). + """ + chunks = [ + '<|tool_call>call:name_a_color{color_hex:<|"|>00ff11<|"|>}', + ] + + results = self._simulate_streaming(parser, mock_request, chunks) + + # Exactly one delta should carry tool_calls, and it must not be + # emitted as plain content (which would yield finish_reason="stop"). + tool_call_deltas = [ + delta for delta, _ in results if delta is not None and delta.tool_calls + ] + assert len(tool_call_deltas) == 1, ( + "Expected exactly one delta carrying the batched tool call" + ) + assert all( + delta.content is None for delta, _ in results if delta is not None + ), "Complete tool call must not leak as content" + + by_index = self._collect_tool_calls_by_index(results) + assert set(by_index) == {0} + assert by_index[0]["name"] == "name_a_color" + assert json.loads(by_index[0]["arguments"]) == {"color_hex": "00ff11"} + + def test_streaming_multi_chunk_batched_tool_calls(self, parser, mock_request): + """A single delta may batch MULTIPLE complete tool calls. + + ``<|tool_call>...<|tool_call>...`` arriving in + one chunk must emit BOTH calls (one DeltaToolCall each, with distinct + indices), not just the first. + """ + chunks = [ + '<|tool_call>call:get_weather{location:<|"|>London<|"|>}' + '<|tool_call>call:get_time{timezone:<|"|>GMT<|"|>}', + ] + + results = self._simulate_streaming(parser, mock_request, chunks) + + by_index = self._collect_tool_calls_by_index(results) + assert set(by_index) == {0, 1}, ( + f"Expected two tool calls (indices 0 and 1), got {sorted(by_index)}" + ) + + assert by_index[0]["name"] == "get_weather" + assert json.loads(by_index[0]["arguments"]) == {"location": "London"} + + assert by_index[1]["name"] == "get_time" + assert json.loads(by_index[1]["arguments"]) == {"timezone": "GMT"} + def test_streaming_trailing_bare_bool_not_duplicated(self, parser, mock_request): """Trailing bare boolean must not be streamed twice.""" chunks = [ diff --git a/tests/v1/cudagraph/test_cudagraph_dispatch.py b/tests/v1/cudagraph/test_cudagraph_dispatch.py index 97b5fd46a2e..c10835821f5 100644 --- a/tests/v1/cudagraph/test_cudagraph_dispatch.py +++ b/tests/v1/cudagraph/test_cudagraph_dispatch.py @@ -49,6 +49,7 @@ def _create_vllm_config( ) mock_config.parallel_config = ParallelConfig() mock_config.speculative_config = None # No speculative decoding + mock_config.num_speculative_tokens = 0 if not lora_config: mock_config.lora_config = None else: diff --git a/tests/v1/kv_connector/unit/test_handshake_pp_aggregation.py b/tests/v1/kv_connector/unit/test_handshake_pp_aggregation.py index 4a2ca6d2721..0c0f9f1f899 100644 --- a/tests/v1/kv_connector/unit/test_handshake_pp_aggregation.py +++ b/tests/v1/kv_connector/unit/test_handshake_pp_aggregation.py @@ -104,7 +104,7 @@ def _run_engine_core_handshake( speculative_config=None, ec_transfer_config=None, max_concurrent_batches=1, - model_config=SimpleNamespace(runner_type="generate"), + model_config=SimpleNamespace(runner_type="generate", is_diffusion=False), cache_config=SimpleNamespace( enable_prefix_caching=False, prefix_caching_hash_algo="builtin", diff --git a/tests/v1/worker/test_gpu_model_runner_v2_eplb.py b/tests/v1/worker/test_gpu_model_runner_v2_eplb.py index 1db07baf93d..9d39621f4fa 100644 --- a/tests/v1/worker/test_gpu_model_runner_v2_eplb.py +++ b/tests/v1/worker/test_gpu_model_runner_v2_eplb.py @@ -70,6 +70,7 @@ def _make_runner(**overrides: Any) -> Any: runner.use_aux_hidden_state_outputs = False runner.speculative_config = None runner.speculator = None + runner.num_speculative_steps = 0 runner.encoder_cache = None runner.is_pooling_model = False runner.is_last_pp_rank = True @@ -102,18 +103,22 @@ def test_v2_load_model_registers_moe_with_eplb(monkeypatch): lambda load_config: SimpleNamespace(load_model=lambda **_: model), ) monkeypatch.setattr(mrv2, "prepare_communication_buffer_for_model", prepared.append) - monkeypatch.setattr(mrv2, "init_model_state", lambda *args: "model-state") + monkeypatch.setattr( + mrv2, + "init_model_state", + lambda *args: SimpleNamespace(num_new_sampled_tokens_per_step=1), + ) monkeypatch.setattr( eplb, "is_mixture_of_experts", lambda loaded_model: getattr(loaded_model, "is_moe", False), ) - runner = _make_runner() + runner = _make_runner(is_last_pp_rank=False) mrv2.GPUModelRunner.load_model(runner) assert runner.model is model - assert runner.model_state == "model-state" + assert runner.model_state is not None assert prepared == [model] assert runner.eplb_state is not None assert runner.eplb_state.add_model_calls == [(model, runner.model_config)] @@ -133,10 +138,14 @@ def test_v2_load_model_with_dummy_weights_skips_eplb_registration(monkeypatch): lambda load_config: SimpleNamespace(load_model=lambda **_: model), ) monkeypatch.setattr(mrv2, "prepare_communication_buffer_for_model", prepared.append) - monkeypatch.setattr(mrv2, "init_model_state", lambda *args: "model-state") + monkeypatch.setattr( + mrv2, + "init_model_state", + lambda *args: SimpleNamespace(num_new_sampled_tokens_per_step=1), + ) monkeypatch.setattr(eplb, "is_mixture_of_experts", lambda *_: True) - runner = _make_runner() + runner = _make_runner(is_last_pp_rank=False) mrv2.GPUModelRunner.load_model(runner, load_dummy_weights=True) assert runner.load_config.load_format == "dummy" diff --git a/vllm/benchmarks/serve.py b/vllm/benchmarks/serve.py index cbf7be44ae9..4d6fdbe22af 100644 --- a/vllm/benchmarks/serve.py +++ b/vllm/benchmarks/serve.py @@ -248,6 +248,68 @@ async def fetch_spec_decode_metrics( return None +@dataclass +class DiffusionMetrics: + """Diffusion (dLLM) decoding metrics from the server's Prometheus endpoint.""" + + num_denoising_steps: int + num_canvas_positions: int + num_committed_tokens: int + + +async def fetch_diffusion_metrics( + base_url: str, session: aiohttp.ClientSession +) -> DiffusionMetrics | None: + """Fetch diffusion decoding metrics from the server's Prometheus endpoint. + + Returns None if the model is not a diffusion model or metrics are not + available. + """ + metrics_url = f"{base_url}/metrics" + try: + async with session.get(metrics_url) as response: + if response.status != 200: + return None + text = await response.text() + + num_denoising_steps = 0 + num_canvas_positions = 0 + num_committed_tokens = 0 + found_diffusion = False + + for line in text.split("\n"): + line = line.strip() + if not line or line.startswith("#"): + continue + + if line.startswith("vllm:diffusion"): + # Extract metric name (before labels) to avoid matching + # substrings inside label values. + parts = line.split(None, 1) + metric_name = parts[0].split("{")[0] + if not metric_name.endswith("_total"): + continue + found_diffusion = True + with contextlib.suppress(ValueError): + if "num_denoising_steps" in metric_name: + num_denoising_steps += int(float(parts[-1])) + elif "num_canvas_positions" in metric_name: + num_canvas_positions += int(float(parts[-1])) + elif "num_committed_tokens" in metric_name: + num_committed_tokens += int(float(parts[-1])) + + if not found_diffusion: + return None + + return DiffusionMetrics( + num_denoising_steps=num_denoising_steps, + num_canvas_positions=num_canvas_positions, + num_committed_tokens=num_committed_tokens, + ) + except (aiohttp.ClientError, asyncio.TimeoutError): + return None + + class TaskType(Enum): GENERATION = "generation" POOLING = "pooling" @@ -887,6 +949,7 @@ async def benchmark( print("Self timing is set, using the timestamps from the trace file.") spec_decode_metrics_before = await fetch_spec_decode_metrics(base_url, session) + diffusion_metrics_before = await fetch_diffusion_metrics(base_url, session) pbar = None if disable_tqdm else tqdm(total=len(input_requests)) @@ -1016,6 +1079,34 @@ async def benchmark( "per_position_acceptance_rates": per_pos_rates, } + diffusion_metrics_after = await fetch_diffusion_metrics(base_url, session) + diffusion_stats: dict[str, Any] | None = None + if diffusion_metrics_before is not None and diffusion_metrics_after is not None: + delta_steps = ( + diffusion_metrics_after.num_denoising_steps + - diffusion_metrics_before.num_denoising_steps + ) + delta_positions = ( + diffusion_metrics_after.num_canvas_positions + - diffusion_metrics_before.num_canvas_positions + ) + delta_committed = ( + diffusion_metrics_after.num_committed_tokens + - diffusion_metrics_before.num_committed_tokens + ) + if delta_steps > 0 and delta_committed > 0: + block_size = delta_positions / delta_steps # canvas length (CL) + num_canvases = delta_committed / block_size # = number of commit steps + denoising_steps = delta_steps - num_canvases # exclude commit steps + diffusion_stats = { + "denoising_steps": denoising_steps, + "canvas_positions": delta_positions, + "committed_tokens": delta_committed, + "committed_throughput": delta_committed / benchmark_duration, + "steps_per_canvas": denoising_steps / num_canvases, + "committed_per_step": delta_committed / denoising_steps, + } + if task_type == TaskType.GENERATION: metrics, actual_output_lens = calculate_metrics( input_requests=input_requests, @@ -1134,6 +1225,16 @@ async def benchmark( "per_position_acceptance_rates", [] ) + if diffusion_stats is not None: + result["diffusion_committed_throughput"] = diffusion_stats[ + "committed_throughput" + ] + result["diffusion_steps_per_canvas"] = diffusion_stats["steps_per_canvas"] + result["diffusion_committed_per_step"] = diffusion_stats["committed_per_step"] + result["diffusion_committed_tokens"] = int(diffusion_stats["committed_tokens"]) + result["diffusion_denoising_steps"] = int(diffusion_stats["denoising_steps"]) + result["diffusion_canvas_positions"] = int(diffusion_stats["canvas_positions"]) + def process_one_metric( # E.g., "ttft" metric_attribute_name: str, @@ -1179,7 +1280,22 @@ async def benchmark( process_one_metric("itl", "ITL", "Inter-token Latency") process_one_metric("e2el", "E2EL", "End-to-end Latency") - if spec_decode_stats is not None: + if diffusion_stats is not None: + print("{s:{c}^{n}}".format(s="Diffusion Decoding", n=50, c="-")) + for label, key, value_fmt in ( + ("Committed throughput (tok/s):", "committed_throughput", "{:<10.2f}"), + ("Denoising steps per canvas:", "steps_per_canvas", "{:<10.2f}"), + ("Committed per denoising step:", "committed_per_step", "{:<10.2f}"), + ("Committed tokens:", "committed_tokens", "{:<10d}"), + ("Denoising steps:", "denoising_steps", "{:<10d}"), + ("Canvas positions evaluated:", "canvas_positions", "{:<10d}"), + ): + value = diffusion_stats[key] + if value_fmt.endswith("d}"): + value = int(value) + print("{:<40} ".format(label) + value_fmt.format(value)) + + if spec_decode_stats is not None and diffusion_stats is None: print("{s:{c}^{n}}".format(s="Speculative Decoding", n=50, c="-")) print( "{:<40} {:<10.2f}".format( diff --git a/vllm/config/__init__.py b/vllm/config/__init__.py index b189c45c8d7..82ab1842fe9 100644 --- a/vllm/config/__init__.py +++ b/vllm/config/__init__.py @@ -10,6 +10,7 @@ from vllm.config.compilation import ( PassConfig, ) from vllm.config.device import DeviceConfig +from vllm.config.diffusion import DiffusionConfig from vllm.config.ec_transfer import ECTransferConfig from vllm.config.kernel import KernelConfig from vllm.config.kv_events import KVEventsConfig @@ -72,6 +73,8 @@ __all__ = [ "PassConfig", # From vllm.config.device "DeviceConfig", + # From vllm.config.diffusion + "DiffusionConfig", # From vllm.config.ec_transfer "ECTransferConfig", # From vllm.config.kernel diff --git a/vllm/config/diffusion.py b/vllm/config/diffusion.py new file mode 100644 index 00000000000..6f59c40a836 --- /dev/null +++ b/vllm/config/diffusion.py @@ -0,0 +1,26 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Configuration for discrete diffusion (dLLM) models.""" + +from pydantic import Field + +from vllm.config.utils import config + + +@config +class DiffusionConfig: + """Configuration for discrete diffusion language models (dLLMs). + + dLLMs generate tokens via iterative denoising over a fixed-length canvas + rather than left-to-right autoregressive decoding. They reuse the + speculative-decoding data path (draft token ids, scheduled spec decode + tokens) with overloaded semantics for block-based generation. + """ + + canvas_length: int = Field(default=None, gt=0) # type: ignore[assignment] + """Length of the denoising canvas (block). Also determines the number of + speculative tokens scheduled per step.""" + + max_denoising_steps: int | None = None + """Maximum number of denoising iterations per canvas block. + If not set, read from the model's generation_config.json.""" diff --git a/vllm/config/model.py b/vllm/config/model.py index 015e75afac2..42c11eacd46 100644 --- a/vllm/config/model.py +++ b/vllm/config/model.py @@ -1546,6 +1546,11 @@ class ModelConfig: """Extract the HF encoder/decoder model flag.""" return is_encoder_decoder(self.hf_config) + @cached_property + def is_diffusion(self) -> bool: + """Detect discrete diffusion (dLLM) models from HF config.""" + return getattr(self.hf_config, "canvas_length", None) is not None + @property def uses_alibi(self) -> bool: cfg = self.hf_text_config diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 86a2f4d09e0..890d2b72e31 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -31,6 +31,7 @@ from .attention import AttentionConfig from .cache import CacheConfig from .compilation import CompilationConfig, CompilationMode, CUDAGraphMode from .device import DeviceConfig +from .diffusion import DiffusionConfig from .ec_transfer import ECTransferConfig from .kernel import KernelConfig from .kv_events import KVEventsConfig @@ -323,6 +324,9 @@ class VllmConfig: """LoRA configuration.""" speculative_config: SpeculativeConfig | None = None """Speculative decoding configuration.""" + diffusion_config: DiffusionConfig | None = None + """Diffusion LLM (dLLM) configuration.""" + structured_outputs_config: StructuredOutputsConfig = Field( default_factory=StructuredOutputsConfig ) @@ -511,6 +515,11 @@ class VllmConfig: and self.speculative_config.num_speculative_tokens is not None ): return self.speculative_config.num_speculative_tokens + if ( + self.diffusion_config is not None + and self.diffusion_config.canvas_length is not None + ): + return self.diffusion_config.canvas_length return 0 @property @@ -519,6 +528,9 @@ class VllmConfig: if use_v2_model_runner is not None: return use_v2_model_runner + if self.model_config is not None and self.model_config.is_diffusion: + return True + if not self._is_default_v2_model_runner_model(): return False @@ -1654,12 +1666,7 @@ class VllmConfig: self.compilation_config.max_cudagraph_capture_size ) if max_cudagraph_capture_size is None: - decode_query_len = 1 - if ( - self.speculative_config - and self.speculative_config.num_speculative_tokens - ): - decode_query_len += self.speculative_config.num_speculative_tokens + decode_query_len = 1 + self.num_speculative_tokens max_cudagraph_capture_size = min( self.scheduler_config.max_num_seqs * decode_query_len * 2, 512 ) diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index f0dade83716..f863fad17de 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -38,6 +38,7 @@ from vllm.config import ( CompilationConfig, ConfigType, DeviceConfig, + DiffusionConfig, ECTransferConfig, EPLBConfig, KernelConfig, @@ -616,6 +617,7 @@ class EngineArgs: spec_method: str | None = None spec_model: str | None = None spec_tokens: int | None = None + diffusion_config: dict[str, Any] | None = None show_hidden_metrics_for_version: str | None = ( ObservabilityConfig.show_hidden_metrics_for_version @@ -1473,6 +1475,10 @@ class EngineArgs: vllm_group.add_argument( "--spec-tokens", **speculative_kwargs["num_speculative_tokens"] ) + vllm_kwargs["diffusion_config"]["type"] = optional_type(json.loads) + vllm_group.add_argument( + "--diffusion-config", "-dc", **vllm_kwargs["diffusion_config"] + ) vllm_group.add_argument( "--kv-transfer-config", **vllm_kwargs["kv_transfer_config"] ) @@ -1702,6 +1708,14 @@ class EngineArgs: ) return SpeculativeConfig(**self.speculative_config) + def create_diffusion_config(self) -> DiffusionConfig | None: + if self.diffusion_config is None: + return None + cfg = self.diffusion_config + if isinstance(cfg, str): + cfg = json.loads(cfg) + return DiffusionConfig(**cfg) + def create_engine_config( self, usage_context: UsageContext | None = None, @@ -2016,6 +2030,7 @@ class EngineArgs: target_model_config=model_config, target_parallel_config=parallel_config, ) + diffusion_config = self.create_diffusion_config() self._set_default_max_num_seqs_and_batched_tokens_args( usage_context, @@ -2243,6 +2258,7 @@ class EngineArgs: kernel_config=kernel_config, lora_config=lora_config, speculative_config=speculative_config, + diffusion_config=diffusion_config, structured_outputs_config=self.structured_outputs_config, observability_config=observability_config, compilation_config=compilation_config, diff --git a/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutlass_moe.py b/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutlass_moe.py index ff259c828f4..76cd15ff5a0 100644 --- a/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutlass_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutlass_moe.py @@ -188,6 +188,7 @@ class FlashInferExperts(mk.FusedMoEExpertsModular): def _supports_activation(activation: MoEActivation) -> bool: return activation in [ MoEActivation.SILU, + MoEActivation.GELU_TANH, MoEActivation.RELU2_NO_MUL, MoEActivation.SWIGLUOAI, ] @@ -267,6 +268,7 @@ class FlashInferExperts(mk.FusedMoEExpertsModular): activation_str_to_value_map = { MoEActivation.SILU: ActivationType.Swiglu, # This is the default + MoEActivation.GELU_TANH: ActivationType.Geglu, MoEActivation.SWIGLUOAI: ActivationType.Swiglu, # gpt-oss alias MoEActivation.RELU2_NO_MUL: ActivationType.Relu2, } diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py index e90c4d6646e..e45fc77ad90 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py @@ -142,6 +142,7 @@ class TrtLlmNvFp4ExpertsBase: MoEActivation.SILU, MoEActivation.RELU2_NO_MUL, MoEActivation.GELU, + MoEActivation.GELU_TANH, ] @staticmethod diff --git a/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py b/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py index 61b52345ab8..26fea5d5244 100644 --- a/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py +++ b/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py @@ -34,6 +34,7 @@ def activation_to_flashinfer_type(activation: MoEActivation) -> "ActivationType" MoEActivation.GELU_NO_MUL: ActivationType.Gelu, MoEActivation.SILU: ActivationType.Swiglu, MoEActivation.GELU: ActivationType.Geglu, + MoEActivation.GELU_TANH: ActivationType.Geglu, MoEActivation.RELU2_NO_MUL: ActivationType.Relu2, } return ACTIVATION_TO_FI_ACTIVATION[activation] diff --git a/vllm/model_executor/models/config.py b/vllm/model_executor/models/config.py index 64d606c2890..7354771764d 100644 --- a/vllm/model_executor/models/config.py +++ b/vllm/model_executor/models/config.py @@ -105,6 +105,60 @@ class Gemma4Config(VerifyAndUpdateConfig): ) +class DiffusionGemmaModelForBlockDiffusionConfig(VerifyAndUpdateConfig): + @classmethod + def verify_and_update_config(cls, vllm_config: "VllmConfig") -> None: + """Set up the diffusion config and defaults for DiffusionGemma. + + Auto-creates DiffusionConfig from the HF config when the user + didn't pass ``--diffusion-config``. Diffusion sampling params are + read straight from generation_config.json at sampler-build time + (see DiffusionGemma's custom_sampler), not injected here. + """ + # Inherit Gemma4's attention backend selection (FA4 on Hopper, + # TRITON_ATTN fallback for heterogeneous head dims). + Gemma4Config.verify_and_update_config(vllm_config) + + from vllm.v1.attention.backends.registry import AttentionBackendEnum + + attention_config = vllm_config.attention_config + if attention_config.backend == AttentionBackendEnum.FLASHINFER: + raise ValueError( + "FlashInfer does not support DiffusionGemma's mixed " + "causal/bidirectional attention. Use --attention-backend " + "FLASH_ATTN or TRITON_ATTN instead." + ) + if attention_config.backend is None and not attention_config.use_non_causal: + attention_config.use_non_causal = True + logger.info( + "DiffusionGemma uses mixed causal/bidirectional attention " + "within a batch; setting use_non_causal=True to exclude " + "FlashInfer from auto-selection." + ) + + # Auto-create DiffusionConfig from HF config if not provided. + if vllm_config.diffusion_config is None: + from vllm.config.diffusion import DiffusionConfig + + hf_config = vllm_config.model_config.hf_config + canvas_length = getattr(hf_config, "canvas_length", 256) + vllm_config.diffusion_config = DiffusionConfig( + canvas_length=canvas_length, + ) + + # The diffusion sampler materializes [num_seqs, canvas_length, vocab] + # fp32 transients, so concurrency is memory-bound (>8 OOMs a single H200). + # Default to 8 when the user didn't pass --max-num-seqs. + # We can't see the original None here (the engine already filled a generic + # default), so use >= DEFAULT_MAX_NUM_SEQS as a proxy, (the default is much + # larger than any deliberate value for this model) + from vllm.config.scheduler import SchedulerConfig + + sc = vllm_config.scheduler_config + if sc is not None and sc.max_num_seqs >= SchedulerConfig.DEFAULT_MAX_NUM_SEQS: + sc.max_num_seqs = 8 + + class DeepseekV4ForCausalLMConfig(VerifyAndUpdateConfig): @staticmethod def verify_and_update_model_config(model_config: "ModelConfig") -> None: @@ -591,6 +645,7 @@ MODELS_CONFIG_MAP: dict[str, type[VerifyAndUpdateConfig]] = { "ColQwen3_5": Qwen3_5ForConditionalGenerationConfig, "DeepseekV4ForCausalLM": DeepseekV4ForCausalLMConfig, "DeepseekV32ForCausalLM": DeepseekV32ForCausalLM, + "DiffusionGemmaForBlockDiffusion": DiffusionGemmaModelForBlockDiffusionConfig, # noqa: E501 "Ernie4_5_VLMoeForConditionalGeneration": Ernie4_5_VLMoeForConditionalGenerationConfig, # noqa: E501 "FalconMambaForCausalLM": MambaModelConfig, "Gemma3TextModel": Gemma3TextModelConfig, diff --git a/vllm/model_executor/models/diffusion_gemma.py b/vllm/model_executor/models/diffusion_gemma.py new file mode 100644 index 00000000000..91dd5e6b6a5 --- /dev/null +++ b/vllm/model_executor/models/diffusion_gemma.py @@ -0,0 +1,1363 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""DiffusionGemma model, ModelState, and Sampler for vLLM. + +Single Gemma4 backbone run in two modes (like YOCO): +- encoder mode: causal attention, writes KV cache +- decoder mode: bidirectional attention, reads encoder KV, doesn't write + +Same weights, same layers. The only decoder-unique component is a +self-conditioning MLP. + +Multimodal support: the model always includes a vision tower (shared with Gemma4). +Images are encoded through the vision tower and projected into the LM embedding space +via Gemma4MultimodalEmbedder. +""" + +from __future__ import annotations + +from collections.abc import Iterable, Mapping +from types import SimpleNamespace +from typing import Any + +import numpy as np +import torch +from torch import nn +from torch.nn import functional as F +from transformers import AutoModel + +from vllm.config import VllmConfig +from vllm.config.compilation import CUDAGraphMode +from vllm.logger import init_logger +from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.vocab_parallel_embedding import ( + ParallelLMHead, +) +from vllm.model_executor.models.gemma4 import Gemma4Model +from vllm.model_executor.models.gemma4_mm import ( + Gemma4DummyInputsBuilder, + Gemma4ForConditionalGeneration, + Gemma4MultimodalEmbedder, + Gemma4MultiModalProcessor, + Gemma4ProcessingInfo, +) +from vllm.model_executor.models.module_mapping import MultiModelKeys +from vllm.model_executor.models.transformers.utils import recursive_replace_linear +from vllm.model_executor.models.utils import WeightsMapper, maybe_prefix +from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.v1.outputs import LogprobsTensors +from vllm.v1.worker.gpu.attn_utils import build_attn_metadata +from vllm.v1.worker.gpu.buffer_utils import UvaBackedTensor, async_copy_to_gpu +from vllm.v1.worker.gpu.model_states.interface import ModelState +from vllm.v1.worker.gpu.sample.logprob import compute_topk_logprobs +from vllm.v1.worker.gpu.sample.output import SamplerOutput +from vllm.v1.worker.gpu.sample.penalties import use_penalty + +from .interfaces import ( + SupportsMultiModal, + SupportsPP, + SupportsQuant, +) + +logger = init_logger(__name__) + + +class DiffusionGemmaSelfConditioning(nn.Module): + """Gated MLP that processes soft embeddings from the previous denoising step. + + Structurally identical to Gemma4MLP but with self_conditioning_size + and post_norm without learned scale. + """ + + def __init__( + self, hidden_size: int, self_conditioning_size: int, eps: float = 1e-6 + ): + super().__init__() + self.pre_norm = RMSNorm(hidden_size, eps=eps) + self.post_norm = RMSNorm(hidden_size, eps=eps, has_weight=False) + self.gate_proj = nn.Linear(hidden_size, self_conditioning_size, bias=False) + self.up_proj = nn.Linear(hidden_size, self_conditioning_size, bias=False) + self.down_proj = nn.Linear(self_conditioning_size, hidden_size, bias=False) + + def forward( + self, + inputs_embeds: torch.Tensor, + soft_embeds: torch.Tensor, + ) -> torch.Tensor: + x = self.pre_norm(soft_embeds) + sc_signal = self.down_proj( + F.gelu(self.gate_proj(x), approximate="tanh") * self.up_proj(x) + ) + return self.post_norm(inputs_embeds + sc_signal) + + +# --------------------------------------------------------------------------- +# Multimodal processing info (overrides Gemma4 config type check) +# --------------------------------------------------------------------------- + + +class DiffusionGemmaProcessingInfo(Gemma4ProcessingInfo): + """Processing info for DiffusionGemma. + + Overrides ``get_hf_config`` to accept ``DiffusionGemmaConfig`` + (which inherits from ``PretrainedConfig``, not ``Gemma4Config``). + Supports image and video modalities. + """ + + def get_hf_config(self): + # DiffusionGemmaConfig doesn't inherit from Gemma4Config, so we + # accept any PretrainedConfig here. + return self.ctx.get_hf_config() + + def get_supported_mm_limits(self) -> Mapping[str, int | None]: + # DiffusionGemma supports image and video inputs. + return {"image": None, "video": None} + + def get_mm_max_tokens_per_item( + self, seq_len: int, mm_counts: Mapping[str, int] + ) -> Mapping[str, int] | None: + return super().get_mm_max_tokens_per_item(seq_len, mm_counts) + + +@torch.compile(dynamic=True) +def _softcap_logits(logits: torch.Tensor, cap: float) -> torch.Tensor: + # fp32 before tanh for numerical stability (matches HF DiffusionGemma). + # Compiling fuses the cast/div/tanh/mul into one elementwise kernel over + # the [num_tokens, vocab] logits instead of four separate passes. + logits = logits.float() + return torch.tanh(logits / cap) * cap + + +@MULTIMODAL_REGISTRY.register_processor( + Gemma4MultiModalProcessor, + info=DiffusionGemmaProcessingInfo, + dummy_inputs=Gemma4DummyInputsBuilder, +) +class DiffusionGemmaForConditionalGeneration( + nn.Module, + SupportsMultiModal, + SupportsQuant, + SupportsPP, +): + """DiffusionGemma for vLLM. + + Single Gemma4 backbone that switches between encoder and decoder mode. + The encoder path uses standard Gemma4 layers (causal attention, KV write). + The decoder path uses the same weights with bidirectional attention and + KV read-only, plus self-conditioning. + + Always includes a vision tower (same as Gemma4) for image understanding. + + In practice, the model's forward() dispatches based on the `mode` kwarg + set by DiffusionGemmaModelState.prepare_inputs(). + """ + + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_prefix={ + "model.decoder.": "model.", + "model.encoder.language_model.": "model.", + "model.encoder.vision_tower.": "vision_tower.", + "model.encoder.embed_vision.": "embed_vision.", + }, + orig_to_new_substr={ + ".experts.": ".moe.experts.", + }, + ) + + packed_modules_mapping = { + "qkv_proj": ["q_proj", "k_proj", "v_proj"], + "gate_up_proj": ["gate_proj", "up_proj"], + } + + @staticmethod + def get_model_state_cls(): + return DiffusionGemmaModelState + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + config = vllm_config.model_config.hf_config + text_config = vllm_config.model_config.hf_text_config + self.config = config + self.model_dtype = vllm_config.model_config.dtype + + # DiffusionGemma's full-attention layers have NO v_proj — V is + # computed from k_proj's output (`value_states = key_states` before + # k_norm in `DiffusionGemmaDecoderTextAttention.forward`). This is + # the "k_eq_v" variant in our Gemma4 backbone. The checkpoint has no + # v_proj weights for full-attention layers; without this flag they + # would silently load with random V projections. + text_config.attention_k_eq_v = True + + # ---- Vision tower ---- + vision_config = getattr(config, "vision_config", None) + if vision_config is not None: + quant_config = vllm_config.quant_config + if quant_config and quant_config.get_name() in [ + "bitsandbytes", + "torchao", + "compressed-tensors", + ]: + tower_quant = quant_config + else: + quantizable = ( + vision_config.hidden_size % 64 == 0 + and vision_config.intermediate_size % 64 == 0 + ) + tower_quant = quant_config if quantizable else None + + with self._mark_tower_model(vllm_config, {"image", "video"}): + self.vision_tower = AutoModel.from_config(config=vision_config) + self.embed_vision = Gemma4MultimodalEmbedder( + vision_config, + text_config, + quant_config=tower_quant, + prefix=maybe_prefix(prefix, "embed_vision"), + ) + recursive_replace_linear( + self.vision_tower, + tower_quant, + prefix=maybe_prefix(prefix, "vision_tower"), + ) + else: + self.vision_tower = None + self.embed_vision = None + + # ---- Language backbone (Gemma4Model) ---- + # Use maybe_prefix to ensure correct weight name prefixes for + # quantization. The quantization config uses hf_to_vllm_mapper to + # match checkpoint weight names to model parameter names. + self.model = Gemma4Model( + vllm_config=vllm_config, + prefix=maybe_prefix(prefix, "model"), + ) + + self.lm_head = ParallelLMHead( + num_embeddings=text_config.vocab_size, + embedding_dim=text_config.hidden_size, + ) + + if text_config.tie_word_embeddings: + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens) + + # HF DiffusionGemma applies the final-logit softcap in fp32, before + # any other processing. Do it manually in `compute_logits` so the + # LogitsProcessor only handles the lm_head GEMM. + self.final_logit_softcapping = getattr( + text_config, "final_logit_softcapping", None + ) + self.logits_processor = LogitsProcessor( + text_config.vocab_size, + soft_cap=None, + ) + + sc_size = ( + getattr(config, "self_conditioning_size", None) + or text_config.intermediate_size + ) + self.self_conditioning = DiffusionGemmaSelfConditioning( + hidden_size=text_config.hidden_size, + self_conditioning_size=sc_size, + eps=getattr(text_config, "rms_norm_eps", 1e-6), + ) + + self.make_empty_intermediate_tensors = ( + self.model.make_empty_intermediate_tensors + ) + + def compute_self_conditioning( + self, + inputs_embeds: torch.Tensor, + probs: torch.Tensor, + ) -> torch.Tensor: + embed_weight = self.model.embed_tokens.weight + soft_embeds = torch.matmul( + probs.to(embed_weight.dtype), embed_weight + ) * self.model.normalizer.to(inputs_embeds.dtype) + return self.self_conditioning(inputs_embeds, soft_embeds) + + # ------------------------------------------------------------------ # + # Multimodal: reuse Gemma4's image parsing, processing & embedding + # ------------------------------------------------------------------ # + # The vision tower, pooler, embed_vision, and their processing logic + # are architecturally identical to Gemma4. Delegate to avoid + # maintaining a duplicate copy. + + _parse_and_validate_image_input = ( + Gemma4ForConditionalGeneration._parse_and_validate_image_input + ) + _parse_and_validate_video_input = ( + Gemma4ForConditionalGeneration._parse_and_validate_video_input + ) + _parse_and_validate_multimodal_inputs = ( + Gemma4ForConditionalGeneration._parse_and_validate_multimodal_inputs + ) + _encoder_chunk = staticmethod(Gemma4ForConditionalGeneration._encoder_chunk) + _process_image_input = Gemma4ForConditionalGeneration._process_image_input + _process_video_input = Gemma4ForConditionalGeneration._process_video_input + embed_multimodal = Gemma4ForConditionalGeneration.embed_multimodal + + def get_mm_mapping(self) -> MultiModelKeys: + """Get the module prefix mapping for multimodal models.""" + return MultiModelKeys.from_string_field( + language_model="model", + connector=["embed_vision"], + tower_model=["vision_tower"], + ) + + # ------------------------------------------------------------------ # + # Forward + # ------------------------------------------------------------------ # + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: Any | None = None, + inputs_embeds: torch.Tensor | None = None, + **kwargs: Any, + ) -> torch.Tensor: + if intermediate_tensors is not None: + inputs_embeds = None + return self.model( + input_ids=input_ids, + positions=positions, + intermediate_tensors=intermediate_tensors, + inputs_embeds=inputs_embeds, + **kwargs, + ) + + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None: + logits = self.logits_processor(self.lm_head, hidden_states) + if logits is not None and self.final_logit_softcapping is not None: + logits = _softcap_logits(logits, self.final_logit_softcapping) + return logits + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): + """Load weights from checkpoint. + + Checkpoint layout (HF DiffusionGemma): + model.encoder.vision_tower.* → vision tower + model.encoder.embed_vision.* → vision embedder + model.encoder.language_model.layers.* → backbone + model.decoder.layers.* → backbone (tied) + model.decoder.embed_tokens.* → embeddings + model.decoder.self_conditioning.* → self-conditioning MLP + lm_head.* → LM head (tied) + + We load encoder weights into our single ``Gemma4Model`` backbone, + skip duplicate decoder backbone weights, handle vision tower and + self-conditioning separately. + """ + + sc_params = dict( + (n, p) + for n, p in self.named_parameters() + if n.startswith("self_conditioning.") + ) + + # Collect vision tower + embedder parameters AND buffers for manual + # loading. The HF vision tower registers std_bias / std_scale as + # buffers (not parameters) when config.standardize is True, so we + # must include named_buffers() to avoid "not found in model" warnings. + vision_params: dict[str, torch.Tensor] = {} + for n, p in self.named_parameters(): + if n.startswith(("vision_tower.", "embed_vision.")): + vision_params[n] = p + for n, b in self.named_buffers(): + if n.startswith(("vision_tower.", "embed_vision.")): + vision_params[n] = b + + def _remap_weights(): + # Use full weight names (including suffixes like .weight_scale, + # .weight_packed) for dedup instead of just the base layer name. Critical + # for quantized checkpoints where each weight has multiple tensors; + # tracking only base names skips scales as duplicates. + seen_weights: set[str] = set() + for name, weight in weights: + # Self-conditioning lives under model.decoder.self_conditioning.* + # in the checkpoint but at self_conditioning.* in our model. + if "self_conditioning" in name: + sc_name = name.split("self_conditioning.", 1)[1] + sc_name = "self_conditioning." + sc_name + if sc_name in sc_params: + sc_params[sc_name].data.copy_(weight) + continue + + # Vision tower: model.encoder.vision_tower.* → vision_tower.* + # In HF, the vision tower is a sibling of language_model + # under the encoder module. + if name.startswith("model.encoder.vision_tower."): + vt_name = name[len("model.encoder.") :] + if vt_name in vision_params: + vision_params[vt_name].data.copy_(weight) + else: + logger.warning( + "Vision tower weight %s (mapped to %s) not found in model", + name, + vt_name, + ) + continue + + # Vision embedder: model.encoder.embed_vision.* → embed_vision.* + if name.startswith("model.encoder.embed_vision."): + ev_name = name[len("model.encoder.") :] + if ev_name in vision_params: + vision_params[ev_name].data.copy_(weight) + else: + logger.warning( + "Embed vision weight %s (mapped to %s) not found in model", + name, + ev_name, + ) + continue + + # Skip vestigial embed_vision.embedding weights. + if "embed_vision.embedding." in name: + continue + + # Encoder backbone → model.* + if name.startswith("model.encoder.language_model."): + name = name.replace("model.encoder.language_model.", "model.") + # Decoder backbone → model.* (skip exact duplicates) + elif name.startswith("model.decoder."): + name = name.replace("model.decoder.", "model.") + + # Skip only if we've seen the exact same weight name (including scales) + if name in seen_weights: + continue + seen_weights.add(name) + yield name, weight + + # Delegate to Gemma4ForCausalLM.load_weights for the backbone, + # which handles stacked params, MoE, k_eq_v, etc. + # Temporarily set self.config to text_config since Gemma4's + # load_weights expects it (e.g. tie_word_embeddings, layer_types). + from vllm.model_executor.models.gemma4 import Gemma4ForCausalLM + + saved_config = self.config + self.config = self.model.config + try: + Gemma4ForCausalLM.load_weights(self, _remap_weights()) + finally: + self.config = saved_config + + @classmethod + def get_placeholder_str(cls, modality: str, i: int) -> str | None: + if modality == "image": + return "" + if modality == "video": + return "<|video|>" + raise ValueError(f"Unsupported modality: {modality}") + + +@torch.compile(dynamic=True) +def _compute_num_rejected( + num_logits: torch.Tensor, + num_sampled: torch.Tensor, + query_start_loc: torch.Tensor, +) -> torch.Tensor: + query_lens = query_start_loc[1:] - query_start_loc[:-1] + num_rejected = num_logits - num_sampled + is_denoise = (num_logits > 0) & (num_sampled == 0) + return torch.where(is_denoise, query_lens, num_rejected) + + +@torch.compile(dynamic=True) +def _compiled_sample_step( + # Logits from the model [num_decode * CL, vocab] + logits: torch.Tensor, + # Request mapping + decode_slots: torch.Tensor, # [num_decode] int64 → slot indices + decode_idx: torch.Tensor, # [num_decode] int64 → position in num_reqs + all_slots: torch.Tensor, # [num_reqs] int64 → all slot indices + valid_canvas_len: torch.Tensor, # [num_decode] int64 → real canvas length (<=CL) + # State tensors (modified in-place) + canvas: torch.Tensor, # [max_num_reqs, CL] + argmax_canvas: torch.Tensor, # [max_num_reqs, CL] + step_tensor: torch.Tensor, # [max_num_reqs] + is_encoder_phase: torch.Tensor, # [max_num_reqs] + confident_tensor: torch.Tensor, # [max_num_reqs] + sc_embeds: torch.Tensor, # [max_num_reqs, CL, hidden] + embed_weight: torch.Tensor, # [vocab, hidden] + normalizer: torch.Tensor, + history: torch.Tensor, # [max_num_reqs, ST, CL] + history_len_tensor: torch.Tensor, # [max_num_reqs] + # Output tensors (modified in-place) + sampled: torch.Tensor, # [num_reqs, CL] + num_sampled: torch.Tensor, # [num_reqs] + draft_tokens: torch.Tensor, # [max_num_reqs, >=CL] + # Scalar config + max_denoising_steps: float, + t_min: float, + t_max: float, + confidence_threshold: float, + vocab_size: int, + CL: int, + ST: int, + # Sampler config + entropy_bound: float, +) -> torch.Tensor: + """Compiled decode step: temperature → Gumbel sample → probs/confidence → + accept/renoise → convergence, all as vectorized PyTorch ops. + + Returns the temperature-scaled logits ``[num_decode, CL, vocab]`` so the + caller can compute logprobs outside the compiled region.""" + num_decode = decode_slots.shape[0] + device = decode_slots.device + + # Clear outputs so prefill / non-decode slots report 0 (decode slots are + # overwritten below). + sampled.zero_() + num_sampled.zero_() + + # ---- Phase 1: Temperature schedule ---- + steps_f = step_tensor[decode_slots].float() + remaining = (max_denoising_steps - steps_f).clamp(min=1.0) + temp = t_min + (t_max - t_min) * (remaining / max_denoising_steps) + + # ---- Phase 2: Temperature scaling + Gumbel-max sampling ---- + logits_3d = logits.reshape(num_decode, CL, -1).float() + scaled = logits_3d / temp[:, None, None].clamp(min=1e-10) + + # Gumbel-max trick: argmax(logits/T + Gumbel) ~ sample from softmax(logits/T) + u = torch.rand_like(scaled).clamp(min=1e-20) + gumbel = -torch.log(-torch.log(u)) + # Zero noise when temp==0 (greedy) + noisy = scaled + gumbel * (temp[:, None, None] > 0).float() + new_tokens = noisy.view(-1, noisy.shape[-1]).argmax(dim=-1).view(num_decode, CL) + argmax_tokens = ( + scaled.view(-1, scaled.shape[-1]).argmax(dim=-1).view(num_decode, CL) + ) + + # ---- Phase 3: Probs, self-conditioning, confidence ---- + log_probs = scaled.log_softmax(dim=-1) + probs = log_probs.exp() + + token_entropy = -(probs * log_probs).sum(dim=-1) # [num_decode, CL] + # A canvas truncated near max_model_len is zero-padded up to CL by the + # caller; those padded rows are uniform (max entropy, argmax 0), so they + # never trigger early convergence and are stable, and only the real + # ``valid_canvas_len`` tokens are committed (num_sampled below). + mean_entropy = token_entropy.mean(dim=-1) # [num_decode] + confident_tensor[decode_slots] = mean_entropy < confidence_threshold + + # ---- Phase 4: Entropy-bound acceptance mask ---- + sorted_ent, sorted_idx = torch.sort(token_entropy, dim=-1) + cumsum_ent = torch.cumsum(sorted_ent, dim=-1) + cummax_ent = torch.cummax(sorted_ent, dim=-1).values + sorted_mask = (cumsum_ent - cummax_ent) <= entropy_bound + eb_mask = torch.zeros_like(sorted_mask) + eb_mask.scatter_(1, sorted_idx, sorted_mask) + + # ---- Phase 5: Post-sample ---- + is_commit = is_encoder_phase[decode_slots] # [num_decode] + is_denoise = ~is_commit + cur_step = step_tensor[decode_slots].float() + + # Step update: +1 for denoise, reset to 0 for commit + new_step_val = torch.where( + is_denoise, + (cur_step + 1).to(step_tensor.dtype), + step_tensor.new_zeros(num_decode), + ) + step_tensor[decode_slots] = new_step_val + + # Random tokens for renoise / canvas reinit + random_tokens = torch.randint( + 0, vocab_size, (num_decode, CL), device=device, dtype=canvas.dtype + ) + + # Compute denoise canvas (accept/renoise) + denoise_canvas = torch.where(eb_mask, new_tokens, random_tokens) + + # Canvas: commit → random reinit, denoise → accept/renoise result + canvas[decode_slots] = torch.where( + is_commit.unsqueeze(1), random_tokens, denoise_canvas + ) + + # History: write argmax_tokens for denoise requests at circular position + hist_len = history_len_tensor[decode_slots] + write_pos = hist_len % ST + for i in range(ST): + write_here = ((write_pos == i) & is_denoise).unsqueeze(1) + history[decode_slots, i] = torch.where( + write_here, argmax_tokens, history[decode_slots, i] + ) + + # Argmax canvas: update for denoise, preserve for commit + argmax_canvas[decode_slots] = torch.where( + is_denoise.unsqueeze(1), argmax_tokens, argmax_canvas[decode_slots] + ) + + # History length: increment for denoise, reset for commit + new_hist_len = torch.where(is_denoise, hist_len + 1, hist_len.new_zeros(num_decode)) + history_len_tensor[decode_slots] = new_hist_len + + # Sampled output: commit → emit argmax_canvas, denoise → 0 (pre-zeroed) + sampled[decode_idx] = argmax_canvas[decode_slots].to( + sampled.dtype + ) * is_commit.unsqueeze(1).to(sampled.dtype) + # Commit only the real canvas length (== CL except for a canvas truncated + # near max_model_len); the padded tail positions are never emitted. + num_sampled[decode_idx] = is_commit.to(num_sampled.dtype) * valid_canvas_len.to( + num_sampled.dtype + ) + + # ---- Phase 6: Stability + convergence ---- + ref = history[decode_slots, 0] + mismatch = torch.zeros(num_decode, device=device, dtype=torch.int32) + for h in range(1, ST): + mismatch = mismatch + (ref != history[decode_slots, h]).sum(dim=-1).int() + stable = mismatch == 0 + + step_after = step_tensor[decode_slots] + converged = (stable & confident_tensor[decode_slots] & (new_hist_len >= ST)) | ( + step_after >= max_denoising_steps + ) + # Commit done → denoise next (False); denoise converged → commit next (True) + is_encoder_phase[decode_slots] = torch.where( + is_commit, is_commit.new_zeros(num_decode), converged + ) + + # SC soft embedding: store ``probs @ embed_weight`` (the value the next step's + # self-conditioning MLP consumes) only for slots that will denoise next — i.e. + # this step denoised AND it isn't about to commit (is_encoder_phase now False). + # Masking here (rather than in the consumer) lets _apply_self_conditioning read + # sc_embeds directly. Storing the [.., hidden] soft embed instead of the full + # [.., vocab] probs avoids a giant persistent buffer. + sc_keep = (is_denoise & ~is_encoder_phase[decode_slots])[:, None, None] + soft_embeds = torch.matmul(probs.to(embed_weight.dtype), embed_weight) * normalizer + sc_embeds[decode_slots] = soft_embeds * sc_keep + + # Overwrite canvas with argmax for newly converged denoise requests + newly_converged = (converged & is_denoise).unsqueeze(1) + canvas[decode_slots] = torch.where( + newly_converged, argmax_canvas[decode_slots], canvas[decode_slots] + ) + + # ---- Phase 7: Copy canvas → draft_tokens for all slots ---- + draft_tokens[all_slots, :CL] = canvas[all_slots] + + return scaled + + +class DiffusionGemmaRequestStates: + """Pre-allocated GPU tensors for DiffusionGemma per-request state. + + Follows the indexed-slot pattern used by ``RequestState``. + """ + + def __init__( + self, + max_num_reqs: int, + canvas_length: int, + vocab_size: int, + max_denoising_steps: int, + device: torch.device, + hidden_size: int, + stability_threshold: int, + ): + self.max_num_reqs = max_num_reqs + self.canvas_length = canvas_length + self.vocab_size = vocab_size + self.max_denoising_steps = max_denoising_steps + self.stability_threshold = stability_threshold + self.device = device + + self.is_encoder_phase = torch.zeros( + max_num_reqs, dtype=torch.bool, device=device + ) + # Canvas tokens [max_num_reqs, canvas_length] + self.canvas = torch.zeros( + max_num_reqs, canvas_length, dtype=torch.int64, device=device + ) + # Step counter (counts up from 0 to max_denoising_steps) + self.step = torch.zeros( + max_num_reqs, + dtype=torch.int32, + device=device, + ) + # Accepted canvas history for stability check + self.accepted_canvas_history = torch.zeros( + max_num_reqs, + stability_threshold, + canvas_length, + dtype=torch.int64, + device=device, + ) + self.accepted_canvas_history_len = torch.zeros( + max_num_reqs, dtype=torch.int32, device=device + ) + # Latest argmax(processed_logits) per slot — what we COMMIT. + # NOT `current_canvas` (which is the post-renoise stochastic input for + # the next denoise step). We keep this separate from `canvas` because + # canvas gets renoised in-place during denoise, while argmax_canvas is + # the deterministic best-guess we ultimately emit. + self.argmax_canvas = torch.zeros( + max_num_reqs, canvas_length, dtype=torch.int64, device=device + ) + + # Per-slot prompt length (set by add_request). + self.prompt_len = torch.zeros( + max_num_reqs, + dtype=torch.int32, + device=device, + ) + + # Per-slot confidence flag, set by the sampler each step. + self.confident = torch.zeros(max_num_reqs, dtype=torch.bool, device=device) + + # Per-slot self-conditioning soft embedding (probs @ embed_weight) from + # the previous denoise step. Storing the [.., hidden] soft embed instead + # of the full [.., vocab] distribution shrinks this buffer by + # vocab/hidden (~170x) and moves the matmul to denoise time; the result + # is identical (SC consumes probs @ embed_weight anyway). + self.self_conditioning_embeds = torch.zeros( + max_num_reqs, canvas_length, hidden_size, dtype=torch.float32, device=device + ) + + def init_canvas(self, slot_indices_np: np.ndarray) -> None: + """Initialize canvas with random tokens for the given slots.""" + n = slot_indices_np.shape[0] + self.canvas[slot_indices_np] = torch.randint( + 0, + self.vocab_size, + (n, self.canvas_length), + dtype=torch.int64, + device=self.device, + ) + + def add_request(self, slot_idx: int) -> None: + self.is_encoder_phase[slot_idx] = True + self.init_canvas(torch.tensor([slot_idx], device=self.device)) + self.step[slot_idx] = 0 + self.accepted_canvas_history_len[slot_idx] = 0 + self.self_conditioning_embeds[slot_idx] = 0 + + def remove_request(self, slot_idx: int) -> None: + self.is_encoder_phase[slot_idx] = False + self.accepted_canvas_history_len[slot_idx] = 0 + self.self_conditioning_embeds[slot_idx] = 0 + + +class DiffusionGemmaModelState(ModelState): + """ModelState for DiffusionGemma. + + Single Gemma4 backbone in two modes: + - encoder mode (num_draft_tokens == 0): causal attention, writes KV + - decoder mode (num_draft_tokens > 0): bidirectional attention, reads KV + """ + + def __init__( + self, + vllm_config: VllmConfig, + model: nn.Module, + encoder_cache: Any, + device: torch.device, + ) -> None: + self.vllm_config = vllm_config + self.model_config = vllm_config.model_config + self.scheduler_config = vllm_config.scheduler_config + self.model = model + self.device = device + + self.supports_mm_inputs = encoder_cache is not None + self.max_num_reqs = self.scheduler_config.max_num_seqs + self.max_num_tokens = self.scheduler_config.max_num_batched_tokens + self.max_model_len = self.model_config.max_model_len + self.inputs_embeds_size = self.model_config.get_inputs_embeds_size() + self.dtype = self.model_config.dtype + + if self.supports_mm_inputs: + from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache + from vllm.v1.worker.gpu.mm.encoder_runner import EncoderRunner + + assert isinstance(encoder_cache, EncoderCache) + self.encoder_cache = encoder_cache + self.encoder_runner = EncoderRunner( + model=self.model, + max_num_tokens=self.max_num_tokens, + hidden_size=self.inputs_embeds_size, + encoder_cache=encoder_cache, + dtype=self.dtype, + device=self.device, + ) + + # Per-step MM data produced by get_mm_embeddings and consumed by + # prepare_inputs. Stored as raw (mm_embeds, is_mm_embed) so that + # prepare_inputs can call embed_input_ids directly into the + # persistent _inputs_embeds_buf, avoiding the intermediate copy + # through encoder_runner.inputs_embeds. + self._pending_mm_embeds: tuple[list[torch.Tensor], torch.Tensor] | None = None + + diffusion_config = vllm_config.diffusion_config + canvas_length = diffusion_config.canvas_length if diffusion_config else 32 + + text_config = self.model_config.hf_text_config + self.gen_config = self.model_config.try_get_generation_config() + max_denoising_steps = ( + diffusion_config.max_denoising_steps if diffusion_config else None + ) or self.gen_config.get("max_denoising_steps", 48) + self.diffusion_states = DiffusionGemmaRequestStates( + max_num_reqs=self.max_num_reqs, + canvas_length=canvas_length, + vocab_size=self.model_config.get_vocab_size(), + max_denoising_steps=max_denoising_steps, + device=device, + hidden_size=text_config.hidden_size, + stability_threshold=self.gen_config["stability_threshold"], + ) + self._req_id_to_index: dict[str, int] = {} + + # Persistent buffer for per-request causal flags, updated in-place + # so FULL CUDA graph replay sees the latest values. + self._causal_buf = torch.zeros( + self.max_num_reqs, dtype=torch.bool, device=device + ) + + # Persistent inputs_embeds buffer — required so FULL CUDA graph + # capture and runtime point at the SAME memory address. + # `prepare_dummy_inputs` (capture path) and `prepare_inputs` (runtime + # path) both must hand the captured graph a tensor at this address. + self._inputs_embeds_buf = torch.zeros( + self.max_num_tokens, + text_config.hidden_size, + dtype=self.model_config.dtype, + device=device, + ) + + def get_supported_generation_tasks(self): + return ("generate",) + + def custom_sampler(self, sampler: Any) -> tuple[Any, Any] | None: + diffusion_config = self.vllm_config.diffusion_config + gen = self.gen_config + sampler_cfg = gen.get("sampler_config") or {} + if "EntropyBound" not in sampler_cfg.get("_cls_name", ""): + raise ValueError("DiffusionGemma requires an EntropyBound sampler_config") + entropy_bound = sampler_cfg.get("entropy_bound") + if entropy_bound is None or entropy_bound <= 0: + raise ValueError( + f"entropy_bound must be a positive float (got {entropy_bound})" + ) + return DiffusionSampler( + sampler=sampler, + diffusion_config=diffusion_config, + vocab_size=self.model_config.get_vocab_size(), + diffusion_states=self.diffusion_states, + t_min=gen["t_min"], + t_max=gen["t_max"], + entropy_bound=entropy_bound, + confidence_threshold=gen["confidence_threshold"], + embed_weight=self.model.model.embed_tokens.weight, + normalizer=self.model.model.normalizer, + ), None + + def apply_staged_writes(self) -> None: + pass + + def add_request(self, req_index: int, new_req_data: Any) -> None: + self._req_id_to_index[new_req_data.req_id] = req_index + self.diffusion_states.add_request(req_index) + if not new_req_data.req_id.startswith("_warmup_"): + prompt_len = len(new_req_data.prompt_token_ids) + self.diffusion_states.prompt_len[req_index] = prompt_len + + def remove_request(self, req_id: str) -> None: + idx = self._req_id_to_index.pop(req_id, None) + if idx is not None: + self.diffusion_states.remove_request(idx) + + def get_mm_embeddings(self, scheduled_encoder_inputs, input_batch): + if not self.supports_mm_inputs: + return None + + mm_hashes, mm_kwargs = self.encoder_runner.prepare_mm_inputs( + scheduled_encoder_inputs + ) + if mm_kwargs: + encoder_outputs = self.encoder_runner.execute_mm_encoder(mm_kwargs) + self.encoder_cache.encoder_outputs.update(zip(mm_hashes, encoder_outputs)) + + mm_embeds, is_mm_embed = self.encoder_runner.gather_mm_embeddings( + input_batch.req_ids, + input_batch.num_tokens, + input_batch.num_scheduled_tokens, + input_batch.query_start_loc_np, + input_batch.prefill_len_np, + input_batch.num_computed_prefill_tokens_np, + ) + + if not mm_embeds: + # No MM tokens in this batch (e.g. all-decode step). + # prepare_inputs will use embed_input_ids (text-only) directly. + self._pending_mm_embeds = None + return None + + # Stash raw MM ingredients for prepare_inputs to merge directly + # into the persistent buffer, avoiding the intermediate copy + # through encoder_runner.inputs_embeds. + self._pending_mm_embeds = (mm_embeds, is_mm_embed) + return None + + def _apply_self_conditioning( + self, + decode_slots_np: np.ndarray, + decode_idx_np: np.ndarray, + query_start_loc_np: np.ndarray, + inputs_embeds: torch.Tensor, + sc_embeds: torch.Tensor, + ) -> None: + # One self-conditioning MLP call per decode request, over that request's + # query span [start, end) = its canvas. The span is the full canvas (CL) + # or, for the final canvas truncated near max_model_len, fewer than CL + # positions. sc_embeds already holds probs @ embed_weight from the prior + # denoise step, masked to zero by the sampler for slots not denoising + # this step; only the MLP runs here. CPU metadata -> no GPU syncs. + for slot, idx in zip(decode_slots_np.tolist(), decode_idx_np.tolist()): + start = int(query_start_loc_np[idx]) + end = int(query_start_loc_np[idx + 1]) + canvas = slice(start, end) + soft = sc_embeds[slot, : end - start] + inputs_embeds[canvas] = self.model.self_conditioning( + inputs_embeds[canvas], soft.to(inputs_embeds.dtype) + ) + + def prepare_inputs(self, input_batch, req_states) -> dict[str, Any]: + states = self.diffusion_states + num_tokens = input_batch.num_tokens + num_reqs = input_batch.num_reqs + + # Write into the PERSISTENT inputs_embeds buffer so FULL CUDA graph + # replay sees the latest values at the captured address. + num_tokens_padded = input_batch.num_tokens_after_padding + inputs_embeds = self._inputs_embeds_buf[:num_tokens_padded] + + # Populate embeddings: merge MM features when available, + # otherwise embed input_ids as text-only. + input_ids = input_batch.input_ids[:num_tokens] + if self._pending_mm_embeds is not None: + mm_embeds, is_mm_embed = self._pending_mm_embeds + self._pending_mm_embeds = None + inputs_embeds[:num_tokens].copy_( + self.model.embed_input_ids( + input_ids, + multimodal_embeddings=mm_embeds, + is_multimodal=is_mm_embed, + ) + ) + else: + inputs_embeds[:num_tokens].copy_(self.model.embed_input_ids(input_ids)) + + # Apply self-conditioning ONLY for denoising decode requests. + if input_batch.num_draft_tokens > 0 and self._req_id_to_index: + slots_np = input_batch.idx_mapping_np[:num_reqs] + num_logits_np = np.diff(input_batch.cu_num_logits_np[: num_reqs + 1]) + is_decode_indices_np = np.where(num_logits_np > 0)[0] + self._apply_self_conditioning( + slots_np[is_decode_indices_np], + is_decode_indices_np, + input_batch.query_start_loc_np, + inputs_embeds, + states.self_conditioning_embeds, + ) + + return {"inputs_embeds": inputs_embeds} + + def prepare_dummy_inputs(self, num_reqs: int, num_tokens: int) -> dict[str, Any]: + # CUDA graph capture path — return a slice of the SAME persistent + # inputs_embeds buffer that `prepare_inputs` writes to at runtime, + # so the captured graph and runtime point to identical addresses. + return {"inputs_embeds": self._inputs_embeds_buf[:num_tokens]} + + def postprocess_state(self, idx_mapping, num_sampled) -> None: + return None + + def prepare_attn( + self, + input_batch, + cudagraph_mode, + block_tables, + slot_mappings, + attn_groups, + kv_cache_config, + for_capture=False, + ) -> dict[str, Any]: + if cudagraph_mode == CUDAGraphMode.FULL: + num_reqs = input_batch.num_reqs_after_padding + num_tokens = input_batch.num_tokens_after_padding + else: + num_reqs = input_batch.num_reqs + num_tokens = input_batch.num_tokens + + query_start_loc_cpu = torch.from_numpy(input_batch.query_start_loc_np) + max_query_len = input_batch.num_scheduled_tokens.max().item() + + # Per-request causal mode: encoder (commit) = causal, + # denoise = bidirectional. Pass GPU tensor so the attention + # backend can handle mixed batches. + actual_num_reqs = input_batch.num_reqs + slots = input_batch.idx_mapping[:actual_num_reqs] + # Invariant: the sampler flips is_encoder_phase to False only after a + # request's FINAL prompt chunk, so a prompt spanning multiple chunks + # (longer than the token budget) stays causal for every chunk. + self._causal_buf[:actual_num_reqs] = self.diffusion_states.is_encoder_phase[ + slots + ] + if actual_num_reqs < num_reqs: + self._causal_buf[actual_num_reqs:num_reqs] = False + causal: bool | torch.Tensor = self._causal_buf[:num_reqs] + + return build_attn_metadata( + attn_groups=attn_groups, + num_reqs=num_reqs, + num_tokens=num_tokens, + query_start_loc_gpu=input_batch.query_start_loc, + query_start_loc_cpu=query_start_loc_cpu, + max_query_len=max_query_len, + seq_lens=input_batch.seq_lens, + max_seq_len=self.max_model_len, + block_tables=block_tables, + slot_mappings=slot_mappings, + kv_cache_config=kv_cache_config, + causal=causal, + ) + + num_new_sampled_tokens_per_step: int = 0 + + +# Penalty stub for the diffusion path: the runner reads +# penalties_state.output_bin_counts, and post_update treats None as +# "no penalty bookkeeping". +_NO_PENALTIES_STATE = SimpleNamespace(output_bin_counts=None) + + +class DiffusionSampler: + """Batched accept/renoise sampler for DiffusionGemma. + + Follows the same structure as ``vllm.v1.worker.gpu.sample.sampler.Sampler``: + decomposed into named methods, all GPU state in pre-allocated buffers, + no GPU→CPU syncs on the hot path. + """ + + def __init__( + self, + sampler: Any, + diffusion_config: Any, + vocab_size: int, + diffusion_states: DiffusionGemmaRequestStates | None = None, + *, + confidence_threshold: float, + t_min: float, + t_max: float, + entropy_bound: float, + embed_weight: torch.Tensor, + normalizer: torch.Tensor, + ): + self.sampling_states = sampler.sampling_states + self.req_states = sampler.req_states + # Self-conditioning soft embed = probs @ embed_weight * normalizer, + # computed in the sampler (see _compiled_sample_step). + self.embed_weight = embed_weight + self.normalizer = normalizer + self.canvas_length = ( + diffusion_config.canvas_length if diffusion_config is not None else 32 + ) + self.t_min = t_min + self.t_max = t_max + self.confidence_threshold = confidence_threshold + self.vocab_size = vocab_size + self.diffusion_states = diffusion_states + self.entropy_bound = entropy_bound + + max_num_reqs = diffusion_states.max_num_reqs + device = diffusion_states.device + self._sampled = torch.zeros( + max_num_reqs, + self.canvas_length, + dtype=torch.int32, + device=device, + ) + self._num_sampled = torch.zeros( + max_num_reqs, + dtype=torch.int32, + device=device, + ) + self._decode_slots = UvaBackedTensor(max_num_reqs, dtype=torch.int64) + self._decode_idx = UvaBackedTensor(max_num_reqs, dtype=torch.int64) + self._query_lens = UvaBackedTensor(max_num_reqs, dtype=torch.int32) + self._num_logits = UvaBackedTensor(max_num_reqs, dtype=torch.int32) + + # Per-slot stash for logprobs computed on the converging denoise step. + # Populated after the post-sample kernel detects convergence; consumed + # on the subsequent commit step when num_sampled=CANVAS_LEN. + self._pending_logprobs: dict[int, LogprobsTensors] = {} + + def add_request(self, req_idx: int, prompt_len: int, sampling_params: Any) -> None: + if use_penalty(sampling_params): + logger.warning_once( + "DiffusionGemma does not support repetition/frequency/presence " + "penalties; ignoring them for this request." + ) + # Purge any stale logprobs stashed under this slot by a prior request + # that was aborted between its converging denoise and commit steps. + self._pending_logprobs.pop(req_idx, None) + self.sampling_states.add_request(req_idx, sampling_params) + + def apply_staged_writes(self) -> None: + self.sampling_states.apply_staged_writes() + + @property + def penalties_state(self): + # Diffusion applies no penalties. The runner reads + # penalties_state.output_bin_counts, so expose a stub holding None; + # post_update treats None bin counts as "no penalty bookkeeping". + return _NO_PENALTIES_STATE + + # ------------------------------------------------------------------ + # Prefill + # ------------------------------------------------------------------ + + def _finish_prefills( + self, input_batch: Any, prefill_indices_np: np.ndarray + ) -> None: + """Transition requests whose prompt completes this step to denoising. + + Initializes their canvas, seeds draft tokens, and flips + is_encoder_phase to False. Mid-chunk requests (prompt longer than the + token budget) are left untouched so is_encoder_phase stays True and + prepare_attn keeps causal attention for their remaining chunks. + """ + states = self.diffusion_states + done_prefill_np = ( + input_batch.num_computed_prefill_tokens_np[prefill_indices_np] + + input_batch.num_scheduled_tokens[prefill_indices_np] + >= input_batch.prefill_len_np[prefill_indices_np] + ) + ps = input_batch.idx_mapping_np[prefill_indices_np[done_prefill_np]] + if len(ps) == 0: + return + states.init_canvas(ps) + self.req_states.draft_tokens[ps, : self.canvas_length] = states.canvas[ps] + ps_gpu = async_copy_to_gpu( + ps.astype(np.int64), device=states.is_encoder_phase.device + ) + states.is_encoder_phase.index_fill_(0, ps_gpu, False) + + def _handle_prefill( + self, + input_batch: Any, + device: torch.device, + ) -> SamplerOutput: + num_reqs = input_batch.num_reqs + self._finish_prefills(input_batch, np.arange(num_reqs)) + sampled = self._sampled[:num_reqs, :1] + sampled.zero_() + num_sampled = self._num_sampled[:num_reqs] + num_sampled.zero_() + return SamplerOutput( + sampled_token_ids=sampled, + logprobs_tensors=None, + num_nans=None, + num_sampled=num_sampled, + num_rejected=num_sampled, + ) + + # ------------------------------------------------------------------ + # Decode helpers + # ------------------------------------------------------------------ + + def _build_output( + self, + input_batch: Any, + sampled: torch.Tensor, + num_sampled: torch.Tensor, + per_req_nlogits_np: np.ndarray, + device: torch.device, + logprobs_tensors: LogprobsTensors | None = None, + ) -> SamplerOutput: + """Compute num_rejected and build SamplerOutput.""" + num_reqs = input_batch.num_reqs + + self._query_lens.np[:num_reqs] = np.diff( + input_batch.query_start_loc_np[: num_reqs + 1] + ) + self._num_logits.np[:num_reqs] = per_req_nlogits_np + self._query_lens.copy_to_uva() + self._num_logits.copy_to_uva() + + num_rejected = _compute_num_rejected( + self._num_logits.gpu[:num_reqs], + num_sampled, + input_batch.query_start_loc[: num_reqs + 1], + ) + + return SamplerOutput( + sampled_token_ids=sampled, + logprobs_tensors=logprobs_tensors, + num_nans=None, + num_sampled=num_sampled, + num_rejected=num_rejected, + ) + + # ------------------------------------------------------------------ + # Main entry point + # ------------------------------------------------------------------ + + def __call__( + self, + logits: torch.Tensor, + input_batch: Any, + draft_logits: torch.Tensor | None = None, + ) -> SamplerOutput: + num_reqs = input_batch.num_reqs + device = logits.device + + if input_batch.num_draft_tokens == 0: + return self._handle_prefill(input_batch, device) + + # --- CPU/NumPy setup (outside compile): split decode vs prefill, init + # canvas for any new prefills, and stage decode slot indices to GPU. --- + states = self.diffusion_states + CL = self.canvas_length + slots_np = input_batch.idx_mapping_np[:num_reqs] + per_req_nlogits_np = np.diff(input_batch.cu_num_logits_np[: num_reqs + 1]) + + decode_indices_np = np.where(per_req_nlogits_np > 0)[0] + prefill_indices_np = np.where(per_req_nlogits_np == 0)[0] + decode_slots_np = slots_np[decode_indices_np] + + if len(prefill_indices_np) > 0: + self._finish_prefills(input_batch, prefill_indices_np) + + num_decode = len(decode_indices_np) + self._decode_slots.np[:num_decode] = decode_slots_np + self._decode_idx.np[:num_decode] = decode_indices_np + self._decode_slots.copy_to_uva() + self._decode_idx.copy_to_uva() + decode_slots = self._decode_slots.gpu[:num_decode] + decode_idx = self._decode_idx.gpu[:num_decode] + + # Real canvas length per decode request. Equals CL except when a canvas + # was truncated near max_model_len, in which case the scheduler gave us + # fewer than CL logits for that request. + valid_canvas_len_np = per_req_nlogits_np[per_req_nlogits_np > 0] + valid_canvas_len = async_copy_to_gpu( + valid_canvas_len_np.astype(np.int64), device=device + ) + + # Pad any truncated canvas back to CL so the uniform-CL sampler math + # holds. Phantom (padded) positions are zeroed → uniform logits → high + # entropy (no premature convergence) and argmax 0 (stable); they are + # never committed (num_sampled == real length). + if num_decode > 0 and valid_canvas_len_np.min() < CL: + ar = torch.arange(CL, device=device) + starts = valid_canvas_len.cumsum(0) - valid_canvas_len # row offset per req + valid = ar.unsqueeze(0) < valid_canvas_len.unsqueeze(1) # [num_decode, CL] + src = (starts.unsqueeze(1) + ar.unsqueeze(0)).clamp_max(logits.shape[0] - 1) + logits = logits[src.reshape(-1)] * valid.reshape(-1, 1).to(logits.dtype) + + # Cleared inside _compiled_sample_step so prefill/non-decode slots stay 0. + sampled = self._sampled[:num_reqs] + num_sampled = self._num_sampled[:num_reqs] + + all_slots = input_batch.idx_mapping[:num_reqs] + + # Snapshot which slots are committing BEFORE the compiled step runs, + # since it mutates is_encoder_phase (commit→False, converge→True). + is_committing = states.is_encoder_phase[decode_slots].clone() + + # --- Single compiled call: temp → sample → probs → post-process --- + scaled = _compiled_sample_step( + logits, + decode_slots, + decode_idx, + all_slots, + valid_canvas_len, + # State + states.canvas, + states.argmax_canvas, + states.step, + states.is_encoder_phase, + states.confident, + states.self_conditioning_embeds, + self.embed_weight, + self.normalizer, + states.accepted_canvas_history, + states.accepted_canvas_history_len, + # Output + sampled, + num_sampled, + self.req_states.draft_tokens, + # Config + max_denoising_steps=float(states.max_denoising_steps), + t_min=self.t_min, + t_max=self.t_max, + confidence_threshold=self.confidence_threshold, + vocab_size=self.vocab_size, + CL=self.canvas_length, + ST=states.stability_threshold, + entropy_bound=self.entropy_bound, + ) + + # --- Logprobs: stash on convergence, return on commit --- + slots_np = input_batch.idx_mapping_np[:num_reqs] + is_decode_np = per_req_nlogits_np > 0 + + logprobs_tensors = None + max_num_logprobs = self.sampling_states.max_num_logprobs(slots_np) + if max_num_logprobs >= 0: + # Denoise steps that just converged: the compiled step flipped + # is_encoder_phase from False→True. Detect as slots where + # is_encoder_phase is now True but is_committing was False. + converged_mask = states.is_encoder_phase[decode_slots] + just_converged = converged_mask & ~is_committing + if just_converged.any(): + flat_logits = scaled.reshape(-1, scaled.shape[-1]) + argmax_tokens = scaled.argmax(dim=-1) + for local_idx in just_converged.nonzero(as_tuple=True)[0]: + li = local_idx.item() + slot = decode_slots[local_idx] + # Stash only the real canvas positions (== CL unless this + # canvas was truncated near max_model_len); padded tail + # positions are never emitted. + k_i = int(valid_canvas_len_np[li]) + start = li * CL + self._pending_logprobs[slot.item()] = compute_topk_logprobs( + flat_logits[start : start + k_i], + max_num_logprobs, + argmax_tokens[local_idx][:k_i], + ) + + # Commit steps: is_committing was True at entry. Reassemble + # previously stashed logprobs and attach to SamplerOutput. + if is_committing.any() and self._pending_logprobs: + parts_ids, parts_lp, parts_ranks = [], [], [] + cu_gen: list[int] = [] + flat_offset = 0 + for i in range(num_reqs): + cu_gen.append(flat_offset) + slot = int(slots_np[i]) + if is_decode_np[i] and slot in self._pending_logprobs: + lp = self._pending_logprobs.pop(slot) + parts_ids.append(lp.logprob_token_ids) + parts_lp.append(lp.logprobs) + parts_ranks.append(lp.selected_token_ranks) + flat_offset += lp.logprobs.shape[0] + if parts_ids: + logprobs_tensors = LogprobsTensors( + logprob_token_ids=torch.cat(parts_ids), + logprobs=torch.cat(parts_lp), + selected_token_ranks=torch.cat(parts_ranks), + cu_num_generated_tokens=cu_gen, + ) + + return self._build_output( + input_batch, + sampled, + num_sampled, + per_req_nlogits_np, + device, + logprobs_tensors=logprobs_tensors, + ) diff --git a/vllm/model_executor/models/gemma4.py b/vllm/model_executor/models/gemma4.py index 45e82c26d95..03e67c4ada7 100644 --- a/vllm/model_executor/models/gemma4.py +++ b/vllm/model_executor/models/gemma4.py @@ -725,10 +725,8 @@ class Gemma4DecoderLayer(nn.Module): 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) + router_logits = self.router(residual) hidden_states_2 = self.moe(hidden_states_2, router_logits) hidden_states_2 = self.post_feedforward_layernorm_2(hidden_states_2) diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index 175f0f2dab2..722ba93d393 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -400,6 +400,10 @@ _MULTIMODAL_MODELS = { "gemma3n_mm", "Gemma3nForConditionalGeneration", ), + "DiffusionGemmaForBlockDiffusion": ( + "diffusion_gemma", + "DiffusionGemmaForConditionalGeneration", + ), "Gemma4ForConditionalGeneration": ("gemma4_mm", "Gemma4ForConditionalGeneration"), "Gemma4UnifiedForConditionalGeneration": ( "gemma4_unified", diff --git a/vllm/tool_parsers/gemma4_tool_parser.py b/vllm/tool_parsers/gemma4_tool_parser.py index 9925284273f..a92ab9bb6cd 100644 --- a/vllm/tool_parsers/gemma4_tool_parser.py +++ b/vllm/tool_parsers/gemma4_tool_parser.py @@ -20,9 +20,11 @@ import json from collections.abc import Sequence import regex as re +from openai.types.responses import ToolChoiceFunction from vllm.entrypoints.chat_utils import make_tool_call_id from vllm.entrypoints.openai.chat_completion.protocol import ( + ChatCompletionNamedToolChoiceParam, ChatCompletionRequest, ) from vllm.entrypoints.openai.engine.protocol import ( @@ -343,6 +345,9 @@ class Gemma4ToolParser(ToolParser): tool parsers. """ + # Gemma4 emits native special-token tool calls, not generic JSON calls. + supports_required_and_named = False + def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None): super().__init__(tokenizer, tools) @@ -390,6 +395,23 @@ class Gemma4ToolParser(ToolParser): def adjust_request( self, request: ChatCompletionRequest | ResponsesRequest ) -> ChatCompletionRequest | ResponsesRequest: + if request.tools: + tc = request.tool_choice + if tc == "required" or isinstance( + tc, + (ChatCompletionNamedToolChoiceParam, ToolChoiceFunction), + ): + # Do NOT call super().adjust_request() for required/named tool + # choice. The base implementation injects a JSON-array + # `structured_outputs` schema and forces xgrammar guided + # decoding, which conflicts with Gemma4's native + # `<|tool_call>call:...` (non-JSON) tool syntax and crashes + # EngineCore under MTP spec decode. The streaming/extraction + # parser already handles the native output, so guided decoding + # is skipped here (mirrors the GLM4 precedent). + if request.tool_choice != "none": + request.skip_special_tokens = False + return request request = super().adjust_request(request) if request.tools and request.tool_choice != "none": # Don't skip special tokens — <|tool_call> etc. are needed for @@ -549,22 +571,40 @@ class Gemma4ToolParser(ToolParser): 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 + # Case 2: One or more new tool calls started in this delta. + # A single delta can batch several complete calls, so advance the + # tool id once per newly-seen start token and allocate a tracking + # slot for each. + if start_count > prev_start_count: + num_new = start_count - prev_start_count + for _ in range(num_new): + self.current_tool_id += 1 + self.streamed_args_for_tool.append("") + self.prev_tool_call_arr.append({}) 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): + logger.debug( + "Started %d new tool call(s); current_tool_id=%d", + num_new, + self.current_tool_id, + ) + # Don't return yet if this delta also contains call payload or + # the end marker; backends can batch one or more complete tool + # calls into a single streaming chunk. Only wait for more text + # when the delta is just the start token itself. + if start_count > end_count and len(delta_text) <= len( + self.tool_call_start_token + ): return None - # Case 3: Tool call just ended + # Case 3: One or more tool calls just ended (possibly several in a + # single batched delta) — drain every newly-completed call. if end_count > prev_end_count: - return self._handle_tool_call_end(current_text) + return self._handle_tool_call_end( + current_text, + prev_end_count=prev_end_count, + end_count=end_count, + start_count=start_count, + ) # Case 4: In the middle of a tool call — parse partial content if start_count > end_count: @@ -652,45 +692,111 @@ class Gemma4ToolParser(ToolParser): return None - def _handle_tool_call_end(self, current_text: str) -> DeltaMessage | None: - """Handle streaming when a tool call has just completed. + def _handle_tool_call_end( + self, + current_text: str, + prev_end_count: int, + end_count: int, + start_count: int, + ) -> DeltaMessage | None: + """Handle streaming when one or more tool calls have just completed. - Performs a final parse of the complete tool call and flushes - any remaining un-streamed argument fragments. + A single streaming delta can batch several complete tool calls + (``<|tool_call>...<|tool_call>...``). Every + call whose ```` end marker arrived in this delta — i.e. + those with index in ``[prev_end_count, end_count)`` — is drained and + emitted, with one ``DeltaToolCall`` per call in a single + ``DeltaMessage`` (this matches the OpenAI streaming wire format, and + the serving layer iterates over ``delta.tool_calls``). + + Per call: + + * If the function name was already streamed incrementally (the + token-by-token path), only the remaining argument fragment is + flushed as a diff. + * If the call is seen complete for the first time in this delta (the + batched-complete path), the id + name + full arguments JSON are + emitted exactly once. """ - 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, - ) + # Parse the complete tool calls using regex for accuracy. + all_matches = self.tool_call_regex.findall(current_text) + if not all_matches: + logger.debug("Tool call end detected but no complete tool call parsed yet.") 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] + deltas: list[DeltaToolCall] = [] + for idx in range(prev_end_count, end_count): + if idx >= len(all_matches): + break + # Ensure the tracking arrays have a slot for this index (defensive; + # Case 2 normally allocates these when the start token arrives). + while len(self.prev_tool_call_arr) <= idx: + self.prev_tool_call_arr.append({}) + self.streamed_args_for_tool.append("") + + func_name, args_str = all_matches[idx] 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 + # The name is sent exactly once per call. We track that via the + # per-call entry in prev_tool_call_arr (set either by the middle + # path or by the batched-complete branch below), which is robust + # even when several calls are drained in one delta. + name_already_sent = bool(self.prev_tool_call_arr[idx].get("name")) - return DeltaMessage( - tool_calls=[ + if not name_already_sent: + # Batched-complete call: emit id + name + full arguments once. + self.streamed_args_for_tool[idx] = final_args_json + self.prev_tool_call_arr[idx] = { + "name": func_name, + "arguments": final_args, + } + deltas.append( + DeltaToolCall( + index=idx, + type="function", + id=make_tool_call_id(), + function=DeltaFunctionCall( + name=func_name, arguments=final_args_json + ).model_dump(exclude_none=True), + ) + ) + else: + # Incrementally-streamed call: flush the remaining argument + # tail that was withheld during the middle phase. + prev_streamed = self.streamed_args_for_tool[idx] + if len(final_args_json) > len(prev_streamed): + diff = final_args_json[len(prev_streamed) :] + self.streamed_args_for_tool[idx] = final_args_json + self.prev_tool_call_arr[idx]["arguments"] = final_args + deltas.append( DeltaToolCall( - index=self.current_tool_id, + index=idx, function=DeltaFunctionCall(arguments=diff).model_dump( exclude_none=True ), ) - ] - ) + ) + # Advance streaming state past the calls completed in this delta. If a + # further tool call is still being accumulated (start without a + # matching end), point current_tool_id at it so the middle path can + # stream its arguments next; otherwise settle on the last completed + # call. + if start_count > end_count: + self.current_tool_id = end_count + while len(self.prev_tool_call_arr) <= self.current_tool_id: + self.prev_tool_call_arr.append({}) + self.streamed_args_for_tool.append("") + self.current_tool_name_sent = bool( + self.prev_tool_call_arr[self.current_tool_id].get("name") + ) + else: + self.current_tool_id = end_count - 1 + self.current_tool_name_sent = True + + if deltas: + return DeltaMessage(tool_calls=deltas) return None def _emit_argument_diff(self, raw_args_str: str) -> DeltaMessage | None: diff --git a/vllm/transformers_utils/config.py b/vllm/transformers_utils/config.py index 427f30b3992..3edfe932e0c 100644 --- a/vllm/transformers_utils/config.py +++ b/vllm/transformers_utils/config.py @@ -87,6 +87,7 @@ _CONFIG_REGISTRY: dict[str, type[PretrainedConfig]] = LazyConfigDict( ops_colqwen3="OpsColQwen3Config", qwen3_vl_nemotron_embed="Qwen3VLNemotronEmbedConfig", cosmos3_omni="Cosmos3Config", + diffusion_gemma="DiffusionGemmaConfig", deepseek_vl_v2="DeepseekVLV2Config", deepseek_v32="DeepseekV3Config", deepseek_v4="DeepseekV4Config", diff --git a/vllm/transformers_utils/configs/__init__.py b/vllm/transformers_utils/configs/__init__.py index 71f7723e4c8..e91f89b2d09 100644 --- a/vllm/transformers_utils/configs/__init__.py +++ b/vllm/transformers_utils/configs/__init__.py @@ -26,6 +26,8 @@ _CLASS_TO_MODULE: dict[str, str] = { "OpsColQwen3Config": "vllm.transformers_utils.configs.colqwen3", "Qwen3VLNemotronEmbedConfig": "vllm.transformers_utils.configs.colqwen3", "Cosmos3Config": "vllm.transformers_utils.configs.cosmos3", + "DiffusionGemmaConfig": "vllm.transformers_utils.configs.diffusion_gemma", + "DiffusionGemmaTextConfig": "vllm.transformers_utils.configs.diffusion_gemma", "DeepseekVLV2Config": "vllm.transformers_utils.configs.deepseek_vl2", "DeepseekV4Config": "vllm.transformers_utils.configs.deepseek_v4", "DotsOCRConfig": "vllm.transformers_utils.configs.dotsocr", @@ -97,6 +99,8 @@ __all__ = [ "OpsColQwen3Config", "Qwen3VLNemotronEmbedConfig", "Cosmos3Config", + "DiffusionGemmaConfig", + "DiffusionGemmaTextConfig", "DeepseekVLV2Config", "DeepseekV3Config", "DeepseekV4Config", diff --git a/vllm/transformers_utils/configs/diffusion_gemma.py b/vllm/transformers_utils/configs/diffusion_gemma.py new file mode 100644 index 00000000000..246a25b32c6 --- /dev/null +++ b/vllm/transformers_utils/configs/diffusion_gemma.py @@ -0,0 +1,44 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from typing import Any + +from transformers import PretrainedConfig +from transformers.models.gemma4.configuration_gemma4 import Gemma4VisionConfig + + +def _init_text_config(self: PretrainedConfig, **kwargs: Any) -> None: + PretrainedConfig.__init__(self, **kwargs) + # DiffusionGemma always uses MoE and K=V sharing for full_attention + # layers. The HF reference removed these config fields entirely. + if getattr(self, "num_experts", None): + self.enable_moe_block = True + self.attention_k_eq_v = True + + +class DiffusionGemmaTextConfig(PretrainedConfig): + model_type = "diffusion_gemma_text" + + def __init__(self, **kwargs: Any): + _init_text_config(self, **kwargs) + + +class DiffusionGemmaConfig(PretrainedConfig): + model_type = "diffusion_gemma" + + def __init__( + self, + text_config: dict[str, Any] | None = None, + canvas_length: int = 256, + self_conditioning_size: int | None = None, + **kwargs: Any, + ): + self.text_config = DiffusionGemmaTextConfig(**(text_config or {})) + self.canvas_length = canvas_length + self.self_conditioning_size = self_conditioning_size + vision_config = kwargs.pop("vision_config", None) + if isinstance(vision_config, dict): + self.vision_config = Gemma4VisionConfig(**vision_config) + else: + self.vision_config = vision_config + self.audio_config = None + PretrainedConfig.__init__(self, **kwargs) diff --git a/vllm/transformers_utils/model_arch_config_convertor.py b/vllm/transformers_utils/model_arch_config_convertor.py index 250aee50378..37402dcaa0b 100644 --- a/vllm/transformers_utils/model_arch_config_convertor.py +++ b/vllm/transformers_utils/model_arch_config_convertor.py @@ -582,6 +582,7 @@ MODEL_ARCH_CONFIG_CONVERTORS = { "cohere_asr": CohereAsrModelArchConfigConvertor, "dbrx": DbrxModelArchConfigConvertor, "deepseek_mtp": DeepSeekMTPModelArchConfigConvertor, + "diffusion_gemma_text": Gemma4ModelArchConfigConvertor, "ernie_mtp": ErnieMTPModelArchConfigConvertor, "falcon": FalconModelArchConfigConvertor, "falcon_mamba": MambaModelArchConfigConvertor, diff --git a/vllm/v1/attention/backend.py b/vllm/v1/attention/backend.py index 32b4b8ab9a0..152178ec2b3 100644 --- a/vllm/v1/attention/backend.py +++ b/vllm/v1/attention/backend.py @@ -387,7 +387,7 @@ class CommonAttentionMetadata: block_table_tensor: torch.Tensor slot_mapping: torch.Tensor - causal: bool = True + causal: bool | torch.Tensor = True # Needed by FastPrefillAttentionBuilder logits_indices_padded: torch.Tensor | None = None @@ -497,7 +497,9 @@ class CommonAttentionMetadata: max_seq_len=self.max_seq_len, block_table_tensor=self.block_table_tensor[:num_actual_reqs], slot_mapping=self.slot_mapping[:num_actual_tokens], - causal=self.causal, + causal=self.causal[:num_actual_reqs] + if isinstance(self.causal, torch.Tensor) + else self.causal, logits_indices_padded=self.logits_indices_padded, num_logits_indices=self.num_logits_indices, encoder_seq_lens=maybe_slice_reqs(self.encoder_seq_lens), diff --git a/vllm/v1/attention/backends/fa_utils.py b/vllm/v1/attention/backends/fa_utils.py index 0d6a3d298b6..474523780ff 100644 --- a/vllm/v1/attention/backends/fa_utils.py +++ b/vllm/v1/attention/backends/fa_utils.py @@ -131,6 +131,12 @@ def get_flash_attn_version( and head_size != head_size_v ): upgrade_reason = "Diff-KV with sinks" + elif ( + vllm_config is not None + and vllm_config.model_config is not None + and vllm_config.model_config.is_diffusion + ): + upgrade_reason = "Per-sequence causal (dynamic_causal) requires FA4" if upgrade_reason: logger.info_once( "%s: upgrading FlashAttention 3 -> 4", diff --git a/vllm/v1/attention/backends/flash_attn.py b/vllm/v1/attention/backends/flash_attn.py index d6774a6eb99..9e33c0d823b 100755 --- a/vllm/v1/attention/backends/flash_attn.py +++ b/vllm/v1/attention/backends/flash_attn.py @@ -267,7 +267,7 @@ class FlashAttentionMetadata: prefix_scheduler_metadata: torch.Tensor | None = None max_num_splits: int = 0 - causal: bool = True + causal: bool | torch.Tensor = True # PrefixLM bidirectional ranges for multimodal tokens. # Shape: (num_seqs, max_ranges, 2) int32, [start, end] per range. @@ -570,6 +570,9 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad self.scheduler_metadata[n:] = 0 scheduler_metadata = self.scheduler_metadata[:n] + if isinstance(causal, torch.Tensor) and causal.dtype != torch.int32: + causal = causal.to(torch.int32) + attn_metadata = FlashAttentionMetadata( num_actual_tokens=num_actual_tokens, max_query_len=max_query_len, @@ -824,18 +827,46 @@ class FlashAttentionImpl(AttentionImpl): if self.sliding_window is not None else None ) + + causal = attn_metadata.causal + is_dynamic_causal = isinstance(causal, torch.Tensor) + + # For non-causal (bidirectional) attention, make the + # sliding window symmetric so queries attend in both + # directions. + if ( + sliding_window_size is not None + and sliding_window_size[1] == 0 + and (is_dynamic_causal or causal is False) + ): + sliding_window_size = [ + sliding_window_size[0], + sliding_window_size[0], + ] + mm_prefix_ranges = attn_metadata.mm_prefix_range_tensor mm_mask_mod = None mm_aux = None if ( mm_prefix_ranges is not None - and attn_metadata.causal + and not is_dynamic_causal + and causal is True and self.vllm_flash_attn_version == 4 ): max_ranges = mm_prefix_ranges.shape[1] mm_mask_mod = _make_mm_prefix_mask_mod(max_ranges) mm_aux = [mm_prefix_ranges] + dynamic_causal = None + if isinstance(causal, torch.Tensor): + if self.vllm_flash_attn_version != 4: + raise NotImplementedError( + "Per-sequence causal requires FA4. Current version: " + f"FA{self.vllm_flash_attn_version}" + ) + dynamic_causal = causal + causal = False + flash_attn_varlen_func( q=query[:num_actual_tokens], k=key_cache, @@ -846,7 +877,7 @@ class FlashAttentionImpl(AttentionImpl): seqused_k=seqused_k, max_seqlen_k=max_seqlen_k, softmax_scale=self.scale, - causal=attn_metadata.causal, + causal=causal, alibi_slopes=self.alibi_slopes, window_size=sliding_window_size, block_table=block_table, @@ -856,6 +887,7 @@ class FlashAttentionImpl(AttentionImpl): q_descale=q_descale, k_descale=k_descale, v_descale=v_descale, + dynamic_causal=dynamic_causal, num_splits=attn_metadata.max_num_splits, s_aux=self.sinks, mask_mod=mm_mask_mod, diff --git a/vllm/v1/attention/backends/triton_attn.py b/vllm/v1/attention/backends/triton_attn.py index 92ff08cc0f3..377e9e7ab1d 100644 --- a/vllm/v1/attention/backends/triton_attn.py +++ b/vllm/v1/attention/backends/triton_attn.py @@ -79,6 +79,8 @@ class TritonAttentionMetadata: softmax_segm_max: torch.Tensor softmax_segm_expsum: torch.Tensor + causal: bool | torch.Tensor + # For cascade attention. use_cascade: bool common_prefix_len: int @@ -219,6 +221,7 @@ class TritonAttentionMetadataBuilder(AttentionMetadataBuilder[TritonAttentionMet seq_lens=seq_lens, block_table=block_table_tensor, slot_mapping=slot_mapping, + causal=common_attn_metadata.causal, use_cascade=use_cascade, common_prefix_len=common_prefix_len, cu_prefix_query_lens=cu_prefix_query_lens, @@ -271,6 +274,10 @@ class TritonAttentionBackend(AttentionBackend): forward_includes_kv_cache_update: bool = False + @classmethod + def supports_non_causal(cls) -> bool: + return True + @staticmethod def get_name() -> str: return "TRITON_ATTN" @@ -619,7 +626,7 @@ class TritonAttentionImpl(AttentionImpl): seqused_k=seqused_k, max_seqlen_k=max_seqlen_k, softmax_scale=self.scale, - causal=True, + causal=attn_metadata.causal, alibi_slopes=self.alibi_slopes, use_alibi_sqrt=self.use_alibi_sqrt, window_size=self.sliding_window, diff --git a/vllm/v1/attention/ops/triton_attention_helpers.py b/vllm/v1/attention/ops/triton_attention_helpers.py index 6ed50f6a2df..ed9a38ad6cd 100644 --- a/vllm/v1/attention/ops/triton_attention_helpers.py +++ b/vllm/v1/attention/ops/triton_attention_helpers.py @@ -153,6 +153,8 @@ def compute_tile_loop_bounds( SLIDING_WINDOW: tl.constexpr, USE_MM_PREFIX: tl.constexpr, IS_3D: tl.constexpr, + USE_CAUSAL: tl.constexpr = True, + USE_PER_SEQ_CAUSAL: tl.constexpr = False, CHUNK_LOOKBACK: tl.constexpr = -1, CHUNK_SIZE: tl.constexpr = -1, ): @@ -163,10 +165,11 @@ def compute_tile_loop_bounds( 1. Longest prefix spanned by any query token in this q-block. Clamped to ``seq_len`` (causal) or extended to it when - mm_prefix is active (bidirectional ranges can reach past the - causal prefix). + mm_prefix is active or non-causal sequences need the full + sequence. 2. Sliding-window pruning: narrows ``[tile_start, tile_end)`` to only tiles that can contain an allowed key under SWA. + For non-causal sequences, the window extends in both directions. 3. 3D scoping: when ``IS_3D`` is True, further narrows to the segment's slice via ``(segm_idx * tiles_per_segment, (segm_idx + 1) * tiles_per_segment)``. @@ -179,9 +182,10 @@ def compute_tile_loop_bounds( + (BLOCK_M - 1) // num_queries_per_kv + 1 ) - if USE_MM_PREFIX: - # image bidirectional attention ranges require a full range - # including q_block padding to make sure doc mask is correct + if USE_MM_PREFIX or USE_PER_SEQ_CAUSAL or (not USE_CAUSAL): + # Non-causal or mixed batches need the full sequence range. + # Per-element masking in compute_kv_seq_mask handles the + # actual causal/non-causal boundary per sequence. max_seq_prefix_len = tl.maximum(max_seq_prefix_len, seq_len) else: max_seq_prefix_len = tl.minimum(max_seq_prefix_len, seq_len) @@ -207,12 +211,17 @@ def compute_tile_loop_bounds( # [context_len + qpos_lo - SLIDING_WINDOW + 1, context_len + qpos_hi] q_abs = context_len + qpos_lo if CHUNK_LOOKBACK > -1: - # Chunked attention: align lower bound to the start of the - # lookback'th previous chunk. first_allowed_key = ((q_abs // CHUNK_SIZE) - CHUNK_LOOKBACK) * CHUNK_SIZE else: first_allowed_key = q_abs - SLIDING_WINDOW + 1 - last_allowed_key = context_len + qpos_hi + if USE_PER_SEQ_CAUSAL or (not USE_CAUSAL): + # Non-causal: keys can be AHEAD of query within the window + last_allowed_key = tl.minimum( + context_len + qpos_hi + SLIDING_WINDOW - 1, + seq_len - 1, + ) + else: + last_allowed_key = context_len + qpos_hi # Convert to tile indices and clamp tile_start = tl.maximum(0, first_allowed_key // TILE_SIZE) tile_end = tl.minimum((last_allowed_key // TILE_SIZE) + 1, num_tiles) @@ -262,10 +271,14 @@ def compute_kv_seq_mask( query_abs_pos, seq_offset, seq_idx, + seq_len, mm_prefix_range_ptr, SLIDING_WINDOW: tl.constexpr, USE_MM_PREFIX: tl.constexpr, MAX_MM_RANGES: tl.constexpr, + USE_CAUSAL: tl.constexpr = True, + USE_PER_SEQ_CAUSAL: tl.constexpr = False, + per_seq_causal_ptr=None, CHUNK_LOOKBACK: tl.constexpr = -1, CHUNK_SIZE: tl.constexpr = -1, ): @@ -279,9 +292,23 @@ def compute_kv_seq_mask( Chunked attention takes precedence over sliding window when both are non-default — the launcher zeros ``CHUNK_LOOKBACK`` whenever sliding window is disabled. + + When ``USE_PER_SEQ_CAUSAL`` is set, each sequence carries its own + causal flag via ``per_seq_causal_ptr``; non-causal sequences use a + simple ``key < seq_len`` bound instead. ``USE_CAUSAL=False`` + disables causal masking entirely. """ - # Compute attention mask: causal by default (key <= query) - seq_mask = seq_offset[None, :] <= query_abs_pos + if USE_PER_SEQ_CAUSAL: + is_causal = tl.load(per_seq_causal_ptr + seq_idx) + seq_mask = tl.where( + is_causal, + seq_offset[None, :] <= query_abs_pos, + seq_offset[None, :] < seq_len, + ) + elif USE_CAUSAL: + seq_mask = seq_offset[None, :] <= query_abs_pos + else: + seq_mask = seq_offset[None, :] < seq_len # Apply sliding window / chunked attention to base mask # BEFORE mm_prefix OR. @@ -293,7 +320,15 @@ def compute_kv_seq_mask( <= CHUNK_LOOKBACK ) elif SLIDING_WINDOW > 0: - seq_mask = seq_mask & ((query_abs_pos - seq_offset) < SLIDING_WINDOW) + sw_left = (query_abs_pos - seq_offset) < SLIDING_WINDOW + if USE_PER_SEQ_CAUSAL: + sw_right = (seq_offset[None, :] - query_abs_pos) < SLIDING_WINDOW + seq_mask = seq_mask & tl.where(is_causal, sw_left, sw_left & sw_right) + elif not USE_CAUSAL: + sw_right = (seq_offset[None, :] - query_abs_pos) < SLIDING_WINDOW + seq_mask = seq_mask & sw_left & sw_right + else: + seq_mask = seq_mask & sw_left # PrefixLM: extend mask with bidirectional ranges for multimodal tokens. # Applied AFTER sliding window so mm_prefix ranges override SW restriction. diff --git a/vllm/v1/attention/ops/triton_unified_attention.py b/vllm/v1/attention/ops/triton_unified_attention.py index 56f1d1c1d08..f39e44286be 100644 --- a/vllm/v1/attention/ops/triton_unified_attention.py +++ b/vllm/v1/attention/ops/triton_unified_attention.py @@ -215,6 +215,9 @@ def kernel_unified_attention( USE_SOFTCAP: tl.constexpr, # bool USE_SINKS: tl.constexpr, # bool SLIDING_WINDOW: tl.constexpr, # int + USE_CAUSAL: tl.constexpr, # bool + USE_PER_SEQ_CAUSAL: tl.constexpr, # bool + per_seq_causal_ptr, # [num_seqs] bool, or None USE_MM_PREFIX: tl.constexpr, # bool MAX_MM_RANGES: tl.constexpr, # int mm_prefix_range_ptr, @@ -389,6 +392,8 @@ def kernel_unified_attention( SLIDING_WINDOW, USE_MM_PREFIX, IS_3D, + USE_CAUSAL, + USE_PER_SEQ_CAUSAL, CHUNK_LOOKBACK, CHUNK_SIZE, ) @@ -493,10 +498,14 @@ def kernel_unified_attention( query_abs_pos, seq_offset, seq_idx, + seq_len, mm_prefix_range_ptr, SLIDING_WINDOW, USE_MM_PREFIX, MAX_MM_RANGES, + USE_CAUSAL, + USE_PER_SEQ_CAUSAL, + per_seq_causal_ptr, CHUNK_LOOKBACK, CHUNK_SIZE, ) @@ -532,11 +541,19 @@ def kernel_unified_attention( if SLIDING_WINDOW: qpos_lo = q_block_local_idx * BLOCK_Q - V = tl.where( - (context_len + qpos_lo - seq_offset[:, None]) < SLIDING_WINDOW, - V, - 0.0, - ) + dist = context_len + qpos_lo - seq_offset[:, None] + if USE_PER_SEQ_CAUSAL: + is_causal_seq = tl.load(per_seq_causal_ptr + seq_idx) + sw_mask_v = tl.where( + is_causal_seq, + dist < SLIDING_WINDOW, + (dist < SLIDING_WINDOW) & (dist > -SLIDING_WINDOW), + ) + elif USE_CAUSAL: + sw_mask_v = dist < SLIDING_WINDOW + else: + sw_mask_v = (dist < SLIDING_WINDOW) & (dist > -SLIDING_WINDOW) + V = tl.where(sw_mask_v, V, 0.0) if USE_PER_TOKEN_HEAD_SCALES: # Per-token-head quant: apply v_scale to P instead of V. P_v = (P * v_token_head_scales[None, :]).to(V.dtype) @@ -802,7 +819,11 @@ def unified_attention( # disabling this flag costs nothing. use_td: bool = False, ): - assert causal, "Only causal attention is supported" + # Resolve causal: bool or per-seq tensor. + use_per_seq_causal = isinstance(causal, torch.Tensor) + use_causal = bool(causal) if not use_per_seq_causal else True + per_seq_causal_ptr = causal if use_per_seq_causal else None + if sinks is not None: assert sinks.shape[0] == q.shape[1], "Sinks must be num_query_heads size" @@ -841,6 +862,26 @@ def unified_attention( ) BLOCK_Q = BLOCK_M // num_queries_per_kv + # Tuned launch parameters; ``None`` lets Triton pick its defaults. + launch_num_warps: int | None = None + launch_num_stages: int | None = None + + # head_size 256 with many query rows per sequence (e.g. diffusion-gemma + # bidirectional canvas passes) is prefill-shaped, but the decode-oriented + # defaults (BLOCK_Q=8, TILE=32, 4 warps) under-tile it. A wider KV tile + + # more query rows per block + 8 warps is ~2x faster on B200. + tuned_large_head = ( + head_size == 256 + and max_seqlen_q > 1 + and num_queries_per_kv <= 16 + and current_platform.is_device_capability_family(100) + ) + if tuned_large_head: + BLOCK_M = 32 + BLOCK_Q = BLOCK_M // num_queries_per_kv + launch_num_warps = 8 + launch_num_stages = 2 + # Ideally we would launch with kernel with: # \sum_i[ceil(query_len[i] / BLOCK_Q)] blocks. # However, it is slow to realize the query_lens on cpu. @@ -869,6 +910,11 @@ def unified_attention( head_size, sliding_window_val, q.element_size(), is_prefill=False ) + # Wider KV tile for the tuned large-head path (see above). Only the 2D + # path (used when max_seqlen_q > 1) reads TILE_SIZE_PREFILL. + if tuned_large_head: + TILE_SIZE_PREFILL = 128 + # USE_TD requires BLOCK_SIZE % TILE_SIZE == 0 (enforced by a # ``tl.static_assert`` in the kernel). The default prefill tile # size (32) is larger than a common ``block_size=16``, so clamp it @@ -964,6 +1010,12 @@ def unified_attention( grid = (total_num_q_blocks, num_kv_heads, num_par_softmax_segments) tile_size = TILE_SIZE_DECODE + launch_kwargs: dict[str, int] = {} + if launch_num_warps is not None: + launch_kwargs["num_warps"] = launch_num_warps + if launch_num_stages is not None: + launch_kwargs["num_stages"] = launch_num_stages + kernel_unified_attention[grid]( output_ptr=out, segm_output_ptr=segm_output_ptr, @@ -1002,10 +1054,13 @@ def unified_attention( USE_QQ_BIAS=use_qq_bias, USE_SOFTCAP=(softcap > 0), USE_SINKS=(sinks is not None), + SLIDING_WINDOW=(1 + window_size[0]), + USE_CAUSAL=use_causal, + USE_PER_SEQ_CAUSAL=use_per_seq_causal, + per_seq_causal_ptr=per_seq_causal_ptr, USE_MM_PREFIX=use_mm_prefix, MAX_MM_RANGES=max_mm_ranges, mm_prefix_range_ptr=mm_prefix_range, - SLIDING_WINDOW=(1 + window_size[0]), stride_k_cache_0=k.stride(0), stride_k_cache_1=k.stride(1), stride_k_cache_2=k.stride(2), @@ -1033,6 +1088,7 @@ def unified_attention( CHUNK_SIZE=chunk_size, USE_TD=use_td, USE_TD_QO=use_td_qo, + **launch_kwargs, ) if use_3d: diff --git a/vllm/v1/attention/ops/triton_unified_attention_diffkv.py b/vllm/v1/attention/ops/triton_unified_attention_diffkv.py index ef4f2835b5c..eaf62b6bce6 100644 --- a/vllm/v1/attention/ops/triton_unified_attention_diffkv.py +++ b/vllm/v1/attention/ops/triton_unified_attention_diffkv.py @@ -226,6 +226,7 @@ def kernel_unified_attention_diffkv( query_abs_pos, seq_offset, seq_idx, + seq_len, None, # mm_prefix_range_ptr SLIDING_WINDOW, False, # USE_MM_PREFIX diff --git a/vllm/v1/core/sched/async_scheduler.py b/vllm/v1/core/sched/async_scheduler.py index 2fd22f4c0cb..a79e84289af 100644 --- a/vllm/v1/core/sched/async_scheduler.py +++ b/vllm/v1/core/sched/async_scheduler.py @@ -27,10 +27,14 @@ class AsyncScheduler(Scheduler): scheduler_output.pending_structured_output_tokens |= ( request.use_structured_output and request.num_output_placeholders > 0 ) - # The request will generate a new token plus num_spec_tokens - # in this scheduling step. + # The request will generate num_sampled_tokens_per_step new tokens + # plus num_spec_tokens in this scheduling step. Diffusion has no AR + # bonus token (num_sampled_tokens_per_step == 0) — only the canvas + # (spec) tokens. cur_num_spec_tokens = len(spec_decode_tokens.get(req_id, ())) - request.num_output_placeholders += 1 + cur_num_spec_tokens + request.num_output_placeholders += ( + self.num_sampled_tokens_per_step + cur_num_spec_tokens + ) # Add placeholders for the new draft/spec tokens. # We will update the actual spec token ids in the worker process. request.spec_token_ids = self._spec_token_placeholders diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 9a3a9ffa7d6..926f406f199 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -113,6 +113,10 @@ class Scheduler(SchedulerInterface): self.kv_events_config is not None and self.kv_events_config.enable_kv_cache_events ) + # Diffusion models may not sample any tokens for a denoising step. + self.num_sampled_tokens_per_step = ( + 1 if not vllm_config.model_config.is_diffusion else 0 + ) # Create KVConnector for the Scheduler. Note that each Worker # will have a corresponding KVConnector with Role=WORKER. @@ -212,9 +216,9 @@ class Scheduler(SchedulerInterface): speculative_config = vllm_config.speculative_config self.use_eagle = False - self.num_spec_tokens = self.num_lookahead_tokens = 0 - if speculative_config: - self.num_spec_tokens = speculative_config.num_speculative_tokens + self.num_spec_tokens = vllm_config.num_speculative_tokens + self.num_lookahead_tokens = 0 + if speculative_config is not None: if speculative_config.use_eagle(): self.use_eagle = True self.num_lookahead_tokens = self.num_spec_tokens @@ -425,7 +429,10 @@ class Scheduler(SchedulerInterface): # Make sure the input position does not exceed the max model len. # This is necessary when using spec decoding. num_new_tokens = min( - num_new_tokens, self.max_model_len - 1 - request.num_computed_tokens + num_new_tokens, + self.max_model_len + - request.num_computed_tokens + - self.num_sampled_tokens_per_step, ) # Schedule encoder inputs. @@ -1473,9 +1480,12 @@ class Scheduler(SchedulerInterface): scheduled_spec_token_ids = ( scheduler_output.scheduled_spec_decode_tokens.get(req_id) ) - if scheduled_spec_token_ids and generated_token_ids: + if scheduled_spec_token_ids and ( + generated_token_ids or self.num_sampled_tokens_per_step == 0 + ): num_draft_tokens = len(scheduled_spec_token_ids) - num_accepted = len(generated_token_ids) - 1 + num_sampled = self.num_sampled_tokens_per_step + num_accepted = max(len(generated_token_ids) - num_sampled, 0) num_rejected = num_draft_tokens - num_accepted # num_computed_tokens represents the number of tokens # processed in the current step, considering scheduled diff --git a/vllm/v1/cudagraph_dispatcher.py b/vllm/v1/cudagraph_dispatcher.py index cf0c1d41772..6a48b6282d4 100644 --- a/vllm/v1/cudagraph_dispatcher.py +++ b/vllm/v1/cudagraph_dispatcher.py @@ -34,11 +34,7 @@ class CudagraphDispatcher: def __init__(self, vllm_config: VllmConfig): self.vllm_config = vllm_config self.compilation_config = vllm_config.compilation_config - self.uniform_decode_query_len = ( - 1 - if not self.vllm_config.speculative_config - else 1 + self.vllm_config.speculative_config.num_speculative_tokens - ) + self.uniform_decode_query_len = 1 + self.vllm_config.num_speculative_tokens # Dict to store valid cudagraph dispatching keys. self.cudagraph_keys: dict[CUDAGraphMode, set[BatchDescriptor]] = { diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index 08c814ab34e..91ca1f30317 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -156,6 +156,9 @@ class EngineCore: hash_block_size=hash_block_size, ) self.use_spec_decode = vllm_config.speculative_config is not None + self.check_for_draft_tokens = ( + self.use_spec_decode or vllm_config.model_config.is_diffusion + ) if self.scheduler.connector is not None: # type: ignore self.model_executor.init_kv_output_aggregator(self.scheduler.connector) # type: ignore @@ -475,8 +478,7 @@ class EngineCore: # When using async scheduling we can't get draft token ids in advance, # so we update draft token ids in the worker process and don't # need to update draft token ids here. - if not self.async_scheduling and self.use_spec_decode and model_executed: - # Take the draft token ids. + if self.check_for_draft_tokens and not self.async_scheduling and model_executed: draft_token_ids = self.model_executor.take_draft_token_ids() if draft_token_ids is not None: self.scheduler.update_draft_token_ids(draft_token_ids) @@ -575,18 +577,17 @@ class EngineCore: # in a field and do it immediately once step_with_batch_queue is # re-called. The latter slightly favors TTFT over TPOT/throughput. if deferred_scheduler_output: - # If we are doing speculative decoding with structured output, - # we need to get the draft token ids from the prior step before - # we can compute the grammar bitmask for the deferred request. - if self.use_spec_decode: + # When draft tokens are used with structured output, validate them + # before computing the grammar bitmask for the deferred request. + if self.check_for_draft_tokens: draft_token_ids = self.model_executor.take_draft_token_ids() - assert draft_token_ids is not None - # Update the draft token ids in the scheduler output to - # filter out the invalid spec tokens, which will be padded - # with -1 and skipped by the grammar bitmask computation. - self.scheduler.update_draft_token_ids_in_output( - draft_token_ids, deferred_scheduler_output - ) + if draft_token_ids is not None: + # Update the draft token ids in the scheduler output to + # filter out the invalid spec tokens, which will be padded + # with -1 and skipped by the grammar bitmask computation. + self.scheduler.update_draft_token_ids_in_output( + draft_token_ids, deferred_scheduler_output + ) # We now have the tokens needed to compute the bitmask for the # deferred request. Get the bitmask and call sample tokens. grammar_output = self.scheduler.get_grammar_bitmask( diff --git a/vllm/v1/metrics/loggers.py b/vllm/v1/metrics/loggers.py index 0052a35366a..021019dc1cd 100644 --- a/vllm/v1/metrics/loggers.py +++ b/vllm/v1/metrics/loggers.py @@ -110,7 +110,9 @@ class LoggingStatLogger(StatLoggerBase): self.connector_prefix_caching_metrics = CachingMetrics() self.mm_caching_metrics = CachingMetrics() - self.spec_decoding_logging = SpecDecodingLogging() + model_config = self.vllm_config.model_config + is_diffusion = model_config is not None and model_config.is_diffusion + self.spec_decoding_logging = SpecDecodingLogging(is_diffusion=is_diffusion) kv_transfer_config = self.vllm_config.kv_transfer_config self.kv_connector_logging = KVConnectorLogging(kv_transfer_config) self.cudagraph_logging = None @@ -436,7 +438,10 @@ class PrometheusStatLogger(AggregateStatLoggerBase): per_engine_labelvalues = self.per_engine_labelvalues self.spec_decoding_prom = self._spec_decoding_cls( - vllm_config.speculative_config, labelnames, per_engine_labelvalues + vllm_config.speculative_config, + labelnames, + per_engine_labelvalues, + is_diffusion=vllm_config.model_config.is_diffusion, ) self.kv_connector_prom = self._kv_connector_cls( vllm_config, labelnames, per_engine_labelvalues diff --git a/vllm/v1/spec_decode/metrics.py b/vllm/v1/spec_decode/metrics.py index 9a41ff5c818..5da41510b4d 100644 --- a/vllm/v1/spec_decode/metrics.py +++ b/vllm/v1/spec_decode/metrics.py @@ -53,7 +53,11 @@ class SpecDecodingLogging: before resetting to zero. """ - def __init__(self): + def __init__(self, is_diffusion: bool = False): + # Diffusion (dLLM) models reuse the spec-decode data path with + # overloaded semantics, so the raw spec-decode framing (drafts, bonus + # token, per-position vector) is logged with diffusion-native terms. + self.is_diffusion = is_diffusion self.reset() def reset(self): @@ -85,6 +89,17 @@ class SpecDecodingLogging: draft_throughput = num_draft_tokens / elapsed_time accepted_throughput = num_accepted_tokens / elapsed_time + if self.is_diffusion: + self._log_diffusion( + log_fn, + num_denoising_steps=num_drafts, + num_canvas_tokens=num_draft_tokens, + num_committed_tokens=num_accepted_tokens, + committed_throughput=accepted_throughput, + ) + self.reset() + return + draft_acceptance_rate = ( num_accepted_tokens / num_draft_tokens * 100 if num_draft_tokens > 0 @@ -117,6 +132,43 @@ class SpecDecodingLogging: ) self.reset() + def _log_diffusion( + self, + log_fn, + num_denoising_steps: int, + num_canvas_tokens: int, + num_committed_tokens: int, + committed_throughput: float, + ): + # Each "draft" is one denoising step that re-evaluates the canvas block + # and finalizes some of its positions. + mean_committed_per_step = ( + num_committed_tokens / num_denoising_steps + if num_denoising_steps > 0 + else float("nan") + ) + mean_steps_per_canvas = ( + num_canvas_tokens / num_committed_tokens + if num_committed_tokens > 0 + else float("nan") + ) + + log_fn( + "DiffusionDecoding metrics: " + "Committed token throughput: %.2f tokens/s, " + "Mean denoising steps per canvas: %.2f, " + "Mean tokens committed per denoising step: %.2f, " + "Committed: %d tokens, " + "Denoising steps: %d, " + "Canvas positions evaluated: %d", + committed_throughput, + mean_steps_per_canvas, + mean_committed_per_step, + num_committed_tokens, + num_denoising_steps, + num_canvas_tokens, + ) + class SpecDecodingProm: """Record spec decoding metrics in Prometheus. @@ -146,56 +198,66 @@ class SpecDecodingProm: speculative_config: SpeculativeConfig | None, labelnames: list[str], per_engine_labelvalues: dict[int, list[object]], + is_diffusion: bool = False, ): - self.spec_decoding_enabled = speculative_config is not None + # Diffusion (dLLM) models reuse the spec-decode counters but expose them + # under diffusion-native names; the per-position acceptance vector does + # not apply, so it is omitted. + self.is_diffusion = is_diffusion + self.spec_decoding_enabled = speculative_config is not None or is_diffusion if not self.spec_decoding_enabled: return - counter_drafts = self._counter_cls( - name="vllm:spec_decode_num_drafts", - documentation="Number of spec decoding drafts.", - labelnames=labelnames, - ) - self.counter_spec_decode_num_drafts = create_metric_per_engine( - counter_drafts, per_engine_labelvalues - ) + if is_diffusion: + counter_specs = [ + ("vllm:diffusion_num_denoising_steps", "Number of denoising steps."), + ( + "vllm:diffusion_num_canvas_positions", + "Number of canvas positions evaluated.", + ), + ( + "vllm:diffusion_num_committed_tokens", + "Number of committed (finalized) tokens.", + ), + ] + else: + counter_specs = [ + ("vllm:spec_decode_num_drafts", "Number of spec decoding drafts."), + ("vllm:spec_decode_num_draft_tokens", "Number of draft tokens."), + ("vllm:spec_decode_num_accepted_tokens", "Number of accepted tokens."), + ] - counter_draft_tokens = self._counter_cls( - name="vllm:spec_decode_num_draft_tokens", - documentation="Number of draft tokens.", - labelnames=labelnames, - ) - self.counter_spec_decode_num_draft_tokens = create_metric_per_engine( - counter_draft_tokens, per_engine_labelvalues - ) + counters = [ + create_metric_per_engine( + self._counter_cls(name=name, documentation=doc, labelnames=labelnames), + per_engine_labelvalues, + ) + for name, doc in counter_specs + ] + # num_drafts/num_draft_tokens/num_accepted_tokens map onto denoising + # steps/canvas positions/committed tokens in the diffusion path. + self.counter_spec_decode_num_drafts = counters[0] + self.counter_spec_decode_num_draft_tokens = counters[1] + self.counter_spec_decode_num_accepted_tokens = counters[2] - counter_accepted_tokens = self._counter_cls( - name="vllm:spec_decode_num_accepted_tokens", - documentation="Number of accepted tokens.", - labelnames=labelnames, - ) - self.counter_spec_decode_num_accepted_tokens = create_metric_per_engine( - counter_accepted_tokens, per_engine_labelvalues - ) - - assert speculative_config is not None - num_spec_tokens = ( - speculative_config.num_speculative_tokens - if self.spec_decoding_enabled - else 0 - ) - pos_labelnames = labelnames + ["position"] - base_counter = self._counter_cls( - name="vllm:spec_decode_num_accepted_tokens_per_pos", - documentation="Accepted tokens per draft position.", - labelnames=pos_labelnames, - ) self.counter_spec_decode_num_accepted_tokens_per_pos: dict[ int, list[prometheus_client.Counter] - ] = { - idx: [base_counter.labels(*lv, str(pos)) for pos in range(num_spec_tokens)] - for idx, lv in per_engine_labelvalues.items() - } + ] = {} + if not is_diffusion: + assert speculative_config is not None + num_spec_tokens = speculative_config.num_speculative_tokens + pos_labelnames = labelnames + ["position"] + base_counter = self._counter_cls( + name="vllm:spec_decode_num_accepted_tokens_per_pos", + documentation="Accepted tokens per draft position.", + labelnames=pos_labelnames, + ) + self.counter_spec_decode_num_accepted_tokens_per_pos = { + idx: [ + base_counter.labels(*lv, str(pos)) for pos in range(num_spec_tokens) + ] + for idx, lv in per_engine_labelvalues.items() + } def observe(self, spec_decoding_stats: SpecDecodingStats, engine_idx: int = 0): if not self.spec_decoding_enabled: @@ -210,6 +272,6 @@ class SpecDecodingProm: spec_decoding_stats.num_accepted_tokens ) for pos, counter in enumerate( - self.counter_spec_decode_num_accepted_tokens_per_pos[engine_idx] + self.counter_spec_decode_num_accepted_tokens_per_pos.get(engine_idx, []) ): counter.inc(spec_decoding_stats.num_accepted_tokens_per_pos[pos]) diff --git a/vllm/v1/structured_output/__init__.py b/vllm/v1/structured_output/__init__.py index 6a4fcbb629f..30921f3d74a 100644 --- a/vllm/v1/structured_output/__init__.py +++ b/vllm/v1/structured_output/__init__.py @@ -211,11 +211,8 @@ class StructuredOutputManager: if not structured_output_request_ids: return None - max_num_spec_tokens = 0 - if self.vllm_config.speculative_config is not None: - max_num_spec_tokens = ( - self.vllm_config.speculative_config.num_speculative_tokens - ) + # Covers both speculative decoding and diffusion LLMs (canvas_length). + max_num_spec_tokens = self.vllm_config.num_speculative_tokens if self._grammar_bitmask is None: assert self.backend is not None @@ -277,7 +274,13 @@ class StructuredOutputManager: state_advancements = 0 req_tokens = scheduled_spec_decode_tokens.get(req_id, ()) - for token in itertools.chain(req_tokens, (-1,)): + if self.vllm_config.model_config.is_diffusion and req_tokens: + # Diffusion LLMs don't sample a bonus token after the + # scheduled positions, so don't append the -1 placeholder. + token_iter: Iterable[int] = req_tokens + else: + token_iter = itertools.chain(req_tokens, (-1,)) + for token in token_iter: self._fill_bitmasks(((grammar, cumulative_index, apply_bitmask),)) if token == -1: # Stop advancing the grammar once we hit a padding token. diff --git a/vllm/v1/worker/gpu/input_batch.py b/vllm/v1/worker/gpu/input_batch.py index f905d09e45f..6b750fe7ebf 100644 --- a/vllm/v1/worker/gpu/input_batch.py +++ b/vllm/v1/worker/gpu/input_batch.py @@ -302,6 +302,7 @@ def _combine_sampled_and_draft_tokens_kernel( cu_num_logits_ptr, logits_indices_ptr, BLOCK_SIZE: tl.constexpr, + NUM_NEW_SAMPLED_TOKENS: tl.constexpr = 1, ): batch_idx = tl.program_id(0) req_state_idx = tl.load(idx_mapping_ptr + batch_idx) @@ -310,7 +311,7 @@ def _combine_sampled_and_draft_tokens_kernel( cu_num_logits_start = tl.load(cu_num_logits_ptr + batch_idx) cu_num_logits_end = tl.load(cu_num_logits_ptr + batch_idx + 1) num_logits = cu_num_logits_end - cu_num_logits_start - num_draft_tokens = num_logits - 1 + num_draft_tokens = num_logits - NUM_NEW_SAMPLED_TOKENS # Compute the logits indices. block = tl.arange(0, BLOCK_SIZE) @@ -328,9 +329,10 @@ def _combine_sampled_and_draft_tokens_kernel( # Handling prefill tokens. No sampled or draft tokens. return - # Write the last sampled token ID to input_ids. - last_token_id = tl.load(last_sampled_tokens_ptr + req_state_idx) - tl.store(input_ids_ptr + query_end - num_logits, last_token_id) + if NUM_NEW_SAMPLED_TOKENS > 0: + # Write the last sampled token ID to input_ids. + last_token_id = tl.load(last_sampled_tokens_ptr + req_state_idx) + tl.store(input_ids_ptr + query_end - num_logits, last_token_id) # Write the draft tokens (if any) to input_ids. if num_draft_tokens > 0: @@ -356,7 +358,11 @@ def combine_sampled_and_draft_tokens( draft_tokens: torch.Tensor, cu_num_logits: torch.Tensor, num_logits: int, + num_new_sampled_tokens: int = 1, # excl accepted draft tokens, a.k.a bonus tokens ) -> torch.Tensor: + assert num_new_sampled_tokens in (0, 1), ( + f"num_new_sampled_tokens must be 0 or 1, got {num_new_sampled_tokens}" + ) # use idx_mapping.shape[0] for actual request count num_reqs = idx_mapping.shape[0] num_speculative_steps = draft_tokens.shape[-1] @@ -377,9 +383,12 @@ def combine_sampled_and_draft_tokens( draft_tokens.stride(0), cu_num_logits, logits_indices, - # NOTE(woosuk): Add 1 to ensure the block can cover the last sampled token - # in addition to all draft tokens. - BLOCK_SIZE=triton.next_power_of_2(num_speculative_steps + 1), + NUM_NEW_SAMPLED_TOKENS=num_new_sampled_tokens, + # NOTE(woosuk): Add num_new_sampled_tokens to ensure the block covers the + # last sampled token in addition to all draft tokens. + BLOCK_SIZE=triton.next_power_of_2( + num_speculative_steps + num_new_sampled_tokens + ), ) return logits_indices diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 7cd1e6c5c86..d269bf25bdb 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -78,7 +78,6 @@ from vllm.v1.worker.gpu.input_batch import ( InputBuffers, combine_sampled_and_draft_tokens, expand_idx_mapping, - get_num_sampled_and_rejected, post_update, post_update_num_computed_tokens, prepare_pos_seq_lens, @@ -185,11 +184,9 @@ class GPUModelRunner(LoRAModelRunnerMixin): # Speculative decoding. self.speculator = None - self.num_speculative_steps = 0 self.use_aux_hidden_state_outputs = False + self.num_speculative_steps = vllm_config.num_speculative_tokens if self.speculative_config is not None: - self.num_speculative_steps = self.speculative_config.num_speculative_tokens - if self.is_last_pp_rank: self.speculator = init_speculator(self.vllm_config, self.device) @@ -204,7 +201,6 @@ class GPUModelRunner(LoRAModelRunnerMixin): # Draft tokens propagation - for spec-dec + struct outputs. self.draft_tokens_handler = DraftTokensHandler(self.device) - self.uniform_decode_query_len = 1 + self.num_speculative_steps # Pooling models. self.is_pooling_model = self.model_config.runner_type == "pooling" @@ -232,38 +228,12 @@ class GPUModelRunner(LoRAModelRunnerMixin): device=self.device, ) + # Samplers and decode_query_len created in load_model() after + # model_state exists (num_new_sampled_tokens_per_step from ModelState). self.sampler: Sampler | None = None self.rejection_sampler: RejectionSampler | None = None self.prompt_logprobs_worker: PromptLogprobsWorker | None = None self.structured_outputs_worker: StructuredOutputsWorker | None = None - if self.is_last_pp_rank and not self.is_pooling_model: - # Initialize sampling-related workers. - # These components are only set up on the last PP rank and - # for generative (non-pooling) models. - self.sampler = Sampler( - max_num_reqs=self.max_num_reqs, - vocab_size=self.vocab_size, - device=self.device, - req_states=self.req_states, - logprobs_mode=self.model_config.logprobs_mode, - num_speculative_tokens=self.num_speculative_steps + 1, - use_fp64_gumbel=self.model_config.use_fp64_gumbel, - ) - if self.speculative_config is not None: - self.rejection_sampler = RejectionSampler( - self.sampler, - self.speculative_config, - self.device, - ) - self.prompt_logprobs_worker = PromptLogprobsWorker(self.max_num_reqs) - self.structured_outputs_worker = StructuredOutputsWorker( - max_num_logits=self.max_num_reqs * (self.num_speculative_steps + 1), - vocab_size=self.vocab_size, - device=self.device, - ) - - # For CUDA graphs, and will init cudagraph_manager after init_attn_backend. - self.decode_query_len = self.num_speculative_steps + 1 self.cudagraph_manager: ModelCudaGraphManager | None = None # LoRA-related workers. self.lora_state = LoraState(max_num_reqs=self.max_num_reqs) @@ -335,6 +305,40 @@ class GPUModelRunner(LoRAModelRunnerMixin): self.model_state = init_model_state( self.vllm_config, self.model, self.encoder_cache, self.device ) + + self.decode_query_len = ( + self.num_speculative_steps + + self.model_state.num_new_sampled_tokens_per_step + ) + + # Initialize samplers. Model states may override via custom_sampler(). + if self.is_last_pp_rank and not self.is_pooling_model: + self.sampler = Sampler( + max_num_reqs=self.max_num_reqs, + vocab_size=self.vocab_size, + device=self.device, + req_states=self.req_states, + logprobs_mode=self.model_config.logprobs_mode, + num_speculative_tokens=self.decode_query_len, + use_fp64_gumbel=self.model_config.use_fp64_gumbel, + ) + custom = self.model_state.custom_sampler(self.sampler) + + if custom: + self.sampler, self.rejection_sampler = custom + elif self.speculative_config is not None: + self.rejection_sampler = RejectionSampler( + self.sampler, + self.speculative_config, + self.device, + ) + self.prompt_logprobs_worker = PromptLogprobsWorker(self.max_num_reqs) + self.structured_outputs_worker = StructuredOutputsWorker( + max_num_logits=self.max_num_reqs * self.decode_query_len, + vocab_size=self.vocab_size, + device=self.device, + ) + if self.is_pooling_model and self.is_last_pp_rank: self.pooling_runner = PoolingRunner(self.model) eplb_models_added |= self.eplb.maybe_register_model( @@ -447,7 +451,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): cudagraph_mode = self.compilation_config.resolve_cudagraph_mode_and_sizes( attn_cg_support.min_cg_support, attn_cg_support.min_cg_attn_backend, - self.uniform_decode_query_len, + self.decode_query_len, self.parallel_config.tensor_parallel_size, self.kv_cache_config, self.max_num_reqs, @@ -710,6 +714,9 @@ class GPUModelRunner(LoRAModelRunnerMixin): return cuda_graph_size def _remove_request(self, req_id: str) -> bool: + # Call model_state.remove_request *before* req_states.remove_request + # so the model_state can still look up the slot index. + self.model_state.remove_request(req_id) req_idx = self.req_states.remove_request(req_id) if req_idx is None: return False @@ -857,16 +864,16 @@ class GPUModelRunner(LoRAModelRunnerMixin): dtype=np.int32, count=num_reqs, ) + num_bonus_tokens = self.model_state.num_new_sampled_tokens_per_step total_num_draft_tokens = int(num_draft_tokens_per_req.sum()) - total_num_logits = num_reqs + total_num_draft_tokens - - num_logits = num_draft_tokens_per_req + 1 + total_num_logits = num_reqs * num_bonus_tokens + total_num_draft_tokens + num_logits = num_draft_tokens_per_req + num_bonus_tokens cu_num_logits_np = np.empty(num_reqs + 1, dtype=np.int32) cu_num_logits_np[0] = 0 np.cumsum(num_logits, out=cu_num_logits_np[1:]) cu_num_logits = async_copy_to_gpu(cu_num_logits_np, device=self.device) - max_expand_len = self.num_speculative_steps + 1 + max_expand_len = self.decode_query_len expanded_idx_mapping, expanded_local_pos = expand_idx_mapping( idx_mapping, total_num_logits, cu_num_logits, max_expand_len ) @@ -935,6 +942,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): self.req_states.draft_tokens, cu_num_logits, total_num_logits, + self.model_state.num_new_sampled_tokens_per_step, ) # CPU upper bound on seq_lens; padded entries left at zero. @@ -1027,8 +1035,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): grammar_output.grammar_bitmask, ) - if input_batch.num_draft_tokens == 0: - # No draft tokens (common case). + if input_batch.num_draft_tokens == 0 or self.rejection_sampler is None: assert self.sampler is not None sampler_output = self.sampler(logits, input_batch) else: @@ -1042,16 +1049,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): self.speculator.draft_logits, ) - # Get the number of sampled and rejected tokens. - # For chunked prefills, num_sampled and num_rejected are both 0. - num_sampled, num_rejected = get_num_sampled_and_rejected( - sampler_output.num_sampled, - input_batch.seq_lens, - input_batch.cu_num_logits, - input_batch.idx_mapping, - self.req_states.prefill_len.gpu, - ) - return sampler_output, num_sampled, num_rejected + return sampler_output, sampler_output.num_sampled, sampler_output.num_rejected def postprocess_sampled( self, @@ -1448,7 +1446,14 @@ class GPUModelRunner(LoRAModelRunnerMixin): mm_inputs=mm_inputs, ) self.req_states.draft_tokens[input_batch.idx_mapping] = draft_tokens - self.draft_tokens_handler.set_draft_tokens(input_batch, draft_tokens) + + if self.num_speculative_steps > 0: + # Spec-decode and diffusion LLMs both use draft tokens but the latter does + # not have a speculator (i.e. self.speculator is None) + self.draft_tokens_handler.set_draft_tokens( + input_batch, + self.req_states.draft_tokens[input_batch.idx_mapping], + ) # Post-step KV connector related operations. kv_connector_output = self.kv_connector.post_forward(finished_req_ids) diff --git a/vllm/v1/worker/gpu/model_states/__init__.py b/vllm/v1/worker/gpu/model_states/__init__.py index b096fcaf5e6..e24c7e9b1cb 100644 --- a/vllm/v1/worker/gpu/model_states/__init__.py +++ b/vllm/v1/worker/gpu/model_states/__init__.py @@ -13,6 +13,11 @@ def init_model_state( encoder_cache: EncoderCache | None, device: torch.device, ): + # Let the model provide its own ModelState if it defines one. + if hasattr(model, "get_model_state_cls"): + cls = model.get_model_state_cls() + return cls(vllm_config, model, encoder_cache, device) + if ( "WhisperForConditionalGeneration" in vllm_config.model_config.architectures or "CohereAsrForConditionalGeneration" in vllm_config.model_config.architectures diff --git a/vllm/v1/worker/gpu/model_states/interface.py b/vllm/v1/worker/gpu/model_states/interface.py index 55bf8d473cc..86f28e08ea9 100644 --- a/vllm/v1/worker/gpu/model_states/interface.py +++ b/vllm/v1/worker/gpu/model_states/interface.py @@ -53,6 +53,9 @@ class ModelState(ABC): def add_request(self, req_index: int, new_req_data: NewRequestData) -> None: return None + def remove_request(self, req_id: str) -> None: + return None + def apply_staged_writes(self) -> None: return None @@ -89,3 +92,16 @@ class ModelState(ABC): for_capture: bool = False, ) -> dict[str, Any]: raise NotImplementedError + + def custom_sampler(self, sampler: Any) -> tuple[Any, Any] | None: + """Wrap or replace the default sampler. + + Called after model loading with the already-constructed base + ``Sampler``. Return ``None`` to keep the defaults, or + ``(sampler, rejection_sampler | None)`` to override. + """ + return None + + num_new_sampled_tokens_per_step: int = 1 + """New tokens sampled on each decode step + (excluding accepted draft tokens, a.k.a num bonus tokens).""" diff --git a/vllm/v1/worker/gpu/sample/output.py b/vllm/v1/worker/gpu/sample/output.py index f38ac8affd8..130f4ddbf8a 100644 --- a/vllm/v1/worker/gpu/sample/output.py +++ b/vllm/v1/worker/gpu/sample/output.py @@ -13,3 +13,4 @@ class SamplerOutput: logprobs_tensors: LogprobsTensors | None num_nans: torch.Tensor | None num_sampled: torch.Tensor | None + num_rejected: torch.Tensor | None = None diff --git a/vllm/v1/worker/gpu/sample/sampler.py b/vllm/v1/worker/gpu/sample/sampler.py index 6b545aef3a2..b269de9eaed 100644 --- a/vllm/v1/worker/gpu/sample/sampler.py +++ b/vllm/v1/worker/gpu/sample/sampler.py @@ -12,7 +12,7 @@ from vllm.v1.sample.ops.topk_topp_sampler import ( flashinfer_sample, flashinfer_sampler_supported, ) -from vllm.v1.worker.gpu.input_batch import InputBatch +from vllm.v1.worker.gpu.input_batch import InputBatch, get_num_sampled_and_rejected from vllm.v1.worker.gpu.metrics.logits import get_num_nans from vllm.v1.worker.gpu.sample.bad_words import BadWordsState from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample @@ -44,6 +44,7 @@ class Sampler: self.compute_nans = envs.VLLM_COMPUTE_NANS_IN_LOGITS # False by default. self.use_fp64_gumbel = use_fp64_gumbel + self.req_states = req_states self.sampling_states = SamplingStates(max_num_reqs, vocab_size) self.penalties_state = PenaltiesState(req_states) self.logit_bias_state = LogitBiasState(max_num_reqs, device) @@ -118,6 +119,17 @@ class Sampler: else: logprobs_tensors = None + # 1 sampled token per request, except chunked-prefill requests + # (seq_len < prefill_len) which aren't done prefilling and produce no + # output token. num_rejected is always 0 here (one logit per request). + num_sampled, num_rejected = get_num_sampled_and_rejected( + input_batch.seq_lens.new_ones(input_batch.num_reqs), + input_batch.seq_lens, + input_batch.cu_num_logits, + input_batch.idx_mapping, + self.req_states.prefill_len.gpu, + ) + # These are GPU tensors. sampler_output = SamplerOutput( # The sampled tokens are expanded to 2D tensor with shape @@ -126,7 +138,8 @@ class Sampler: sampled_token_ids=sampled.view(-1, 1), logprobs_tensors=logprobs_tensors, num_nans=num_nans, - num_sampled=input_batch.seq_lens.new_ones(input_batch.num_reqs), + num_sampled=num_sampled, + num_rejected=num_rejected, ) return sampler_output diff --git a/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py b/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py index 1fe079a43e7..3868604d3ae 100644 --- a/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py +++ b/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py @@ -6,7 +6,10 @@ from vllm.config import SpeculativeConfig from vllm.triton_utils import tl, triton from vllm.v1.outputs import LogprobsTensors from vllm.v1.spec_decode.utils import unconditional_to_conditional_rates -from vllm.v1.worker.gpu.input_batch import InputBatch +from vllm.v1.worker.gpu.input_batch import ( + InputBatch, + get_num_sampled_and_rejected, +) from vllm.v1.worker.gpu.metrics.logits import get_num_nans from vllm.v1.worker.gpu.sample.logprob import compute_topk_logprobs from vllm.v1.worker.gpu.sample.output import SamplerOutput @@ -136,9 +139,18 @@ class RejectionSampler: else logits, ) + num_sampled, num_rejected = get_num_sampled_and_rejected( + num_sampled, + input_batch.seq_lens, + input_batch.cu_num_logits, + input_batch.idx_mapping, + self.sampler.req_states.prefill_len.gpu, + ) + return SamplerOutput( sampled_token_ids=sampled, logprobs_tensors=logprobs_tensors, num_nans=num_nans, num_sampled=num_sampled, + num_rejected=num_rejected, ) diff --git a/vllm/v1/worker/gpu/spec_decode/utils.py b/vllm/v1/worker/gpu/spec_decode/utils.py index 7bfd981ee0c..4ab45b2ae27 100644 --- a/vllm/v1/worker/gpu/spec_decode/utils.py +++ b/vllm/v1/worker/gpu/spec_decode/utils.py @@ -35,6 +35,10 @@ class DraftTokensHandler: self.copy_stream.wait_stream(current_stream) with torch.cuda.stream(self.copy_stream): self.draft_tokens_np = async_copy_to_np(draft_tokens) + # draft_tokens is a temporary allocation on the main stream and read here on + # copy_stream; without record_stream, the caching allocator may reuse its + # memory before the async copy executes. + draft_tokens.record_stream(self.copy_stream) self.copy_event.record() def get_draft_tokens(self) -> DraftTokenIds | None: diff --git a/vllm/v1/worker/gpu/warmup.py b/vllm/v1/worker/gpu/warmup.py index 83d87c74a4a..0da845a0673 100644 --- a/vllm/v1/worker/gpu/warmup.py +++ b/vllm/v1/worker/gpu/warmup.py @@ -30,17 +30,18 @@ def warmup_kernels( pipeline parallel coordination. The first iteration simulates a prefill with requests of - 2 + num_spec_steps prompt tokens each. The second iteration simulates - a decode step with all requests generating 1 + num_spec_steps tokens. + decode_query_len + 1 prompt tokens each. The second iteration simulates + a decode step with all requests generating decode_query_len tokens. """ num_spec_steps = model_runner.num_speculative_steps - # Use 1 + num_spec_steps + 1 tokens so the prefill batch's per-request - # query length exceeds decode_query_len (= 1 + num_spec_steps), preventing - # it from being misclassified as a uniform decode batch. - prompt_len = 2 + num_spec_steps + decode_query_len = model_runner.decode_query_len + # Use decode_query_len + 1 tokens so the prefill batch's per-request query + # length exceeds decode_query_len, preventing it from being misclassified as + # a uniform decode batch. + prompt_len = decode_query_len + 1 prompt_token_ids = list(range(prompt_len)) - # After prefill, decode generates 1 verified + num_spec_steps draft tokens. - decode_len = prompt_len + 1 + num_spec_steps + # After prefill, decode generates decode_query_len tokens. + decode_len = prompt_len + decode_query_len kv_cache_groups = model_runner.kv_cache_config.kv_cache_groups num_kv_cache_groups = len(kv_cache_groups) @@ -57,7 +58,7 @@ def warmup_kernels( num_reqs = min( model_runner.scheduler_config.max_num_seqs, model_runner.scheduler_config.max_num_batched_tokens - // max(prompt_len, 1 + num_spec_steps), + // max(prompt_len, decode_query_len), # Reserve block 0 (null block) and ensure we have enough blocks. max(1, (model_runner.kv_cache_config.num_blocks - 1) // max_blocks_per_req), ) @@ -79,7 +80,7 @@ def warmup_kernels( nonlocal next_block_id return list(range(next_block_id, next_block_id := next_block_id + num_blocks)) - # Step 1: Prefill all requests with 2 + num_spec_steps prompt tokens each. + # Step 1: Prefill all requests with 1 + decode_query_len prompt tokens each. new_reqs = [ NewRequestData.from_request( Request(req_ids[i], prompt_token_ids, sampling_params, pooling_params), @@ -117,7 +118,7 @@ def warmup_kernels( worker_sample_tokens(grammar_output) - # Step 2: Decode all requests with 1 + num_spec_steps tokens each. + # Step 2: Decode all requests with decode_query_len tokens each. cached_req_data = CachedRequestData.make_empty() cached_req_data.req_ids = list(req_ids) cached_req_data.num_computed_tokens = [prompt_len] * num_reqs @@ -131,7 +132,7 @@ def warmup_kernels( decode_output = SchedulerOutput.make_empty() decode_output.scheduled_cached_reqs = cached_req_data decode_output.num_scheduled_tokens = { - req_id: 1 + num_spec_steps for req_id in req_ids + req_id: decode_query_len for req_id in req_ids } if num_spec_steps > 0: decode_output.scheduled_spec_decode_tokens = { diff --git a/vllm/vllm_flash_attn/flash_attn_interface.py b/vllm/vllm_flash_attn/flash_attn_interface.py index 5004ba9c8f2..276b9b4250f 100644 --- a/vllm/vllm_flash_attn/flash_attn_interface.py +++ b/vllm/vllm_flash_attn/flash_attn_interface.py @@ -209,6 +209,7 @@ def flash_attn_varlen_func( # FA4 only mask_mod=None, aux_tensors=None, + dynamic_causal: "torch.Tensor | None" = None, ): """dropout_p should be set to 0.0 during evaluation Supports multi-query and grouped-query attention (MQA/GQA) by passing in K, V with fewer heads @@ -392,6 +393,7 @@ def flash_attn_varlen_func( page_table=block_table, softmax_scale=softmax_scale, causal=causal, + dynamic_causal=dynamic_causal, softcap=softcap, window_size_left=real_window_size[0] if real_window_size[0] >= 0 else None, window_size_right=real_window_size[1] if real_window_size[1] >= 0 else None, From 39dee1114a2cd183a9fb72b561808b385b6c9daa Mon Sep 17 00:00:00 2001 From: allgather Date: Thu, 11 Jun 2026 22:17:55 -0700 Subject: [PATCH 30/52] [MM][Perf][CG] Support ViT full cudagraphs for mllama4 (#40660) Signed-off-by: allgather Co-authored-by: Isotr0py --- docs/design/cuda_graphs_multimodal.md | 9 + .../multimodal/vision_language_offline.py | 1 + .../generation/test_vit_cudagraph.py | 20 ++ tests/models/utils.py | 5 +- vllm/model_executor/models/mllama4.py | 172 ++++++++++++++++-- 5 files changed, 193 insertions(+), 14 deletions(-) diff --git a/docs/design/cuda_graphs_multimodal.md b/docs/design/cuda_graphs_multimodal.md index 8cbbedf9d0b..dd0e47a1950 100644 --- a/docs/design/cuda_graphs_multimodal.md +++ b/docs/design/cuda_graphs_multimodal.md @@ -82,6 +82,7 @@ Models opt-in to encoder CUDA Graphs by implementing the [SupportsEncoderCudaGra | Architecture | Models | CG for Image | CG for Video | | ------------ | ------ | ------------ | ------------ | +| `Llama4ForConditionalGeneration` | `Llama 4` | ✅︎ | - | | `InternVLChatModel` | `InternVL3.5`, `InternVL3`, `InternVL2.5`, `InternVL2` | ✅︎ | ✅︎ | | `Qwen2VLForConditionalGeneration` | `Qwen2-VL` | ✅︎ | ✅︎ | | `Qwen2_5_VLForConditionalGeneration` | `Qwen2.5-VL` | ✅︎ | ✅︎ | @@ -114,6 +115,14 @@ vllm serve Qwen/Qwen3-VL-32B \ --compilation-config '{"cudagraph_mm_encoder": true}' ``` +For `Llama 4` (image only): + +```bash +vllm serve meta-llama/Llama-4-Scout-17B-16E-Instruct \ + --limit-mm-per-prompt '{"image": 1}' \ + --compilation-config '{"cudagraph_mm_encoder": true}' +``` + With explicit budgets: ```bash diff --git a/examples/generate/multimodal/vision_language_offline.py b/examples/generate/multimodal/vision_language_offline.py index 40a4b8ae6d1..a7df5b00c3b 100644 --- a/examples/generate/multimodal/vision_language_offline.py +++ b/examples/generate/multimodal/vision_language_offline.py @@ -2532,6 +2532,7 @@ MODELS_NEED_VIDEO_METADATA = [ MODELS_SUPPORT_VIT_CUDA_GRAPH = [ + "llama4", "internvl_chat", "qwen2_5_vl", "qwen3_vl", diff --git a/tests/models/multimodal/generation/test_vit_cudagraph.py b/tests/models/multimodal/generation/test_vit_cudagraph.py index f781caf492b..a1dc4e5bdd8 100644 --- a/tests/models/multimodal/generation/test_vit_cudagraph.py +++ b/tests/models/multimodal/generation/test_vit_cudagraph.py @@ -55,6 +55,26 @@ def step3_vl_chat_template(content: str) -> str: MODEL_CONFIGS: dict[str, VitCudagraphTestConfig] = { + "llama4": VitCudagraphTestConfig( + model="meta-llama/Llama-4-Scout-17B-16E-Instruct", + modalities=["image"], + image_prompt=( + "<|begin_of_text|><|header_start|>user<|header_end|>\n\n" + "<|image|>What is in this image?<|eot|>" + "<|header_start|>assistant<|header_end|>\n\n" + ), + max_model_len=4096, + max_tokens=32, + max_num_seqs=2, + vllm_runner_kwargs={ + "load_format": "dummy", + "hf_overrides": partial( + dummy_hf_overrides, + model_arch="Llama4ForConditionalGeneration", + ), + }, + marks=[pytest.mark.core_model], + ), "internvl": VitCudagraphTestConfig( model="OpenGVLab/InternVL3-1B", num_video_frames=8, diff --git a/tests/models/utils.py b/tests/models/utils.py index 259cdac13c0..8a629552131 100644 --- a/tests/models/utils.py +++ b/tests/models/utils.py @@ -507,12 +507,13 @@ def dummy_hf_overrides( # Only set MoE related config when the model has MoE layers. # Otherwise all models detected as MoE by _get_transformers_backend_cls. if model_arch_config.num_experts > 0: + num_experts_per_tok = 1 if model_arch == "Llama4ForConditionalGeneration" else 2 update_dict.update( { "num_experts": num_experts, - "num_experts_per_tok": 2, + "num_experts_per_tok": num_experts_per_tok, # Kimi uses `num_experts_per_token`. - "num_experts_per_token": 2, + "num_experts_per_token": num_experts_per_tok, "num_local_experts": num_experts, # Otherwise there will not be any expert layers "first_k_dense_replace": 0, diff --git a/vllm/model_executor/models/mllama4.py b/vllm/model_executor/models/mllama4.py index 742dccc36f1..797826c6bf5 100644 --- a/vllm/model_executor/models/mllama4.py +++ b/vllm/model_executor/models/mllama4.py @@ -19,7 +19,7 @@ import math from collections.abc import Iterable, Mapping from itertools import tee -from typing import Annotated, Literal +from typing import Annotated, Any, Literal import torch from torch import nn @@ -78,6 +78,7 @@ from .interfaces import ( MixtureOfExperts, MultiModalEmbeddings, SupportsEagle3, + SupportsEncoderCudaGraph, SupportsLoRA, SupportsMultiModal, SupportsPP, @@ -105,7 +106,7 @@ class Llama4ImagePatchInputs(TensorSchema): patches_per_image: Annotated[torch.Tensor, TensorShape("batch_size")] """ - The number of total patches for each image in the batch. + The number of chunked image tiles for each image in the batch. This is used to split the embeddings which has the first two dimensions flattened just like `pixel_values`. @@ -731,6 +732,7 @@ class Llama4ForConditionalGeneration( SupportsMultiModal, SupportsPP, MixtureOfExperts, + SupportsEncoderCudaGraph, SupportsEagle3, SupportsLoRA, ): @@ -828,10 +830,161 @@ class Llama4ForConditionalGeneration( num_physical_experts, num_local_physical_experts ) + def get_image_patches_per_chunk(self) -> int: + return Mllama4ProcessingInfo.get_patch_per_chunk(self.config.vision_config) + + def encode_image_chunks( + self, + pixel_values: torch.Tensor, + *, + use_data_parallel: bool, + ) -> torch.Tensor: + if use_data_parallel: + vision_embeddings = run_dp_sharded_vision_model( + pixel_values, self.vision_model + ) + else: + vision_embeddings = self.vision_model(pixel_values) + + return self.multi_modal_projector(vision_embeddings) + + def get_encoder_cudagraph_config(self): + from vllm.v1.worker.encoder_cudagraph_defs import ( + EncoderCudaGraphConfig, + ) + + return EncoderCudaGraphConfig( + modalities=["image"], + buffer_keys=["pixel_values"], + out_hidden_size=self.config.text_config.hidden_size, + ) + + def get_input_modality( + self, + mm_kwargs: dict[str, Any], + ) -> str: + return "image" + + def get_encoder_cudagraph_budget_range( + self, + vllm_config: VllmConfig, + ) -> tuple[int, int]: + min_budget = self.get_image_patches_per_chunk() + max_budget = min( + vllm_config.scheduler_config.max_num_batched_tokens, + self.vllm_config.model_config.max_model_len, + ) + return (min_budget, max_budget) + + def get_encoder_cudagraph_item_specs( + self, + mm_kwargs: dict[str, Any], + ): + from vllm.v1.worker.encoder_cudagraph_defs import EncoderItemSpec + + patches_per_chunk = self.get_image_patches_per_chunk() + return [ + EncoderItemSpec( + input_size=num_chunks, + output_tokens=num_chunks * patches_per_chunk, + ) + for num_chunks in mm_kwargs["patches_per_image"].tolist() + ] + + def select_encoder_cudagraph_items( + self, + mm_kwargs: dict[str, Any], + indices: list[int], + ) -> dict[str, Any]: + pixel_values = mm_kwargs["pixel_values"] + patches_per_image = mm_kwargs["patches_per_image"] + + if len(indices) == 0: + return { + "pixel_values": pixel_values[:0], + "patches_per_image": patches_per_image[:0], + } + + cum_chunks = [0] + for num_chunks in patches_per_image.tolist(): + cum_chunks.append(cum_chunks[-1] + num_chunks) + + selected_pixel_values = torch.cat( + [pixel_values[cum_chunks[i] : cum_chunks[i + 1]] for i in indices], + dim=0, + ) + + return { + "pixel_values": selected_pixel_values, + "patches_per_image": patches_per_image[indices], + } + + def prepare_encoder_cudagraph_capture_inputs( + self, + token_budget: int, + max_batch_size: int, + max_frames_per_batch: int, + device: torch.device, + dtype: torch.dtype, + ): + from vllm.v1.worker.encoder_cudagraph_defs import ( + EncoderCudaGraphCaptureInputs, + ) + + vision_config = self.config.vision_config + patches_per_chunk = self.get_image_patches_per_chunk() + chunks_per_capture = max( + 1, (token_budget + patches_per_chunk - 1) // patches_per_chunk + ) + dummy_pixel_values = torch.randn( + chunks_per_capture, + vision_config.num_channels, + vision_config.image_size, + vision_config.image_size, + device=device, + dtype=dtype, + ) + + return EncoderCudaGraphCaptureInputs( + values={"pixel_values": dummy_pixel_values}, + ) + + def prepare_encoder_cudagraph_replay_buffers( + self, + mm_kwargs: dict[str, Any], + max_batch_size: int, + max_frames_per_batch: int, + ): + from vllm.v1.worker.encoder_cudagraph_defs import ( + EncoderCudaGraphReplayBuffers, + ) + + return EncoderCudaGraphReplayBuffers( + values={"pixel_values": mm_kwargs["pixel_values"]}, + ) + + def encoder_cudagraph_forward( + self, + inputs: dict[str, torch.Tensor], + ) -> torch.Tensor: + return self.encode_image_chunks( + inputs["pixel_values"], + use_data_parallel=False, + ).flatten(0, 1) + + def encoder_eager_forward( + self, + mm_kwargs: dict[str, Any], + ) -> torch.Tensor: + return self.encode_image_chunks( + mm_kwargs["pixel_values"], + use_data_parallel=False, + ).flatten(0, 1) + def _parse_and_validate_image_input( self, **kwargs: object ) -> Llama4ImagePatchInputs | None: - # num_images, 1, num_chunks, channel, image_size, image_size + # total_num_chunks, channel, image_size, image_size pixel_values = kwargs.pop("pixel_values", None) if pixel_values is None: return None @@ -853,15 +1006,10 @@ class Llama4ForConditionalGeneration( pixel_values = image_input["pixel_values"] patches_per_image = image_input["patches_per_image"].tolist() - # shard image input - if self.use_data_parallel: - vision_embeddings_flat = run_dp_sharded_vision_model( - pixel_values, self.vision_model - ) - else: - vision_embeddings_flat = self.vision_model(pixel_values) - - vision_embeddings_flat = self.multi_modal_projector(vision_embeddings_flat) + vision_embeddings_flat = self.encode_image_chunks( + pixel_values, + use_data_parallel=self.use_data_parallel, + ) return [ img.flatten(0, 1) From fe042382925000e5adfe530a1cc2b91d7a125fd5 Mon Sep 17 00:00:00 2001 From: Rohan Potdar <66227218+Rohan138@users.noreply.github.com> Date: Fri, 12 Jun 2026 01:02:04 -0500 Subject: [PATCH 31/52] [ROCm][gpt-oss] Pass GateMode.INTERLEAVE for MXFP4 W4A16 fused MoE (#44893) Signed-off-by: Rohan Potdar Signed-off-by: Rohan138 Signed-off-by: Rohan Potdar <66227218+Rohan138@users.noreply.github.com> --- vllm/_aiter_ops.py | 24 +++++++++++++++++++ .../fused_moe/experts/rocm_aiter_moe.py | 16 +++++++++++++ 2 files changed, 40 insertions(+) diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index 1d75b7c7628..d744da0b89b 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -167,6 +167,7 @@ def _rocm_aiter_fused_moe_impl( output_dtype: torch.dtype | None = None, hidden_pad: int = 0, intermediate_pad: int = 0, + gate_mode: str = "", bias1: torch.Tensor | None = None, bias2: torch.Tensor | None = None, moe_sorting_dispatch_policy: int = 0, @@ -177,6 +178,10 @@ def _rocm_aiter_fused_moe_impl( activation = ActivationType(activation_method) quant_type = QuantType(quant_method) + extra_kwargs: dict = {} + if gate_mode and rocm_aiter_ops.fused_moe_supports_gate_mode(): + extra_kwargs["gate_mode"] = gate_mode + return fused_moe( hidden_states, w1, @@ -198,6 +203,7 @@ def _rocm_aiter_fused_moe_impl( bias1=bias1, bias2=bias2, moe_sorting_dispatch_policy=moe_sorting_dispatch_policy, + **extra_kwargs, ) @@ -219,6 +225,7 @@ def _rocm_aiter_fused_moe_fake( output_dtype: torch.dtype | None = None, hidden_pad: int = 0, intermediate_pad: int = 0, + gate_mode: str = "", bias1: torch.Tensor | None = None, bias2: torch.Tensor | None = None, moe_sorting_dispatch_policy: int = 0, @@ -1804,6 +1811,21 @@ class rocm_aiter_ops: except (ImportError, ModuleNotFoundError): return False + @classmethod + @if_aiter_supported + @functools.cache + def fused_moe_supports_gate_mode(cls) -> bool: + """Probe whether the installed aiter.fused_moe accepts `gate_mode`. + + Added in https://github.com/ROCm/aiter/pull/3123 (>=0.1.14). + Builds with older AITER must omit this argument. + """ + import inspect + + from aiter.fused_moe import fused_moe + + return "gate_mode" in inspect.signature(fused_moe).parameters + @staticmethod @if_aiter_supported def register_ops_once() -> None: @@ -2172,6 +2194,7 @@ class rocm_aiter_ops: output_dtype: torch.dtype | None = None, hidden_pad: int = 0, intermediate_pad: int = 0, + gate_mode: str = "", bias1: torch.Tensor | None = None, bias2: torch.Tensor | None = None, moe_sorting_dispatch_policy: int = 0, @@ -2194,6 +2217,7 @@ class rocm_aiter_ops: output_dtype, hidden_pad, intermediate_pad, + gate_mode, bias1, bias2, moe_sorting_dispatch_policy, diff --git a/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py b/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py index 5c2aa455600..bd9b285fe74 100644 --- a/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py @@ -351,6 +351,21 @@ def rocm_aiter_fused_experts( intermediate_pad // 64 * 64 * (2 if moe_config.tp_size == 1 else 1) ) + # https://github.com/ROCm/aiter/pull/3123 specialized the AITER stage1 GEMMs + # for interleaved vs separated gate and up weights. + # For gpt-oss i.e. use_mxfp4_w4a16=True, the weights are shuffled by + # `rocm_aiter_ops.shuffle_weight_a16w4` in `oracle/mxfp4.py`, + # which always sets `is_guinterleave=True`. + # Hence, we pass in GateMode.INTERLEAVE to match the weight shuffling. + gate_mode = "" + if quant_config.use_mxfp4_w4a16: + try: + from aiter.ops.flydsl.moe_common import GateMode + + gate_mode = GateMode.INTERLEAVE.value + except ImportError: + pass + return rocm_aiter_ops.fused_moe( hidden_states, w1, @@ -369,6 +384,7 @@ def rocm_aiter_fused_experts( output_dtype=output_dtype, hidden_pad=hidden_pad, intermediate_pad=intermediate_pad, + gate_mode=gate_mode, bias1=quant_config.w1_bias if quant_config.use_mxfp4_w4a16 else None, bias2=quant_config.w2_bias if quant_config.use_mxfp4_w4a16 else None, moe_sorting_dispatch_policy=moe_sorting_dispatch_policy, From a2c72d43883e21f3e36f3b970008d2394a714282 Mon Sep 17 00:00:00 2001 From: Isotr0py Date: Fri, 12 Jun 2026 15:10:18 +0800 Subject: [PATCH 32/52] [Bugfix] Fix Dockerfile dependency graph pre-commit error (#45374) Signed-off-by: Isotr0py --- .../dockerfile-stages-dependency.png | Bin 382338 -> 396782 bytes 1 file changed, 0 insertions(+), 0 deletions(-) diff --git a/docs/assets/contributing/dockerfile-stages-dependency.png b/docs/assets/contributing/dockerfile-stages-dependency.png index 0c7a8ab246ec7b5b49516b34a4d464228ab56dba..90aaf01a0b7e5a1ffc3af57e218efb517737f037 100644 GIT binary patch literal 396782 zcmZ_1XFycv)&)H3RTJx*8%wN29UF)!O#vw;K{F`DLa!=KiYSPRw24N&*M=Dc1f;}< z6a}e*G&L4LiXb4Mf=Uq)P!Q?xt$k*I!*{=L-XCw|!pu3(d7i!3UTf_=|JBjjJay9C zNgNJmD)0B-wsScD{)5B${_}Se@RQ$OsW#*PP1wC<^KYC%`oHju2wx88XAbYTjXRHA z{`A`IVwvmf!7t^{_$$_|dUtY5z2v^XZCg@L)QokmK76Ihm zAf2+4F!#5w~W&cH5(nA#(MAke?%Xpx5d=dbeY!W zOPj?!WlL^o=6!7}Y5y1muI+lYTY~3e7b;j^Wop` zyr^+~`|}@|caN#?mo)Z2CFZsM@M~B7$7kpGy0zEO4Ea})-N6qQV+si5Q z*ROG`N|sdoe&mw)kGJg+EwzN~;g(mgobx2R>!iE`1I)eLQr&R}L1r08t=jX~X*U#v z>@YQbo&Fa`Qq9rPvA!@YY+bd73|;-FiiCiBvF@EM3f;B7?YEpfC4wE|c{2K}3od#p zdU!RycJ}rz{T<~ELH4(EUZuLs+kmA?`Px)2HEzO`r?lp`t`API%n`J8dv?BFm-70M@u8QyPK|9Ec@7){P0`-G z@XC9zpVl#9|Jw&`saMY{iU*h;GE9DbsA_B;#~`%a%w1BoJ4f63$%4bR=7X=wvElmq z8dM7#Bf`w~FJQf8^i$z(( zKQ6DoxI-(ay-?L$7elhw~{NxeKxfH ze)Q{yZTshpyd6ivLb#RH>}BlBV(92drGNfKVEJ(Bx@Vu?K1{&U%`};i-5|H;@!9!! z0&d-P0lwVeeQA3ipPLuZoKl_nat-V4{GSZ1pCvxG@%6pKp>Fko+(D)Jmd`G&tvK*x z!Q|EB!tbQx;Wn1Xn%~e43*}ZgB;-pUc=wM(YMuZ10M^F*yjIRneKzy%MH_GeO0v8M z+c-BDDj&6MtIe69HYRZmKA^Dc%lm|?efO5~uvX{3wpOdU)IVE!cgo0B@qJ{3ku7}Y z3O@~nFaP`~v$yyfH|Iw1zM#zu-MGdMchv7&ab!<)l$)zw-Rx6vsLC;b8@P7Qrt=%x zp1vq8#R`nPTlVHT!Z24f$;7=m$*ty(i2|kDZI!D!rbApOb&f9ls?ViV$ z6a@rZFRfYB5a#W5Uw7nfI37DGz|UO3+^||n;F?{(^n7a$PD@MMsi$Y%_4%#Br?`g* zVBvvV1ErTvoWEa9rK{%n&Ye4J2YOoz|JuCJ@Zc=g+wnB5o(rR|Y{;WCvu!@t+>eZo z-XbtL_*|WqO4_|8)#moSmNYIea%b+|h!Fq7Z`>kqts^giDAe+2ecG(rv(38Tav+x; z;p)t;cgN0fRy>$vbNR2UJ}Qd5ww8OX^`8<7SKs|^fb-8WR(9Uq~<6KuP>>ALsez<7Puvrn`edG@je zs)Kzm%FCDYyn0H)=5odt({Dgq#)|Ru}mz&j{@8pbAJj>)?BX?9}SGseS zDf4f&_xCkDbY2lO_oN)NFN;N8tNhr11>m0z9(uan}@qddPrX_m>j?8+-!0#@?|T0L3&m|rgfENYYQ z*gn6)<>4K-8#b0_eU|$5)HF?X^+(QGUf2n-Ie)6mtCaeW1>gmejGxuxsAHV|cImK7w`Tniu$HxtJj9$34B5Y~ z8un|O{+7=LO0Bgu$ERz#y1G(ceNk4n`fzogRF=s=rc`RJVA)=`f$oO9&h>+R?ZLGR zN1isPXuD{=IC1{njoRKpclJE7ett#W?eqPK?d@TM=^wbP->#7vLbfaZQE0gAd;0Wg zgwbo7!Sd=2F}fd{n&!ncef&6yjfpra(?uJZQzkvz+~dH}&X$mG%~dwp{4w1)k;VuA z&H8}PbGi&1nSb+>wsy9>mhtNAx+25+1%6><{0G**MgPe!6Tyv9|I-MLw!ef$#W7YT zl=pa^(`DNJ@6W9(_+j(>za}m`^0S!7=Z%X@&&-x82s!*FaEZ?bY2~9|f5&pBI%f5M z%X4c830NNHA&o%(f#!OHwgb$9teX z%;jTbg!R4V-QEykp5>*`^WmzD-HV&)h|$PJc@*>~ig~Qryufv53qY%hOTL^$dO_G=Uqjs+ z_nPcK@?KXXLHHf=S)e3e<(OhhTcz*qbnmCn0%a1gahqf8o}Kx?TOsqD z(%0hl2mW!T-)-}Z&G=T?){Q_i(@XW&lT*4sseaat!(UFBE4zX)J$>v3Y$wAwi;c4+ zjV{hMJF*Cf+~r%Pee7kgo<@oJ*`@PipBOu3Bt3Bdlv3?=?pu4;&C1Hk`QCkBDu6Ow zdYe-kTD=F|fk$M#2l{|&me3h$eUsr*%v&zm*}+N$9F-7J5QsZzfMSO>IS|DX^$6L| z|NN+JiOu8&ytf0V*U7WVI8b))_WJzb06+lT@bash&gCvTHDAE(ZcrTzOej#4IQ;6s z=KFOL`r>}hU!QMoARGV$8ydMI!usul6M^fT%ogG@QM1IQw_0+50 zoeSk147wGb;19&tO`juuZ9~_G(-OCzYOZ^IF{AU1Qi#%FEnQvRG}k8of&N|_tS!)N zW_x7$T*ZU7sc(C$2feyI@C)62PqOxe13l5mer}XKbG%3U%tmzQ-1Gfm&U4;)+q0*VfQO+!a~0P zaoNT4auXNLmeht&pc@sO5QwOPoa}uc zV>Y$pGhdvm%Po6j3O%nED(!uJ5CMmWHlw>GFk5C?dbTac{MDI*N~;8`Dr2 zHW=AkZfFzY(GhL#n0sy0VXTDL$M-X5&Yb!U0IUHq$l+tmFQzD&SNdtxbs)f`6gOP< z9_ZE#9JL;T;q~}w!;$?_(b4){UBJV%Z*o^g$QthmSZz8tMk?cL@q~we_&4ApG=WMl zEc!J_`qnn!-f0y_n{Mf+#R3-Ytx9%u^DI>Eyq1}nxw2%^4rAlsT@O$Cm17}|uHQGK zvY(Qc{4RD~lej^&XS3(N3qL+hT=DIXN56dvBz&Hpi4r%yHM6(=?^6pkKq~xUb#6~X z$YHbVo_I?QpM}b~O(!PLuE!>@kBa$4Gel7iDOED|aOK{J3FG2aG-J&&g2X4YSC~7z z!gu<=Rm~3h{5^``RQzWG(x-o+W{~X4aTBLK+ceMd;`{STNOGQssl*;Dp?a_7^}X1( zH$Kz;&=541M65ehF!wF?qGQ|}T7d3UsYsZA`=%3BufP+ao`X$w_p!`Bk+LWFMfB!fhYZV83hK(7PJ_jr8VqyFq{FPD|W-fv&!(A2=KZ z@2-bpZP%qdBm0gBH>EUKp2Fas@B4c?wnXkUYAm1b_4&A;S=0ODAu9QO-R<*Y%FVs| zOZS=oRoHt}_Qt-~X>taI!Rws-!i@_5|%2zs`a8VuQAEEv3=r&Yd&+k9wGi4Q2um zr{)%#hJ-C=eU7NC=freWrZ`nwWLVrk=BLu|Xtue@b^Gb_L@orrP@{!4Mlz zT=}rBcPQ(Vvj?R`Sg*`^oaG5%r#iZjjS)&36@AJb6K84t6Sf$9DH=lN9c4m}*G}R?wg>bIos*Qb7s&8&;rmw4O z_4L9biC5M0V>jE!CCqEO0G1O?m7o##)@X8$+9~%o(+fW|mifZquoS zM_0tAx8*u_YN67n1SDVOkWiXg7-$n3NYDL*j}&SbVz6?XP@l+y49K3;>4Z9NYMGeW z&vI6|7u&vQArPjeMqk4w(l_^X1wE7UKpFUE7cnqxJ+E8}H48(P_3vh-n)%#FLBAJjccowXWviK!1BitEW`qhW4$qq|7c;{06o76baY`S9P&d4A^Jt1k+o;yQDAF zY$h7@kBJKn%)&r|V3X~U1|hTxyDBLuDKauL2Px3+k8u;NQEC^gjJUkrVlC?##(0h{ zTft-0J9+mSgUY)IYhxPzx&QnUhz0Fd<~{P>;o3zkV!b;2$9YzZRfe`@Dyt z{t>9fafb50JQ2&%=rjuj&09aOy6dsC*u7f87Ux$7o>&!yZ*&8q*E@0!;pQ^3Z|axF z^Gd)A=%zXC$y(@?(V>Sd0cI{=L?x!;S+!@6n{HW|xYrE;9|Wgkvf;E1pR=bd+;AjGDd=5??LLY= zRArTDDM=FOY>!awUzuQ?uVH0nMKrAk!b<@ln;+D-RxC%I@EuRS$ei zIr#C#%JzmrDaDEWq(3`Zn2rD7#Srnt+ISpJ(cGcv#{RJmXD9@Q&y?`fP-o|XL?aOn zO8mbCiU)repvIp2VNABTaiZ-u!8KGGHLOCGUp7SF&SF>u0SYPn0YfY7ev|I@|4%DC zhosZ?1<+$MQ5@S)1FD~ywRCUSXcAzUJ$bqNjAZT;sm!k340~ec@#PUH0q)f5&8+*f zT{wK{8Am(B#k7OvZ)N|RQ#MP?Gpp-EFN^Qv&*>$t_Siv`UeMq{HQ+B3Gp`Z_a;=LeKBa>a?5bnV%m_o*aWqXn!a+um<6i zCqPaq1Q|K};*TB0f!@Z({RB8l9KN8I@)R-opnO+U{;0j-uUM~+Xw)x_Z!+9z>p}4Q z8;d&p;+I9;i2op5ley-&jMZ6vZEZaO)raZ}LX0hb;HVqNTkB~ zmh!2yW@z8sL=veda4M(;Sbc^{c*lPHbx-TyKr7J@Jb~x8isjXP_3I|DW<7)OD_>)M zdYE?%f8+)Up(Mek5X9g*o-^u{`agfvT5)$0zg`|R6A|0AKY+=6u@%1pK#jhY+rzgK zqy=iKJt3?=CVikRCDRWO;CA+6PEb`-Y*y&T;}fofr{}c+V7?r+vU0zT(kT>)CctCl z@*4CUSX)VPnGX;g#iZ$~baEr+Vxb|5Z7Z z5F$d4iu2;Z3Rr@nLe0JU?iBtF(iJ_L?alK(JiqG8HQtwM6?*hzsTzFLZr8#qz+<+takc-$Ss7QyX* z8hU8|MP`%sQT}EiUuk=ln*+fk$g97{YQSS48{tXl1?cOBQdM+`kRlb`s7BlcT%ZNB zh0xt-=ZzM!4h-G%VgNhMfvzX2H&Ai$1W&}gH-N{pAV7L7k7GM!@~lG~)fjk>E08PTDDI{d=L%6;T_z#|7>bh41 z#S2kN8;uec!yb!5vYs8{YCZ^i7vQ?(JX$gjXfRq}%-V4>Ginq$dXG>O`u(W2ei1f- z1DG{m?!8a(Td1yt27eX5G4j$IeqRUDC~t1N>iov7;rf*;i+R$!9{kBc<~ATfynnU-;Z!0&C&#cKnp~+U zA*dO`3*dz!vmOeh1`3YuKI>pd{~cNiOK}wJ2>g@K9)WCSgV03Yf(nbS9Y|4H8@~KQ zqN5>dsmZHYjLR^^&DTB(&c^w<`fRJ0al5)d#Qs&4xYrkfuAzFczgos3n-A&zYPJU7 zzd5PgEbr-GOJsKc^PjUIT!~9Q{$#r1hj0_G?z)OSKFX8C*`MGdN~8Rl!H~6X-+gy_ zikQj=cSC zm8W$}24%ABUSYgd9^{Y;%kqgM?-M>obhr(a!4*FJ&kswaAqsBj|5_3Vd=4c^QVX7d z1YCXuE_u!x3pFy5KqQcDR2IXDF+u8_qkQxrWT{uzI$4qW=l_bbrH2_I=}?4oqHf#k z^YhZ(+ND7dUr8u40DY7n9h(_MJ_PgbT>WCe?<>9re^}-D*PWPc_H>7~vL78J6oR(T zd5JpETCi)kpt=@-OSo=R7+L|q6qE!iyzAEqQ|Dg9+FDcT1WJv7u(eQp5Z$I5EkP??Q1VRLM!9rBHYAN=es@$+K5Kn`P|q7xDE_un8Z;qdJ%!GG6l? zsDfu>CeFW*^5&=nuQ)?{gIif~ZMM_!Uf!?JTxF;dY)HxSb<(8g>?K>(HJTxMuf z_x6-n0jZHFYhH0~CaGc5bh{FB=}cOt$-gvhfKiw_kQ8`c+si#ZR*(Y-=fle|@@cje z?8?iaDg%@#xg<9EEDSmD{_*wXx|q_iY6yUmaa$VSMAsP5yRZ<>dD zG|K^1WPEUzY~acrH)_ES<~}$v*_TVy3j#zYbm|-I9K)X#DT+Cqt|?H9+Mmf3T0);E z>6WOPc=J?ZkU&~AR%dz3=obG(j|g=1I$m3}dB6DaiPNEz2bKA1s|5Z!U0L?bEJ+Ok z=?TOqi%Uqr5y85}14BoixMB7Ul2UlKbyEhds0DPdW;^K*sgsix_Nyb0$2OmkMqMk7Op-d{;AkOyPc!z zOMA}q1Y|krhD2$y{|DJH^S*bp=SW-r?&g^mTYPpTdzmCW_Xv_v~ta(VnI(3@!&?dR}0=)Q$V5Iv-g|i-AO@(;GY6oJzL#*0@VJmKqxD8 zyx0LSKJpO;(Zdc0PW~;h?Q}K1)eX=E*Kwl8N5;fFLx3_cGY3o_U#8Yh3|ISAvHlgj zVi3R#8e2T_Zu}&3^~gxo`uWU{%G3zpLF@U8%4JRGeyW^PUO8{^H3@#V3Tsa}o+2hU z2jou)Xv$;J;clk;?~f%p&=7JgvQq~3pbXUp32fa*SNE?U{)!0m zQbal7@S-W+`Y!5u8QY?8ZVnEMFPvJVTpgU4>vrO2Rx@wSzXHT8L-gNOvjhmj-OWzk z_MDxQzmbB2GJUIH1JJ^pb^A<4nKZI1b4ALhQBo_OEm?KuHXj9Ve2Cq3-deO^F7U5mp@!`yD$XVXD;99#82!Jk~5&YDV$l}pL4+M2ch)o z3-cO?ioTtTYfKz_hV@iAwoPY+;$mE`dwrlyJc^1sY~=okP}SNCzdC1+nHV$5q{OLA zGs2oQ3d{c$Ryf*o6r3%|Hpt}xceEjd-}w9Qx8szva+U>Ke*08m2sdG{%#W-O