forked from Karylab-cklius/vllm
Compare commits
43
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
5ad0151d64 | ||
|
|
8019b6ec63 | ||
|
|
eb3ee8afc3 | ||
|
|
77648ab261 | ||
|
|
fd3e0cac12 | ||
|
|
82d1ddf39e | ||
|
|
69bdd34542 | ||
|
|
ec67d7ae61 | ||
|
|
7341ff152f | ||
|
|
ecf9d83520 | ||
|
|
c9135db27c | ||
|
|
2a6c6b9429 | ||
|
|
ab66606993 | ||
|
|
9ea3a4015b | ||
|
|
560fb8b867 | ||
|
|
675cd5d228 | ||
|
|
7f616c327d | ||
|
|
c3c6d723fd | ||
|
|
41dcf49ca5 | ||
|
|
35e4dd4a69 | ||
|
|
4ce2d01453 | ||
|
|
16908e132e | ||
|
|
225936a1dd | ||
|
|
f6ba720963 | ||
|
|
b53b1c7ffe | ||
|
|
79ca54d221 | ||
|
|
09f3cd5c10 | ||
|
|
ea6078fe6a | ||
|
|
a0df04e477 | ||
|
|
e2352c2974 | ||
|
|
25faa1f4cc | ||
|
|
4583630b56 | ||
|
|
21da47dabe | ||
|
|
6c379b9e54 | ||
|
|
5099474633 | ||
|
|
058cc0a8b6 | ||
|
|
837db7605e | ||
|
|
383a950d04 | ||
|
|
ef8a54a77e | ||
|
|
c0089373bb | ||
|
|
c8f8f1951a | ||
|
|
a50380e5d2 | ||
|
|
04009ff40b |
@@ -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
@@ -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]
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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]);
|
||||
|
||||
@@ -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**!
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
@@ -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
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
@@ -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)
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
@@ -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 = [
|
||||
|
||||
@@ -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
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -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),
|
||||
|
||||
@@ -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
@@ -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(
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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),
|
||||
)
|
||||
|
||||
+2
-2
@@ -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)
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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,
|
||||
(),
|
||||
|
||||
@@ -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)
|
||||
@@ -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,
|
||||
(),
|
||||
|
||||
@@ -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):
|
||||
"""
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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):
|
||||
"""
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
@@ -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
|
||||
@@ -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,
|
||||
|
||||
@@ -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)
|
||||
)
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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"))
|
||||
|
||||
Reference in New Issue
Block a user