Compare commits

...
Author SHA1 Message Date
Tyler Michael SmithandClaude 5ad0151d64 Reduce SP correctness test matrix from 32 to 4 cases
SPTestSettings.fast() was identical to detailed(), generating the full
cross-product of eager/compiled × chunked/no-chunk × pp1/pp2 × mp/ray
(8 setups × 2 backends = 16 combos, ×2 for inductor = 32 tests).

Slim it down to 2 representative setups (compiled + chunked prefill,
with pp=1 and pp=2) and a single backend (mp). Remove the unused
detailed() method entirely.

Co-authored-by: Claude <noreply@anthropic.com>
Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>
2026-06-20 15:02:07 -04:00
Tyler Michael Smith 8019b6ec63 factor out _build_anthropic_usage helper
Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>
2026-06-20 14:39:13 -04:00
mistral0105 eb3ee8afc3 Merge branch 'main' into anthropic-cache-usage
Signed-off-by: mistral0105 <zhangshuoming17@mails.ucas.ac.cn>
2026-06-19 12:44:27 +00:00
mistral0105 77648ab261 Merge fork branch updates 2026-06-19 12:40:02 +00:00
mistral0105 fd3e0cac12 Address review on Anthropic cache usage reporting
- api_router: stop silently overriding --enable-prompt-tokens-details for
  AnthropicServingMessages; pass through the user's CLI setting like the
  other serving objects.
- _compute_cache_usage: rewrite docstring to document where
  prompt_tokens_details is attached in vLLM's OpenAI streaming path
  (terminal include_usage chunk only), why message_start cannot populate
  cache fields today, and why cache_creation_input_tokens defaults to 0
  rather than None when cache info is present.
- AnthropicUsage construction: omit cache fields entirely when the
  underlying cache info is unknown (cache_read is None), rather than
  emitting null. Applied uniformly to non-streaming responses,
  message_start, and message_delta so "unknown" is signaled by key
  absence rather than null, distinguishing it from a real zero.
- Tests: add TestStreamingCacheUsageSemantics covering the three usage
  states (cache hit, cache miss with details, no details at all) for
  both message_start and message_delta.

Signed-off-by: mistral0105 <zhangshuoming17@mails.ucas.ac.cn>
2026-06-19 12:40:01 +00:00
shuoming zhangandmistral0105 82d1ddf39e Merge branch 'main' into anthropic-cache-usage 2026-06-19 12:40:01 +00:00
69bdd34542 [Bugfix] Fall back to Pydantic loc for param in validation errors (#46038)
Signed-off-by: professorsab <135441198+professorsab@users.noreply.github.com>
Co-authored-by: Mahad Durrani <114791389+mahadrehmann@users.noreply.github.com>
2026-06-19 19:11:11 +08:00
Kunshang JiandGitHub ec67d7ae61 [xpu] bump up vllm-xpu-kernels v0.1.10 and upgrade 2618 umd (#40367)
Signed-off-by: Kunshang Ji <jikunshang95@gmail.com>
Signed-off-by: Kunshang Ji <kunshang.ji@intel.com>
2026-06-19 15:37:20 +08:00
mistral0105 7341ff152f Merge branch 'main' into anthropic-cache-usage 2026-06-19 04:33:55 +00:00
ecf9d83520 [AMD][CI] Fix Language Models Test (Extended Generation) failures (#45509)
Signed-off-by: Oxana Korzh <okorzh@amd.com>
Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
2026-06-19 12:06:56 +08:00
Samuel ShenandGitHub c9135db27c [Docs] Update stale LMCache examples (#45762)
Signed-off-by: Samuel Shen <slshen@tensormesh.ai>
2026-06-19 03:21:36 +00:00
Jeff (Junze) MaGitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
2a6c6b9429 [DeepSeek-V4] Support TEP=16 for the block-FP8 shared expert (#46001)
Signed-off-by: Jeff Ma <jeffjma@umich.edu>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
2026-06-18 20:10:12 -07:00
Jared WenandGitHub ab66606993 [bugfix]Indexer init skip and MTP TopK share for iteration (#45895)
Signed-off-by: JaredforReal <w13431838023@gmail.com>
2026-06-19 09:57:51 +08:00
9ea3a4015b [Bugfix] Fix corrupt outputs in MoE FP8 LoRA responses and MoE base model responses when LoRAs are loaded (#42120)
Signed-off-by: Nicholas Edelman <nedelman@nvidia.com>
Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai>
Co-authored-by: Jee Jee Li <jeejeelee@inferact.ai>
Co-authored-by: Jee Jee Li <pandaleefree@gmail.com>
2026-06-18 18:26:09 -07:00
Flora FengandGitHub 560fb8b867 [Cohere] Remove dead prepare_structured_tag override in Cohere parser (#46099)
Signed-off-by: sfeng33 <4florafeng@gmail.com>
2026-06-19 01:02:11 +00:00
Wentao YeandGitHub 675cd5d228 [Model Runner V2] Fix MRv2 memory leak test (#46095)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
2026-06-19 00:36:40 +00:00
7f616c327d [Bugfix] [Parser] Fix empty tool block silently dropping subsequent content (#46091)
Signed-off-by: Ben Browning <bbrownin@redhat.com>
Co-authored-by: Flora Feng <4florafeng@gmail.com>
2026-06-18 23:17:18 +00:00
Ivy XuandGitHub c3c6d723fd [Perf] Remove unused loggers in reasoning/ (#45988)
Signed-off-by: Ivy <fakeshadow1337@gmail.com>
2026-06-18 22:24:29 +00:00
Yifan QiaoGitHubJingyi Yangmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
41dcf49ca5 [Bugfix][KV Connector] Disable Mooncake TP put-striding when DCP > 1 (#45371)
Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai>
Co-authored-by: Jingyi Yang <girasoleyang@gmail.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
2026-06-18 15:13:44 -07:00
35e4dd4a69 [KV Connector][Mooncake] Async lookup to reduce scheduler overhead (#45659)
Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai>
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Co-authored-by: Nick Hill <nickhill123@gmail.com>
2026-06-18 21:44:02 +00:00
4ce2d01453 fix(anthropic): auto-detect template support for mid-conversation system messages (#46025)
Signed-off-by: felix0080 <felix0080@users.noreply.github.com>
Signed-off-by: Ben Browning <bbrownin@redhat.com>
Co-authored-by: felix0080 <felix0080@users.noreply.github.com>
Co-authored-by: Ben Browning <bbrownin@redhat.com>
2026-06-18 16:19:11 -04:00
Woosuk KwonandGitHub 16908e132e [MRV2] Make FP32 Gumbel sampling more accurate (#45996)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
2026-06-18 19:42:09 +00:00
Wentao YeandGitHub 225936a1dd [CI Bug] Revert #42379 to fix CI Multi-Modal Models (Extended Generation 1) (#46070)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
2026-06-18 12:37:39 -07:00
f6ba720963 (security) Upgrade Starlette to >= 1.0.1 to fix CVE-2026-48710 (#45675)
Signed-off-by: jperezde <jperezde@redhat.com>
Co-authored-by: Isotr0py <mozf@mail2.sysu.edu.cn>
2026-06-18 12:35:13 -07:00
Wentao YeandGitHub b53b1c7ffe [Model Runner V2] Migration to support quantized model by default [5/N] (#44446)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
2026-06-18 12:20:44 -07:00
Ting SUNGitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
79ca54d221 [Bugfix][Quantization] Don't reject fp8_e5m2 KV cache for non-fp8 quantized checkpoints (#45040)
Signed-off-by: Ting Sun <suntcrick@gmail.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
2026-06-18 14:18:25 -04:00
Ben BrowningandGitHub 09f3cd5c10 [Bugfix] [Parser] Fix Qwen3 latent bug in partial params dropping values containing < (#46047)
Signed-off-by: Ben Browning <bbrownin@redhat.com>
2026-06-18 18:04:06 +00:00
ea6078fe6a [KV Connector][Offloading] Disable parallel-agnostic fs-tier cache on V2 model runner (#46044)
Signed-off-by: Itay Etelis <etelis2019@gmail.com>
Co-authored-by: Itay Etelis <etelis2019@gmail.com>
2026-06-18 20:43:35 +03:00
Palaiologos1453andGitHub a0df04e477 [Tests] Add Qwen3 streaming parser delta boundary cases (#45708)
Signed-off-by: test test <2260891073@qq.com>
2026-06-18 17:37:39 +00:00
stefankoncarevicandGitHub e2352c2974 [ROCm][Spec Decode] Fix probabilistic draft probs test attention backend (#45706)
Signed-off-by: Stefan Koncarevic <stefan.koncarevic@amd.com>
2026-06-18 11:59:37 -05:00
qli88andGitHub 25faa1f4cc [CI]Enable mxfp4 lora test for ROCm platform (#43802)
Signed-off-by: Qiang Li <qiang.li2@amd.com>
2026-06-18 16:59:09 +00:00
HumphreyandGitHub 4583630b56 [Bugfix][Kernel] Check output alignment in vectorize_with_alignment (fixes misaligned-address crash for non-multiple-of-8 head sizes) (#45466)
Signed-off-by: HumphreySun98 <humphreysun98@gmail.com>
2026-06-18 16:58:22 +00:00
Divakar VermaandGitHub 21da47dabe [ROCm][CI] move lora%N test to mi300 and gate (#45970)
Signed-off-by: Divakar Verma <divakar.verma@amd.com>
2026-06-19 00:50:32 +08:00
ChaunceyGitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
6c379b9e54 [Frontend] Add Streaming Parser Engine and new GLM4.7/GLM5.1/GLM5.2 Parser (#45915)
Signed-off-by: chaunceyjiang <chaunceyjiang@gmail.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
2026-06-19 00:42:10 +08:00
Rohan PotdarandGitHub 5099474633 [Bugfix][ROCm] Fix rocm_aiter_per_tensor_quant custom op aliasing (#45747)
Signed-off-by: Rohan138 <rohanpotdar138@gmail.com>
2026-06-18 11:30:21 -05:00
Yuwen ZhouandGitHub 058cc0a8b6 [Bugfix] Restore is_sym guard for zp in GPTQ/CT MoE to fix symmetric quant regression (#45656)
Signed-off-by: yuwenzho <yuwen.zhou@intel.com>
2026-06-18 16:20:29 +00:00
837db7605e [Bugfix][Tool Parser] Handle non-finite numbers in coerce_to_schema_type (#43984)
Signed-off-by: ashishpatel26 <shriganesh.patel@gmail.com>
Co-authored-by: Ben Browning <bbrownin@redhat.com>
2026-06-18 16:00:20 +00:00
mistral0105 383a950d04 Merge branch 'main' into anthropic-cache-usage 2026-06-03 16:29:46 +00:00
shuoming zhangandGitHub ef8a54a77e Merge branch 'main' into anthropic-cache-usage 2026-06-02 12:29:42 +08:00
mistral0105 c0089373bb Merge branch 'main' into anthropic-cache-usage 2026-06-02 04:20:05 +00:00
shuoming zhangandGitHub c8f8f1951a Merge branch 'main' into anthropic-cache-usage 2026-04-27 01:17:42 +08:00
shuoming zhangandGitHub a50380e5d2 Merge branch 'main' into anthropic-cache-usage 2026-04-26 19:42:04 +08:00
mistral0105andClaude 04009ff40b [Frontend] Report cache usage in Anthropic /v1/messages API
Populate cache_read_input_tokens and cache_creation_input_tokens in
the Anthropic Messages API response, which were previously always None.

Key changes:
- Add _get_cached_tokens() and _compute_cache_usage() helpers to map
  vLLM's prefix cache hits to Anthropic's usage format
- Fix input_tokens semantics: Anthropic defines total_input =
  input_tokens + cache_read + cache_creation, so input_tokens must
  exclude cached tokens (previously it included them)
- Set cache_creation_input_tokens to 0 when cache info is available
  (vLLM's prefix caching only tracks cache reads, not writes)
- Force enable_prompt_tokens_details=True for AnthropicServingMessages
  so cache fields are always populated regardless of CLI flag
- Cover all three AnthropicUsage construction sites: non-streaming
  full response, streaming message_start, and streaming message_delta

Fixes #33923

Co-authored-by: Claude
Signed-off-by: mistral0105 <zhangshuoming17@mails.ucas.ac.cn>
2026-04-26 11:31:12 +00:00
94 changed files with 2596 additions and 2500 deletions
+1 -1
View File
@@ -67,7 +67,7 @@ steps:
pytest -v -s v1/worker --ignore=v1/worker/test_gpu_model_runner.py --ignore=v1/worker/test_worker_memory_snapshot.py &&
pytest -v -s v1/structured_output &&
pytest -v -s v1/test_serial_utils.py &&
pytest -v -s v1/spec_decode --ignore=v1/spec_decode/test_max_len.py --ignore=v1/spec_decode/test_speculators_eagle3.py --ignore=v1/spec_decode/test_acceptance_length.py &&
pytest -v -s v1/spec_decode --ignore=v1/spec_decode/test_max_len.py --ignore=v1/spec_decode/test_speculators_eagle3.py --ignore=v1/spec_decode/test_acceptance_length.py --ignore=v1/spec_decode/test_speculators_correctness.py &&
pytest -v -s v1/kv_connector/unit --ignore=v1/kv_connector/unit/test_multi_connector.py --ignore=v1/kv_connector/unit/test_example_connector.py --ignore=v1/kv_connector/unit/test_lmcache_integration.py --ignore=v1/kv_connector/unit/test_hf3fs_client.py --ignore=v1/kv_connector/unit/test_hf3fs_connector.py --ignore=v1/kv_connector/unit/test_hf3fs_metadata_server.py --ignore=v1/kv_connector/unit/test_offloading_connector.py'
- label: "XPU server test"
depends_on:
+14 -16
View File
@@ -415,22 +415,6 @@ steps:
commands:
- pytest -v -s kernels/mamba
#----------------------------------------------------------- mi250 · lora ------------------------------------------------------------#
- label: LoRA %N # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
agent_pool: mi250_1
parallelism: 4
optional: true
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/lora
- tests/lora
- vllm/platforms/rocm.py
commands:
- pytest -v -s lora --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --ignore=lora/test_chatglm3_tp.py --ignore=lora/test_llama_tp.py --ignore=lora/test_qwen3_with_multi_loras.py --ignore=lora/test_olmoe_tp.py --ignore=lora/test_deepseekv2_tp.py --ignore=lora/test_gptoss_tp.py --ignore=lora/test_qwen3moe_tp.py --ignore=lora/test_qwen35_densemodel_lora.py
#------------------------------------------------------ mi250 · models / basic -------------------------------------------------------#
- label: Basic Models Test (Other CPU) # TBD
@@ -1699,6 +1683,20 @@ steps:
#----------------------------------------------------------- mi300 · lora ------------------------------------------------------------#
- label: LoRA %N # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
agent_pool: mi300_1
parallelism: 4
optional: true
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/lora
- tests/lora
- vllm/platforms/rocm.py
commands:
- pytest -v -s lora --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --ignore=lora/test_chatglm3_tp.py --ignore=lora/test_llama_tp.py --ignore=lora/test_qwen3_with_multi_loras.py --ignore=lora/test_olmoe_tp.py --ignore=lora/test_deepseekv2_tp.py --ignore=lora/test_gptoss_tp.py --ignore=lora/test_qwen3moe_tp.py --ignore=lora/test_qwen35_densemodel_lora.py
- label: LoRA TP (Distributed) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
+11
View File
@@ -12,6 +12,17 @@ steps:
commands:
- pytest -v -s lora --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --ignore=lora/test_chatglm3_tp.py --ignore=lora/test_llama_tp.py --ignore=lora/test_qwen3_with_multi_loras.py --ignore=lora/test_olmoe_tp.py --ignore=lora/test_deepseekv2_tp.py --ignore=lora/test_gptoss_tp.py --ignore=lora/test_qwen3moe_tp.py --ignore=lora/test_qwen35_densemodel_lora.py
parallelism: 4
mirror:
amd:
device: mi325_1
working_dir: "/vllm-workspace/tests"
timeout_in_minutes: 60
source_file_dependencies:
- vllm/lora
- tests/lora
- vllm/platforms/rocm.py
depends_on:
- image-build-amd
- label: LoRA TP (Distributed)
+6
View File
@@ -21,6 +21,12 @@ steps:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
# TODO: create another `optional` test group for slow tests
- pytest -v -s -m 'not slow_test' v1/spec_decode
mirror:
amd:
device: mi300_1
timeout_in_minutes: 65
depends_on:
- image-build-amd
- label: V1 Sample + Logits
key: v1-sample-logits
+7 -6
View File
@@ -81,11 +81,11 @@ __global__ void rms_norm_kernel(
#pragma unroll
for (int j = 0; j < VEC_SIZE; j++) {
float x = static_cast<float>(src1.val[j]);
scalar_t normalized = static_cast<scalar_t>(x * s_variance);
if constexpr (HasWeight) {
dst.val[j] = normalized * src2.val[j];
float w = static_cast<float>(src2.val[j]);
dst.val[j] = static_cast<scalar_t>(x * s_variance * w);
} else {
dst.val[j] = normalized;
dst.val[j] = static_cast<scalar_t>(x * s_variance);
}
}
v_out[i] = dst;
@@ -151,7 +151,8 @@ fused_add_rms_norm_kernel(
#pragma unroll
for (int j = 0; j < width; ++j) {
float x = Converter::convert(res.data[j]);
out.data[j] = Converter::convert(x * s_variance) * w.data[j];
float wf = Converter::convert(w.data[j]);
out.data[j] = Converter::convert(x * s_variance * wf);
}
} else {
#pragma unroll
@@ -198,8 +199,8 @@ fused_add_rms_norm_kernel(
for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) {
float x = (float)residual[blockIdx.x * hidden_size + idx];
if constexpr (HasWeight) {
input[blockIdx.x * input_stride + idx] =
(scalar_t)(x * s_variance) * weight[idx];
float w = (float)weight[idx];
input[blockIdx.x * input_stride + idx] = (scalar_t)(x * s_variance * w);
} else {
input[blockIdx.x * input_stride + idx] = (scalar_t)(x * s_variance);
}
@@ -66,8 +66,13 @@ __global__ void rms_norm_static_fp8_quant_kernel(
#pragma unroll
for (int j = 0; j < VEC_SIZE; j++) {
float x = static_cast<float>(src1.val[j]);
// Multiply in weight's native dtype to match rms_norm_kernel.
scalar_t out_norm = static_cast<scalar_t>(x * s_variance) * src2.val[j];
float w = static_cast<float>(src2.val[j]);
// Round normalized result through scalar_t to match the precision of the
// unfused composite (rms_norm writes scalar_t, then
// static_scaled_fp8_quant re-loads it as float before FP8 conversion).
// Without this round, the fused path is strictly more accurate and
// disagrees with the composite at exact E4M3 quantization tie boundaries.
scalar_t out_norm = static_cast<scalar_t>(x * s_variance * w);
out[blockIdx.x * hidden_size + idx * VEC_SIZE + j] =
scaled_fp8_conversion<true, fp8_type>(static_cast<float>(out_norm),
scale_inv);
@@ -137,8 +142,12 @@ fused_add_rms_norm_static_fp8_quant_kernel(
#pragma unroll
for (int i = 0; i < width; ++i) {
float x = Converter::convert(res.data[i]);
// Multiply in weight's native dtype to match fused_add_rms_norm_kernel.
HipT out_norm_h = Converter::convert(x * s_variance) * w.data[i];
float wf = Converter::convert(w.data[i]);
// See note in rms_norm_static_fp8_quant_kernel: round through scalar_t
// to match the unfused composite path at FP8 boundaries. We use the
// backend's hip_type for the intermediate since c10::Half/BFloat16 has
// ambiguous conversions on CUDA and no implicit conversion on ROCm.
HipT out_norm_h = Converter::convert(x * s_variance * wf);
out[id * width + i] = scaled_fp8_conversion<true, fp8_type>(
Converter::convert(out_norm_h), scale_inv);
}
@@ -183,8 +192,10 @@ fused_add_rms_norm_static_fp8_quant_kernel(
for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) {
float x = (float)residual[blockIdx.x * hidden_size + idx];
// Multiply in weight's native dtype to match fused_add_rms_norm_kernel.
scalar_t out_norm = static_cast<scalar_t>(x * s_variance) * weight[idx];
float w = (float)weight[idx];
// See note in rms_norm_static_fp8_quant_kernel: round through scalar_t
// to match the unfused composite path at FP8 boundaries.
scalar_t out_norm = static_cast<scalar_t>(x * s_variance * w);
out[blockIdx.x * hidden_size + idx] = scaled_fp8_conversion<true, fp8_type>(
static_cast<float>(out_norm), scale_inv);
}
@@ -24,13 +24,21 @@ __device__ inline void vectorize_with_alignment(
ScaOp&& scalar_op) { // InT -> OutT
static_assert(VEC_SIZE > 0 && (VEC_SIZE & (VEC_SIZE - 1)) == 0,
"VEC_SIZE must be a positive power-of-two");
constexpr int WIDTH = VEC_SIZE * sizeof(InT); // eg: 64 B
constexpr int WIDTH = VEC_SIZE * sizeof(InT); // eg: 16 B
constexpr int OUT_WIDTH = VEC_SIZE * sizeof(OutT); // eg: 16 B
uintptr_t addr = reinterpret_cast<uintptr_t>(in);
uintptr_t out_addr = reinterpret_cast<uintptr_t>(out);
// fast path when the whole region is already aligned
// Note: currently the output is guaranteed to be same as the input, so we
// don't check it here, comments here just for future reference.
bool can_vec = ((addr & (WIDTH - 1)) == 0) && ((len & (VEC_SIZE - 1)) == 0);
// fast path when input and output are both fully aligned. The vector
// load/store below go through vec_n_t<T, VEC_SIZE>, declared
// __align__(VEC_SIZE * sizeof(T)), so each side must be aligned to its
// own vector width. out is NOT generally co-aligned with in: e.g.
// reshape_and_cache_flash writes KV-cache rows whose byte offset is a
// multiple of head_size, which for head sizes that are not a multiple
// of VEC_SIZE puts some rows off the vector-width boundary.
bool can_vec = ((addr & (WIDTH - 1)) == 0) &&
((out_addr & (OUT_WIDTH - 1)) == 0) &&
((len & (VEC_SIZE - 1)) == 0);
if (can_vec) {
int num_vec = len / VEC_SIZE;
@@ -55,6 +63,16 @@ __device__ inline void vectorize_with_alignment(
prefix_elems /= sizeof(InT);
prefix_elems = min(prefix_elems, len); // 0 ≤ prefix < 16
// the prefix below aligns in; if that does not also align out (their
// addresses differ modulo the vector width), vectorizing is impossible
// and the whole copy must stay scalar.
if (((out_addr + prefix_elems * sizeof(OutT)) & (OUT_WIDTH - 1)) != 0) {
for (int i = tid; i < len; i += stride) {
scalar_op(out[i], in[i]);
}
return;
}
// 1. prefill the when it is unsafe to vectorize
for (int i = tid; i < prefix_elems; i += stride) {
scalar_op(out[i], in[i]);
+7 -7
View File
@@ -75,13 +75,13 @@ RUN wget -O- https://apt.repos.intel.com/intel-gpg-keys/GPG-PUB-KEY-INTEL-SW-PRO
# Install UMD
RUN mkdir neo && \
cd neo && \
wget https://github.com/intel/intel-graphics-compiler/releases/download/v2.24.8/intel-igc-core-2_2.24.8+20344_amd64.deb && \
wget https://github.com/intel/intel-graphics-compiler/releases/download/v2.24.8/intel-igc-opencl-2_2.24.8+20344_amd64.deb && \
wget https://github.com/intel/compute-runtime/releases/download/25.48.36300.8/intel-ocloc_25.48.36300.8-0_amd64.deb && \
wget https://github.com/intel/compute-runtime/releases/download/25.48.36300.8/intel-opencl-icd_25.48.36300.8-0_amd64.deb && \
wget https://github.com/intel/compute-runtime/releases/download/25.48.36300.8/libigdgmm12_22.8.2_amd64.deb && \
wget https://github.com/intel/compute-runtime/releases/download/25.48.36300.8/libze-intel-gpu1_25.48.36300.8-0_amd64.deb && \
wget https://github.com/oneapi-src/level-zero/releases/download/v1.26.0/level-zero_1.26.0+u24.04_amd64.deb && \
wget https://github.com/intel/intel-graphics-compiler/releases/download/v2.34.4/intel-igc-core-2_2.34.4+21428_amd64.deb && \
wget https://github.com/intel/intel-graphics-compiler/releases/download/v2.34.4/intel-igc-opencl-2_2.34.4+21428_amd64.deb && \
wget https://github.com/intel/compute-runtime/releases/download/26.18.38308.1/intel-ocloc_26.18.38308.1-0_amd64.deb && \
wget https://github.com/intel/compute-runtime/releases/download/26.18.38308.1/intel-opencl-icd_26.18.38308.1-0_amd64.deb && \
wget https://github.com/intel/compute-runtime/releases/download/26.18.38308.1/libigdgmm12_22.10.0_amd64.deb && \
wget https://github.com/intel/compute-runtime/releases/download/26.18.38308.1/libze-intel-gpu1_26.18.38308.1-0_amd64.deb && \
wget https://github.com/oneapi-src/level-zero/releases/download/v1.28.2/level-zero_1.28.2+u24.04_amd64.deb && \
dpkg -i *.deb && \
cd .. && \
rm -rf neo
@@ -4,7 +4,7 @@ Deploying vLLM on Kubernetes is a scalable and efficient way to serve machine le
* **Upstream vLLM compatibility** It wraps around upstream vLLM without modifying its code.
* **Ease of use** Simplified deployment via Helm charts and observability through Grafana dashboards.
* **High performance** Optimized for LLM workloads with features like multimodel support, model-aware and prefix-aware routing, fast vLLM bootstrapping, and KV cache offloading with [LMCache](https://github.com/LMCache/LMCache), among others.
* **High performance** Optimized for LLM workloads with features like multimodel support, model-aware and prefix-aware routing, fast vLLM bootstrapping, and KV cache offloading with [LMCache](https://github.com/LMCache/LMCache) (wired up in vLLM via `--kv-offloading-backend lmcache`; see the [LMCache examples](https://github.com/vllm-project/vllm/tree/main/examples/disaggregated/lmcache) and [docs.lmcache.ai](https://docs.lmcache.ai)), among others.
If you are new to Kubernetes, don't worry: in the vLLM production stack [repo](https://github.com/vllm-project/production-stack), we provide a step-by-step [guide](https://github.com/vllm-project/production-stack/blob/main/tutorials/00-install-kubernetes-env.md) and a [short video](https://www.youtube.com/watch?v=EsTJbQtzj0g) to set up everything and get started in **4 minutes**!
+1 -1
View File
@@ -20,7 +20,7 @@ Two main reasons:
Now supports 9 types of connectors:
- **ExampleConnector**: refer to [examples/disaggregated/example_connector/run.sh](../../examples/disaggregated/example_connector/run.sh) for the example usage of ExampleConnector disaggregated prefilling.
- **LMCacheConnectorV1**: refer to [examples/disaggregated/lmcache/disagg_prefill_lmcache_v1/disagg_example_nixl.sh](../../examples/disaggregated/lmcache/disagg_prefill_lmcache_v1/disagg_example_nixl.sh) for the example usage of LMCacheConnectorV1 disaggregated prefilling which uses NIXL as the underlying KV transmission.
- **LMCacheConnectorV1**: refer to [examples/disaggregated/lmcache/disagg_prefill_lmcache_v1/disagg_example_nixl.sh](../../examples/disaggregated/lmcache/disagg_prefill_lmcache_v1/disagg_example_nixl.sh) for the example usage of LMCacheConnectorV1 disaggregated prefilling which uses NIXL as the underlying KV transmission. LMCache also offers a multi-process (MP) mode via `LMCacheMPConnector`, where a standalone `lmcache server` holds the KV cache shared by one or more vLLM instances; see the [LMCache examples](../../examples/disaggregated/lmcache/README.md) and the [LMCache docs](https://docs.lmcache.ai) for setup.
- **NixlConnector**: refer to [tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh](../../tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh) for the example usage of NixlConnector disaggregated prefilling which support fully async send/recv. For detailed usage guide, see [NixlConnector Usage Guide](nixl_connector_usage.md). For feature compatibility details, see [NixlConnector Compatibility Matrix](nixl_connector_compatibility.md). You may specify one or multiple NIXL transfer backends, such as:
```bash
@@ -203,6 +203,7 @@ the vLLM JSON config.
### kv_connector_extra_config
- `load_async` (bool): Enable asynchronous loading for better compute-I/O overlap. Default: `true`.
- `lookup_async` (bool): Run the external prefix-cache lookup on a background thread so it never blocks the scheduler step. The request is held until the in-flight lookup completes, then resumed on a later step. Default: `false`.
- `enable_cross_layers_blocks` (bool): Enable cross-layer block packing for reduced store operations. Default: `false`.
- `lookup_rpc_port` (int): Custom port for the ZMQ lookup RPC socket. Default: `0`.
- `cache_prefix` (str): Namespace prepended to every store key. Lets separate deployments share one Mooncake master without polluting each other — instances configured with different prefixes never see each other's cached blocks, even for identical prompts. All instances that should share a prefix cache must use the same value. Default: `""` (no prefix; keys are byte-identical to the unprefixed format).
@@ -27,6 +27,7 @@ Currently, there are no pre-built XPU wheels.
- First, install required [driver](https://dgpu-docs.intel.com/driver/installation.html#installing-gpu-drivers).
- Second, install Python packages for vLLM XPU backend building (Intel OneAPI dependencies are installed automatically as part of `torch-xpu`, see [PyTorch XPU get started](https://docs.pytorch.org/docs/stable/notes/get_start_xpu.html)):
- Start from vllm-xpu-kernels v0.1.10, we recommend user upgrade driver to [compute runtime 26.18](https://github.com/intel/compute-runtime/releases/tag/26.14.37833.4) release, to avoid potential compatibility issue.
```bash
git clone https://github.com/vllm-project/vllm.git
+34 -14
View File
@@ -1,10 +1,38 @@
# LMCache Examples
This folder demonstrates how to use LMCache for disaggregated prefilling, CPU offloading and KV cache sharing.
This folder demonstrates how to use LMCache with vLLM v1 for KV cache
offloading, disaggregated prefilling, and KV cache sharing.
## 1. Disaggregated Prefill in vLLM v1
## Integration modes
This example demonstrates how to run LMCache with disaggregated prefill using NIXL on a single node.
LMCache integrates with vLLM v1 in two ways:
- **In-process mode** (`LMCacheConnectorV1`): LMCache runs inside the vLLM
process and is configured through environment variables or a YAML config
file (`LMCACHE_CONFIG_FILE`). This is the simplest way to add single-node
CPU/disk offloading.
- **Multi-process (MP) mode** (`LMCacheMPConnector`): LMCache runs as a
standalone server (`lmcache server`) that owns the KV cache storage; one or
more vLLM instances connect to it. This is the recommended mode for
distributed KV storage and for sharing KV cache across instances. See the
[LMCache docs](https://docs.lmcache.ai) for the full MP setup.
## 1. CPU offload (in-process)
- `python cpu_offload_lmcache.py` - CPU offloading with `LMCacheConnectorV1`
for vLLM v1.
## 2. CPU offload (multi-process)
- `bash cpu_offload_lmcache_mp.sh` - CPU offloading with `LMCacheMPConnector`,
using a standalone `lmcache server`. vLLM provides a built-in shortcut for
this setup via `--kv-offloading-backend lmcache` and
`--kv-offloading-size <GiB>`.
## 3. Disaggregated Prefill in vLLM v1
This example demonstrates how to run LMCache with disaggregated prefill using
NIXL on a single node.
### Prerequisites
@@ -46,15 +74,7 @@ The main script generates several log files:
- `decoder.log` - Logs from the decode server
- `proxy.log` - Logs from the proxy server
## 2. CPU Offload Examples
## 4. KV Cache Sharing
- `python cpu_offload_lmcache.py -v v0` - CPU offloading implementation for vLLM v0
- `python cpu_offload_lmcache.py -v v1` - CPU offloading implementation for vLLM v1
## 3. KV Cache Sharing
The `kv_cache_sharing_lmcache_v1.py` example demonstrates how to share KV caches between vLLM v1 instances.
## 4. Disaggregated Prefill in vLLM v0
The `disaggregated_prefill_lmcache_v0.py` provides an example of how to run disaggregated prefill in vLLM v0.
The `kv_cache_sharing_lmcache_v1.py` example demonstrates how to share KV
caches between vLLM v1 instances through a centralized LMCache server.
@@ -1,20 +1,8 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
This file demonstrates the example usage of cpu offloading
with LMCache in vLLM v1 or v0.
Usage:
Specify vLLM version
-v v0 : Use LMCacheConnector
model = mistralai/Mistral-7B-Instruct-v0.2
(Includes enable_chunked_prefill = True)
-v v1 : Use LMCacheConnectorV1 (default)
model = meta-llama/Meta-Llama-3.1-8B-Instruct
(Without enable_chunked_prefill)
This file demonstrates the example usage of CPU offloading
with LMCache in vLLM v1.
Note that `lmcache` is needed to run this example.
Requirements:
@@ -23,7 +11,6 @@ Learn more about LMCache environment setup, please refer to:
https://docs.lmcache.ai/getting_started/installation.html
"""
import argparse
import contextlib
import os
import time
@@ -39,8 +26,6 @@ from vllm.engine.arg_utils import EngineArgs
def setup_environment_variables():
# LMCache-related environment variables
# Use experimental features in LMCache
os.environ["LMCACHE_USE_EXPERIMENTAL"] = "True"
# LMCache is set to use 256 tokens per chunk
os.environ["LMCACHE_CHUNK_SIZE"] = "256"
# Enable local CPU backend in LMCache
@@ -50,9 +35,9 @@ def setup_environment_variables():
@contextlib.contextmanager
def build_llm_with_lmcache(lmcache_connector: str, model: str):
def build_llm_with_lmcache(model: str):
ktc = KVTransferConfig(
kv_connector=lmcache_connector,
kv_connector="LMCacheConnectorV1",
kv_role="kv_both",
)
# Set GPU memory utilization to 0.8 for an A40 GPU with 40GB
@@ -92,23 +77,10 @@ def print_output(
print("-" * 50)
def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument(
"-v",
"--version",
choices=["v0", "v1"],
default="v1",
help="Specify vLLM version (default: v1)",
)
return parser.parse_args()
def main():
lmcache_connector = "LMCacheConnectorV1"
model = "meta-llama/Meta-Llama-3.1-8B-Instruct"
setup_environment_variables()
with build_llm_with_lmcache(lmcache_connector, model) as llm:
with build_llm_with_lmcache(model) as llm:
# This example script runs two requests with a shared prefix.
# Define the shared prompt and specific prompts
shared_prompt = "Hello, how are you?" * 1000
+43
View File
@@ -0,0 +1,43 @@
#!/bin/bash
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
#
# CPU offloading with LMCache in multi-process (MP) mode.
#
# In MP mode, LMCache runs as a standalone server process (`lmcache server`)
# that owns the KV cache storage. One or more vLLM instances connect to it via
# the `LMCacheMPConnector`. This is the recommended way to run LMCache for
# distributed KV storage and for sharing KV cache across vLLM instances.
#
# vLLM ships a built-in shortcut for this setup: pass `--kv-offloading-backend
# lmcache` together with `--kv-offloading-size <GiB>` and vLLM wires up the
# `LMCacheMPConnector` for you (it defaults to the LMCache server at
# tcp://localhost:5555, matching the `lmcache server` default).
#
# Requires `lmcache` to be installed (`pip install lmcache`).
# Learn more: https://docs.lmcache.ai
set -euo pipefail
MODEL=${MODEL:-meta-llama/Meta-Llama-3.1-8B-Instruct}
# 1. Launch the standalone LMCache server (binds tcp://localhost:5555 by
# default). `--l1-size-gb` sets the CPU memory budget for the L1 cache.
echo "Starting LMCache server..."
lmcache server --host localhost --port 5555 --l1-size-gb 5 &
LMCACHE_SERVER_PID=$!
trap 'kill $LMCACHE_SERVER_PID 2>/dev/null || true' EXIT
# 2. Launch vLLM and offload KV cache to the LMCache server.
# The MP connector currently requires the non-hybrid KV cache manager.
echo "Starting vLLM server with LMCache MP offloading..."
vllm serve "$MODEL" \
--port 8000 \
--kv-offloading-size 5 \
--kv-offloading-backend lmcache \
--disable-hybrid-kv-cache-manager
# Equivalent explicit configuration (instead of the two flags above):
# --kv-transfer-config \
# '{"kv_connector":"LMCacheMPConnector","kv_role":"kv_both",
# "kv_connector_extra_config":{"lmcache.mp.host":"tcp://localhost",
# "lmcache.mp.port":5555}}'
@@ -1,144 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
This file demonstrates the example usage of disaggregated prefilling
with LMCache.
We will launch 2 vllm instances (GPU 0 for prefill and GPU 1 for decode),
and launch an additional LMCache server.
KV cache is transferred in the following manner:
vLLM prefill node -> LMCache server -> vLLM decode node.
Note that `pip install lmcache` is needed to run this example.
Learn more about LMCache in https://github.com/LMCache/LMCache.
"""
import os
import subprocess
import time
from multiprocessing import Event, Process
from lmcache.experimental.cache_engine import LMCacheEngineBuilder
from lmcache.integration.vllm.utils import ENGINE_NAME
from vllm import LLM, SamplingParams
from vllm.config import KVTransferConfig
# LMCache-related environment variables
# The port to start LMCache server
port = 8100
# Use experimental features in LMCache
os.environ["LMCACHE_USE_EXPERIMENTAL"] = "True"
# LMCache is set to use 256 tokens per chunk
os.environ["LMCACHE_CHUNK_SIZE"] = "256"
# Disable local CPU backend in LMCache
os.environ["LMCACHE_LOCAL_CPU"] = "False"
# Set local CPU memory buffer limit to 5.0 GB
os.environ["LMCACHE_MAX_LOCAL_CPU_SIZE"] = "5.0"
# Set the remote URL for LMCache server
os.environ["LMCACHE_REMOTE_URL"] = f"lm://localhost:{port}"
# Set the serializer/deserializer between vllm and LMCache server
# `naive` indicates using raw bytes of the tensor without any compression
os.environ["LMCACHE_REMOTE_SERDE"] = "naive"
prompts = [
"Hello, how are you?" * 1000,
]
def run_prefill(prefill_done, prompts):
# We use GPU 0 for prefill node.
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
sampling_params = SamplingParams(temperature=0, top_p=0.95, max_tokens=1)
ktc = KVTransferConfig(
kv_connector="LMCacheConnector",
kv_role="kv_producer",
kv_rank=0,
kv_parallel_size=2,
)
# Set GPU memory utilization to 0.8 for an A40 GPU with 40GB
# memory. Reduce the value if your GPU has less memory.
llm = LLM(
model="mistralai/Mistral-7B-Instruct-v0.2",
kv_transfer_config=ktc,
max_model_len=8000,
gpu_memory_utilization=0.8,
enforce_eager=True,
)
# llm.generate(prompts, sampling_params)
outputs = llm.generate(prompts, sampling_params)
for output in outputs:
generated_text = output.outputs[0].text
print(f"Generated text: {generated_text!r}")
print("Prefill node is finished.")
prefill_done.set()
# Clean up lmcache backend
LMCacheEngineBuilder.destroy(ENGINE_NAME)
def run_decode(prefill_done, prompts, timeout=1):
# We use GPU 1 for decode node.
os.environ["CUDA_VISIBLE_DEVICES"] = "1"
sampling_params = SamplingParams(temperature=0, top_p=0.95, max_tokens=10)
ktc = KVTransferConfig(
kv_connector="LMCacheConnector",
kv_role="kv_consumer",
kv_rank=1,
kv_parallel_size=2,
)
# Set GPU memory utilization to 0.8 for an A40 GPU with 40GB
# of memory. Reduce the value if your GPU has less memory.
llm = LLM(
model="mistralai/Mistral-7B-Instruct-v0.2",
kv_transfer_config=ktc,
max_model_len=8000,
gpu_memory_utilization=0.8,
enforce_eager=True,
)
print("Waiting for prefill node to finish...")
prefill_done.wait()
time.sleep(timeout)
outputs = llm.generate(prompts, sampling_params)
for output in outputs:
generated_text = output.outputs[0].text
print(f"Generated text: {generated_text!r}")
# Clean up lmcache backend
LMCacheEngineBuilder.destroy(ENGINE_NAME)
def run_lmcache_server(port):
server_proc = subprocess.Popen(
["python", "-m", "lmcache.experimental.server", "localhost", str(port)]
)
return server_proc
def main():
prefill_done = Event()
prefill_process = Process(target=run_prefill, args=(prefill_done, prompts))
decode_process = Process(target=run_decode, args=(prefill_done, prompts))
lmcache_server_process = run_lmcache_server(port)
# Start prefill node
prefill_process.start()
# Start decode node
decode_process.start()
# Clean up the processes
decode_process.join()
prefill_process.terminate()
lmcache_server_process.terminate()
lmcache_server_process.wait()
if __name__ == "__main__":
main()
@@ -30,7 +30,6 @@ if [[ $1 == "prefiller" ]]; then
UCX_TLS=cuda_ipc,cuda_copy,tcp \
LMCACHE_CONFIG_FILE=$prefill_config_file \
LMCACHE_USE_EXPERIMENTAL=True \
VLLM_ENABLE_V1_MULTIPROCESSING=1 \
VLLM_WORKER_MULTIPROC_METHOD=spawn \
CUDA_VISIBLE_DEVICES=0 \
@@ -47,7 +46,6 @@ elif [[ $1 == "decoder" ]]; then
UCX_TLS=cuda_ipc,cuda_copy,tcp \
LMCACHE_CONFIG_FILE=$decode_config_file \
LMCACHE_USE_EXPERIMENTAL=True \
VLLM_ENABLE_V1_MULTIPROCESSING=1 \
VLLM_WORKER_MULTIPROC_METHOD=spawn \
CUDA_VISIBLE_DEVICES=1 \
@@ -26,8 +26,6 @@ from vllm.config import KVTransferConfig
# LMCache-related environment variables
# The port to start LMCache server
port = 8100
# Use experimental features in LMCache
os.environ["LMCACHE_USE_EXPERIMENTAL"] = "True"
# LMCache is set to use 256 tokens per chunk
os.environ["LMCACHE_CHUNK_SIZE"] = "256"
# Disable local CPU backend in LMCache
+3 -2
View File
@@ -11,13 +11,14 @@ transformers >= 5.5.3
tokenizers >= 0.21.1 # Required for fast incremental detokenization.
safetensors >= 0.6.2 # MXFP4/MXFP6 dtype support (F8_E8M0, F4) added in 0.6.0: https://github.com/huggingface/safetensors/pull/611
protobuf >= 5.29.6, !=6.30.*, !=6.31.*, !=6.32.*, !=6.33.0.*, !=6.33.1.*, !=6.33.2.*, !=6.33.3.*, !=6.33.4.* # Required by LlamaTokenizer, gRPC. CVE-2026-0994
fastapi[standard] >= 0.115.0 # Required by FastAPI's form models in the OpenAI API server's audio transcriptions endpoint.
fastapi[standard] >= 0.133.0, < 0.137.0 # First version supporting Starlette 1.0; < 0.137.0 avoids route-tree change that breaks model-hosting-container-standards handler overrides.
starlette >= 1.0.1 # CVE-2026-48710: Host header injection in < 1.0.1
aiohttp >= 3.13.3
openai >= 2.0.0 # For Responses API with reasoning content
pydantic >= 2.12.0
prometheus_client >= 0.18.0
pillow # Required for image processing
prometheus-fastapi-instrumentator >= 7.0.0
prometheus-fastapi-instrumentator >= 8.0.0 # v8 unblocks starlette >= 1.0
tiktoken >= 0.6.0 # Required for DBRX tokenizer
lm-format-enforcer == 0.11.3
llguidance >= 1.7.0, < 1.8.0; platform_machine == "x86_64" or platform_machine == "arm64" or platform_machine == "aarch64" or platform_machine == "ppc64le"
+23 -48
View File
@@ -35,14 +35,11 @@ arctic-inference==0.1.1
# via -r requirements/test/cuda.in
argcomplete==3.5.1
# via datamodel-code-generator
arrow==1.3.0
# via isoduration
attrs==24.2.0
# via
# aiohttp
# hypothesis
# jsonschema
# pytest-subtests
# referencing
audioread==3.0.1
# via librosa
@@ -57,9 +54,7 @@ azure-identity==1.25.2
azure-storage-blob==12.28.0
# via runai-model-streamer-azure
backoff==2.2.1
# via
# -r requirements/test/cuda.in
# schemathesis
# via -r requirements/test/cuda.in
bitsandbytes==0.49.2
# via -r requirements/test/cuda.in
black==24.10.0
@@ -110,7 +105,6 @@ colorama==0.4.6
# via
# perceptron
# sacrebleu
# schemathesis
colorful==0.5.6
# via ray
colorlog==6.10.1
@@ -183,7 +177,7 @@ et-xmlfile==2.0.0
# via openpyxl
evaluate==0.4.3
# via lm-eval
fastapi==0.128.0
fastapi==0.136.3
# via
# -c requirements/common.txt
# gpt-oss
@@ -206,8 +200,6 @@ filelock==3.16.1
# virtualenv
fonttools==4.55.0
# via matplotlib
fqdn==1.5.1
# via jsonschema
frozendict==2.4.6
# via einx
frozenlist==1.5.0
@@ -269,7 +261,7 @@ h11==0.14.0
# uvicorn
h2==4.3.0
# via httpx
harfile==0.3.0
harfile==0.5.0
# via schemathesis
hf-xet==1.4.3
# via huggingface-hub
@@ -309,7 +301,7 @@ hypothesis==6.131.0
# hypothesis-graphql
# hypothesis-jsonschema
# schemathesis
hypothesis-graphql==0.11.1
hypothesis-graphql==0.13.0
# via schemathesis
hypothesis-jsonschema==0.23.1
# via schemathesis
@@ -318,7 +310,6 @@ idna==3.10
# anyio
# email-validator
# httpx
# jsonschema
# requests
# yarl
imagehash==4.3.2
@@ -335,8 +326,6 @@ instanttensor==0.1.5
# via -r requirements/test/cuda.in
isodate==0.7.2
# via azure-storage-blob
isoduration==20.11.0
# via jsonschema
isort==5.13.2
# via datamodel-code-generator
jinja2==3.1.6
@@ -356,15 +345,14 @@ joblib==1.4.2
# librosa
# nltk
# scikit-learn
jsonpointer==3.0.0
# via jsonschema
jsonschema==4.23.0
# via
# -c requirements/common.txt
# hypothesis-jsonschema
# mistral-common
# ray
# schemathesis
jsonschema-rs==0.46.5
# via schemathesis
jsonschema-specifications==2024.10.1
# via jsonschema
junit-xml==1.9
@@ -715,18 +703,20 @@ pydantic-core==2.41.1
pydantic-extra-types==2.10.5
# via mistral-common
pygments==2.18.0
# via rich
# via
# pytest
# rich
pyjwt==2.11.0
# via msal
pyparsing==3.2.0
# via matplotlib
pyrate-limiter==3.7.0
pyrate-limiter==4.4.0
# via schemathesis
pystemmer==3.0.0
# via mteb
pytablewriter==1.2.0
# via lm-eval
pytest==8.3.5
pytest==9.1.0
# via
# -r requirements/test/cuda.in
# buildkite-test-collector
@@ -737,10 +727,9 @@ pytest==8.3.5
# pytest-mock
# pytest-rerunfailures
# pytest-shard
# pytest-subtests
# pytest-timeout
# schemathesis
pytest-asyncio==0.24.0
pytest-asyncio==1.4.0
# via -r requirements/test/cuda.in
pytest-cov==6.3.0
# via -r requirements/test/cuda.in
@@ -752,13 +741,10 @@ pytest-rerunfailures==14.0
# via -r requirements/test/cuda.in
pytest-shard==0.1.2
# via -r requirements/test/cuda.in
pytest-subtests==0.14.1
# via schemathesis
pytest-timeout==2.3.1
# via -r requirements/test/cuda.in
python-dateutil==2.9.0.post0
# via
# arrow
# botocore
# matplotlib
# pandas
@@ -829,15 +815,12 @@ requests==2.32.3
# tiktoken
responses==0.25.3
# via genai-perf
rfc3339-validator==0.1.4
# via jsonschema
rfc3987==1.3.8
# via jsonschema
rich==13.9.4
# via
# genai-perf
# mteb
# perceptron
# schemathesis
# typer
rouge-score==0.1.2
# via lm-eval
@@ -868,7 +851,7 @@ safetensors==0.7.0
# segmentation-models-pytorch
# timm
# transformers
schemathesis==3.39.15
schemathesis==4.21.6
# via -r requirements/test/cuda.in
scikit-image==0.25.2
# via albumentations
@@ -912,7 +895,6 @@ six==1.16.0
# junit-xml
# opencensus
# python-dateutil
# rfc3339-validator
# rouge-score
smart-open==7.1.0
# via ray
@@ -938,10 +920,10 @@ sqlalchemy==2.0.41
# optuna
sqlitedict==2.1.0
# via lm-eval
starlette==0.50.0
starlette==1.3.1
# via
# -c requirements/common.txt
# fastapi
# schemathesis
# starlette-testclient
starlette-testclient==0.4.1
# via schemathesis
@@ -966,6 +948,7 @@ tenacity==9.1.2
# gpt-oss
# lm-eval
# plotly
# schemathesis
tensorizer==2.10.1
# via -r requirements/test/cuda.in
termcolor==3.1.0
@@ -990,10 +973,6 @@ tokenizers==0.22.2
# -c requirements/common.txt
# -r requirements/test/cuda.in
# transformers
tomli==2.2.1
# via schemathesis
tomli-w==1.2.0
# via schemathesis
torch==2.11.0+cu130
# via
# -c requirements/cuda.txt
@@ -1066,8 +1045,6 @@ typer==0.15.2
# huggingface-hub
# perceptron
# transformers
types-python-dateutil==2.9.0.20241206
# via arrow
typing-extensions==4.15.0
# via
# -c requirements/common.txt
@@ -1092,6 +1069,8 @@ typing-extensions==4.15.0
# pydantic
# pydantic-core
# pydantic-extra-types
# pytest-asyncio
# schemathesis
# sentence-transformers
# sqlalchemy
# starlette
@@ -1099,11 +1078,11 @@ typing-extensions==4.15.0
# typer
# typing-inspection
typing-inspection==0.4.2
# via pydantic
# via
# fastapi
# pydantic
tzdata==2024.2
# via pandas
uri-template==1.3.0
# via jsonschema
urllib3==2.2.3
# via
# blobfile
@@ -1122,8 +1101,6 @@ vocos==0.1.0
# via -r requirements/test/cuda.in
wcwidth==0.2.13
# via ftfy
webcolors==24.11.1
# via jsonschema
werkzeug==3.1.3
# via schemathesis
word2number==1.1
@@ -1135,8 +1112,6 @@ xxhash==3.5.0
# datasets
# evaluate
yarl==1.17.1
# via
# aiohttp
# schemathesis
# via aiohttp
zipp==3.23.0
# via importlib-metadata
+21 -47
View File
@@ -51,15 +51,12 @@ arctic-inference==0.1.1
# via -r requirements/test/rocm.in
argcomplete==3.6.3
# via datamodel-code-generator
arrow==1.4.0
# via isoduration
astor==0.8.1
# via depyf
attrs==26.1.0
# via
# aiohttp
# jsonschema
# pytest-subtests
# referencing
audioread==3.0.1
# via librosa
@@ -74,9 +71,7 @@ azure-identity==1.25.3
azure-storage-blob==12.28.0
# via runai-model-streamer-azure
backoff==2.2.1
# via
# -r requirements/test/rocm.in
# schemathesis
# via -r requirements/test/rocm.in
bitsandbytes==0.49.2
# via -r requirements/test/rocm.in
black==26.3.1
@@ -139,7 +134,6 @@ colorama==0.4.6
# via
# perceptron
# sacrebleu
# schemathesis
colorful==0.5.8
# via ray
colorlog==6.10.1
@@ -258,8 +252,6 @@ filelock==3.25.2
# virtualenv
fonttools==4.62.1
# via matplotlib
fqdn==1.5.1
# via jsonschema
frozendict==2.4.7
# via einx
frozenlist==1.8.0
@@ -328,7 +320,7 @@ h11==0.16.0
# uvicorn
h2==4.3.0
# via httpx
harfile==0.4.0
harfile==0.5.0
# via schemathesis
hf-xet==1.4.3
# via huggingface-hub
@@ -378,7 +370,7 @@ hypothesis==6.151.9
# hypothesis-graphql
# hypothesis-jsonschema
# schemathesis
hypothesis-graphql==0.12.0
hypothesis-graphql==0.13.0
# via schemathesis
hypothesis-jsonschema==0.23.1
# via schemathesis
@@ -387,7 +379,6 @@ idna==3.11
# anyio
# email-validator
# httpx
# jsonschema
# requests
# yarl
ijson==3.5.0
@@ -408,8 +399,6 @@ interegular==0.3.3
# via lm-format-enforcer
isodate==0.7.2
# via azure-storage-blob
isoduration==20.11.0
# via jsonschema
isort==8.0.1
# via datamodel-code-generator
jinja2==3.1.6
@@ -435,8 +424,6 @@ joblib==1.5.3
# librosa
# nltk
# scikit-learn
jsonpointer==3.1.0
# via jsonschema
jsonschema==4.26.0
# via
# -c requirements/common.txt
@@ -445,7 +432,8 @@ jsonschema==4.26.0
# mcp
# mistral-common
# ray
# schemathesis
jsonschema-rs==0.46.5
# via schemathesis
jsonschema-specifications==2025.9.1
# via jsonschema
junit-xml==1.9
@@ -792,7 +780,7 @@ prometheus-client==0.24.1
# opentelemetry-exporter-prometheus
# prometheus-fastapi-instrumentator
# ray
prometheus-fastapi-instrumentator==7.1.0
prometheus-fastapi-instrumentator==8.0.0
# via
# -c requirements/common.txt
# -r requirements/test/../common.txt
@@ -876,20 +864,22 @@ pydantic-settings==2.13.1
# fastapi
# mcp
pygments==2.19.2
# via rich
# via
# pytest
# rich
pyjwt==2.12.1
# via
# mcp
# msal
pyparsing==3.3.2
# via matplotlib
pyrate-limiter==3.9.0
pyrate-limiter==4.4.0
# via schemathesis
pystemmer==3.0.0
# via mteb
pytablewriter==1.2.1
# via lm-eval
pytest==8.3.5
pytest==9.1.0
# via
# -r requirements/test/rocm.in
# buildkite-test-collector
@@ -900,10 +890,9 @@ pytest==8.3.5
# pytest-mock
# pytest-rerunfailures
# pytest-shard
# pytest-subtests
# pytest-timeout
# schemathesis
pytest-asyncio==0.24.0
pytest-asyncio==1.4.0
# via -r requirements/test/rocm.in
pytest-cov==6.3.0
# via -r requirements/test/rocm.in
@@ -915,13 +904,10 @@ pytest-rerunfailures==14.0
# via -r requirements/test/rocm.in
pytest-shard==0.1.2
# via -r requirements/test/rocm.in
pytest-subtests==0.14.2
# via schemathesis
pytest-timeout==2.3.1
# via -r requirements/test/rocm.in
python-dateutil==2.9.0.post0
# via
# arrow
# botocore
# matplotlib
# pandas
@@ -1016,16 +1002,13 @@ requests==2.32.5
# tiktoken
responses==0.26.0
# via genai-perf
rfc3339-validator==0.1.4
# via jsonschema
rfc3987==1.3.8
# via jsonschema
rich==14.3.3
# via
# genai-perf
# mteb
# perceptron
# rich-toolkit
# schemathesis
# typer
rich-toolkit==0.19.7
# via
@@ -1063,7 +1046,7 @@ safetensors==0.7.0
# segmentation-models-pytorch
# timm
# transformers
schemathesis==3.39.15
schemathesis==4.21.6
# via -r requirements/test/rocm.in
scikit-image==0.26.0
# via albumentations
@@ -1120,7 +1103,6 @@ six==1.17.0
# junit-xml
# opencensus
# python-dateutil
# rfc3339-validator
# rouge-score
smart-open==7.5.1
# via ray
@@ -1149,13 +1131,14 @@ sqlitedict==2.1.0
# via lm-eval
sse-starlette==3.3.4
# via mcp
starlette==0.52.1
starlette==1.3.1
# via
# -c requirements/common.txt
# -r requirements/test/../common.txt
# fastapi
# mcp
# model-hosting-container-standards
# prometheus-fastapi-instrumentator
# schemathesis
# sse-starlette
# starlette-testclient
starlette-testclient==0.4.1
@@ -1182,6 +1165,7 @@ tenacity==9.1.4
# via
# gpt-oss
# lm-eval
# schemathesis
tensorizer==2.10.1
# via
# -c requirements/rocm.txt
@@ -1215,10 +1199,6 @@ tokenizers==0.22.2
# -r requirements/test/../common.txt
# -r requirements/test/rocm.in
# transformers
tomli==2.4.0
# via schemathesis
tomli-w==1.2.0
# via schemathesis
torch-c-dlpack-ext==0.1.5
# via tilelang
tqdm==4.67.3
@@ -1301,8 +1281,10 @@ typing-extensions==4.15.0
# pydantic
# pydantic-core
# pydantic-extra-types
# pytest-asyncio
# referencing
# rich-toolkit
# schemathesis
# sentence-transformers
# sqlalchemy
# starlette
@@ -1317,10 +1299,6 @@ typing-inspection==0.4.2
# mcp
# pydantic
# pydantic-settings
tzdata==2025.3
# via arrow
uri-template==1.3.0
# via jsonschema
urllib3==2.6.3
# via
# blobfile
@@ -1351,8 +1329,6 @@ watchfiles==1.1.1
# uvicorn
wcwidth==0.6.0
# via ftfy
webcolors==25.10.0
# via jsonschema
websockets==16.0
# via uvicorn
werkzeug==3.1.6
@@ -1370,9 +1346,7 @@ xxhash==3.6.0
# datasets
# evaluate
yarl==1.23.0
# via
# aiohttp
# schemathesis
# via aiohttp
z3-solver==4.15.4.0
# via tilelang
zipp==3.23.0
+2 -1
View File
@@ -593,8 +593,9 @@ soxr==0.5.0.post1
# mistral-common
sqlitedict==2.1.0
# via lm-eval
starlette==1.0.0
starlette==1.3.1
# via
# -c requirements/common.txt
# fastapi
# starlette-testclient
starlette-testclient==0.4.1
+1 -1
View File
@@ -17,4 +17,4 @@ torchaudio
torchvision
auto_round_lib>=0.13.3
vllm_xpu_kernels @ https://github.com/vllm-project/vllm-xpu-kernels/releases/download/v0.1.9.1/vllm_xpu_kernels-0.1.9.1-cp38-abi3-manylinux_2_28_x86_64.whl
vllm_xpu_kernels @ https://github.com/vllm-project/vllm-xpu-kernels/releases/download/v0.1.10/vllm_xpu_kernels-0.1.10-cp38-abi3-manylinux_2_28_x86_64.whl
@@ -53,38 +53,6 @@ class SPTestSettings:
runner: RunnerOption
test_options: SPTestOptions
@staticmethod
def detailed(
*,
tp_base: int = 2,
pp_base: int = 1,
multi_node_only: bool = False,
runner: RunnerOption = "auto",
load_format: str | None = None,
):
parallel_setups = []
for eager_mode_val in [False, True]:
for pp_multiplier in [1, 2]:
for chunked_prefill_val in [False, True]:
parallel_setups.append(
ParallelSetup(
tp_size=tp_base,
pp_size=pp_multiplier * pp_base,
fuse_norm_quant=False,
fuse_act_quant=False,
eager_mode=eager_mode_val,
chunked_prefill=chunked_prefill_val,
)
)
return SPTestSettings(
parallel_setups=parallel_setups,
distributed_backends=["mp", "ray"],
runner=runner,
test_options=SPTestOptions(
multi_node_only=multi_node_only, load_format=load_format
),
)
@staticmethod
def fast(
*,
@@ -94,23 +62,26 @@ class SPTestSettings:
multi_node_only: bool = False,
load_format: str | None = None,
):
parallel_setups = []
for eager_mode_val in [False, True]:
for pp_multiplier in [1, 2]:
for chunked_prefill_val in [False, True]:
parallel_setups.append(
ParallelSetup(
tp_size=tp_base,
pp_size=pp_multiplier * pp_base,
fuse_norm_quant=False,
fuse_act_quant=False,
eager_mode=eager_mode_val,
chunked_prefill=chunked_prefill_val,
)
)
return SPTestSettings(
parallel_setups=parallel_setups,
distributed_backends=["mp", "ray"],
parallel_setups=[
ParallelSetup(
tp_size=tp_base,
pp_size=pp_base,
fuse_norm_quant=False,
fuse_act_quant=False,
eager_mode=False,
chunked_prefill=True,
),
ParallelSetup(
tp_size=tp_base,
pp_size=2 * pp_base,
fuse_norm_quant=False,
fuse_act_quant=False,
eager_mode=False,
chunked_prefill=True,
),
],
distributed_backends=["mp"],
runner=runner,
test_options=SPTestOptions(
multi_node_only=multi_node_only, load_format=load_format
@@ -8,6 +8,8 @@ AnthropicServingMessages._convert_anthropic_to_openai_request().
Also covers extended-thinking edge cases such as ``redacted_thinking``
blocks echoed back by Anthropic clients, and streaming conversion in
``message_stream_converter``.
Also covers cache usage computation in ``_build_anthropic_usage``.
"""
import json
@@ -18,7 +20,11 @@ import pytest
from vllm.entrypoints.anthropic.protocol import (
AnthropicMessagesRequest,
)
from vllm.entrypoints.anthropic.serving import AnthropicServingMessages
from vllm.entrypoints.anthropic.serving import (
AnthropicServingMessages,
_build_anthropic_usage,
_get_cached_tokens,
)
from vllm.entrypoints.openai.chat_completion.protocol import (
ChatCompletionResponseStreamChoice,
ChatCompletionStreamResponse,
@@ -27,6 +33,7 @@ from vllm.entrypoints.openai.engine.protocol import (
DeltaFunctionCall,
DeltaMessage,
DeltaToolCall,
PromptTokenUsageInfo,
UsageInfo,
)
@@ -653,6 +660,108 @@ class TestThinkingBlockConversion:
assert asst.get("content") == "Hi!"
# ======================================================================
# Cache usage computation
# ======================================================================
class TestGetCachedTokens:
"""Tests for _get_cached_tokens helper."""
def test_none_usage(self):
assert _get_cached_tokens(None) is None
def test_no_prompt_tokens_details(self):
usage = UsageInfo(prompt_tokens=100, completion_tokens=10)
assert _get_cached_tokens(usage) is None
def test_cached_tokens_present(self):
usage = UsageInfo(
prompt_tokens=100,
completion_tokens=10,
prompt_tokens_details=PromptTokenUsageInfo(cached_tokens=80),
)
assert _get_cached_tokens(usage) == 80
def test_cached_tokens_zero(self):
"""Zero cached tokens should return 0, not None."""
usage = UsageInfo(
prompt_tokens=100,
completion_tokens=10,
prompt_tokens_details=PromptTokenUsageInfo(cached_tokens=0),
)
assert _get_cached_tokens(usage) == 0
def test_cached_tokens_none_in_details(self):
usage = UsageInfo(
prompt_tokens=100,
completion_tokens=10,
prompt_tokens_details=PromptTokenUsageInfo(cached_tokens=None),
)
assert _get_cached_tokens(usage) is None
class TestBuildAnthropicUsage:
"""Tests for _build_anthropic_usage helper.
Anthropic defines: total_input = input_tokens + cache_read + cache_creation
vLLM's prompt_tokens is the total.
"""
def test_no_cache_info(self):
"""When cache info is unavailable, return raw prompt_tokens."""
result = _build_anthropic_usage(100, 10, None)
assert result.input_tokens == 100
assert result.output_tokens == 10
assert result.cache_read_input_tokens is None
assert result.cache_creation_input_tokens is None
def test_cache_hit(self):
"""When cache is hit, input_tokens excludes cached tokens."""
usage = UsageInfo(
prompt_tokens=100,
completion_tokens=10,
prompt_tokens_details=PromptTokenUsageInfo(cached_tokens=80),
)
result = _build_anthropic_usage(100, 10, usage)
assert result.input_tokens == 20 # 100 - 80
assert result.output_tokens == 10
assert result.cache_read_input_tokens == 80
assert result.cache_creation_input_tokens == 0
def test_zero_cached_tokens(self):
"""Zero cached tokens should still set cache_creation to 0."""
usage = UsageInfo(
prompt_tokens=100,
completion_tokens=10,
prompt_tokens_details=PromptTokenUsageInfo(cached_tokens=0),
)
result = _build_anthropic_usage(100, 10, usage)
assert result.input_tokens == 100 # 100 - 0
assert result.cache_read_input_tokens == 0
assert result.cache_creation_input_tokens == 0
def test_all_tokens_cached(self):
"""When all tokens are cached, input_tokens should be 0."""
usage = UsageInfo(
prompt_tokens=100,
completion_tokens=10,
prompt_tokens_details=PromptTokenUsageInfo(cached_tokens=100),
)
result = _build_anthropic_usage(100, 10, usage)
assert result.input_tokens == 0
assert result.cache_read_input_tokens == 100
assert result.cache_creation_input_tokens == 0
def test_no_prompt_tokens_details(self):
"""UsageInfo without prompt_tokens_details returns no cache info."""
usage = UsageInfo(prompt_tokens=100, completion_tokens=10)
result = _build_anthropic_usage(100, 10, usage)
assert result.input_tokens == 100
assert result.cache_read_input_tokens is None
assert result.cache_creation_input_tokens is None
class TestInlineSystemMessageInMessagesArray:
"""Verify that ``role: system`` messages embedded inside the ``messages``
array are preserved in their original position.
@@ -1096,3 +1205,179 @@ class TestMessageStartIncludesTypeAndRole:
message = events[0][1]["message"]
assert message["type"] == "message"
assert message["role"] == "assistant"
class TestStreamingCacheUsageSemantics:
"""Locks in the documented streaming behavior of cache usage fields.
vLLM's OpenAI chat completion streaming only attaches
``prompt_tokens_details`` to the terminal usage chunk. The Anthropic layer
mirrors that contract: cache fields are omitted on ``message_start`` (key
absence signals "unknown") and populated on ``message_delta`` (the final
cumulative count). This is intentionally consistent with vLLM's OpenAI
behavior, even though Anthropic's upstream API populates cache fields on
``message_start``; closing that gap requires plumbing cache info into the
first chunk at the OpenAI layer, which is out of scope here.
"""
@pytest.mark.asyncio
async def test_streaming_cache_fields_absent_then_populated(self):
"""First chunk lacks prompt_tokens_details (vLLM contract);
message_start omits cache fields. The final chunk carries
prompt_tokens_details, so message_delta carries resolved values."""
async def sse_input():
yield _make_stream_chunk(
delta=DeltaMessage(role="assistant", content="hi"),
usage=UsageInfo(prompt_tokens=100, total_tokens=100),
)
yield _make_stream_chunk(finish_reason="stop")
yield _make_stream_chunk(
choices=[],
usage=UsageInfo(
prompt_tokens=100,
completion_tokens=5,
total_tokens=105,
prompt_tokens_details=PromptTokenUsageInfo(cached_tokens=80),
),
)
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)
# message_start: cache fields unknown → omitted from JSON entirely.
start_usage = events[0][1]["message"]["usage"]
assert events[0][0] == "message_start"
assert start_usage["input_tokens"] == 100
assert "cache_read_input_tokens" not in start_usage
assert "cache_creation_input_tokens" not in start_usage
# message_delta: authoritative usage with cache fields populated.
delta_usage = next(
data["usage"] for ev, data in events if ev == "message_delta"
)
assert delta_usage["input_tokens"] == 20 # 100 - 80
assert delta_usage["cache_read_input_tokens"] == 80
assert delta_usage["cache_creation_input_tokens"] == 0
@pytest.mark.asyncio
async def test_streaming_no_cache_hit(self):
"""When the final chunk reports cached_tokens=0, message_delta carries
cache fields = 0 (cache miss); message_start still omits them."""
async def sse_input():
yield _make_stream_chunk(
delta=DeltaMessage(role="assistant"),
usage=UsageInfo(prompt_tokens=50, total_tokens=50),
)
yield _make_stream_chunk(finish_reason="stop")
yield _make_stream_chunk(
choices=[],
usage=UsageInfo(
prompt_tokens=50,
completion_tokens=5,
total_tokens=55,
prompt_tokens_details=PromptTokenUsageInfo(cached_tokens=0),
),
)
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)
start_usage = events[0][1]["message"]["usage"]
delta_usage = next(
data["usage"] for ev, data in events if ev == "message_delta"
)
assert start_usage["input_tokens"] == 50
assert "cache_read_input_tokens" not in start_usage
assert "cache_creation_input_tokens" not in start_usage
assert delta_usage["input_tokens"] == 50 # 50 - 0
assert delta_usage["cache_read_input_tokens"] == 0
assert delta_usage["cache_creation_input_tokens"] == 0
@pytest.mark.asyncio
async def test_streaming_no_prompt_tokens_details_at_all(self):
"""If --enable-prompt-tokens-details is off, no chunk carries cache
info; both message_start and message_delta omit cache fields."""
async def sse_input():
yield _make_stream_chunk(
delta=DeltaMessage(role="assistant"),
usage=UsageInfo(prompt_tokens=30, total_tokens=30),
)
yield _make_stream_chunk(finish_reason="stop")
yield _make_stream_chunk(
choices=[],
usage=UsageInfo(prompt_tokens=30, completion_tokens=2, total_tokens=32),
)
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)
start_usage = events[0][1]["message"]["usage"]
delta_usage = next(
data["usage"] for ev, data in events if ev == "message_delta"
)
assert "cache_read_input_tokens" not in start_usage
assert "cache_creation_input_tokens" not in start_usage
assert "cache_read_input_tokens" not in delta_usage
assert "cache_creation_input_tokens" not in delta_usage
# ======================================================================
# Auto-detection of system-first template requirement
# ======================================================================
Q35_TEMPLATE = (
"{%- for message in messages %}"
"{%- if message.role == 'system' %}"
"{%- if not loop.first %}"
"{{- raise_exception('System message must be at the beginning.') }}"
"{%- endif %}"
"{%- endif %}"
"{%- endfor %}"
)
class TestDetectMergeInlineSystem:
"""Verify _detect_merge_inline_system auto-detection.
Tests three scenarios:
1. Template with system-first guard (e.g. Qwen) → merge needed
2. Template without restrictions → no merge, cache-friendly
3. No template provided → safe default: merge
"""
def test_qwen_template_requires_merge(self):
"""Template with loop.first guard rejects mid-conversation system."""
assert (
AnthropicServingMessages._detect_merge_inline_system(Q35_TEMPLATE) is True
)
def test_no_restriction_no_merge(self):
"""Template without restriction accepts mid-conversation system."""
assert (
AnthropicServingMessages._detect_merge_inline_system(
"{%- for message in messages %}"
"{{- message.role }}: {{ message.content }}\n"
"{%- endfor %}"
)
is False
)
def test_no_template_defaults_merge(self):
"""No chat_template → conservative default: merge."""
assert AnthropicServingMessages._detect_merge_inline_system(None) is True
@@ -0,0 +1,65 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Tests that validation_exception_handler populates the `param` field
in its error response using the Pydantic error's `loc`, even when no
custom VLLMValidationError context is present.
Previously, `param` was only populated for errors carrying a custom
VLLMValidationError in their Pydantic `ctx`. Plain validation failures
(missing fields, wrong types) left `param` as None, even though the
field name was readily available from `error['loc']`.
"""
import json
from types import SimpleNamespace
import pytest
from fastapi.exceptions import RequestValidationError
from vllm.entrypoints.serve.utils.server_utils import validation_exception_handler
def _fake_request(log_error_stack: bool = False) -> SimpleNamespace:
"""Minimal stand-in for a FastAPI Request - just enough for the
handler to read req.app.state.args.log_error_stack."""
return SimpleNamespace(
app=SimpleNamespace(
state=SimpleNamespace(args=SimpleNamespace(log_error_stack=log_error_stack))
),
state=SimpleNamespace(), # no request_metadata -> hasattr(...) is False
)
class TestValidationErrorParamFallback:
"""Ensure `param` falls back to the Pydantic error's `loc` when no
custom VLLMValidationError context is present."""
@pytest.mark.parametrize(
("error_type", "msg"),
[
("missing", "Field required"),
("list_type", "Input should be a valid list"),
],
ids=["missing-field", "wrong-type"],
)
@pytest.mark.asyncio
async def test_param_falls_back_to_loc(self, error_type: str, msg: str):
errors = [{"type": error_type, "loc": ("body", "messages"), "msg": msg}]
exc = RequestValidationError(errors)
response = await validation_exception_handler(_fake_request(), exc)
body = json.loads(response.body)
assert body["error"]["param"] == "body.messages"
@pytest.mark.asyncio
async def test_param_fallback_does_not_crash_on_non_dict_error(self):
"""Schemathesis fuzzing found that errors[0] isn't always a dict.
The fallback must not crash in that case - it should just leave
param as None instead of raising."""
exc = RequestValidationError(["some unexpected non-dict error"])
response = await validation_exception_handler(_fake_request(), exc)
body = json.loads(response.body)
assert body["error"]["param"] is None
+37
View File
@@ -428,6 +428,43 @@ def test_reshape_and_cache_flash(
torch.testing.assert_close(value_cache_compact, cloned_value_cache)
@pytest.mark.parametrize("dtype", DTYPES)
@pytest.mark.parametrize("kv_cache_dtype", KV_CACHE_DTYPE)
@pytest.mark.parametrize("kv_cache_layout", CACHE_LAYOUTS)
@pytest.mark.parametrize("implementation", RESHAPE_FLASH_IMPLEMENTATIONS)
@torch.inference_mode()
def test_reshape_and_cache_flash_unaligned_rows(
kv_cache_factory_flashinfer,
dtype: torch.dtype,
kv_cache_dtype: str,
kv_cache_layout: str,
implementation: str,
) -> None:
"""Regression test for https://github.com/vllm-project/vllm/issues/41257.
head_size=46 with num_heads=13 places KV-cache rows at byte offsets
that are not a multiple of the vector width (NHD row pitch
13*46*itemsize, HND head pitch 46*itemsize), unlike HEAD_SIZES above
which are all 16-byte multiples. The CUDA kernel used to issue
vectorized stores to those rows -> CUDA misaligned address.
"""
test_reshape_and_cache_flash(
kv_cache_factory_flashinfer,
num_tokens=42,
num_heads=13,
head_size=46,
block_size=16,
num_blocks=128,
dtype=dtype,
seed=0,
device=CUDA_DEVICES[0],
kv_cache_dtype=kv_cache_dtype,
kv_cache_layout=kv_cache_layout,
kv_scale_type="tensor",
implementation=implementation,
)
@pytest.mark.parametrize("direction", COPYING_DIRECTION)
@pytest.mark.parametrize("num_mappings", NUM_MAPPINGS)
@pytest.mark.parametrize("num_heads", NUM_HEADS)
+25 -11
View File
@@ -70,17 +70,20 @@ def generate_and_test(llm: vllm.LLM, lora_path: str, lora_id: int) -> None:
assert generated_texts[i].startswith(EXPECTED_LORA_OUTPUT[i])
@pytest.mark.skipif(
not current_platform.is_cuda(),
reason=(
"Mxfp4 LoRA on ROCm is blocked by a spawn compatibility issue. "
"The fused_moe_lora Triton kernel crashes in spawned subprocesses, "
"and vLLM forces spawn mode when HIP is initialized before "
"multiprocessing. Fixing this requires either making the LoRA "
"Triton kernel spawn-safe or pre-warming the kernel cache."
),
# TODO: make the Mxfp4MoeBackend.TRITON spawn-safe.
# For now just use TRITON_UNFUSED kernel
@pytest.mark.parametrize(
"mxfp4_use_marlin",
[
False,
pytest.param(
True,
marks=pytest.mark.skipif(
current_platform.is_rocm(), reason="marlin not supported"
),
),
],
)
@pytest.mark.parametrize("mxfp4_use_marlin", [True, False])
@pytest.mark.parametrize("specialize_active_lora", [True, False])
def test_gpt_oss_lora(
gptoss20b_lora_files,
@@ -109,7 +112,18 @@ def test_gpt_oss_lora(
@multi_gpu_test(num_gpus=2)
@pytest.mark.parametrize("fully_sharded_loras", [False, True])
@pytest.mark.parametrize("mxfp4_use_marlin", [True, False])
@pytest.mark.parametrize(
"mxfp4_use_marlin",
[
False,
pytest.param(
True,
marks=pytest.mark.skipif(
current_platform.is_rocm(), reason="marlin not supported"
),
),
],
)
def test_gpt_oss_lora_tp2(
gptoss20b_lora_files,
fully_sharded_loras,
+124
View File
@@ -482,3 +482,127 @@ def test_kernels_hidden_size(
seq_length=128,
add_inputs=True,
)
@pytest.mark.parametrize("device", DEVICES)
def test_add_lora_fused_moe_early_exit(device):
"""
Ensures add_lora_fused_moe does not invoke the LoRA kernel or
modify the output tensor when no_lora_flag_cpu is True
"""
from types import SimpleNamespace
from vllm.lora.punica_wrapper.punica_gpu import PunicaWrapperGPU
torch.set_default_device(device)
torch.accelerator.set_device_index(device)
max_loras, num_tokens = 4, 16
num_experts, top_k, max_lora_rank = 8, 2, 16
K, N = 256, 128
# build PunicaWrapperGPU with minimal lora_config mock
lora_config = SimpleNamespace(
max_loras=max_loras,
specialize_active_lora=False,
)
wrapper = PunicaWrapperGPU(
max_num_batched_tokens=num_tokens,
max_batches=num_tokens,
device=device,
lora_config=lora_config,
)
# simulate a prior LoRA batch so the internal mapping is
# populated with stale LoRA IDs
lora_mapping = torch.zeros(
num_tokens,
dtype=torch.int32,
device=device,
)
lora_mapping[:8] = 1
lora_mapping[8:] = 2
wrapper.token_mapping_meta.prepare_tensors(lora_mapping)
# simulate a base-model batch (all -1)
base_mapping = torch.full(
(num_tokens,),
-1,
dtype=torch.int32,
device=device,
)
wrapper.token_mapping_meta.prepare_tensors(base_mapping)
assert wrapper.token_mapping_meta.no_lora_flag_cpu[0].item() is True
# dummy tensors for add_lora_fused_moe
y = torch.rand(num_tokens, top_k, N, dtype=torch.bfloat16, device=device)
y_snapshot = y.clone()
x = torch.rand(num_tokens, K, dtype=torch.bfloat16, device=device)
lora_a_stacked = (
torch.rand(
max_loras,
num_experts,
max_lora_rank,
K,
dtype=torch.bfloat16,
device=device,
),
)
lora_b_stacked = (
torch.rand(
max_loras,
num_experts,
N,
max_lora_rank,
dtype=torch.bfloat16,
device=device,
),
)
topk_weights = torch.ones(
num_tokens,
top_k,
dtype=torch.float32,
device=device,
)
adapter_enabled = torch.ones(
max_loras + 1,
dtype=torch.int32,
device=device,
)
shrink_config = expand_config = {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 32,
"BLOCK_SIZE_K": 64,
"GROUP_SIZE_M": 1,
"NUM_WARPS": 4,
"NUM_STAGES": 3,
"SPLIT_K": 1,
}
# call add_lora_fused_moe - the early exit should prevent any
# modification to the output
wrapper.add_lora_fused_moe(
y=y,
x=x,
lora_a_stacked=lora_a_stacked,
lora_b_stacked=lora_b_stacked,
topk_weights=topk_weights,
sorted_token_ids=None,
expert_ids=torch.zeros(
num_tokens * top_k,
dtype=torch.int32,
device=device,
),
num_tokens_post_padded=None,
max_lora_rank=max_lora_rank,
top_k_num=top_k,
shrink_config=shrink_config,
expand_config=expand_config,
adapter_enabled=adapter_enabled,
)
assert torch.equal(y, y_snapshot), (
"add_lora_fused_moe modified output tensor despite no_lora_flag_cpu=True"
)
@@ -130,8 +130,12 @@ def test_models(
monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1")
if model == "TitanML/tiny-mixtral":
# Untrained model: near-uniform logits make argmax sensitive to
# AITER's bfloat16 rounding error in plain rms_norm.
# AITER's bfloat16 rounding error. Route the plain rms_norm and the
# fused MoE (whose near-uniform router logits flip expert selection
# under ~1 ULP drift) through the native kernels for this model.
# See ROCm/aiter#3806 for the tracking issue and minimal repro.
monkeypatch.setenv("VLLM_ROCM_USE_AITER_RMSNORM", "0")
monkeypatch.setenv("VLLM_ROCM_USE_AITER_MOE", "0")
elif use_rocm_aiter and model not in AITER_MODEL_LIST:
# Skip model that are not using AITER tests.
# When more AITER kernels are added, this list will not be
@@ -25,7 +25,7 @@ TEST_IMAGE_NAMES = [
]
MAX_MODEL_LEN = 8192
REQUESTS_PER_ROUND = 4
WARMUP_ROUNDS = 1
WARMUP_ROUNDS = 2
MEASURED_ROUNDS = 16
GPU_GROWTH_THRESHOLD_MIB = 0
CPU_PEAK_GROWTH_THRESHOLD_MIB = 0
+18
View File
@@ -615,6 +615,24 @@ class TestArgConverter:
assert result["command"] == "ls -la"
assert result["desc"] == "\npartial value"
def test_partial_value_with_angle_bracket(self):
from vllm.parser.qwen3 import (
_qwen3_arg_converter,
)
raw = "<parameter=expr>x<5"
result = json.loads(_qwen3_arg_converter(raw, partial=True))
assert result == {"expr": "x<5"}
def test_partial_value_with_angle_bracket_and_complete_param(self):
from vllm.parser.qwen3 import (
_qwen3_arg_converter,
)
raw = "<parameter=city>Tokyo</parameter>\n<parameter=expr>x<5"
result = json.loads(_qwen3_arg_converter(raw, partial=True))
assert result == {"city": "Tokyo", "expr": "x<5"}
class TestSchemaAwareTypeCoercion:
"""Verify that _fix_arg_types corrects miscoerced values using the
+93 -2
View File
@@ -30,6 +30,7 @@ from vllm.entrypoints.openai.chat_completion.protocol import (
)
from vllm.parser.engine.registered_adapters import (
Gemma4Parser,
Glm47MoeParser,
MinimaxM2Parser,
NemotronV3Parser,
Qwen3Parser,
@@ -142,6 +143,12 @@ SCENARIOS: list[Scenario] = [
tool_calls=[_READ_TOOL],
after_tool_response=True,
),
Scenario(
id="empty-tool-block",
description="Empty tool block followed by content (edge case recovery)",
content="Content after empty tools.",
tool_calls=[],
),
]
@@ -343,8 +350,11 @@ def _qwen3_segments(scenario: Scenario) -> list[tuple[str, bool]]:
segs: list[tuple[str, bool]] = []
if scenario.reasoning is not None:
segs.append((scenario.reasoning, False))
if scenario.content is not None or scenario.tool_calls:
if scenario.content is not None or scenario.tool_calls is not None:
segs.append(("</think>", True))
if scenario.tool_calls is not None and not scenario.tool_calls:
segs.append(("<tool_call>", True))
segs.append(("</tool_call>", True))
if scenario.content is not None:
segs.append((scenario.content, False))
if scenario.tool_calls:
@@ -436,8 +446,11 @@ def _minimax_m2_segments(scenario: Scenario) -> list[tuple[str, bool]]:
segs: list[tuple[str, bool]] = []
if scenario.reasoning is not None:
segs.append((scenario.reasoning, False))
if scenario.content is not None or scenario.tool_calls:
if scenario.content is not None or scenario.tool_calls is not None:
segs.append(("</think>", True))
if scenario.tool_calls is not None and not scenario.tool_calls:
segs.append(("<minimax:tool_call>", True))
segs.append(("</minimax:tool_call>", True))
if scenario.content is not None:
segs.append((scenario.content, False))
if scenario.tool_calls:
@@ -533,6 +546,9 @@ def _gemma4_segments(scenario: Scenario) -> list[tuple[str, bool]]:
segs.append((_GEMMA4_THOUGHT_PREFIX, False))
segs.append((scenario.reasoning, False))
segs.append(("<channel|>", True))
if scenario.tool_calls is not None and not scenario.tool_calls:
segs.append(("<|tool_call>", True))
segs.append(("<tool_call|>", True))
if scenario.content is not None:
segs.append((scenario.content, False))
if scenario.tool_calls:
@@ -571,6 +587,80 @@ def _build_nemotron_v3(scenario: Scenario, validate: bool = True) -> Sample:
)
# ── GLM-4.7 MoE (XML tool format, starts in REASONING) ──────────────
_GLM47_MOE_VOCAB: dict[str, int] = {
"<think>": 50,
"</think>": 51,
"<tool_call>": 60,
"</tool_call>": 61,
"<arg_key>": 62,
"</arg_key>": 63,
"<arg_value>": 64,
"</arg_value>": 65,
}
def _glm47_moe_arg_value(value: Any) -> str:
if isinstance(value, bool):
return "true" if value else "false"
if isinstance(value, (int, float)):
return str(value)
if isinstance(value, str):
return value
return json.dumps(value, ensure_ascii=False)
def _glm47_moe_tool_segments(tc: ToolCallSpec) -> list[tuple[str, bool]]:
segs: list[tuple[str, bool]] = [
("<tool_call>", True),
(tc.name, False),
]
for key, value in tc.arguments.items():
segs.extend(
[
("<arg_key>", True),
(key, False),
("</arg_key>", True),
("<arg_value>", True),
(_glm47_moe_arg_value(value), False),
("</arg_value>", True),
]
)
segs.append(("</tool_call>", True))
return segs
def _glm47_moe_segments(scenario: Scenario) -> list[tuple[str, bool]]:
segs: list[tuple[str, bool]] = []
if scenario.reasoning is not None:
segs.append((scenario.reasoning, False))
if scenario.content is not None or scenario.tool_calls:
segs.append(("</think>", True))
if scenario.content is not None:
segs.append((scenario.content, False))
if scenario.tool_calls:
for tc in scenario.tool_calls:
segs.extend(_glm47_moe_tool_segments(tc))
return segs
def _build_glm47_moe(scenario: Scenario, validate: bool = True) -> Sample:
sample = _make_sample(
sample_id=f"glm47_moe-{scenario.id}",
description=scenario.description,
vocab=_GLM47_MOE_VOCAB,
segments=_glm47_moe_segments(scenario),
expected_reasoning=scenario.reasoning if scenario.reasoning is not None else "",
expected_content=_qwen3_expected_content(scenario),
expected_tool_calls=_expected_tc(scenario),
tools=_expected_tools(scenario),
)
if validate:
_validate_sample(sample, Glm47MoeParser)
return sample
# ── Registry and public API ──────────────────────────────────────────
_BUILDERS: dict[str, Any] = {
@@ -578,6 +668,7 @@ _BUILDERS: dict[str, Any] = {
"gemma4": _build_gemma4,
"minimax_m2": _build_minimax_m2,
"nemotron_v3": _build_nemotron_v3,
"glm47_moe": _build_glm47_moe,
}
@@ -11,7 +11,7 @@ parser_name = "glm45"
start_token = "<think>"
end_token = "</think>"
REASONING_MODEL_NAME = "zai-org/GLM-4.5"
REASONING_MODEL_NAME = "zai-org/GLM-4.7"
@pytest.fixture(scope="module")
@@ -35,18 +35,32 @@ WITH_THINK_STREAM = {
WITHOUT_THINK = {
"output": "This is the rest",
"reasoning": None,
"content": "This is the rest",
"reasoning": "This is the rest",
"content": None,
"is_reasoning_end": False,
}
WITHOUT_THINK_STREAM = {
"output": "This is the rest",
"reasoning": None,
"content": "This is the rest",
"reasoning": "This is the rest",
"content": None,
"is_reasoning_end": False,
}
WITHOUT_OPEN_THINK = {
"output": "This is a reasoning section</think>This is the rest",
"reasoning": "This is a reasoning section",
"content": "This is the rest",
"is_reasoning_end": True,
}
WITHOUT_OPEN_THINK_STREAM = {
"output": "This is a reasoning section</think>This is the rest",
"reasoning": "This is a reasoning section",
"content": "This is the rest",
"is_reasoning_end": True,
}
COMPLETE_REASONING = {
"output": "<think>This is a reasoning section</think>",
"reasoning": "This is a reasoning section",
@@ -61,8 +75,8 @@ MULTILINE_REASONING = {
}
ONLY_OPEN_TAG = {
"output": "<think>This is a reasoning section",
"reasoning": None,
"content": "<think>This is a reasoning section",
"reasoning": "This is a reasoning section",
"content": None,
"is_reasoning_end": False,
}
@@ -94,6 +108,16 @@ TEST_CASES = [
WITHOUT_THINK_STREAM,
id="without_think_stream",
),
pytest.param(
False,
WITHOUT_OPEN_THINK,
id="without_open_think",
),
pytest.param(
True,
WITHOUT_OPEN_THINK_STREAM,
id="without_open_think_stream",
),
pytest.param(
False,
COMPLETE_REASONING,
+145
View File
@@ -0,0 +1,145 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# Schema/aliasing tests for the AITER FP8 quantization custom ops.
#
# These use torch.library.opcheck, whose test_schema check catches custom ops
# whose implementation aliases an input that the registered schema declares as
# non-aliasing -- the failure mode behind the rocm_aiter_per_tensor_quant
# regression (a returned scale that aliased the input scale).
#
# Skipped if AITER is not installed or the platform is not ROCm.
import importlib.util
import pytest
import torch
# this import statement is needed to ensure the ops are registered
from vllm._aiter_ops import rocm_aiter_ops
from vllm.platforms import current_platform
aiter_available = importlib.util.find_spec("aiter") is not None
pytestmark = pytest.mark.skipif(
not (current_platform.is_rocm() and aiter_available),
reason="AITER ops are only available on ROCm with aiter package installed",
)
FP8_DTYPE = current_platform.fp8_dtype()
def _x(M=128, N=4096):
return torch.randn((M, N), dtype=torch.float16, device="cuda")
# The in-place per-tensor op takes the fp8 output buffer as an input, which
# opcheck's test_schema cannot exercise ("mul_cuda" is unimplemented for fp8),
# so restrict to the utils that run on fp8 inputs. The aliasing contract for
# this op is instead covered by test_per_tensor_quant_torch_compile below.
_INPLACE_OPCHECK_UTILS = (
"test_faketensor",
"test_aot_dispatch_dynamic",
"test_autograd_registration",
)
def test_per_tensor_quant_static_schema():
"""Static per-tensor: caller provides scale (the aliasing regression)."""
x = _x()
out = torch.empty_like(x, dtype=FP8_DTYPE)
scale = torch.ones(1, dtype=torch.float32, device="cuda")
torch.library.opcheck(
torch.ops.vllm.rocm_aiter_per_tensor_quant,
(out, x, scale, False),
test_utils=_INPLACE_OPCHECK_UTILS,
)
def test_per_tensor_quant_dynamic_schema():
"""Dynamic per-tensor: op computes scale into the caller's buffer."""
x = _x()
out = torch.empty_like(x, dtype=FP8_DTYPE)
scale = torch.empty(1, dtype=torch.float32, device="cuda")
torch.library.opcheck(
torch.ops.vllm.rocm_aiter_per_tensor_quant,
(out, x, scale, True),
test_utils=_INPLACE_OPCHECK_UTILS,
)
def test_per_token_quant_dynamic_schema():
"""Dynamic per-token: op computes scale into a freshly allocated buffer."""
x = _x()
torch.library.opcheck(
torch.ops.vllm.rocm_aiter_per_token_quant,
(x, FP8_DTYPE, None),
)
def test_group_fp8_quant_schema():
"""Dynamic per-token-group quant."""
x = _x()
torch.library.opcheck(
torch.ops.vllm.rocm_aiter_group_fp8_quant,
(x, 128),
)
@pytest.mark.parametrize("dynamic", [True, False])
def test_per_tensor_quant_matches_native(dynamic):
"""Wrapper output matches the native scaled_fp8_quant reference."""
from vllm import _custom_ops as ops
torch.manual_seed(0)
x = _x()
if dynamic:
scale_in = None
else:
scale_in = torch.tensor([0.5], dtype=torch.float32, device="cuda")
out, scale = rocm_aiter_ops.per_tensor_quant(x, FP8_DTYPE, scale_in)
ref_out, ref_scale = ops.scaled_fp8_quant(x, scale_in)
assert out.shape == x.shape
assert out.dtype == FP8_DTYPE
assert scale.shape == ref_scale.shape
if not dynamic:
# static scale is passed through unchanged
assert torch.equal(scale, scale_in)
# Compare dequantized values to be robust to 1-ULP fp8 boundary flips.
deq = out.to(torch.float32) * scale
ref_deq = ref_out.to(torch.float32) * ref_scale
torch.testing.assert_close(deq, ref_deq, rtol=2e-2, atol=2e-2)
@pytest.mark.parametrize("dynamic", [True, False])
def test_per_tensor_quant_torch_compile(monkeypatch, dynamic):
"""per_tensor_quant compiles under inductor without an aliasing error.
Forces the custom-op aliasing check to error (it is otherwise only a
warning outside CI), so a regression that returns an input-aliasing
scale fails here regardless of the CI env var.
"""
aliasing_cfg = pytest.importorskip("torch._functorch.config")
monkeypatch.setattr(
aliasing_cfg, "error_on_custom_op_aliasing", True, raising=False
)
x = _x()
scale = None if dynamic else torch.tensor([0.5], dtype=torch.float32, device="cuda")
def fn(x, s):
return rocm_aiter_ops.per_tensor_quant(x, FP8_DTYPE, s)
compiled = torch.compile(fn, fullgraph=True, backend="inductor", dynamic=False)
out_eager, scale_eager = fn(x, scale)
out_compiled, scale_compiled = compiled(x, scale)
assert out_compiled.shape == out_eager.shape
torch.testing.assert_close(
out_compiled.to(torch.float32) * scale_compiled,
out_eager.to(torch.float32) * scale_eager,
rtol=2e-2,
atol=2e-2,
)
+1 -1
View File
@@ -188,7 +188,7 @@ def test_v2_model_runner_env_tri_state(monkeypatch, env_value, expected):
is_moe=False,
is_quantized=True,
),
False,
True,
),
(
SimpleNamespace(
@@ -16,7 +16,7 @@ from vllm.entrypoints.openai.chat_completion.protocol import (
from vllm.tokenizers import get_tokenizer
from vllm.tool_parsers.glm47_moe_tool_parser import Glm47MoeModelToolParser
MODEL = "zai-org/GLM-4.5"
MODEL = "zai-org/GLM-4.7"
@pytest.fixture(scope="module")
@@ -136,9 +136,10 @@ class TestGlm47Streaming:
_reset(glm47_tool_parser)
chunks = ["<tool_call>", "get_current_date", "</tool_call>"]
current_text = ""
deltas = []
for chunk in chunks:
current_text += chunk
glm47_tool_parser.extract_tool_calls_streaming(
delta = glm47_tool_parser.extract_tool_calls_streaming(
previous_text="",
current_text=current_text,
delta_text=chunk,
@@ -147,7 +148,23 @@ class TestGlm47Streaming:
delta_token_ids=[],
request=mock_request,
)
assert len(glm47_tool_parser.prev_tool_call_arr) >= 1
if delta:
deltas.append(delta)
tool_calls = [
tool_call for delta in deltas for tool_call in (delta.tool_calls or [])
]
names = [
tool_call.function.name
for tool_call in tool_calls
if tool_call.function and tool_call.function.name
]
arguments = [
tool_call.function.arguments
for tool_call in tool_calls
if tool_call.function and tool_call.function.arguments
]
assert names == ["get_current_date"]
assert "".join(arguments) == "{}"
def test_with_args(self, glm47_tool_parser, mock_request):
_reset(glm47_tool_parser)
@@ -161,9 +178,10 @@ class TestGlm47Streaming:
"</tool_call>",
]
current_text = ""
deltas = []
for chunk in chunks:
current_text += chunk
glm47_tool_parser.extract_tool_calls_streaming(
delta = glm47_tool_parser.extract_tool_calls_streaming(
previous_text="",
current_text=current_text,
delta_text=chunk,
@@ -172,5 +190,13 @@ class TestGlm47Streaming:
delta_token_ids=[],
request=mock_request,
)
args = json.loads(glm47_tool_parser.prev_tool_call_arr[0]["arguments"])
if delta:
deltas.append(delta)
arguments = [
tool_call.function.arguments
for delta in deltas
for tool_call in (delta.tool_calls or [])
if tool_call.function and tool_call.function.arguments
]
args = json.loads("".join(arguments))
assert args["city"] == "Beijing"
File diff suppressed because it is too large Load Diff
@@ -1300,6 +1300,73 @@ def test_streaming_multi_param_single_chunk(qwen3_tool_parser, qwen3_tokenizer):
assert args["unit"] == "fahrenheit"
def test_streaming_complete_tool_call_single_delta(qwen3_tool_parser):
"""Regression: one delta may contain a complete tool call."""
request = ChatCompletionRequest(model=MODEL, messages=[])
from tests.tool_parsers.utils import (
run_tool_extraction_streaming,
)
reconstructor = run_tool_extraction_streaming(
qwen3_tool_parser,
[
(
"<tool_call>\n"
"<function=get_current_weather>\n"
"<parameter=city>\nDallas\n</parameter>\n"
"<parameter=state>\nTX\n</parameter>\n"
"</function>\n"
"</tool_call>"
)
],
request,
assert_one_tool_per_delta=False,
)
assert len(reconstructor.tool_calls) == 1
assert reconstructor.tool_calls[0].function.name == "get_current_weather"
args = json.loads(reconstructor.tool_calls[0].function.arguments)
assert args == {"city": "Dallas", "state": "TX"}
def test_streaming_next_tool_call_starts_in_close_delta(qwen3_tool_parser):
"""Regression: a close delta may also contain the next tool call."""
request = ChatCompletionRequest(model=MODEL, messages=[])
from tests.tool_parsers.utils import (
run_tool_extraction_streaming,
)
reconstructor = run_tool_extraction_streaming(
qwen3_tool_parser,
[
"<tool_call>\n",
"<function=get_current_weather>\n",
"<parameter=city>\nDallas\n</parameter>\n",
"<parameter=state>\nTX\n</parameter>\n",
"</function>",
(
"\n</tool_call>\n"
"<tool_call>\n"
"<function=get_current_weather>\n"
"<parameter=city>\nOrlando\n</parameter>\n"
"<parameter=state>\nFL\n</parameter>\n"
"</function>\n"
"</tool_call>"
),
],
request,
assert_one_tool_per_delta=False,
)
assert len(reconstructor.tool_calls) == 2
first_args = json.loads(reconstructor.tool_calls[0].function.arguments)
second_args = json.loads(reconstructor.tool_calls[1].function.arguments)
assert first_args == {"city": "Dallas", "state": "TX"}
assert second_args == {"city": "Orlando", "state": "FL"}
def test_no_double_serialization_string_args(qwen3_tool_parser):
"""Regression: string arguments must not be double-serialized (PR #35615)."""
tools = [
+67
View File
@@ -1,6 +1,8 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import json
import pytest
from vllm.tool_parsers.utils import (
@@ -91,6 +93,71 @@ class TestCoerceToSchemaType:
def test_invalid_number_fallback(self):
assert coerce_to_schema_type("abc", "number") == "abc"
class TestNonFiniteNumbers:
"""Non-finite numeric strings must not crash and must coerce to a
JSON-serializable value.
Regression: ``int(float("inf"))`` raised an uncaught ``OverflowError``
(only ``ValueError``/``TypeError`` were handled), and ``"1e999"``
round-tripped through ``json.loads`` to a float ``inf`` that
``json.dumps`` renders as invalid JSON ``Infinity``.
"""
@pytest.mark.parametrize(
"value", ["inf", "-inf", "Infinity", "1e999", "nan", "-nan"]
)
def test_non_finite_number_does_not_crash(self, value):
# Must not raise (previously OverflowError for inf/1e999/Infinity).
result = coerce_to_schema_type(value, "number")
# Result must serialize to valid, finite JSON and round-trip.
assert json.loads(json.dumps(result)) == result
@pytest.mark.parametrize("value", ["inf", "-inf", "1e999"])
def test_non_finite_number_preserved_as_string(self, value):
assert coerce_to_schema_type(value, "number") == value
@pytest.mark.parametrize("value", ["inf", "1e999", "Infinity"])
def test_non_finite_integer_not_float_inf(self, value):
result = coerce_to_schema_type(value, "integer")
assert isinstance(result, str)
assert result == value
class TestNonFiniteContainers:
"""Non-finite floats nested in object/array values must not produce
invalid JSON.
Regression: the ``object``/``array`` branch returned
``json.loads(value)`` directly, so ``"[1e999]"`` became ``[inf]`` and
``'{"x": Infinity}'`` became ``{"x": inf}`` -- values that
``json.dumps`` later renders as invalid JSON (``Infinity``/``NaN``).
"""
@pytest.mark.parametrize(
"value", ["[1e999]", "[1, 2, 1e999]", "[NaN]", "[-Infinity]"]
)
def test_array_with_non_finite_preserved_as_string(self, value):
result = coerce_to_schema_type(value, "array")
assert result == value
assert json.loads(json.dumps(result)) == result
@pytest.mark.parametrize(
"value", ['{"x": 1e999}', '{"x": Infinity}', '{"a": [1e999, 2]}']
)
def test_object_with_non_finite_preserved_as_string(self, value):
result = coerce_to_schema_type(value, "object")
assert result == value
assert json.loads(json.dumps(result)) == result
def test_finite_array_still_coerced(self):
assert coerce_to_schema_type("[1, 2, 3]", "array") == [1, 2, 3]
def test_finite_object_still_coerced(self):
assert coerce_to_schema_type('{"a": 1}', "object") == {"a": 1}
def test_unknown_type_non_finite_falls_back_to_string(self):
# Exercises the final json.loads fallback path.
assert coerce_to_schema_type("1e999", "unknown_type") == "1e999"
class TestBooleanType:
def test_true(self):
assert coerce_to_schema_type("true", "boolean") is True
@@ -1,6 +1,8 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import threading
import time
from unittest.mock import MagicMock, patch
from vllm.config import set_current_vllm_config
@@ -406,7 +408,9 @@ def test_lookup_key_client_lookup_prepends_typed_tag():
fake_socket = mock_make_socket.return_value
fake_socket.recv.return_value = (5).to_bytes(4, "big")
assert client.lookup(token_len=128, block_hashes=[]) == 5
# Blocking lookup (non_block defaults to False) runs on the executor and
# returns the resolved hit length.
assert client.lookup("req0", token_len=128, block_hashes=[]) == 5
sent_frames = fake_socket.send_multipart.call_args[0][0]
assert sent_frames[0] == protocol.LOOKUP_MSG
@@ -435,6 +439,127 @@ def test_lookup_key_client_reset_uses_typed_protocol():
assert client.reset() is False
def _poll_lookup(client, req_id, token_len=128, block_hashes=(), timeout=5.0):
"""Drive non-blocking lookup until the executor completes it."""
deadline = time.monotonic() + timeout
while time.monotonic() < deadline:
result = client.lookup(req_id, token_len, list(block_hashes), non_block=True)
if result is not None:
return result
time.sleep(0.005)
return None
def _gated_recv(gate: threading.Event, value: int):
"""Mock recv side-effect that blocks until ``gate`` is set, so the
executor's lookup can be held pending deterministically."""
def recv():
gate.wait()
return value.to_bytes(4, "big")
return recv
def test_lookup_key_client_non_block_lookup_async():
"""Non-blocking lookup defers to the executor: None first, hit once the
Future resolves."""
vllm_config = _make_vllm_config()
with patch(
"vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
"worker.make_zmq_socket"
) as mock_make_socket:
client = worker.LookupKeyClient(vllm_config)
fake_socket = mock_make_socket.return_value
# Hold the executor's lookup pending until we release the gate.
gate = threading.Event()
fake_socket.recv.side_effect = _gated_recv(gate, 7)
# First query submits the lookup and returns None while it is in flight.
assert client.lookup("req1", 128, [], non_block=True) is None
# Release the executor; a later poll returns the hit length.
gate.set()
assert _poll_lookup(client, "req1") == 7
# Future is consumed (popped) on read.
assert "req1" not in client.futures
def test_lookup_key_client_discard_clears_state():
"""discard() drops a completed lookup Future so it is not served stale."""
vllm_config = _make_vllm_config()
with patch(
"vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
"worker.make_zmq_socket"
) as mock_make_socket:
client = worker.LookupKeyClient(vllm_config)
fake_socket = mock_make_socket.return_value
gate = threading.Event()
fake_socket.recv.side_effect = _gated_recv(gate, 9)
# Submit while gated so the call returns None and the Future stays in
# `futures` (unconsumed) once it resolves.
assert client.lookup("req2", 128, [], non_block=True) is None
gate.set()
deadline = time.monotonic() + 5.0
while time.monotonic() < deadline:
if client.futures["req2"].done():
break
time.sleep(0.005)
# discard() drops the completed result before any lookup consumes it.
client.discard("req2")
assert "req2" not in client.futures
# A fresh query re-submits rather than returning a stale value: hold the
# gate so the resubmitted lookup stays in flight.
gate.clear()
assert client.lookup("req2", 128, [], non_block=True) is None
gate.set() # release the executor so the worker thread can drain
def test_get_num_new_matched_tokens_async_defers_then_reports():
"""Async lookup returns (None, False) until ready, then the hit count."""
vllm_config = create_vllm_config(
kv_connector="MooncakeStoreConnector",
kv_role="kv_both",
kv_connector_extra_config={"lookup_async": True},
)
kv_cache_config = _make_kv_cache_config()
with (
set_current_vllm_config(vllm_config),
patch(
"vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
"scheduler.LookupKeyClient"
) as mock_client_cls,
):
sched = scheduler.MooncakeStoreScheduler(vllm_config, kv_cache_config)
assert sched.lookup_async is True
mock_client = mock_client_cls.return_value
block_size = sched._block_size
request = MagicMock()
request.request_id = "r1"
request.num_tokens = 4 * block_size
request.block_hashes = []
# Lookup not ready -> defer.
mock_client.lookup.return_value = None
assert sched.get_num_new_matched_tokens(request, 0) == (None, False)
assert "r1" not in sched.load_specs
# Lookup ready with a hit -> report need_to_allocate + async-load flag.
hit = 3 * block_size
mock_client.lookup.return_value = hit
need, load_async = sched.get_num_new_matched_tokens(request, 0)
assert need == hit
assert load_async == sched.load_async
assert sched.load_specs["r1"].kvpool_cached_tokens == hit
def test_protocol_tags_are_distinct_and_non_empty():
"""Protocol tags must be unique and non-empty to avoid collision."""
tags = {protocol.LOOKUP_MSG, protocol.RESET_MSG}
@@ -16,6 +16,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.scheduler impor
def _make_bare_scheduler() -> MooncakeStoreScheduler:
scheduler = object.__new__(MooncakeStoreScheduler)
scheduler.kv_role = "kv_both"
scheduler.lookup_async = False
scheduler._block_size = 16
scheduler.load_specs = {}
scheduler._preempted_req_ids = set()
@@ -405,7 +406,13 @@ class _StubLookupClient:
def __init__(self, hit_tokens: int) -> None:
self._hit_tokens = hit_tokens
def lookup(self, token_len: int, block_hashes: list[bytes]) -> int:
def lookup(
self,
req_id: str,
token_len: int,
block_hashes: list[bytes],
non_block: bool = False,
) -> int:
return self._hit_tokens
@@ -175,14 +175,17 @@ class _FakeModelConfig:
def _make_vllm_config(
*, extra_config: dict[str, object] | None = None
*,
extra_config: dict[str, object] | None = None,
rank: int = 0,
decode_context_parallel_size: int = 1,
) -> SimpleNamespace:
return SimpleNamespace(
model_config=_FakeModelConfig(),
parallel_config=SimpleNamespace(
pipeline_parallel_size=1,
rank=0,
decode_context_parallel_size=1,
rank=rank,
decode_context_parallel_size=decode_context_parallel_size,
prefill_context_parallel_size=1,
),
kv_transfer_config=_FakeKVTransferConfig(extra_config=extra_config),
@@ -231,13 +234,23 @@ def _install_fake_mooncake(monkeypatch, store_instance: MagicMock):
return FakeReplicateConfig
def _patch_worker_runtime(monkeypatch, *, local_ip: str = "10.0.0.7") -> None:
def _patch_worker_runtime(
monkeypatch,
*,
local_ip: str = "10.0.0.7",
tp_rank: int = 0,
tp_size: int = 1,
dcp_size: int = 1,
) -> None:
single_rank_group = SimpleNamespace(world_size=1, rank_in_group=0)
# DCP groups are contiguous splits of the TP group (see
# parallel_state.py), so dcp_rank == tp_rank % dcp_size.
dcp_group = SimpleNamespace(world_size=dcp_size, rank_in_group=tp_rank % dcp_size)
monkeypatch.setattr(worker, "get_mooncake_dp_engine_index", lambda _: 0)
monkeypatch.setattr(worker, "get_tensor_model_parallel_rank", lambda: 0)
monkeypatch.setattr(worker, "get_tensor_model_parallel_world_size", lambda: 1)
monkeypatch.setattr(worker, "get_tensor_model_parallel_rank", lambda: tp_rank)
monkeypatch.setattr(worker, "get_tensor_model_parallel_world_size", lambda: tp_size)
monkeypatch.setattr(worker, "get_pcp_group", lambda: single_rank_group)
monkeypatch.setattr(worker, "get_dcp_group", lambda: single_rank_group)
monkeypatch.setattr(worker, "get_dcp_group", lambda: dcp_group)
monkeypatch.setattr(worker, "get_ip", lambda: local_ip)
@@ -884,6 +897,66 @@ def test_requester_worker_init_builds_replicate_config_for_preferred_segment(
assert w.store_replicate_config.preferred_segment == "10.0.0.7:50053"
@pytest.mark.parametrize("dcp_size", [1, 4])
def test_worker_put_striding_covers_every_rank_get_namespace(
tmp_path, monkeypatch, dcp_size
):
"""Every key a rank GETs must have been PUT by some rank.
When num_kv_head < tp_size, ranks holding the same KV heads stripe
their PUTs across one shared key namespace. That dedup is only valid
when those ranks really share a namespace: with DCP > 1 each rank GETs
every key from its own ``@dcpN`` namespace, so striding must be
disabled.
"""
tp_size = 4
store = MagicMock()
store.setup.return_value = 0
_install_fake_mooncake(monkeypatch, store)
monkeypatch.setenv(
"MOONCAKE_CONFIG_PATH",
_write_mooncake_config(
tmp_path,
{
"metadata_server": "http://metadata/endpoint",
"protocol": "tcp",
"device_name": "",
"master_server_address": "10.0.0.7:50051",
},
),
)
# _FakeModelConfig has num_kv_head=1 < tp_size, which enables striding.
block_hashes = [f"hash-{i}".encode() for i in range(4)]
put_keys: set[str] = set()
get_keys_per_rank: dict[int, set[str]] = {}
for tp_rank in range(tp_size):
_patch_worker_runtime(
monkeypatch, tp_rank=tp_rank, tp_size=tp_size, dcp_size=dcp_size
)
w = worker.MooncakeStoreWorker(
_make_vllm_config(rank=tp_rank, decode_context_parallel_size=dcp_size),
_make_kv_cache_config(),
)
db = w.token_dbs[0]
token_len = len(block_hashes) * db.block_size
keys = [
key.to_string() for _, _, key in db.process_tokens(token_len, block_hashes)
]
assert len(keys) == len(block_hashes)
# PUT side: mirrors KVCacheStoreSendingThread's striding slice.
put_keys.update(keys[w.tp_rank % w.put_step :: w.put_step])
# GET side: KVCacheStoreRecvingThread fetches every key.
get_keys_per_rank[tp_rank] = set(keys)
for tp_rank, rank_keys in get_keys_per_rank.items():
missing = rank_keys - put_keys
assert not missing, (
f"tp_rank={tp_rank} would GET {len(missing)}/{len(rank_keys)} keys "
f"that no rank PUT (Mooncake OBJECT_NOT_FOUND): {sorted(missing)}"
)
# ---------------------------------------------------------------------------
# Helpers for register_kv_caches tests
# ---------------------------------------------------------------------------
+14
View File
@@ -64,6 +64,7 @@ def make_mapper_from_offloading_spec(**kwargs) -> FileMapper:
"dcp_size", 1
)
mock_vllm_config.parallel_config.rank = kwargs.get("rank", 0)
mock_vllm_config.use_v2_model_runner = kwargs.get("use_v2_model_runner", False)
mock_kv_cache_config = MagicMock()
mock_kv_cache_config.kv_cache_groups = kwargs.get("kv_cache_groups", [])
@@ -210,3 +211,16 @@ def test_parallel_agnostic_excludes_mla():
)
assert fm.fields["tp_size"] == 2
assert fm.rank == 1
def test_parallel_agnostic_disabled_on_v2_model_runner():
# V2's KV layout is not known to be parallelism-invariant: don't collapse.
fm = make_mapper_from_offloading_spec(
tp_size=2,
rank=1,
kv_cache_groups=[_full_attention_group()],
use_v2_model_runner=True,
parallel_agnostic=True,
)
assert fm.fields["tp_size"] == 2
assert fm.rank == 1
@@ -37,6 +37,7 @@ def _make_vllm_config():
decode_context_parallel_size=1,
rank=0,
),
use_v2_model_runner=False,
)
+6 -2
View File
@@ -1002,7 +1002,11 @@ def test_propose(method, attn_backend, num_speculative_tokens, monkeypatch):
assert torch.equal(result, expected_tokens)
def test_propose_stores_probabilistic_draft_probs(monkeypatch):
@pytest.mark.parametrize(
"attn_backend",
["ROCM_ATTN", "TRITON_ATTN"] if current_platform.is_rocm() else ["FLASH_ATTN"],
)
def test_propose_stores_probabilistic_draft_probs(attn_backend, monkeypatch):
device = torch.device(DEVICE_TYPE)
batch_size = 2
seq_lens = [5, 3]
@@ -1053,7 +1057,7 @@ def test_propose_stores_probabilistic_draft_probs(monkeypatch):
)
attn_metadata_builder_cls, _ = try_get_attention_backend(
AttentionBackendEnum.FLASH_ATTN
AttentionBackendEnum[attn_backend]
)
attn_metadata_builder = attn_metadata_builder_cls(
kv_cache_spec=create_standard_kv_cache_spec(proposer.vllm_config),
+227
View File
@@ -0,0 +1,227 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Tests for the Model Runner V2 Gumbel-max sampling kernel.
Accuracy: define a target categorical distribution as a non-negative int64
count tensor summing to N, turn it into logits (= log(count)), sample many
times with `gumbel_sample`, and check the empirical distribution matches.
The count tensor is deliberately heavy-tailed (one dominant token, the rest
~18 logits below). That tail is the sensitive part: the fp32 Gumbel noise must
reach ~18 to ever sample it. A flat distribution would keep every token within
a few logits of the top and would not exercise the noise tail at all.
"""
import math
import pytest
import torch
pytest.importorskip("triton")
if not torch.cuda.is_available():
pytest.skip("CUDA required for Gumbel sampler tests", allow_module_level=True)
from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample
DEVICE = "cuda"
VOCAB_SIZE = 200_000
NUM_SAMPLES = 500_000
# Dominant token is exp(HEAD_LOG_GAP)x larger than the unit-count tail, so the
# tail sits ~HEAD_LOG_GAP logits below the top.
HEAD_LOG_GAP = 18.0
# 10-sigma band: a correct sampler effectively never trips it.
Z_TOLERANCE = 10.0
def _make_heavy_tailed_counts(seed: int = 1234) -> torch.Tensor:
"""Non-negative int64 counts of shape [VOCAB_SIZE]; target prob = counts/N."""
gen = torch.Generator(device=DEVICE).manual_seed(seed)
counts = torch.randint(
1, 4, (VOCAB_SIZE,), generator=gen, dtype=torch.int64, device=DEVICE
)
counts[0] = round(math.exp(HEAD_LOG_GAP)) # dominant token
return counts
def _counts_to_logits(counts: torch.Tensor) -> torch.Tensor:
# softmax(log(count)) == count / sum(count); count 0 -> logit -inf -> prob 0.
return counts.double().log().to(torch.float32)
def _sample(
logits_1d: torch.Tensor,
num_samples: int,
*,
use_fp64: bool = False,
temperature: float = 1.0,
) -> torch.Tensor:
"""Sample `num_samples` tokens from one logit vector.
Fixed seed with a distinct `pos` per sample gives independent draws; the
logits are broadcast with a 0-stride view to avoid materializing
[num_samples, vocab_size].
"""
vocab_size = logits_1d.shape[0]
logits = logits_1d.unsqueeze(0).expand(num_samples, vocab_size)
idx_mapping = torch.zeros(num_samples, dtype=torch.int32, device=DEVICE)
temp = torch.tensor([temperature], dtype=torch.float32, device=DEVICE)
seed = torch.tensor([0xABCD], dtype=torch.int64, device=DEVICE)
pos = torch.arange(num_samples, dtype=torch.int64, device=DEVICE)
return gumbel_sample(
logits,
idx_mapping,
temp,
seed,
pos,
apply_temperature=True,
use_fp64=use_fp64,
)
def _z_score(observed: int, expected: float, num_trials: int) -> float:
p = expected / num_trials
return (observed - expected) / math.sqrt(num_trials * p * (1 - p))
def _sample_histogram(
logits_1d: torch.Tensor, num_samples: int, *, chunk: int = 1_000_000
) -> torch.Tensor:
"""Histogram of `num_samples` draws, accumulated in chunks.
Chunking keeps the kernel's per-sample scratch ([chunk, num_blocks]) bounded
so a large sample count does not blow up memory.
"""
vocab_size = logits_1d.shape[0]
hist = torch.zeros(vocab_size, dtype=torch.float64, device=DEVICE)
for start in range(0, num_samples, chunk):
size = min(chunk, num_samples - start)
logits = logits_1d.unsqueeze(0).expand(size, vocab_size)
idx_mapping = torch.zeros(size, dtype=torch.int32, device=DEVICE)
temp = torch.tensor([1.0], dtype=torch.float32, device=DEVICE)
seed = torch.tensor([0xABCD], dtype=torch.int64, device=DEVICE)
pos = torch.arange(start, start + size, dtype=torch.int64, device=DEVICE)
out = gumbel_sample(
logits, idx_mapping, temp, seed, pos, apply_temperature=True
)
hist += torch.bincount(out, minlength=vocab_size).double()
return hist
# ----------------------------- Accuracy ------------------------------------
@pytest.mark.parametrize("use_fp64", [False, True])
def test_sampling_matches_target_distribution(use_fp64: bool):
counts = _make_heavy_tailed_counts()
total = counts.sum().item()
logits = _counts_to_logits(counts)
sampled = _sample(logits, NUM_SAMPLES, use_fp64=use_fp64)
assert sampled.min() >= 0 and sampled.max() < VOCAB_SIZE
# The dominant token (index 0) and the aggregate tail are the two
# statistically resolvable bins (individual tail tokens are far below the
# ~5/N detectability floor). The tail mass is small but well above noise,
# and it lives beyond the fp32 Gumbel cap -- the regime sensitive to noise
# precision -- so matching it is the meaningful check.
tail_prob = (total - counts[0].item()) / total
tail_count = (sampled != 0).sum().item()
z = _z_score(tail_count, NUM_SAMPLES * tail_prob, NUM_SAMPLES)
assert abs(z) < Z_TOLERANCE, (
f"sampled tail mass {tail_count / NUM_SAMPLES:.3e} != target "
f"{tail_prob:.3e} (z={z:.2f})"
)
def test_full_vocab_distribution_fidelity():
"""The sampled distribution matches the target across the WHOLE vocab.
A near-flat count tensor makes every one of the 200K bins individually
measurable. With ~20 samples/bin, a goodness-of-fit over all bins checks
that no part of the vocab is over- or under-represented (the heavy-tailed
test above only resolves head vs aggregate tail). Empirically the fp32
sampler is as faithful here as torch.multinomial; the residual error is the
multinomial sampling-noise floor, not the kernel.
"""
gen = torch.Generator(device=DEVICE).manual_seed(2024)
counts = torch.randint(
500, 1500, (VOCAB_SIZE,), generator=gen, dtype=torch.int64, device=DEVICE
)
total = counts.sum().item()
logits = _counts_to_logits(counts)
num_samples = 4_000_000
hist = _sample_histogram(logits, num_samples)
# Diversity: essentially every token must be reachable (no starved region).
coverage = (hist > 0).sum().item() / VOCAB_SIZE
assert coverage > 0.99, f"only {coverage:.4f} of the vocab was ever sampled"
# Goodness-of-fit across all bins (each has expected count >= ~10).
expected = (counts.double() / total) * num_samples
chi2 = (((hist - expected) ** 2) / expected).sum().item()
df = VOCAB_SIZE - 1
assert chi2 < df + 10 * math.sqrt(2 * df), f"chi2={chi2:.0f}, df={df}"
# ----------------------------- Edge cases ----------------------------------
def test_greedy_temperature_zero_returns_argmax():
"""temperature == 0 skips Gumbel noise and returns the exact argmax."""
torch.manual_seed(0)
num_reqs = 128
logits = torch.randn(num_reqs, VOCAB_SIZE, device=DEVICE, dtype=torch.float32)
idx_mapping = torch.arange(num_reqs, dtype=torch.int32, device=DEVICE)
temp = torch.zeros(num_reqs, dtype=torch.float32, device=DEVICE)
seed = torch.arange(num_reqs, dtype=torch.int64, device=DEVICE)
pos = torch.arange(num_reqs, dtype=torch.int64, device=DEVICE)
sampled = gumbel_sample(
logits, idx_mapping, temp, seed, pos, apply_temperature=True
)
assert torch.equal(sampled, logits.argmax(dim=-1))
def test_zero_count_tokens_are_never_sampled():
"""Count 0 -> -inf logit -> probability 0; must never be selected."""
counts = _make_heavy_tailed_counts(seed=7)
zeroed = torch.arange(1, VOCAB_SIZE, 2, device=DEVICE) # odd indices (not head)
counts[zeroed] = 0
logits = _counts_to_logits(counts)
sampled = _sample(logits, NUM_SAMPLES)
assert sampled.min() >= 0 and sampled.max() < VOCAB_SIZE
assert not torch.isin(sampled, zeroed).any(), "sampled a zero-probability token"
def test_single_nonzero_token_is_always_sampled():
"""A lone finite logit must win every draw, regardless of its index."""
counts = torch.zeros(VOCAB_SIZE, dtype=torch.int64, device=DEVICE)
counts[123_456] = 1000
logits = _counts_to_logits(counts)
sampled = _sample(logits, 10_000)
assert (sampled == 123_456).all()
@pytest.mark.parametrize("vocab_size", [1, 999, 1024, 4097])
def test_vocab_size_not_multiple_of_block(vocab_size: int):
"""Per-block tail masking for non-block-aligned vocab; all bins measurable."""
gen = torch.Generator(device=DEVICE).manual_seed(vocab_size)
counts = torch.randint(
20, 200, (vocab_size,), generator=gen, dtype=torch.int64, device=DEVICE
)
total = counts.sum().item()
logits = _counts_to_logits(counts)
num_samples = max(40 * vocab_size, 50_000)
sampled = _sample(logits, num_samples)
assert sampled.min() >= 0 and sampled.max() < vocab_size
observed = torch.bincount(sampled, minlength=vocab_size).double()
expected = (counts.double() / total) * num_samples
chi2 = (((observed - expected) ** 2) / expected).sum().item()
df = vocab_size - 1
if df >= 1:
assert chi2 < df + 10 * math.sqrt(2 * df), f"chi2={chi2:.1f}, df={df}"
+21 -13
View File
@@ -1019,23 +1019,26 @@ def _rocm_aiter_fused_allreduce_rmsnorm_quant_per_group_with_bf16_norm_fake(
def _rocm_aiter_per_tensor_quant_impl(
out: torch.Tensor,
x: torch.Tensor,
quant_dtype: torch.dtype,
scale: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
from aiter.ops.quant import per_tensor_quant_hip
scale: torch.Tensor,
is_dynamic: bool,
) -> None:
from aiter.ops.quant import dynamic_per_tensor_quant, static_per_tensor_quant
return per_tensor_quant_hip(x, scale, quant_dtype)
if is_dynamic:
dynamic_per_tensor_quant(out, x, scale)
else:
static_per_tensor_quant(out, x, scale)
def _rocm_aiter_per_tensor_quant_fake(
out: torch.Tensor,
x: torch.Tensor,
quant_dtype: torch.dtype,
scale: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
return torch.empty_like(x, dtype=quant_dtype), torch.empty(
1, dtype=torch.float32, device=x.device
)
scale: torch.Tensor,
is_dynamic: bool,
) -> None:
pass
def _rocm_aiter_per_token_quant_impl(
@@ -1979,7 +1982,7 @@ class rocm_aiter_ops:
direct_register_custom_op(
op_name="rocm_aiter_per_tensor_quant",
op_func=_rocm_aiter_per_tensor_quant_impl,
mutates_args=[],
mutates_args=["out", "scale"],
fake_impl=_rocm_aiter_per_tensor_quant_fake,
dispatch_key=current_platform.dispatch_key,
)
@@ -2392,7 +2395,12 @@ class rocm_aiter_ops:
quant_dtype: torch.dtype,
scale: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
return torch.ops.vllm.rocm_aiter_per_tensor_quant(x, quant_dtype, scale)
out = torch.empty_like(x, dtype=quant_dtype)
is_dynamic = scale is None
if is_dynamic:
scale = torch.empty(1, dtype=torch.float32, device=x.device)
torch.ops.vllm.rocm_aiter_per_tensor_quant(out, x, scale, is_dynamic)
return out, scale
@staticmethod
def per_token_quant(
-3
View File
@@ -555,9 +555,6 @@ class VllmConfig:
if model_config.runner_type != "generate":
return False
if model_config.is_quantized:
return False
architectures = getattr(model_config, "architectures", [])
return any(
arch in DEFAULT_V2_MODEL_RUNNER_ARCHITECTURES for arch in architectures
@@ -176,7 +176,7 @@ class MooncakeStoreConnector(KVConnectorBase_V1, SupportsHMA):
self,
request: Request,
num_computed_tokens: int,
) -> tuple[int, bool]:
) -> tuple[int | None, bool]:
assert self.connector_scheduler is not None
return self.connector_scheduler.get_num_new_matched_tokens(
request, num_computed_tokens
@@ -54,9 +54,9 @@ class MooncakeStoreScheduler:
):
assert vllm_config.kv_transfer_config is not None
self.kv_role = vllm_config.kv_transfer_config.kv_role
self.load_async = vllm_config.kv_transfer_config.kv_connector_extra_config.get(
"load_async", True
)
kvc_extra_config = vllm_config.kv_transfer_config.kv_connector_extra_config
self.load_async = kvc_extra_config.get("load_async", True)
self.lookup_async = kvc_extra_config.get("lookup_async", False)
self.client = LookupKeyClient(vllm_config)
# Align with the engine's own scheduler_block_size and hash_block_size.
@@ -75,14 +75,26 @@ class MooncakeStoreScheduler:
self,
request: Request,
num_computed_tokens: int,
) -> tuple[int, bool]:
"""Check for external KV cache hit."""
) -> tuple[int | None, bool]:
"""Check for external KV cache hit.
Returns ``(None, False)`` when an async lookup is still in flight,
signaling the scheduler to retry this request on a later step.
"""
# Look up against the full prefill range, not just the prompt.
token_len = request.num_tokens // self._block_size * self._block_size
if token_len < self._block_size:
return 0, False
num_external_hit_tokens = self.client.lookup(token_len, request.block_hashes)
num_external_hit_tokens = self.client.lookup(
request.request_id,
token_len,
request.block_hashes,
non_block=self.lookup_async,
)
if num_external_hit_tokens is None:
# Lookup not ready yet; scheduler will retry on a later step.
return None, False
if num_external_hit_tokens == request.num_tokens:
# Leave a sub-block tail uncomputed for sampling, on a block
@@ -158,6 +170,7 @@ class MooncakeStoreScheduler:
force_skip_save = self.kv_role == "kv_consumer"
for finished_req_id in scheduler_output.finished_req_ids:
self.client.discard(finished_req_id)
self.load_specs.pop(finished_req_id, None)
self._request_trackers.pop(finished_req_id, None)
self._unfinished_requests.pop(finished_req_id, None)
@@ -19,6 +19,7 @@ import threading
import time
from collections import defaultdict
from collections.abc import Callable
from concurrent.futures import Future, ThreadPoolExecutor
from dataclasses import dataclass
from typing import Any, Literal, TypeVar
@@ -971,7 +972,13 @@ class MooncakeStoreWorker:
else:
self.num_kv_head = model_config.get_total_num_kv_heads()
if self.num_kv_head < self.tp_size:
if self.num_kv_head < self.tp_size and self.dcp_size <= 1:
# Dedup: TP ranks holding the same KV heads stripe PUTs across
# one shared key namespace. DCP splits the TP group, so with
# DCP>1 those ranks have different `@dcpN` namespaces and
# striping would leave keys unwritten (OBJECT_NOT_FOUND on
# GET). PCP is outer to TP (pcp_rank is constant within a TP
# group), so it needs no guard.
self.put_step = self.tp_size // self.num_kv_head
self.head_or_tp_rank = self.tp_rank // self.put_step
else:
@@ -1560,7 +1567,13 @@ class LookupKeyClient:
bind=False,
)
def lookup(self, token_len: int, block_hashes: list[BlockHash]) -> int:
# Async lookup support
self.executor = ThreadPoolExecutor(
max_workers=1, thread_name_prefix="MooncakeLookupClient"
)
self.futures: dict[str, Future[int]] = {}
def _lookup(self, token_len: int, block_hashes: list[BlockHash]) -> int:
hash_strs = [h.hex() for h in block_hashes]
hash_frames = self.encoder.encode(hash_strs)
token_len_bytes = token_len.to_bytes(4, byteorder="big")
@@ -1570,7 +1583,36 @@ class LookupKeyClient:
result = int.from_bytes(resp, "big")
return result
def reset(self) -> bool:
def lookup(
self,
req_id: str,
token_len: int,
block_hashes: list[BlockHash],
non_block: bool = False,
) -> int | None:
"""If non_block is True, will return None until the result is ready,
so the caller retries on a later step."""
future = self.futures.get(req_id)
if future is None:
future = self.executor.submit(self._lookup, token_len, list(block_hashes))
self.futures[req_id] = future
if non_block and not future.done():
return None
try:
return future.result()
except Exception as e:
logger.error("Async Mooncake lookup failed for %s: %s", req_id, e)
return 0
finally:
del self.futures[req_id]
def discard(self, req_id: str) -> None:
"""Drop any cached/in-flight lookup for ``req_id`` (e.g. on abort)."""
future = self.futures.pop(req_id, None)
if future is not None:
future.cancel()
def _reset(self) -> bool:
"""Trigger ``store.remove_all(force=True)`` on worker rank 0.
Ordering assumption: caller MUST ensure no in-flight Mooncake
@@ -1582,7 +1624,11 @@ class LookupKeyClient:
resp = self.socket.recv()
return bytes(resp) == RESP_OK
def reset(self) -> bool:
return self.executor.submit(self._reset).result()
def close(self):
self.executor.shutdown(wait=False, cancel_futures=True)
self.socket.close(linger=0)
+125 -15
View File
@@ -12,6 +12,7 @@ import uuid
from collections.abc import AsyncGenerator
from typing import TYPE_CHECKING, Any
import jinja2
from fastapi import Request
from vllm.engine.protocol import EngineClient
@@ -42,6 +43,7 @@ from vllm.entrypoints.openai.engine.protocol import (
JsonSchemaResponseFormat,
ResponseFormat,
StreamOptions,
UsageInfo,
)
from vllm.entrypoints.openai.models.serving import OpenAIServingModels
from vllm.entrypoints.serve.utils.api_utils import sanitize_message
@@ -53,6 +55,45 @@ if TYPE_CHECKING:
logger = logging.getLogger(__name__)
def _get_cached_tokens(usage: UsageInfo | None) -> int | None:
"""Extract cached token count from OpenAI UsageInfo."""
if usage is None or usage.prompt_tokens_details is None:
return None
return usage.prompt_tokens_details.cached_tokens
def _build_anthropic_usage(
prompt_tokens: int,
completion_tokens: int,
usage: UsageInfo | None,
) -> AnthropicUsage:
"""Build an AnthropicUsage from OpenAI-style token counts.
Anthropic defines ``total_input == input_tokens + cache_read +
cache_creation``. vLLM's ``prompt_tokens`` is the total, so
``input_tokens = prompt_tokens - cached_tokens``.
OpenAI usage only exposes ``cached_tokens`` (hits); there is no
cache-creation analog, so ``cache_creation_input_tokens`` is ``0``
when cache info is present. When cache info is absent (e.g.
``--enable-prompt-tokens-details`` off, or a streaming chunk that
hasn't carried it yet), cache fields are left **unset** so
``exclude_unset=True`` serialization omits them entirely.
"""
cached = _get_cached_tokens(usage)
if cached is not None:
return AnthropicUsage(
input_tokens=prompt_tokens - cached,
output_tokens=completion_tokens,
cache_read_input_tokens=cached,
cache_creation_input_tokens=0,
)
return AnthropicUsage(
input_tokens=prompt_tokens,
output_tokens=completion_tokens,
)
def wrap_data_with_event(data: str, event: str):
return f"event: {event}\ndata: {data}\n\n"
@@ -99,6 +140,36 @@ class AnthropicServingMessages(OpenAIServingChat):
"length": "max_tokens",
"tool_calls": "tool_use",
}
self._merge_inline_system = self._detect_merge_inline_system(chat_template)
@staticmethod
def _detect_merge_inline_system(chat_template: str | None) -> bool:
"""Auto-detect whether the chat template requires system-first ordering.
Renders a [system, user, system, user] conversation against the
template; if it raises (e.g. Qwen's ``loop.first`` guard), the
model needs inline system messages merged into the leading block.
"""
if not chat_template:
return True
try:
env = jinja2.sandbox.ImmutableSandboxedEnvironment(
trim_blocks=True,
lstrip_blocks=True,
extensions=[jinja2.ext.loopcontrols],
)
env.from_string(chat_template).render(
messages=[
{"role": "system", "content": "t"},
{"role": "user", "content": "t"},
{"role": "system", "content": "t"},
{"role": "user", "content": "t"},
],
add_generation_prompt=False,
)
return False
except jinja2.TemplateError:
return True
@staticmethod
def _convert_image_source_to_url(source: dict[str, Any]) -> str:
@@ -123,13 +194,24 @@ class AnthropicServingMessages(OpenAIServingChat):
@classmethod
def _convert_anthropic_to_openai_request(
cls, anthropic_request: AnthropicMessagesRequest | AnthropicCountTokensRequest
cls,
anthropic_request: AnthropicMessagesRequest | AnthropicCountTokensRequest,
*,
merge_inline_system: bool = False,
) -> ChatCompletionRequest:
"""Convert Anthropic message format to OpenAI format"""
openai_messages: list[dict[str, Any]] = []
cls._convert_system_message(anthropic_request, openai_messages)
cls._convert_messages(anthropic_request.messages, openai_messages)
cls._convert_system_message(
anthropic_request,
openai_messages,
merge_inline_system=merge_inline_system,
)
cls._convert_messages(
anthropic_request.messages,
openai_messages,
merge_inline_system=merge_inline_system,
)
req = cls._build_base_request(anthropic_request, openai_messages)
cls._handle_streaming_options(req, anthropic_request)
cls._handle_output_config(req, anthropic_request)
@@ -142,6 +224,8 @@ class AnthropicServingMessages(OpenAIServingChat):
cls,
anthropic_request: AnthropicMessagesRequest | AnthropicCountTokensRequest,
openai_messages: list[dict[str, Any]],
*,
merge_inline_system: bool = False,
) -> None:
"""Convert Anthropic system message to OpenAI format"""
system_parts: list[str] = []
@@ -159,6 +243,17 @@ class AnthropicServingMessages(OpenAIServingChat):
continue
system_parts.append(block.text)
# When the template requires system-first ordering, extract inline
# system messages from the messages array and merge them into the
# top-level block so the template doesn't reject them.
if merge_inline_system:
for msg in anthropic_request.messages:
if msg.role != "system":
continue
text = cls._extract_system_text(msg)
if text:
system_parts.append(text)
if system_parts:
openai_messages.append({"role": "system", "content": "".join(system_parts)})
@@ -180,7 +275,11 @@ class AnthropicServingMessages(OpenAIServingChat):
@classmethod
def _convert_messages(
cls, messages: list, openai_messages: list[dict[str, Any]]
cls,
messages: list,
openai_messages: list[dict[str, Any]],
*,
merge_inline_system: bool = False,
) -> None:
"""Convert Anthropic messages to OpenAI format"""
for msg in messages:
@@ -190,6 +289,8 @@ class AnthropicServingMessages(OpenAIServingChat):
# doesn't strip billing headers and may produce messages with
# no "content" key.
if msg.role == "system":
if merge_inline_system:
continue # already merged into top-level by _convert_system_message
text = cls._extract_system_text(msg)
if text:
openai_messages.append({"role": "system", "content": text})
@@ -497,7 +598,10 @@ class AnthropicServingMessages(OpenAIServingChat):
"""
if logger.isEnabledFor(logging.DEBUG):
logger.debug("Received messages request %s", request.model_dump_json())
chat_req = self._convert_anthropic_to_openai_request(request)
chat_req = self._convert_anthropic_to_openai_request(
request,
merge_inline_system=self._merge_inline_system,
)
if logger.isEnabledFor(logging.DEBUG):
logger.debug("Convert to OpenAI request %s", chat_req.model_dump_json())
generator = await self.create_chat_completion(chat_req, raw_request)
@@ -518,9 +622,10 @@ class AnthropicServingMessages(OpenAIServingChat):
id=generator.id,
content=[],
model=generator.model,
usage=AnthropicUsage(
input_tokens=generator.usage.prompt_tokens,
output_tokens=generator.usage.completion_tokens,
usage=_build_anthropic_usage(
generator.usage.prompt_tokens,
generator.usage.completion_tokens,
generator.usage,
),
kv_transfer_params=generator.kv_transfer_params,
)
@@ -701,11 +806,12 @@ class AnthropicServingMessages(OpenAIServingChat):
model=origin_chunk.model,
stop_reason=None,
stop_sequence=None,
usage=AnthropicUsage(
input_tokens=origin_chunk.usage.prompt_tokens
usage=_build_anthropic_usage(
origin_chunk.usage.prompt_tokens
if origin_chunk.usage
else 0,
output_tokens=0,
0,
origin_chunk.usage,
),
),
)
@@ -724,13 +830,14 @@ class AnthropicServingMessages(OpenAIServingChat):
chunk = AnthropicStreamEvent(
type="message_delta",
delta=AnthropicDelta(stop_reason=stop_reason),
usage=AnthropicUsage(
input_tokens=origin_chunk.usage.prompt_tokens
usage=_build_anthropic_usage(
origin_chunk.usage.prompt_tokens
if origin_chunk.usage
else 0,
output_tokens=origin_chunk.usage.completion_tokens
origin_chunk.usage.completion_tokens
if origin_chunk.usage
else 0,
origin_chunk.usage,
),
)
data = chunk.model_dump_json(exclude_unset=True)
@@ -905,7 +1012,10 @@ class AnthropicServingMessages(OpenAIServingChat):
raw_request: Request | None = None,
) -> AnthropicCountTokensResponse | ErrorResponse:
"""Implements Anthropic's messages.count_tokens endpoint."""
chat_req = self._convert_anthropic_to_openai_request(request)
chat_req = self._convert_anthropic_to_openai_request(
request,
merge_inline_system=self._merge_inline_system,
)
result = await self.render_chat_request(chat_req)
if isinstance(result, ErrorResponse):
return result
@@ -427,6 +427,12 @@ async def validation_exception_handler(req: Request, exc: RequestValidationError
param = ctx_error.parameter
break
if param is None and errors:
first_error = errors[0]
loc = first_error.get("loc") if isinstance(first_error, dict) else None
if loc:
param = ".".join(str(part) for part in loc)
exc_str = str(exc)
errors_str = str(errors)
+7 -1
View File
@@ -446,11 +446,17 @@ class PunicaWrapperGPU(PunicaWrapperBase):
_,
_,
lora_ids,
_,
no_lora_flag,
num_active_loras,
) = self.token_mapping_meta.meta_args(
x.size(0), self.lora_config.specialize_active_lora
)
assert no_lora_flag.numel() == 1
if no_lora_flag.item():
# None of the inputs require LoRA.
return
if token_lora_mapping is None:
token_lora_mapping = token_lora_mapping_meta
fused_moe_lora(
@@ -1,7 +1,7 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from typing import TYPE_CHECKING, Any
from typing import TYPE_CHECKING, Any, cast
import torch
import torch.nn as nn
@@ -166,7 +166,21 @@ def _init_kv_cache_quant(
# TODO (mgoin): kv cache dtype should be specified in the FP8
# checkpoint config and become the "auto" behavior
if layer.kv_cache_dtype == "fp8_e5m2":
raise ValueError("fp8_e5m2 kv-cache is not supported with fp8 checkpoints.")
# A compressed-tensors checkpoint stores fp8 KV scales only when it
# declares a kv_cache_scheme; weight-only ones declare none and must
# keep fp8_e5m2, the only fp8 KV dtype usable on Ampere.
from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import ( # noqa: E501
CompressedTensorsConfig,
CompressedTensorsKVCacheMethod,
)
if not isinstance(quant_method, CompressedTensorsKVCacheMethod) or (
cast(CompressedTensorsConfig, quant_method.quant_config).kv_cache_scheme
is not None
):
raise ValueError(
"fp8_e5m2 kv-cache is not supported with fp8 checkpoints."
)
# If quantization is enabled, we make "k_scale" and "v_scale"
# parameters so that it can be loaded from the model checkpoint.
# The k/v_scale will then be converted back to native float32
@@ -349,6 +349,7 @@ class MLAAttention(nn.Module, AttentionLayerBase):
attn_backend: type[AttentionBackend] | None = None,
use_sparse: bool = False,
indexer: object | None = None,
topk_indices_buffer: torch.Tensor | None = None,
**extra_impl_args,
):
super().__init__()
@@ -437,6 +438,11 @@ class MLAAttention(nn.Module, AttentionLayerBase):
)
cache_config.enable_prefix_caching = False
# Sparse MLA reads top-k indices from a shared buffer. Pass it
# explicitly so backbone "skip" layers (indexer=None) still find it.
if use_sparse:
extra_impl_args["topk_indices_buffer"] = topk_indices_buffer
impl_cls = cast(type[MLAAttentionImpl], self.attn_backend.get_impl_cls())
self.impl = impl_cls( # type: ignore[assignment] # impl_cls always returns an MLAAttentionImpl subclass
num_heads=self.num_heads,
@@ -59,3 +59,10 @@ class MoELoRAContext:
# None means no dispatch happened (non-EP path), in which case callers
# fall back to punica_wrapper.token_mapping_meta.
local_token_lora_mapping: torch.Tensor | None = None
# Original unquantized hidden states, stashed by the modular kernel
# before the prepare step potentially quantizes them. Used by
# apply_w13_lora so the LoRA kernel sees correct-magnitude activations
# instead of raw quantized values that are missing the activation scale.
# Set per forward pass; None until the modular kernel writes it.
original_hidden_states: torch.Tensor | None = None
@@ -77,6 +77,16 @@ class TritonExperts(LoRAExpertsMixin, mk.FusedMoEExpertsModular):
def activation_format() -> mk.FusedMoEActivationFormat:
return mk.FusedMoEActivationFormat.Standard
@property
def expects_unquantized_inputs(self) -> bool:
# Defer activation quantization to apply() only when LoRA is active AND
# tokens are dispatched across ranks (DP+EP all2all).
return (
self._lora_context is not None
and self.quant_dtype is not None
and self.moe_config.moe_parallel_config.use_all2all_kernels
)
@staticmethod
def _supports_current_device() -> bool:
return current_platform.is_cuda_alike() or current_platform.is_xpu()
@@ -223,6 +233,25 @@ class TritonExperts(LoRAExpertsMixin, mk.FusedMoEExpertsModular):
torch.float8_e4m3fnuz,
]
# We declared expects_unquantized_inputs (LoRA + DP/EP all2all), so the
# prepare step deferred activation quantization to this kernel:
# `hidden_states` arrives unquantized. Keep the unquantized tensor for
# the LoRA shrink input and quantize a copy here for the base GEMM
# (mirrors what the prepare step would have done, but after the
# all-gather so the layout matches the gathered topk_ids / token map).
lora_unquantized_hidden_states: torch.Tensor | None = None
if self.expects_unquantized_inputs:
assert a1q_scale is None
lora_unquantized_hidden_states = hidden_states
hidden_states, a1q_scale = moe_kernel_quantize_input(
hidden_states,
self.a1_scale,
self.quant_dtype,
self.per_act_token_quant,
self.block_shape,
quantization_emulation=self.quantization_emulation,
)
E, num_tokens, N, K, top_k_num = self.moe_problem_size(
hidden_states, w1, w2, topk_ids
)
@@ -280,12 +309,28 @@ class TritonExperts(LoRAExpertsMixin, mk.FusedMoEExpertsModular):
# GEMM on the default stream and the LoRA fast-path on aux_stream;
# the LoRA writes its delta into a fresh zero buffer (add_inputs=
# False) and we sum it into intermediate_cache1 after both finish.
#
# The LoRA shrink kernel needs unquantized, gathered-layout
# activations. When activation quant was deferred to this kernel
# (expects_unquantized_inputs), the input we quantized above is exactly
# that, so use it directly. Otherwise fall back to the context stash
# (e.g. weight-only quant), guarding on a row-count match so a
# DP-gathered layout never indexes a local stash out of bounds.
sorted_token_ids_lora = None
expert_ids_lora = None
num_tokens_post_padded_lora = None
token_lora_mapping = None
lora_context = self._lora_context
if lora_unquantized_hidden_states is not None:
lora_x = lora_unquantized_hidden_states
elif (
lora_context is not None
and lora_context.original_hidden_states is not None
and lora_context.original_hidden_states.shape[0] == hidden_states.shape[0]
):
lora_x = lora_context.original_hidden_states
else:
lora_x = hidden_states
def _base_w13_fn():
invoke_fused_moe_triton_kernel(
@@ -322,7 +367,7 @@ class TritonExperts(LoRAExpertsMixin, mk.FusedMoEExpertsModular):
return self.apply_w13_lora(
lora_context,
y=lora_delta_w13,
x=hidden_states,
x=lora_x,
topk_ids=topk_ids,
topk_weights=topk_weights,
expert_map=expert_map,
@@ -359,7 +404,7 @@ class TritonExperts(LoRAExpertsMixin, mk.FusedMoEExpertsModular):
) = self.apply_w13_lora(
lora_context,
y=intermediate_cache1,
x=hidden_states,
x=lora_x,
topk_ids=topk_ids,
topk_weights=topk_weights,
expert_map=expert_map,
@@ -1407,6 +1407,13 @@ class FusedMoEKernelModularImpl:
apply_router_weight_on_input,
)
# Stash the original unquantized hidden states on the LoRA context
# so apply_w13_lora sees correct-magnitude activations instead of
# the potentially quantized values produced by _prepare().
lora_ctx = getattr(self.fused_experts, "_lora_context", None)
if lora_ctx is not None:
lora_ctx.original_hidden_states = hidden_states
fused_out = self._fused_experts(
in_dtype=hidden_states.dtype,
a1q=a1q,
@@ -1424,6 +1431,9 @@ class FusedMoEKernelModularImpl:
output_alias=output,
)
if lora_ctx is not None:
lora_ctx.original_hidden_states = None
return self._finalize(
output,
fused_out,
+1
View File
@@ -112,6 +112,7 @@ class MultiHeadLatentAttentionWrapper(PluggableLayer):
kv_b_proj=self.kv_b_proj,
use_sparse=self.is_sparse,
indexer=self.indexer,
topk_indices_buffer=mla_modules.topk_indices_buffer,
)
self.prefix = prefix
@@ -25,6 +25,7 @@ from vllm.model_executor.layers.fused_moe import (
UnquantizedFusedMoEMethod,
)
from vllm.model_executor.layers.fused_moe.oracle.int_wna16 import (
WNA16MoEBackend,
convert_to_wna16_moe_kernel_format,
make_wna16_moe_kernel,
select_wna16_moe_backend,
@@ -753,13 +754,18 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase):
gptq_marlin_moe_quant_config,
)
# CPU fused_experts_cpu requires zero points even for symmetric quant
use_zp = (
not self.quant_config.is_sym
or self.wna16_moe_backend == WNA16MoEBackend.CPU
)
return gptq_marlin_moe_quant_config(
w1_scale=layer.w13_scales,
w2_scale=layer.w2_scales,
weight_bits=self.quant_config.weight_bits,
group_size=self.quant_config.group_size,
w1_zp=getattr(layer, "w13_qzeros", None),
w2_zp=getattr(layer, "w2_qzeros", None),
w1_zp=getattr(layer, "w13_qzeros", None) if use_zp else None,
w2_zp=getattr(layer, "w2_qzeros", None) if use_zp else None,
w1_bias=getattr(layer, "w13_bias", None),
w2_bias=getattr(layer, "w2_bias", None),
)
@@ -415,9 +415,9 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
replace_parameter(layer, "w13_weight_scale", w13_scales)
replace_parameter(layer, "w2_weight_scale", w2_scales)
if w13_qzeros is not None:
# CPU fused_experts_cpu requires zero points even for symmetric quant
if not self.symmetric or self.wna16_backend == WNA16MoEBackend.CPU:
replace_parameter(layer, "w13_weight_zero_point", w13_qzeros)
if w2_qzeros is not None:
replace_parameter(layer, "w2_weight_zero_point", w2_qzeros)
# Marlin-specific parameters (not needed for Flashinfer)
+6 -2
View File
@@ -119,8 +119,12 @@ class DeepSeekMultiTokenPredictorLayer(nn.Module):
hidden_states=hidden_states,
residual=None,
)
hidden_states = residual + hidden_states
return hidden_states
hidden_states = residual + hidden_states # pre-final-norm (logits hidden)
# Recycle the post-final-norm hidden into the next draft step.
# compute_logits applies shared_head (== final norm) to the pre-norm
# element, so logits and the recycle each get exactly one final-norm.
# Matches SGLang's deepseek_nextn.
return hidden_states, self.shared_head(hidden_states)
class DeepSeekMultiTokenPredictor(nn.Module):
+22 -17
View File
@@ -998,8 +998,29 @@ class DeepseekV2MLAAttention(nn.Module):
self.is_v32 = hasattr(config, "index_topk")
# IndexCache config
# Refer: https://arxiv.org/abs/2603.12201 for more details.
_skip_topk = False
if self.is_v32:
_index_topk_freq = getattr(config, "index_topk_freq", 1)
_index_topk_pattern = getattr(config, "index_topk_pattern", None)
_index_skip_topk_offset = getattr(config, "index_skip_topk_offset", 2)
layer_id = extract_layer_index(prefix)
if _index_topk_pattern is None:
_skip_topk = (
max(layer_id - _index_skip_topk_offset + 1, 0) % _index_topk_freq != 0
)
elif 0 <= layer_id < len(_index_topk_pattern):
_skip_topk = _index_topk_pattern[layer_id] == "S"
# The skip pattern only governs backbone layers. MTP/nextn layers
# (layer_id >= num_hidden_layers) always build a full indexer: they
# compute indices at draft step 0 and toggle at runtime via
# set_skip_topk (index_share_for_mtp_iteration).
_num_hidden_layers = getattr(config, "num_hidden_layers", None)
is_mtp_layer = _num_hidden_layers is not None and layer_id >= _num_hidden_layers
if self.is_v32 and (not _skip_topk or is_mtp_layer):
self.indexer_rope_emb = get_rope(
qk_rope_head_dim,
max_position=max_position_embeddings,
@@ -1017,22 +1038,6 @@ class DeepseekV2MLAAttention(nn.Module):
f"{prefix}.indexer",
is_inplace_rope=self.indexer_rope_emb.enabled(),
)
# IndexCache config
# Refer: https://arxiv.org/abs/2603.12201 for more details.
_index_topk_freq = getattr(config, "index_topk_freq", 1)
_index_topk_pattern = getattr(config, "index_topk_pattern", None)
_index_skip_topk_offset = getattr(config, "index_skip_topk_offset", 2)
layer_id = extract_layer_index(prefix)
if _index_topk_pattern is None:
_skip_topk = (
max(layer_id - _index_skip_topk_offset + 1, 0) % _index_topk_freq
!= 0
)
elif 0 <= layer_id < len(_index_topk_pattern):
_skip_topk = _index_topk_pattern[layer_id] == "S"
else:
self.indexer_rope_emb = None
self.indexer = None
+44
View File
@@ -64,6 +64,7 @@ from vllm.models.deepseek_v4.nvidia.flashinfer_sparse import (
from vllm.models.deepseek_v4.nvidia.flashmla import DeepseekV4FlashMLAAttention
from vllm.models.deepseek_v4.nvidia.ops.prepare_megamoe import prepare_megamoe_inputs
from vllm.sequence import IntermediateTensors
from vllm.utils.math_utils import cdiv
from vllm.v1.attention.backends.registry import AttentionBackendEnum
@@ -85,6 +86,15 @@ class DeepseekV4MLP(nn.Module):
# across the ranks within the tp_group. In this case the weights are
# replicated and no collective ops are needed.
# Otherwise we use standard TP with an allreduce at the end.
#
# Block-FP8 shards in whole 128-blocks; cdiv rounds the per-rank block
# count up so the linear's even TP split stays block-aligned, with the
# trailing ranks zero-filled by load_weights.
block_size = getattr(quant_config, "weight_block_size", None)
if block_size is not None and not is_sequence_parallel:
tp_size = get_tensor_model_parallel_world_size()
n_local = cdiv(intermediate_size // block_size[0], tp_size)
intermediate_size = n_local * block_size[0] * tp_size
self.gate_up_proj = MergedColumnParallelLinear(
hidden_size,
[intermediate_size] * 2,
@@ -892,6 +902,8 @@ class DeepseekV4Model(nn.Module):
config = vllm_config.model_config.hf_config
quant_config = vllm_config.quant_config
self.config = config
self.quant_config = quant_config
self.parallel_config = vllm_config.parallel_config
self.use_mega_moe = (
vllm_config.kernel_config.moe_backend == "deep_gemm_mega_moe"
)
@@ -1080,7 +1092,17 @@ class DeepseekV4Model(nn.Module):
# Pre-compute expert mapping ONCE.
expert_mapping = self.get_expert_mapping()
# Block-FP8 shared experts: pad the intermediate up to the TP-uniform
# block count so the standard loaders below slice it evenly (trailing
# ranks land on the zero pad). SP / unquantized ones need no padding.
pad_shared_expert = (
getattr(self.quant_config, "weight_block_size", None) is not None
and not self.parallel_config.use_sequence_parallel_moe
)
for name, loaded_weight in weights:
if pad_shared_expert and ".shared_experts." in name:
loaded_weight = self._pad_shared_expert_weight(name, loaded_weight)
for param_name, weight_name, shard_id in stacked_params_mapping:
# Skip non-stacked layers and experts (experts handled below).
if ".experts." in name:
@@ -1155,6 +1177,28 @@ class DeepseekV4Model(nn.Module):
return loaded_params
def _pad_shared_expert_weight(
self, name: str, loaded_weight: torch.Tensor
) -> torch.Tensor:
"""Zero-pad a block-FP8 shared-expert weight/scale on its intermediate
axis so the standard TP loaders split it into even, block-aligned shards
(trailing ranks get the zero pad). gate (w1)/up (w3) [I, H] pad dim 0;
down (w2 -> down_proj) [H, I] pads dim 1.
"""
block_size = getattr(self.quant_config, "weight_block_size", None)
assert block_size is not None
# Round the intermediate axis up to a whole number of TP shards. The axis
# is in elements for weights (step = block) and in blocks for scales.
step = 1 if name.endswith("weight_scale_inv") else block_size[0]
dim = 1 if ".down_proj." in name else 0
mult = get_tensor_model_parallel_world_size() * step
pad = cdiv(loaded_weight.shape[dim], mult) * mult - loaded_weight.shape[dim]
if pad == 0:
return loaded_weight
pad_shape = list(loaded_weight.shape)
pad_shape[dim] = pad
return torch.cat([loaded_weight, loaded_weight.new_zeros(pad_shape)], dim=dim)
def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
first_layer = next(iter(islice(self.layers, self.start_layer, self.end_layer)))
if first_layer.ffn.use_mega_moe:
+1 -1
View File
@@ -672,7 +672,7 @@ class ParserEngine(Parser):
if len(tool_call_deltas) > 1:
tool_call_deltas = self._coalesce_tool_call_deltas(tool_call_deltas)
if self._deferred_content and not seen_tool_event:
if self._deferred_content and (not seen_tool_event or not tool_call_deltas):
content_parts.insert(0, self._deferred_content)
self._deferred_content = ""
@@ -9,6 +9,7 @@ names so that :class:`ReasoningParserManager` and
from vllm.parser.engine.adapters import make_adapters
from vllm.parser.gemma4 import Gemma4Parser
from vllm.parser.glm47_moe import Glm47MoeParser
from vllm.parser.minimax_m2 import MinimaxM2Parser
from vllm.parser.nemotron_v3 import NemotronV3Parser
from vllm.parser.qwen3 import Qwen3Parser
@@ -32,3 +33,8 @@ from vllm.parser.qwen3 import Qwen3Parser
Qwen3ParserReasoningAdapter,
Qwen3ParserToolAdapter,
) = make_adapters(Qwen3Parser)
(
Glm47MoeParserReasoningAdapter,
Glm47MoeParserToolAdapter,
) = make_adapters(Glm47MoeParser)
+4
View File
@@ -375,6 +375,10 @@ def gemma4_config() -> ParserEngineConfig:
ParserState.TOOL_PREAMBLE,
(EventType.REASONING_END, EventType.TOOL_CALL_START),
),
(ParserState.TOOL_PREAMBLE, "TOOL_END"): Transition(
ParserState.CONTENT,
(EventType.TOOL_CALL_END,),
),
(ParserState.TOOL_PREAMBLE, "CALL_PREFIX"): Transition(
ParserState.TOOL_NAME,
(),
+226
View File
@@ -0,0 +1,226 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""GLM-4.7 parser for reasoning and tool calls.
GLM-4.7 uses XML-like tool calls::
<tool_call>func_name<arg_key>key</arg_key><arg_value>value</arg_value></tool_call>
The function name can be followed directly by the first ``<arg_key>`` tag,
and tool calls may have no arguments.
"""
from __future__ import annotations
import functools
import json
from typing import TYPE_CHECKING
import regex as re
from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
from vllm.parser.engine.events import EventType
from vllm.parser.engine.parser_engine import ParserEngine
from vllm.parser.engine.parser_engine_config import (
ParserEngineConfig,
ParserState,
Transition,
)
if TYPE_CHECKING:
from vllm.tokenizers import TokenizerLike
from vllm.tool_parsers.abstract_tool_parser import Tool
THINK_START = "<think>"
THINK_END = "</think>"
TOOL_CALL_START = "<tool_call>"
TOOL_CALL_END = "</tool_call>"
ARG_KEY_START = "<arg_key>"
ARG_KEY_END = "</arg_key>"
ARG_VALUE_START = "<arg_value>"
ARG_VALUE_END = "</arg_value>"
_ARG_RE = re.compile(
r"<arg_key>(?P<key>.*?)</arg_key>\s*"
r"<arg_value>(?P<value>.*?)</arg_value>",
re.DOTALL,
)
_PARTIAL_ARG_RE = re.compile(
r"<arg_key>(?P<key>.*?)</arg_key>\s*"
r"<arg_value>(?P<value>.*)$",
re.DOTALL,
)
def _glm47_arg_converter(raw_args: str, partial: bool) -> str:
params: dict[str, object] = {}
for match in _ARG_RE.finditer(raw_args):
params[match.group("key").strip()] = match.group("value")
if partial:
remaining = _ARG_RE.sub("", raw_args)
match = _PARTIAL_ARG_RE.search(remaining)
if match:
key = match.group("key").strip()
if key:
params[key] = match.group("value")
return json.dumps(params, ensure_ascii=False)
@functools.cache
def glm47_moe_config(thinking: bool = True) -> ParserEngineConfig:
arg_tag_transitions = {
(ParserState.TOOL_ARGS, terminal): Transition(
ParserState.TOOL_ARGS,
(EventType.ARG_VALUE_CHUNK,),
)
for terminal in (
"ARG_KEY_START",
"ARG_KEY_END",
"ARG_VALUE_START",
"ARG_VALUE_END",
)
}
reasoning_terminals = (
{
"THINK_START": THINK_START,
"THINK_END": THINK_END,
}
if thinking
else {}
)
reasoning_token_id_terminals = (
{
"THINK_START": THINK_START,
"THINK_END": THINK_END,
}
if thinking
else {}
)
reasoning_transitions = (
{
(ParserState.CONTENT, "THINK_START"): Transition(
ParserState.REASONING,
(EventType.REASONING_START,),
),
(ParserState.REASONING, "THINK_END"): Transition(
ParserState.CONTENT,
(EventType.REASONING_END,),
),
(ParserState.CONTENT, "THINK_END"): Transition(
ParserState.CONTENT,
(),
),
}
if thinking
else {}
)
return ParserEngineConfig(
name="glm47_moe",
initial_state=ParserState.REASONING if thinking else ParserState.CONTENT,
terminals={
**reasoning_terminals,
"TOOL_START": TOOL_CALL_START,
"TOOL_END": TOOL_CALL_END,
"ARG_KEY_START": ARG_KEY_START,
"ARG_KEY_END": ARG_KEY_END,
"ARG_VALUE_START": ARG_VALUE_START,
"ARG_VALUE_END": ARG_VALUE_END,
},
token_id_terminals={
**reasoning_token_id_terminals,
"TOOL_START": TOOL_CALL_START,
"TOOL_END": TOOL_CALL_END,
},
transitions={
**reasoning_transitions,
(ParserState.REASONING, "THINK_START"): Transition(
ParserState.REASONING,
(),
),
(ParserState.REASONING, "TOOL_START"): Transition(
ParserState.TOOL_NAME,
(EventType.REASONING_END, EventType.TOOL_CALL_START),
),
(ParserState.CONTENT, "TOOL_START"): Transition(
ParserState.TOOL_NAME,
(EventType.TOOL_CALL_START,),
),
(ParserState.TOOL_NAME, "ARG_KEY_START"): Transition(
ParserState.TOOL_ARGS,
(EventType.ARG_VALUE_CHUNK,),
),
(ParserState.TOOL_NAME, "TOOL_END"): Transition(
ParserState.CONTENT,
(EventType.TOOL_CALL_END,),
),
(ParserState.TOOL_ARGS, "TOOL_END"): Transition(
ParserState.CONTENT,
(EventType.TOOL_CALL_END,),
),
**arg_tag_transitions,
},
arg_converter=_glm47_arg_converter,
stream_arg_deltas=True,
tool_args_json=False,
validate_tool_names=True,
)
class Glm47MoeParser(ParserEngine):
"""GLM-4.7 parser backed by the declarative parser engine."""
def __init__(
self,
tokenizer: TokenizerLike,
tools: list[Tool] | None = None,
**kwargs,
) -> None:
chat_kwargs = kwargs.get("chat_template_kwargs", {}) or {}
thinking = chat_kwargs.get("thinking", None)
enable_thinking = chat_kwargs.get("enable_thinking", None)
self.thinking_enabled = (
True
if thinking is None and enable_thinking is None
else bool(thinking) or bool(enable_thinking)
)
kwargs.setdefault(
"parser_engine_config",
glm47_moe_config(thinking=self.thinking_enabled),
)
super().__init__(tokenizer, tools, **kwargs)
def _emit_name_delta(self, idx: int, deltas, name: str | None) -> None:
if name is not None:
name = name.strip()
super()._emit_name_delta(idx, deltas, name)
def _handle_tool_end(self, event, deltas) -> None:
idx = event.tool_index
if 0 <= idx < len(self._tool_slots):
self._tool_slots[idx].name = self._tool_slots[idx].name.strip()
super()._handle_tool_end(event, deltas)
def is_reasoning_end(self, input_ids: list[int]) -> bool:
if not self.thinking_enabled:
return True
return super().is_reasoning_end(input_ids)
def extract_content_ids(self, input_ids: list[int]) -> list[int]:
if not self.thinking_enabled:
return input_ids
return super().extract_content_ids(input_ids)
def extract_reasoning(
self,
model_output: str,
request: ChatCompletionRequest | ResponsesRequest,
) -> tuple[str | None, str | None]:
if not self.thinking_enabled:
return None, model_output
return super().extract_reasoning(model_output, request)
+5 -1
View File
@@ -49,7 +49,7 @@ _PARAM_RE = re.compile(
r"(?:<\s*/\s*parameter\s*>|(?=<\s*parameter\s*=))",
re.DOTALL,
)
_PARTIAL_PARAM_RE = re.compile(r"<\s*parameter\s*=\s*([^>]+)>([^<]*)$", re.DOTALL)
_PARTIAL_PARAM_RE = re.compile(r"<\s*parameter\s*=\s*([^>]+)>(.*)$", re.DOTALL)
def _qwen3_arg_converter(raw_args: str, partial: bool) -> str:
@@ -125,6 +125,10 @@ def qwen3_config(thinking: bool = True) -> ParserEngineConfig:
ParserState.TOOL_NAME,
(EventType.TOOL_CALL_START,),
),
(ParserState.TOOL_PREAMBLE, "TOOL_END"): Transition(
ParserState.CONTENT,
(EventType.TOOL_CALL_END,),
),
(ParserState.TOOL_PREAMBLE, "FUNC_PREFIX"): Transition(
ParserState.TOOL_NAME,
(),
+6 -2
View File
@@ -53,8 +53,12 @@ _REASONING_PARSERS_TO_REGISTER = {
"Gemma4ParserReasoningAdapter",
),
"glm45": (
"deepseek_v3_reasoning_parser",
"DeepSeekV3ReasoningWithThinkingParser",
"glm47_moe_reasoning_parser",
"Glm47MoeParserReasoningAdapter",
),
"glm47": (
"glm47_moe_reasoning_parser",
"Glm47MoeParserReasoningAdapter",
),
"openai_gptoss": (
"gptoss_reasoning_parser",
@@ -20,7 +20,6 @@ except ImportError as e:
) from e
from vllm.entrypoints.mcp.tool_server import ToolServer
from vllm.entrypoints.openai.chat_completion.protocol import (
ChatCompletionRequest,
)
@@ -481,15 +480,6 @@ class BaseCohereCommandReasoningParser(ReasoningParser):
def is_reasoning_end(self, input_ids: Sequence[int]) -> bool:
return any(tid == self.end_token_id for tid in reversed(input_ids))
def prepare_structured_tag(
self, original_tag: str | None, tool_server: ToolServer | None
) -> str | None:
# Responses API replaces ``structural_tag`` via the reasoning parser.
# Default ``ReasoningParser.prepare_structured_tag`` returns None, which
# would clear a Cohere tag produced in ``adjust_request`` and break
# ``StructuredOutputsParams`` validation. Preserve the existing tag.
return original_tag
def adjust_request(
self, request: ChatCompletionRequest | ResponsesRequest
) -> ChatCompletionRequest | ResponsesRequest:
@@ -6,7 +6,6 @@ from typing import TYPE_CHECKING
from transformers import PreTrainedTokenizerBase
from vllm.logger import init_logger
from vllm.reasoning import ReasoningParser
from vllm.reasoning.deepseek_r1_reasoning_parser import DeepSeekR1ReasoningParser
@@ -17,8 +16,6 @@ if TYPE_CHECKING:
from vllm.entrypoints.openai.engine.protocol import DeltaMessage
from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
logger = init_logger(__name__)
class DeepSeekV3ReasoningParser(ReasoningParser):
"""
@@ -7,15 +7,12 @@ from typing import TYPE_CHECKING
from transformers import PreTrainedTokenizerBase
from vllm.entrypoints.openai.engine.protocol import DeltaMessage
from vllm.logger import init_logger
from vllm.reasoning.basic_parsers import BaseThinkingReasoningParser
if TYPE_CHECKING:
from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
logger = init_logger(__name__)
class Ernie45ReasoningParser(BaseThinkingReasoningParser):
"""
@@ -0,0 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from vllm.parser.engine.registered_adapters import Glm47MoeParserReasoningAdapter
__all__ = ["Glm47MoeParserReasoningAdapter"]
@@ -8,15 +8,12 @@ import regex as re
from transformers import PreTrainedTokenizerBase
from vllm.entrypoints.openai.engine.protocol import DeltaMessage
from vllm.logger import init_logger
from vllm.reasoning import ReasoningParser
if TYPE_CHECKING:
from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
logger = init_logger(__name__)
class GraniteReasoningParser(ReasoningParser):
"""
@@ -8,15 +8,12 @@ import regex as re
from transformers import PreTrainedTokenizerBase
from vllm.entrypoints.openai.engine.protocol import DeltaMessage
from vllm.logger import init_logger
from vllm.reasoning import ReasoningParser
if TYPE_CHECKING:
from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
logger = init_logger(__name__)
class HunyuanA13BReasoningParser(ReasoningParser):
"""
@@ -7,15 +7,12 @@ from typing import TYPE_CHECKING
from transformers import PreTrainedTokenizerBase
from vllm.entrypoints.openai.engine.protocol import DeltaMessage
from vllm.logger import init_logger
from vllm.reasoning import ReasoningParser
if TYPE_CHECKING:
from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
logger = init_logger(__name__)
class IdentityReasoningParser(ReasoningParser):
"""
@@ -7,7 +7,6 @@ from typing import TYPE_CHECKING
from vllm.entrypoints.openai.engine.protocol import (
DeltaMessage,
)
from vllm.logger import init_logger
from vllm.parser.engine.registered_adapters import MinimaxM2ParserReasoningAdapter
from vllm.reasoning.abs_reasoning_parsers import ReasoningParser
from vllm.tokenizers import TokenizerLike
@@ -16,8 +15,6 @@ if TYPE_CHECKING:
from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
logger = init_logger(__name__)
class MiniMaxM2ReasoningParser(MinimaxM2ParserReasoningAdapter): # type: ignore[valid-type, misc]
"""
@@ -5,7 +5,6 @@ from collections.abc import Iterable, Sequence
from functools import cached_property
from typing import TYPE_CHECKING
from vllm.logger import init_logger
from vllm.reasoning import ReasoningParser
from vllm.reasoning.basic_parsers import BaseThinkingReasoningParser
from vllm.tokenizers.mistral import MistralTokenizer
@@ -14,8 +13,6 @@ if TYPE_CHECKING:
from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
logger = init_logger(__name__)
class MistralReasoningParser(BaseThinkingReasoningParser):
"""
-3
View File
@@ -9,7 +9,6 @@ from typing import TYPE_CHECKING
import regex as re
from vllm.entrypoints.openai.engine.protocol import DeltaMessage
from vllm.logger import init_logger
from vllm.reasoning import ReasoningParser
if TYPE_CHECKING:
@@ -17,8 +16,6 @@ if TYPE_CHECKING:
from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
from vllm.tokenizers import TokenizerLike
logger = init_logger(__name__)
class Olmo3ReasoningState(enum.Enum):
REASONING = 1
-3
View File
@@ -9,15 +9,12 @@ import regex as re
from transformers import PreTrainedTokenizerBase
from vllm.entrypoints.openai.engine.protocol import DeltaMessage
from vllm.logger import init_logger
from vllm.reasoning import ReasoningParser
if TYPE_CHECKING:
from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
logger = init_logger(__name__)
class Step3ReasoningParser(ReasoningParser):
"""
+2 -2
View File
@@ -51,8 +51,8 @@ _TOOL_PARSERS_TO_REGISTER = {
"Ernie45ToolParser",
),
"glm45": (
"glm4_moe_tool_parser",
"Glm4MoeModelToolParser",
"glm47_moe_tool_parser",
"Glm47MoeModelToolParser",
),
"glm47": (
"glm47_moe_tool_parser",
+3 -33
View File
@@ -1,41 +1,11 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
GLM-4.7 Tool Call Parser.
GLM-4.7 uses a slightly different tool call format compared to GLM-4.5:
- The function name may appear on the same line as ``<tool_call>`` without
a newline separator before the first ``<arg_key>``.
- Tool calls may have zero arguments
(e.g. ``<tool_call>func</tool_call>``).
from __future__ import annotations
This parser overrides the parent regex patterns to handle both formats.
"""
import regex as re
from vllm.logger import init_logger
from vllm.tokenizers import TokenizerLike
from vllm.tool_parsers.abstract_tool_parser import Tool
from vllm.tool_parsers.glm4_moe_tool_parser import Glm4MoeModelToolParser
logger = init_logger(__name__)
from vllm.parser.engine.registered_adapters import Glm47MoeParserToolAdapter
class Glm47MoeModelToolParser(Glm4MoeModelToolParser):
class Glm47MoeModelToolParser(Glm47MoeParserToolAdapter): # type: ignore[valid-type, misc]
supports_required_and_named = False
structural_tag_model = "glm_4_7"
def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None):
super().__init__(tokenizer, tools)
# GLM-4.7 format: <tool_call>func_name[<arg_key>...]*</tool_call>
# The function name can be followed by a newline, whitespace, or
# directly by <arg_key> tags (no separator). The arg section is
# optional so that zero-argument calls are supported.
self.func_detail_regex = re.compile(
r"<tool_call>\s*(\S+?)\s*(<arg_key>.*)?</tool_call>", re.DOTALL
)
self.func_arg_regex = re.compile(
r"<arg_key>(.*?)</arg_key>\s*<arg_value>(.*?)</arg_value>",
re.DOTALL,
)
-495
View File
@@ -1,495 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
GLM-4 Tool Call Parser with incremental string streaming support.
This parser fixes the streaming issue reported in Issue #32829 where long string
parameters (e.g., file content with 4000+ characters of code) are buffered until
complete, causing multi-second delays before the user sees any content.
The fix streams string values incrementally as they arrive, providing a true
streaming experience for long content.
"""
import json
from collections.abc import Sequence
from typing import Any
import regex as re
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 (
DeltaFunctionCall,
DeltaMessage,
DeltaToolCall,
ExtractedToolCallInformation,
FunctionCall,
ToolCall,
)
from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
from vllm.logger import init_logger
from vllm.tokenizers import TokenizerLike
from vllm.tool_parsers.abstract_tool_parser import (
Tool,
ToolParser,
)
from vllm.tool_parsers.utils import (
extract_types_from_schema,
find_tool_properties,
partial_tag_overlap,
safe_literal_eval,
)
logger = init_logger(__name__)
class Glm4MoeModelToolParser(ToolParser):
"""Tool parser for GLM-4 models with incremental string streaming.
On every streaming call the parser re-parses ``current_text`` to find
``<tool_call>`` regions, builds the JSON arguments string for each tool
call, and diffs against what was previously sent to emit only new content.
"""
supports_required_and_named = False
def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None):
super().__init__(tokenizer, tools)
# Stateful streaming fields
self.current_tool_name_sent: bool = False
self.prev_tool_call_arr: list[dict[str, Any]] = []
self.current_tool_id: int = -1
self.streamed_args_for_tool: list[str] = []
self.tool_call_start_token: str = "<tool_call>"
self.tool_call_end_token: str = "</tool_call>"
self.arg_key_start: str = "<arg_key>"
self.arg_key_end: str = "</arg_key>"
self.arg_val_start: str = "<arg_value>"
self.arg_val_end: str = "</arg_value>"
self.tool_calls_start_token = self.tool_call_start_token
self.func_call_regex = re.compile(r"<tool_call>.*?</tool_call>", re.DOTALL)
self.func_detail_regex = re.compile(
r"<tool_call>([^\n]*)\n(.*)</tool_call>", re.DOTALL
)
self.func_arg_regex = re.compile(
r"<arg_key>(.*?)</arg_key>\s*<arg_value>(.*?)</arg_value>", re.DOTALL
)
if not self.model_tokenizer:
raise ValueError(
"The model tokenizer must be passed to the ToolParser "
"constructor during construction."
)
self.tool_call_start_token_id = self.vocab.get(self.tool_call_start_token)
self.tool_call_end_token_id = self.vocab.get(self.tool_call_end_token)
# Pre-compiled pattern for finding the last <arg_key>...</arg_key>
# before a partial <arg_value> (used in _build_args_json_so_far).
self._arg_key_pattern = re.compile(
re.escape(self.arg_key_start) + r"(.*?)" + re.escape(self.arg_key_end),
re.DOTALL,
)
# Streaming state for re-parse-and-diff approach
self._sent_content_idx: int = 0
self._tool_call_ids: list[str] = []
@staticmethod
def _deserialize(value: str) -> Any:
try:
return json.loads(value)
except json.JSONDecodeError:
pass
try:
return safe_literal_eval(value)
except (ValueError, SyntaxError):
pass
return value
@staticmethod
def _json_escape_string_content(s: str) -> str:
"""JSON-escape string content for incremental streaming.
This escapes the content that goes INSIDE a JSON string (between quotes),
not including the surrounding quotes themselves.
"""
if not s:
return ""
return json.dumps(s, ensure_ascii=False)[1:-1]
def _is_string_type(self, tool_name: str, arg_name: str) -> bool:
tool_properties = find_tool_properties(self.tools, tool_name)
param_schema = tool_properties.get(arg_name)
if param_schema is None:
return False
param_types = extract_types_from_schema(param_schema)
return set(param_types) - {"null"} == {"string"}
@staticmethod
def _tools_enabled(request: ChatCompletionRequest) -> bool:
"""Return whether tool parsing should be applied for this request."""
try:
tools = getattr(request, "tools", None)
tool_choice = getattr(request, "tool_choice", None)
return bool(tools) and tool_choice != "none"
except Exception:
logger.exception("Failed to determine if tools are enabled.")
return False
def adjust_request(
self, request: ChatCompletionRequest | ResponsesRequest
) -> ChatCompletionRequest | ResponsesRequest:
"""Adjust request parameters for tool call token handling.
For required/named tool_choice, skip setting structured_outputs
because GLM models output tool calls in XML format (per chat
template). Guided decoding would force JSON output, conflicting
with the XML format and causing parsing failures.
"""
if request.tools:
tc = request.tool_choice
if tc == "required" or isinstance(tc, ChatCompletionNamedToolChoiceParam):
# Do NOT call super().adjust_request() for required/named,
# because it would set structured_outputs and force JSON
# output via guided decoding. GLM models use XML tool-call
# syntax (defined in the chat template), so guided decoding
# must be skipped to let the model output XML freely.
# The tool_parser handles extraction from XML output.
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":
# Ensure tool call tokens (<tool_call>, </tool_call>) are not skipped
# during decoding. Even though they are not marked as special tokens,
# setting skip_special_tokens=False ensures proper handling in
# transformers 5.x where decoding behavior may have changed.
request.skip_special_tokens = False
return request
def extract_tool_calls(
self,
model_output: str,
request: ChatCompletionRequest,
) -> ExtractedToolCallInformation:
matched_tool_calls = self.func_call_regex.findall(model_output)
logger.debug("model_output: %s", model_output)
try:
tool_calls: list[ToolCall] = []
for match in matched_tool_calls:
tc_detail = self.func_detail_regex.search(match)
if not tc_detail:
logger.warning(
"Failed to parse tool call details from: %s",
match,
)
continue
tc_name = tc_detail.group(1).strip()
tc_args = tc_detail.group(2)
pairs = self.func_arg_regex.findall(tc_args) if tc_args else []
arg_dct: dict[str, Any] = {}
for key, value in pairs:
arg_key = key.strip()
if self._is_string_type(tc_name, arg_key):
arg_val = value
else:
arg_val = self._deserialize(value.strip())
logger.debug("arg_key = %s, arg_val = %s", arg_key, arg_val)
arg_dct[arg_key] = arg_val
tool_calls.append(
ToolCall(
type="function",
function=FunctionCall(
name=tc_name,
arguments=json.dumps(arg_dct, ensure_ascii=False),
),
)
)
except Exception:
logger.exception("Failed to extract tool call spec")
return ExtractedToolCallInformation(
tools_called=False, tool_calls=[], content=model_output
)
else:
if len(tool_calls) > 0:
content: str | None = model_output[
: model_output.find(self.tool_calls_start_token)
]
# Normalize empty/whitespace-only content to None
if not content or not content.strip():
content = None
return ExtractedToolCallInformation(
tools_called=True, tool_calls=tool_calls, content=content
)
return ExtractedToolCallInformation(
tools_called=False, tool_calls=[], content=model_output
)
def _extract_content(self, current_text: str) -> str | None:
"""Return unsent non-tool-call text, or None.
Collects all text outside ``<tool_call>...</tool_call>`` regions,
including text between consecutive tool calls. Holds back any
suffix that could be a partial ``<tool_call>`` tag.
"""
# Build the "sendable index" — the furthest point we can send
# content up to. We scan through the text collecting segments
# that are outside tool-call regions.
content_segments: list[str] = []
pos = self._sent_content_idx
while pos < len(current_text):
start = current_text.find(self.tool_call_start_token, pos)
if start == -1:
# No more tool calls — send up to (len - partial-tag overlap)
tail = current_text[pos:]
overlap = partial_tag_overlap(tail, self.tool_call_start_token)
sendable = tail[: len(tail) - overlap] if overlap else tail
if sendable:
content_segments.append(sendable)
pos = len(current_text) - overlap
break
# Text before this <tool_call>
if start > pos:
content_segments.append(current_text[pos:start])
# Skip past the </tool_call> (or to end if incomplete)
end = current_text.find(self.tool_call_end_token, start)
if end != -1:
pos = end + len(self.tool_call_end_token)
else:
# Incomplete tool call — nothing more to send
pos = start
break
if content_segments:
self._sent_content_idx = pos
return "".join(content_segments)
# Even if no content, advance past completed tool-call regions
if pos > self._sent_content_idx:
self._sent_content_idx = pos
return None
def _extract_tool_call_regions(self, text: str) -> list[tuple[str, bool]]:
"""Extract ``(inner_text, is_complete)`` for each ``<tool_call>`` region."""
results: list[tuple[str, bool]] = []
pos = 0
while True:
start = text.find(self.tool_call_start_token, pos)
if start == -1:
break
inner_start = start + len(self.tool_call_start_token)
end = text.find(self.tool_call_end_token, inner_start)
if end != -1:
results.append((text[inner_start:end], True))
pos = end + len(self.tool_call_end_token)
else:
# Incomplete tool call — strip partial </tool_call> suffix
raw = text[inner_start:]
overlap = partial_tag_overlap(raw, self.tool_call_end_token)
if overlap:
raw = raw[:-overlap]
results.append((raw, False))
break
return results
def _extract_tool_name_from_region(self, inner_text: str) -> str | None:
"""Extract the tool name from the beginning of a tool-call region.
The name is everything before the first ``\\n`` or ``<arg_key>``.
Returns ``None`` if the name hasn't fully arrived yet.
"""
nl = inner_text.find("\n")
ak = inner_text.find(self.arg_key_start)
candidates = [i for i in [nl, ak] if i != -1]
if not candidates:
return None
cut = min(candidates)
name = inner_text[:cut].strip()
return name if name else None
def _build_args_json_so_far(
self,
tool_name: str,
inner_text: str,
is_complete: bool,
) -> str:
"""Build the JSON arguments string from the XML pairs seen so far.
For complete ``<arg_key>/<arg_value>`` pairs the value is fully
formatted. For the last argument whose ``<arg_value>`` has been
opened but not closed, the partial string content is included
(JSON-escaped, with an opening ``"`` but no closing ``"``).
The closing ``}`` is only appended when ``is_complete`` is True
(i.e. the ``</tool_call>`` tag has arrived).
"""
# Find all complete arg pairs
pairs = self.func_arg_regex.findall(inner_text)
parts: list[str] = []
for key, value in pairs:
key = key.strip()
key_json = json.dumps(key, ensure_ascii=False)
if self._is_string_type(tool_name, key):
# Don't strip string values — whitespace is significant
# and must match the partial-value path for diffing.
val_json = json.dumps(value, ensure_ascii=False)
else:
val_json = json.dumps(
self._deserialize(value.strip()), ensure_ascii=False
)
parts.append(f"{key_json}: {val_json}")
# Check for a partial (incomplete) arg value
# Find the last <arg_value> that isn't closed
last_val_start = inner_text.rfind(self.arg_val_start)
last_val_end = inner_text.rfind(self.arg_val_end)
has_partial_value = last_val_start != -1 and (
last_val_end == -1 or last_val_end < last_val_start
)
if has_partial_value:
# Find the key for this partial value
# Look for the last <arg_key>...</arg_key> before this <arg_value>
last_key_match = None
for m in self._arg_key_pattern.finditer(inner_text[:last_val_start]):
last_key_match = m
if last_key_match:
partial_key = last_key_match.group(1).strip()
partial_content_start = last_val_start + len(self.arg_val_start)
partial_content = inner_text[partial_content_start:]
# Hold back any partial </arg_value> suffix
overlap = partial_tag_overlap(partial_content, self.arg_val_end)
if overlap:
partial_content = partial_content[:-overlap]
key_json = json.dumps(partial_key, ensure_ascii=False)
if is_complete:
# Tool call finished but </arg_value> is missing
# (malformed output). Treat partial as complete value
# so the diff naturally closes any open quotes.
if self._is_string_type(tool_name, partial_key):
val_json = json.dumps(partial_content, ensure_ascii=False)
else:
val_json = json.dumps(
self._deserialize(partial_content.strip()),
ensure_ascii=False,
)
parts.append(f"{key_json}: {val_json}")
elif self._is_string_type(tool_name, partial_key):
escaped = self._json_escape_string_content(partial_content)
# Open quote but no close — more content may arrive
parts.append(f'{key_json}: "{escaped}')
else:
# Non-string partial: include raw content, no wrapping
parts.append(f"{key_json}: {partial_content}")
if not parts:
return "{}" if is_complete else ""
joined = "{" + ", ".join(parts)
if is_complete:
joined += "}"
return joined
def _compute_args_diff(self, index: int, args_so_far: str) -> str | None:
"""Return new argument text not yet sent for tool *index*, or None."""
if not args_so_far or len(args_so_far) <= len(
self.streamed_args_for_tool[index]
):
return None
diff = args_so_far[len(self.streamed_args_for_tool[index]) :]
self.streamed_args_for_tool[index] = args_so_far
self.prev_tool_call_arr[index]["arguments"] = args_so_far
return diff
def _ensure_tool_state_for(self, index: int) -> None:
"""Grow state arrays so that *index* is valid."""
while len(self._tool_call_ids) <= index:
self._tool_call_ids.append(
make_tool_call_id(id_type="random", func_name=None, idx=None)
)
while len(self.streamed_args_for_tool) <= index:
self.streamed_args_for_tool.append("")
while len(self.prev_tool_call_arr) <= index:
self.prev_tool_call_arr.append({})
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:
if not self._tools_enabled(request):
return DeltaMessage(content=delta_text) if delta_text else None
content = self._extract_content(current_text)
regions = self._extract_tool_call_regions(current_text)
tool_call_deltas: list[DeltaToolCall] = []
for i, (inner_text, is_complete) in enumerate(regions):
self._ensure_tool_state_for(i)
# Extract tool name
tool_name = self._extract_tool_name_from_region(inner_text)
if not tool_name:
break
# Emit tool name (once per tool call)
if "name" not in self.prev_tool_call_arr[i]:
self.prev_tool_call_arr[i]["name"] = tool_name
tool_call_deltas.append(
DeltaToolCall(
index=i,
id=self._tool_call_ids[i],
type="function",
function=DeltaFunctionCall(
name=tool_name,
arguments="",
).model_dump(exclude_none=True),
)
)
# Build args JSON so far, diff, emit
args_so_far = self._build_args_json_so_far(
tool_name, inner_text, is_complete
)
diff = self._compute_args_diff(i, args_so_far)
if diff:
tool_call_deltas.append(
DeltaToolCall(
index=i,
function=DeltaFunctionCall(arguments=diff).model_dump(
exclude_none=True
),
)
)
# Update current_tool_id for serving layer compatibility
if regions:
self.current_tool_id = len(regions) - 1
if content or tool_call_deltas:
return DeltaMessage(
content=content,
tool_calls=tool_call_deltas,
)
return None
+35 -3
View File
@@ -3,6 +3,7 @@
import ast
import json
import math
import warnings
from json import JSONDecodeError, JSONDecoder
from typing import Any, TypeAlias
@@ -145,6 +146,20 @@ def is_complete_json(input_str: str) -> bool:
return False
def _is_json_finite(obj: Any) -> bool:
"""Whether *obj* can be serialized to valid JSON.
``json.dumps(..., allow_nan=False)`` raises ``ValueError`` on any
non-finite float (``inf``/``-inf``/``nan``) anywhere in the value, so this
detects non-finite floats nested inside parsed lists/dicts too.
"""
try:
json.dumps(obj, allow_nan=False)
return True
except (ValueError, TypeError):
return False
def consume_space(i: int, s: str) -> int:
while i < len(s) and s[i].isspace():
i += 1
@@ -601,9 +616,15 @@ def coerce_to_schema_type(value: str, schema_type: str | list[str]) -> Any:
if candidate_type == "number":
try:
val = float(value)
return val if val != int(val) else int(val)
except (ValueError, TypeError):
continue
if not math.isfinite(val):
# inf/-inf/nan are not valid JSON numbers. Fall through so
# the value is preserved as a string instead of crashing
# (int(float("inf")) raises OverflowError) or emitting
# invalid JSON (json.dumps(inf) -> "Infinity").
continue
return val if val != int(val) else int(val)
if candidate_type == "boolean":
lower_val = value.lower().strip()
if lower_val in ("true", "1"):
@@ -613,14 +634,25 @@ def coerce_to_schema_type(value: str, schema_type: str | list[str]) -> Any:
continue
if candidate_type in ("object", "array"):
try:
return json.loads(value)
parsed = json.loads(value)
except (json.JSONDecodeError, ValueError, TypeError):
continue
if _is_json_finite(parsed):
return parsed
# Non-finite floats (e.g. "[1e999]" -> [inf]) cannot be
# serialized back to valid JSON; preserve the raw string.
continue
try:
return json.loads(value)
parsed = json.loads(value)
except (json.JSONDecodeError, ValueError):
return value
# Reject non-finite results (e.g. json.loads("1e999") -> inf, or nested
# inf/nan inside a parsed list/dict) which json.dumps would render as
# invalid JSON (Infinity/NaN). Preserve the raw string instead.
if not _is_json_finite(parsed):
return value
return parsed
def compute_tool_delta(
@@ -271,7 +271,7 @@ class FlashInferMLASparseImpl(SparseMLAAttentionImpl[FlashInferMLASparseMetadata
attn_type: str,
kv_sharing_target_layer_name: str | None,
# MLA Specific Arguments
topk_indice_buffer: torch.Tensor | None = None,
topk_indices_buffer: torch.Tensor | None = None,
indexer: "Indexer | None" = None,
**mla_args,
) -> None:
@@ -301,8 +301,12 @@ class FlashInferMLASparseImpl(SparseMLAAttentionImpl[FlashInferMLASparseMetadata
self.qk_nope_head_dim: int = mla_args["qk_nope_head_dim"]
self.qk_rope_head_dim: int = mla_args["qk_rope_head_dim"]
assert indexer is not None, "Indexer required for sparse MLA"
self.topk_indices_buffer: torch.Tensor | None = indexer.topk_indices_buffer
# The indexer carries the shared buffer for normal layers and tests;
# the explicitly-passed buffer covers backbone skip layers, whose
# indexer is not constructed (see deepseek_v2.py).
self.topk_indices_buffer: torch.Tensor | None = (
indexer.topk_indices_buffer if indexer is not None else topk_indices_buffer
)
self._workspace_buffer: torch.Tensor | None = None
self.bmm1_scale: float | None = None
@@ -568,8 +568,12 @@ class FlashMLASparseImpl(SparseMLAAttentionImpl[FlashMLASparseMetadata]):
self.kv_cache_dtype = kv_cache_dtype
self.kv_lora_rank: int = mla_args["kv_lora_rank"]
self.softmax_scale = scale
assert indexer is not None
self.topk_indices_buffer: torch.Tensor | None = indexer.topk_indices_buffer
# The indexer carries the shared buffer for normal layers and tests;
# the explicitly-passed buffer covers backbone skip layers, whose
# indexer is not constructed (see deepseek_v2.py).
self.topk_indices_buffer: torch.Tensor | None = (
indexer.topk_indices_buffer if indexer is not None else topk_indices_buffer
)
# Prefill BF16 kernel requires 64 on Hopper, 128 on Blackwell
self.prefill_padding = (
128 if current_platform.is_device_capability_family(100) else 64
@@ -629,7 +629,7 @@ class ROCMAiterMLASparseImpl(SparseMLAAttentionImpl[ROCMAiterMLASparseMetadata])
attn_type: str,
kv_sharing_target_layer_name: str | None,
# MLA Specific Arguments
topk_indice_buffer: torch.Tensor | None = None,
topk_indices_buffer: torch.Tensor | None = None,
indexer: "Indexer | None" = None,
**mla_args,
) -> None:
@@ -642,8 +642,12 @@ class ROCMAiterMLASparseImpl(SparseMLAAttentionImpl[ROCMAiterMLASparseMetadata])
self.kv_cache_dtype = kv_cache_dtype
self.kv_lora_rank: int = mla_args["kv_lora_rank"]
self.softmax_scale = scale
assert indexer is not None
self.topk_indices_buffer: torch.Tensor | None = indexer.topk_indices_buffer
# The indexer carries the shared buffer for normal layers and tests;
# the explicitly-passed buffer covers backbone skip layers, whose
# indexer is not constructed (see deepseek_v2.py).
self.topk_indices_buffer: torch.Tensor | None = (
indexer.topk_indices_buffer if indexer is not None else topk_indices_buffer
)
vllm_config = get_current_vllm_config()
max_tokens = vllm_config.scheduler_config.max_num_batched_tokens
@@ -184,7 +184,7 @@ class XPUMLASparseImpl(SparseMLAAttentionImpl[XPUMLASparseMetadata]):
attn_type: str,
kv_sharing_target_layer_name: str | None,
# MLA Specific Arguments
topk_indice_buffer: torch.Tensor | None = None,
topk_indices_buffer: torch.Tensor | None = None,
indexer: Optional["Indexer"] = None,
**mla_args,
) -> None:
@@ -195,8 +195,12 @@ class XPUMLASparseImpl(SparseMLAAttentionImpl[XPUMLASparseMetadata]):
self.kv_cache_dtype = kv_cache_dtype
self.kv_lora_rank: int = mla_args["kv_lora_rank"]
self.softmax_scale = scale
assert indexer is not None
self.topk_indices_buffer: torch.Tensor | None = indexer.topk_indices_buffer
# The indexer carries the shared buffer for normal layers and tests;
# the explicitly-passed buffer covers backbone skip layers, whose
# indexer is not constructed (see deepseek_v2.py).
self.topk_indices_buffer: torch.Tensor | None = (
indexer.topk_indices_buffer if indexer is not None else topk_indices_buffer
)
def _forward_bf16_kv(
self,
+3
View File
@@ -84,10 +84,13 @@ class FileMapper:
]
# Only a single full-attention group is parallelism-invariant. MLA is
# excluded: its latent KV is replicated per rank, never head-sharded.
# The V2 model runner is excluded: its KV layout is not known to be
# parallelism-invariant.
groups = kv_cache_config.kv_cache_groups
spec = groups[0].kv_cache_spec if len(groups) == 1 else None
parallel_agnostic = (
parallel_agnostic
and not vllm_config.use_v2_model_runner
and isinstance(spec, FullAttentionSpec)
and not isinstance(spec, MLAAttentionSpec)
)
+7
View File
@@ -918,6 +918,13 @@ class SpecDecodeBaseProposer:
return per_group_attn_metadata, per_layer_attn_metadata
def model_returns_tuple(self) -> bool:
if self.method == "mtp":
# DeepSeek-family MTP (deepseek_mtp.py) recycles the post-final-
# norm hidden, so its forward returns (logit_hidden,
# recycle_hidden). Other MTP families return a single tensor.
return "DeepSeekMTPModel" in (
self.draft_model_config.hf_config.architectures or []
)
return self.method not in ("mtp", "draft_model", "dflash")
def prepare_next_token_ids_cpu(
+13 -8
View File
@@ -2,18 +2,16 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import torch
from vllm.triton_utils import HAS_TRITON, tl, triton
from vllm.triton_utils import HAS_TRITON, tl, tldevice, triton
# Smallest positive normal fp32 value. Used to clamp the uniform draw so that
# `log(u)` cannot produce -inf (and thus `-log(-log(u))` stays finite).
# Smallest positive value produced by Triton's fp32 `tl.rand`. Used to clamp
# zero draws before the flipped Gumbel transform below.
#
# Triton requires globals accessed from `@triton.jit` functions to be wrapped
# in `tl.constexpr(...)`. We can only do that when Triton is actually
# available — on the CPU worker path `tl` is a placeholder whose `constexpr`
# attribute is `None`, and `tl.constexpr(...)` would crash at import time.
_FP32_TINY = (
tl.constexpr(float.fromhex("0x1p-126")) if HAS_TRITON else float.fromhex("0x1p-126")
)
_TL_RAND_MIN = tl.constexpr(4.6566127342e-10) if HAS_TRITON else 4.6566127342e-10
@triton.jit
@@ -131,10 +129,17 @@ def gumbel_block_argmax(
if USE_FP64:
u = tl_rand64(gumbel_seed, block, includes_zero=False)
gumbel_noise = -tl.log(-tl.log(u))
else:
u = tl.rand(gumbel_seed, block)
u = tl.maximum(u, _FP32_TINY)
gumbel_noise = -tl.log(-tl.log(u))
u = tl.maximum(u, _TL_RAND_MIN)
# Draw the large-noise tail (which decides the argmax winner) from u -> 0,
# where fp32 has fine resolution, instead of u -> 1, where fp32 spacing is
# ~2**-24. The naive `-log(-log(u))` puts the winning tail at u -> 1,
# hard-capping the noise at ~16.6 and coarsely quantizing it; using
# `log1p(-u)` == `log(1 - u)` keeps the tail in the well-resolved region.
# Note `1 - u` would lose precision for small u, so `log1p` is required.
gumbel_noise = -tl.log(-tldevice.log1p(-u))
# Apply gumbel noise.
logits = tl.where(mask, logits + gumbel_noise, float("-inf"))